Improve sampling performance

[tovrstra]

A few ideas (can be combined)

• As mentioned in Algorithm for sampling power-law covariance kernels #20, we can use scipy.signal.lfilter to speed up the sampling. As far as I understand its implementation, the loop over time steps is implemented with specialized low-level code in SciPy, which should be faster than our Python implementation.
• We can vectorize all operations over different independent sequences (parameter $M$ in the documentation).

Some basic testing would be useful to verify that such changes do actually improve performance. There may be other bottlenecks.

[RobbeBohy]

After some basic testing using timeit, the use of scipy.signal.lfilter leads to a speedup of $\approx 2.7 x$ (with nstep= 65536 and nseed= 16) for the exponential kernels. Based on my understanding of this function, this approach cannot be applied to any of the SHO-based kernels, as these are no longer defined by scalar SDEs but using matrix representations.

Regarding the vectorization over the different independent sequences $M$, the current implementation already satisfies this. For example, in the exponential kernel only a single for-loop is used and all nseq repetitions are handled simultaneously:

def sample(self, nseq: int, nstep: int, rng: np.random.Generator) -> NDArray[float]:
    traj = np.zeros((nseq, nstep))
    var = self.a0 / (2 * self.tau)
    noise = rng.normal(0, np.sqrt(var * (1 - np.exp(-2 / self.tau))), size=(nseq, nstep))
    traj[:, 0] = noise[:, 0]
    prop = np.exp(-1 / self.tau)
    for istep in range(1, nstep):
        traj[:, istep] = noise[:, istep] + prop * traj[:, istep - 1]

    return traj

[tovrstra]

I'm looking into this and there may be a way out with lfilter for the SHO with some non-trivial tricks. Before digging deeper, it may be worth checking the following first:

Cache memory efficiency can be improved by swapping the order of the indexes: (nstep, nseq). You can still transpose in the end, so the arrays get stored in the ZIP file as before. The reason this might speed up the calculation (also for the SHO) is that all the data accessed in one iteration is located in nearby memory. (NumPy uses row-major storage by default.)

[tovrstra]

I've asked Claude how to rewrite the SHO implementation with lfilter. It came up with something way too complicated, but it had one crucial idea that I've isolated and that alone was sufficient. This is how it works:

Notation:

• $x(t)$ is the sampled sequences (2D time-dependent vector)
• $n(t)$ is the pre-generated noise (also 2D time-dependent vector)
• $A$ is the matrix exponential

You can combine two iterations into one by substituting the recurrence relation into itself, and then you can decouple them using the Cayley–Hamilton theorem:

$$ \begin{aligned} x(t) &= A x(t-1) + n(t) \\ x(t) &= A (A x(t-2) + n(t-1)) + n(t) \\ x(t) &= (\mathrm{Tr}(A) A - |A| I) x(t-2) + A n(t-1) + n(t) \\ x(t) &= -|A| x(t-2) + \mathrm{Tr}(A) x(t-1) + (A - \mathrm{Tr}(A) I) n(t-1) + n(t) \end{aligned} $$

This can be implemented with a single call to scipy.signal.lfilter.

Demo implementation with verification: (I've already swapped the indexes)

import numpy as np
import scipy as sp


cases = [
    (1.2, 0.15, 0.2),  # overdamped
    (1.2, 0.04, 0.5),  # critically damped
    (1.2, 0.03, 1.4),  # underdamped
]

for a0, f0, q in cases:
    # Copied from kernel.py:
    omega0 = 2 * np.pi * f0
    theta = np.array([[0, 1], [-(omega0**2), -omega0 / q]])
    b = np.array([[0, 0], [0, 0.5 * a0 * omega0**4]])
    matexp = sp.linalg.expm(theta)
    upper = 2 * np.dot(matexp, np.dot(b, matexp.T))
    lower = 2 * b
    covar = sp.linalg.solve_continuous_lyapunov(theta, upper - lower)

    # Prepare some noise. Notice the change in indexes.
    # Time step index goes first for cache efficiency.
    nstep = 10
    nseq = 3
    rng = np.random.default_rng(0)
    noise = rng.multivariate_normal(np.zeros(2), covar, size=(nstep, nseq))

    # Preprocess the noise
    trace = np.trace(matexp)
    det = np.linalg.det(matexp)
    noise_tf = noise.copy()
    noise_tf[1:] -= trace * noise[:-1]
    noise_tf[1:] += np.einsum("ic,tsc->tsi", matexp, noise[:-1])

    # Use lfilter to compute the second-order autoregressive process.
    traj_vec = sp.signal.lfilter([1], [1, -trace, det], noise_tf, axis=0)

    # Compute with the explicit Python loop for validation.
    traj_loop = np.zeros((nstep, nseq, 2))
    traj_loop[0] = noise[0]
    for istep in range(1, nstep):
        traj_loop[istep] = noise[istep] + np.einsum("ic,sc->si", matexp, traj_loop[istep - 1])

    # Compare the two trajectories.
    print(traj_loop.reshape(nstep, nseq * 2))
    print(traj_vec.reshape(nstep, nseq * 2))
    print(np.allclose(traj_vec, traj_loop))

We should check whether this has been published already somewhere. I could not immediately find a good reference with this idea, so I'm not sure how new it is. We'll also have to test it more carefullly, just to be sure.

[tovrstra]

I've gone through several references and most treat more general cases, without the simplicity of what I posted above. In time series modeling, this would be called the transformation of a 2D VAR(1) = vector autoregressive process into a ARMA(2,1) = autoregressive moving average process.

This thesis reviews the basic theory: https://pure.tue.nl/ws/files/67736084/759210-1.pdf The code above is essentially a special case of Eq. (1.15) in this thesis. As far as I can tell, the original reference for this equation is:

Arnold Zellner, Franz Palm, Time series analysis and simultaneous equation econometric models, Journal of Econometrics, Volume 2, Issue 1, 1974, Pages 17-54, https://doi.org/10.1016/0304-4076(74)90028-1

We could cite it for completeness. It would still make sense to just include the derivation above, because it is much simpler.