Avoid SamplerReport methods in convergence checks

The goal was to uncouple sampling functions
from `MultiTrace` and `SamplerReport`.

Some calls to `SamplerReport._log_summary()` were unnecessary because
`MultiTrace._add_warnings()` was never called inbetween instantiation
and `_log_summary()`, therefore the traces never contained warnings.

Running convergence checks and logging the warnings can also be done
without needing `MultiTrace` or `SamplerReport` instances/methods.

Author: michaelosthege

pymc/backends/report.py

@@ -17,14 +17,7 @@
 
 from typing import Dict, List, Optional
 
-import arviz
-
-from pymc.stats.convergence import (
-    _LEVELS,
-    SamplerWarning,
-    log_warnings,
-    run_convergence_checks,
-)
+from pymc.stats.convergence import _LEVELS, SamplerWarning
 
 logger = logging.getLogger("pymc")
 
@@ -73,22 +66,13 @@ def raise_ok(self, level="error"):
         if errors:
             raise ValueError("Serious convergence issues during sampling.")
 
-    def _run_convergence_checks(self, idata: arviz.InferenceData, model):
-        warnings = run_convergence_checks(idata, model)
-        self._add_warnings(warnings)
-
     def _add_warnings(self, warnings, chain=None):
         if chain is None:
             warn_list = self._global_warnings
         else:
             warn_list = self._chain_warnings.setdefault(chain, [])
         warn_list.extend(warnings)
 
-    def _log_summary(self):
-        for chain, warns in self._chain_warnings.items():
-            log_warnings(warns)
-        log_warnings(self._global_warnings)
-
     def _slice(self, start, stop, step):
         report = SamplerReport()
 

pymc/sampling/mcmc.py

@@ -40,7 +40,11 @@
 from pymc.model import Model, modelcontext
 from pymc.sampling.parallel import Draw, _cpu_count
 from pymc.sampling.population import _sample_population
-from pymc.stats.convergence import log_warning_stats, run_convergence_checks
+from pymc.stats.convergence import (
+    log_warning_stats,
+    log_warnings,
+    run_convergence_checks,
+)
 from pymc.step_methods import NUTS, CompoundStep
 from pymc.step_methods.arraystep import BlockedStep, PopulationArrayStepShared
 from pymc.step_methods.hmc import quadpotential
@@ -602,7 +606,6 @@ def sample(
         f"({n_tune*n_chains:_d} + {n_draws*n_chains:_d} draws total) "
         f"took {t_sampling:.0f} seconds."
     )
-    mtrace.report._log_summary()
 
     idata = None
     if compute_convergence_checks or return_inferencedata:
@@ -612,14 +615,9 @@ def sample(
         idata = pm.to_inference_data(mtrace, **ikwargs)
 
         if compute_convergence_checks:
-            if draws - tune < 100:
-                warnings.warn(
-                    "The number of samples is too small to check convergence reliably.",
-                    stacklevel=2,
-                )
-            else:
-                convergence_warnings = run_convergence_checks(idata, model)
-                mtrace.report._add_warnings(convergence_warnings)
+            warns = run_convergence_checks(idata, model)
+            mtrace.report._add_warnings(warns)
+            log_warnings(warns)
 
         if return_inferencedata:
             # By default we drop the "warning" stat which contains `SamplerWarning`
@@ -925,9 +923,6 @@ def _mp_sample(
             strace = traces[error._chain]
             for strace in traces:
                 strace.close()
-
-            multitrace = MultiTrace(traces)
-            multitrace._report._log_summary()
             raise
     except KeyboardInterrupt:
         pass

pymc/smc/sampling.py

@@ -34,6 +34,7 @@
 from pymc.model import Model, modelcontext
 from pymc.sampling.parallel import _cpu_count
 from pymc.smc.kernels import IMH
+from pymc.stats.convergence import log_warnings, run_convergence_checks
 from pymc.util import RandomState, _get_seeds_per_chain
 
 
@@ -237,7 +238,11 @@ def sample_smc(
     )
 
     if compute_convergence_checks:
-        _compute_convergence_checks(idata, draws, model, trace)
+        if idata is None:
+            idata = to_inference_data(trace, log_likelihood=False)
+        warns = run_convergence_checks(idata, model)
+        trace.report._add_warnings(warns)
+        log_warnings(warns)
 
     if return_inferencedata:
         assert idata is not None
@@ -298,21 +303,6 @@ def _save_sample_stats(
     return sample_stats, idata
 
 
-def _compute_convergence_checks(
-    idata: Optional[InferenceData], draws: int, model: Model, trace: MultiTrace
-):
-    if draws < 100:
-        warnings.warn(
-            "The number of samples is too small to check convergence reliably.",
-            stacklevel=2,
-        )
-    else:
-        if idata is None:
-            idata = to_inference_data(trace, log_likelihood=False)
-        trace.report._run_convergence_checks(idata, model)
-    trace.report._log_summary()
-
-
 def _sample_smc_int(
     draws,
     kernel,

pymc/stats/convergence.py

@@ -53,10 +53,13 @@ def run_convergence_checks(idata: arviz.InferenceData, model) -> List[SamplerWar
         warn = SamplerWarning(WarningType.BAD_PARAMS, msg, "info", None, None, None)
         return [warn]
 
+    if idata["posterior"].sizes["draw"] < 100:
+        msg = "The number of samples is too small to check convergence reliably."
+        warn = SamplerWarning(WarningType.BAD_PARAMS, msg, "info", None, None, None)
+        return [warn]
+
     if idata["posterior"].sizes["chain"] == 1:
-        msg = (
-            "Only one chain was sampled, this makes it impossible to " "run some convergence checks"
-        )
+        msg = "Only one chain was sampled, this makes it impossible to run some convergence checks"
         warn = SamplerWarning(WarningType.BAD_PARAMS, msg, "info")
         return [warn]
 

pymc/tests/smc/test_smc.py

@@ -11,6 +11,7 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
+import logging
 import warnings
 
 import numpy as np
@@ -215,13 +216,11 @@ def test_return_datatype(self, chains):
         assert mt.nchains == chains
         assert mt["x"].size == chains * draws
 
-    def test_convergence_checks(self):
-        with self.fast_model:
-            with pytest.warns(
-                UserWarning,
-                match="The number of samples is too small",
-            ):
+    def test_convergence_checks(self, caplog):
+        with caplog.at_level(logging.INFO):
+            with self.fast_model:
                 pm.sample_smc(draws=99)
+        assert "The number of samples is too small" in caplog.text
 
     def test_deprecated_parallel_arg(self):
         with self.fast_model: