Merge pull request #1158 from rhayes777/feature/build_fixes

Feature/build fixes

Author: Jammy2211

autofit/__init__.py

@@ -94,7 +94,6 @@
 from .non_linear.result import ResultsCollection
 from .non_linear.settings import SettingsSearch
 from .non_linear.samples.pdf import marginalize
-from .example.model import Gaussian, Exponential
 from .text import formatter
 from .text import samples_text
 from .visualise import VisualiseGraph
@@ -129,6 +128,7 @@
     "LogUniform",
     "Gaussian",
     "LogGaussian",
+    "TruncatedGaussian",
     "compound",
     "Constant",
 ):
@@ -142,4 +142,4 @@ def save_abc(pickler, obj):
 
 
 
-__version__ = "2025.5.10.1"
+__version__ = "2025.10.16.2"

autofit/config/general.yaml

@@ -1,10 +1,12 @@
 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.
+updates:
+  iterations_per_quick_update: 1e99 # Non-linear search iterations between every quick update, which just displays the maximum likelihood model fit.
+  iterations_per_full_update: 1e99  # Non-linear search iterations between every full update, which outputs all visuals and result fits (e.g. model.result, search.summary), this exits the search and can be slow.  
 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.
+  iterations_per_quick_update: 1e99 #  Non-linear search iterations between every quick update, which just displays the maximum likelihood model fit.
+  iterations_per_full_update: 1e99  # Non-linear search iterations between every full update, which outputs all visuals and result fits (e.g. model.result, search.summary), this exits the search and can be slow.
 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.

autofit/config/non_linear/mcmc.yaml

@@ -26,9 +26,6 @@ Emcee:
     number_of_cores: 1              # The number of cores the search is parallelized over by default, using Python multiprocessing.
   printing:
     silence: false                  # If True, the default print output of the non-linear search is silcened and not printed by the Python interpreter.
-  updates:
-    iterations_per_update: 500      # The number of iterations of the non-linear search performed between every 'update', where an update performs tasks like outputting model.results.
-    remove_state_files_at_end: true # Whether to remove the savestate of the seach (e.g. the Emcee hdf5 file) at the end to save hard-disk space (results are still stored as PyAutoFit pickles and loadable).
 Zeus:
   run:
     check_walkers: true
@@ -58,6 +55,6 @@ Zeus:
   printing:
     silence: false                  # If True, the default print output of the non-linear search is silenced and not printed by the Python interpreter.
 
-  updates:
-    iterations_per_update: 500      # The number of iterations of the non-linear search performed between every 'update', where an update performs tasks like outputting model.results.
+    iterations_per_full_update: 500      # Non-linear search iterations between every full update, which outputs all visuals and result fits (e.g. model.result, search.summary), this exits the search and can be slow.
+    iterations_per_quick_update: 500 #  Non-linear search iterations between every quick update, which just displays the maximum likelihood model fit.
     remove_state_files_at_end: true # Whether to remove the savestate of the seach (e.g. the Emcee hdf5 file) at the end to save hard-disk space (results are still stored as PyAutoFit pickles and loadable).

autofit/config/non_linear/mle.yaml

@@ -24,8 +24,8 @@ PySwarmsGlobal:
     number_of_cores: 1              # The number of cores the search is parallelized over by default, using Python multiprocessing.
   printing:
     silence: false                  # If True, the default print output of the non-linear search is silcened and not printed by the Python interpreter.
-  updates:
-    iterations_per_update: 500      # The number of iterations of the non-linear search performed between every 'update', where an update performs tasks like outputting model.results.
+    iterations_per_full_update: 500      # Non-linear search iterations between every full update, which outputs all visuals and result fits (e.g. model.result, search.summary), this exits the search and can be slow.
+    iterations_per_quick_update: 500 #  Non-linear search iterations between every quick update, which just displays the maximum likelihood model fit.
     remove_state_files_at_end: true # Whether to remove the savestate of the seach (e.g. the Emcee hdf5 file) at the end to save hard-disk space (results are still stored as PyAutoFit pickles and loadable).
 PySwarmsLocal:
   run:
