diff --git a/autolens/__init__.py b/autolens/__init__.py index 5b32e20ec..0d1ca96db 100644 --- a/autolens/__init__.py +++ b/autolens/__init__.py @@ -127,6 +127,7 @@ from .point.solver.shape_solver import ShapeSolver from .weak.dataset import WeakDataset from .weak.fit import FitWeak +from .weak.model.analysis import AnalysisWeak from .weak.simulator import SimulatorShearYX from . import exc diff --git a/autolens/config/visualize/plots.yaml b/autolens/config/visualize/plots.yaml index 0768fbdfd..f49a599a2 100644 --- a/autolens/config/visualize/plots.yaml +++ b/autolens/config/visualize/plots.yaml @@ -57,6 +57,11 @@ point_dataset: # Settings for plots of point source fit_point_dataset: {} # Settings for plots of fits to point source datasets (e.g. FitPointDatasetPlotter). +weak_dataset: # Settings for plots of weak lensing shear catalogues (e.g. PlotterWeak). + subplot_dataset: true # Plot subplot containing all dataset quantities (e.g. the shear field, noise-map, etc.)? + +fit_weak: {} # Settings for plots of fits to weak lensing shear catalogues (e.g. PlotterWeak). + fit_ellipse: # Settings for plots of ellipse fitting fits (e.g. FitEllipse) data : true # Plot the data of the ellipse fit? data_no_ellipse: true # Plot the data without the black data ellipses, which obscure noisy data? diff --git a/autolens/weak/model/__init__.py b/autolens/weak/model/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/autolens/weak/model/analysis.py b/autolens/weak/model/analysis.py new file mode 100644 index 000000000..874e7d0b8 --- /dev/null +++ b/autolens/weak/model/analysis.py @@ -0,0 +1,157 @@ +""" +Analysis class for fitting a ``Tracer`` model to a weak-lensing shear catalogue. + +``AnalysisWeak`` implements the ``log_likelihood_function`` called by a ``PyAutoFit`` +non-linear search at each iteration. It: + +1. Constructs a ``Tracer`` from the current model instance. +2. Calls ``FitWeak`` to compare the tracer's model shear field (evaluated at the + catalogue's galaxy positions via ``LensCalc.shear_yx_2d_via_hessian_from``) against + the observed ``WeakDataset``. +3. Returns the fit's log likelihood as the figure of merit. + +It also manages result output (``ResultWeak``) and on-the-fly visualisation +(``VisualizerWeak``). +""" +import autofit as af +import autogalaxy as ag + +from autogalaxy.analysis.analysis.analysis import Analysis as AgAnalysis + +from autolens.analysis.analysis.lens import AnalysisLens +from autolens.weak.dataset import WeakDataset +from autolens.weak.fit import FitWeak +from autolens.weak.model.result import ResultWeak +from autolens.weak.model.visualizer import VisualizerWeak + + +class AnalysisWeak(AgAnalysis, AnalysisLens): + Visualizer = VisualizerWeak + Result = ResultWeak + + def __init__( + self, + dataset: WeakDataset, + cosmology: ag.cosmo.LensingCosmology = None, + title_prefix: str = None, + use_jax: bool = False, + **kwargs, + ): + """ + Fits a lens model to a weak-lensing shear catalogue via a non-linear search. + + The `Analysis` class defines the `log_likelihood_function` which fits the model to the dataset and returns the + log likelihood value defining how well the model fitted the data. + + It handles many other tasks, such as visualization, outputting results to hard-disk and storing results in + a format that can be loaded after the model-fit is complete. + + This class is used for model-fits which fit lens mass models to `WeakDataset` shear catalogues — the + weak-lensing analogue of `AnalysisImaging` / `AnalysisPoint`. Each background galaxy in the catalogue + contributes two independent shear measurements (gamma_1 and gamma_2), which `FitWeak` compares against + the model shear field of the `Tracer`. + + `use_jax` defaults to `False` because `FitWeak` is a NumPy-only fit (its `model_shear` is cached via + `functools.cached_property` and its statistics use `np.asarray`); JAX support requires pytree + registration of `FitWeak` and an `xp`-threaded fit path, which is deliberate future work. + + Parameters + ---------- + dataset + The `WeakDataset` that is fitted by the model, containing the observed per-galaxy shear + measurements, their positions and the per-galaxy noise. + cosmology + The Cosmology assumed for this analysis. + title_prefix + A string that is added before the title of all figures output by visualization, for example to + put the name of the dataset and galaxy in the title. + """ + super().__init__(cosmology=cosmology, use_jax=use_jax, **kwargs) + + AnalysisLens.__init__(self=self, cosmology=cosmology, use_jax=use_jax) + + self.dataset = dataset + + self.title_prefix = title_prefix + + def log_likelihood_function(self, instance): + """ + Given an instance of the model, where the model parameters are set via a non-linear search, fit the model + instance to the weak-lensing shear catalogue. + + This function returns a log likelihood which is used by the non-linear search to guide the model-fit. + + For this analysis class, this function performs the following steps: + + 1) Extracts all galaxies from the model instance and sets up a `Tracer`, which includes ordering the galaxies + by redshift to set up each `Plane`. + + 2) Uses the `Tracer` to create a `FitWeak` object, which evaluates the tracer's shear field at the + catalogue's galaxy positions (via the same `LensCalc.shear_yx_2d_via_hessian_from` primitive the + `SimulatorShearYX` uses) and compares it to the observed shears. + + 3) Returns the fit's log likelihood — a Gaussian likelihood over the N x 2 independent shear components. + + Parameters + ---------- + instance + An instance of the model that is being fitted to the data by this analysis (whose parameters have been set + via a non-linear search). + + Returns + ------- + float + The log likelihood indicating how well this model instance fitted the weak-lensing data. + """ + return self.fit_from(instance=instance).log_likelihood + + def fit_from(self, instance) -> FitWeak: + """ + Given a model instance create a `FitWeak` object. + + This function is used in the `log_likelihood_function` to fit the model to the weak-lensing data and + compute the log likelihood. + + Parameters + ---------- + instance + An instance of the model that is being fitted to the data by this analysis (whose parameters have been set + via a non-linear search). + + Returns + ------- + The fit of the lens model to the weak-lensing shear catalogue. + """ + tracer = self.tracer_via_instance_from( + instance=instance, + ) + + return FitWeak( + dataset=self.dataset, + tracer=tracer, + ) + + def save_attributes(self, paths: af.DirectoryPaths): + """ + Before the non-linear search begins, this routine saves attributes of the `Analysis` object to the `files` + folder such that they can be loaded after the analysis using PyAutoFit's database and aggregator tools. + + For this analysis, it outputs the following: + + - The weak-lensing shear catalogue as a readable .json file. + + It is common for these attributes to be loaded by many of the template aggregator functions given in the + `aggregator` modules. For example, when using the database tools to perform a fit, the default behaviour is for + the dataset, settings and other attributes necessary to perform the fit to be loaded via the pickle files + output by this function. + + Parameters + ---------- + paths + The paths object which manages all paths, e.g. where the non-linear search outputs are stored, + visualization, and the pickled objects used by the aggregator output by this function. + """ + ag.output_to_json( + obj=self.dataset, + file_path=paths._files_path / "dataset.json", + ) diff --git a/autolens/weak/model/plotter.py b/autolens/weak/model/plotter.py new file mode 100644 index 000000000..485dd9cd4 --- /dev/null +++ b/autolens/weak/model/plotter.py @@ -0,0 +1,69 @@ +from autolens.analysis.plotter import Plotter +from autolens.analysis.plotter import plot_setting + +from autolens.weak.dataset import WeakDataset +from autolens.weak.fit import FitWeak +from autolens.weak.plot.weak_dataset_plots import subplot_weak_dataset +from autolens.weak.plot.fit_weak_plots import subplot_fit_weak +from autolens.weak.plot.fit_weak_plots import subplot_fit_quick as subplot_fit_quick_weak + + +class PlotterWeak(Plotter): + def dataset_weak(self, dataset: WeakDataset): + """ + Output visualization of a `WeakDataset` shear catalogue. + + Parameters + ---------- + dataset + The weak-lensing dataset which is visualized. + """ + + def should_plot(name): + return plot_setting(section=["weak_dataset"], name=name) + + output_path = str(self.image_path) + fmt = self.fmt + + if should_plot("subplot_dataset"): + subplot_weak_dataset( + dataset, + output_path=output_path, + output_format=fmt, + title_prefix=self.title_prefix, + ) + + def fit_weak(self, fit: FitWeak, quick_update: bool = False): + """ + Visualizes a `FitWeak` object. + + Parameters + ---------- + fit + The maximum log likelihood `FitWeak` of the non-linear search. + quick_update + If `True`, a lighter-weight quick-update subplot is output instead of the full fit subplot. + """ + + def should_plot(name): + return plot_setting(section=["fit", "fit_weak"], name=name) + + output_path = str(self.image_path) + fmt = self.fmt + + if quick_update: + subplot_fit_quick_weak( + fit, + output_path=output_path, + output_format=fmt, + title_prefix=self.title_prefix, + ) + return + + if should_plot("subplot_fit"): + subplot_fit_weak( + fit, + output_path=output_path, + output_format=fmt, + title_prefix=self.title_prefix, + ) diff --git a/autolens/weak/model/result.py b/autolens/weak/model/result.py new file mode 100644 index 000000000..5ffef4508 --- /dev/null +++ b/autolens/weak/model/result.py @@ -0,0 +1,13 @@ +import autoarray as aa + +from autolens.analysis.result import Result + + +class ResultWeak(Result): + @property + def grid(self): + return aa.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.1) + + @property + def max_log_likelihood_fit(self): + return self.analysis.fit_from(instance=self.instance) diff --git a/autolens/weak/model/visualizer.py b/autolens/weak/model/visualizer.py new file mode 100644 index 000000000..90a3e0e3a --- /dev/null +++ b/autolens/weak/model/visualizer.py @@ -0,0 +1,91 @@ +import autofit as af +import autogalaxy as ag + +from autolens.weak.model.plotter import PlotterWeak + + +class VisualizerWeak(af.Visualizer): + @staticmethod + def visualize_before_fit( + analysis, + paths: af.AbstractPaths, + model: af.AbstractPriorModel, + ): + """ + PyAutoFit calls this function immediately before the non-linear search begins. + + It visualizes objects which do not change throughout the model fit like the dataset. + + Parameters + ---------- + paths + The paths object which manages all paths, e.g. where the non-linear search outputs are stored, + visualization and the pickled objects used by the aggregator output by this function. + model + The model object, which includes model components representing the galaxies that are fitted to + the weak-lensing data. + """ + + plotter = PlotterWeak( + image_path=paths.image_path, title_prefix=analysis.title_prefix + ) + + plotter.dataset_weak(dataset=analysis.dataset) + + @staticmethod + def visualize( + analysis, + paths: af.DirectoryPaths, + instance: af.ModelInstance, + during_analysis: bool, + quick_update: bool = False, + ): + """ + Output images of the maximum log likelihood model inferred by the model-fit. This function is called throughout + the non-linear search at regular intervals, and therefore provides on-the-fly visualization of how well the + model-fit is going. + + The visualization performed by this function includes: + + - Images of the best-fit `Tracer`, including the convergence and potential of its mass profiles over + the extent of the shear catalogue. + + - Images of the best-fit `FitWeak`, including the data, model and residual shear fields and the + chi-squared map of its fit to the weak-lensing data. + + The images output by this function are customized using the file `config/visualize/plots.yaml`. + + Parameters + ---------- + paths + The paths object which manages all paths, e.g. where the non-linear search outputs are stored, + visualization, and the pickled objects used by the aggregator output by this function. + instance + An instance of the model that is being fitted to the data by this analysis (whose parameters have been set + via a non-linear search). + """ + fit = analysis.fit_for_visualization(instance=instance) + + plotter = PlotterWeak( + image_path=paths.image_path, title_prefix=analysis.title_prefix + ) + + plotter.fit_weak(fit=fit, quick_update=quick_update) + + if quick_update: + return + + tracer = fit.tracer + + grid = ag.Grid2D.from_extent( + extent=fit.dataset.extent_from(), shape_native=(100, 100) + ) + + plotter.tracer( + tracer=tracer, + grid=grid, + ) + plotter.galaxies( + galaxies=tracer.galaxies, + grid=grid, + ) diff --git a/test_autolens/weak/test_analysis.py b/test_autolens/weak/test_analysis.py new file mode 100644 index 000000000..4f91414ec --- /dev/null +++ b/test_autolens/weak/test_analysis.py @@ -0,0 +1,112 @@ +import os + +import numpy as np + +import autoarray as aa +import autofit as af +import autolens as al + +from autolens.weak.fit import FitWeak +from autolens.weak.model.result import ResultWeak + + +def _lens_model(einstein_radius=1.6, ell_comps=(0.0, 0.05)): + return af.Collection( + galaxies=af.Collection( + lens=al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal( + centre=(0.0, 0.0), + ell_comps=ell_comps, + einstein_radius=einstein_radius, + ), + ), + source=al.Galaxy(redshift=1.0), + ) + ) + + +def _make_dataset(noise_sigma=0.3, seed=1, einstein_radius=1.6): + grid = aa.Grid2DIrregular( + values=[(0.7, 0.5), (1.0, 1.0), (-0.3, 0.6), (-1.1, -0.8)] + ) + instance = _lens_model(einstein_radius=einstein_radius).instance_from_unit_vector( + [] + ) + tracer = al.Tracer(galaxies=[instance.galaxies.lens, instance.galaxies.source]) + simulator = al.SimulatorShearYX(noise_sigma=noise_sigma, seed=seed) + return simulator.via_tracer_from(tracer=tracer, grid=grid, name="test") + + +def test__log_likelihood__matches_fit_weak(): + dataset = _make_dataset() + + model = _lens_model() + instance = model.instance_from_unit_vector([]) + + analysis = al.AnalysisWeak(dataset=dataset) + + analysis_log_likelihood = analysis.log_likelihood_function(instance=instance) + + tracer = analysis.tracer_via_instance_from(instance=instance) + + fit = FitWeak(dataset=dataset, tracer=tracer) + + assert fit.log_likelihood == analysis_log_likelihood + + +def test__log_likelihood__noise_free_round_trip_gives_zero_chi_squared(): + """ + `FitWeak.model_shear` uses the same `LensCalc.shear_yx_2d_via_hessian_from` primitive as + `SimulatorShearYX`, so fitting a noise-free dataset with its own truth tracer must give + chi_squared == 0 through the analysis path too. + """ + dataset = _make_dataset(noise_sigma=0.0) + dataset.noise_map = aa.ArrayIrregular(values=[0.3, 0.3, 0.3, 0.3]) + + analysis = al.AnalysisWeak(dataset=dataset) + + instance = _lens_model().instance_from_unit_vector([]) + + fit = analysis.fit_from(instance=instance) + + assert fit.chi_squared == 0.0 + np.testing.assert_allclose(fit.residual_map, 0.0, atol=1e-12) + + +def test__log_likelihood__changes_when_model_wrong(): + dataset = _make_dataset(noise_sigma=0.0) + dataset.noise_map = aa.ArrayIrregular(values=[0.3, 0.3, 0.3, 0.3]) + + analysis = al.AnalysisWeak(dataset=dataset) + + truth = _lens_model(einstein_radius=1.6).instance_from_unit_vector([]) + wrong = _lens_model(einstein_radius=1.0).instance_from_unit_vector([]) + + assert analysis.log_likelihood_function( + instance=wrong + ) < analysis.log_likelihood_function(instance=truth) + + +def test__result_class_is_result_weak(): + assert al.AnalysisWeak.Result is ResultWeak + + +def test__save_attributes__dataset_json_round_trips(): + dataset = _make_dataset() + + analysis = al.AnalysisWeak(dataset=dataset) + + paths = af.DirectoryPaths() + + analysis.save_attributes(paths=paths) + + loaded = al.from_json(file_path=paths._files_path / "dataset.json") + + assert isinstance(loaded, al.WeakDataset) + assert loaded.name == "test" + np.testing.assert_allclose( + np.asarray(loaded.shear_yx), np.asarray(dataset.shear_yx) + ) + + os.remove(paths._files_path / "dataset.json")