Add Case Study 3: ODE-based SIR model with deterministic and stochastic variants #90

seabbs · 2025-12-02T19:12:24Z

Summary

Adds Case Study 3 demonstrating ODE-based SIR compartmental modelling using EpiAware
Shows two observation model approaches: deterministic (direct Poisson link) and stochastic (AR(1) time-varying ascertainment)
Uses TransformObservationModel to compose population scaling with observation models
Creates 2x3 comparison figure panel following established case study format
Updates index.qmd with figure reference

Key implementation details

ODEProcess with SciML ODE solvers for SIR dynamics
TransformObservationModel wraps observation models with softplus population scaling
Ascertainment with AR(1) latent process for time-varying ascertainment in stochastic model
Based on Chatzilena et al. boarding school influenza dataset

Test plan

Run task render-case-studies to verify case study renders
Run task render-index to verify main document renders with figure
Check generated figures in figures/ directory

Add ODE-based SIR compartmental model case study demonstrating: - ODEProcess with SciML ODE solvers for SIR dynamics - TransformObservationModel for population scaling - Deterministic model with direct Poisson link - Stochastic model with AR(1) time-varying ascertainment - 2x3 comparison figure panel following established format Update index.qmd with figure reference and description.

@set

- Move ODE packages to case study section (not global setup) - Restructure model section with equations together - Introduce Ascertainment as new observation model modifier - Use @set from Accessors.jl for cleaner model composition - Simplify sol2infs: ODE returns proportions, obs model handles counts - Update index.qmd to reflect Ascertainment introduction

seabbs · 2025-12-02T19:34:58Z

Comparison with Draft PR #79

This PR builds on the work in draft PR #79, addressing the outstanding tasks whilst making structural changes to better fit the established case study format.

Key Differences from PR #79

Structural Changes

Aspect	PR #79	This PR (#90)
Location	All code in `index.qmd`	Code in `case-studies.qmd`, figure reference in `index.qmd`
Figure format	2x2 panel	2x3 panel (matching Case Studies 1 & 2)
Sample counts	`ndraws=100`, `nadapts=250` (temporary)	Uses standard `ndraws=1000`, `nadapts=1000`
YAML changes	Added `julia: exeflags: ["--threads=4", "-O3"]`	No YAML changes

Code Architecture

Aspect	PR #79	This PR (#90)
sol2infs	`sol -> softplus.(N .* sol[2, :])` in ODEProcess	`sol -> sol[2, :]` - ODE returns proportions
Population scaling	Inside infection process	Inside observation model via `TransformObservationModel`
Accessors usage	Used for `stochastic_sir` only	Used for both `transform_varying` and `stochastic_sir`
ODE packages	Not explicitly shown where loaded	Loaded at start of Case Study 3 section

Content

Aspect	PR #79	This PR (#90)
Equations	ODE equations + observation model scattered	All equations together in Model section
Ascertainment	Introduced without noting it's new	Explicitly notes `Ascertainment` is a new observation model modifier
Writing style	More verbose explanations	Condensed to match other case studies

Tasks from PR #79 Addressed

~~Go back through to reduce amount of writing~~ - Condensed significantly
~~Go back through with emphasis on more reusability (mainly plotting functions)~~ - Helper functions created
~~Consider updating figure from 2x2 to 2x3 as per other examples~~ - Now 2x3
~~Put nadapts and ndraws back to standard values (1000 each)~~ - Uses standard values

What This PR Preserves from PR #79

Both deterministic and stochastic models
Same priors for SIR parameters (SIRParams)
Same AR(1) priors for ascertainment process
@set from Accessors.jl for model composition
Posterior R0 comparison between models

Design Decision: sol2infs in Observation Model

The key architectural change is moving population scaling from the infection process to the observation model:

# PR #79: Scaling in ODEProcess
sir_process = ODEProcess(
    params = sir_params,
    sol2infs = sol -> softplus.(N .* sol[2, :]),  # Combined here
    ...
)

# This PR: Scaling in TransformObservationModel
sir_process = ODEProcess(
    params = sir_params,
    sol2infs = sol -> sol[2, :],  # Returns proportions
    ...
)
transform_pois = TransformObservationModel(pois, x -> log1pexp.(N .* x))

This separation means the ODE model returns epidemiologically meaningful quantities (infected proportions), whilst the observation model handles the transformation to expected counts - demonstrating cleaner compositional design.

- Load ODE packages inline with data, not separate chunk - Consolidate equations: deterministic and stochastic together - Add figure references when defining components (like other case studies) - Explain proportions choice: numerical stability across population sizes - Rename ascert_ar to ascertainment_ar/ascertainment_model for clarity - Reference inference_method from sec-example1 - Define both EpiProblems together, then fit sequentially - Hide R0 computation code (echo: false) - Reword index.qmd intro to match other case study style - Add discussion note: ODE composition via AlgebraicPetri/Catalyst

