Stan Coding Corner: O(N) Change-Point Program with Clever Forward-Backward Calculation

(This article was originally published at Statistical Modeling, Causal Inference, and Social Science, and syndicated at StatsBlogs.)

It’s so much fun to work in open source. Luke Wiklendt sent along this improved code for a change-point model calculation in Stan. With N data points in the time series, the version in the manual is O(N²), whereas the improved version is O(N). In practice, Luke says

[the new code] results in a dramatic speed up. On my machine 10k samples of the original took 8m53s and the new one took 33s.

Now that’s the kind of improvement I can get behind, both in theory and in practice. As soon as I understood his code, I slapped my head and yelled “D’oh!”. Of course there’s a linear algorithm, because it can be encoded as a simple state-space HMM. Everyone (who spent 25 years working on natural language processing and speech recognition) knows the HMM forward-backward algorithm is linear in N.

Quadratic version from Stan manual v2.9

Here’s my original version.

transformed parameters {
  vector[T] lp;
  lp <- rep_vector(log_unif, T);
  for (s in 1:T)
    for (t in 1:T)
      lp[s] <- lp[s] + poisson_log(D[t], if_else(t < s, e, l));
}

Linear version from Luke Wiklendt

And here's Luke's improved version, which he prefaces by saying

Given a vector v and two indices i and j, the sum from i to j is performed by v[i] + v[i+1] + ... + v[j-1] + v[j]. It can also be performed by precalculating the cumulative sum of v (call it cv), and then the sum from i to j is cv[j] - cv[i].

That's what the engineers call "dynamic programming."

transformed parameters {
    vector[T+1] lp_e;
    vector[T+1] lp_l;
    vector[T] lp;

    lp_e[1] <- 0;
    lp_l[1] <- 0;
    for (t in 1:T) {
        lp_e[t + 1] <- lp_e[t] + poisson_log(D[t], e);
        lp_l[t + 1] <- lp_l[t] + poisson_log(D[t], l);
    }

    lp <- rep_vector(log_unif + lp_l[T + 1], T) + head(lp_e, T) - head(lp_l, T);
}

The key for me was realizing this is just the forward-backward algorithm of HMMs. The forward values are lp_e and the backward values are lp_l[T+1] - lp_l, so lp is the complete forward+backward values (plus log_unif for the prior).

Thanks, Luke!

The post Stan Coding Corner: O(N) Change-Point Program with Clever Forward-Backward Calculation appeared first on Statistical Modeling, Causal Inference, and Social Science.

Please comment on the article here: Statistical Modeling, Causal Inference, and Social Science

Tags: Uncategorized