Merge pull request #1155 from rhayes777/feature/jax_wrapper

Jax to Main!

Author: Jammy2211

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