MAINT: rework stopping criteria

Author: mj-will

nessai/samplers/importancesampler.py

@@ -276,6 +276,93 @@ def __getstate__(self):
         return state
 
 
+class StoppingCriterion:
+    """Object for storing information about a stopping criterion.
+
+    Includes the tolerance, current value and whether the tolerance
+    has been reached based on how it should be checked.
+
+    Parameters
+    ----------
+    name : str
+        Name of the stopping criterion
+    check : {"leq", "geq"}
+        Indicates whether to check if the value is less than or equal (leq) or
+        or greater than or equal (geq) than the tolerance.
+    aliases : List[str]
+        List of aliases (alternative names) for the criterion.
+    tolerance : float
+        Tolerance for the criterion.
+    value : Optional[None]
+        Current value. Does not have to be specified.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        check: str,
+        aliases: list,
+        tolerance: float = None,
+        value: float = None,
+    ):
+        if check.lower() not in ["leq", "geq"]:
+            raise ValueError(
+                f"Invalid value for `check`: {check}. "
+                f"Choose from ['leq', 'geq']."
+            )
+        self._name = name
+        self._check = check
+        self._aliases = aliases
+        self._tolerance = tolerance
+        self._value = value
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def check(self) -> str:
+        return self._check
+
+    @property
+    def aliases(self) -> List[str]:
+        return self._aliases
+
+    @property
+    def tolerance(self) -> float:
+        return self._tolerance
+
+    @property
+    def value(self) -> float:
+        return self._value
+
+    def update_tolerance(self, tolerance) -> None:
+        """Update the tolerance"""
+        self._tolerance = tolerance
+
+    def update_value(self, value) -> None:
+        """Update the current value."""
+        self._value = value
+
+    def update_value_from_sampler(self, sampler) -> None:
+        """Udate"""
+        value = getattr(sampler, self.name, None)
+        if value is None:
+            raise RuntimeError(f"{self.name} has not been computed!")
+        self.update_value(value)
+
+    @property
+    def reached_tolerance(self) -> bool:
+        """Indicates if the stopping criterion has been reached"""
+        if self.check == "leq":
+            return self._value <= self._tolerance
+        else:
+            return self._value >= self._tolerance
+
+    def summary(self) -> str:
+        return f"{self.name}: {self.value:.4g} ({self.tolerance:.2g})"
+
+
 class ImportanceNestedSampler(BaseNestedSampler):
     """
 
@@ -326,33 +413,46 @@ class ImportanceNestedSampler(BaseNestedSampler):
         If False, this can help reduce the disk usage.
     """
 
-    stopping_criterion_config = dict(
-        ratio=dict(
-            type="min",
+    stopping_criteria = {
+        "ratio": StoppingCriterion(
+            name="ratio",
+            check="leq",
             aliases=["ratio", "ratio_all"],
+            tolerance=0.0,
+            value=np.inf,
         ),
-        ratio_ns=dict(
-            type="min",
+        "ratio_ns": StoppingCriterion(
+            name="ratio_ns",
+            check="leq",
             aliases=["ratio_ns"],
+            tolerance=0.0,
+            value=np.inf,
         ),
-        Z_err=dict(
-            type="min",
+        "Z_err": StoppingCriterion(
+            name="Z_err",
+            check="leq",
             aliases=["Z_err", "evidence_error"],
+            value=np.inf,
         ),
-        log_dZ=dict(
-            type="min",
+        "log_dZ": StoppingCriterion(
+            name="log_dZ",
+            check="leq",
             aliases=["log_dZ", "log_evidence"],
+            value=np.inf,
         ),
-        ess=dict(
-            type="max",
+        "ess": StoppingCriterion(
+            name="ess",
+            check="geq",
             aliases=["ess"],
+            value=0.0,
         ),
-        fractional_error=dict(
-            type="min",
+        "fractional_error": StoppingCriterion(
+            name="fractional_error",
+            check="leq",
             aliases=["fractional_error"],
+            value=0.0,
         ),
-    )
-    """Dictionary of available stopping criteria and their aliases."""
+    }
 
     def __init__(
         self,
@@ -378,7 +478,7 @@ def __init__(
         min_remove: int = 1,
         max_samples: Optional[int] = None,
         stopping_criterion: str = "ratio",
-        tolerance: float = 0.0,
+        tolerance: Optional[float] = None,
         n_update: Optional[int] = None,
         plot_pool: bool = False,
         plot_level_cdf: bool = False,
@@ -667,16 +767,14 @@ def reached_tolerance(self) -> bool:
         Checks if any or all of the criteria have been met, this depends on the
         value of :code:`check_criteria`.
         """
-        flags = {}
-        for sc, value in self.criterion.items():
-            if self.stopping_criterion_config[sc]["type"] == "min":
-                flags[sc] = value <= self.tolerance[sc]
-            else:
-                flags[sc] = value > self.tolerance[sc]
+        flags = [
+            self.stopping_criteria[name].reached_tolerance
+            for name in self.criteria_to_check
+        ]
         if self._stop_any:
