speed up notebooks on github actions by lower pop size (#501)

* speed up notebooks on github actions by lower pop size

* test sth

* test sth else

* test sth else

* test sth else

* test sth else

* test sth else

* test sth else

* test sth else

* test sth else

* add nbs

* add nbs

* add nbs

* refactor placement

* fix rename

pyabc/inference/smc.py

@@ -167,21 +167,21 @@ class ABCSMC:
         https://doi.org/10.1093/bioinformatics/btp619
     """
     def __init__(
-            self,
-            models: Union[List[Model], Model, Callable],
-            parameter_priors: Union[List[Distribution],
-                                    Distribution, Callable],
-            distance_function: Union[Distance, Callable] = None,
-            population_size: Union[PopulationStrategy, int] = 100,
-            summary_statistics: Callable[[model_output], dict] = identity,
-            model_prior: RV = None,
-            model_perturbation_kernel: ModelPerturbationKernel = None,
-            transitions: Union[List[Transition], Transition] = None,
-            eps: Epsilon = None,
-            sampler: Sampler = None,
-            acceptor: Acceptor = None,
-            stop_if_only_single_model_alive: bool = False,
-            max_nr_recorded_particles: int = np.inf):
+        self,
+        models: Union[List[Model], Model, Callable],
+        parameter_priors: Union[List[Distribution], Distribution, Callable],
+        distance_function: Union[Distance, Callable] = None,
+        population_size: Union[PopulationStrategy, int] = 100,
+        summary_statistics: Callable[[model_output], dict] = identity,
+        model_prior: RV = None,
+        model_perturbation_kernel: ModelPerturbationKernel = None,
+        transitions: Union[List[Transition], Transition] = None,
+        eps: Epsilon = None,
+        sampler: Sampler = None,
+        acceptor: Acceptor = None,
+        stop_if_only_single_model_alive: bool = False,
+        max_nr_recorded_particles: int = np.inf,
+    ):
         if not isinstance(models, list):
             models = [models]
         models = list(map(SimpleModel.assert_model, models))
@@ -216,7 +216,7 @@ def __init__(
                            for _ in self.models]
         if not isinstance(transitions, list):
             transitions = [transitions]
-        self.transitions = transitions  # type: List[Transition]
+        self.transitions: List[Transition] = transitions
 
         if eps is None:
             eps = MedianEpsilon(median_multiplier=1)
@@ -255,9 +255,11 @@ def __init__(
 
     def _sanity_check(self):
         # check stochastic setting
-        stochastics = [isinstance(self.acceptor, StochasticAcceptor),
-                       isinstance(self.eps, TemperatureBase),
-                       isinstance(self.distance_function, StochasticKernel)]
+        stochastics = [
+            isinstance(self.acceptor, StochasticAcceptor),
+            isinstance(self.eps, TemperatureBase),
+            isinstance(self.distance_function, StochasticKernel),
+        ]
         # check if usage is consistent
         if not all(stochastics) and any(stochastics):
             raise ValueError(
@@ -275,12 +277,15 @@ def __getstate__(self):
         del state_red_dict['sampler']
         return state_red_dict
 
-    def new(self, db: str,
-            observed_sum_stat: dict = None,
-            *,
-            gt_model: int = None,
-            gt_par: dict = None,
-            meta_info=None) -> History:
+    def new(
+        self,
+        db: str,
+        observed_sum_stat: dict = None,
+        *,
+        gt_model: int = None,
+        gt_par: dict = None,
+        meta_info=None,
+    ) -> History:
         """
         Make a new ABCSMC run.
 
@@ -372,9 +377,12 @@ def new(self, db: str,
         # contains id generated in store_initial_data
         return self.history
 
-    def load(self, db: str,
-             abc_id: int = 1,
-             observed_sum_stat: dict = None) -> History:
+    def load(
+        self,
+        db: str,
+        abc_id: int = 1,
+        observed_sum_stat: dict = None,
+    ) -> History:
         """
         Load an ABC-SMC run for continuation.
 
@@ -587,12 +595,14 @@ def _create_transition_pdf(self, t: int, transitions):
             model_perturbation_kernel=self.model_perturbation_kernel)
 
     @run_cleanup
-    def run(self,
-            minimum_epsilon: float = None,
-            max_nr_populations: int = np.inf,
-            min_acceptance_rate: float = 0.,
-            max_total_nr_simulations: int = np.inf,
-            max_walltime: timedelta = None) -> History:
+    def run(
+        self,
+        minimum_epsilon: float = None,
+        max_nr_populations: int = np.inf,
+        min_acceptance_rate: float = 0.,
+        max_total_nr_simulations: int = np.inf,
+        max_walltime: timedelta = None,
+    ) -> History:
         """
         Run the ABCSMC model selection until either of the stopping
         criteria is met.

pyabc/populationstrategy.py

@@ -17,10 +17,23 @@
 from pyabc.cv.bootstrap import calc_cv
 from .transition import Transition
 from .transition.predict_population_size import predict_population_size
+from .util import bound_pop_size_from_env
 
 logger = logging.getLogger("ABC.Adaptation")
 
 
+def dec_bound_pop_size_from_env(fun):
+    """Bound population size"""
+
+    def bounded_fun(*args, **kwargs):
+        pop_size: int = fun(*args, **kwargs)
+        # bound population size potentially by environment variable
+        pop_size = bound_pop_size_from_env(pop_size)
+        return pop_size
+
+    return bounded_fun
+
+
 class PopulationStrategy(ABC):
     """
     Strategy to select the sizes of the populations.
@@ -109,6 +122,7 @@ def __init__(
             nr_calibration_particles=nr_calibration_particles)
         self.nr_particles = nr_particles
 
+    @dec_bound_pop_size_from_env
     def __call__(self, t: int = None) -> int:
         if t == -1 and self.nr_calibration_particles is not None:
             return self.nr_calibration_particles
@@ -151,13 +165,14 @@ class AdaptivePopulationSize(PopulationStrategy):
         Number of calibration particles.
 
 
-    .. [#klingerhasenaueradaptive] Klinger, Emmanuel, and Jan Hasenauer.
-            “A Scheme for Adaptive Selection of Population Sizes in "
-            Approximate Bayesian Computation - Sequential Monte Carlo."
-            Computational Methods in Systems Biology, 128-44.
-            Lecture Notes in Computer Science.
-            Springer, Cham, 2017.
-            https://doi.org/10.1007/978-3-319-67471-1_8.
+    .. [#klingerhasenaueradaptive]
+        Klinger, Emmanuel, and Jan Hasenauer.
+        “A Scheme for Adaptive Selection of Population Sizes in "
+        Approximate Bayesian Computation - Sequential Monte Carlo."
+        Computational Methods in Systems Biology, 128-44.
+        Lecture Notes in Computer Science.
+        Springer, Cham, 2017.
+        https://doi.org/10.1007/978-3-319-67471-1_8.
     """
 
     def __init__(self,
@@ -210,6 +225,7 @@ def update(self, transitions: List[Transition],
         logger.info("Change nr particles {} -> {}"
                     .format(reference_nr_part, self.nr_particles))
 
+    @dec_bound_pop_size_from_env
     def __call__(self, t: int = None) -> int:
         if t == -1 and self.nr_calibration_particles is not None:
             return self.nr_calibration_particles
@@ -223,25 +239,28 @@ class ListPopulationSize(PopulationStrategy):
 
     Parameters
     ----------
-    values: List[float]
+    values:
         List of population size values.
         ``values[t]`` is the value for population t.
     nr_calibration_particles:
         Number of calibration particles.
     """
 
-    def __init__(self,
-                 values: Union[List[int], Dict[int, int]],
-                 nr_calibration_particles: int = None):
+    def __init__(
+        self,
+        values: Union[List[int], Dict[int, int]],
+        nr_calibration_particles: int = None,
+    ):
         super().__init__(
             nr_calibration_particles=nr_calibration_particles)
         self.values = values
 
     def get_config(self) -> dict:
         config = super().get_config()
-        config["population_values"] = self.population_values
+        config["values"] = self.values
         return config
 
+    @dec_bound_pop_size_from_env
     def __call__(self, t: int = None) -> int:
         if t == -1 and self.nr_calibration_particles is not None:
             return self.nr_calibration_particles

pyabc/util/__init__.py

@@ -19,6 +19,10 @@
     ParTrafo,
 )
 
+from .test import (
+    bound_pop_size_from_env,
+)
+
 from .log import (
     log_samples,
 )

pyabc/util/test.py

@@ -0,0 +1,32 @@
+"""Util functions for tests."""
+
+import os
+import logging
+
+# maximum population size environment variable
+PYABC_MAX_POP_SIZE = "PYABC_MAX_POP_SIZE"
+
+logger = logging.getLogger("ABC.Util")
+
+
+def bound_pop_size_from_env(pop_size: int):
+    """Bound population size if corresponding environment variable set.
+
+    Parameters
+    ----------
+    pop_size: Intended population size
+
+    Returns
+    -------
+    bounded_pop_size:
+        Minimum of `pop_size` and environment variable `PYABC_MAX_POP_SIZE`.
+    """
+    if PYABC_MAX_POP_SIZE not in os.environ:
+        return pop_size
+    pop_size = min(pop_size, int(os.environ[PYABC_MAX_POP_SIZE]))
+
+    logger.warning(
+        f"Bounding population size to {pop_size} via environment variable "
+        f"{PYABC_MAX_POP_SIZE}")
+
+    return pop_size

test/run_notebooks.sh

@@ -2,14 +2,18 @@
 
 # Run selected notebooks, on error show output and return with error code.
 
+# set environment
+export PYABC_MAX_POP_SIZE=20
 
 # Notebooks to run
 nbs_1=(
   "adaptive_distances" "wasserstein" "conversion_reaction"
   "early_stopping" "model_selection" "noise"
   "parameter_inference" "resuming")
 nbs_2=(
-  "external_simulators" "using_R")
+  "external_simulators" "using_R"
+  "aggregated_distances"
+  "data_plots" "discrete_parameters")
 
 # All notebooks
 nbs_all=("${nbs_1[@]}" "${nbs_2[@]}")