diff --git a/pymc3/distributions/continuous.py b/pymc3/distributions/continuous.py index 75c214c3e8..171e30b7b3 100644 --- a/pymc3/distributions/continuous.py +++ b/pymc3/distributions/continuous.py @@ -129,6 +129,7 @@ def __init__(self, lower=0, upper=1, transform='interval', *args, **kwargs): super(Uniform, self).__init__(*args, **kwargs) + self.param_names = ['lower', 'upper'] self.lower = lower = tt.as_tensor_variable(lower) self.upper = upper = tt.as_tensor_variable(upper) self.mean = (upper + lower) / 2. @@ -204,6 +205,7 @@ class Normal(Continuous): """ def __init__(self, mu=0, sd=None, tau=None, **kwargs): + self.param_names = ['mu', 'sd', 'tau'] tau, sd = get_tau_sd(tau=tau, sd=sd) self.sd = tt.as_tensor_variable(sd) self.tau = tt.as_tensor_variable(tau) @@ -258,6 +260,7 @@ class HalfNormal(PositiveContinuous): def __init__(self, sd=None, tau=None, *args, **kwargs): super(HalfNormal, self).__init__(*args, **kwargs) + self.param_names = ['sd', 'tau'] tau, sd = get_tau_sd(tau=tau, sd=sd) self.sd = sd = tt.as_tensor_variable(sd) @@ -341,6 +344,7 @@ class Wald(PositiveContinuous): def __init__(self, mu=None, lam=None, phi=None, alpha=0., *args, **kwargs): super(Wald, self).__init__(*args, **kwargs) + self.param_names = ['mu', 'lam', 'phi', 'alpha'] mu, lam, phi = self.get_mu_lam_phi(mu, lam, phi) self.alpha = alpha = tt.as_tensor_variable(alpha) self.mu = mu = tt.as_tensor_variable(mu) @@ -451,6 +455,7 @@ def __init__(self, alpha=None, beta=None, mu=None, sd=None, *args, **kwargs): super(Beta, self).__init__(*args, **kwargs) + self.param_names = ['alpha', 'beta'] alpha, beta = self.get_alpha_beta(alpha, beta, mu, sd) self.alpha = alpha = tt.as_tensor_variable(alpha) self.beta = beta = tt.as_tensor_variable(beta) @@ -514,6 +519,7 @@ class Exponential(PositiveContinuous): def __init__(self, lam, *args, **kwargs): super(Exponential, self).__init__(*args, **kwargs) + self.param_names = ['lam'] self.lam = lam = tt.as_tensor_variable(lam) self.mean = 1. / self.lam self.median = self.mean * tt.log(2) @@ -559,6 +565,7 @@ class Laplace(Continuous): def __init__(self, mu, b, *args, **kwargs): super(Laplace, self).__init__(*args, **kwargs) + self.param_names = ['mu', 'b'] self.b = b = tt.as_tensor_variable(b) self.mean = self.median = self.mode = self.mu = mu = tt.as_tensor_variable(mu) @@ -612,6 +619,7 @@ def __init__(self, mu=0, sd=None, tau=None, *args, **kwargs): super(Lognormal, self).__init__(*args, **kwargs) tau, sd = get_tau_sd(tau=tau, sd=sd) + self.param_names = ['mu', 'tau'] self.mu = mu = tt.as_tensor_variable(mu) self.tau = tau = tt.as_tensor_variable(tau) self.sd = sd = tt.as_tensor_variable(sd) @@ -674,6 +682,7 @@ class StudentT(Continuous): def __init__(self, nu, mu=0, lam=None, sd=None, *args, **kwargs): super(StudentT, self).__init__(*args, **kwargs) + self.param_names = ['nu', 'mu', 'lam'] self.nu = nu = tt.as_tensor_variable(nu) lam, sd = get_tau_sd(tau=lam, sd=sd) self.lam = lam = tt.as_tensor_variable(lam) @@ -844,6 +853,7 @@ class HalfCauchy(PositiveContinuous): def __init__(self, beta, *args, **kwargs): super(HalfCauchy, self).__init__(*args, **kwargs) + self.param_names = ['beta'] self.mode = tt.as_tensor_variable(0) self.median = tt.as_tensor_variable(beta) self.beta = tt.as_tensor_variable(beta) @@ -1148,6 +1158,7 @@ class ExGaussian(Continuous): def __init__(self, mu, sigma, nu, *args, **kwargs): super(ExGaussian, self).__init__(*args, **kwargs) + self.param_names = ['mu', 'sigma', 'nu'] self.mu = mu = tt.as_tensor_variable(mu) self.sigma = sigma = tt.as_tensor_variable(sigma) self.nu = nu = tt.as_tensor_variable(nu) @@ -1210,6 +1221,7 @@ class VonMises(Continuous): def __init__(self, mu=0.0, kappa=None, transform='circular', *args, **kwargs): super(VonMises, self).__init__(*args, **kwargs) + self.param_names = ['mu', 'kappa'] self.mean = self.median = self.mode = self.mu = mu = tt.as_tensor_variable(mu) self.kappa = kappa = tt.as_tensor_variable(kappa) self.variance = 1 - i1(kappa) / i0(kappa) @@ -1273,6 +1285,7 @@ class SkewNormal(Continuous): """ def __init__(self, mu=0.0, sd=None, tau=None, alpha=1, *args, **kwargs): super(SkewNormal, self).__init__(*args, **kwargs) + self.param_names = ['mu', 'sd', 'tau', 'alpha'] tau, sd = get_tau_sd(tau=tau, sd=sd) self.mu = mu = tt.as_tensor_variable(mu) self.tau = tt.as_tensor_variable(tau) diff --git a/pymc3/distributions/discrete.py b/pymc3/distributions/discrete.py index f759665887..6c891bbeae 100644 --- a/pymc3/distributions/discrete.py +++ b/pymc3/distributions/discrete.py @@ -1,4 +1,4 @@ -from functools import partial +from functools import partial import numpy as np import theano import theano.tensor as tt @@ -39,6 +39,7 @@ class Binomial(Discrete): def __init__(self, n, p, *args, **kwargs): super(Binomial, self).__init__(*args, **kwargs) + self.param_names = ['n', 'p'] self.n = n = tt.as_tensor_variable(n) self.p = p = tt.as_tensor_variable(p) self.mode = tt.cast(tt.round(n * p), self.dtype) @@ -90,6 +91,7 @@ class BetaBinomial(Discrete): def __init__(self, alpha, beta, n, *args, **kwargs): super(BetaBinomial, self).__init__(*args, **kwargs) + self.param_names = ['n', 'alpha', 'beta'] self.alpha = alpha = tt.as_tensor_variable(alpha) self.beta = beta = tt.as_tensor_variable(beta) self.n = n = tt.as_tensor_variable(n) @@ -146,6 +148,7 @@ class Bernoulli(Discrete): def __init__(self, p, *args, **kwargs): super(Bernoulli, self).__init__(*args, **kwargs) + self.param_names = ['p'] self.p = p = tt.as_tensor_variable(p) self.mode = tt.cast(tt.round(p), 'int8') @@ -179,16 +182,16 @@ class DiscreteWeibull(Discrete): """ def __init__(self, q, beta, *args, **kwargs): super(DiscreteWeibull, self).__init__(*args, defaults=['median'], **kwargs) - + self.q = q self.beta = beta self.median = self._ppf(0.5) - + def logp(self, value): q = self.q beta = self.beta - + return bound(tt.log(tt.power(q, tt.power(value, beta)) - tt.power(q, tt.power(value + 1, beta))), 0 <= value, 0 < q, q < 1, @@ -246,6 +249,7 @@ class Poisson(Discrete): def __init__(self, mu, *args, **kwargs): super(Poisson, self).__init__(*args, **kwargs) + self.param_names = ['mu'] self.mu = mu = tt.as_tensor_variable(mu) self.mode = tt.floor(mu).astype('int32') @@ -293,6 +297,7 @@ class NegativeBinomial(Discrete): def __init__(self, mu, alpha, *args, **kwargs): super(NegativeBinomial, self).__init__(*args, **kwargs) + self.param_names = ['mu', 'alpha'] self.mu = mu = tt.as_tensor_variable(mu) self.alpha = alpha = tt.as_tensor_variable(alpha) self.mode = tt.floor(mu).astype('int32') @@ -342,6 +347,7 @@ class Geometric(Discrete): def __init__(self, p, *args, **kwargs): super(Geometric, self).__init__(*args, **kwargs) + self.param_names = ['p'] self.p = p = tt.as_tensor_variable(p) self.mode = 1 @@ -379,6 +385,7 @@ class DiscreteUniform(Discrete): def __init__(self, lower, upper, *args, **kwargs): super(DiscreteUniform, self).__init__(*args, **kwargs) + self.param_names = ['lower', 'upper'] self.lower = tt.floor(lower).astype('int32') self.upper = tt.floor(upper).astype('int32') self.mode = tt.maximum( @@ -426,6 +433,7 @@ class Categorical(Discrete): def __init__(self, p, *args, **kwargs): super(Categorical, self).__init__(*args, **kwargs) + self.param_names = ['p'] try: self.k = tt.shape(p)[-1].tag.test_value except AttributeError: @@ -475,12 +483,13 @@ class Constant(Discrete): Parameters ---------- - value : float or int + c : float or int Constant parameter. """ def __init__(self, c, *args, **kwargs): super(Constant, self).__init__(*args, **kwargs) + self.param_names = ['c'] self.mean = self.median = self.mode = self.c = c = tt.as_tensor_variable(c) def random(self, point=None, size=None, repeat=None): @@ -537,6 +546,7 @@ class ZeroInflatedPoisson(Discrete): def __init__(self, theta, psi, *args, **kwargs): super(ZeroInflatedPoisson, self).__init__(*args, **kwargs) + self.param_names = ['theta', 'psi'] self.theta = theta = tt.as_tensor_variable(theta) self.psi = psi = tt.as_tensor_variable(psi) self.pois = Poisson.dist(theta) @@ -589,6 +599,7 @@ class ZeroInflatedNegativeBinomial(Discrete): def __init__(self, mu, alpha, psi, *args, **kwargs): super(ZeroInflatedNegativeBinomial, self).__init__(*args, **kwargs) + self.param_names = ['mu', 'alpha', 'psi'] self.mu = mu = tt.as_tensor_variable(mu) self.alpha = alpha = tt.as_tensor_variable(alpha) self.psi = psi = tt.as_tensor_variable(psi) diff --git a/pymc3/distributions/timeseries.py b/pymc3/distributions/timeseries.py index c1d9705af4..7bc5ca1e68 100644 --- a/pymc3/distributions/timeseries.py +++ b/pymc3/distributions/timeseries.py @@ -29,6 +29,7 @@ class AR1(distribution.Continuous): def __init__(self, k, tau_e, *args, **kwargs): super(AR1, self).__init__(*args, **kwargs) + self.param_names = ['k', 'tau_e'] self.k = k self.tau_e = tau_e self.tau = tau_e * (1 - k ** 2) @@ -65,6 +66,7 @@ class GaussianRandomWalk(distribution.Continuous): def __init__(self, tau=None, init=continuous.Flat.dist(), sd=None, mu=0., *args, **kwargs): super(GaussianRandomWalk, self).__init__(*args, **kwargs) + self.param_names = ['tau', 'sd', 'mu'] self.tau = tau self.sd = sd self.mu = mu @@ -109,6 +111,7 @@ class GARCH11(distribution.Continuous): def __init__(self, omega=None, alpha_1=None, beta_1=None, initial_vol=None, *args, **kwargs): super(GARCH11, self).__init__(*args, **kwargs) + self.param_names = ['omega', 'alpha_1', 'beta_1', 'initial_vol'] self.omega = omega self.alpha_1 = alpha_1 @@ -149,6 +152,7 @@ class EulerMaruyama(distribution.Continuous): """ def __init__(self, dt, sde_fn, sde_pars, *args, **kwds): super(EulerMaruyama, self).__init__(*args, **kwds) + self.param_names = ['dt'] self.dt = dt self.sde_fn = sde_fn self.sde_pars = sde_pars diff --git a/pymc3/plots.py b/pymc3/plots.py index 2e04524e38..17016d426f 100644 --- a/pymc3/plots.py +++ b/pymc3/plots.py @@ -3,12 +3,13 @@ import numpy as np from scipy.stats import kde, mode import matplotlib.pyplot as plt +import networkx as nx import pymc3 as pm from .stats import quantiles, hpd from scipy.signal import gaussian, convolve __all__ = ['traceplot', 'kdeplot', 'kde2plot', - 'forestplot', 'autocorrplot', 'plot_posterior'] + 'forestplot', 'autocorrplot', 'plot_posterior', 'plot_network'] def traceplot(trace, varnames=None, transform=lambda x: x, figsize=None, @@ -854,3 +855,131 @@ def fast_kde(x): grid /= norm_factor return grid, xmin, xmax + + +def get_inputs(node, variables): + if node.owner is not None: + for inp in node.owner.inputs: + if inp in variables: + yield inp + elif inp.owner is not None: + for i in get_inputs(inp, variables): + yield i + + +def tree_layout(graph): + if not graph.is_directed(): + raise ValueError('graph must be directed') + + if not nx.is_directed_acyclic_graph(graph): + raise ValueError('graph must not contain cycles') + + roots = [node for node, degree in graph.out_degree_iter() if degree < 1] + + # Find the tree width at every depth in order to layout + # the nodes in a justified manner + depths = {} + for node in graph.nodes(): + depth = 0 + parents = graph.successors(node) + while len(parents) > 0: + _node = parents[0] + parents = graph.successors(_node) + depth += 1 + depths[node] = depth + depths_list = depths.values() + + max_depth = max(depths_list) + widths = [depths_list.count(i) for i in range(max_depth+1)] + max_width = max(widths) + nodes_positioned_at_depth = [0] * len(widths) + + # breadth first tree transversal + midpoint = max_width / 2.0 + depth_positions = [] + for width in widths: + pos = (np.linspace(0, width, width) - + (width / 2.0) + + midpoint).tolist() + depth_positions.append(pos) + positions = {} + assert len(roots) == len(depth_positions[0]) + for root in roots: + root_depth = depths[root] + positions[root] = (depth_positions[root_depth].pop(0), root_depth) + for node in graph.nodes(): + if node not in positions: + node_depth = depths[node] + positions[node] = ( + depth_positions[node_depth].pop(0), node_depth + ) + + return positions + + +def plot_network(model, layout='tree', plot=True, ax=None): + """ + Create and optionally draw a directed graph (network) representation of a + provided model. + + Parameters + ---------- + model : pymc3.Model + The Model object to be coerced into a graph. + layout : str, default='tree' + Layout to use when plotting the graph. Available options are: + 'circular', 'shell', 'spectral', 'spring', 'force-directed', 'random', + 'tree'. + plot : bool + Plots the graph using Matplotlib if True. If False, no plot is created + and only the graph object is returned. + ax : Matplotlib Axes object, optional + Draw the graph in specified Matplotlib axes. + + Returns + ------- + graph : networkx.DiGraph + The graph object extracted from the information in the model. + """ + + available_layouts = { + "circular": nx.circular_layout, + "shell": nx.shell_layout, + "spectral": nx.spectral_layout, + "spring": nx.spring_layout, + "force-directed": nx.fruchterman_reingold_layout, + "random": nx.random_layout, + "tree": tree_layout, + } + + try: + layout_func = available_layouts[layout] + except KeyError: + msg = (u"The selected layout option ({}) is not available. Please " + u"choose from the following options: {}".format( + layout, available_layouts.keys() + )) + raise KeyError(msg) + + graph = nx.DiGraph() + + variables = model.named_vars.values() + for var in variables: + graph.add_node(var) + dist = var.distribution + for param_name in getattr(dist, "param_names", []): + param_value = getattr(dist, param_name) + owner = getattr(param_value, "owner", None) + if param_value in variables: + graph.add_edge(param_value, var) + elif owner is not None: + for parent in get_inputs(param_value, variables): + if parent in variables: + graph.add_edge(parent, var) + + if plot: + pos = layout_func(graph) + nx.draw_networkx(graph, pos, arrows=True, node_size=1000, node_color='w', font_size=8) + plt.axis('off') + + return graph diff --git a/requirements.txt b/requirements.txt index 81b5f4a401..6ae82f430b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,6 +6,7 @@ pandas>=0.18.0 patsy>=0.4.0 joblib>=0.9 tqdm>=4.8.4 +networkx==1.11.0 CommonMark==0.5.4 recommonmark sphinx diff --git a/scripts/create_testenv.sh b/scripts/create_testenv.sh index 89131a6338..16c9867be4 100755 --- a/scripts/create_testenv.sh +++ b/scripts/create_testenv.sh @@ -23,9 +23,10 @@ then source activate testenv fi + pip install jupyter conda install --yes pyqt matplotlib --channel conda-forge -conda install --yes pyzmq numpy scipy nose pandas Cython patsy statsmodels joblib coverage mkl-service +conda install --yes pyzmq numpy scipy nose pandas Cython patsy statsmodels networkx joblib coverage mkl-service if [ ${PYTHON_VERSION} == "2.7" ]; then conda install --yes mock enum34; fi