Merge pull request #361 from Jammy2211/feature/jax_remove_profiling

Feature/jax remove profiling

Author: Jammy2211

autolens/analysis/analysis/lens.py

@@ -46,14 +46,10 @@ def __init__(
     def tracer_via_instance_from(
         self,
         instance: af.ModelInstance,
-        run_time_dict: Optional[Dict] = None,
     ) -> Tracer:
         """
         Create a `Tracer` from the galaxies contained in a model instance.
 
-        If PyAutoFit's profiling tools are used with the analysis class, this function may receive a `run_time_dict`
-        which times how long each set of the model-fit takes to perform.
-
         Parameters
         ----------
         instance
@@ -90,13 +86,11 @@ def tracer_via_instance_from(
             if getattr(instance, "extra_galaxies", None) is not None:
                 return Tracer(
                     galaxies=instance.galaxies + instance.extra_galaxies,
-                    run_time_dict=run_time_dict,
                 )
 
         return Tracer(
             galaxies=instance.galaxies,
             cosmology=cosmology,
-            run_time_dict=run_time_dict,
         )
 
     def log_likelihood_penalty_from(

autolens/config/general.yaml

@@ -1,3 +1,5 @@
+jax:
+  use_jax: true                       # If True, uses JAX internally, whereas False uses normal Numpy.
 output:
   fit_dill: false
 test:

autolens/imaging/fit_imaging.py

@@ -23,7 +23,6 @@ def __init__(
         dataset_model : Optional[aa.DatasetModel] = None,
         adapt_images: Optional[ag.AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
-        run_time_dict: Optional[Dict] = None,
     ):
         """
         Fits an imaging dataset using a `Tracer` object.
@@ -62,12 +61,9 @@ def __init__(
             reconstructed galaxy's morphology.
         settings_inversion
             Settings controlling how an inversion is fitted for example which linear algebra formalism is used.
-        run_time_dict
-            A dictionary which if passed to the fit records how long function calls which have the `profile_func`
-            decorator take to run.
         """
 
-        super().__init__(dataset=dataset, dataset_model=dataset_model, run_time_dict=run_time_dict)
+        super().__init__(dataset=dataset, dataset_model=dataset_model)
         AbstractFitInversion.__init__(
             self=self, model_obj=tracer, settings_inversion=settings_inversion
         )
@@ -120,7 +116,6 @@ def tracer_to_inversion(self) -> TracerToInversion:
             tracer=self.tracer,
             adapt_images=self.adapt_images,
             settings_inversion=self.settings_inversion,
-            run_time_dict=self.run_time_dict
         )
 
     @cached_property

autolens/imaging/model/analysis.py

@@ -92,34 +92,15 @@ def log_likelihood_function(self, instance: af.ModelInstance) -> float:
             The log likelihood indicating how well this model instance fitted the imaging data.
         """
 
-        try:
-            log_likelihood_penalty = self.log_likelihood_penalty_from(
-                instance=instance
-            )
-        except Exception as e:
-            raise e
-
-        try:
-            return self.fit_from(instance=instance).figure_of_merit + log_likelihood_penalty
-        except (
-            PixelizationException,
-            exc.PixelizationException,
-            exc.InversionException,
-            exc.GridException,
-            exc.MeshException,
-            ValueError,
-            TypeError,
-            np.linalg.LinAlgError,
-            OverflowError,
-        ) as e:
-            print(e)
-            fggdfg
-            raise exc.FitException from e
+        log_likelihood_penalty = self.log_likelihood_penalty_from(
+            instance=instance
+        )
+
+        return self.fit_from(instance=instance).figure_of_merit + log_likelihood_penalty
 
     def fit_from(
         self,
         instance: af.ModelInstance,
-        run_time_dict: Optional[Dict] = None,
     ) -> FitImaging:
         """
         Given a model instance create a `FitImaging` object.
@@ -135,8 +116,6 @@ def fit_from(
         check_positions
             Whether the multiple image positions of the lensed source should be checked, i.e. whether they trace
             within the position threshold of one another in the source plane.
-        run_time_dict
-            A dictionary which times functions called to fit the model to data, for profiling.
 
         Returns
         -------
@@ -145,7 +124,7 @@ def fit_from(
         """
 
         tracer = self.tracer_via_instance_from(
-            instance=instance, run_time_dict=run_time_dict
+            instance=instance,
         )
 
         dataset_model = self.dataset_model_via_instance_from(instance=instance)
@@ -158,7 +137,6 @@ def fit_from(
             dataset_model=dataset_model,
             adapt_images=adapt_images,
             settings_inversion=self.settings_inversion,
-            run_time_dict=run_time_dict,
         )
 
     def save_attributes(self, paths: af.DirectoryPaths):
@@ -201,42 +179,3 @@ def save_attributes(self, paths: af.DirectoryPaths):
         )
 
         analysis.save_attributes(paths=paths)
-
-    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 `profiling` 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 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

autolens/interferometer/fit_interferometer.py

@@ -20,7 +20,6 @@ def __init__(
         dataset_model: Optional[aa.DatasetModel] = None,
         adapt_images: Optional[ag.AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
-        run_time_dict: Optional[Dict] = None,
     ):
         """
         Fits an interferometer dataset using a `Tracer` object.
@@ -60,9 +59,6 @@ def __init__(
             reconstructed galaxy's morphology.
         settings_inversion
             Settings controlling how an inversion is fitted for example which linear algebra formalism is used.
-        run_time_dict
-            A dictionary which if passed to the fit records how long function calls which have the `profile_func`
-            decorator take to run.
         """
 
         try:
@@ -76,10 +72,9 @@ def __init__(
 
         self.settings_inversion = settings_inversion
 
-        self.run_time_dict = run_time_dict
-
         super().__init__(
-            dataset=dataset, dataset_model=dataset_model, run_time_dict=run_time_dict
+            dataset=dataset,
+            dataset_model=dataset_model,
         )
         AbstractFitInversion.__init__(
             self=self, model_obj=tracer, settings_inversion=settings_inversion

autolens/interferometer/model/analysis.py

@@ -151,35 +151,13 @@ def log_likelihood_function(self, instance):
             The log likelihood indicating how well this model instance fitted the interferometer data.
         """
 
-        try:
-            log_likelihood_penalty = self.log_likelihood_penalty_from(instance=instance)
-        except Exception as e:
-            raise e
-
-        try:
-            return (
-                self.fit_from(instance=instance).figure_of_merit
-                + log_likelihood_penalty
-            )
-        except (
-            PixelizationException,
-            exc.PixelizationException,
-            exc.InversionException,
-            exc.GridException,
-            exc.MeshException,
-            ValueError,
-            TypeError,
-            np.linalg.LinAlgError,
-            OverflowError,
-        ) as e:
-            print(e)
-            fggdfg
-            raise exc.FitException from e
+        log_likelihood_penalty = self.log_likelihood_penalty_from(instance=instance)
+
+        return self.fit_from(instance=instance).figure_of_merit + log_likelihood_penalty
 
     def fit_from(
         self,
         instance: af.ModelInstance,
-        run_time_dict: Optional[Dict] = None,
     ) -> FitInterferometer:
         """
         Given a model instance create a `FitInterferometer` object.
@@ -205,7 +183,7 @@ def fit_from(
         """
 
         tracer = self.tracer_via_instance_from(
-            instance=instance, run_time_dict=run_time_dict
+            instance=instance,
         )
 
         adapt_images = self.adapt_images_via_instance_from(instance=instance)
@@ -215,7 +193,6 @@ def fit_from(
             tracer=tracer,
             adapt_images=adapt_images,
             settings_inversion=self.settings_inversion,
-            run_time_dict=run_time_dict,
         )
 
     def save_attributes(self, paths: af.DirectoryPaths):
@@ -257,45 +234,3 @@ def save_attributes(self, paths: af.DirectoryPaths):
         )
 
         analysis.save_attributes(paths=paths)
-
-    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 `profiling` 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 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["number_of_visibilities"] = self.dataset.data.shape[0]
-        info_dict["transformer_cls"] = self.dataset.transformer.__class__.__name__
-
-        self.output_profiling_info(
-            paths=paths, run_time_dict=run_time_dict, info_dict=info_dict
-        )
-
-        return run_time_dict, info_dict

autolens/lens/mock/mock_to_inversion.py

@@ -6,13 +6,10 @@ def __init__(
         self,
         tracer,
         image_plane_mesh_grid_pg_list=None,
-        run_time_dict: Optional[Dict] = None,
     ):
         self.tracer = tracer
 
         self.image_plane_mesh_grid_pg_list = image_plane_mesh_grid_pg_list
 
-        self.run_time_dict = run_time_dict
-
     def image_plane_mesh_grid_pg_list(self):
         return self.image_plane_mesh_grid_pg_list