Merge pull request #1181 from pymc-devs/stan_model_arma

Stan model arma

Author: springcoil

pymc3/examples/arma_example.py

@@ -0,0 +1,99 @@
+
+from pymc3 import Normal, sample, Model, traceplot, plots, NUTS, Potential, variational, Cauchy, find_MAP, Slice, HalfCauchy
+from theano import scan, shared
+from scipy import optimize
+import theano.tensor as T
+
+
+import numpy as np
+"""
+ARMA example
+It is interesting to note just how much more compact this is that the original STAN example
+
+The original implementation is in the STAN documentation by Gelman et al and is reproduced below
+
+
+Example from STAN- slightly altered
+
+data {
+  int<lower=1> T;
+  real y[T];
+}
+parameters {
+    // assume err[0] == 0
+}
+nu[t] <- mu + phi * y[t-1] + theta * err[t-1];
+  err[t] <- y[t] - nu[t];
+}
+mu ~ normal(0,10);
+phi ~ normal(0,2);
+theta ~ normal(0,2);
+  real mu;
+  real phi;
+  real theta;
+  real<lower=0> sigma;
+} model {
+  vector[T] nu;
+  vector[T] err;
+  nu[1] <- mu + phi * mu;
+  err[1] <- y[1] - nu[1];
+  for (t in 2:T) {
+    // num observations
+    // observed outputs
+    // mean coeff
+    // autoregression coeff
+    // moving avg coeff
+    // noise scale
+    // prediction for time t
+    // error for time t
+sigma ~ cauchy(0,5);
+err ~ normal(0,sigma);
+// priors
+// likelihood
+Ported to PyMC3 by Peadar Coyle and Chris Fonnesbeck (c) 2016.
+"""
+
+t = np.array([1, 2, 4,5,6,8, 19, 18, 12])
+
+y = shared(np.array([15, 10, 16, 11, 9, 11, 10, 18], dtype=np.float32))
+
+
+with Model() as arma_model:
+
+    sigma = HalfCauchy('sigma', 5)
+    theta = Normal('theta', 0, sd=2)
+    phi = Normal('phi', 0, sd=2)
+    mu = Normal('mu', 0, sd=10)
+
+    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])
+                  
+    like = Potential('like', Normal.dist(0, sd=sigma).logp(err))
+
+with arma_model:
+    mu, sds, elbo = variational.advi(n=2000)
+
+
+def run(n=1000):
+    if n == "short":
+        n = 50
+    with arma_model:
+
+        trace = sample(1000)
+
+    burn = n/10
+
+    traceplot(trace[burn:])
+    plots.summary(trace[burn:])
+
+
+if __name__ == '__main__':
+    run()