Merge pull request #110 from Jammy2211/feature/profiling_lh

Feature/profiling lh

autogalaxy/__init__.py

@@ -109,4 +109,4 @@
 
 conf.instance.register(__file__)
 
-__version__ = "2023.3.27.1"
+__version__ = "2023.7.7.2"

autogalaxy/analysis/analysis.py

@@ -1,9 +1,10 @@
 import json
 import logging
 import numpy as np
-from typing import Optional, Union
+from typing import Callable, Dict, Optional, Tuple, Union
 from os import path
 import os
+import time
 
 from autoconf import conf
 import autofit as af
@@ -42,7 +43,9 @@ def __init__(self, cosmology: LensingCosmology = Planck15):
         """
         self.cosmology = cosmology
 
-    def plane_via_instance_from(self, instance: af.ModelInstance) -> Plane:
+    def plane_via_instance_from(
+        self, instance: af.ModelInstance, run_time_dict: Optional[Dict] = None
+    ) -> Plane:
         """
         Create a `Plane` from the galaxies contained in a model instance.
 
@@ -57,8 +60,134 @@ def plane_via_instance_from(self, instance: af.ModelInstance) -> Plane:
         An instance of the Plane class that is used to then fit the dataset.
         """
         if hasattr(instance, "clumps"):
-            return Plane(galaxies=instance.galaxies + instance.clumps)
-        return Plane(galaxies=instance.galaxies)
+            return Plane(
+                galaxies=instance.galaxies + instance.clumps,
+                run_time_dict=run_time_dict,
+            )
+        return Plane(galaxies=instance.galaxies, run_time_dict=run_time_dict)
+
+    @property
+    def fit_func(self) -> Callable:
+        raise NotImplementedError
+
+    def profile_log_likelihood_function(
+        self, instance: af.ModelInstance, paths: Optional[af.DirectoryPaths] = None
+    ) -> Tuple[Dict, Dict]:
+        """
+        This function is optionally called throughout a model-fit to profile the log likelihood function.
+
+        All function calls inside the `log_likelihood_function` that are decorated with the `profile_func` are timed
+        with their times stored in a dictionary called the `run_time_dict`.
+
+        An `info_dict` is also created which stores information on aspects of the model and dataset that dictate
+        run times, so the profiled times can be interpreted with this context.
+
+        The results of this profiling are then output to hard-disk in the `preloads` folder of the model-fit results,
+        which they can be inspected to ensure run-times are as expected.
+
+        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).
+        paths
+            The PyAutoFit 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.
+
+        Returns
+        -------
+        Two dictionaries, the profiling dictionary and info dictionary, which contain the profiling times of the
+        `log_likelihood_function` and information on the model and dataset used to perform the profiling.
+        """
+        run_time_dict = {}
+        info_dict = {}
+
+        repeats = conf.instance["general"]["profiling"]["repeats"]
+        info_dict["repeats"] = repeats
+
+        fit = self.fit_func(instance=instance)
+        fit.figure_of_merit
+
+        start = time.time()
+
+        for i in range(repeats):
+            try:
+                fit = self.fit_func(instance=instance)
+                fit.figure_of_merit
+            except Exception:
+                logger.info(
+                    "Profiling failed. Returning without outputting information."
+                )
+                return
+
+        fit_time = (time.time() - start) / repeats
+
+        run_time_dict["fit_time"] = fit_time
+
+        fit = self.fit_func(instance=instance, run_time_dict=run_time_dict)
+        fit.figure_of_merit
+
+        try:
+            info_dict["image_pixels"] = self.dataset.grid.sub_shape_slim
+            info_dict["sub_size_light_profiles"] = self.dataset.grid.sub_size
+        except AttributeError:
+            pass
+
+        if fit.model_obj.has(cls=aa.Pixelization):
+            info_dict["use_w_tilde"] = fit.inversion.settings.use_w_tilde
+            info_dict["sub_size_pixelization"] = self.dataset.grid_pixelization.sub_size
+            info_dict[
+                "use_positive_only_solver"
+            ] = fit.inversion.settings.use_positive_only_solver
+            info_dict[
+                "force_edge_pixels_to_zeros"
+            ] = fit.inversion.settings.force_edge_pixels_to_zeros
+            info_dict["use_w_tilde_numpy"] = fit.inversion.settings.use_w_tilde_numpy
+            info_dict["source_pixels"] = len(fit.inversion.reconstruction)
+
+            if hasattr(fit.inversion, "w_tilde"):
+                info_dict[
+                    "w_tilde_curvature_preload_size"
+                ] = fit.inversion.w_tilde.curvature_preload.shape[0]
+
+        self.output_profiling_info(
+            paths=paths, run_time_dict=run_time_dict, info_dict=info_dict
+        )
+
+        return run_time_dict, info_dict
+
+    def output_profiling_info(
+        self, paths: Optional[af.DirectoryPaths], run_time_dict: Dict, info_dict: Dict
+    ):
+        """
+        Output the log likelihood function profiling information to hard-disk as a json file.
+
+        This function is separate from the `profile_log_likelihood_function` function above such that it can be
+        called by children `Analysis` classes that profile additional aspects of the model-fit and therefore add
+        extra information to the `run_time_dict` and `info_dict`.
+
+        Parameters
+        ----------
+        paths
+            The PyAutoFit 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.
+        run_time_dict
+            A dictionary containing the profiling times of the functions called by the `log_likelihood_function`.
+        info_dict
+            A dictionary containing information on the model and dataset used to perform the profiling, where these
+            settings typically control the overall run-time.
+        """
+
+        if paths is None:
+            return
+
+        os.makedirs(paths.profile_path, exist_ok=True)
+
+        with open(path.join(paths.profile_path, "run_time_dict.json"), "w+") as f:
+            json.dump(run_time_dict, f, indent=4)
+
+        with open(path.join(paths.profile_path, "info_dict.json"), "w+") as f:
+            json.dump(info_dict, f, indent=4)
 
 
 class AnalysisDataset(Analysis):
@@ -419,7 +548,3 @@ def output_or_check_figure_of_merit_sanity(
                         f"Old Figure of Merit = {figure_of_merit_sanity}\n"
                         f"New Figure of Merit = {figure_of_merit}"
                     )
-
-    @property
-    def fit_func(self):
-        raise NotImplementedError

autogalaxy/config/general.yaml

@@ -11,4 +11,4 @@ test:
   bypass_figure_of_merit_sanity: false
   check_preloads: false
   exception_override: false
-  preloads_check_threshold: 0.1     # If the figure of merit of a fit with and without preloads is greater than this threshold, the check preload test fails and an exception raised for a model-fit. 
+  preloads_check_threshold: 1.0     # If the figure of merit of a fit with and without preloads is greater than this threshold, the check preload test fails and an exception raised for a model-fit. 

autogalaxy/imaging/fit_imaging.py

@@ -25,7 +25,7 @@ def __init__(
         settings_pixelization: aa.SettingsPixelization = aa.SettingsPixelization(),
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
         preloads: aa.Preloads = Preloads(),
-        profiling_dict: Optional[Dict] = None,
+        run_time_dict: Optional[Dict] = None,
     ):
         """
         Fits an imaging dataset using a `Plane` object.
