minor changes

Author: Jammy2211

autofit/mapper/prior/abstract.py

@@ -5,7 +5,6 @@
 from typing import Union, Tuple, Optional, Dict
 
 from autoconf import conf
-from autoconf import jax_wrapper
 
 from autofit.mapper.prior.arithmetic import ArithmeticMixin
 from autofit.mapper.prior.constant import Constant

autofit/mapper/prior_model/abstract.py

@@ -1,7 +1,5 @@
 import copy
 import inspect
-import jax.numpy as jnp
-import jax
 import json
 import logging
 import random

autofit/non_linear/analysis/analysis.py

@@ -58,7 +58,7 @@ def method(*args, **kwargs):
 
         return method
 
-    def compute_latent_samples(self, samples: Samples) -> Optional[Samples]:
+    def compute_latent_samples(self, samples: Samples, batch_size : Optional[int] = None) -> Optional[Samples]:
         """
         Compute latent variables from a model instance.
 
@@ -91,11 +91,19 @@ def compute_latent_samples(self, samples: Samples) -> Optional[Samples]:
             `(intensity_total, magnitude, angle)`. Each entry may be NaN if the corresponding component
             of the model is not present.
         """
+
+        if use_jax:
+            xp = jnp
+        else:
+            xp = np
+
+        batch_size = batch_size or 10
+
         try:
 
             start_latent = time.time()
 
-            compute_latent_for_model = functools.partial(self.compute_latent_variables, model=samples.model)
+            compute_latent_for_model = functools.partial(self.compute_latent_variables, model=samples.model, xp=xp)
 
             if use_jax:
                 start = time.time()
@@ -107,7 +115,6 @@ def batched_compute_latent(x):
                     return np.array([compute_latent_for_model(xx) for xx in x])
 
             parameter_array = np.array(samples.parameter_lists)
-            batch_size = 50
             latent_samples = []
 
             # process in batches

autofit/non_linear/fitness.py

@@ -43,6 +43,7 @@ def __init__(
         convert_to_chi_squared: bool = False,
         store_history: bool = False,
         use_jax_vmap : bool = False,
+        batch_size : Optional[int] = None,
         iterations_per_quick_update: Optional[int] = None,
     ):
         """
@@ -123,6 +124,7 @@ def __init__(
             if self.use_jax_vmap:
                 self._call = self._vmap
 
+        self.batch_size = batch_size
         self.iterations_per_quick_update = iterations_per_quick_update
         self.quick_update_max_lh_parameters = None
         self.quick_update_max_lh = -xp.inf
@@ -152,7 +154,7 @@ def call(self, parameters):
         instance = self.model.instance_from_vector(vector=parameters)
 
         # Evaluate log likelihood (must be side-effect free and exception-free)
-        log_likelihood = self.analysis.log_likelihood_function(instance=instance)
+        log_likelihood = self.analysis.log_likelihood_function(instance=instance, xp=xp)
 
         # Penalize NaNs in the log-likelihood
         log_likelihood = xp.where(xp.isnan(log_likelihood), self.resample_figure_of_merit, log_likelihood)

autofit/non_linear/search/abstract_search.py

@@ -939,7 +939,10 @@ def perform_update(
 
                 latent_samples = samples_save
 
-            latent_samples = analysis.compute_latent_samples(latent_samples)
+            latent_samples = analysis.compute_latent_samples(
+                latent_samples,
+                batch_size=fitness.batch_size
+            )
 
             if latent_samples:
                 if not conf.instance["output"]["latent_draw_via_pdf"]:

autofit/non_linear/search/nest/nautilus/search.py

@@ -139,6 +139,7 @@ def _fit(self, model: AbstractPriorModel, analysis):
                 fom_is_log_likelihood=True,
                 resample_figure_of_merit=-1.0e99,
                 use_jax_vmap=True,
+                batch_size=self.config_dict_search["n_batch"],
                 iterations_per_quick_update=self.iterations_per_quick_update
 
             )