refactor: search hygiene and dead code (search-extensibility A0b)

autofit/non_linear/parallel/__init__.py

@@ -4,4 +4,3 @@
 from .process import Process
 from .sneaky import SneakyJob
 from .sneaky import SneakyPool
-from .sneaky import SneakierPool

autofit/non_linear/parallel/sneaky.py

@@ -306,136 +306,3 @@ def __del__(self):
                 process.join(0.5)
             except Exception as e:
                 logger.exception(e)
-
-
-class FunctionCache:
-    """
-    Singleton class to cache the functions and optional arguments between calls
-    """
-
-
-def initializer(
-    fitness,
-    prior_transform,
-    fitness_args,
-    fitness_kwargs,
-    prior_transform_args,
-    prior_transform_kwargs,
-):
-    """
-    Initialized function used to initialize the
-    singleton object inside each worker of the pool
-    """
-    FunctionCache.fitness = fitness
-    FunctionCache.prior_transform = prior_transform
-    FunctionCache.fitness_args = fitness_args
-    FunctionCache.fitness_kwargs = fitness_kwargs
-    FunctionCache.prior_transform_args = prior_transform_args
-    FunctionCache.prior_transform_kwargs = prior_transform_kwargs
-
-
-def fitness_cache(x):
-    """
-    Likelihood function call
-    """
-    return FunctionCache.fitness(
-        x, *FunctionCache.fitness_args, **FunctionCache.fitness_kwargs
-    )
-
-
-def prior_transform_cache(x):
-    """
-    Prior transform call
-    """
-
-    return FunctionCache.prior_transform(
-        x, *FunctionCache.prior_transform_args, **FunctionCache.prior_transform_kwargs
-    )
-
-
-class SneakierPool:
-    def __init__(
-        self,
-        processes: int,
-        fitness: Callable,
-        prior_transform: Optional[Callable] = None,
-        fitness_args: Optional[Iterable] = None,
-        fitness_kwargs: Optional[dict] = None,
-        prior_transform_args: Optional[Iterable] = None,
-        prior_transform_kwargs: Optional[dict] = None,
-    ):
-        self.fitness_init = fitness
-        self.prior_transform_init = prior_transform
-        self.fitness = fitness_cache
-        self.prior_transform = prior_transform_cache
-        self.fitness_args = fitness_args or ()
-        self.fitness_kwargs = fitness_kwargs or {}
-        self.prior_transform_args = prior_transform_args or ()
-        self.prior_transform_kwargs = prior_transform_kwargs or {}
-        self.processes = processes
-        self.pool = None
-
-        init_args = (
-            self.fitness_init,
-            self.prior_transform_init,
-            self.fitness_args,
-            self.fitness_kwargs,
-            self.prior_transform_args,
-            self.prior_transform_kwargs,
-        )
-        initializer(*init_args)
-
-    def __enter__(self):
-        """
-        Activate the mp / mpi pool
-        """
-
-        logger.info("... using multiprocessing")
-
-        self.pool = fork_context().Pool(
-            processes=self.processes,
-        )
-
-        return self
-
-    def map(self, function: Callable, iterable: Iterable):
-        """
-        Map a function over an iterable using the map method
-        of the initialized pool.
-
-        Parameters
-        ----------
-        function
-            A function to map
-        iterable
-            An iterable to map over
-
-        """
-        return self.pool.map(function, iterable)
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        try:
-            self.pool.terminate()
-        except:  # noqa
-            pass
-        try:
-            del (
-                FunctionCache.fitness,
-                FunctionCache.prior_transform,
-                FunctionCache.fitness_args,
-                FunctionCache.fitness_kwargs,
-                FunctionCache.prior_transform_args,
-                FunctionCache.prior_transform_kwargs,
-            )
-        except:  # noqa
-            pass
-
-    @property
-    def size(self):
-        return self.njobs
-
-    def close(self):
-        self.pool.close()
-
-    def join(self):
-        self.pool.join()

autofit/non_linear/search/abstract_search.py

@@ -13,7 +13,6 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Optional, Union, Tuple, List, Dict
 
-import psutil
 
 if TYPE_CHECKING:
     from autofit.database.sqlalchemy_ import sa