@@ -46,8 +46,8 @@ PySwarmsLocal:
     number_of_cores: 1              # The number of cores the search is parallelized over by default, using Python multiprocessing.
   printing:
     silence: false                  # If True, the default print output of the non-linear search is silcened and not printed by the Python interpreter.
-  updates:
-    iterations_per_update: 500      # The number of iterations of the non-linear search performed between every 'update', where an update performs tasks like outputting model.results.
+    iterations_per_full_update: 500      # Non-linear search iterations between every full update, which outputs all visuals and result fits (e.g. model.result, search.summary), this exits the search and can be slow.
+    iterations_per_quick_update: 500 #  Non-linear search iterations between every quick update, which just displays the maximum likelihood model fit.
     remove_state_files_at_end: true # Whether to remove the savestate of the seach (e.g. the Emcee hdf5 file) at the end to save hard-disk space (results are still stored as PyAutoFit pickles and loadable).
 BFGS:
   search:
@@ -70,8 +70,8 @@ BFGS:
     number_of_cores: 1              # The number of cores the search is parallelized over by default, using Python multiprocessing.
   printing:
     silence: false                  # If True, the default print output of the non-linear search is silcened and not printed by the Python interpreter.
-  updates:
-    iterations_per_update: 500      # The number of iterations of the non-linear search performed between every 'update', where an update performs tasks like outputting model.results.
+    iterations_per_full_update: 500      # Non-linear search iterations between every full update, which outputs all visuals and result fits (e.g. model.result, search.summary), this exits the search and can be slow.
+    iterations_per_quick_update: 500 #  Non-linear search iterations between every quick update, which just displays the maximum likelihood model fit.
     remove_state_files_at_end: true # Whether to remove the savestate of the seach (e.g. the Emcee hdf5 file) at the end to save hard-disk space (results are still stored as PyAutoFit pickles and loadable).
 LBFGS:
   search:
@@ -94,8 +94,8 @@ LBFGS:
     number_of_cores: 1              # The number of cores the search is parallelized over by default, using Python multiprocessing.
   printing:
     silence: false                  # If True, the default print output of the non-linear search is silcened and not printed by the Python interpreter.
-  updates:
-    iterations_per_update: 500      # The number of iterations of the non-linear search performed between every 'update', where an update performs tasks like outputting model.results.
+    iterations_per_full_update: 500      # Non-linear search iterations between every full update, which outputs all visuals and result fits (e.g. model.result, search.summary), this exits the search and can be slow.
+    iterations_per_quick_update: 500 #  Non-linear search iterations between every quick update, which just displays the maximum likelihood model fit.
     remove_state_files_at_end: true # Whether to remove the savestate of the seach (e.g. the Emcee hdf5 file) at the end to save hard-disk space (results are still stored as PyAutoFit pickles and loadable).
 Drawer:
   search:
@@ -108,6 +108,6 @@ Drawer:
     number_of_cores: 1              # The number of cores the search is parallelized over by default, using Python multiprocessing.
   printing:
     silence: false                  # If True, the default print output of the non-linear search is silcened and not printed by the Python interpreter.
-  updates:
-    iterations_per_update: 500      # The number of iterations of the non-linear search performed between every 'update', where an update performs tasks like outputting model.results.
+    iterations_per_full_update: 500      # Non-linear search iterations between every full update, which outputs all visuals and result fits (e.g. model.result, search.summary), this exits the search and can be slow.
+    iterations_per_quick_update: 500 #  Non-linear search iterations between every quick update, which just displays the maximum likelihood model fit.
     remove_state_files_at_end: true # Whether to remove the savestate of the seach (e.g. the Emcee hdf5 file) at the end to save hard-disk space (results are still stored as PyAutoFit pickles and loadable).

autofit/config/non_linear/nest.yaml

@@ -33,9 +33,6 @@ DynestyStatic:
     force_x1_cpu: false             # Force Dynesty to not use Python multiprocessing Pool, which can fix issues on certain operating systems.
   printing:
     silence: false                  # If True, the default print output of the non-linear search is silenced and not printed by the Python interpreter.
