Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
35 commits
Select commit Hold shift + click to select a range
48849a4
Merge branch 'feature/declarative_deterministic'
Oct 6, 2025
85bc594
fix delcarative thing brekaing unit tests
Oct 6, 2025
5038d4f
merge
Oct 6, 2025
e782a43
reinstate plot module
Oct 7, 2025
15dde06
fix imports
Oct 7, 2025
4293d90
try except on latent variable PDF drawingh
Oct 7, 2025
93ac19e
add Truncated Gaussian to from dict list
Oct 7, 2025
c19ca06
af.Gaussian -> af.ex.Gaussian
Oct 7, 2025
7c5e38b
fix TruncatedGaussian ids
Oct 7, 2025
9d279d1
Extend fitness API for better JAX support
Oct 8, 2025
02e576e
update fitness for more flexibility
Oct 8, 2025
74a920a
remove MPI
Oct 8, 2025
33b9871
all tests pass
Oct 8, 2025
5c347e8
err
Oct 8, 2025
b64291d
remove prints
Oct 8, 2025
5c6d71a
more fixes
Oct 14, 2025
e104303
add model centred to autofit result
Oct 16, 2025
d5d939a
'Updated version in __init__ to 2025.10.16.1
rhayes777 Oct 16, 2025
89d6774
'Updated version in __init__ to 2025.10.16.2
rhayes777 Oct 16, 2025
c67705e
quick update visualization works nice
Oct 19, 2025
2a5f5d7
update text_util to be used for quick update
Oct 19, 2025
d92f4ea
implementation seems to work, now to fix unit tests
Oct 19, 2025
14f58ac
fremove Maximum Log Posterior
Oct 19, 2025
950b333
ierations per update i configs
Oct 19, 2025
44dabba
remove analysos n cores
Oct 19, 2025
cc1a2a1
more iterations update fixes
Oct 19, 2025
0202b6e
add unit tests pass
Oct 19, 2025
fb76039
remove print
Oct 19, 2025
941d405
Merge pull request #1159 from rhayes777/feature/quick_update
Jammy2211 Oct 19, 2025
ddf4916
minor edits
Oct 19, 2025
99dcbd3
float conversion
Oct 19, 2025
eb44458
simpify update message
Oct 19, 2025
416bf1a
disable graph visualization
Oct 19, 2025
00087f7
enviroment variable JAX fast time check
Oct 20, 2025
13abefd
final fixes
Oct 20, 2025
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions autofit/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -129,6 +128,7 @@
"LogUniform",
"Gaussian",
"LogGaussian",
"TruncatedGaussian",
"compound",
"Constant",
):
Expand All @@ -142,4 +142,4 @@ def save_abc(pickler, obj):



__version__ = "2025.5.10.1"
__version__ = "2025.10.16.2"
8 changes: 5 additions & 3 deletions autofit/config/general.yaml
Original file line number Diff line number Diff line change
@@ -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.
Expand Down
7 changes: 2 additions & 5 deletions autofit/config/non_linear/mcmc.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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).
20 changes: 10 additions & 10 deletions autofit/config/non_linear/mle.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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).
12 changes: 0 additions & 12 deletions autofit/config/non_linear/nest.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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).
12 changes: 6 additions & 6 deletions autofit/example/analysis.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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:
Expand Down
21 changes: 11 additions & 10 deletions autofit/example/model.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion autofit/fixtures.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
def make_model_gaussian_x1():

return af.Model(
af.Gaussian
af.ex.Gaussian
)


Expand Down
38 changes: 36 additions & 2 deletions autofit/jax_wrapper.py
Original file line number Diff line number Diff line change
@@ -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(
Expand Down
1 change: 1 addition & 0 deletions autofit/mapper/prior/abstract.py
Original file line number Diff line number Diff line change
Expand Up @@ -233,6 +233,7 @@ def from_dict(
)
if id_ is not None:
loaded_ids[id_] = prior

return prior

def dict(self) -> dict:
Expand Down
11 changes: 10 additions & 1 deletion autofit/mapper/prior/vectorized.py
Original file line number Diff line number Diff line change
Expand Up @@ -135,7 +135,12 @@ def __call__(self, cube: np.ndarray) -> np.ndarray:
Transformed parameters of shape (n_samples, n_priors).
"""

n_samples, n_priors = cube.shape
cube_reshaped = False

if len(cube.shape) == 1:
cube = cube[None, :]
cube_reshaped = True

out = np.empty_like(cube)

# 2) Batch‐process all UniformPriors
Expand Down Expand Up @@ -177,4 +182,8 @@ def __call__(self, cube: np.ndarray) -> np.ndarray:
+ subcube * (self.loguniform_log_uppers - self.loguniform_log_lowers)
)

if cube_reshaped:

return out[0]

return out
4 changes: 2 additions & 2 deletions autofit/mapper/prior_model/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
from autoconf.dictable import from_dict
from .abstract import AbstractPriorModel
from autofit.mapper.prior.abstract import Prior
from autofit.jax_wrapper import numpy as jnp, use_jax
from autofit.jax_wrapper import numpy as xp, use_jax
import numpy as np

from autofit.jax_wrapper import register_pytree_node_class
Expand Down Expand Up @@ -77,7 +77,7 @@ def _instance_for_arguments(
-------
The array with the priors replaced.
"""
array = jnp.zeros(self.shape)
array = xp.zeros(self.shape)
for index in self.indices:
value = self[index]
try:
Expand Down
4 changes: 1 addition & 3 deletions autofit/messages/truncated_normal.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,14 +74,11 @@ def __init__(
----------
mean
The mean (μ) of the normal distribution.

sigma
The standard deviation (σ) of the distribution. Must be non-negative.

log_norm
An additive constant to the log probability of the message. Used internally for message-passing normalization.
Default is 0.0.

id_
An optional unique identifier used to track the message in larger probabilistic graphs or models.
"""
Expand All @@ -96,6 +93,7 @@ def __init__(
log_norm=log_norm,
id_=id_,
)

self.mean, self.sigma, self.lower_limit, self.upper_limit = self.parameters

def cdf(self, x: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
Expand Down
Loading
Loading