remove MPI

Author: Jammy2211

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
 
@@ -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/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/non_linear/fitness.py

@@ -10,7 +10,7 @@
 from autoconf import cached_property
 
 from autofit import jax_wrapper
-from autofit.jax_wrapper import numpy as np
+from autofit.jax_wrapper import numpy as xp
 from autofit import exc
 
 
@@ -36,10 +36,10 @@ def __init__(
         analysis : Analysis,
         paths : Optional[AbstractPaths] = None,
         fom_is_log_likelihood: bool = True,
-        resample_figure_of_merit: float = -np.inf,
+        resample_figure_of_merit: float = -xp.inf,
         convert_to_chi_squared: bool = False,
         store_history: bool = False,
-        use_jax_vmap : bool =  True
+        use_jax_vmap : bool = False
     ):
         """
         Interfaces with any non-linear search to fit the model to the data and return a log likelihood via
@@ -147,15 +147,15 @@ def call(self, parameters):
         log_likelihood = self.analysis.log_likelihood_function(instance=instance)
 
         # Penalize NaNs in the log-likelihood
-        log_likelihood = np.where(np.isnan(log_likelihood), self.resample_figure_of_merit, log_likelihood)
+        log_likelihood = xp.where(xp.isnan(log_likelihood), self.resample_figure_of_merit, log_likelihood)
 
         # Determine final figure of merit
         if self.fom_is_log_likelihood:
             figure_of_merit = log_likelihood
         else:
             # Ensure prior list is compatible with JAX (must return a JAX array, not list)
-            log_prior_array = np.array(self.model.log_prior_list_from_vector(vector=parameters))
-            figure_of_merit = log_likelihood + np.sum(log_prior_array)
+            log_prior_array = xp.array(self.model.log_prior_list_from_vector(vector=parameters))
+            figure_of_merit = log_likelihood + xp.sum(log_prior_array)
 
         # Convert to chi-squared scale if requested
         if self.convert_to_chi_squared:
@@ -194,8 +194,8 @@ def call_wrap(self, parameters):
         if self.fom_is_log_likelihood:
             log_likelihood = figure_of_merit
         else:
-            log_prior_list = np.array(self.model.log_prior_list_from_vector(vector=parameters))
-            log_likelihood = figure_of_merit - np.sum(log_prior_list)
+            log_prior_list = xp.array(self.model.log_prior_list_from_vector(vector=parameters))
+            log_likelihood = figure_of_merit - xp.sum(log_prior_list)
 
         if self.store_history:
 
@@ -322,6 +322,7 @@ def check_log_likelihood(self, fitness):
         result
             The result containing the maximum log likelihood fit of the model.
         """
+        import numpy as np
 
         if os.environ.get("PYAUTOFIT_TEST_MODE") == "1":
             return

autofit/non_linear/parallel/sneaky.py

@@ -373,13 +373,6 @@ def __init__(
         self.prior_transform_kwargs = prior_transform_kwargs or {}
         self.processes = processes
         self.pool = None
-        try:
-            from mpi4py import MPI
-
-            self.comm = MPI.COMM_WORLD
-            self._processes = self.comm.size
-        except ModuleNotFoundError:
-            self._processes = 1
 
         init_args = (
             self.fitness_init,
@@ -391,46 +384,16 @@ def __init__(
         )
         initializer(*init_args)
 
-    def check_if_mpi(self):
-        return self._processes > 1
-
-    def is_master(self):
-        is_mpi = self.check_if_mpi()
-        if is_mpi:
-            return_value = self.comm.rank == 0
-        else:
-            return_value = True
-
-        return return_value
-
-    def wait(self):
-        is_mpi = self.check_if_mpi()
-        if is_mpi:
-            self.pool.wait()
-        else:
-            pass
-            warnings.warn("Cannot wait for pool to finish if not using MPI")
-
     def __enter__(self):
         """
         Activate the mp / mpi pool
         """
 
-        use_mpi = self.check_if_mpi()
-
-        if use_mpi:
-            from schwimmbad import MPIPool
-
-            if self.is_master():
-                logger.info("... using Schwimmbad MPIPool")
-            self.pool = MPIPool(use_dill=True)
+        logger.info("... using multiprocessing")
 
-        else:
-            if self.is_master():
-                logger.info("... using multiprocessing")
-            self.pool = mp.Pool(
-                processes=self.processes,
-            )
+        self.pool = mp.Pool(
+            processes=self.processes,
+        )
 
         return self