-  updates:
-    iterations_per_update: 500      # The number of iterations of the non-linear search performed between every 'update', where an update performs tasks like outputting model.results.
-    remove_state_files_at_end: true # Whether to remove the savestate of the seach (e.g. the Emcee hdf5 file) at the end to save hard-disk space (results are still stored as PyAutoFit pickles and loadable).
 DynestyDynamic:
   search:
     bootstrap: null
@@ -64,9 +61,6 @@ DynestyDynamic:
     force_x1_cpu: false             # Force Dynesty to not use Python multiprocessing Pool, which can fix issues on certain operating systems.
   printing:
     silence: false                  # If True, the default print output of the non-linear search is silenced and not printed by the Python interpreter.
-  updates:
-    iterations_per_update: 500      # The number of iterations of the non-linear search performed between every 'update', where an update performs tasks like outputting model.results.
-    remove_state_files_at_end: true # Whether to remove the savestate of the seach (e.g. the Emcee hdf5 file) at the end to save hard-disk space (results are still stored as PyAutoFit pickles and loadable).
 Nautilus:
   search:
     n_live: 3000                    # Number of so-called live points. New bounds are constructed so that they encompass the live points.
@@ -93,9 +87,6 @@ Nautilus:
     force_x1_cpu: false             # Force Dynesty to not use Python multiprocessing Pool, which can fix issues on certain operating systems.
   printing:
     silence: false                  # If True, the default print output of the non-linear search is silenced and not printed by the Python interpreter.
-  updates:
-    iterations_per_update: 500      # The number of iterations of the non-linear search performed between every 'update', where an update performs tasks like outputting model.results.
-    remove_state_files_at_end: true # Whether to remove the savestate of the seach (e.g. the Emcee hdf5 file) at the end to save hard-disk space (results are still stored as PyAutoFit pickles and loadable).
 UltraNest:
   search:
     draw_multiple: true
@@ -140,6 +131,3 @@ UltraNest:
     number_of_cores: 1              # The number of cores the search is parallelized over by default, using Python multiprocessing.
   printing:
     silence: false                  # If True, the default print output of the non-linear search is silenced and not printed by the Python interpreter.
-  updates:
-    iterations_per_update: 500      # The number of iterations of the non-linear search performed between every 'update', where an update performs tasks like outputting model.results.
-    remove_state_files_at_end: true # Whether to remove the savestate of the seach (e.g. the Emcee hdf5 file) at the end to save hard-disk space (results are still stored as PyAutoFit pickles and loadable).

autofit/example/analysis.py

@@ -1,7 +1,7 @@
+import numpy as np
+from typing import Dict, Optional
 
-from typing import Dict, List, Optional
-
-from autofit.jax_wrapper import numpy as np
+from autofit.jax_wrapper import numpy as xp
 
 import autofit as af
 
@@ -36,7 +36,7 @@ class Analysis(af.Analysis):
     """
     Result = ResultExample
 
-    LATENT_KEYS = ["fwhm"]
+    LATENT_KEYS = ["gaussian.fwhm"]
 
     def __init__(self, data: np.ndarray, noise_map: np.ndarray):
         """
@@ -98,8 +98,8 @@ def model_data_1d_from(self, instance: af.ModelInstance) -> np.ndarray:
         The model data of the profiles.
         """
 
-        xvalues = np.arange(self.data.shape[0])
-        model_data_1d = np.zeros(self.data.shape[0])
+        xvalues = xp.arange(self.data.shape[0])
+        model_data_1d = xp.zeros(self.data.shape[0])
 
         try:
             for profile in instance:

autofit/example/model.py

@@ -1,7 +1,8 @@
 import math
+import numpy as np
 from typing import Tuple
 
-from autofit.jax_wrapper import numpy as np
+from autofit.jax_wrapper import numpy as xp
 
 """
 The `Gaussian` class in this module is the model components that is fitted to data using a non-linear search. The
