Added live_traceplot function (#1934)

davidbrochart · twiecki · commit d2bca9045ffd · 2017-03-27T11:36:12.000+02:00
* Added live_traceplot function * Cosmetic change * Changed the API to pm.sample(..., live_plot=True) * Don't include `-np.inf` in calculating average ELBO (#1880) * Adds an infmean for advi reporting * fixing typo * Add tutorial to detect sampling problems (#1866) * Expand sampler-stats.ipynb example include model diagnose from case study example in Stan http://mc-stan.org/documentation/case-studies/divergences_and_bias.html * Sampler Diagnose for NUTS * descriptive annotation and axis labels * Fix typos * PEP8 styling * minor updates 1, add example to examples.rst 2, original content in Markdown code block * Make install scripts idempotent (#1879) * DOC Change heading names. * Add examples of censored data models (#1870) * Raise TypeError on non-data values of observed (#1872) * Raise TypeError on non-data values of observed * Added check for observed TypeError * Make exponential mode have the correct shape * Fix support of LKJCorr * Added tutorial notebook on updating priors * Fixed y-axis bug in forestplot; added transform argument to summary * Style cleanup * Made small changes and executed the notebook * Added probit and invprobit functions * Added carriage return to end of file * Fixed indentation * Changed probit test to use assert_allclose * Fix tests for LKJCorr * Added warning for ignoring init arguments in sample * Kill stray tab * Improve performance of transformations * DOC Add new features * Bump version. * Added docs and scripts to MANIFEST * WIP: Implement opvi (#1694) * migrate useful functions from previous PR (cherry picked from commit 9f61ab4) * opvi draft (cherry picked from commit d0997ff) * made some test work (cherry picked from commit b1a87d5) * refactored approximation to support aevb (without test) * refactor opvi delete unnecessary methods from operator, change method order * change log_q_local computation * add full rank approximation * add more_params argument to ObjectiveFunction.updates (aevb case) * refactor density computation in full rank approximation * typo: cast dict values to list * typo: cast dict values to list * typo: undefined T in dist_math * refactor gradient scaling as suggested in approximateinference.org/accepted/RoederEtAl2016.pdf * implement Langevin-Stein (LS) operator * fix docstring * add blank line in docs * refactor ObjectiveFunction * add not working LS Op test * experiments with not working LS Op * change activations * refactor networks * add step_function * remove Langevin Stein, done refactoring * remove Langevin Stein, done refactoring * change optimizers * refactor init params * implement tests * implement Inference * code style * test fix * add minibatch test (fails now) * add more tests for minibatch training * add logdet to FullRank approximation * add conversion of arrays to floatX * tiny changes * change number of iterations * fix test and pylint check * memoize functions in Objective function * Optimize code a lot * a bit more efficient pickling * add docs * Add MeanField -> FullRank parameter transfer * refactor MeanField and FullRank a bit * fix FullRank bug with shapes in random * refactor Model.flatten (CC @taku-y) * add `approximate` to inference * rename approximate->fit * change abbreviations * Fix bug with scaling input variable in aevb * fix theano bottleneck in graph * more efficient scaling for local vars * fix typo in local Q * add aevb test * refactor memoize to work with my objects * add tests for numpy view usage * pickle-hash fix * pickle-hash fix again * add node sampling + make up some code * add notebook with example * sample_proba explained * Revert "small fix for multivariate mixture models" * Added message about init only working with auto-assigned step methods * doc(DiagInferDiv): formatting fix in blog post quote. Closes #1895. (#1909) * delete unnecessary text and add some benchmarks (#1901) * Add LKJCholeskyCov * Added newline to MANIFEST * Replaced package list with find_packages in setup.py; removed examples/data/__init__.py * Fix log jacobian in LKJCholeskyCov * Updated version to rc2 * Fixed stray version string * Fix indexing traces with steps greater one * refactor variational module, add histogram approximation (#1904) * refactor module, add histogram * add more tests * refactor some code concerning AEVB histogram * fix test for histogram * use mean as deterministic point in Histogram * remove unused import * change names of shortcuts * add names to shared params * add new line at the end of `approximations.py` * Add documentation for LKJCholeskyCov * SVGD problems (#1916) * fix some svgd problems * switch -> ifelse * except in record * Histogram docs (#1914) * add docs * delete redundant code * add usage example * remove unused import * Add expand_packed_triangular * improve aesthetics * Bump theano to 0.9.0rc4 (#1921) * Add tests for LKJCholeskyCov * Histogram: use only free RVs from trace (#1926) * use only free RVs from trace * use memoize in Histogram.histogram_logp * Change tests for histogram * Bump theano to be at least 0.9.0 * small fix to prevent a TypeError with the ufunc true_divide * Fix tests for py2 * Add floatX wrappers in test_advi * Changed the API to pm.sample(..., live_plot=True) * Better formatting
diff --git a/docs/source/notebooks/live_sample_plots.ipynb b/docs/source/notebooks/live_sample_plots.ipynb
@@ -0,0 +1,122 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "deletable": true,
+    "editable": true
+   },
+   "source": [
+    "# Live sample plots"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "deletable": true,
+    "editable": true
+   },
+   "source": [
+    "This notebook illustrates how we can have live sample plots when calling the `sample` function with `live_plot=True`. It is based on the \"Coal mining disasters\" case study in the [Getting started notebook](https://github.com/pymc-devs/pymc3/blob/master/docs/source/notebooks/getting_started.ipynb)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false,
+    "deletable": true,
+    "editable": true
+   },
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "from pymc3 import Model, Exponential, DiscreteUniform, Poisson, sample\n",
+    "from pymc3.math import switch\n",
+    "\n",
+    "%matplotlib notebook"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false,
+    "deletable": true,
+    "editable": true
+   },
+   "outputs": [],
+   "source": [
+    "disaster_data = np.ma.masked_values([4, 5, 4, 0, 1, 4, 3, 4, 0, 6, 3, 3, 4, 0, 2, 6,\n",
+    "                            3, 3, 5, 4, 5, 3, 1, 4, 4, 1, 5, 5, 3, 4, 2, 5,\n",
+    "                            2, 2, 3, 4, 2, 1, 3, -999, 2, 1, 1, 1, 1, 3, 0, 0,\n",
+    "                            1, 0, 1, 1, 0, 0, 3, 1, 0, 3, 2, 2, 0, 1, 1, 1,\n",
+    "                            0, 1, 0, 1, 0, 0, 0, 2, 1, 0, 0, 0, 1, 1, 0, 2,\n",
+    "                            3, 3, 1, -999, 2, 1, 1, 1, 1, 2, 4, 2, 0, 0, 1, 4,\n",
+    "                            0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1], value=-999)\n",
+    "year = np.arange(1851, 1962)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false,
+    "deletable": true,
+    "editable": true
+   },
+   "outputs": [],
+   "source": [
+    "with Model() as disaster_model:\n",
+    "\n",
+    "    switchpoint = DiscreteUniform('switchpoint', lower=year.min(), upper=year.max(), testval=1900)\n",
+    "\n",
+    "    # Priors for pre- and post-switch rates number of disasters\n",
+    "    early_rate = Exponential('early_rate', 1)\n",
+    "    late_rate = Exponential('late_rate', 1)\n",
+    "\n",
+    "    # Allocate appropriate Poisson rates to years before and after current\n",
+    "    rate = switch(switchpoint >= year, early_rate, late_rate)\n",
+    "\n",
+    "    disasters = Poisson('disasters', rate, observed=disaster_data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false,
+    "deletable": true,
+    "editable": true,
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "with disaster_model:\n",
+    "    trace = sample(10000, live_plot=True, skip_first=100, refresh_every=300, roll_over=1000)"
+   ]
+  }
+ ],
+ "metadata": {
+  "anaconda-cloud": {},
+  "kernelspec": {
+   "display_name": "Python [default]",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.5.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/pymc3/plots/traceplot.py b/pymc3/plots/traceplot.py
@@ -7,7 +7,8 @@
 
 def traceplot(trace, varnames=None, transform=identity_transform, figsize=None, lines=None,
               combined=False, plot_transformed=False, grid=False, alpha=0.35, priors=None,
-              prior_alpha=1, prior_style='--', ax=None):
+              prior_alpha=1, prior_style='--', ax=None, live_plot=False,
+              skip_first=0, refresh_every=100, roll_over=1000):
     """Plot samples histograms and values.
 
     Parameters
@@ -45,6 +46,16 @@ def traceplot(trace, varnames=None, transform=identity_transform, figsize=None,
         Line style for prior plot. Defaults to '--' (dashed line).
     ax : axes
         Matplotlib axes. Accepts an array of axes, e.g.:
+    live_plot: bool
+        Flag for updating the current figure while sampling
+    skip_first : int
+        Number of first samples not shown in plots (burn-in). This affects
+        frequency and stream plots.
+    refresh_every : int
+        Period of plot updates (in sample number)
+    roll_over : int
+        Width of the sliding window for the sample stream plots: last roll_over
+        samples are shown (no effect on frequency plots).
 
         >>> fig, axs = plt.subplots(3, 2) # 3 RVs
         >>> pymc3.traceplot(trace, ax=axs)
@@ -57,6 +68,8 @@ def traceplot(trace, varnames=None, transform=identity_transform, figsize=None,
     ax : matplotlib axes
 
     """
+    trace = trace[skip_first:]
+
     if varnames is None:
         varnames = get_default_varnames(trace, plot_transformed)
 
@@ -70,9 +83,23 @@ def traceplot(trace, varnames=None, transform=identity_transform, figsize=None,
             prior = priors[i]
         else:
             prior = None
+        first_time = True
         for d in trace.get_values(v, combine=combined, squeeze=False):
             d = np.squeeze(transform(d))
             d = make_2d(d)
+            d_stream = d
+            x0 = 0
+            if live_plot:
+                x0 = skip_first
+                if first_time:
+                    ax[i, 0].cla()
+                    ax[i, 1].cla()
+                    first_time = False
+                if roll_over is not None:
+                    if len(d) >= roll_over:
+                        x0 = len(d) - roll_over + skip_first
+                    d_stream = d[-roll_over:]
+            width = len(d_stream)
             if d.dtype.kind == 'i':
                 hist_objs = histplot_op(ax[i, 0], d, alpha=alpha)
                 colors = [h[-1][0].get_facecolor() for h in hist_objs]
@@ -82,7 +109,7 @@ def traceplot(trace, varnames=None, transform=identity_transform, figsize=None,
             ax[i, 0].set_title(str(v))
             ax[i, 0].grid(grid)
             ax[i, 1].set_title(str(v))
-            ax[i, 1].plot(d, alpha=alpha)
+            ax[i, 1].plot(range(x0, x0 + width), d_stream, alpha=alpha)
 
             ax[i, 0].set_ylabel("Frequency")
             ax[i, 1].set_ylabel("Sample value")
@@ -103,6 +130,13 @@ def traceplot(trace, varnames=None, transform=identity_transform, figsize=None,
                                          lw=1.5, alpha=alpha)
                 except KeyError:
                     pass
+        if live_plot:
+            for j in [0, 1]:
+                ax[i, j].relim()
+                ax[i, j].autoscale_view(True, True, True)
+            ax[i, 1].set_xlim(x0, x0 + width)
         ax[i, 0].set_ylim(ymin=0)
+    if live_plot:
+        ax[0, 0].figure.canvas.draw()
     plt.tight_layout()
     return ax
diff --git a/pymc3/sampling.py b/pymc3/sampling.py
@@ -11,6 +11,8 @@
 from .step_methods import (NUTS, HamiltonianMC, Metropolis, BinaryMetropolis,
                            BinaryGibbsMetropolis, CategoricalGibbsMetropolis,
                            Slice, CompoundStep)
+from .plots.utils import identity_transform
+from .plots.traceplot import traceplot
 from tqdm import tqdm
 
 import warnings
@@ -85,7 +87,7 @@ def assign_step_methods(model, step=None, methods=(NUTS, HamiltonianMC, Metropol
 
 def sample(draws, step=None, init='ADVI', n_init=200000, start=None,
            trace=None, chain=0, njobs=1, tune=None, progressbar=True,
-           model=None, random_seed=-1):
+           model=None, random_seed=-1, live_plot=False, **kwargs):
     """Draw samples from the posterior using the given step methods.
 
     Multiple step methods are supported via compound step methods.
@@ -141,6 +143,8 @@ def sample(draws, step=None, init='ADVI', n_init=200000, start=None,
     model : Model (optional if in `with` context)
     random_seed : int or list of ints
         A list is accepted if more if `njobs` is greater than one.
+    live_plot: bool
+        Flag for live plotting the trace while sampling
 
     Returns
     -------
@@ -175,7 +179,9 @@ def sample(draws, step=None, init='ADVI', n_init=200000, start=None,
                    'tune': tune,
                    'progressbar': progressbar,
                    'model': model,
-                   'random_seed': random_seed}
+                   'random_seed': random_seed,
+                   'live_plot': live_plot,
+                   **kwargs}
 
     if njobs > 1:
         sample_func = _mp_sample
@@ -187,15 +193,27 @@ def sample(draws, step=None, init='ADVI', n_init=200000, start=None,
 
 
 def _sample(draws, step=None, start=None, trace=None, chain=0, tune=None,
-            progressbar=True, model=None, random_seed=-1):
+            progressbar=True, model=None, random_seed=-1, live_plot=False,
+            **kwargs):
+    live_plot_args = {'skip_first': 0, 'refresh_every': 100}
+    live_plot_args = {arg: kwargs[arg] if arg in kwargs else live_plot_args[arg] for arg in live_plot_args}
+    skip_first = live_plot_args['skip_first']
+    refresh_every = live_plot_args['refresh_every']
+
     sampling = _iter_sample(draws, step, start, trace, chain,
                             tune, model, random_seed)
     if progressbar:
         sampling = tqdm(sampling, total=draws)
     try:
         strace = None
-        for strace in sampling:
-            pass
+        for it, strace in enumerate(sampling):
+            if live_plot:
+                if it >= skip_first:
+                    trace = MultiTrace([strace])
+                    if it == skip_first:
+                        ax = traceplot(trace, live_plot=False, **kwargs)
+                    elif (it - skip_first) % refresh_every == 0 or it == draws - 1:
+                        traceplot(trace, ax=ax, live_plot=True, **kwargs)
     except KeyboardInterrupt:
         pass
     finally:
