Add warnings in new multiprocessing sampling

Author: aseyboldt

pymc3/parallel_sampling.py

@@ -98,8 +98,12 @@ def _start_loop(self):
             elif msg[0] == 'write_next':
                 self._write_point(point)
                 is_last = draw + 1 == self._draws + self._tune
+                if is_last:
+                    warns = self._collect_warnings()
+                else:
+                    warns = None
                 self._msg_pipe.send(
-                    ('writing_done', is_last, draw, tuning, stats))
+                    ('writing_done', is_last, draw, tuning, stats, warns))
                 draw += 1
             else:
                 raise ValueError('Unknown message ' + msg[0])
@@ -112,6 +116,12 @@ def _compute_point(self):
             stats = None
         return point, stats
 
+    def _collect_warnings(self):
+        if hasattr(self._step_method, 'warnings'):
+            return self._step_method.warnings()
+        else:
+            return []
+
 
 class ProcessAdapter(object):
     """Control a Chain process from the main thread."""
@@ -218,7 +228,7 @@ def terminate_all(processes, patience=2):
 
 Draw = namedtuple(
     'Draw',
-    ['chain', 'is_last', 'draw_idx', 'tuning', 'stats', 'point']
+    ['chain', 'is_last', 'draw_idx', 'tuning', 'stats', 'point', 'warnings']
 )
 
 
@@ -274,7 +284,7 @@ def __iter__(self):
 
         while self._active:
             draw = ProcessAdapter.recv_draw(self._active)
-            proc, is_last, draw, tuning, stats = draw
+            proc, is_last, draw, tuning, stats, warns = draw
             if self._progress is not None:
                 self._progress[proc.chain - self._start_chain_num].update()
 
@@ -299,7 +309,7 @@ def __iter__(self):
             if not is_last:
                 proc.write_next()
 
-            yield Draw(proc.chain, is_last, draw, tuning, stats, point)
+            yield Draw(proc.chain, is_last, draw, tuning, stats, point, warns)
 
     def __enter__(self):
         self._in_context = True

pymc3/sampling.py

@@ -663,7 +663,7 @@ def _iter_sample(draws, step, start=None, trace=None, chain=0, tune=None,
     except KeyboardInterrupt:
         strace.close()
         if hasattr(step, 'warnings'):
-            warns = step.warnings(strace)
+            warns = step.warnings()
             strace._add_warnings(warns)
         raise
     except BaseException:
@@ -672,7 +672,7 @@ def _iter_sample(draws, step, start=None, trace=None, chain=0, tune=None,
     else:
         strace.close()
         if hasattr(step, 'warnings'):
-            warns = step.warnings(strace)
+            warns = step.warnings()
             strace._add_warnings(warns)
 
 
@@ -1002,6 +1002,8 @@ def _mp_sample(draws, tune, step, chains, cores, chain, random_seed,
                         trace.record(draw.point)
                     if draw.is_last:
                         trace.close()
+                        if draw.warnings is not None:
+                            trace._add_warnings(draw.warnings)
             return MultiTrace(traces)
         except KeyboardInterrupt:
             traces, length = _choose_chains(traces, tune)
@@ -1013,12 +1015,14 @@ def _mp_sample(draws, tune, step, chains, cores, chain, random_seed,
     else:
         chain_nums = list(range(chain, chain + chains))
         pbars = [progressbar] + [False] * (chains - 1)
-        jobs = (delayed(_sample)(
-                    chain=args[0], progressbar=args[1], random_seed=args[2],
-                    start=args[3], draws=draws, step=step, trace=trace,
-                    tune=tune, model=model, **kwargs
-                )
-                for args in zip(chain_nums, pbars, random_seed, start))
+        jobs = (
+            delayed(_sample)(
+                chain=args[0], progressbar=args[1], random_seed=args[2],
+                start=args[3], draws=draws, step=step, trace=trace,
+                tune=tune, model=model, **kwargs
+            )
+            for args in zip(chain_nums, pbars, random_seed, start)
+        )
         if use_mmap:
             traces = Parallel(n_jobs=cores)(jobs)
         else:

pymc3/step_methods/compound.py

@@ -33,11 +33,11 @@ def step(self, point):
                 point = method.step(point)
             return point
 
-    def warnings(self, strace):
+    def warnings(self):
         warns = []
         for method in self.methods:
             if hasattr(method, 'warnings'):
-                warns.extend(method.warnings(strace))
+                warns.extend(method.warnings())
         return warns
 
     def stop_tuning(self):

pymc3/step_methods/hmc/base_hmc.py

@@ -164,7 +164,7 @@ def reset(self, start=None):
         self.tune = True
         self.potential.reset()
 
-    def warnings(self, strace):
+    def warnings(self):
         # list.copy() is not available in python2
         warnings = self._warnings[:]
 

pymc3/step_methods/hmc/nuts.py

@@ -184,8 +184,8 @@ def competence(var, has_grad):
             return Competence.IDEAL
         return Competence.INCOMPATIBLE
 
-    def warnings(self, strace):
-        warnings = super(NUTS, self).warnings(strace)
+    def warnings(self):
+        warnings = super(NUTS, self).warnings()
         n_samples = self._samples_after_tune
         n_treedepth = self._reached_max_treedepth