Merge pull request #241 from Jammy2211/feature/jax_remove_profiling

Feature/jax remove profiling

Author: Jammy2211

autogalaxy/__init__.py

@@ -1,3 +1,8 @@
+from autoconf.dictable import register_parser
+from autofit import conf
+
+conf.instance.register(__file__)
+
 from autoconf.dictable import from_dict, from_json, output_to_json, to_dict
 from autoarray.dataset import preprocess  # noqa
 from autoarray.dataset.imaging.dataset import Imaging  # noqa
@@ -112,6 +117,4 @@
 from autoconf.fitsable import output_to_fits
 from autoconf.fitsable import hdu_list_for_output_from
 
-conf.instance.register(__file__)
-
 __version__ = "2025.5.10.1"

autogalaxy/aggregator/agg_util.py

@@ -140,7 +140,7 @@ def adapt_images_from(
 
             galaxy_name_image_dict[value.header["EXTNAME"].lower()] = adapt_image
 
-        instance = fit.model.instance_from_prior_medians(ignore_prior_limits=True)
+        instance = fit.model.instance_from_prior_medians(ignore_assertions=True)
 
         adapt_images = AdaptImages(galaxy_name_image_dict=galaxy_name_image_dict)
 

autogalaxy/analysis/analysis/analysis.py

@@ -37,7 +37,8 @@ def __init__(self, cosmology: LensingCosmology = Planck15):
         self.cosmology = cosmology
 
     def galaxies_via_instance_from(
-        self, instance: af.ModelInstance, run_time_dict: Optional[Dict] = None
+        self,
+        instance: af.ModelInstance,
     ) -> List[Galaxy]:
         """
         Create a list of galaxies from a model instance, which is used to fit the dataset.
@@ -60,10 +61,9 @@ def galaxies_via_instance_from(
             if getattr(instance, "extra_galaxies", None) is not None:
                 return Galaxies(
                     galaxies=instance.galaxies + instance.extra_galaxies,
-                    run_time_dict=run_time_dict,
                 )
 
-        return Galaxies(galaxies=instance.galaxies, run_time_dict=run_time_dict)
+        return Galaxies(galaxies=instance.galaxies)
 
     def dataset_model_via_instance_from(
         self, instance: af.ModelInstance
@@ -126,135 +126,3 @@ def make_result(
             search_internal=search_internal,
             analysis=self,
         )
-
-    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.
-        """
-
-        if isinstance(paths, af.DatabasePaths):
-            return
-
-        run_time_dict = {}
-        info_dict = {}
-
-        repeats = conf.instance["general"]["profiling"]["repeats"]
-        info_dict["repeats"] = repeats
-
-        # Ensure numba functions are compiled before profiling begins.
-
-        fit = self.fit_from(instance=instance)
-        fit.figure_of_merit
-
-        start = time.time()
-
-        for _ in range(repeats):
-            try:
-                fit = self.fit_from(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_from(instance=instance, run_time_dict=run_time_dict)
-        fit.figure_of_merit
-
-        try:
-            info_dict["image_pixels"] = self.dataset.grids.lp.shape_slim
-            info_dict["sub_total_light_profiles"] = (
-                self.dataset.grids.lp.over_sampler.sub_total
-            )
-        except AttributeError:
-            pass
-
-        if fit.model_obj.has(cls=aa.Pixelization):
-            info_dict["use_w_tilde"] = fit.inversion.settings.use_w_tilde
-            try:
-                info_dict["sub_total_pixelization"] = (
-                    self.dataset.grids.pixelization.over_sampler.sub_total
-                )
-            except AttributeError:
-                pass
-            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 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)

autogalaxy/config/general.yaml

@@ -1,3 +1,5 @@
+jax:
+  use_jax: true                       # If True, uses JAX internally, whereas False uses normal Numpy.
 fits:
   flip_for_ds9: true
 grid:

autogalaxy/config/priors/README.rst

@@ -13,7 +13,7 @@ They appear as follows:
         width_modifier:
           type: Absolute
           value: 20.0
-        gaussian_limits:
+        limits:
           lower: -inf
           upper: inf
 
@@ -27,9 +27,9 @@ The sections of this example config set the following:
     width_modifier
         When the results of a search are passed to a subsequent search to set up the priors of its non-linear search,
         this entry describes how the Prior is passed.
-    gaussian_limits
+    limits
         When the results of a search are passed to a subsequent search, they are passed using a GaussianPrior. The
-        gaussian_limits set the physical lower and upper limits of this GaussianPrior, such that parameter samples
+        limits set the physical lower and upper limits of this GaussianPrior, such that parameter samples
         can not go beyond these limits.
 
 The files ``template_module.yaml`` and ``TemplateObject.yaml`` give templates one can use to set up prior default

autogalaxy/config/priors/ellipse/ellipse.yaml

@@ -3,47 +3,39 @@ Ellipse:
     type: Gaussian
     mean: 0.0
     sigma: 0.3
-    lower_limit: -inf
-    upper_limit: inf
     width_modifier:
       type: Absolute
       value: 0.05
-    gaussian_limits:
+    limits:
       lower: -inf
       upper: inf
   centre_1:
     type: Gaussian
     mean: 0.0
     sigma: 0.3
-    lower_limit: -inf
-    upper_limit: inf
     width_modifier:
       type: Absolute
       value: 0.05
-    gaussian_limits:
+    limits:
       lower: -inf
       upper: inf
   ell_comps_0:
     type: Gaussian
     mean: 0.0
     sigma: 0.3
-    lower_limit: -1.0
-    upper_limit: 1.0
     width_modifier:
       type: Absolute
       value: 0.2
-    gaussian_limits:
+    limits:
       lower: -1.0
       upper: 1.0
   ell_comps_1:
     type: Gaussian
     mean: 0.0
     sigma: 0.3
-    lower_limit: -1.0
-    upper_limit: 1.0
     width_modifier:
       type: Absolute
       value: 0.2
-    gaussian_limits:
+    limits:
       lower: -1.0
       upper: 1.0

autogalaxy/config/priors/ellipse/ellipse_multipole.yaml

@@ -6,7 +6,7 @@ EllipseMultipole:
     width_modifier:
       type: Absolute
       value: 0.05
-    gaussian_limits:
+    limits:
       lower: -inf
       upper: inf
   multipole_comps_1:
@@ -16,6 +16,6 @@ EllipseMultipole:
     width_modifier:
       type: Absolute
       value: 0.05
-    gaussian_limits:
+    limits:
       lower: -inf
       upper: inf

autogalaxy/config/priors/galaxy/redshift.yaml

@@ -6,6 +6,6 @@ Redshift:
     width_modifier:
       type: Absolute
       value: 1.0
-    gaussian_limits:
+    limits:
       lower: 0.0
       upper: inf

autogalaxy/config/priors/image_mesh/hilbert.yaml

@@ -6,7 +6,7 @@ Hilbert:
     width_modifier:
       type: Absolute
       value: 100.0
-    gaussian_limits:
+    limits:
       lower: 50.0
       upper: inf
   weight_floor:
@@ -16,7 +16,7 @@ Hilbert:
     width_modifier:
       type: Absolute
       value: 0.1
-    gaussian_limits:
+    limits:
       lower: 0.0
       upper: inf
   weight_power:
@@ -26,6 +26,6 @@ Hilbert:
     width_modifier:
       type: Absolute
       value: 5.0
-    gaussian_limits:
+    limits:
       lower: 0.0
       upper: inf

autogalaxy/config/priors/image_mesh/kmeans.yaml

@@ -6,7 +6,7 @@ KMeans:
     width_modifier:
       type: Absolute
       value: 100.0
-    gaussian_limits:
+    limits:
       lower: 50.0
       upper: inf
   weight_floor:
@@ -16,7 +16,7 @@ KMeans:
     width_modifier:
       type: Absolute
       value: 0.1
-    gaussian_limits:
+    limits:
       lower: 0.0
       upper: inf
   weight_power:
@@ -26,6 +26,6 @@ KMeans:
     width_modifier:
       type: Absolute
       value: 5.0
-    gaussian_limits:
+    limits:
       lower: 0.0
       upper: inf