Change SamplerWarning.kind to enum

Author: aseyboldt

pymc3/backends/base.py

@@ -464,7 +464,7 @@ def get_sampler_stats(self, varname, burn=0, thin=1, combine=True,
         return _squeeze_cat(results, combine, squeeze)
 
     def _slice(self, slice):
-        """Return a new MultiTrace object sliced according to `idx`."""
+        """Return a new MultiTrace object sliced according to `slice`."""
         new_traces = [trace._slice(slice) for trace in self._straces.values()]
         trace = MultiTrace(new_traces)
         idxs = slice.indices(len(self))

pymc3/backends/report.py

@@ -1,10 +1,25 @@
 from collections import namedtuple
 import logging
+import enum
 
 
 logger = logging.getLogger('pymc3')
 
 
+@enum.unique
+class WarningType(enum.Enum):
+    # For HMC and NUTS
+    DIVERGENCE = 1
+    TUNING_DIVERGENCE = 2
+    DIVERGENCES = 3
+    TREEDEPTH = 4
+    # Problematic sampler parameters
+    BAD_PARAMS = 5
+    # Indications that chains did not converge, eg Rhat
+    CONVERGENCE = 6
+    BAD_ACCEPTANCE = 7
+
+
 SamplerWarning = namedtuple(
     'SamplerWarning',
     "kind, message, level, step, exec_info, extra")
@@ -47,7 +62,8 @@ def _run_convergence_checks(self, trace):
         if trace.nchains == 1:
             msg = ("Only one chain was sampled, this makes it impossible to "
                    "run some convergence checks")
-            warn = SamplerWarning('bad-params', msg, 'info', None, None, None)
+            warn = SamplerWarning(WarningType.BAD_PARAMS, msg, 'info',
+                                  None, None, None)
             self._add_warnings([warn])
             return
 
@@ -62,34 +78,35 @@ def _run_convergence_checks(self, trace):
             msg = ("The gelman-rubin statistic is larger than 1.4 for some "
                    "parameters. The sampler did not converge.")
             warn = SamplerWarning(
-                'convergence', msg, 'error', None, None, gelman_rubin)
+                WarningType.CONVERGENCE, msg, 'error', None, None, gelman_rubin)
             warnings.append(warn)
         elif rhat_max > 1.2:
             msg = ("The gelman-rubin statistic is larger than 1.2 for some "
                    "parameters.")
             warn = SamplerWarning(
-                'convergence', msg, 'warn', None, None, gelman_rubin)
+                WarningType.CONVERGENCE, msg, 'warn', None, None, gelman_rubin)
             warnings.append(warn)
         elif rhat_max > 1.05:
             msg = ("The gelman-rubin statistic is larger than 1.05 for some "
                    "parameters. This indicates slight problems during "
                    "sampling.")
             warn = SamplerWarning(
-                'convergence', msg, 'info', None, None, gelman_rubin)
+                WarningType.CONVERGENCE, msg, 'info', None, None, gelman_rubin)
             warnings.append(warn)
 
         eff_min = min(val.min() for val in effective_n.values())
-        if eff_min < 100:
+        n_samples = len(trace) * trace.nchains
+        if eff_min < 200 and n_samples >= 500:
             msg = ("The estimated number of effective samples is smaller than "
-                   "100 for some parameters.")
+                   "200 for some parameters.")
             warn = SamplerWarning(
-                'convergence', msg, 'error', None, None, effective_n)
+                WarningType.CONVERGENCE, msg, 'error', None, None, effective_n)
             warnings.append(warn)
-        elif eff_min < 300:
-            msg = ("The estimated number of effective samples is smaller than "
-                   "300 for some parameters.")
+        elif eff_min / n_samples < 0.25:
+            msg = ("The number of effective samples is smaller than "
+                   "25% for some parameters.")
             warn = SamplerWarning(
-                'convergence', msg, 'warn', None, None, effective_n)
+                WarningType.CONVERGENCE, msg, 'warn', None, None, effective_n)
             warnings.append(warn)
 
         self._add_warnings(warnings)

pymc3/sampling.py

@@ -18,7 +18,6 @@
 from .util import update_start_vals
 from .vartypes import discrete_types
 from pymc3.step_methods.hmc import quadpotential
-from pymc3.backends.report import SamplerWarning
 from pymc3 import plots
 import pymc3 as pm
 from tqdm import tqdm

pymc3/step_methods/hmc/base_hmc.py