case-studies.qmd: - Add complete mathematical model with R0 equation - Expand explanation of continuous/discrete time connection - Add softplus explanation (numerical stability for small negative values) - Add OU to AR(1) connection note - Add SciML compositional approach context - Add SciMLSensitivity requirement note - Update figure caption with detailed panel descriptions - Add inline computed R0 values with 95% CIs - Add calibration comparison between models index.qmd: - Add OU to AR(1) derivation note - Add softplus stability note - Add calibration comparison finding - Add submodel conditioning flexibility note in discussion

- Fix I_t_sol field access error by solving ODE directly from parameters - Add solve_sir helper to solve ODE with given β, γ, I₀ - Add posterior_sir_solutions to generate ODE solutions from chain samples - Add compartment_quantiles for computing quantile ribbons - Reuse post_predictive_yt! from Case Study 1 instead of duplicate - Simplify figure code by reusing existing helper functions

…ction

Restructure figure layout to 2x3 grid with stacked S/I compartment subplots for clearer presentation. Fix R0 outlier issues caused by numerical instability and improve statistical reporting. Changes: - Create stacked S/I subplots for prior (Panel A) and posterior (Panel D) compartments - Filter extreme R0 outliers (>100) from γ≈0 numerical instability - Change R0 statistics from mean to median for proper Bayesian reporting - Update credible interval notation from CI to CrI - Combine both models in posterior panels D, E, F with blue/orange color scheme - Add posterior Rt over time plot (Panel F) with Rt=1 reference line - Modify post_predictive_yt! to support multiple models on same plot - Remove all plot titles, relying on caption for descriptions - Update figure caption to reflect new 6-panel layout

Correct Rt formula to Rt = (β/γ) × S(t) = R0 × S(t) instead of just β × S(t). Create helper function calculate_Rt to eliminate code duplication between deterministic and stochastic model calculations.

Keep only the R0 legend in Panel C since the blue/orange color scheme for deterministic/stochastic is established there and consistent across all posterior panels.

Plot individual posterior trajectories for S and I compartments to match Panel A style, replacing median lines and credible interval bands with transparent trajectory lines.

Add details about deterministic vs stochastic SIR formulations, parameter inference, and model comparison. Also expand explanation of SciML ecosystem advantages.

- Update equations to show S(t), I(t), R(t) are proportions of population N - Fix infection rate to r_SI(t) = β*S(t)*I(t) (proportion-based) - Fix expected observations to softplus(N·I(t)) to show scaling step - Condense softplus explanation to single sentence mentioning N scaling

- Remove redundant latent_prefix parameter (uses default) - Clarify ascertainment is log-scale AR(1) process applied multiplicatively via exp(κ_t) - Add context about AR process prior specifications (damping, initial state, innovation std)

Explain that AR latent model type is reused from earlier case studies but for observation process rather than infection dynamics. Clarify Ascertainment takes latent_model as argument (unlike earlier direct generate_latent calls) and handles generation internally.

Replace log1pexp with softplus from LogExpFunctions to match mathematical notation in text. Both are equivalent but softplus is more explicit.

Show generated code for generate_latent_infs (ODEProcess) and generate_observations (observation model composition), matching pattern from case study 2.

seabbs · 2025-12-03T21:20:50Z

Also not I have tried both the solver used in the original and the default EpiAware solver and both have issues.

- Panel B: Use longer time scale (70 days) with new infection pattern - Panel C: Swap density plot order (stochastic first, deterministic second) - Panels D, E, F: Maintain blue first, orange second for layering

Use ylims!() function instead of invalid ylims attribute for Makie Axis

Change from Beta(0.5, 0.5) to Beta(1, 99) for initial_prop_infected. This concentrates prior mass on small initial outbreaks, ensuring S₀ stays near 1 as expected epidemiologically.

Change from LogNormal(0, 1.0) to LogNormal(0, 0.5) for infectiousness prior. This reduces prior mass on extreme β values while maintaining flexibility (95% CI: [0.37, 2.7] vs [0.14, 7.1]).

Plot deterministic density first, then stochastic, ensuring orange appears on top of blue.

Change from 0.1 to 0.05 to reduce visual density while maintaining same number of lines.

SamuelBrand1 · 2025-12-03T22:24:40Z

@SamuelBrand1 this is just a first pass aiming to get it to okay. So I think we will need to iterate a few times in issues. Note that the model both here and in your PR has some major convergence issues for the detministic model especially.

I've just run this branch locally and not found a convergence problem... thats a bit odd.

EDIT: running cell by cell gets a different result to running the task which calls quarto render

Add xticklabelrotation=π/4 (45 degrees) to date tick labels in case studies 2 and 3 to prevent month/year labels from overlapping.

Change from (1 + cospi(...)) to (1 + 0.3 * cospi(...)) so infections oscillate between 350-650 instead of 0-1000, making the ascertainment variation more visible.

case-studies.qmd

seabbs · 2025-12-04T09:35:05Z

I've just run this branch locally and not found a convergence problem... thats a bit odd.

I tuned the priors since posting to get this to work

SamuelBrand1 · 2025-12-04T10:25:19Z

I tuned the priors since posting to get this to work

From the rendered artefact downloaded from CI, this has done the job.

Remove duplicate SIR solving code by using internal _generate_ode_solution model function. Add sir_R0 helper function to avoid code duplication.

Extract calculate_Rt as reusable helper function that uses sir_R0 internally.

…e names - Reference @fig-epinow2 D when introducing delay distributions in case study 2 - Fix deterministic_R0/stochastic_R0 variable names (remove _raw suffix)

Modify EpiNow2 case study to use more informative prior (Normal(0.3, 0.3)) for initial log reproduction number, favouring outbreak scenarios. Add explanations of what initialization priors mean on the natural Rt scale for both Mishra and EpiNow2 examples. Include transitional text before broadcasting to improve flow.

Create ar2_mishra with ρ₁ ~ N(0.8, 0.05) as specified in Mishra et al., reflecting strong autocorrelation in log R_t process. Previous ar2 used ρ₁ ~ N(0.2, 0.2) which did not match paper. Update case-studies.qmd to use ar2_mishra in Mishra case study. Fix misleading claim in index.qmd that illustration priors were based on Mishra specification.

SamuelBrand1

This looks good. In particular I like the pattern of recursively using @SubModel inside the ODE model.

SamuelBrand1 and others added 4 commits November 27, 2025 13:00

Add SIR case study with time-varying ascertainment

4d41eb0

Merge main into case-study-3-cleanup

0098b07

seabbs added 25 commits December 2, 2025 20:08

Fix awkward OU/AR(1) phrasing

32c587b

Fix incorrect claim about swapping in previous case studies

df1a96f

Clarify swapping is at EpiProblem level not lower modules

e6c2554

Expand figure discussion in main text

e313b8c

Clarify EpiProblem swapping mirrors lower-level component swapping

9605204

Update figure caption in main text to match case study

e763ead

Fix submodel conditioning point: about reducing DSL layer work

056cc19

Soften submodel conditioning statement

5fdb952

Define SIR vector field inline instead of using internal EpiAware fun…

9094d89

…ction

Remove redundant legend from stacked I compartment subplot

c06e118

Fix Rt calculation and reduce code duplication

1f32d62

Correct Rt formula to Rt = (β/γ) × S(t) = R0 × S(t) instead of just β × S(t). Create helper function calculate_Rt to eliminate code duplication between deterministic and stochastic model calculations.

Add missing citations for DifferentialEquations.jl and outbreaks package

c7b58d9

Change R0 density plots to transparent fills with solid outlines

9a91692

Remove redundant legends from panels D, E, F

8e953d2

Keep only the R0 legend in Panel C since the blue/orange color scheme for deterministic/stochastic is established there and consistent across all posterior panels.

Change Panel D to plot posterior trajectories instead of ribbons

75b7f4e

Plot individual posterior trajectories for S and I compartments to match Panel A style, replacing median lines and credible interval bands with transparent trajectory lines.

Expand case study 3 opening with more context about Chatzilena et al.

c929ac3

Add details about deterministic vs stochastic SIR formulations, parameter inference, and model comparison. Also expand explanation of SciML ecosystem advantages.

Use softplus function instead of log1pexp for clarity

a8e36bd

Replace log1pexp with softplus from LogExpFunctions to match mathematical notation in text. Both are equivalent but softplus is more explicit.

Add code printing for key compositional functions in case study 3

1202f1c

Show generated code for generate_latent_infs (ODEProcess) and generate_observations (observation model composition), matching pattern from case study 2.

Remove old case-study-3.qmd file

ce12f98

seabbs removed the request for review from SamuelBrand1 December 3, 2025 21:20

seabbs added 6 commits December 3, 2025 21:26

Ensure consistent plotting order and extend Panel B time scale

8cc3a77

- Panel B: Use longer time scale (70 days) with new infection pattern - Panel C: Swap density plot order (stochastic first, deterministic second) - Panels D, E, F: Maintain blue first, orange second for layering

Fix ylims syntax for Panel F axis

0c07ece

Use ylims!() function instead of invalid ylims attribute for Makie Axis

Use more informative prior for initial proportion infected

1fdbb3c

Change from Beta(0.5, 0.5) to Beta(1, 99) for initial_prop_infected. This concentrates prior mass on small initial outbreaks, ensuring S₀ stays near 1 as expected epidemiologically.

Tighten transmission rate prior to reduce β outliers

66e3daa

Change from LogNormal(0, 1.0) to LogNormal(0, 0.5) for infectiousness prior. This reduces prior mass on extreme β values while maintaining flexibility (95% CI: [0.37, 2.7] vs [0.14, 7.1]).

Fix Panel C plotting order so stochastic appears on top

88d61a4

Plot deterministic density first, then stochastic, ensuring orange appears on top of blue.

Reduce alpha transparency of Panel D trajectory lines

167eb01

Change from 0.1 to 0.05 to reduce visual density while maintaining same number of lines.

seabbs requested a review from SamuelBrand1 December 3, 2025 22:23

seabbs added 2 commits December 3, 2025 22:25

Rotate x-axis tick labels to prevent overlap

138b83f

Add xticklabelrotation=π/4 (45 degrees) to date tick labels in case studies 2 and 3 to prevent month/year labels from overlapping.

Reduce Panel B latent infection oscillation amplitude

1737560

Change from (1 + cospi(...)) to (1 + 0.3 * cospi(...)) so infections oscillate between 350-650 instead of 0-1000, making the ascertainment variation more visible.

SamuelBrand1 reviewed Dec 3, 2025

View reviewed changes

case-studies.qmd Outdated Show resolved Hide resolved

seabbs added 6 commits December 4, 2025 11:00

Refactor SIR figure code to re-solve ODE from posterior samples

6a5aae9

Remove duplicate SIR solving code by using internal _generate_ode_solution model function. Add sir_R0 helper function to avoid code duplication.

Soften submodel conditioning statement

00191dc

Extract calculate_Rt as reusable helper function that uses sir_R0 internally.

Reduce pathfinder runs from 5 to 3 for faster inference

1a84abc

Add figure panel reference for delay distributions and fix R0 variabl…

92cbcbd

…e names - Reference @fig-epinow2 D when introducing delay distributions in case study 2 - Fix deterministic_R0/stochastic_R0 variable names (remove _raw suffix)

add new internals for SIR

b724d5d

seabbs mentioned this pull request Dec 4, 2025

Add ODE case study #79

Closed

5 tasks

seabbs added 2 commits December 4, 2025 13:31

update toml

c9bdaf8

This was referenced Dec 4, 2025

Case studies: Figure 4 #14

Closed

Case studies: Example 3: change to use EpiAware functionality directly #24

Closed

SamuelBrand1 approved these changes Dec 4, 2025

View reviewed changes

SamuelBrand1 merged commit a388a9a into main Dec 4, 2025
2 checks passed

SamuelBrand1 deleted the case-study-3-cleanup branch December 4, 2025 16:58

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add Case Study 3: ODE-based SIR model with deterministic and stochastic variants #90

Add Case Study 3: ODE-based SIR model with deterministic and stochastic variants #90

seabbs commented Dec 2, 2025 •

edited

Loading

Uh oh!

seabbs commented Dec 2, 2025

Uh oh!

seabbs commented Dec 3, 2025

Uh oh!

SamuelBrand1 commented Dec 3, 2025 •

edited

Loading

Uh oh!

Uh oh!

seabbs commented Dec 4, 2025

Uh oh!

SamuelBrand1 commented Dec 4, 2025

Uh oh!

SamuelBrand1 left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Add Case Study 3: ODE-based SIR model with deterministic and stochastic variants #90

Add Case Study 3: ODE-based SIR model with deterministic and stochastic variants #90

Conversation

seabbs commented Dec 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Key implementation details

Test plan

Uh oh!

seabbs commented Dec 2, 2025

Comparison with Draft PR #79

Key Differences from PR #79

Structural Changes

Code Architecture

Content

Tasks from PR #79 Addressed

What This PR Preserves from PR #79

Design Decision: sol2infs in Observation Model

Uh oh!

seabbs commented Dec 3, 2025

Uh oh!

SamuelBrand1 commented Dec 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

seabbs commented Dec 4, 2025

Uh oh!

SamuelBrand1 commented Dec 4, 2025

Uh oh!

SamuelBrand1 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

seabbs commented Dec 2, 2025 •

edited

Loading

SamuelBrand1 commented Dec 3, 2025 •

edited

Loading