@@ -38,7 +37,8 @@
 # `ITERATIONS_NEVER` is defined in `fitness.py` (so `Fitness.manage_quick_update` can
 # use it) and re-exported here for existing importers, e.g. `multi_start_gradient`.
 from autofit.non_linear.fitness import ITERATIONS_NEVER  # noqa: F401
-from autofit.non_linear.parallel import SneakyPool, SneakierPool, fork_context
+from autofit.non_linear.parallel import SneakyPool, fork_context
+from autofit.non_linear.search.updater import log_process_state
 from autofit.non_linear.paths.abstract import AbstractPaths
 from autofit.non_linear.paths.database import DatabasePaths
 from autofit.non_linear.paths.directory import DirectoryPaths
@@ -144,6 +144,13 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
     _visualize_fit = True
     _visualize_before_fit = True
 
+    # Keyword arguments the base class expects in ``**kwargs``: ``save_all_samples`` is
+    # read by ``__init__``, and ``initial_values`` / ``inplace`` are the
+    # ``AbstractFactorOptimiser`` arguments ``search.json`` records, so a search loaded
+    # from it passes them back. Any other key reaching the base class is logged as
+    # unknown (a warning, never an error).
+    _known_kwargs = ("save_all_samples", "initial_values", "inplace")
+
     def __init__(
         self,
         name: Optional[str] = None,
@@ -259,15 +266,21 @@ def __init__(
                 "iterations_per_full_update"
             ])
 
-        self.iterations = 0
-
         self.silence = silence
 
         if conf.instance["general"]["hpc"]["hpc_mode"]:
             self.silence = True
 
         self.kwargs = kwargs
 
