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(N2), 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!
Hey, fantastic! Has this been pushed to stan-dev/stan repo already?
The doc update has been merged with the develop branch. I didn’t recall there being a code example checked in anywhere.
Hi!
Thank you a lot for sharing this massive computational improvement! Can this framework be extended for multiple change-points?