Latex model (#2450)

* Add __latex__ alias, latex support for models

* Guard for missing latex

* Add LaTeX for deterministics, test for more complicated model

Author: ColCarroll

pymc3/distributions/distribution.py

@@ -88,8 +88,11 @@ def getattr_value(self, val):
         return val
 
     def _repr_latex_(self, name=None, dist=None):
+        """Magic method name for IPython to use for LaTeX formatting."""
         return None
 
+    __latex__ = _repr_latex_
+
 
 def TensorType(dtype, shape, broadcastable=None):
     if broadcastable is None:

pymc3/model.py

@@ -1,4 +1,6 @@
 import collections
+import functools
+import itertools
 import threading
 import six
 
@@ -892,6 +894,14 @@ def flatten(self, vars=None, order=None, inputvar=None):
         flat_view = FlatView(inputvar, replacements, view)
         return flat_view
 
+    def _repr_latex_(self, name=None, dist=None):
+        tex_vars = []
+        for rv in itertools.chain(self.unobserved_RVs, self.observed_RVs):
+            tex_vars.append(rv.__latex__())
+        return u'$${}$$'.format('\\\\'.join([tex.strip('$') for tex in tex_vars if tex is not None]))
+
+    __latex__ = _repr_latex_
+
 
 def fn(outs, mode=None, model=None, *args, **kwargs):
     """Compiles a Theano function which returns the values of `outs` and
@@ -1073,6 +1083,8 @@ def _repr_latex_(self, name=None, dist=None):
             dist = self.distribution
         return self.distribution._repr_latex_(name=name, dist=dist)
 
+    __latex__ = _repr_latex_
+
     @property
     def init_value(self):
         """Convenience attribute to return tag.test_value"""
@@ -1176,6 +1188,8 @@ def _repr_latex_(self, name=None, dist=None):
             dist = self.distribution
         return self.distribution._repr_latex_(name=name, dist=dist)
 
+    __latex__ = _repr_latex_
+
     @property
     def init_value(self):
         """Convenience attribute to return tag.test_value"""
@@ -1212,6 +1226,26 @@ def __init__(self, name, data, distribution, total_size=None, model=None):
         self.scaling = _get_scaling(total_size, self.logp_elemwiset.shape, self.logp_elemwiset.ndim)
 
 
+def _walk_up_rv(rv):
+    """Walk up theano graph to get inputs for deterministic RV."""
+    all_rvs = []
+    parents = list(itertools.chain(*[j.inputs for j in rv.get_parents()]))
+    if parents:
+        for parent in parents:
+            all_rvs.extend(_walk_up_rv(parent))
+    else:
+        if rv.name:
+            all_rvs.append(rv.name)
+        else:
+            all_rvs.append(r'\text{Constant}')
+    return all_rvs
+
+
+def _latex_repr_rv(rv):
+    """Make latex string for a Deterministic variable"""
+    return r'${} \sim \text{{Deterministic}}({})$'.format(rv.name, r', '.join(_walk_up_rv(rv)))
+
+
 def Deterministic(name, var, model=None):
     """Create a named deterministic variable
 
@@ -1228,6 +1262,8 @@ def Deterministic(name, var, model=None):
     var.name = model.name_for(name)
     model.deterministics.append(var)
     model.add_random_variable(var)
+    var._repr_latex_ = functools.partial(_latex_repr_rv, var)
+    var.__latex__ = var._repr_latex_
     return var
 
 
@@ -1301,6 +1337,8 @@ def _repr_latex_(self, name=None, dist=None):
             dist = self.distribution
         return self.distribution._repr_latex_(name=name, dist=dist)
 
+    __latex__ = _repr_latex_
+
     @property
     def init_value(self):
         """Convenience attribute to return tag.test_value"""

pymc3/tests/test_distributions.py

@@ -4,17 +4,17 @@
 
 from .helpers import SeededTest, select_by_precision
 from ..vartypes import continuous_types
-from ..model import Model, Point, Potential
+from ..model import Model, Point, Potential, Deterministic
 from ..blocking import DictToVarBijection, DictToArrayBijection, ArrayOrdering
 from ..distributions import (DensityDist, Categorical, Multinomial, VonMises, Dirichlet,
-                             MvStudentT, MvNormal, ZeroInflatedPoisson, GaussianRandomWalk,
+                             MvStudentT, MvNormal, ZeroInflatedPoisson,
                              ZeroInflatedNegativeBinomial, Constant, Poisson, Bernoulli, Beta,
-                             BetaBinomial, HalfStudentT, StudentT, Weibull, Pareto, NormalMixture,
+                             BetaBinomial, HalfStudentT, StudentT, Weibull, Pareto,
                              InverseGamma, Gamma, Cauchy, HalfCauchy, Lognormal, Laplace,
                              NegativeBinomial, Geometric, Exponential, ExGaussian, Normal,
                              Flat, LKJCorr, Wald, ChiSquared, HalfNormal, DiscreteUniform,
-                             Bound, Uniform, Triangular, Binomial, SkewNormal, DiscreteWeibull, Gumbel,
-                             Interpolated, ZeroInflatedBinomial, HalfFlat)
+                             Bound, Uniform, Triangular, Binomial, SkewNormal, DiscreteWeibull,
+                             Gumbel, Interpolated, ZeroInflatedBinomial, HalfFlat)
 from ..distributions import continuous
 from pymc3.theanof import floatX
 from numpy import array, inf, log, exp
