Composable recurrences

Build renewal processes, autoregressions and reporting delays from two fast, differentiable operators that compose.
Two operators, many models
Built for fast reverse-mode gradients
No epidemiology built in
renewal_then_delay.jl 
using ComposableRecurrences
using ComposableRecurrences: Depletion, Floor, Add, with_state

# Generation interval, lag 1 first: weights on y[t-1], y[t-2], y[t-3]
g = [0.6, 0.3, 0.1]
# Two strata mixed by a contact matrix
K = [0.9 0.1; 0.2 0.8]
# Modifiers are structs in a tuple, applied in order after each step
N = PerStratum([1_000.0, 500.0])
infections = Recurrence(g; coupling = K,
    modifiers = (Depletion(N, Floor()), Add(0.5)))

# Reporting delay, lag 0 first: P(delay = 0, 1, 2, 3)
reports = Convolution([0.0, 0.5, 0.3, 0.2])

# Operators are callable. R_t is the gain, one column per day
R = fill(1.1, 2, 10)
y = infections(R; history = ones(2, 3))
cases = reports(y)

# Keep the state after day 5, then resume for a forecast
y1, st = with_state(infections, R; history = ones(2, 3), stop = 5)
y2 = infections(R; state = st)
y ≈ hcat(y1, y2)

The example uses API v0.1 of ComposableRecurrences.jl, which is landing on the main branch now. The package is not registered yet, so install it from the repository to try it.

The approach

Renewal processes, random walks, autoregressions and reporting delays share one structure. Each step reads a window of past values, weights it by a kernel, and produces the next value. This approach builds all of them from two operators rather than one bespoke loop per model.

A Recurrence feeds its output back into its own history. At each step it convolves the last L values with its kernel, mixes strata through a coupling, scales by a gain, adds an input, and passes the result through a tuple of modifiers. An operator is built from structs, and those structs compose. A renewal process is a Recurrence with the generation interval as kernel and R_t as gain. A random walk is a Recurrence([1.0]) with the noise as the added input.

A Convolution has no feedback. It weights current and past inputs by a kernel, which covers reporting delays, moving averages and occupancy.

The package has no epidemiological assumptions. It takes numbers in (kernels, gains, histories) and returns series out. Modelling packages build their own named steps on top.

Lag-first kernels, callable operators

Every kernel is written lag first, the way a generation interval or AR coefficients are written down. A Recurrence kernel starts at lag 1 and a Convolution kernel at lag 0, so nothing is ever reversed by hand. Operators are built once, outside the model, and called with the per-day inputs that are usually sampled. Chaining is function composition, as in reports(infections(R; history)). Every time-indexed input is read at absolute time, so with_state returns the output with its state and a later call resumes from state.

Strata, time variation and coupling

Wrappers name the extra axis of a coefficient or parameter. PerStratum gives each stratum its own value, and TimeVarying lets a kernel, coupling or parameter change by day. TimeVarying(P, Primary()) indexes a delay by the day of the primary event rather than the day it is observed. Pairwise is a kernel that gives each pair of strata its own lags, so it replaces the coupling. Otherwise the coupling mixes strata before the gain, as a contact or mobility matrix does. Couplings can be dense, sparse or diagonal matrices, and any real element type is accepted, so Float32 and dual numbers pass through.

Modifiers

Modifiers are structs composed in a tuple. They are applied in order after each step and carry their own state. They handle what the linear core cannot. Depletion(N) depletes susceptibles by a hazard, and Depletion(N, Floor()) by a floored fraction. Add(b) adds an input at that point in the tuple, such as imports after depletion. Redistribute moves cases between strata and Clamp bounds them. Variants such as Floor() are structs too. A new modifier or variant is a struct with a forward method for its step, so a modelling package can add its own.

Plug-in adjoints

Reverse-mode AD is slow when each step copies its lag window into a new state. The operators own a history buffer instead, which avoids that copy. Native adjoints will plug in per operator, modifier or coupling by adding a pullback! method, with no AD declarations from the user. NoAdjoint(op) turns an adjoint off, so its gain can be timed and the adjoint dropped if a backend catches up. A benchmark script in the repository times plain AD gradients of the operators under Mooncake and Enzyme.

How it fits with the other approaches

The operators are a foundation for the other approaches. Composable Turing models will build its models on them. Its renewal, delay and latent process steps become calls to Recurrence and Convolution, with the priors and observation models staying in the Turing layer. The composed distributions stack will use them inside its packages. ComposedDistributions can run a renewal across a distribution tree, taking each step’s kernel from the tree (#82). CensoredDistributions may build on them in the same way. The operators depend on none of these packages, so they can also be used in other models.

Package map

Arrows point from a package to the packages that use it. The dashed grey arrow is a test and documentation dependency, and it is the only link in place today. The dashed teal arrows are planned and not yet implemented.

None of the three packages depends on ComposableRecurrences yet. Today each runs its own renewal and delay loops. Moving them onto one set of operators gives one forward kernel and one adjoint to maintain and test.

Status

This approach is under development. The core operators and the built-in modifiers are merged on the main branch of ComposableRecurrences.jl. These are Recurrence, Convolution, the TimeVarying, PerStratum and Pairwise wrappers, the coupling and modifier interfaces, and Depletion, Imports, Redistribute and Clamp. The test suite differentiates them with ForwardDiff, ReverseDiff, Mooncake and Enzyme through plain AD. Use-case tests check them against copies of the ComposableTuringIDModels and other model steps they will replace.

The API v0.1 consolidation is in progress. It makes every option a struct, renames Imports to Add, makes Pairwise a kernel, and replaces return_state with with_state. The native Mooncake and Enzyme adjoints are also in progress. A first registered release follows them, and then the port of ComposableTuringIDModels onto the operators. An FFT method for long fixed kernels is a later extension.

Where next