Jax to Main!

README.rst

@@ -17,6 +17,18 @@ PyAutoFit: Classy Probabilistic Programming
 .. |JOSS| image:: https://joss.theoj.org/papers/10.21105/joss.02550/status.svg
    :target: https://doi.org/10.21105/joss.02550
 
+.. image:: https://www.repostatus.org/badges/latest/active.svg
+    :target: https://www.repostatus.org/#active
+    :alt: Project Status: Active
+
+.. image:: https://img.shields.io/pypi/pyversions/autofit
+    :target: https://pypi.org/project/autofit/
+    :alt: Python Versions
+
+.. image:: https://img.shields.io/pypi/v/autofit.svg
+    :target: https://pypi.org/project/autofit/
+    :alt: PyPI Version
+
 |binder| |Tests| |Build| |RTD| |JOSS|
 
 `Installation Guide <https://pyautofit.readthedocs.io/en/latest/installation/overview.html>`_ |

autofit/__init__.py

@@ -1,13 +1,15 @@
+from autoconf.dictable import register_parser
+from . import conf
+
+conf.instance.register(__file__)
+
 import abc
 import pickle
-
 from dill import register
 
-from autoconf.dictable import register_parser
-from .non_linear.grid.grid_search import GridSearch as SearchGridSearch
-from . import conf
 from . import exc
 from . import mock as m
+from .non_linear.grid.grid_search import GridSearch as SearchGridSearch
 from .aggregator.base import AggBase
 from .database.aggregator.aggregator import GridSearchAggregator
 from .graphical.expectation_propagation.history import EPHistory
@@ -54,6 +56,8 @@
 from .mapper.prior import GaussianPrior
 from .mapper.prior import LogGaussianPrior
 from .mapper.prior import LogUniformPrior
+from .mapper.prior import TruncatedGaussianPrior
+from .mapper.prior.vectorized import PriorVectorized
 from .mapper.prior.abstract import Prior
 from .mapper.prior.tuple_prior import TuplePrior
 from .mapper.prior import UniformPrior
@@ -136,6 +140,6 @@ def save_abc(pickler, obj):
     pickle._Pickler.save_type(pickler, obj)
 
 
-conf.instance.register(__file__)
+
 
 __version__ = "2025.5.10.1"

autofit/aggregator/summary/aggregate_fits.py

@@ -2,8 +2,6 @@
 from enum import Enum
 from typing import Dict, List, Union
 
-from astropy.table import Table
-from astropy.io import fits
 from pathlib import Path
 
 from autofit.aggregator.search_output import SearchOutput
@@ -53,7 +51,7 @@ def __init__(self, aggregator: Union[Aggregator, List[SearchOutput]]):
     def _hdus(
         result: SearchOutput,
         hdus: List[Enum],
-    ) -> List[fits.ImageHDU]:
+    ) -> "List[fits.ImageHDU]":
         """
         Extract the HDUs from a given fits for a given search.
 
@@ -68,6 +66,8 @@ def _hdus(
         -------
         The extracted HDUs.
         """
+        from astropy.io import fits
+
         row = []
         for hdu in hdus:
             source = result.value(subplot_filename(hdu))
@@ -80,7 +80,7 @@ def _hdus(
             )
         return row
 
-    def extract_fits(self, hdus: List[Enum]) -> fits.HDUList:
+    def extract_fits(self, hdus: List[Enum]) -> "fits.HDUList":
         """
         Extract the HDUs from the fits files for every search in the aggregator.
 
@@ -95,6 +95,8 @@ def extract_fits(self, hdus: List[Enum]) -> fits.HDUList:
         -------
         The extracted HDUs.
         """
+        from astropy.io import fits
+
         output = [fits.PrimaryHDU()]
         for result in self.aggregator:
             output.extend(self._hdus(result, hdus))
