blacken examples (#4118)

Author: MarcoGorelli

pymc3/examples/GHME_2013.py

@@ -5,11 +5,11 @@
 from pymc3 import HalfCauchy, Model, Normal, get_data, sample
 from pymc3.distributions.timeseries import GaussianRandomWalk
 
-data = pd.read_csv(get_data('pancreatitis.csv'))
-countries = ['CYP', 'DNK', 'ESP', 'FIN', 'GBR', 'ISL']
+data = pd.read_csv(get_data("pancreatitis.csv"))
+countries = ["CYP", "DNK", "ESP", "FIN", "GBR", "ISL"]
 data = data[data.area.isin(countries)]
 
-age = data['age'] = np.array(data.age_start + data.age_end) / 2
+age = data["age"] = np.array(data.age_start + data.age_end) / 2
 rate = data.value = data.value * 1000
 group, countries = pd.factorize(data.area, order=countries)
 
@@ -20,7 +20,7 @@
     plt.subplot(2, 3, i + 1)
     plt.title(country)
     d = data[data.area == country]
-    plt.plot(d.age, d.value, '.')
+    plt.plot(d.age, d.value, ".")
 
     plt.ylim(0, rate.max())
 
@@ -43,33 +43,33 @@ def interpolate(x0, y0, x, group):
 
 
 with Model() as model:
-    coeff_sd = HalfCauchy('coeff_sd', 5)
+    coeff_sd = HalfCauchy("coeff_sd", 5)
 
-    y = GaussianRandomWalk('y', sigma=coeff_sd, shape=(nknots, ncountries))
+    y = GaussianRandomWalk("y", sigma=coeff_sd, shape=(nknots, ncountries))
 
     p = interpolate(knots, y, age, group)
 
-    sd = HalfCauchy('sd', 5)
+    sd = HalfCauchy("sd", 5)
 
-    vals = Normal('vals', p, sigma=sd, observed=rate)
+    vals = Normal("vals", p, sigma=sd, observed=rate)
 
 
 def run(n=3000):
     if n == "short":
         n = 150
     with model:
-        trace = sample(n, tune=int(n/2), init='advi+adapt_diag')
+        trace = sample(n, tune=int(n / 2), init="advi+adapt_diag")
 
     for i, country in enumerate(countries):
         plt.subplot(2, 3, i + 1)
         plt.title(country)
 
         d = data[data.area == country]
-        plt.plot(d.age, d.value, '.')
-        plt.plot(knots, trace[y][::5, :, i].T, color='r', alpha=.01)
+        plt.plot(d.age, d.value, ".")
+        plt.plot(knots, trace[y][::5, :, i].T, color="r", alpha=0.01)
 
         plt.ylim(0, rate.max())
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     run()

pymc3/examples/LKJ_correlation.py

@@ -14,41 +14,47 @@
 stds = np.ones(4) / 2.0
 
 # Correlation matrix of 4 variables:
-corr_r = np.array([[1.,  0.75,  0.,  0.15],
-                   [0.75,  1., -0.06,  0.19],
-                   [0., -0.06,  1., -0.04],
-                   [0.15,  0.19, -0.04,  1.]])
+corr_r = np.array(
+    [
+        [1.0, 0.75, 0.0, 0.15],
+        [0.75, 1.0, -0.06, 0.19],
+        [0.0, -0.06, 1.0, -0.04],
+        [0.15, 0.19, -0.04, 1.0],
+    ]
+)
 cov_matrix = np.diag(stds).dot(corr_r.dot(np.diag(stds)))
 
 dataset = multivariate_normal(mu_r, cov_matrix, size=n_obs)
 
 with pm.Model() as model:
 
-    mu = pm.Normal('mu', mu=0, sigma=1, shape=n_var)
+    mu = pm.Normal("mu", mu=0, sigma=1, shape=n_var)
 
     # Note that we access the distribution for the standard
     # deviations, and do not create a new random variable.
     sd_dist = pm.HalfCauchy.dist(beta=2.5)
-    packed_chol = pm.LKJCholeskyCov('chol_cov', n=n_var, eta=1, sd_dist=sd_dist)
+    packed_chol = pm.LKJCholeskyCov("chol_cov", n=n_var, eta=1, sd_dist=sd_dist)
     # compute the covariance matrix
     chol = pm.expand_packed_triangular(n_var, packed_chol, lower=True)
     cov = tt.dot(chol, chol.T)
 
     # Extract the standard deviations etc
-    sd = pm.Deterministic('sd', tt.sqrt(tt.diag(cov)))
-    corr = tt.diag(sd**-1).dot(cov.dot(tt.diag(sd**-1)))
-    r = pm.Deterministic('r', corr[np.triu_indices(n_var, k=1)])
+    sd = pm.Deterministic("sd", tt.sqrt(tt.diag(cov)))
+    corr = tt.diag(sd ** -1).dot(cov.dot(tt.diag(sd ** -1)))
+    r = pm.Deterministic("r", corr[np.triu_indices(n_var, k=1)])
 
-    like = pm.MvNormal('likelihood', mu=mu, chol=chol, observed=dataset)
+    like = pm.MvNormal("likelihood", mu=mu, chol=chol, observed=dataset)
 
 
 def run(n=1000):
     if n == "short":
         n = 50
     with model:
         trace = pm.sample(n)
-    pm.traceplot(trace, varnames=['mu', 'r'],
-                 lines={'mu': mu_r, 'r': corr_r[np.triu_indices(n_var, k=1)]})
+    pm.traceplot(
+        trace, varnames=["mu", "r"], lines={"mu": mu_r, "r": corr_r[np.triu_indices(n_var, k=1)]}
+    )
 
-if __name__ == '__main__':
+
+if __name__ == "__main__":
     run()

pymc3/examples/arbitrary_stochastic.py

@@ -10,15 +10,13 @@ def logp(failure, lam, value):
 
 def build_model():
     # data
-    failure = np.array([0., 1.])
-    value = np.array([1., 0.])
+    failure = np.array([0.0, 1.0])
+    value = np.array([1.0, 0.0])
 
     # model
     with pm.Model() as model:
-        lam = pm.Exponential('lam', 1.)
-        pm.DensityDist('x', logp, observed={'failure': failure,
-                                            'lam': lam,
-                                            'value': value})
+        lam = pm.Exponential("lam", 1.0)
+        pm.DensityDist("x", logp, observed={"failure": failure, "lam": lam, "value": value})
     return model
 
 
@@ -28,5 +26,6 @@ def run(n_samples=3000):
         trace = pm.sample(n_samples)
     return trace
 
+
 if __name__ == "__main__":
     run()

pymc3/examples/arma_example.py

@@ -2,6 +2,7 @@
 from theano import scan, shared
 
 import numpy as np
+
 """
 ARMA example
 It is interesting to note just how much more compact this is than the original Stan example
@@ -53,36 +54,36 @@
 def build_model():
     y = shared(np.array([15, 10, 16, 11, 9, 11, 10, 18], dtype=np.float32))
     with pm.Model() as arma_model:
-        sigma = pm.HalfNormal('sigma', 5.)
-        theta = pm.Normal('theta', 0., sigma=1.)
-        phi = pm.Normal('phi', 0., sigma=2.)
-        mu = pm.Normal('mu', 0., sigma=10.)
+        sigma = pm.HalfNormal("sigma", 5.0)
+        theta = pm.Normal("theta", 0.0, sigma=1.0)
+        phi = pm.Normal("phi", 0.0, sigma=2.0)
+        mu = pm.Normal("mu", 0.0, sigma=10.0)
 
         err0 = y[0] - (mu + phi * mu)
 
         def calc_next(last_y, this_y, err, mu, phi, theta):
             nu_t = mu + phi * last_y + theta * err
             return this_y - nu_t
 
-        err, _ = scan(fn=calc_next,
-                      sequences=dict(input=y, taps=[-1, 0]),
-                      outputs_info=[err0],
-                      non_sequences=[mu, phi, theta])
+        err, _ = scan(
+            fn=calc_next,
+            sequences=dict(input=y, taps=[-1, 0]),
+            outputs_info=[err0],
+            non_sequences=[mu, phi, theta],
+        )
 
-        pm.Potential('like', pm.Normal.dist(0, sigma=sigma).logp(err))
+        pm.Potential("like", pm.Normal.dist(0, sigma=sigma).logp(err))
     return arma_model
 
 
 def run(n_samples=1000):
     model = build_model()
     with model:
-        trace = pm.sample(draws=n_samples,
-                          tune=1000,
-                          target_accept=.99)
+        trace = pm.sample(draws=n_samples, tune=1000, target_accept=0.99)
 
     pm.plots.traceplot(trace)
     pm.plots.forestplot(trace)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     run()

pymc3/examples/baseball.py

@@ -6,31 +6,35 @@
 import pymc3 as pm
 import numpy as np
 
+
 def build_model():
-    data = np.loadtxt(pm.get_data('efron-morris-75-data.tsv'), delimiter="\t", 
-                      skiprows=1, usecols=(2,3))
-    
-    atbats = pm.floatX(data[:,0])
-    hits = pm.floatX(data[:,1])
-    
+    data = np.loadtxt(
+        pm.get_data("efron-morris-75-data.tsv"), delimiter="\t", skiprows=1, usecols=(2, 3)
+    )
+
+    atbats = pm.floatX(data[:, 0])
+    hits = pm.floatX(data[:, 1])
+
     N = len(hits)
-    
+
     # we want to bound the kappa below
     BoundedKappa = pm.Bound(pm.Pareto, lower=1.0)
-    
+
     with pm.Model() as model:
-        phi = pm.Uniform('phi', lower=0.0, upper=1.0)
-        kappa = BoundedKappa('kappa', alpha=1.0001, m=1.5)
-        thetas = pm.Beta('thetas', alpha=phi*kappa, beta=(1.0-phi)*kappa, shape=N)
-        ys = pm.Binomial('ys', n=atbats, p=thetas, observed=hits)
+        phi = pm.Uniform("phi", lower=0.0, upper=1.0)
+        kappa = BoundedKappa("kappa", alpha=1.0001, m=1.5)
+        thetas = pm.Beta("thetas", alpha=phi * kappa, beta=(1.0 - phi) * kappa, shape=N)
+        ys = pm.Binomial("ys", n=atbats, p=thetas, observed=hits)
     return model
 
+
 def run(n=2000):
     model = build_model()
     with model:
         trace = pm.sample(n, target_accept=0.99)
 
     pm.traceplot(trace)
 
-if __name__ == '__main__':
+
+if __name__ == "__main__":
     run()

pymc3/examples/custom_dists.py

@@ -22,23 +22,25 @@
 # add scatter to points
 xdata = np.random.normal(xdata, 10)
 ydata = np.random.normal(ydata, 10)
-data = {'x': xdata, 'y': ydata}
+data = {"x": xdata, "y": ydata}
 
 # define loglikelihood outside of the model context, otherwise cores wont work:
 # Lambdas defined in local namespace are not picklable (see issue #1995)
 def loglike1(value):
-    return -1.5 * tt.log(1 + value**2)
+    return -1.5 * tt.log(1 + value ** 2)
+
+
 def loglike2(value):
     return -tt.log(tt.abs_(value))
 
+
 with pm.Model() as model:
-    alpha = pm.Normal('intercept', mu=0, sigma=100)
+    alpha = pm.Normal("intercept", mu=0, sigma=100)
     # Create custom densities
-    beta = pm.DensityDist('slope', loglike1, testval=0)
-    sigma = pm.DensityDist('sigma', loglike2, testval=1)
+    beta = pm.DensityDist("slope", loglike1, testval=0)
+    sigma = pm.DensityDist("sigma", loglike2, testval=1)
     # Create likelihood
-    like = pm.Normal('y_est', mu=alpha + beta *
-                        xdata, sigma=sigma, observed=ydata)
+    like = pm.Normal("y_est", mu=alpha + beta * xdata, sigma=sigma, observed=ydata)
 
     trace = pm.sample(2000, cores=2)
 
@@ -47,10 +49,11 @@ def loglike2(value):
 # Create some convenience routines for plotting
 # All functions below written by Jake Vanderplas
 
+
 def compute_sigma_level(trace1, trace2, nbins=20):
     """From a set of traces, bin by number of standard deviations"""
     L, xbins, ybins = np.histogram2d(trace1, trace2, nbins)
-    L[L == 0] = 1E-16
+    L[L == 0] = 1e-16
 
     shape = L.shape
     L = L.ravel()
@@ -73,37 +76,36 @@ def plot_MCMC_trace(ax, xdata, ydata, trace, scatter=False, **kwargs):
     xbins, ybins, sigma = compute_sigma_level(trace[0], trace[1])
     ax.contour(xbins, ybins, sigma.T, levels=[0.683, 0.955], **kwargs)
     if scatter:
-        ax.plot(trace[0], trace[1], ',k', alpha=0.1)
-    ax.set_xlabel(r'$\alpha$')
-    ax.set_ylabel(r'$\beta$')
+        ax.plot(trace[0], trace[1], ",k", alpha=0.1)
+    ax.set_xlabel(r"$\alpha$")
+    ax.set_ylabel(r"$\beta$")
 
 
 def plot_MCMC_model(ax, xdata, ydata, trace):
     """Plot the linear model and 2sigma contours"""
-    ax.plot(xdata, ydata, 'ok')
+    ax.plot(xdata, ydata, "ok")
 
     alpha, beta = trace[:2]
     xfit = np.linspace(-20, 120, 10)
     yfit = alpha[:, None] + beta[:, None] * xfit
     mu = yfit.mean(0)
     sig = 2 * yfit.std(0)
 
-    ax.plot(xfit, mu, '-k')
-    ax.fill_between(xfit, mu - sig, mu + sig, color='lightgray')
+    ax.plot(xfit, mu, "-k")
+    ax.fill_between(xfit, mu - sig, mu + sig, color="lightgray")
 
-    ax.set_xlabel('x')
-    ax.set_ylabel('y')
+    ax.set_xlabel("x")
+    ax.set_ylabel("y")
 
 
-def plot_MCMC_results(xdata, ydata, trace, colors='k'):
+def plot_MCMC_results(xdata, ydata, trace, colors="k"):
     """Plot both the trace and the model together"""
     _, ax = plt.subplots(1, 2, figsize=(10, 4))
     plot_MCMC_trace(ax[0], xdata, ydata, trace, True, colors=colors)
     plot_MCMC_model(ax[1], xdata, ydata, trace)
 
-pymc3_trace = [trace['intercept'],
-               trace['slope'],
-               trace['sigma']]
+
+pymc3_trace = [trace["intercept"], trace["slope"], trace["sigma"]]
 
 plot_MCMC_results(xdata, ydata, pymc3_trace)
 plt.show()