-            return any(flags.values())
+            return any(flags)
         else:
-            return all(flags.values())
+            return all(flags)
 
     @staticmethod
     def add_fields():
@@ -698,41 +796,17 @@ def configure_stopping_criterion(
         else:
             tolerance = [float(tolerance)]
 
-        self.stopping_criterion = []
-        for c in stopping_criterion:
-            for criterion, cfg in self.stopping_criterion_config.items():
-                if c in cfg["aliases"]:
-                    self.stopping_criterion.append(criterion)
-        if len(self.stopping_criterion) != len(stopping_criterion):
-            raise ValueError(
-                f"Unknown stopping criterion: {stopping_criterion}"
-            )
-        for c, c_use in zip(stopping_criterion, self.stopping_criterion):
-            if c != c_use:
-                logger.info(
-                    f"Stopping criterion specified ({c}) is "
-                    f"an alias for {c_use}. Using {c_use}."
-                )
-
-        self.tolerance = {
-            sc: t for sc, t in zip(stopping_criterion, tolerance)
-        }
-        if len(self.stopping_criterion) != len(self.tolerance):
-            raise ValueError(
-                "Number of stopping criteria must match tolerances"
-            )
-
-        types = {
-            sc: self.stopping_criterion_config[sc]["type"]
-            for sc in stopping_criterion
-        }
-        self.criterion = {
-            sc: np.inf if t == "min" else -np.inf for sc, t in types.items()
-        }
-
-        logger.info(f"Stopping criteria: {self.stopping_criterion}")
-        logger.info(f"Tolerance: {self.tolerance}")
+        self.criteria_to_check = []
+        for name, tol in zip(stopping_criterion, tolerance):
+            for sc in self.stopping_criteria.values():
+                if name in sc.aliases:
+                    sc.update_tolerance(tol)
+                    self.criteria_to_check.append(name)
+                    break
+            else:
+                raise ValueError(f"Unknown stopping criterion: {name}")
 
+        logger.info(f"Stopping criteria to check: {self.criteria_to_check}")
         if check_criteria not in {"any", "all"}:
             raise ValueError("check_criteria must be any or all")
         if check_criteria == "any":
@@ -870,7 +944,7 @@ def initialise_history(self) -> None:
                     samples_entropy=[],
                     proposal_entropy=[],
                     stopping_criteria={
-                        k: [] for k in self.stopping_criterion_config.keys()
+                        name: [] for name in self.stopping_criteria.keys()
                     },
                 )
             )
@@ -900,10 +974,8 @@ def update_history(self) -> None:
             self.model.likelihood_evaluations
         )
 
-        for k in self.stopping_criterion_config.keys():
-            self.history["stopping_criteria"][k].append(
-                getattr(self, k, np.nan)
-            )
+        for name, sc in self.stopping_criteria.items():
+            self.history["stopping_criteria"][name].append(sc.value)
 
     def determine_threshold_quantile(
         self,
@@ -1456,11 +1528,18 @@ def compute_stopping_criterion(self) -> List[float]:
         self.ess = self.state.effective_n_posterior_samples
         self.Z_err = np.exp(self.log_evidence_error)
         self.fractional_error = self.state.evidence_error / self.state.evidence
-        cond = {sc: getattr(self, sc) for sc in self.stopping_criterion}
 
-        logger.info(
-            f"Stopping criteria: {cond} " f"- Tolerance: {self.tolerance}"
-        )
+        cond = {}
+        for name, sc in self.stopping_criteria.items():
+            sc.update_value_from_sampler(self)
+            if name in self.criteria_to_check:
+                cond[name] = sc.value
+
+        status = [
+            self.stopping_criteria[sc].summary()
+            for sc in self.criteria_to_check
+        ]
+        logger.info(f"Stopping criteria: {status}")
         return cond
 
     def checkpoint(self, periodic: bool = False, force: bool = False):
@@ -2049,17 +2128,17 @@ def plot_state(
         ax[m].legend()
         m += 1
 
-        for (i, sc), tol in zip(
-            enumerate(self.stopping_criterion), self.tolerance
-        ):
+        for i, sc_name in enumerate(self.criteria_to_check):
             ax[m].plot(
                 its,
-                self.history["stopping_criteria"][sc],
-                label=sc,
+                self.history["stopping_criteria"][sc_name],
+                label=sc_name,
                 c=f"C{i}",
                 ls=config.plotting.line_styles[i],
             )
-            ax[m].axhline(tol, ls=":", c=f"C{i}")
+            ax[m].axhline(
+                self.stopping_criteria[sc_name].tolerance, ls=":", c=f"C{i}"
+            )
         ax[m].legend()
         ax[m].set_ylabel("Stopping criterion")