@@ -65,7 +65,7 @@ def __init__(
         preloads
             Contains preloaded calculations (e.g. linear algebra matrices) which can skip certain calculations in
             the fit.
-        profiling_dict
+        run_time_dict
             A dictionary which if passed to the fit records how long fucntion calls which have the `profile_func`
             decorator take to run.
         """
@@ -75,7 +75,7 @@ def __init__(
 
         super().__init__(
             dataset=dataset,
-            profiling_dict=profiling_dict,
+            run_time_dict=run_time_dict,
         )
         AbstractFitInversion.__init__(
             self=self, model_obj=plane, settings_inversion=settings_inversion
@@ -145,6 +145,7 @@ def plane_to_inversion(self) -> PlaneToInversion:
             settings_pixelization=self.settings_pixelization,
             settings_inversion=self.settings_inversion,
             preloads=self.preloads,
+            run_time_dict=self.run_time_dict,
         )
 
     @cached_property
@@ -304,7 +305,7 @@ def refit_with_new_preloads(
         -------
         A new fit which has used new preloads input into this function but the same dataset, plane and other settings.
         """
-        profiling_dict = {} if self.profiling_dict is not None else None
+        run_time_dict = {} if self.run_time_dict is not None else None
 
         settings_inversion = (
             self.settings_inversion
@@ -318,5 +319,5 @@ def refit_with_new_preloads(
             settings_pixelization=self.settings_pixelization,
             settings_inversion=settings_inversion,
             preloads=preloads,
-            profiling_dict=profiling_dict,
+            run_time_dict=run_time_dict,
         )

autogalaxy/imaging/model/analysis.py

@@ -1,6 +1,6 @@
 import numpy as np
-import os
-from typing import Dict, Optional
+
+from typing import Dict, Optional, Tuple
 
 import autofit as af
 import autoarray as aa
@@ -151,7 +151,7 @@ def fit_imaging_via_instance_from(
         self,
         instance: af.ModelInstance,
         preload_overwrite: Optional[Preloads] = None,
-        profiling_dict: Optional[Dict] = None,
+        run_time_dict: Optional[Dict] = None,
     ) -> FitImaging:
         """
         Given a model instance create a `FitImaging` object.
@@ -166,7 +166,7 @@ def fit_imaging_via_instance_from(
             via a non-linear search).
         preload_overwrite
             If a `Preload` object is input this is used instead of the preloads stored as an attribute in the analysis.
-        profiling_dict
+        run_time_dict
             A dictionary which times functions called to fit the model to data, for profiling.
 
         Returns
@@ -176,19 +176,21 @@ def fit_imaging_via_instance_from(
         """
         instance = self.instance_with_associated_adapt_images_from(instance=instance)
 
-        plane = self.plane_via_instance_from(instance=instance)
+        plane = self.plane_via_instance_from(
+            instance=instance, run_time_dict=run_time_dict
+        )
 
         return self.fit_imaging_via_plane_from(
             plane=plane,
             preload_overwrite=preload_overwrite,
-            profiling_dict=profiling_dict,
+            run_time_dict=run_time_dict,
         )
 
     def fit_imaging_via_plane_from(
         self,
         plane: Plane,
         preload_overwrite: Optional[Preloads] = None,
-        profiling_dict: Optional[Dict] = None,
+        run_time_dict: Optional[Dict] = None,
     ) -> FitImaging:
         """
         Given a `Plane`, which the analysis constructs from a model instance, create a `FitImaging` object.
@@ -202,7 +204,7 @@ def fit_imaging_via_plane_from(
             The plane of galaxies whose model images are used to fit the imaging data.
         preload_overwrite
             If a `Preload` object is input this is used instead of the preloads stored as an attribute in the analysis.
-        profiling_dict
+        run_time_dict
             A dictionary which times functions called to fit the model to data, for profiling.
 
         Returns
@@ -219,7 +221,7 @@ def fit_imaging_via_plane_from(
             settings_pixelization=self.settings_pixelization,
             settings_inversion=self.settings_inversion,
             preloads=preloads,
-            profiling_dict=profiling_dict,
+            run_time_dict=run_time_dict,
         )
 
     @property
@@ -400,3 +402,44 @@ def save_attributes_for_aggregator(self, paths: af.DirectoryPaths):
 
         paths.save_object("psf", self.dataset.psf)
         paths.save_object("mask", self.dataset.mask)
+
+    def profile_log_likelihood_function(
+        self, instance: af.ModelInstance, paths: Optional[af.DirectoryPaths] = None
+    ) -> Tuple[Dict, Dict]:
+        """
+        This function is optionally called throughout a model-fit to profile the log likelihood function.
+
+        All function calls inside the `log_likelihood_function` that are decorated with the `profile_func` are timed
+        with their times stored in a dictionary called the `run_time_dict`.
+
+        An `info_dict` is also created which stores information on aspects of the model and dataset that dictate
+        run times, so the profiled times can be interpreted with this context.
+
+        The results of this profiling are then output to hard-disk in the `preloads` folder of the model-fit results,
+        which they can be inspected to ensure run-times are as expected.
+
+        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).
+        paths
+            The PyAutoFit 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.
+
+        Returns
+        -------
+        Two dictionaries, the profiling dictionary and info dictionary, which contain the profiling times of the
+        `log_likelihood_function` and information on the model and dataset used to perform the profiling.
+        """
+        run_time_dict, info_dict = super().profile_log_likelihood_function(
+            instance=instance,
+        )
+
+        info_dict["psf_shape_2d"] = self.dataset.psf.shape_native
+
+        self.output_profiling_info(
+            paths=paths, run_time_dict=run_time_dict, info_dict=info_dict
+        )
+
+        return run_time_dict, info_dict

autogalaxy/interferometer/fit_interferometer.py

@@ -20,7 +20,7 @@ def __init__(
         settings_pixelization: aa.SettingsPixelization = aa.SettingsPixelization(),
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
         preloads: aa.Preloads = Preloads(),
-        profiling_dict: Optional[Dict] = None,
+        run_time_dict: Optional[Dict] = None,
     ):
         """
         Fits an interferometer dataset using a `Plane` object.
@@ -61,7 +61,7 @@ def __init__(
         preloads
             Contains preloaded calculations (e.g. linear algebra matrices) which can skip certain calculations in
             the fit.
-        profiling_dict
+        run_time_dict
             A dictionary which if passed to the fit records how long fucntion calls which have the `profile_func`
             decorator take to run.
         """
@@ -72,7 +72,7 @@ def __init__(
             settings_inversion.use_w_tilde = False
 
         super().__init__(
-            dataset=dataset, use_mask_in_fit=False, profiling_dict=profiling_dict
+            dataset=dataset, use_mask_in_fit=False, run_time_dict=run_time_dict
         )
         AbstractFitInversion.__init__(
             self=self, model_obj=plane, settings_inversion=settings_inversion
@@ -241,10 +241,10 @@ def refit_with_new_preloads(
         -------
         A new fit which has used new preloads input into this function but the same dataset, plane and other settings.
         """
-        if self.profiling_dict is not None:
-            profiling_dict = {}
+        if self.run_time_dict is not None:
+            run_time_dict = {}
         else:
-            profiling_dict = None
+            run_time_dict = None
 
         if settings_inversion is None:
             settings_inversion = self.settings_inversion
@@ -255,5 +255,5 @@ def refit_with_new_preloads(
             settings_pixelization=self.settings_pixelization,
             settings_inversion=settings_inversion,
             preloads=preloads,
-            profiling_dict=profiling_dict,
+            run_time_dict=run_time_dict,
         )