Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
187 changes: 87 additions & 100 deletions approaches/composable-recurrences.qmd
Original file line number Diff line number Diff line change
Expand Up @@ -76,12 +76,16 @@ page-layout: full

```{.julia .code-overflow-scroll}
using ComposableRecurrences
using ComposableRecurrences: Depletion, Floor, Add, with_state

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion The import line brings in Depletion, Floor, Add and with_state but not PerStratum. The example also uses Recurrence, Convolution and PerStratum. If these are not exported, the example fails as written, and if they are exported it is unclear why the four listed names are not. Confirm the export list and import consistently.


# 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]
infections = Recurrence(g; coupling = K)
# 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])
Expand All @@ -91,10 +95,9 @@ R = fill(1.1, 2, 10)
y = infections(R; history = ones(2, 3))
cases = reports(y)

# Stop, keep the state, and resume for a forecast
y1, state = infections(R[:, 1:5]; history = ones(2, 3),
return_state = true)
y2 = infections(R[:, 6:end]; history = state)
# 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)
```

Expand All @@ -103,7 +106,7 @@ y ≈ hcat(y1, y2)

:::

The example runs on the main branch of [ComposableRecurrences.jl](https://github.com/EpiAware/ComposableRecurrences.jl).
The example uses API v0.1 of [ComposableRecurrences.jl](https://github.com/EpiAware/ComposableRecurrences.jl), which is landing on the main branch now.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion "Which is landing on the main branch now" is time-relative and will go stale. It also conflicts with the Status section, which says the operators and modifiers are already merged. Say one thing in both places, for example "The example uses the v0.1 API of ComposableRecurrences.jl. Check the repository for which commit or release carries it."

The package is not registered yet, so install it from the repository to try it.

## The approach
Expand All @@ -114,6 +117,7 @@ This approach builds all of them from two operators rather than one bespoke loop

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.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion "An operator is built from structs, and those structs compose." says nothing the next paragraph and the modifiers section do not. Cut it.

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.

Expand All @@ -130,150 +134,133 @@ Every kernel is written lag first, the way a generation interval or AR coefficie
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`.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion This sentence joins two claims with "so". Absolute-time indexing does not explain with_state returning its state. Split it: "Time-indexed inputs are read at absolute time, so a resumed call needs the full series and not a slice. with_state returns the state to resume from." The example passes the full R to the resumed call, and the text should say why.


### Strata, time variation and coupling

Wrappers name the extra axis of a coefficient.
`PerStratum` gives each stratum its own kernel, `TimeVarying` lets a kernel or coupling change by day, and `Pairwise` gives each pair of strata its own lag kernel.
The coupling mixes strata before the gain, as a contact or mobility matrix does.
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 applied in order after each step and carry their own state.
They handle what the linear core cannot, such as susceptible depletion or moving cases between strata.
A modifier is a plain type with an `apply!` method, so a modelling package can add its own.
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 method, with no AD declarations from the user.
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](https://github.com/EpiAware/ComposableRecurrences.jl/blob/main/benchmark/gradients.jl) in the repository times plain AD gradients of the operators under Mooncake and Enzyme.

### How it fits with the other approaches

[Composable Turing models](composable-turing-models.qmd) will build on these operators.
The operators are a foundation for the other approaches.
[Composable Turing models](composable-turing-models.qmd) 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](composed-distributions.qmd) stack supplies the delay pmfs that become `Convolution` kernels.
The operators have no dependency on either, so they can also be used in other models.
The [composed distributions](composed-distributions.qmd) 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](https://github.com/EpiAware/ComposedDistributions.jl/issues/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

Solid arrows are data or code that will flow once the package is adopted.
The dashed grey arrow is a test-time link.
The dashed amber arrow is an optional future backend, not a dependency.
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.

```{=html}
<figure class="cr-map" role="img" aria-labelledby="cr-map-title">
<svg viewBox="0 0 960 560" xmlns="http://www.w3.org/2000/svg">
<title id="cr-map-title">EpiAware packages around ComposableRecurrences</title>
<svg viewBox="0 0 960 400" xmlns="http://www.w3.org/2000/svg">
<title id="cr-map-title">ComposableRecurrences and the EpiAware packages planned to build on it</title>
<defs>
<marker id="cr-map-a" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto-start-reverse"><path class="cr-ah" d="M0,0 L10,5 L0,10 z"/></marker>
<marker id="cr-map-am" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto-start-reverse"><path class="cr-ah cr-muted" d="M0,0 L10,5 L0,10 z"/></marker>
<marker id="cr-map-aw" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto-start-reverse"><path class="cr-ah cr-warn" d="M0,0 L10,5 L0,10 z"/></marker>
</defs>
<!-- bands -->
<rect class="cr-band" x="12" y="14" width="936" height="150" rx="8"/>
<text class="cr-band-t" x="24" y="32">DISTRIBUTIONS · DELAY PMFS</text>
<rect class="cr-band" x="12" y="196" width="936" height="130" rx="8"/>
<text class="cr-band-t" x="24" y="214">OPERATORS · TOOLING</text>
<rect class="cr-band" x="12" y="400" width="936" height="146" rx="8"/>
<text class="cr-band-t" x="24" y="418">MODELS · BACKENDS</text>

<!-- distribution layer -->
<rect class="cr-node" x="36" y="48" width="196" height="96" rx="6"/>
<text x="50" y="72" font-size="14" font-weight="600">CensoredDistributions</text>
<text class="cr-sub" x="50" y="92">censored delay distributions</text>
<text class="cr-sub" x="50" y="108">discretised to a delay pmf</text>
<text class="cr-sub" x="50" y="124">lag 0 first</text>

<rect class="cr-node" x="280" y="48" width="208" height="96" rx="6"/>
<text x="294" y="72" font-size="14" font-weight="600">ConvolvedDistributions</text>
<text class="cr-sub" x="294" y="92">sums of delay distributions</text>
<text class="cr-sub" x="294" y="108">convolved series</text>
<text class="cr-sub" x="294" y="124">time-varying delays</text>

<rect class="cr-node" x="536" y="48" width="196" height="96" rx="6"/>
<text x="550" y="72" font-size="14" font-weight="600">Composed / Modified</text>
<text class="cr-sub" x="550" y="92">chained and modified delays</text>
<text class="cr-sub" x="550" y="108">discretise, then convolve</text>

<rect class="cr-node" x="772" y="48" width="160" height="96" rx="6"/>
<text x="786" y="72" font-size="13.5" font-weight="600">Lowered / Reparam.</text>
<text class="cr-sub" x="786" y="92">DistributionsInference</text>
<text class="cr-sub" x="786" y="116">no touch point</text>

<!-- operator layer -->
<rect class="cr-node cr-focus" x="300" y="224" width="360" height="88" rx="6"/>
<text x="316" y="250" font-size="15" font-weight="700">ComposableRecurrences</text>
<text class="cr-sub" x="316" y="270">Recurrence · Convolution · TimeVarying · PerStratum · Pairwise</text>
<text class="cr-sub" x="316" y="286">in progress: Depletion · Imports · Redistribute · Clamp</text>
<text class="cr-sub" x="316" y="302">in progress: native Mooncake and Enzyme adjoints</text>

<rect class="cr-node" x="36" y="224" width="200" height="88" rx="6"/>
<text x="50" y="250" font-size="14" font-weight="600">EpiAwarePackageTools</text>
<text class="cr-sub" x="50" y="270">scaffold, Taskfile, CI</text>
<text class="cr-sub" x="50" y="286">AD test harness</text>
<text class="cr-sub" x="50" y="302">AD benchmark pages</text>

<rect class="cr-node" x="724" y="224" width="208" height="88" rx="6"/>
<text x="738" y="250" font-size="14" font-weight="600">EpiAwareADTools</text>
<text class="cr-sub" x="738" y="270">shared AD workarounds</text>
<text class="cr-sub" x="738" y="286">used where a backend needs one</text>
<text class="cr-sub" x="738" y="302">no dependency either way</text>

<!-- models and backends -->
<rect class="cr-node" x="160" y="434" width="300" height="96" rx="6"/>
<text x="174" y="458" font-size="14" font-weight="600">ComposableTuringIDModels</text>
<text class="cr-sub" x="174" y="478">renewal, susceptible depletion, imports</text>
<text class="cr-sub" x="174" y="494">mixing, latent delays, AR, MA, random walk</text>
<text class="cr-sub" x="174" y="510">own loops → Recurrence, Convolution</text>

<rect class="cr-node cr-opt" x="520" y="434" width="300" height="96" rx="6"/>
<text x="534" y="458" font-size="14" font-weight="600">Reactant</text>
<text class="cr-sub" x="534" y="478">compiled backend for the operators</text>
<text class="cr-sub" x="534" y="494">optional, not a dependency</text>
<text class="cr-sub" x="534" y="510">later extension</text>
<rect class="cr-band" x="12" y="14" width="936" height="130" rx="8"/>
<text class="cr-band-t" x="24" y="32">FOUNDATION · TOOLING</text>
<rect class="cr-band" x="12" y="236" width="936" height="150" rx="8"/>
<text class="cr-band-t" x="24" y="254">PLANNED USES</text>

<!-- foundation -->
<rect class="cr-node cr-focus" x="300" y="42" width="360" height="88" rx="6"/>
<text x="316" y="68" font-size="15" font-weight="700">ComposableRecurrences</text>
<text class="cr-sub" x="316" y="88">Recurrence · Convolution · TimeVarying · PerStratum · Pairwise</text>
<text class="cr-sub" x="316" y="104">Depletion · Add · Redistribute · Clamp</text>
<text class="cr-sub" x="316" y="120">no runtime EpiAware dependency</text>

<rect class="cr-node" x="36" y="42" width="200" height="88" rx="6"/>
<text x="50" y="68" font-size="14" font-weight="600">EpiAwarePackageTools</text>
<text class="cr-sub" x="50" y="88">scaffold, Taskfile, CI</text>
<text class="cr-sub" x="50" y="104">AD test harness</text>
<text class="cr-sub" x="50" y="120">tests and docs only</text>

<!-- planned uses -->
<rect class="cr-node" x="36" y="270" width="280" height="96" rx="6"/>
<text x="50" y="294" font-size="14" font-weight="600">ComposableTuringIDModels</text>
<text class="cr-sub" x="50" y="314">models built on the operators</text>
<text class="cr-sub" x="50" y="330">renewal, delay and latent process steps</text>
<text class="cr-sub" x="50" y="346">runs its own loops today</text>

<rect class="cr-node" x="340" y="270" width="280" height="96" rx="6"/>
<text x="354" y="294" font-size="14" font-weight="600">ComposedDistributions</text>
<text class="cr-sub" x="354" y="314">renewal across a distribution tree</text>
<text class="cr-sub" x="354" y="330">per-step kernels from the tree</text>
<text class="cr-sub" x="354" y="346">ComposedDistributions.jl#82</text>

<rect class="cr-node" x="644" y="270" width="280" height="96" rx="6"/>
<text x="658" y="294" font-size="14" font-weight="600">CensoredDistributions</text>
<text class="cr-sub" x="658" y="314">may build on the operators</text>
<text class="cr-sub" x="658" y="330">in the same way</text>

<!-- edges -->
<path class="cr-e" d="M134,144 C134,190 330,190 360,222" marker-end="url(#cr-map-a)"/>
<path class="cr-e" d="M384,146 L420,222" marker-end="url(#cr-map-a)"/>
<text class="cr-lbl" x="420" y="186">delay pmf → Convolution kernel</text>
<path class="cr-e cr-dash" d="M238,268 L298,268" marker-end="url(#cr-map-am)"/>
<text class="cr-lbl" x="242" y="258">tests</text>
<path class="cr-e" d="M400,314 L330,432" marker-end="url(#cr-map-a)"/>
<text class="cr-lbl" x="196" y="384">Recurrence, Convolution</text>
<path class="cr-e cr-warn" d="M560,314 L640,432" marker-end="url(#cr-map-aw)"/>
<text class="cr-lbl" x="624" y="384">optional backend</text>
<path class="cr-e cr-dash" d="M238,86 L298,86" marker-end="url(#cr-map-am)"/>
<text class="cr-lbl" x="242" y="76">tests, docs</text>
<path class="cr-e cr-plan" d="M420,132 L200,268" marker-end="url(#cr-map-a)"/>
<path class="cr-e cr-plan" d="M480,132 L480,268" marker-end="url(#cr-map-a)"/>
<path class="cr-e cr-plan" d="M540,132 L760,268" marker-end="url(#cr-map-a)"/>
<text class="cr-lbl" x="490" y="196">planned</text>
</svg>
</figure>
```

The package has no runtime dependency on any other EpiAware package.
Today ComposableTuringIDModels runs its own renewal, delay and latent process loops.
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 are merged on the main branch of ComposableRecurrences.jl.
These are `Recurrence`, `Convolution`, the `TimeVarying`, `PerStratum` and `Pairwise` wrappers, and the coupling and modifier interfaces.
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`.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion The page names Imports as merged here, but the example, the modifiers section and the package map all use Add. Line 252 then says Imports is renamed to Add. A reader meeting Imports here will not find it in the v0.1 API the example uses. List Add here and drop the rename from line 252.

The test suite differentiates them with ForwardDiff, ReverseDiff, Mooncake and Enzyme through plain AD.
Use-case tests against copies of the ComposableTuringIDModels steps they will replace are written but not yet switched on.
Use-case tests check them against copies of the ComposableTuringIDModels and other model steps they will replace.

The built-in modifiers (`Depletion`, `Imports`, `Redistribute`, `Clamp`) and the native Mooncake and Enzyme adjoints are in progress.
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`.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion This paragraph is a changelog ("renames Imports to Add", "replaces return_state with with_state"). Readers of the page never saw the old names. Say what the API is, for example: "v0.1 makes every option a struct and makes Pairwise a kernel. A release follows once it and the native adjoints are done."

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 and a Reactant backend are later extensions.
An FFT method for long fixed kernels is a later extension.

## Where next

- The [ComposableRecurrences.jl repository](https://github.com/EpiAware/ComposableRecurrences.jl)
tracks development, and its documentation will follow the first release.
- Building a full model with priors and observations? See the
[composable Turing models approach](composable-turing-models.qmd).
- Need the delay distributions that feed a `Convolution`? See the
- Composing delay distributions? See the
[composed distributions approach](composed-distributions.qmd).

::::::
6 changes: 1 addition & 5 deletions styles.css
Original file line number Diff line number Diff line change
Expand Up @@ -330,17 +330,13 @@ h1, h2, h3 { letter-spacing: -0.02em; }
.cr-map .cr-band-t {
fill-opacity: 0.6; font-size: 10.5px; letter-spacing: 0.08em;
}
.cr-map .cr-node.cr-opt {
stroke: var(--epi-accent); stroke-dasharray: 5 4;
}
.cr-map .cr-e { stroke: var(--epi-primary); stroke-width: 1.5; fill: none; }
.cr-map .cr-e.cr-dash {
stroke: currentColor; stroke-opacity: 0.5; stroke-dasharray: 5 4;
}
.cr-map .cr-e.cr-warn { stroke: var(--epi-accent); stroke-dasharray: 5 4; }
.cr-map .cr-e.cr-plan { stroke-dasharray: 6 4; }
.cr-map .cr-ah { fill: var(--epi-primary); }
.cr-map .cr-ah.cr-muted { fill: currentColor; fill-opacity: 0.5; }
.cr-map .cr-ah.cr-warn { fill: var(--epi-accent); }

/* Value propositions + foundations grid ----------------------------------- */
.feature-grid {
Expand Down
Loading