Merge pull request #580 from PyAutoLabs/feature/weak-modeling

feat: AnalysisWeak — weak lensing modeling (weak series step 4)

Author: Jammy2211

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

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?

autolens/weak/model/__init__.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",
+        )

autolens/weak/model/analysis.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,
+            )

autolens/weak/model/plotter.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)

autolens/weak/model/result.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,
+        )