+        for key in kwargs:
+            if key not in self._known_kwargs:
+                logger.warning(
+                    f"{type(self).__name__} received the unknown keyword argument "
+                    f"{key!r}, which is ignored. Check its spelling against the "
+                    f"search's constructor."
+                )
+
         self.number_of_cores = number_of_cores
 
         if number_of_cores > 1 and any(
@@ -655,7 +668,7 @@ class represented by model M and gives a score for their fitness.
         else:
             logger.info(f"Starting non-linear search with {self.number_of_cores} cores.")
         logger.info(self.quick_update_message)
-        self._log_process_state()
+        log_process_state()
 
         model = analysis.modify_model(model)
         self.paths.model = model
@@ -712,28 +725,6 @@ class represented by model M and gives a score for their fitness.
 
         return result
 
-    @staticmethod
-    def _log_process_state():
-        total_files = 0
-
-        for process in psutil.process_iter(attrs=["pid"]):
-            try:
-                proc_info = process.as_dict(attrs=["pid"])
-                logger.debug(
-                    f"Process ID: {proc_info['pid']} has the following open files:"
-                )
-
-                open_files = process.open_files()
-                for file in open_files:
-                    logger.debug(file)
-                    total_files += 1
-
-            except (psutil.NoSuchProcess, psutil.AccessDenied):
-                pass
-
-        if conf.instance["logging"]["total_files_open"]:
-            logger.info(f"Total Files Open: {total_files}")
-
     def pre_fit_output(
         self, analysis: Analysis, model: AbstractPriorModel, info: Optional[Dict] = None
     ):
@@ -1385,9 +1376,10 @@ def apply_test_mode(self):
         """
         Override in subclasses to reduce sampler iterations for test mode.
 
-        Called during __init__ when test mode is active (level 1).
-        Subclasses should directly mutate instance attributes to minimize
-        the number of iterations the sampler performs.
+        The base class never calls this method: each search that overrides it
+        calls it at the end of its own ``__init__`` when ``is_test_mode()`` is
+        true (test mode level 1). Overrides directly mutate instance attributes
+        to minimize the number of iterations the sampler performs.
         """
         pass
 
@@ -1537,6 +1529,7 @@ def perform_visualization(
         instance: Optional[ModelInstance] = None,
         paths_override: Optional[AbstractPaths] = None,
         search_internal=None,
+        plot_search: bool = True,
     ):
         """
         Perform visualization of the non-linear search's model-fitting results.
@@ -1551,6 +1544,7 @@ def perform_visualization(
             instance=instance,
             paths_override=paths_override,
             search_internal=search_internal,
+            plot_search=plot_search,
         )
 
     @property
@@ -1591,12 +1585,15 @@ def plot_start_point(
         paths = copy.copy(self.paths)
         paths.image_path_suffix = "_start"
 
+        # The search has no samples before it starts, so only the analysis visuals
+        # of the starting point are output, not the search's own plots.
         self.perform_visualization(
             model=model,
             analysis=analysis,
             instance=instance,
             during_analysis=False,
             paths_override=paths,
+            plot_search=False,
         )
 
     def samples_from(self, model: AbstractPriorModel, search_internal=None) -> Samples:
@@ -1620,7 +1617,12 @@ def samples_from(self, model: AbstractPriorModel, search_internal=None) -> Sampl
             return self.samples_via_internal_from(
                 model=model, search_internal=search_internal
             )
-        except (FileNotFoundError, NotImplementedError, AttributeError):
+        except (FileNotFoundError, NotImplementedError) as e:
+            logger.warning(
+                f"The samples of {type(self).__name__} could not be loaded from its "
+                f"internal results ({type(e).__name__}: {e}), so they are loaded from "
+                f"the samples.csv in the output folder instead."
+            )
             return self.paths.samples
 
     def samples_via_internal_from(
@@ -1669,21 +1671,6 @@ def make_sneaky_pool(self, fitness: Fitness) -> Optional[SneakyPool]:
             processes=self.number_of_cores, paths=self.paths, fitness=fitness
         )
 
-    def make_sneakier_pool(self, fitness_function: Fitness, **kwargs) -> SneakierPool:
-
-        self.logger.info(f"number of cores == {self.number_of_cores}")
-
-        if self.number_of_cores > 1:
-            self.logger.info("Creating SneakierPool...")
-        else:
-            self.logger.info("Creating multiprocessing Pool of size 1...")
-
-        pool = SneakierPool(
-            processes=self.number_of_cores, fitness=fitness_function, **kwargs
-        )
-
-        return pool
-
     def __eq__(self, other):
         return isinstance(other, NonLinearSearch) and self.__dict__ == other.__dict__
 

autofit/non_linear/search/mcmc/blackjax/chains.py

@@ -22,6 +22,25 @@
 
 _VALID_STRING_KINDS = ("diagonal", "dense")
 
+# The kinds `inverse_mass_matrix_kind_from` returns for a `None`, array or `Result`
+# specification. `search.json` records the kind string, not the specification, so a
+# search loaded from it receives one of these: "none" is equivalent to `None`, while
+# "array" / "result" name a seed whose values were not recorded and cannot be resolved.
+SERIALISED_KINDS = ("none", "array", "result")
+
+
+def inverse_mass_matrix_spec_from(spec: InverseMassMatrixSpec) -> InverseMassMatrixSpec:
+    """
+    Normalise an ``inverse_mass_matrix`` specification passed to a constructor.
+
+    The kind string ``"none"`` that ``search.json`` records for a ``None`` specification is mapped back to
+    ``None``, so a search loaded from ``search.json`` behaves as the one that wrote it. Every other value is
+    returned unchanged.
+    """
+    if isinstance(spec, str) and spec == "none":
+        return None
+    return spec
+
 
 def inverse_mass_matrix_kind_from(spec: InverseMassMatrixSpec) -> str:
     """
@@ -42,7 +61,7 @@ def inverse_mass_matrix_kind_from(spec: InverseMassMatrixSpec) -> str:
         return "none"
 
     if isinstance(spec, str):
-        if spec not in _VALID_STRING_KINDS:
+        if spec not in _VALID_STRING_KINDS + SERIALISED_KINDS:
             raise ValueError(
                 f"inverse_mass_matrix string must be one of {_VALID_STRING_KINDS}, got {spec!r}"
             )
@@ -96,10 +115,16 @@ def resolve_inverse_mass_matrix(
         return True, None
 
     if isinstance(spec, str):
-        if spec == "diagonal":
+        if spec in ("diagonal", "none"):
             return True, None
         if spec == "dense":
             return False, None
+        if spec in SERIALISED_KINDS:
+            raise ValueError(
+                f"inverse_mass_matrix={spec!r} is the kind recorded in search.json for a seed whose values "
+                "were not saved, so this search (loaded from search.json) cannot be fitted with it. Pass "
+                "the original seed array or Result as inverse_mass_matrix instead."
+            )
         raise ValueError(
             f"inverse_mass_matrix string must be one of {_VALID_STRING_KINDS}, got {spec!r}"
         )