spawn parallel

Author: Jammy2211

autofit/non_linear/fitness.py

@@ -1,3 +1,4 @@
+import numpy as np_excplicit
 import logging
 import os
 from typing import Optional
@@ -188,10 +189,17 @@ def __setstate__(self, state):
 
     @cached_property
     def _call(self):
-        logger.info("Compiling fitness function for JAX...")
-        debug.print("aaa")
+        debug.print(str(jax_wrapper.use_jax))
+        debug.print("Compiling fitness function for JAX...")
         return jax_wrapper.jit(self.call)
 
+    def call_numpy_wrapper(self, parameters):
+
+        figure_of_merit = self.__call__(parameters=np_excplicit.array(parameters))
+
+        return figure_of_merit.item()
+
+
     @cached_property
     def _grad(self):
         return jax_wrapper.grad(self._call)

autofit/non_linear/search/abstract_search.py

@@ -3,6 +3,7 @@
 import gc
 import logging
 import multiprocessing as mp
+import numpy as np
 import os
 import time
 import warnings
@@ -31,7 +32,6 @@
 )
 from autofit.graphical.utils import Status
 from autofit.mapper.prior_model.abstract import AbstractPriorModel
-from autofit.mapper.prior_model.collection import Collection
 from autofit.mapper.model import ModelInstance
 from autofit.non_linear.initializer import Initializer
 from autofit.non_linear.fitness import Fitness
@@ -995,7 +995,11 @@ def perform_update(
                 parameters = samples.max_log_likelihood(as_instance=False)
 
                 start = time.time()
-                fitness(parameters)
+                figure_of_merit = fitness(parameters)
+
+                # account for asynchronous JAX calls
+                np.array(figure_of_merit)
+
                 log_likelihood_function_time = time.time() - start
 
                 if jax_wrapper.use_jax:

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

@@ -19,12 +19,22 @@
 
 logger = logging.getLogger(__name__)
 
+import time
 
 def prior_transform(cube, model):
     return model.vector_from_unit_vector(unit_vector=cube)
 
 def prior_transform_vectorized(cube, model):
-    return np.array([model.vector_from_unit_vector(row) for row in cube])
+
+    start = time.time()
+
+    trans = np.array([model.vector_from_unit_vector(row) for row in cube])
+
+    end = time.time()
+
+    print(f"Time taken for prior transform vectorized: {end - start:.4f} seconds")
+
+    return trans
 
 class Nautilus(abstract_nest.AbstractNest):
     __identifier_fields__ = (
@@ -141,7 +151,7 @@ def _fit(self, model: AbstractPriorModel, analysis):
         if (
             self.config_dict.get("force_x1_cpu")
             or self.kwargs.get("force_x1_cpu")
-            or jax_wrapper.use_jax
+      #      or jax_wrapper.use_jax
         ):
             search_internal = self.fit_x1_cpu(
                 fitness=fitness,
@@ -227,7 +237,7 @@ def fit_x1_cpu(self, fitness, model, analysis):
             func = jax.vmap(fitness)
             prior_t = prior_transform_vectorized
         else:
-            func = fitness.__call__
+            func = fitness.call_numpy_wrapper
             prior_t = prior_transform
 
         search_internal = self.sampler_cls(
@@ -264,12 +274,16 @@ def fit_multiprocessing(self, fitness, model, analysis):
             the log likelihood the search maximizes.
         """
 
+        # from dask.distributed import Client
+        # client = Client(processes=True, n_workers=4, threads_per_worker=1)
+
         search_internal = self.sampler_cls(
             prior=prior_transform,
-            likelihood=fitness.__call__,
+            likelihood=fitness.call_numpy_wrapper,
             n_dim=model.prior_count,
             prior_kwargs={"model": model},
             filepath=self.checkpoint_file,
+          #  pool=client,
             pool=self.number_of_cores,
             **self.config_dict_search,
         )