@@ -46,7 +47,7 @@ def fwhm(self) -> float:
         the free parameters of the model which we are interested and may want to store the full samples information
         on (e.g. to create posteriors).
         """
-        return 2 * np.sqrt(2 * np.log(2)) * self.sigma
+        return 2 * xp.sqrt(2 * xp.log(2)) * self.sigma
 
     def _tree_flatten(self):
         return (self.centre, self.normalization, self.sigma), None
@@ -76,16 +77,16 @@ def model_data_from(self, xvalues: np.ndarray) -> np.ndarray:
         """
         transformed_xvalues = xvalues - self.centre
 
-        return np.multiply(
-            np.divide(self.normalization, self.sigma * np.sqrt(2.0 * np.pi)),
-            np.exp(-0.5 * np.square(np.divide(transformed_xvalues, self.sigma))),
+        return xp.multiply(
+            xp.divide(self.normalization, self.sigma * xp.sqrt(2.0 * xp.pi)),
+            xp.exp(-0.5 * xp.square(xp.divide(transformed_xvalues, self.sigma))),
         )
 
     def f(self, x: float):
         return (
             self.normalization
-            / (self.sigma * np.sqrt(2 * math.pi))
-            * np.exp(-0.5 * ((x - self.centre) / self.sigma) ** 2)
+            / (self.sigma * xp.sqrt(2 * math.pi))
+            * xp.exp(-0.5 * ((x - self.centre) / self.sigma) ** 2)
         )
 
     def __call__(self, xvalues: np.ndarray) -> np.ndarray:
@@ -147,9 +148,9 @@ def model_data_from(self, xvalues: np.ndarray) -> np.ndarray:
         values
             The x coordinates in the original reference frame of the grid.
         """
-        transformed_xvalues = np.subtract(xvalues, self.centre)
-        return self.normalization * np.multiply(
-            self.rate, np.exp(-1.0 * self.rate * abs(transformed_xvalues))
+        transformed_xvalues = xp.subtract(xvalues, self.centre)
+        return self.normalization * xp.multiply(
+            self.rate, xp.exp(-1.0 * self.rate * abs(transformed_xvalues))
         )
 
     def __call__(self, xvalues: np.ndarray) -> np.ndarray:

autofit/fixtures.py

@@ -7,7 +7,7 @@
 def make_model_gaussian_x1():
 
     return af.Model(
-        af.Gaussian
+        af.ex.Gaussian
     )
 
 

autofit/jax_wrapper.py

@@ -1,16 +1,50 @@
+import logging
+
+logger = logging.getLogger(__name__)
+
 """
 Allows the user to switch between using NumPy and JAX for linear algebra operations.
 
 If USE_JAX=true in general.yaml then JAX's NumPy is used, otherwise vanilla NumPy is used.
 """
-import jax
-
 from autoconf import conf
 
 use_jax = conf.instance["general"]["jax"]["use_jax"]
 
 if use_jax:
 
+    import os
+
+    xla_env = os.environ.get("XLA_FLAGS")
+
+    xla_env_set = True
+
+    if xla_env is None:
+        xla_env_set = False
+    elif isinstance(xla_env, str):
+        xla_env_set = not "--xla_disable_hlo_passes=constant_folding" in xla_env
+
+
+    if not xla_env_set:
+        logger.info(
+            """
+            For fast JAX compile times, the envirment variable XLA_FLAGS must be set to "--xla_disable_hlo_passes=constant_folding",
+            which is currently not.
+            
+            In Python, to do this manually, use the code: 
+            
+            import os
+            os.environ["XLA_FLAGS"] = "--xla_disable_hlo_passes=constant_folding"
+            
+            The environment variable has been set automatically for you now, however if JAX has already been imported, 
+            this change will not take effect and JAX function compiling times may be slow. 
+            
+            Therefore, it is recommended to set this environment variable before running your script, e.g. in your terminal.
+            """)
+
+        os.environ['XLA_FLAGS'] = "--xla_disable_hlo_passes=constant_folding"
+
+    import jax
     from jax import numpy
 
     print(

autofit/mapper/prior/abstract.py

@@ -233,6 +233,7 @@ def from_dict(
         )
         if id_ is not None:
             loaded_ids[id_] = prior
+
         return prior
 
     def dict(self) -> dict: