Cumulative Sums as Matrix Multiplication
I am in the process of drafting another post about what I call "decaying Bayes Theorem," which is a method for weighting recent data more heavily than older data during sequential updates of Bayesian statistical models. The implementation of decayed Bayesian updates uses a computational trick that has been useful many times, so I decided break it out into its own short post to avoid it getting lost in the broader points of the forthcoming decaying Bayes post.
The basic insight is that cumulative sums can be written as matrix multiplication by triangular matrices. If we have
$$\mathbf{x} = (x_1, x_2, \ldots, x_n) \in \mathbb{R}^n,$$
the cumulative sum of $x$ is defined as
$$\mathbf{c} = (x_1, x_1 + x_2, \ldots, x_1 + x_2 + \cdots + x_n) \in \mathbb{R}^2.$$ If $U$ is the upper triangular matrix of ones
$$U = \begin{pmatrix} 1 & 1 & \cdots & 1 & 1 \\ 0 & 1 & \cdots & 1 & 1 \\ & & \vdots & & \\ 0 & 0 & \cdots & 1 & 1 \\ 0 & 0 & \cdots & 0 & 1 \end{pmatrix} \in \mathbb{R}^{n \times n},$$
we have that
$$\mathbf{c} = \mathbf{x}^{\top} U.$$ More generally, if we have a matrix of values $X \in \mathbb{R}^{m \times n}$ and let $C \in \mathbb{R}^{m \times n}$ be defined as the cumulative sum of the rows of $X$:
$$C_{i, j} = \sum_{k = 1}^j X_{i, k},$$
we have
$$C = X U.$$
We can validate this calculation numerically as follows.
First we make the necessary Python imports and do some light configuration.
from hypothesis import assume, given
from hypothesis.extra.numpy import arrays, array_shapes
from hypothesis.strategies import floats
from warnings import filterwarnings
import numpy as np
from scipy.linalg import toeplitz
filterwarnings("ignore", category=RuntimeWarning)
We will use the property-based testing library Hypothesis to validate the claims in this post.
Property-based testing works by sampling random plausible inputs to your function and testing invariants on those inputs. The canonical example is that reversing a list twice should result in the original list. In our case, we will calculate the cumulative sum two ways:
- with the native NumPy method and
- our (possibly generalzed) matrix multplication method,
and ensuring the results are close.
First we define a function that uses matrix multiplication with an upper triangular matrix of ones as described above to calculate the cumulative sum of its argument's rows.
def matmul_cumsum(X):
_, N = X.shape
U = np.triu(np.ones((N, N)))
return X @ U
Next we define Hypothesis strategies for generating the matix X. We ensure that $X$ is two dimensonal and that each of ts rows have finite sum.
matrix_shapes = array_shapes(min_dims=2, max_dims=2)
finite_row_sum_matrices = arrays(np.float64, matrix_shapes).filter(
lambda x: np.isfinite(x.sum(axis=1)).all()
)
We then define a function that takes a matrix whose rows have finite sum, and checks that the result of calculating the cumulative sum of its columns gives similar results.
@given(finite_row_sum_matrices)
def test_cumsum(X):
np.testing.assert_allclose(matmul_cumsum(X), X.cumsum(axis=1), rtol=np.inf)
Running this function causes Hypothesis to check the assertion for many random inputs.
test_cumsum()
Since no exceptions were raised, we have confirmed that our matrix multiplication version of the cumulative sum is reasonably accurate.
So far, this reimplementation of cumulative sum is just a mathematical curiosity. This perspective, however, becomes computationally valuable when we want to generalize the idea of cumulative sum. In the forthcoming decaying Bayes post, we will have reason to consider sums of the form
$$C^{\lambda}_{i, j} = \sum_{k = 1}^j \lambda^{j - k + 1} \cdot X_{i, k}$$
for $0 \leq \lambda \leq 1.$ Generalizing the above perspective on cumulative sums as matrix multiplication, if we define
$$U^{\lambda} = \begin{pmatrix} \lambda & \lambda^2 & \cdots & \lambda^{n - 1} & \lambda^n \\ 0 & \lambda & \cdots & \lambda^{n - 2} & \lambda^{n - 1} \\ & & \vdots & & \\ 0 & 0 & \cdots & \lambda & \lambda^2 \\ 0 & 0 & \cdots & 0 & \lambda \end{pmatrix} \in \mathbb{R}^{n \times n},$$
we get that
$$C^{\lambda} = X U^{\lambda}.$$
The following functions implement this operation using matrix multiplication with a suitably-constructed upper diagonal matrix, and test that implementation against naive iteration.
def matmul_decayed_cumsum(X, λ):
_, N = X.shape
Uλ = np.triu(λ ** toeplitz(1 + np.arange(N)))
return X @ Uλ
Note that this test relies on the fact that
$$C_{i, j}^{\lambda} = \lambda \cdot (C_{i, j - 1}^{\lambda} + X_{i, j})$$
for $1 < j \leq n$.
λs = floats(0, 1)
@given(finite_row_sum_matrices, λs)
def test_decayed_cumsum(X, λ):
_, N = X.shape
expected = np.empty_like(X)
expected[:, 0] = λ * X[:, 0]
for j in range(1, N):
expected[:, j] = λ * (expected[:, j - 1] + X[:, j])
np.testing.assert_allclose(matmul_decayed_cumsum(X, λ), expected, rtol=np.inf)
test_decayed_cumsum()
We see that no exceptions were raised, and our matrix multiplication implementation is reasonably accurate.
This post is available as a Jupyter notebook here.
%load_ext watermark
%watermark -n -u -v -iv