Feature/build fixes

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",
 ):

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

@@ -4,6 +4,7 @@
 If USE_JAX=true in general.yaml then JAX's NumPy is used, otherwise vanilla NumPy is used.
 """
 import jax
+import os
 
 from autoconf import conf
 

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:

autofit/mapper/prior/vectorized.py

@@ -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
@@ -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

autofit/mapper/prior_model/array.py

@@ -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
@@ -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:

autofit/messages/truncated_normal.py

@@ -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.
         """
@@ -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]: