Repository navigation
docs: add composable recurrences as the third approach #31
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,279 @@ | ||
| --- | ||
| title: "Composable recurrences" | ||
| description: "Build renewal processes, autoregressions and reporting delays from two fast, differentiable operators that compose." | ||
| toc: true | ||
| toc-location: left | ||
| page-layout: full | ||
| --- | ||
|
|
||
| :::::: {.side-space} | ||
|
|
||
| ::: {.what-is-turing-grid .approach-hero .section-start-space .section-end-space} | ||
|
|
||
| ::: {.feature-text-column} | ||
|
|
||
| ::: {.feature-item} | ||
| ::: {.feature-item-header} | ||
|
|
||
| ```{=html} | ||
| <span class="feature-icon icon-composable"></span> | ||
| ``` | ||
|
|
||
| ##### Two operators, many models | ||
|
|
||
| ::: | ||
|
|
||
| ::: | ||
|
|
||
| ::: {.feature-item} | ||
| ::: {.feature-item-header} | ||
|
|
||
| ```{=html} | ||
| <span class="feature-icon icon-uncertainty"></span> | ||
| ``` | ||
|
|
||
| ##### Built for fast reverse-mode gradients | ||
|
|
||
| ::: | ||
|
|
||
| ::: | ||
|
|
||
| ::: {.feature-item} | ||
| ::: {.feature-item-header} | ||
|
|
||
| ```{=html} | ||
| <span class="feature-icon icon-response"></span> | ||
| ``` | ||
|
|
||
| ##### No epidemiology built in | ||
|
|
||
| ::: | ||
|
|
||
| ::: | ||
|
|
||
| ::: | ||
|
|
||
| ::: {.code-example-column} | ||
| ::: {.code-window .box} | ||
| ::: {.code-window-header} | ||
| ::: {.code-window-dots} | ||
|
|
||
| ```{=html} | ||
| <span class="dot red"></span><span class="dot yellow"></span><span class="dot green"></span> | ||
| ``` | ||
|
|
||
| ::: | ||
|
|
||
| ::: {.code-window-title} | ||
|
|
||
| ```{=html} | ||
| <span>renewal_then_delay.jl</span> | ||
| ``` | ||
|
|
||
| ::: | ||
|
|
||
| ::: | ||
|
|
||
| ```{.julia .code-overflow-scroll} | ||
| using ComposableRecurrences | ||
|
|
||
| # 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) | ||
|
|
||
| # 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) | ||
|
|
||
| # 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) | ||
| y ≈ hcat(y1, y2) | ||
| ``` | ||
|
|
||
| ::: | ||
| ::: | ||
|
|
||
| ::: | ||
|
|
||
| The example runs on the main branch of [ComposableRecurrences.jl](https://github.com/EpiAware/ComposableRecurrences.jl). | ||
| 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. | ||
| 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))`. | ||
|
|
||
| ### Strata, time variation and coupling | ||
|
|
||
| Wrappers name the extra axis of a coefficient. | ||
|
There was a problem hiding this comment. suggestion The 1.75 to 15 times and 8 to 25 times speed-ups come from prototypes, and the page links to no benchmark or repository script for them. The Status section says the native adjoints are still only planned, so the 8 to 25 times figure describes code that is not in the package. Link the benchmark source, or cut the figures. If they stay, reword the second one to say it is a prototype result that is not yet reproducible from the package.
Collaborator
Author
There was a problem hiding this comment. Fixed in 8ee3780. The prototype figures are cut. The section now links benchmark/gradients.jl on main and says only what it does, which is time plain AD gradients of the operators under Mooncake and Enzyme. It also no longer claims plain AD is already fast. |
||
| `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. | ||
| 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. | ||
|
|
||
| ### 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. | ||
| `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. | ||
| 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. | ||
|
|
||
| ## Package map | ||
|
There was a problem hiding this comment. suggestion The package map takes about 90 lines of hand-placed SVG coordinates. Most of it repeats what |
||
|
|
||
| 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. | ||
|
|
||
| ```{=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> | ||
| <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> | ||
|
There was a problem hiding this comment. suggestion The box title is 'Lowered / Reparam.', but its subtitle says 'DistributionsInference' and 'no touch point'. The title and subtitle name different packages, and the packages table lists them separately. Name one package, or split the box. If the point is that these packages do not touch ComposableRecurrences, drop the box. The diagram is already crowded. |
||
| <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> | ||
|
|
||
| <!-- edges --> | ||
| <path class="cr-e" d="M134,144 C134,190 330,190 360,222" marker-end="url(#cr-map-a)"/> | ||
|
There was a problem hiding this comment. suggestion The edge from EpiAwareADTools to ComposableRecurrences is drawn with the dashed grey style, which the legend defines as a 'test-time link'. Its node text says 'no dependency either way'. It has no arrowhead and no label, so the drawing contradicts the node text. Delete this
Collaborator
Author
There was a problem hiding this comment. Fixed in 8ee3780. The unlabelled dashed edge from EpiAwareADTools is removed. The node text still says there is no dependency either way, and the legend now names the one test-time link. |
||
| <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> | ||
| </svg> | ||
| </figure> | ||
| ``` | ||
|
|
||
| The package has no runtime dependency on any other EpiAware package. | ||
| Today ComposableTuringIDModels runs its own renewal, delay and latent process loops. | ||
|
There was a problem hiding this comment. suggestion This section says ComposableTuringIDModels 'runs its own' loops, and the next sentence says moving to shared operators gives 'one adjoint to maintain'. The adjoint does not exist yet, and the port is not scheduled. Say 'would give' rather than 'gives', or cut the sentence. The Status section already covers the plan. |
||
| Moving them onto one set of operators gives one forward kernel and one adjoint to maintain and test. | ||
|
|
||
|
There was a problem hiding this comment. suggestion 'Tests check them against copies of the ComposableTuringIDModels steps' is a claim about a development branch that readers cannot inspect. The package is unreleased, and the code sample at the top uses
Collaborator
Author
There was a problem hiding this comment. Fixed in 8ee3780, with a changed premise. The core operators are now merged on ComposableRecurrences.jl main (PR #5), so the note under the hero says the example runs on main and the package is not yet registered. The use-case tests against the ComposableTuringIDModels copies exist on main but are tagged pending and skipped by default, so the Status section now says they are written but not yet switched on. The packages row is updated to match. |
||
| ## 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 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. | ||
|
|
||
| The built-in modifiers (`Depletion`, `Imports`, `Redistribute`, `Clamp`) and the native Mooncake and Enzyme adjoints are 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. | ||
|
|
||
| ## 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 | ||
| [composed distributions approach](composed-distributions.qmd). | ||
|
|
||
| :::::: | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
suggestion This hero heading claims the package is built for fast reverse-mode gradients, but the Status section says the native Mooncake and Enzyme adjoints are still in progress. The 'Plug-in adjoints' section also says the gain is still to be timed. Reword to what exists today, for example 'Differentiable by design', or 'Tested with four AD backends'. That matches the Status paragraph.