@@ -9,7 +9,7 @@
 from pymc3.tuning import guess_scaling
 from .quadpotential import quad_potential, QuadPotentialDiagAdapt
 from pymc3.step_methods import step_sizes
-from pymc3.backends.report import SamplerWarning
+from pymc3.backends.report import SamplerWarning, WarningType
 
 
 HMCStepData = namedtuple(
@@ -107,7 +107,7 @@ def _hamiltonian_step(self, start, p0, step_size):
         raise NotImplementedError("Abstract method")
 
     def astep(self, q0):
-        """Perform a single NUTS iteration."""
+        """Perform a single HMC iteration."""
         p0 = self.potential.random()
         start = self.integrator.compute_state(q0, p0)
 
@@ -130,10 +130,10 @@ def astep(self, q0):
         if hmc_step.divergence_info:
             info = hmc_step.divergence_info
             if self.tune:
-                kind = 'tuning-divergence'
+                kind = WarningType.TUNING_DIVERGENCE
                 point = None
             else:
-                kind = 'divergence'
+                kind = WarningType.DIVERGENCE
                 self._num_divs_sample += 1
                 # We don't want to fill up all memory with divergence info
                 if self._num_divs_sample < 100:
@@ -174,25 +174,25 @@ def warnings(self, strace):
             msg = ('The chain contains only diverging samples. The model is '
                    'probably misspecified.')
             warning = SamplerWarning(
-                'divergences', msg, 'error', None, None, None)
+                WarningType.DIVERGENCES, msg, 'error', None, None, None)
             warnings.append(warning)
         elif n_divs > 0:
             message = ('Divergences after tuning. Increase `target_accept` or '
                        'reparameterize.')
             warning = SamplerWarning(
-                'divergences', message, 'error', None, None, None)
+                WarningType.DIVERGENCES, message, 'error', None, None, None)
             warnings.append(warning)
 
         # small trace
         if self._samples_after_tune == 0:
             msg = "Tuning was enabled throughout the whole trace."
             warning = SamplerWarning(
-                'bad-params', msg, 'error', None, None, None)
+                WarningType.BAD_PARAMS, msg, 'error', None, None, None)
             warnings.append(warning)
         elif self._samples_after_tune < 500:
             msg = "Only %s samples in chain." % self._samples_after_tune
             warning = SamplerWarning(
-                'bad-params', msg, 'error', None, None, None)
+                WarningType.BAD_PARAMS, msg, 'error', None, None, None)
             warnings.append(warning)
 
         warnings.extend(self.step_adapt.warnings())

pymc3/step_methods/hmc/nuts.py

@@ -8,7 +8,7 @@
 from ..arraystep import Competence
 from .base_hmc import BaseHMC, HMCStepData, DivergenceInfo
 from .integration import IntegrationError
-from pymc3.backends.report import SamplerWarning
+from pymc3.backends.report import SamplerWarning, WarningType
 from pymc3.theanof import floatX
 from pymc3.vartypes import continuous_types
 
@@ -189,7 +189,8 @@ def warnings(self, strace):
         if np.mean(self._reached_max_treedepth) > 0.05:
             msg = ('The chain reached the maximum tree depth. Increase '
                    'max_treedepth, increase target_accept or reparameterize.')
-            warn = SamplerWarning('treedepth', msg, 'warn', None, None, None)
+            warn = SamplerWarning(WarningType.TREEDEPTH, msg, 'warn',
+                                  None, None, None)
             warnings.append(warn)
         return warnings
 

pymc3/step_methods/step_sizes.py

@@ -3,7 +3,7 @@
 import numpy as np
 from scipy import stats
 
-from pymc3.backends.report import SamplerWarning
+from pymc3.backends.report import SamplerWarning, WarningType
 
 
 class DualAverageAdaptation(object):
@@ -61,7 +61,7 @@ def warnings(self):
                    % (mean_accept, target_accept))
             info = {'target': target_accept, 'actual': mean_accept}
             warning = SamplerWarning(
-                'bad-acceptance', msg, 'warn', None, None, info)
+                WarningType.BAD_ACCEPTANCE, msg, 'warn', None, None, info)
             return [warning]
         else:
             return []

pymc3/tests/test_step.py

@@ -11,7 +11,6 @@
                                 MultivariateNormalProposal, HamiltonianMC,
                                 EllipticalSlice, smc, DEMetropolis)
 from pymc3.theanof import floatX
-from pymc3 import SamplingError
 from pymc3.distributions import (
     Binomial, Normal, Bernoulli, Categorical, Beta, HalfNormal)