implement fitting sequence

jaxtronomy/Data/image_noise.py

@@ -49,7 +49,7 @@ def __init__(
                 )
         else:
             # make sure no negative exposure values are present no dividing by zero
-            self.exp_map = jnp.where(
+            self.exposure_map = jnp.where(
                 exposure_time <= 10 ** (-10), 10 ** (-10), exposure_time
             )
 
@@ -88,7 +88,7 @@ def __init__(
             self.C_D = covariance_matrix(
                 self.data,
                 self.background_rms,
-                self.exp_map,
+                self.exposure_map,
             )
 
         if gradient_boost_factor is not None:
@@ -107,7 +107,7 @@ def C_D_model(self, model):
         if self._noise_map is not None:
             return self._noise_map**2
         else:
-            return covariance_matrix(model, self.background_rms, self.exp_map)
+            return covariance_matrix(model, self.background_rms, self.exposure_map)
 
 
 @export

jaxtronomy/Data/imaging_data.py

@@ -153,7 +153,7 @@ def update_data(self, image_data):
             self.C_D = covariance_matrix(
                 self.data,
                 self.background_rms,
-                self.exp_map,
+                self.exposure_map,
             )
 
     def _log_likelihood(self, model, mask, additional_error_map=0):

jaxtronomy/Sampling/likelihood.py

@@ -10,6 +10,7 @@
 import jax
 from jax import jit, lax, numpy as jnp
 from functools import partial
+import numpy as np
 
 __all__ = ["LikelihoodModule"]
 
@@ -32,7 +33,7 @@ def __init__(
         kwargs_model,
         param_class,
         image_likelihood=True,
-        check_bounds=False,
+        check_bounds=True,
         astrometric_likelihood=False,
         image_position_likelihood=False,
         source_position_likelihood=False,
@@ -417,8 +418,10 @@ def effective_num_data_points(self, **kwargs):
         num_param, param_names = self.param.num_param()
         return self.num_data - num_param
 
+    # This function should be used to convert the jax type to a normal float
+    # Required for samplers e.g. Cobaya which do not work with jax types
     def likelihood(self, a):
-        return self.logL(a)
+        return np.float64(self.logL(a))
 
     def negativelogL(self, a):
         """For minimizer function, the negative value of the logl value is requested.