@@ -117,6 +119,8 @@ def extract_csv(self, filename: str) -> List[Dict]:
         -------
         The extracted HDUs.
         """
+        from astropy.table import Table
+
         output = []
         for result in self.aggregator:
             output.append(Table.read(result.value(filename), format="fits"))
@@ -144,6 +148,8 @@ def output_to_folder(
             The name of the fits file. This is the attribute of the search output that is used to name the file.
             OR a list of names for each HDU.
         """
+        from astropy.io import fits
+
         folder.mkdir(parents=True, exist_ok=True)
 
         for i, result in enumerate(self.aggregator):

autofit/config/general.yaml

@@ -1,13 +1,13 @@
+jax:
+  use_jax: false                    # If True, PyAutoFit uses JAX internally, whereas False uses normal Numpy.
 analysis:
   n_cores: 1                        # The number of cores a parallelized sum of Analysis classes uses by default.
 hpc:
   hpc_mode: false                   # If True, use HPC mode, which disables GUI visualization, logging to screen and other settings which are not suited to running on a super computer.
   iterations_per_update: 5000       # The number of iterations between every update (visualization, results output, etc) in HPC mode.
 inversion:
   check_reconstruction: true        # If True, the inversion's reconstruction is checked to ensure the solution of a meshs's mapper is not an invalid solution where the values are all the same.
-  reconstruction_vmax_factor: 0.5   # Plots of an Inversion's reconstruction use the reconstructed data's bright value multiplied by this factor.  
-model:
-  ignore_prior_limits: false        # If ``True`` the limits applied to priors will be ignored, where limits set upper / lower limits. This stops PriorLimitException's from being raised.
+  reconstruction_vmax_factor: 0.5   # Plots of an Inversion's reconstruction use the reconstructed data's bright value multiplied by this factor.
 output:
   force_pickle_overwrite: false     # If True pickle files output by a search (e.g. samples.pickle) are recreated when a new model-fit is performed.
   force_visualize_overwrite: false  # If True, visualization images output by a search (e.g. subplots of the fit) are recreated when a new model-fit is performed.

autofit/config/output.yaml

@@ -86,8 +86,10 @@ start_point: true
 # manually after completing a successful model-fit. This will save computational run time by not computing latent
 # variables during a any model-fit which is unsuccessful.
 
-latent_during_fit: true # Whether to output the `latent.csv`, `latent.results` and `latent_summary.json` files during the fit when it performs on-the-fly output.
+latent_during_fit: false # Whether to output the `latent.csv`, `latent.results` and `latent_summary.json` files during the fit when it performs on-the-fly output.
 latent_after_fit: true # If `latent_during_fit` is False, whether to output the `latent.csv`, `latent.results` and `latent_summary.json` files after the fit is complete.
+latent_draw_via_pdf : true # Whether to draw latent variable values via the PDF of every sample, which uses fewer samples to estimate latent variable errors. If False, latent variable values are drawn from every sample.
+latent_draw_via_pdf_size : 100 # The number of samples drawn to estimate latent variable errors if `latent_draw_via_pdf` is True.
 latent_csv: true # Whether to ouptut the `latent.csv` file.
 latent_results: true # Whether to output the `latent.results` file.
 

autofit/config/priors/Exponential.yaml

@@ -1,5 +1,5 @@
 centre:
-  gaussian_limits:
+  limits:
     lower: -inf
     upper: inf
   lower_limit: 0.0
@@ -9,7 +9,7 @@ centre:
     type: Absolute
     value: 20.0
 normalization:
-  gaussian_limits:
+  limits:
     lower: 0.0
     upper: inf
   lower_limit: 1.0e-06
@@ -19,7 +19,7 @@ normalization:
     type: Relative
     value: 0.5
 rate:
-  gaussian_limits:
+  limits:
     lower: 0.0
     upper: inf
   lower_limit: 0.0

autofit/config/priors/Gaussian.yaml

@@ -1,5 +1,5 @@
 centre:
-  gaussian_limits:
+  limits:
     lower: -inf
     upper: inf
   lower_limit: 0.0
@@ -9,7 +9,7 @@ centre:
     type: Absolute
     value: 20.0
 normalization:
-  gaussian_limits:
+  limits:
     lower: 0.0
     upper: inf
   lower_limit: 1.0e-06
@@ -19,7 +19,7 @@ normalization:
     type: Relative
     value: 0.5
 sigma:
-  gaussian_limits:
+  limits:
     lower: 0.0
     upper: inf
   lower_limit: 0.0

autofit/config/priors/Gaussian2D.yaml

@@ -1,5 +1,5 @@
 centre_0:
-  gaussian_limits:
+  limits:
     lower: -inf
     upper: inf
   lower_limit: 0.0
@@ -9,7 +9,7 @@ centre_0:
     type: Absolute
     value: 20.0
 centre_1:
-  gaussian_limits:
+  limits:
     lower: -inf
     upper: inf
   lower_limit: 0.0
@@ -19,7 +19,7 @@ centre_1:
     type: Absolute
     value: 20.0
 normalization:
-  gaussian_limits:
+  limits:
     lower: 0.0
     upper: inf
   lower_limit: 1.0e-06
@@ -29,7 +29,7 @@ normalization:
     type: Relative
     value: 0.5
 sigma:
-  gaussian_limits:
+  limits:
     lower: 0.0
     upper: inf
   lower_limit: 0.0

autofit/config/priors/GaussianKurtosis.yaml

@@ -1,12 +1,5 @@
-GaussianPrior:
-  lower_limit:
-    type: Constant
-    value: -inf
-  upper_limit:
-    type: Constant
-    value: inf
 centre:
-  gaussian_limits:
+  limits:
     lower: -inf
     upper: inf
   lower_limit: 0.0
@@ -16,7 +9,7 @@ centre:
     type: Absolute
     value: 20.0
 kurtosis:
-  gaussian_limits:
+  limits:
     lower: -inf
     upper: inf
   lower_limit: 0.0
@@ -26,7 +19,7 @@ kurtosis:
     type: Absolute
     value: 20.0
 normalization:
-  gaussian_limits:
+  limits:
     lower: 0.0
     upper: inf
   lower_limit: 1.0e-06
@@ -36,7 +29,7 @@ normalization:
     type: Relative
     value: 0.5
 sigma:
-  gaussian_limits:
+  limits:
     lower: 0.0
     upper: inf
   lower_limit: 0.0

autofit/config/priors/MultiLevelGaussians.yaml

@@ -1,5 +1,5 @@
 higher_level_centre:
-  gaussian_limits:
+  limits:
     lower: -inf
     upper: inf
   lower_limit: 0.0

autofit/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
 
@@ -28,9 +28,9 @@ The sections of this example config set the following:
         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. For a full description of prior passing, checkout the examples
         in 'autolens_workspace/examples/complex/linking'.
-    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

autofit/config/priors/model.yaml

@@ -1,6 +1,6 @@
 Exponential:
   centre:
-    gaussian_limits:
+    limits:
       lower: -inf
       upper: inf
     lower_limit: 0.0
@@ -10,7 +10,7 @@ Exponential:
       type: Absolute
       value: 20.0
   normalization:
-    gaussian_limits:
+    limits:
       lower: 0.0
       upper: inf
     lower_limit: 1.0e-06
@@ -20,7 +20,7 @@ Exponential:
       type: Relative
       value: 0.5
   rate:
-    gaussian_limits:
+    limits:
       lower: 0.0
       upper: inf
     lower_limit: 0.0
@@ -31,7 +31,7 @@ Exponential:
       value: 0.5
 Gaussian:
   centre:
-    gaussian_limits:
+    limits:
       lower: -inf
       upper: inf
     lower_limit: 0.0
@@ -41,7 +41,7 @@ Gaussian:
       type: Absolute
       value: 20.0
   normalization:
-    gaussian_limits:
+    limits:
       lower: 0.0
       upper: inf
     lower_limit: 1.0e-06
@@ -51,7 +51,7 @@ Gaussian:
       type: Relative
       value: 0.5
   sigma:
-    gaussian_limits:
+    limits:
       lower: 0.0
       upper: inf
     lower_limit: 0.0