@@ -929,27 +929,54 @@ def test_bound():
     assert rand >= 5 and rand <= 8
 
 
-def test_repr_latex_():
-    with Model():
-        x0 = Binomial('Discrete', p=.5, n=10)
-        x1 = Normal('Continuous', mu=0., sd=1.)
-        x2 = GaussianRandomWalk('Timeseries', mu=x1, sd=1., shape=2)
-        x3 = MvStudentT('Multivariate', nu=5, mu=x2, Sigma=np.diag(np.ones(2)), shape=2)
-        x4 = NormalMixture('Mixture', w=np.array([.5, .5]), mu=x3, sd=x0)
-
-    assert x0._repr_latex_() == '$Discrete \\sim \\text{Binomial}' \
-                                '(\\mathit{n}=10, \\mathit{p}=0.5)$'
-    assert x1._repr_latex_() == '$Continuous \\sim \\text{Normal}' \
-                                '(\\mathit{mu}=0.0, \\mathit{sd}=1.0)$'
-    assert x2._repr_latex_() == '$Timeseries \\sim \\text' \
-                                '{GaussianRandomWalk}(\\mathit{mu}=Continuous, ' \
-                                '\\mathit{sd}=1.0)$'
-    assert x3._repr_latex_() == '$Multivariate \\sim \\text{MvStudentT}' \
-                                '(\\mathit{nu}=5, \\mathit{mu}=Timeseries, ' \
-                                '\\mathit{cov}=array)$'
-    assert x4._repr_latex_() == '$Mixture \\sim \\text{NormalMixture}' \
-                                '(\\mathit{w}=array, \\mathit{mu}=Multivariate, ' \
-                                '\\mathit{sigma}=f(Discrete))$'
+class TestLatex(object):
+
+    def setup_class(self):
+        # True parameter values
+        alpha, sigma = 1, 1
+        beta = [1, 2.5]
+
+        # Size of dataset
+        size = 100
+
+        # Predictor variable
+        X = np.random.normal(size=(size, 2)).dot(np.array([[1, 0], [0, 0.2]]))
+
+        # Simulate outcome variable
+        Y = alpha + X.dot(beta) + np.random.randn(size)*sigma
+        with Model() as self.model:
+            # Priors for unknown model parameters
+            alpha = Normal('alpha', mu=0, sd=10)
+            b = Normal('beta', mu=0, sd=10, shape=(2,), observed=beta)
+            sigma = HalfNormal('sigma', sd=1)
+
+            # Expected value of outcome
+            mu = Deterministic('mu', alpha + tt.dot(X, b))
+
+            # Likelihood (sampling distribution) of observations
+            Y_obs = Normal('Y_obs', mu=mu, sd=sigma, observed=Y)
+        self.distributions = [alpha, sigma, mu, b, Y_obs]
+        self.expected = (
+            '$alpha \\sim \\text{Normal}(\\mathit{mu}=0, \\mathit{sd}=10.0)$',
+            '$sigma \\sim \\text{HalfNormal}(\\mathit{sd}=1.0)$',
+            '$mu \\sim \\text{Deterministic}(alpha, \\text{Constant}, beta)$',
+            '$beta \\sim \\text{Normal}(\\mathit{mu}=0, \\mathit{sd}=10.0)$',
+            '$Y_obs \\sim \\text{Normal}(\\mathit{mu}=mu, \\mathit{sd}=f(sigma))$'
+        )
+
+    def test__repr_latex_(self):
+        for distribution, tex in zip(self.distributions, self.expected):
+            assert distribution._repr_latex_() == tex
+
+        model_tex = self.model._repr_latex_()
+
+        for tex in self.expected:  # make sure each variable is in the model
+            assert tex.strip('$') in model_tex
+
+    def test___latex__(self):
+        for distribution, tex in zip(self.distributions, self.expected):
+            assert distribution._repr_latex_() == distribution.__latex__()
+        assert self.model._repr_latex_() == self.model.__latex__()
 
 
 def test_discrete_trafo():