diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml new file mode 100644 index 0000000..643f97e --- /dev/null +++ b/.github/workflows/test.yaml @@ -0,0 +1,33 @@ +name: Julia tests + +on: + pull_request: + push: + branches: [main] + +permissions: + actions: write + contents: read + +jobs: + test: + name: Julia ${{ matrix.julia-version }} + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + julia-version: ['1.11', '1.12', '1.13'] + + steps: + - uses: actions/checkout@v7 + - uses: julia-actions/setup-julia@v3 + with: + version: ${{ matrix.julia-version }} + project: AlgebraicEpiMech + - uses: julia-actions/cache@v3 + - uses: julia-actions/julia-buildpkg@v1 + with: + project: AlgebraicEpiMech + - uses: julia-actions/julia-runtest@v1 + with: + project: AlgebraicEpiMech diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7d4f327 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +.DS_Store + +Manifest* diff --git a/AlgebraicEpiMech/Project.toml b/AlgebraicEpiMech/Project.toml new file mode 100644 index 0000000..0feadfd --- /dev/null +++ b/AlgebraicEpiMech/Project.toml @@ -0,0 +1,18 @@ +name = "AlgebraicEpiMech" +uuid = "a6e39d6b-8b2d-49f3-aca0-f4154643f040" +version = "0.1.0" +authors = ["Samuel Brand "] + +[workspace] +projects = ["test"] + +[deps] +AlgebraicPetri = "4f99eebe-17bf-4e98-b6a1-2c4f205a959b" +Catlab = "134e5e36-593f-5add-ad60-77f754baafbe" +DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" + +[compat] +AlgebraicPetri = "0.10.0" +Catlab = "0.17.2" +DocStringExtensions = "0.9.3" +julia = "1.11, 1.12, 1.13" diff --git a/AlgebraicEpiMech/src/AlgebraicEpiMech.jl b/AlgebraicEpiMech/src/AlgebraicEpiMech.jl new file mode 100644 index 0000000..2012833 --- /dev/null +++ b/AlgebraicEpiMech/src/AlgebraicEpiMech.jl @@ -0,0 +1,60 @@ +module AlgebraicEpiMech + +using AlgebraicPetri, AlgebraicPetri.TypedPetri +using Catlab +using DocStringExtensions +import Base: * + +# Epidemiological typing strategies and their Petri-net type systems +export EpidemiologicalTyping, OnePopulationTyping, UninfectedInfectedTyping, type_system + +# Epi mechanism functions +export add_infection!, add_disease_progression!, + add_uninfected_density_progression!, add_reversion_progression! + +# Observation chain +# Observation by pushout (attached after composition) +export ObservationTarget, AtEvent, AtCompartment, attach_observation + +export ObservationChainLayout, ObservationLayout, observation_layout + +# Compartmental models types +export CompartmentalModel, SI, SEI, SIS, SIR, SEIR, SEIRS, SEIS + +# Multistrain model types +export MultiStrainModel, NoCrossImmunity, CompleteCrossImmunity + +# Immune-history model types +export ImmuneHistory, ImmuneHistoryMode, FullHistory, LatestInfection + +# Stratification types +export Stratification, ContactStratification, AgeStratification, GeographicStratification, + compose_stratifications + +# Model functions +export EpiMechModel, create_model_uwd, create_model, generate_transition_names + +# Vectorfield functions +export flatten_symbols, vectorfield_flat + +# Documentation strings +include("docstrings.jl") + +# Epidemiological dynamics +include("typing.jl") +include("mechanisms.jl") +include("CompartmentalModel.jl") +include("MultiStrainModel.jl") +include("ImmuneHistory.jl") +include("Stratification.jl") +include("construction_helpers.jl") +include("create_model_uwd.jl") +include("generate_transition_names.jl") +include("create_model.jl") +include("vectorfield_flat.jl") + +# Observation attachment. +include("ObservationTarget.jl") +include("ObservationLayout.jl") + +end diff --git a/AlgebraicEpiMech/src/CompartmentalModel.jl b/AlgebraicEpiMech/src/CompartmentalModel.jl new file mode 100644 index 0000000..a6596fe --- /dev/null +++ b/AlgebraicEpiMech/src/CompartmentalModel.jl @@ -0,0 +1,212 @@ +""" +Abstract base type for all epidemiological-mechanical models in the AlgebraicEpiMech framework. + +All concrete epidemiological-mechanical model types should be subtypes of `EpiMechModel`. +This abstract type serves as the root of the type hierarchy for models that combine +epidemiological dynamics with mechanical or algebraic structures. + +# Extended help + +Subtypes of `EpiMechModel` should implement the necessary interface methods for their +specific model formulation. + +# See also +- Related concrete model types (define as needed) +""" +abstract type EpiMechModel end + +""" +Abstract base type for compartmental epidemiological models. + +Compartmental models define the specific disease dynamics and transitions +between epidemiological states (e.g., S→I→R) as undirected wiring diagrams. +""" +abstract type CompartmentalModel <: EpiMechModel end + +function _validate_stage_count(name::Symbol, count::Int) + count >= 1 || throw(ArgumentError("$name must be at least 1")) + return count +end + +# Base compartmental models - concrete structs for clean dispatch + +""" +Susceptible-Infected (SI) compartmental model. + +The SI model includes: +- Susceptible (S): Individuals who can become infected +- Infected (I): Individuals who are infectious + +Transitions: +- S + I → I + I (infection/transmission) + +The SI model represents endemic diseases with no recovery. +""" +struct SI <: CompartmentalModel + number_I_stages::Int + number_of_states::Int + + function SI(; number_I_stages::Int = 1) + _validate_stage_count(:number_I_stages, number_I_stages) + return new(number_I_stages, 1 + number_I_stages) + end +end + +""" +Susceptible-Exposed-Infected (SEI) compartmental model. + +The SEI model extends SI with an exposed compartment: +- Susceptible (S): Individuals who can become infected +- Exposed (E): Individuals who are infected but not yet infectious +- Infected (I): Individuals who are infectious + +Transitions: +- S + I → E + I (exposure/transmission) +- E → I (progression to infectious) + +The SEI model represents diseases with an incubation period but no recovery. +""" +struct SEI <: CompartmentalModel + number_E_stages::Int + number_I_stages::Int + number_of_states::Int + + function SEI(; number_E_stages::Int = 1, number_I_stages::Int = 1) + _validate_stage_count(:number_E_stages, number_E_stages) + _validate_stage_count(:number_I_stages, number_I_stages) + return new(number_E_stages, number_I_stages, 1 + number_E_stages + number_I_stages) + end +end + +# Extended compartmental models - built by extending base models + +""" +Susceptible-Infected-Recovered (SIR) compartmental model. + +The SIR model extends SI with recovery: +- Susceptible (S): Individuals who can become infected +- Infected (I): Individuals who are infectious +- Recovered (R): Individuals who have recovered and gained immunity + +Transitions: +- S + I → I + I (infection/transmission) [from SI] +- I → R (recovery) [extension] + +Built by extending SI model with recovery transition. +""" +struct SIR <: CompartmentalModel + number_I_stages::Int + number_of_states::Int + + function SIR(; number_I_stages::Int = 1) + _validate_stage_count(:number_I_stages, number_I_stages) + return new(number_I_stages, 2 + number_I_stages) + end +end + +""" +Susceptible-Exposed-Infected-Recovered (SEIR) compartmental model. + +The SEIR model extends SEI with recovery: +- Susceptible (S): Individuals who can become infected +- Exposed (E): Individuals who are infected but not yet infectious +- Infected (I): Individuals who are infectious +- Recovered (R): Individuals who have recovered and gained immunity + +Transitions: +- S + I → E + I (exposure/transmission) [from SEI] +- E → I (progression to infectious) [from SEI] +- I → R (recovery) [extension] + +Built by extending SEI model with recovery transition. +""" +struct SEIR <: CompartmentalModel + number_E_stages::Int + number_I_stages::Int + number_of_states::Int + + function SEIR(; number_E_stages::Int = 1, number_I_stages::Int = 1) + _validate_stage_count(:number_E_stages, number_E_stages) + _validate_stage_count(:number_I_stages, number_I_stages) + return new(number_E_stages, number_I_stages, 2 + number_E_stages + number_I_stages) + end +end + +""" +Susceptible-Infected-Susceptible (SIS) compartmental model. + +The SIS model extends SI with waning immunity (reversion): +- Susceptible (S): Individuals who can become infected +- Infected (I): Individuals who are infectious + +Transitions: +- S + I → I + I (infection/transmission) [from SI] +- I → S (waning immunity/reversion) [extension] + +The SIS model represents diseases with no lasting immunity. +""" +struct SIS <: CompartmentalModel + number_I_stages::Int + number_of_states::Int + + function SIS(; number_I_stages::Int = 1) + _validate_stage_count(:number_I_stages, number_I_stages) + return new(number_I_stages, 1 + number_I_stages) + end +end + +""" +Susceptible-Exposed-Infected-Susceptible (SEIS) compartmental model. + +The SEIS model extends SEI with waning immunity (reversion): +- Susceptible (S): Individuals who can become infected +- Exposed (E): Individuals who are infected but not yet infectious +- Infected (I): Individuals who are infectious + +Transitions: +- S + I → E + I (exposure/transmission) [from SEI] +- E → I (progression to infectious) [from SEI] +- I → S (waning immunity/reversion) [extension] + +The SEIS model represents diseases with incubation period but no lasting immunity. +""" +struct SEIS <: CompartmentalModel + number_E_stages::Int + number_I_stages::Int + number_of_states::Int + + function SEIS(; number_E_stages::Int = 1, number_I_stages::Int = 1) + _validate_stage_count(:number_E_stages, number_E_stages) + _validate_stage_count(:number_I_stages, number_I_stages) + return new(number_E_stages, number_I_stages, 1 + number_E_stages + number_I_stages) + end +end + +""" +Susceptible-Exposed-Infected-Recovered-Susceptible (SEIRS) compartmental model. + +The SEIRS model extends SEIR with waning immunity: +- Susceptible (S): Individuals who can become infected +- Exposed (E): Individuals who are infected but not yet infectious +- Infected (I): Individuals who are infectious +- Recovered (R): Individuals who have recovered but may lose immunity + +Transitions: +- S + I → E + I (exposure/transmission) [from SEIR] +- E → I (progression to infectious) [from SEIR] +- I → R (recovery) [from SEIR] +- R → S (waning immunity) [extension] + +Built by extending SEIR model with waning transition. +""" +struct SEIRS <: CompartmentalModel + number_E_stages::Int + number_I_stages::Int + number_of_states::Int + + function SEIRS(; number_E_stages::Int = 1, number_I_stages::Int = 1) + _validate_stage_count(:number_E_stages, number_E_stages) + _validate_stage_count(:number_I_stages, number_I_stages) + return new(number_E_stages, number_I_stages, 2 + number_E_stages + number_I_stages) + end +end diff --git a/AlgebraicEpiMech/src/ImmuneHistory.jl b/AlgebraicEpiMech/src/ImmuneHistory.jl new file mode 100644 index 0000000..5a0bfe9 --- /dev/null +++ b/AlgebraicEpiMech/src/ImmuneHistory.jl @@ -0,0 +1,134 @@ +# Type system for Immune-history stratification +# +# The behaviour encoded is that for immune-history classes `U_h`; the infected side +# carries `(h, i)` — the prior history `h` and the strain/subtype `i` of infection. This +# is enough to determine how infection recovery increments the immune history. + +""" +Abstract base type for selections of how much immune history is retained. +""" +abstract type ImmuneHistoryMode end + +""" +Full (unordered) immune history: uninfected individuals are indexed by the **set** of +strains they are immune to (`2^n` classes). Immunity accumulates — reversion routes to +`h ∪ {i}`. +""" +struct FullHistory <: ImmuneHistoryMode end + +""" +Latest-infection status: uninfected individuals are indexed by their **most recent** +infection (`n + 1` classes). Reversion overwrites immune status — routes to `{i}`. +""" +struct LatestInfection <: ImmuneHistoryMode end + +""" + ImmuneHistory{M<:ImmuneHistoryMode} <: MultiStrainModel + +Immune-history combination of the available strain/subtype names and `ImmuneHistoryMode`. + +# Fields +- `strain_names::Vector{Symbol}`: strain names (e.g. `[:current, :invader]`) +- `mode::M`: `FullHistory()` or `LatestInfection()` + +# Examples +```julia +typing = UninfectedInfectedTyping() + +seirs = create_model(typing, SEIRS()) +history = create_model(typing, ImmuneHistory([:current, :invader])) + +model = typed_product(seirs, history) # immune-history-resolved SEIRS +# States: S_naive, S_current, …, E_invader_from_current, I_…, R_… +``` +""" +struct ImmuneHistory{M <: ImmuneHistoryMode} <: MultiStrainModel + strain_names::Vector{Symbol} + mode::M + + function ImmuneHistory( + strain_names::Vector{Symbol}; mode::M = FullHistory() + ) where {M <: ImmuneHistoryMode} + isempty(strain_names) && + throw(ArgumentError("ImmuneHistory requires at least one strain")) + length(strain_names) == length(unique(strain_names)) || + throw(ArgumentError("strain names must be unique")) + return new{M}(strain_names, mode) + end +end + +""" + ImmuneHistory(n::Int; kwargs...) + +Auto-name `n` strains as `strain_1 … strain_n`. +""" +function ImmuneHistory(n::Int; kwargs...) + n >= 1 || throw(ArgumentError("number of strains must be >= 1")) + return ImmuneHistory([Symbol("strain_$(i)") for i in 1:n]; kwargs...) +end + +# ============================================================================ +# Immune-class + wiring helpers (internal; not exported) +# +# An immune class is a sorted `Tuple` of strains immune to; `()` is naive. +# ============================================================================ + +# Immune class: a sorted tuple of strain symbols (`()` is naive). +const _ImmuneClass = Tuple{Vararg{Symbol}} + +function _powerset(v::AbstractVector{T}) where {T} + subsets = Vector{Vector{T}}([T[]]) + for x in v + # `append!` grows in place; the comprehension reads the current `subsets` + # fully before appending, so we avoid recopying the accumulated powerset. + append!(subsets, [vcat(s, x) for s in subsets]) + end + return subsets +end + +_immune_classes(::FullHistory, strains::Vector{Symbol}) = + [Tuple(sort(s; by = string)) for s in _powerset(strains)] + +function _immune_classes(::LatestInfection, strains::Vector{Symbol}) + classes = Vector{Tuple{Vararg{Symbol}}}() + push!(classes, ()) + for s in strains + push!(classes, (s,)) + end + return classes +end + +# Strains a class immune to `class` is susceptible to. +_susceptible_strains(class, strains::Vector{Symbol}) = + [i for i in strains if !(i in class)] + +# Immune class reached after recovering from strain `i` (the :reversion target). +_recover_to(::FullHistory, class, i::Symbol) = Tuple(sort(unique((class..., i)); by = string)) +_recover_to(::LatestInfection, ::Any, i::Symbol) = (i,) + +_htag(class) = isempty(class) ? "naive" : join(string.(class), "_") + +_uninfected_name(class) = Symbol("U_", _htag(class)) + +# Infected junction for "history h, currently infected with strain i": `_from_`. +_infected_name(class, i::Symbol) = Symbol(string(i), "_from_", _htag(class)) + +# The infected `(h, i)` junctions to create: every (class, susceptible strain) pair. +function _infected_pairs(mode::ImmuneHistoryMode, strains::Vector{Symbol}) + pairs = Tuple{_ImmuneClass, Symbol}[] + for h in _immune_classes(mode, strains), i in _susceptible_strains(h, strains) + push!(pairs, (h, i)) + end + return pairs +end + +# Identify the strain fought by an infected junction `_from_` (match longest +# strain name first so prefix names are safe). +function _strain_of(var::Symbol, strain_names) + s = string(var) + for cand in sort(strain_names; by = x -> -length(string(x))) + cs = string(cand) + startswith(s, cs * "_from_") && return cand + end + return error("could not identify strain in $var among $strain_names") +end diff --git a/AlgebraicEpiMech/src/MultiStrainModel.jl b/AlgebraicEpiMech/src/MultiStrainModel.jl new file mode 100644 index 0000000..a524efc --- /dev/null +++ b/AlgebraicEpiMech/src/MultiStrainModel.jl @@ -0,0 +1,112 @@ +# Type system for multistrain epidemiological models + +""" +Abstract base type for multistrain epidemiological models. + +Multistrain models define how multiple pathogen strains interact through +different assumptions about cross-immunity. These compose with compartmental +models via typed_product to create strain-structured disease dynamics. +""" +abstract type MultiStrainModel <: EpiMechModel end + +""" +No cross-immunity multistrain model. + +In the no cross-immunity model: +- Each strain operates independently with no interaction +- Strains act as independent strata (similar to age groups) +- Typing: OnePopulationTyping (all compartments have the same type) +- Composition: Creates parallel disease dynamics per strain via typed_product + +# Fields +- `number_of_strains::Int`: Number of independent strains in the model +- `strain_names::Vector{Symbol}`: Names for each strain (e.g., [:h1n1, :h3n2, :b]) + +# Examples +```julia +# Create a 3-strain model with custom names +multistrain = NoCrossImmunity([:h1n1, :h3n2, :b]) + +# Create a 2-strain model with auto-generated names +multistrain = NoCrossImmunity(2) # Creates [:strain_1, :strain_2] + +# Compose to create strain-structured SIR +typing = OnePopulationTyping() +strain_typed = create_model(typing, multistrain) +sir_typed = create_model(typing, SIR()) +combined = typed_product(sir_typed, strain_typed) +``` +""" +struct NoCrossImmunity <: MultiStrainModel + number_of_strains::Int + strain_names::Vector{Symbol} + + function NoCrossImmunity(strain_names::Vector{Symbol}) + if length(strain_names) < 1 + throw(ArgumentError("strain_names must contain at least 1 strain")) + end + if length(strain_names) != length(unique(strain_names)) + throw(ArgumentError("strain_names must be unique")) + end + return new(length(strain_names), strain_names) + end + + function NoCrossImmunity(number_of_strains::Int) + if number_of_strains < 1 + throw(ArgumentError("number_of_strains must be at least 1")) + end + strain_names = [Symbol("strain_$(i)") for i in 1:number_of_strains] + return new(number_of_strains, strain_names) + end +end + +""" +Complete cross-immunity multistrain model. + +In the complete cross-immunity model: +- All strains share a common susceptible pool +- Infection by any strain confers immunity to all strains +- Typing: UninfectedInfectedTyping (typed S vs I/R compartments) +- Composition: Strains compete for susceptibles via shared S depletion + +# Fields +- `number_of_strains::Int`: Number of competing strains in the model +- `strain_names::Vector{Symbol}`: Names for each strain (e.g., [:h1n1, :h3n2]) + +# Examples +```julia +# Create a 2-strain model with custom names +multistrain = CompleteCrossImmunity([:wild_type, :variant]) + +# Create a 3-strain model with auto-generated names +multistrain = CompleteCrossImmunity(3) # Creates [:strain_1, :strain_2, :strain_3] + +# Compose to create competing strain SIR +typing = UninfectedInfectedTyping() +strain_typed = create_model(typing, multistrain) +sir_typed = create_model(typing, SIR()) +combined = typed_product(sir_typed, strain_typed) +``` +""" +struct CompleteCrossImmunity <: MultiStrainModel + number_of_strains::Int + strain_names::Vector{Symbol} + + function CompleteCrossImmunity(strain_names::Vector{Symbol}) + if length(strain_names) < 1 + throw(ArgumentError("strain_names must contain at least 1 strain")) + end + if length(strain_names) != length(unique(strain_names)) + throw(ArgumentError("strain_names must be unique")) + end + return new(length(strain_names), strain_names) + end + + function CompleteCrossImmunity(number_of_strains::Int) + if number_of_strains < 1 + throw(ArgumentError("number_of_strains must be at least 1")) + end + strain_names = [Symbol("strain_$(i)") for i in 1:number_of_strains] + return new(number_of_strains, strain_names) + end +end diff --git a/AlgebraicEpiMech/src/ObservationLayout.jl b/AlgebraicEpiMech/src/ObservationLayout.jl new file mode 100644 index 0000000..e3c3749 --- /dev/null +++ b/AlgebraicEpiMech/src/ObservationLayout.jl @@ -0,0 +1,156 @@ +# Observation-chain metadata for Petri nets. +# +# Observation models are attached to the petri net representing the epidemiological dynamics by +# pushout, that is assembling the joint dynamics and observation metadata on shared components. +# Here we define the shared components. + +""" + ObservationChainLayout(source_name, obs_names) + +Metadata for one observation chain. + +# Fields +- `source_name`: Flattened name of the event or compartment being observed +- `obs_names`: Observation-state names in stage order +- `cumulative_name`: Terminal observation-state name, derived from `obs_names` +""" +struct ObservationChainLayout{N} + source_name::Symbol + obs_names::NTuple{N, Symbol} + cumulative_name::Symbol + + function ObservationChainLayout( + source_name::Symbol, obs_names::NTuple{N, Symbol} + ) where {N} + N > 0 || throw(ArgumentError("An observation chain must contain at least one stage")) + return new{N}(source_name, obs_names, last(obs_names)) + end +end + +""" + ObservationLayout(obs_names, chains) + +Metadata for all observation chains in an augmented Petri net. + +# Fields +- `obs_names`: All flattened observation-state names in Petri-net order +- `chains`: Observation chains in first-encounter order +- `cumulative_names`: Terminal state of each chain, derived from `chains` +""" +struct ObservationLayout{N, C <: Tuple, M} + obs_names::NTuple{N, Symbol} + chains::C + cumulative_names::NTuple{M, Symbol} + + function ObservationLayout( + obs_names::NTuple{N, Symbol}, chains::C + ) where {N, C <: Tuple} + all(chain -> chain isa ObservationChainLayout, chains) || throw( + ArgumentError("All chains must be ObservationChainLayout values") + ) + cumulative_names = Tuple(chain.cumulative_name for chain in chains) + return new{N, C, length(cumulative_names)}( + obs_names, chains, cumulative_names + ) + end +end + +""" + observation_layout(pn::LabelledPetriNet; prefix::Symbol = :O) + +Return deterministic observation-chain metadata for an augmented Petri net. + +The result is an [`ObservationLayout`](@ref) with fields: +- `obs_names`: all flattened observation state names in Petri-net order +- `chains`: ordered [`ObservationChainLayout`](@ref) values +- `cumulative_names`: flattened terminal observation state for each chain + +This centralizes the observation naming/ordering contract, so downstream packages do not each +re-implement observation-state parsing. `prefix` must match the prefix passed to +[`attach_observation`](@ref). The function recognises the `__` leaf that +`attach_observation` produces; a chain named otherwise is not recognised as one. +""" +function observation_layout(pn; prefix::Symbol = :O) + obs_names = Symbol[] + chain_order = Symbol[] + chain_stages = Dict{Symbol, Vector{Tuple{Int, Symbol}}}() + + for raw_name in snames(pn) + meta = _observation_species_metadata(raw_name, prefix) + isnothing(meta) && continue + + push!(obs_names, meta.obs_name) + + if !haskey(chain_stages, meta.source_name) + chain_stages[meta.source_name] = Tuple{Int, Symbol}[] + push!(chain_order, meta.source_name) + end + + push!(chain_stages[meta.source_name], (meta.stage, meta.obs_name)) + end + + chains = map(chain_order) do source_name + stage_entries = sort(chain_stages[source_name], by = first) + stage_numbers = first.(stage_entries) + expected_stages = collect(1:length(stage_numbers)) + stage_numbers == expected_stages || throw( + ArgumentError( + "Observation chain for :$source_name has non-consecutive stages $stage_numbers" + ) + ) + + stage_names = Tuple(last.(stage_entries)) + ObservationChainLayout(source_name, stage_names) + end + + return ObservationLayout(Tuple(obs_names), Tuple(chains)) +end + +function _observation_species_metadata(name, prefix) + leaves = _flatten_symbol_leaves(name) + parsed = map(leaf -> _parse_observation_stage(leaf, prefix), leaves) + matches = findall(!isnothing, parsed) + + isempty(matches) && return nothing + length(matches) == 1 || throw( + ArgumentError("Expected at most one observation-stage leaf in $name") + ) + + match_idx = only(matches) + stage_info = something(parsed[match_idx]) + + source_leaves = copy(leaves) + source_leaves[match_idx] = stage_info.source + + return ( + obs_name = flatten_symbols(name), + source_name = Symbol(join(string.(source_leaves), "_")), + stage = stage_info.stage, + ) +end + +_flatten_symbol_leaves(name::Symbol) = [name] + +function _flatten_symbol_leaves(name::Tuple) + leaves = Symbol[] + for item in name + append!(leaves, _flatten_symbol_leaves(item)) + end + return leaves +end + +function _parse_observation_stage(name::Symbol, prefix::Symbol) + prefix_marker = string(prefix, "_") + raw_name = String(name) + startswith(raw_name, prefix_marker) || return nothing + + parts = split(chopprefix(raw_name, prefix_marker), "_") + length(parts) >= 2 || return nothing + stage = tryparse(Int, parts[end]) + isnothing(stage) && return nothing + + source = Symbol(join(parts[1:(end - 1)], "_")) + isempty(String(source)) && return nothing + + return (source = source, stage = stage) +end diff --git a/AlgebraicEpiMech/src/ObservationTarget.jl b/AlgebraicEpiMech/src/ObservationTarget.jl new file mode 100644 index 0000000..8a26137 --- /dev/null +++ b/AlgebraicEpiMech/src/ObservationTarget.jl @@ -0,0 +1,260 @@ +# Attaching observation to a composed model, by pushout. +# +# The typed petri net models compose via their limit / categorical pullback operation, that is +# identically typed transitions are generated from the set of all tensor products of the component transitions. +# +# Observation models also introduce structure, via a different operation: a colimit / categorical pushout operation. +# These join the component models along shared structure, in this case either a species or transitions that on one hand are part of +# the petri net representing the epidemiological dynamics and on the other hand represent generating new observations +# in the future. +# +# The main types of observation are "incidence"-like and "prevalence"-like. +# Both kinds of observation are the same colimit. Only the rule "L -> R" differs: +# +# incidence L = S + I -> E + I R = S + I -> E + I + O +# the SAME event, additionally recorded. One transition, one rate, so "one +# observation per infection" is structural rather than a convention. +# +# prevalence L = I R = I -> I + O +# a new catalytic mechanism reading occupancy, with its own detection rate. +# Rapid antigen testing is the motivating case. +# +# +# Because this runs after the pullback compositions, stratification factors need to know nothing whatsoever +# about observation: no per-stratum accumulator, no reflexive observation box etc. +# +# Optimisation +# An earlier version of this file materialized the observations literally using Catlab functions one pushout per matched transition, +# re-matching against the grown net after every step. That was O(L^4) in the number of strata — +# ~L^2 matched routes, each paying a `homomorphism` search plus a colimit on a net that is itself +# ~L^2 transitions. For example, on a 50-state geographic model it meant hours of model construction +# before any model running (1.2 s at L = 6, 109 s at L = 20, extrapolated 1-2 h at L = 50; 0.01 s +# now). Both rules are monotone additions — nothing is deleted, nothing is merged beyond gluing +# onto a chain that is findable by label — so the same colimit is materialized directly below in +# one pass, appending chains and arcs in first-encounter order. + +""" + ObservationTarget + +What an observation attaches to. [`AtEvent`](@ref) records a transition (incidence); +[`AtCompartment`](@ref) samples a species (prevalence). +""" +abstract type ObservationTarget end + +""" + AtEvent(transition) + +Record a TRANSITION: incidence. Every transition whose (flattened) name begins with `transition` +gains an output arc into an accumulator, so it fires at that transition's own rate and counts +exactly one observation per occurrence. + +`AtEvent(:transmission)` counts infections as they happen, which is what decouples a reporting +delay from the latent period — the delay is then whatever the chain adds and nothing else. +""" +struct AtEvent <: ObservationTarget + transition::Symbol +end + +""" + AtCompartment(species) + +Sample a SPECIES: prevalence. Every species whose (flattened) name begins with `species` gains a +catalytic tap `X -> X + O`, carrying its own detection rate — whoever is in the compartment is +currently detectable (although detection does not cause removal) +""" +struct AtCompartment <: ObservationTarget + species::Symbol +end + +# Labels in a composed net are nested tuples: `(:S, :a)` for one stratification, `((:S, :a), :b)` +# for two. An accumulator belongs to the same stratum as the thing it observes, so it takes that +# label with the COMPARTMENT part swapped out and the stratification parts kept. +_relabel_head(label::Symbol, head::Symbol) = head +_relabel_head(label::Tuple, head::Symbol) = (_relabel_head(label[1], head), Base.tail(label)...) + +_flat(x) = flatten_symbols(x) +_startswith(label, prefix::Symbol) = startswith(string(_flat(label)), string(prefix)) + +# The species a transition NETS — outputs that are not also inputs. For `S + I -> E + I` that is +# `E`, so grouping infection routes by it gives one accumulator per infectee stratum, which is +# what "incidence in location x" means. Structural, so it needs no name parsing. +function _net_product(pn, t) + ins = Set(AlgebraicPetri.inputs(pn, t)) + outs = [s for s in AlgebraicPetri.outputs(pn, t) if !(s in ins)] + return isempty(outs) ? first(AlgebraicPetri.outputs(pn, t)) : first(outs) +end + +""" + attach_observation(pn, target; n_stages = 1, prefix = :O) -> LabelledPetriNet + +Attach an observation chain to `pn`, returning a new net. The result is the colimit described at +the top of this file, materialized in one direct pass (see the note there for why it is not +built by pushout). + +`n_stages` Erlang delay stages are added per chain: the first receives the observation, each +subsequent one takes flow from the last, and the final stage is cumulative (no outflow), i.e. the +accumulator the observation model reads. + +Reapplying the same target, prefix, and stage count is idempotent. Reusing a target and prefix +for an incompatible chain shape throws an `ArgumentError`; choose a different prefix when both +observation layouts are required. + +Chains are grouped by stratum. For [`AtEvent`](@ref) the group is the transition's net product — +for infection routes `S_x + I_y -> E_x + I_y` that is `E_x`, so all routes infecting `x` count +into `x`'s chain. For [`AtCompartment`](@ref) each matched species gets its own chain. + +Apply this after all `typed_product` composition so chains are grouped over the final strata. +""" +function attach_observation(pn, target::ObservationTarget; n_stages::Int = 1, prefix::Symbol = :O) + n_stages >= 1 || throw(ArgumentError("n_stages must be >= 1")) + net = deepcopy(pn) + _attach_all!(net, target, n_stages, prefix) + return net +end + +# One pass over the matches, appending directly to `net`. Matches are taken +# against the ORIGINAL parts only: everything appended carries the `prefix` head, which the name +# filters cannot match, and freezing the bound makes that explicit. +function _attach_all!(net, target::AtEvent, n_stages, prefix) + chain_heads = Dict{Any, Int}() # chain head label -> species index of stage 1 + for t in 1:nt(net) + _startswith(tname(net, t), target.transition) || continue + key = _net_product(net, t) + chain, _ = _chain_labels(net, key, n_stages, prefix, target.transition) + head = get!(chain_heads, chain[1]) do + # A chain may pre-exist from an earlier `attach_observation` call: validate and glue. + first(_ensure_chain!(net, key, n_stages, prefix, target.transition)) + end + # An identical re-application already has this output arc. + any(==(head), AlgebraicPetri.outputs(net, t)) || + add_outputs!(net, 1, [t], [head]) # the recorded event + end + return net +end + +function _attach_all!(net, target::AtCompartment, n_stages, prefix) + for s in 1:ns(net) + _startswith(sname(net, s), target.species) || continue + src = sname(net, s) + source = _head_symbol(src) + chain, _ = _chain_labels(net, s, n_stages, prefix, source) + tap = _relabel_head(src, Symbol("obs_inflow_", _flat(src))) + head, built = _ensure_chain!(net, s, n_stages, prefix, source) + taps = [t for t in 1:nt(net) if tname(net, t) == tap] + if built + isempty(taps) || _incompatible_chain(chain[1], n_stages) + add_transitions!(net, 1; tname = [tap]) + t = nt(net) + add_inputs!(net, 1, [t], [s]) + add_outputs!(net, 1, [t], [s]) # catalytic: source returned + add_outputs!(net, 1, [t], [head]) # into the chain + elseif length(taps) != 1 || !_is_observation_tap(net, only(taps), s, head) + _incompatible_chain(chain[1], n_stages) + end + end + return net +end + +_stage_head(prefix::Symbol, source::Symbol, stage::Int) = Symbol(prefix, "_", source, "_", stage) + +# Chain labels for the group keyed by species `key_idx`: one per stage, plus the delay transitions. +# +# `source` is part of the name, not decoration. Dropping it breaks two contracts at once: +# every `I` stage of a multi-stage compartment would key to the SAME chain label and all but the +# first would be skipped as already-tapped; and `_parse_observation_stage` (observation_chain.jl) +# requires a `__` leaf, so `observation_layout` — and hence +# `ConfigurableEpi`'s `StateLayout` — would not recognise the accumulators as observation states +# at all and would treat them as core model state. +function _chain_labels(net, key_idx, n_stages, prefix, source) + key = sname(net, key_idx) + species = [_relabel_head(key, _stage_head(prefix, source, i)) for i in 1:n_stages] + delays = [ + _relabel_head( + key, + Symbol( + prefix, "_", source, "_", i, "_to_", prefix, "_", source, "_", i + 1 + ), + ) for i in 1:(n_stages - 1) + ] + return species, delays +end + +function _ensure_chain!(net, key_idx, n_stages, prefix, source) + chain, delays = _chain_labels(net, key_idx, n_stages, prefix, source) + extended_chain, extended_delays = _chain_labels( + net, key_idx, n_stages + 1, prefix, source + ) + + species_matches = [ + [s for s in 1:ns(net) if sname(net, s) == label] for label in chain + ] + delay_matches = [ + [t for t in 1:nt(net) if tname(net, t) == label] for label in delays + ] + next_species = [s for s in 1:ns(net) if sname(net, s) == extended_chain[end]] + next_delay = [t for t in 1:nt(net) if tname(net, t) == extended_delays[end]] + + has_existing_parts = + any(matches -> !isempty(matches), species_matches) || + any(matches -> !isempty(matches), delay_matches) || + !isempty(next_species) || !isempty(next_delay) + + if !has_existing_parts + first_new = ns(net) + 1 + _build_chain!(net, chain, delays, first_new) + return first_new, true + end + + valid = + all(matches -> length(matches) == 1, species_matches) && + all(matches -> length(matches) == 1, delay_matches) && + isempty(next_species) && isempty(next_delay) + + if valid + stages = only.(species_matches) + for (i, matches) in enumerate(delay_matches) + transition = only(matches) + valid &= AlgebraicPetri.inputs(net, transition) == [stages[i]] + valid &= AlgebraicPetri.outputs(net, transition) == [stages[i + 1]] + end + end + + valid || _incompatible_chain(chain[1], n_stages) + return only(first(species_matches)), false +end + +function _incompatible_chain(head, n_stages) + throw( + ArgumentError( + "observation chain $head already exists with structure incompatible with " * + "n_stages=$n_stages; reuse its original stage count or choose a different prefix", + ), + ) +end + +function _is_observation_tap(net, transition, source, head) + inputs = AlgebraicPetri.inputs(net, transition) + outputs = AlgebraicPetri.outputs(net, transition) + return inputs == [source] && length(outputs) == 2 && + count(==(source), outputs) == 1 && count(==(head), outputs) == 1 +end + +# The compartment part of a (possibly nested) label: `(:I1, :a)` -> `:I1`, `:I1` -> `:I1`. +_head_symbol(label::Symbol) = label +_head_symbol(label::Tuple) = _head_symbol(label[1]) + +# Append a chain's species and delay transitions to `P`, the species starting at index +# `first_new`. Returns nothing; mutates `P`. +function _build_chain!(P, chain, delays, first_new) + add_species!(P, length(chain); sname = chain) + if !isempty(delays) + add_transitions!(P, length(delays); tname = delays) + end + for (i, _) in enumerate(delays) + # Delay transitions are numbered after whatever transitions the rule already carries. + t = nt(P) - length(delays) + i + add_inputs!(P, 1, [t], [first_new + i - 1]) + add_outputs!(P, 1, [t], [first_new + i]) + end + return nothing +end diff --git a/AlgebraicEpiMech/src/Stratification.jl b/AlgebraicEpiMech/src/Stratification.jl new file mode 100644 index 0000000..fc3eafe --- /dev/null +++ b/AlgebraicEpiMech/src/Stratification.jl @@ -0,0 +1,130 @@ +# Typing for stratification, in the context meaning a set of static labels for stratum +# carried by population individuals + +""" +Represents a stratification structure for epidemiological models. +""" +abstract type Stratification <: EpiMechModel end + +""" +Contact-based stratification for epidemiological models. + +Represents stratifications where populations are divided into strata with +contact-based interactions (full contact matrix between all pairs). This +generalizes age groups, geographic regions modelled as cross-region contact (rather than explicit movement), +risk groups, etc. + +Multiple contact stratifications can be stacked via `compose_stratifications` to +create product strata (e.g., age × geography). + +# Fields +- `stratum_names::Vector{Symbol}`: Names of strata (e.g., [:child, :adult] or [:urban, :rural]) +- `label::Symbol`: Semantic label for the stratification (e.g., :age, :geography, :risk) + +Reflexive boxes (disease, reversion, and waning) are controlled via the +`include_reflexives` keyword on `create_model`, not on the struct itself. + +# Examples +```julia +# Age stratification +age = ContactStratification([:child, :adult], :age) + +# Geographic stratification +geo = ContactStratification([:urban, :rural], :geography) + +# Compose to create age × geography strata +typing = OnePopulationTyping() +sir = create_model(typing, SIR()) +age_sir = typed_product(sir, create_model(typing, age)) +age_geo_sir = typed_product(age_sir, create_model(typing, geo)) +# Result: 4 strata (childxurban, childxrural, adultxurban, adultxrural) +``` + +See also: [`AgeStratification`](@ref), [`GeographicStratification`](@ref) +""" +struct ContactStratification <: Stratification + stratum_names::Vector{Symbol} + label::Symbol + + function ContactStratification( + stratum_names::Vector{Symbol}, label::Symbol + ) + if isempty(stratum_names) + error("ContactStratification requires at least one stratum") + end + if length(stratum_names) != length(unique(stratum_names)) + error("Stratum names must be unique") + end + return new(stratum_names, label) + end +end + +""" +Compose two `ContactStratification`s via Cartesian product. +# Arguments +- `a::ContactStratification`: First contact stratification +- `b::ContactStratification`: Second contact stratification +# Returns +- `ContactStratification`: New contact stratification representing the product of `a` and + `b`, with combined stratum names and a composite label. +# Examples +```julia +age = ContactStratification([:child, :adult], :age) +geo = ContactStratification([:urban, :rural], :geography) +age_geo = compose_stratifications(age, geo) +# Resulting stratum names: [:childxurban, :childxrural, :adultxurban, :adultxrural] +# Resulting label: :age_x_geography +``` +""" +function compose_stratifications(a::ContactStratification, b::ContactStratification) + combined_names = [ + Symbol(string(sa), "x", string(sb)) for sa in a.stratum_names + for sb in b.stratum_names + ] + combined_label = Symbol(string(a.label), "_x_", string(b.label)) + return ContactStratification(combined_names, combined_label) +end + +*(a::ContactStratification, b::ContactStratification) = compose_stratifications(a, b) + +#------------------------------------- +# Convenience constructors +#------------------------------------- + +""" + AgeStratification(age_group_names::Vector{Symbol}) + AgeStratification(names...) + +Convenience constructor for age-based contact stratification. + +Creates a `ContactStratification` with label `:age`. Provided for backward +compatibility and semantic clarity. + +# Examples +```julia +age = AgeStratification([:child, :adult, :elderly]) +age = AgeStratification(:child, :adult, :elderly) # splatting syntax +``` +""" +AgeStratification(names::Vector{Symbol}) = ContactStratification(names, :age) +AgeStratification(names...) = ContactStratification(collect(Symbol.(names)), :age) + +""" + GeographicStratification(location_names::Vector{Symbol}) + GeographicStratification(names...) + +Convenience constructor for geography-based contact stratification. + +Creates a `ContactStratification` with label `:geography`. Represents +geographic regions with cross-region contact (not movement-based models). + +# Examples +```julia +geo = GeographicStratification([:urban, :suburban, :rural]) +geo = GeographicStratification(:urban, :suburban, :rural) # splatting syntax +``` +""" +GeographicStratification(names::Vector{Symbol}) = ContactStratification(names, :geography) +function GeographicStratification(names...) + return ContactStratification(collect(Symbol.(names)), :geography) +end diff --git a/AlgebraicEpiMech/src/construction_helpers.jl b/AlgebraicEpiMech/src/construction_helpers.jl new file mode 100644 index 0000000..0aa6af3 --- /dev/null +++ b/AlgebraicEpiMech/src/construction_helpers.jl @@ -0,0 +1,192 @@ +# Helpers for constructing compartmental model UWDs +# Not exported - used internally by create_model_uwd methods + +""" +Create the Susceptible (S) compartment junction. + +The returned junction ID is used directly when adding infection and reversion mechanisms +(for example, I→S in SIS). +""" +function set_S_junction!(uwd::RelationDiagram, typing::EpidemiologicalTyping) + error("set_S_junction! not implemented for typing type $(typeof(typing)).") +end + +function set_S_junction!(uwd::RelationDiagram, typing::OnePopulationTyping) + pop_type = typing.population_type + return add_junction!(uwd, pop_type, variable = :S) +end + +function set_S_junction!(uwd::RelationDiagram, typing::UninfectedInfectedTyping) + uninfected_type = typing.uninfected_type + return add_junction!(uwd, uninfected_type, variable = :S) +end + +""" +Create stratum junctions: one per stratum for `OnePopulationTyping`, or an uninfected and +infected pair per stratum for `UninfectedInfectedTyping`. + +Returns a tuple of junction IDs for the stratum: +- For `OnePopulationTyping`: (stratum_junction, stratum_junction) +- For `UninfectedInfectedTyping`: (uninfected_stratum_junction, infected_stratum_junction) +""" +function set_stratum_junction!( + uwd::RelationDiagram, typing::OnePopulationTyping, stratum::Symbol + ) + pop_type = typing.population_type + stratum_junction = add_junction!(uwd, pop_type, variable = stratum) + return (stratum_junction, stratum_junction) +end + +function set_stratum_junction!( + uwd::RelationDiagram, typing::UninfectedInfectedTyping, stratum::Symbol + ) + uninfected_type = typing.uninfected_type + infected_type = typing.infected_type + # Use unique variable names by prefixing with type + uninfected_stratum_junction = add_junction!( + uwd, uninfected_type, variable = Symbol( + string(stratum) * + "_U" + ) + ) + infected_stratum_junction = add_junction!( + uwd, infected_type, variable = Symbol( + string(stratum) * + "_I" + ) + ) + return (uninfected_stratum_junction, infected_stratum_junction) +end + +""" +Generate junction variable names for multi-stage compartments. + +When `number_of_stages == 1`, return the conventional base name (`:E` or `:I`). +For multiple stages, append the stage number to produce stable, readable names such as +`:E1`, `:E2`, and `:I1`. These names aid visualization, inspection, and downstream lookup; +progression mechanisms themselves are connected using junction IDs. + +# Examples +```julia +variable_name(:E, 1, 1) # Returns :E (single stage) +variable_name(:E, 2, 3) # Returns :E2 (stage 2 of 3) +variable_name(:I, 1, 4) # Returns :I1 (stage 1 of 4) +``` +""" +function variable_name(var_name, stage, number_of_stages) + if number_of_stages == 1 + return var_name + else + return Symbol("$(var_name)$(stage)") + end +end + +""" +Create a sequence of compartment stages with progression transitions between them. + +This is the workhorse for multi-stage compartments, enabling Erlang-distributed dwell times. +It creates N junctions (for example, E1→E2→E3 or I1→I2→I3) and chains them with +progression mechanisms. + +# Construction behavior +- Returns the vector of ALL junction IDs (not just first/last) for flexible mechanism attachment +- First stage receives incoming transitions (e.g., S+I→E1), last stage feeds next compartment + +# Arguments +- `uwd`: The undirected wiring diagram to modify (follows convention of other mutation functions) +- `variable_symbol`: Base compartment name (:E or :I) +- `number_of_stages`: Number of sequential stages to create +- `pop_type`: The population type symbol for the junctions +- `typing`: The epidemiological typing for mechanism dispatch +""" +function add_stages!( + uwd::RelationDiagram, variable_symbol::Symbol, number_of_stages::Int, + pop_type::Symbol, typing::EpidemiologicalTyping + ) + number_of_stages >= 1 || throw(ArgumentError("number_of_stages must be at least 1")) + + # Create junctions for each stage + junctions = [ + add_junction!( + uwd, + pop_type, + variable = variable_name(variable_symbol, stage, number_of_stages) + ) + for stage in 1:number_of_stages + ] + # Add progression between stages + for stage in 1:(number_of_stages - 1) + add_disease_progression!(uwd, junctions[stage], junctions[stage + 1], typing) # e.g. E_stage → E_(stage+1) or I_stage → I_(stage+1) + end + return junctions +end + +""" +Populate an SI model UWD with junctions and infection mechanisms. + +# Construction Role +This is the base layer for all direct-infection models (SI, SIR, SIS). It establishes: +1. S compartment (via `set_S_junction!`) +2. I compartment chain (potentially multi-stage via `add_stages!`) +3. Infection dynamics: every I stage can infect S, but new infections enter I1 only + +Returns (uwd, S_junction, last_I_junction) to enable extensions: +- SIR adds last_I_junction → R +- SIS adds last_I_junction → S_junction (reversion) + +The multi-stage capability means SI can represent gamma-distributed infectious periods +even before extending to SIR. +""" +function setup_basic!(uwd::RelationDiagram, typing::EpidemiologicalTyping, model::SI) + # Add junctions for S compartment + S_junction = set_S_junction!(uwd, typing) + pop_type = get_infected_type(typing) + + I_junctions = add_stages!(uwd, :I, model.number_I_stages, pop_type, typing) + + # Add infection to first I stage + for I_junction in I_junctions + add_infection!(uwd, S_junction, I_junction, I_junctions[1], typing) # S + I → I + I + end + + return (uwd, S_junction, I_junctions[end]) +end + +""" +Populate an SEI model UWD with junctions, exposure, and progression mechanisms. + +# Construction Role +This is the base layer for all exposure-based models (SEI, SEIR, SEIS, SEIRS). It establishes: +1. S compartment (via `set_S_junction!`) +2. E compartment chain (potentially multi-stage for gamma-distributed latent period) +3. I compartment chain (potentially multi-stage for gamma-distributed infectious period) +4. Exposure dynamics: every I stage can expose S, but new exposures enter E1 only +5. Progression: E_last → I_1 connects the latent and infectious stages + +Returns (uwd, S_junction, last_I_junction) to enable extensions: +- SEIR adds last_I_junction → R +- SEIS adds last_I_junction → S_junction (reversion) +- SEIRS adds both (recovery then waning) + +The dual multi-stage capability enables independent control of latent and infectious period +distributions - critical for realistic disease modeling. +""" +function setup_basic!(uwd::RelationDiagram, typing::EpidemiologicalTyping, model::SEI) + # Add junctions for S compartment + S_junction = set_S_junction!(uwd, typing) + pop_type = get_infected_type(typing) + + # For multiple E and I stages, create junctions for each stage + E_junctions = add_stages!(uwd, :E, model.number_E_stages, pop_type, typing) + I_junctions = add_stages!(uwd, :I, model.number_I_stages, pop_type, typing) + + # add infection to first E stage from any infection stage + for I_junction in I_junctions + add_infection!(uwd, S_junction, I_junction, E_junctions[1], typing) # S + I → E + I + end + + # Add progression from last E stage to first I stage + add_disease_progression!(uwd, E_junctions[end], I_junctions[1], typing) # E_last → I_1 + + return (uwd, S_junction, I_junctions[end]) +end diff --git a/AlgebraicEpiMech/src/create_model.jl b/AlgebraicEpiMech/src/create_model.jl new file mode 100644 index 0000000..b07a834 --- /dev/null +++ b/AlgebraicEpiMech/src/create_model.jl @@ -0,0 +1,111 @@ +# Unified create_model function for all model types +# Multiple dispatch on model type parameter handles different model types +# (CompartmentalModel, AgeStratification, MultiStrainModel) + +""" +Create a typed Petri net model from an epidemiological typing strategy and model specification. + +This is the unified entry point for creating all model types in AlgebraicEpiMech. Through +multiple dispatch on the `model` parameter, it supports any subtype of `EpiMechModel`. +The undirected wiring diagram (UWD) provides an intermediate structural representation, recording +how compartments and mechanisms connect independently of the final typed Petri net constructed by +`oapply_typed`. + +# Functionality Overview +The `create_model` function builds the typed +Petri net representation of the epidemiological-mechanical model in four steps: + +1. **Build Type System**: Materializes the typing strategy as a `LabelledPetriNet` +2. **Create UWD**: Builds an undirected wiring diagram with model structure and box names +3. **Generate Transition Names**: Creates unique names for ODE parameters via dispatch +4. **Apply Typing**: Uses `oapply_typed` to create a typed Petri net + +The resulting typed Petri net can be used directly for ODE simulation or composed with other +typed Petri nets via `typed_product` to build complex hierarchical models. + +# Arguments +- `typing::EpidemiologicalTyping`: The strategy defining possible compartment and transition types. +- `model::EpiMechModel`: The model specification (uses dispatch for different types) + +# Keyword Arguments +Keyword arguments are forwarded to `create_model_uwd`. Model-specific keywords include: +- `include_reflexives::Bool=true` (for `ContactStratification`): Controls whether per-stratum + reflexive boxes (disease, reversion, and waning) are added. Required for `typed_product` + composition; disable for standalone use to avoid zero-effect transitions. + +# Returns +- `ACSetTransformation`: A typed Petri net with two components accessible via: + - `dom(typed_model)`: Extract the underlying `LabelledPetriNet` for ODE solving + - `codom(typed_model)`: Access the Petri-net type system + +# Model-Specific Transition Naming +Transition names are generated automatically based on model type: +- **Compartmental**: `transmission_S_I`, `E_to_I`, `I_to_R` (describes flow between compartments) +- **Age Stratification**: `child_child`, `child_adult`, `adult_child` (infectee_infector pattern) +- **Multistrain**: Strain identifiers like `h1n1`, `h3n2` (enables clean composition) + +These names become ODE parameters and differ from box names (`:transmission`, `:density`) used for composition. + +# Examples +```julia +# Basic compartmental model +typing = OnePopulationTyping() +seir = create_model(typing, SEIR()) +seir_pn = dom(seir) # Extract for ODE solving +# Transitions: :transmission_S_I, :E_to_I, :I_to_R + +# Multi-stage compartments (gamma-distributed delays) +si_multi = create_model(typing, SI(number_I_stages=3)) +# Transitions: :transmission_S_I1, :transmission_S_I2, :transmission_S_I3, :I1_to_I2, :I2_to_I3 + +# Age-structured model +age_strat = AgeStratification([:child, :adult, :elderly]) +age_model = create_model(typing, age_strat) +# Transitions: :child_child, :child_adult, :child_elderly, :adult_child, ... + +# Multistrain model (no cross-immunity) +strains = NoCrossImmunity([:h1n1, :h3n2, :seasonal_b]) +strain_model = create_model(typing, strains) + +# Compose models via typed_product +sir = create_model(typing, SIR()) +age = create_model(typing, AgeStratification([:child, :adult])) +age_sir = typed_product(sir, age) +age_sir_pn = dom(age_sir) +# Result: S_child, I_child, R_child, S_adult, I_adult, R_adult compartments +# Transition names become tuples: (:transmission_S_I, :child_adult) +``` + +# Compositional Design +Models expressed as typed Petri nets compose via `typed_product` using their common +type-system codomain when box names match. The `vectorfield_flat` function automatically handles tuple transition names +in composed models for ODE parameter assignment. + +See also: [`create_model_uwd`](@ref), [`generate_transition_names`](@ref), [`type_system`](@ref), + `typed_product` (from `AlgebraicPetri.TypedPetri`), [`vectorfield_flat`](@ref) +""" +function create_model( + typing::EpidemiologicalTyping, + model::EpiMechModel; + include_reflexives::Bool = true + ) + # Materialize the typing strategy as the Petri-net codomain for the typed petri net. + codomain = type_system(typing) + + # Create the UWD for the model (dispatch on include_reflexives via positional arg) + uwd = _create_model_uwd(typing, model, include_reflexives) + + # Generate transition names (dispatch on model type) + transition_names = generate_transition_names(uwd, model) + + # Apply oapply_typed to create the typed Petri net + typed_model = oapply_typed(codomain, uwd, transition_names) + + return typed_model +end + +# Intermediate layer: converts include_reflexives kwarg to positional for dispatch +_create_model_uwd(typing, model, ::Bool) = create_model_uwd(typing, model) +function _create_model_uwd(typing, model::ContactStratification, include_reflexives::Bool) + return create_model_uwd(typing, model; include_reflexives) +end diff --git a/AlgebraicEpiMech/src/create_model_uwd.jl b/AlgebraicEpiMech/src/create_model_uwd.jl new file mode 100644 index 0000000..f980467 --- /dev/null +++ b/AlgebraicEpiMech/src/create_model_uwd.jl @@ -0,0 +1,528 @@ +# Constructors for compartmental epidemiological models +# Uses functions from construction_helpers.jl and mechanisms.jl +# Dispatch methods to create compartmental models based on typing and model instances + +""" +Create an undirected wiring diagram (UWD) for a compartmental model within a specific typing. + +This function returns the undirected wiring diagram that defines the +compartmental structure and transitions, using the type names from the typing. +The compartments are typed according to the typing's population structure. +The returned UWD does not declare an outer boundary/ports; `oapply_typed` derives the +boundary from its junctions when materializing the typed Petri net. + +# Arguments +- `typing::EpidemiologicalTyping`: The population typing defining type names +- `model::CompartmentalModel`: The compartmental model instance (SI(), SEI(), SIR(), SEIR(), etc.) + +# Returns +- Undirected wiring diagram defining the compartmental model structure with typing-specific types + +# Examples +```julia +# Create SIR model UWD with single population typing +typing = OnePopulationTyping(population_type = :Individual) +sir_uwd = create_model_uwd(typing, SIR()) # Uses Individual type + +# Create SEIR model UWD with uninfected/infected typing +typing = UninfectedInfectedTyping(uninfected_type = :Susceptible, infected_type = :Infectious) +seir_uwd = create_model_uwd(typing, SEIR()) # S::Susceptible, E,I,R::Infectious +``` +""" +function create_model_uwd(typing::EpidemiologicalTyping, model::CompartmentalModel) + error("create_model_uwd not implemented for typing type $(typeof(typing)) and model type $(typeof(model)). Implement a method for this combination.") +end + +""" +Construct a complete undirected wiring diagram (UWD) for SI or SEI models. + +Populates an initially empty UWD with compartments and mechanisms via `setup_basic!`. +These are the base models that other compartmental models extend from. + +# Usage with oapply_typed +The returned UWD can be applied to the typing using `oapply_typed` to create a typed +Petri net, enabling composition with other typed Petri nets (e.g., demographic processes, +interventions). +""" +function create_model_uwd(typing::EpidemiologicalTyping, model::Union{SI, SEI}) + uwd = RelationDiagram(Symbol[]) + (uwd, _, _) = setup_basic!(uwd, typing, model) + return uwd +end + +""" +Construct a complete SIR model UWD by extending SI with recovery. + +Builds on the SI base model (S→I dynamics with multi-stage I support) and adds: +- R (recovered) compartment +- Recovery transition: I_last → R + +The multi-stage I capability from SI is preserved, enabling gamma-distributed infectious +periods. Recovery always occurs from the last I stage, ensuring proper sequencing. + +# Usage with oapply_typed +Apply to typing with `oapply_typed` to create a typed Petri net for composition. +""" +function create_model_uwd(typing::EpidemiologicalTyping, model::SIR) + uwd = RelationDiagram(Symbol[]) + + # Get SI base model + uwd, _, + last_I_junction = setup_basic!(uwd, typing, SI(number_I_stages = model.number_I_stages)) + + pop_type = get_infected_type(typing) + + # Add R compartment junction + R_junction = add_junction!(uwd, pop_type, variable = :R) + + # Add recovery: last I stage → R + add_disease_progression!(uwd, last_I_junction, R_junction, typing) + + return uwd +end + +""" +Construct a complete SIS model UWD by extending SI with waning immunity. + +Builds on the SI base model (S→I dynamics with multi-stage I support) and adds: +- Reversion transition: I_last → S (waning immunity) + +Represents diseases where recovery returns individuals to susceptible state. The multi-stage +I capability enables realistic infectious period distributions. Reversion occurs from the +last I stage back to S. + +# Usage with oapply_typed +Apply to typing with `oapply_typed` to create a typed Petri net for composition. +""" +function create_model_uwd(typing::EpidemiologicalTyping, model::SIS) + uwd = RelationDiagram(Symbol[]) + + # Get SI base model + uwd, S_junction, + last_I_junction = setup_basic!(uwd, typing, SI(number_I_stages = model.number_I_stages)) + + # Add reversion: last I stage → S (waning immunity) + add_reversion_progression!(uwd, last_I_junction, S_junction, typing) + + return uwd +end + +""" +Construct a complete SEIR model UWD by extending SEI with recovery. + +Builds on the SEI base model (S→E→I dynamics with multi-stage E and I support) and adds: +- R (recovered) compartment +- Recovery transition: I_last → R + +The dual multi-stage capability (independent E and I stages) enables independent control +of latent and infectious period distributions. Recovery always occurs from the last I stage. + +# Usage with oapply_typed +Apply to typing with `oapply_typed` to create a typed Petri net for composition. +""" +function create_model_uwd(typing::EpidemiologicalTyping, model::SEIR) + uwd = RelationDiagram(Symbol[]) + + # Get SEI base model + uwd, _, + last_I_junction = setup_basic!( + uwd, + typing, + SEI( + number_E_stages = model.number_E_stages, + number_I_stages = model.number_I_stages + ) + ) + + pop_type = get_infected_type(typing) + + # Add R compartment junction + R_junction = add_junction!(uwd, pop_type, variable = :R) + + # Add recovery: last I stage → R + add_disease_progression!(uwd, last_I_junction, R_junction, typing) + + return uwd +end + +""" +Construct a complete SEIS model UWD by extending SEI with waning immunity. + +Builds on the SEI base model (S→E→I dynamics with multi-stage E and I support) and adds: +- Reversion transition: I_last → S (waning immunity) + +Represents diseases with latent periods but no lasting immunity. The dual multi-stage +capability enables independent control of latent and infectious period distributions. +Reversion occurs from the last I stage back to S. + +# Usage with oapply_typed +Apply to typing with `oapply_typed` to create a typed Petri net for composition. +""" +function create_model_uwd(typing::EpidemiologicalTyping, model::SEIS) + uwd = RelationDiagram(Symbol[]) + + # Get SEI base model + uwd, S_junction, + last_I_junction = setup_basic!( + uwd, + typing, + SEI( + number_E_stages = model.number_E_stages, + number_I_stages = model.number_I_stages + ) + ) + + # Add reversion: last I stage → S (waning immunity) + add_reversion_progression!(uwd, last_I_junction, S_junction, typing) + + return uwd +end + +""" +Construct a complete SEIRS model UWD by extending SEI with recovery and waning immunity. + +Builds on the SEI base model (S→E→I dynamics with multi-stage E and I support) and adds: +- R (recovered) compartment +- Recovery transition: I_last → R +- Waning immunity transition: R → S + +Represents diseases with latent periods and temporary immunity. The dual multi-stage capability +enables independent control of latent and infectious period distributions. Recovery occurs from +the last I stage, and waning immunity returns individuals from R to S. + +# Usage with oapply_typed +Apply to typing with `oapply_typed` to create a typed Petri net for composition. +""" +function create_model_uwd(typing::EpidemiologicalTyping, model::SEIRS) + uwd = RelationDiagram(Symbol[]) + + # Get SEIR model (which already calls SEI and adds R) + # SEIR returns just uwd since it's already complete + uwd, S_junction, + last_I_junction = setup_basic!( + uwd, + typing, + SEI( + number_E_stages = model.number_E_stages, + number_I_stages = model.number_I_stages + ) + ) + + # Add R compartment junction + pop_type = get_infected_type(typing) + R_junction = add_junction!(uwd, pop_type, variable = :R) + + # Add recovery: last I stage → R + add_disease_progression!(uwd, last_I_junction, R_junction, typing) + + # Add waning immunity: R → S + add_reversion_progression!(uwd, R_junction, S_junction, typing) + + return uwd +end + +# ============================================================================ +# Multistrain Model UWD Construction +# ============================================================================ + + +""" +Construct an undirected wiring diagram (UWD) for a no cross-immunity multistrain model. + +Creates a strain-stratified UWD where each strain operates independently. The strains +are represented as separate junctions, and each strain has its own set of disease +transition boxes (`:transmission`, `:disease`, `:reversion`) that will +compose with corresponding boxes from a compartmental model via `typed_product`. + +All box types are always included; `typed_product` will only compose boxes that exist in +both UWDs, naturally filtering out non-matching transitions. That filtering is silent — a +transition present in one factor and absent from the other is dropped without error — which is +why each factor carries a reflexive box per transition type it wants preserved. + +There is deliberately no reflexive observation box. Observation is attached to the COMPOSED net +by `attach_observation`, so a stratification factor never has to know it exists. + +# Arguments +- `typing::OnePopulationTyping`: The required single-population typing +- `multistrain::NoCrossImmunity`: The multistrain model configuration with strain names + +# Returns +- Undirected wiring diagram with strain junctions and transition boxes for composition + +# Examples +```julia +# Create a 3-strain model with custom names +typing = OnePopulationTyping() +multistrain = NoCrossImmunity([:h1n1, :h3n2, :b]) +strain_uwd = create_model_uwd(typing, multistrain) + +# Compose with compartmental model +sir_typed = create_model(typing, SIR()) +strain_typed = create_model(typing, multistrain) +combined = typed_product(sir_typed, strain_typed) +``` +""" +function create_model_uwd( + typing::OnePopulationTyping, + multistrain::NoCrossImmunity + ) + strain_names = multistrain.strain_names + + # For NoCrossImmunity, use OnePopulationTyping type system + pop_type = get_infected_type(typing) + + uwd = RelationDiagram(Symbol[]) + + # Create junction for each strain + strain_junctions = [ + add_junction!(uwd, pop_type, variable = strain_name) + for strain_name in strain_names + ] + + # Add all transition boxes for each strain + # typed_product will only compose boxes that match between UWDs + for strain_junction in strain_junctions + # Transmission: strain + strain → strain + strain (like S + I → I + I) + add_infection!(uwd, strain_junction, strain_junction, strain_junction, typing) + + # Disease progression: strain → strain (like E → I or I → R) + add_disease_progression!(uwd, strain_junction, strain_junction, typing) + + # Reversion/waning: strain → strain (like I → S or R → S) + add_reversion_progression!(uwd, strain_junction, strain_junction, typing) + + end + + return uwd +end + +""" +Construct an undirected wiring diagram (UWD) for a complete cross-immunity multistrain model. + +Creates a strain-structured UWD where all strains share a common susceptible pool but have +separate infected compartments. Infection by any strain depletes the shared susceptible pool +and confers immunity to all strains. + +The UWD has: +- One shared uninfected/susceptible junction +- N infected strain junctions +- Transmission boxes connecting shared susceptible to each strain's infected +- Disease progression boxes for each strain's infected compartments +- Reversion boxes to return to shared susceptible pool + +# Arguments +- `typing::UninfectedInfectedTyping`: The required uninfected/infected typing +- `multistrain::CompleteCrossImmunity`: The multistrain model configuration with strain names + +# Returns +- Undirected wiring diagram with shared susceptible and strain-specific infected junctions + +# Examples +```julia +# Create a 2-strain competing model +typing = UninfectedInfectedTyping() +multistrain = CompleteCrossImmunity([:wild_type, :variant]) +strain_uwd = create_model_uwd(typing, multistrain) + +# Compose with compartmental model +sir_typed = create_model(typing, SIR()) +strain_typed = create_model(typing, multistrain) +combined = typed_product(sir_typed, strain_typed) +``` +""" +function create_model_uwd( + typing::UninfectedInfectedTyping, + multistrain::CompleteCrossImmunity + ) + strain_names = multistrain.strain_names + + # For CompleteCrossImmunity, use UninfectedInfectedTyping type system + uninfected_type = get_uninfected_type(typing) + infected_type = get_infected_type(typing) + + uwd = RelationDiagram(Symbol[]) + + # Create shared uninfected/susceptible junction + shared_uninfected = add_junction!(uwd, uninfected_type, variable = :susceptible) + + # Create infected junction for each strain + infected_junctions = [ + add_junction!(uwd, infected_type, variable = strain_name) + for strain_name in strain_names + ] + + # Add transition boxes for each strain + for infected_junction in infected_junctions + # Transmission: shared_S + strain_I → strain_I + strain_I + # This connects the shared susceptible pool to each strain's infected compartment + add_infection!(uwd, shared_uninfected, infected_junction, infected_junction, typing) + + # Disease progression: strain_I → strain_I (e.g., I → R within strain) + add_disease_progression!(uwd, infected_junction, infected_junction, typing) + + # Reversion/waning: strain_I → shared_S (e.g., R → S or I → S) + # This returns individuals to the shared susceptible pool + add_reversion_progression!(uwd, infected_junction, shared_uninfected, typing) + + end + + return uwd +end + +""" +Construct an undirected wiring diagram (UWD) for contact strata. + +Creates an UWD where transmission boxes represent contact patterns between different strata. +Example strata include age groups, risk levels, or demographic divisions, but only covers instantaneous +transmission dynamics, e.g. movement between geographic regions would require additional demographic modeling. + +The UWD structure depends on the typing: +- `OnePopulationTyping`: Single junction per age group (combined susceptible/infected) +- `UninfectedInfectedTyping`: Two junctions per age group (uninfected and infected) + +All transmission boxes are created with the `:transmission` name, enabling composition +with compartmental models that have `:transmission` boxes (SI, SIR, SEIR, etc.). + +# Arguments +- `typing::EpidemiologicalTyping`: The population typing defining type system +- `model::ContactStratification`: The contact stratification configuration with stratum names + +# Returns +- Undirected wiring diagram with stratum junctions and transmission boxes for composition + +# Mathematical Structure +For n strata, creates n² transmission boxes representing the contact matrix: +- Diagonal boxes (i→i): within-stratum transmission +- Off-diagonal boxes (i→j, i≠j): between-stratum transmission + + +# Composition Behavior +When composed with a compartmental model UWD via `typed_product`, the transmission boxes +will align based on the `:transmission` name, allowing the contact structure to modulate +the disease dynamics defined in the compartmental model. +""" +function create_model_uwd( + typing::EpidemiologicalTyping, model::ContactStratification; + include_reflexives::Bool = true + ) + uwd = RelationDiagram(Symbol[]) + + # Create junction for each stratum + # These are tuples (uninfected_stratum_junction, infected_stratum_junction) in case needed for more complex typings + stratum_junction_tuples = [ + set_stratum_junction!(uwd, typing, stratum) + for stratum in model.stratum_names + ] + + for stratum_infectee in stratum_junction_tuples + for stratum_infector in stratum_junction_tuples + # Add transmission between strata (full contact matrix) + add_infection!( + uwd, + stratum_infectee[1], # infectee junction (the stratum being infected) + stratum_infector[2], # infector junction (the stratum causing infection) + stratum_infectee[2], # first infected output junction (the stratum being infected is unchanged for ContactStratification) + typing + ) + end + end + + # Add per-stratum reflexive boxes for non-transmission transitions. + # These ensure typed_product preserves disease, reversion, and waning transitions + # when composing with compartmental models. + # Can be disabled via include_reflexives=false for standalone use. + if include_reflexives + for stratum in stratum_junction_tuples + add_disease_progression!(uwd, stratum[2], stratum[2], typing) + add_reversion_progression!(uwd, stratum[2], stratum[1], typing) + add_uninfected_density_progression!(uwd, stratum[1], stratum[1], typing) + end + end + + return uwd +end + +# ============================================================================ +# Immune-history model UWD construction (issue #279) +# ============================================================================ + +""" +Immune-history stratification requires the uninfected/infected type split. +""" +function create_model_uwd(::OnePopulationTyping, ::ImmuneHistory) + return error( + "ImmuneHistory requires UninfectedInfectedTyping — the uninfected/infected " * + "split is what carries immune history; OnePopulationTyping has no uninfected type." + ) +end + +""" +Construct the UWD for the `ImmuneHistory` stratification **factor** over a +`UninfectedInfectedTyping`. + +This is a stratification factor (like `ContactStratification`), meant to be +`typed_product`-composed with a disease model — the disease model supplies +`S→E→I→R`, while this factor supplies the immune-status structure. + +- Uninfected junctions `U_h`, one per immune-history class. +- Infected junctions `(h,i)` = "history `h`, currently fighting strain `i`", one per + (class, susceptible-strain) pair. +- **`:transmission`** (escape-selective): `U_h + (h',i) → (h,i)` for every susceptible + class `h` (`i ∉ h`) and every infector `(h',i)`. A class immune to `i` simply has no + such box, so the pullback drops that infection — the escape. +- **`:disease`** is reflexive on each `(h,i)`, so the base model's progression runs + inside a fixed `(h,i)`. +- **`:reversion`** is off-diagonal: `(h,i) → U_{recover}` with `recover = h∪{i}` + (`FullHistory`) or `{i}` (`LatestInfection`). This is the one stratum-changing move + — recovering from `i` folds it into the immune history. +- **`:waning`** reflexive on each `U_h` (inert with SEIRS; lets a waning-bearing base + compose). + +# Examples +```julia +typing = UninfectedInfectedTyping() +history = create_model(typing, ImmuneHistory([:current, :invader])) +model = typed_product(create_model(typing, SEIRS()), history) +``` +""" +function create_model_uwd(typing::UninfectedInfectedTyping, model::ImmuneHistory) + strains = model.strain_names + mode = model.mode + uninfected_type = get_uninfected_type(typing) + infected_type = get_infected_type(typing) + + uwd = RelationDiagram(Symbol[]) + + # Uninfected immune-history classes. + U = Dict{_ImmuneClass, Int}() + for h in _immune_classes(mode, strains) + U[h] = add_junction!(uwd, uninfected_type, variable = _uninfected_name(h)) + end + + # Infected (history, current-strain) junctions; index the infectors of each strain. + pairs = _infected_pairs(mode, strains) + Iof = Dict{Tuple{_ImmuneClass, Symbol}, Int}() + infectors = Dict{Symbol, Vector{Int}}(s => Int[] for s in strains) + for (h, i) in pairs + Iof[(h, i)] = add_junction!(uwd, infected_type, variable = _infected_name(h, i)) + push!(infectors[i], Iof[(h, i)]) + end + + for (h, i) in pairs + # Escape-selective infection: U_h + (any strain-i infector) → (h,i). + for inf_j in infectors[i] + add_infection!(uwd, U[h], inf_j, Iof[(h, i)], typing) + end + # Reflexive :disease so the base's E→I→R progression runs within this (h,i). + add_disease_progression!(uwd, Iof[(h, i)], Iof[(h, i)], typing) + # Off-diagonal :reversion — the history-incrementing move. + add_reversion_progression!(uwd, Iof[(h, i)], U[_recover_to(mode, h, i)], typing) + end + + # Reflexive :waning on each uninfected class (inert with SEIRS). + for h in keys(U) + add_uninfected_density_progression!(uwd, U[h], U[h], typing) + end + + return uwd +end diff --git a/AlgebraicEpiMech/src/docstrings.jl b/AlgebraicEpiMech/src/docstrings.jl new file mode 100644 index 0000000..d4c3083 --- /dev/null +++ b/AlgebraicEpiMech/src/docstrings.jl @@ -0,0 +1,24 @@ +@template (FUNCTIONS, METHODS, MACROS) = """ + $(TYPEDSIGNATURES) +$(DOCSTRING) +""" + +@template (TYPES) = """ + $(TYPEDEF) +$(DOCSTRING) + +--- +## Fields +$(TYPEDFIELDS) +""" + +@template MODULES = """ +$(DOCSTRING) + +--- +## Exports +$(EXPORTS) +--- +## Imports +$(IMPORTS) +""" diff --git a/AlgebraicEpiMech/src/generate_transition_names.jl b/AlgebraicEpiMech/src/generate_transition_names.jl new file mode 100644 index 0000000..f1647bf --- /dev/null +++ b/AlgebraicEpiMech/src/generate_transition_names.jl @@ -0,0 +1,235 @@ +# Functions for generating transition names from UWD structures +# Used by create_model to generate unique transition names for different model types + +""" +Generate transition names from UWD structure for compartmental models. + +For boxes with 2 ports (one input, one output), generates names in the format `input_to_output`. +For boxes with 4 ports (two inputs, two outputs), generates names in the format `box_name_input1_input2` +where inputs are sorted in reverse alphabetical order for consistency (e.g., `transmission_S_I`). + +# Arguments +- `uwd`: The undirected wiring diagram with species names in junction :variable attributes +- `model::CompartmentalModel`: The compartmental model (used for dispatch) + +# Returns +- `Vector{Symbol}`: Vector of unique transition names, one per box in the UWD + +# Examples +```julia +# Used internally by create_model +typing = OnePopulationTyping() +uwd = create_model_uwd(typing, SEIR()) +names = generate_transition_names(uwd, SEIR()) +# Returns: [:transmission_S_I, :E_to_I, :I_to_R] +``` + +See also: [`create_model`](@ref) +""" +function generate_transition_names(uwd, model::CompartmentalModel) + transition_names = Symbol[] + + for box_id in boxes(uwd) + box_ports = ports(uwd, box_id) + num_ports = length(box_ports) + box_name = subpart(uwd, box_id, :name) + + # Get junctions connected to this box's ports + box_junctions = [junction(uwd, port) for port in box_ports] + + if num_ports == 2 + # One input, one output: name as "input_to_output" + input_junction = box_junctions[1] + output_junction = box_junctions[2] + input_name = subpart(uwd, input_junction, :variable) + output_name = subpart(uwd, output_junction, :variable) + transition_name = Symbol(input_name, "_to_", output_name) + push!(transition_names, transition_name) + elseif num_ports == 4 + # Two inputs, two outputs: name as "box_name_input1_input2" + # First two ports are inputs, last two are outputs (by convention in mechanisms.jl) + input1_junction = box_junctions[1] + input2_junction = box_junctions[2] + input1_name = subpart(uwd, input1_junction, :variable) + input2_name = subpart(uwd, input2_junction, :variable) + # Sort input names in reverse alphabetical order for consistency (e.g., S_I not I_S) + input_names_sorted = sort( + [string(input1_name), string(input2_name)], rev = true + ) + transition_name = Symbol(box_name, "_", join(input_names_sorted, "_")) + push!(transition_names, transition_name) + else + # Fallback: just use the box name + # This shouldn't happen in standard compartmental models + transition_name = box_name + push!(transition_names, transition_name) + end + end + + return transition_names +end + +""" +Generate transition names from UWD structure for age stratification models. + +For transmission boxes with 4 ports (two inputs, two outputs), generates names capturing +the age-to-age contact pattern in the format `infectee_infector`. This naming +convention clearly identifies which age group is being infected (infectee) by which age +group (infector), enabling interpretation of age-structured contact matrices. + +For boxes with 2 ports (one input, one output), generates names in the format `input_to_output`. +This fallback handles any non-standard box types that might appear in extended models. + +# Arguments +- `uwd`: The undirected wiring diagram with age group names in junction :variable attributes +- `model::AgeStratification`: The age stratification model (used for dispatch) + +# Returns +- `Vector{Symbol}`: Vector of transition names, one per box in the UWD + +# Examples +```julia +# Used internally by create_model +typing = OnePopulationTyping() +age_strat = AgeStratification([:child, :adult]) +uwd = create_model_uwd(typing, age_strat) +names = generate_transition_names(uwd, age_strat) +# Returns: [:child_child, :child_adult, :adult_child, :adult_adult] +# Representing: child←child, child←adult, adult←child, adult←adult transmission +``` + +See also: [`create_model`](@ref), [`ContactStratification`](@ref), [`AgeStratification`](@ref) +""" +function generate_transition_names(uwd, model::ContactStratification) + transition_names = Symbol[] + + for box_id in boxes(uwd) + box_ports = ports(uwd, box_id) + num_ports = length(box_ports) + box_name = subpart(uwd, box_id, :name) + + # Get junctions connected to this box's ports + box_junctions = [junction(uwd, port) for port in box_ports] + + if num_ports == 4 + # Transmission: infectee_infector (2 inputs, 2 outputs) + input1_name = subpart(uwd, box_junctions[1], :variable) + input2_name = subpart(uwd, box_junctions[2], :variable) + input_names = [string(input1_name), string(input2_name)] + push!(transition_names, Symbol(join(input_names, "_"))) + elseif num_ports == 2 || num_ports == 3 + # Per-stratum reflexive boxes (disease, reversion, and waning): + # Use first port's junction variable as the stratum name. + # After typed_product, these compose as e.g. (:I_to_R, :child) → I_to_R_child + first_name = subpart(uwd, box_junctions[1], :variable) + push!(transition_names, first_name) + else + push!(transition_names, box_name) + end + end + + return transition_names +end + + +""" +Generate transition names for a multistrain model from UWD structure. + +For multistrain models, transition names are just the strain names. This creates cleaner +composed transition names like `(:transmission_S_I, :h1n1)` rather than +`(:transmission_S_I, :trans_h1n1)`. The strain model UWD is never used alone, only for composition. + +For NoCrossImmunity: All box ports connect to the same strain junction. +For CompleteCrossImmunity: Boxes connect to both shared susceptible and strain-specific infected junctions. + +# Arguments +- `uwd`: The undirected wiring diagram with strain junctions +- `multistrain::MultiStrainModel`: The multistrain model containing strain names + +# Returns +- `Vector{Symbol}`: Vector of transition names (just strain names), one per box in the UWD + +# Examples +```julia +# Used internally by create_model +typing = OnePopulationTyping() +multistrain = NoCrossImmunity([:h1n1, :h3n2]) +uwd = create_model_uwd(typing, multistrain) +names = generate_transition_names(uwd, multistrain) +# Returns one strain name per box; 3 boxes per strain +# (transmission, disease, reversion), so 2 strains → 6 names. +``` +""" +function _strain_junction(uwd, box_id, ::NoCrossImmunity) + box_junctions = unique([junction(uwd, port) for port in ports(uwd, box_id)]) + return only(box_junctions) +end + +function _strain_junction(uwd, box_id, ::CompleteCrossImmunity) + box_junctions = unique([junction(uwd, port) for port in ports(uwd, box_id)]) + strain_junctions = filter(box_junctions) do junction_id + return subpart(uwd, junction_id, :variable) != :susceptible + end + return only(strain_junctions) +end + +function generate_transition_names(uwd, model::MultiStrainModel) + transition_names = Symbol[] + + for box_id in boxes(uwd) + strain_junction = _strain_junction(uwd, box_id, model) + strain_variable = subpart(uwd, strain_junction, :variable) + + # Transition name is just the strain name + # This creates clean composed names like (:transmission_S_I, :h1n1) + push!(transition_names, strain_variable) + end + + return transition_names +end + +""" +Generate transition names for the `ImmuneHistory` stratification factor. These become +the second component of the composed `typed_product` transition names. + +- **`:transmission`** (4-port) → `infect___by_`, + exposing the infecting strain (recovered from the infector junction) so the composed + name stays keyable per strain. +- **reflexive / reversion boxes** (2–3 port) → the first junction's variable, mirroring + `ContactStratification` — e.g. the `:disease` reflexives and the off-diagonal + `:reversion` on `(h,i)` are named by that `(h,i)`, and `:waning` by its + `U_h`. Composition with the base disambiguates them (base name differs). + +# Examples +```julia +factor = create_model_uwd(UninfectedInfectedTyping(), ImmuneHistory([:current, :invader])) +names = generate_transition_names(factor, ImmuneHistory([:current, :invader])) +# transmission names begin `infect_current…` / `infect_invader…` +``` +""" +function generate_transition_names(uwd, model::ImmuneHistory) + strains = model.strain_names + transition_names = Symbol[] + + for box_id in boxes(uwd) + box_ports = ports(uwd, box_id) + box_name = subpart(uwd, box_id, :name) + vars = [subpart(uwd, junction(uwd, port), :variable) for port in box_ports] + + if length(box_ports) == 4 && box_name == :transmission + # Ports: [infectee U_h, infector (h',i), first-infected (h,i), infector]. + strain = _strain_of(vars[2], strains) + susceptible = replace(string(vars[1]), r"^U_" => "") + infector = replace(string(vars[2]), r"^.*_from_" => "") + push!( + transition_names, + Symbol("infect_", strain, "_", susceptible, "_by_", infector) + ) + else + # Reflexive / reversion boxes: name by the first junction (ContactStratification style). + push!(transition_names, vars[1]) + end + end + + return transition_names +end diff --git a/AlgebraicEpiMech/src/mechanisms.jl b/AlgebraicEpiMech/src/mechanisms.jl new file mode 100644 index 0000000..433cca9 --- /dev/null +++ b/AlgebraicEpiMech/src/mechanisms.jl @@ -0,0 +1,120 @@ +# Helper functions for epidemiological mechanisms added to undirected wiring diagram (UWD) +# This file give primitives for programmatic construction of a UWDs + +""" +Add infection mechanism: Infectee + Infector → first_infected + Infector + +General infection mechanism where an infectee (typically S) interacts with an infector (typically I) +to produce a newly infected individual that enters the first_infected compartment, while the infector +remains unchanged. The `first_infected_junction` can be: +- An I stage for SIR-like models (direct infection) +- An E stage for SEIR-like models (exposure before infectiousness) + +The `typing` argument determines the population types used. +""" +function add_infection!( + uwd, infectee_junction, infector_junction, + first_infected_junction, typing::OnePopulationTyping + ) + pop_type = typing.population_type + infect_box = add_box!( + uwd, [pop_type, pop_type, pop_type, pop_type], name = :transmission + ) + ports_infect = ports(uwd, infect_box) + set_junction!(uwd, ports_infect[1], infectee_junction) # Infectee input + set_junction!(uwd, ports_infect[2], infector_junction) # Infector input + set_junction!(uwd, ports_infect[3], first_infected_junction) # First infected output + return set_junction!(uwd, ports_infect[4], infector_junction) # Infector output (unchanged) +end + +function add_infection!( + uwd, infectee_junction, infector_junction, + first_infected_junction, typing::UninfectedInfectedTyping + ) + uninfected_type = typing.uninfected_type + infected_type = typing.infected_type + infect_box = add_box!( + uwd, [uninfected_type, infected_type, infected_type, infected_type], + name = :transmission + ) + ports_infect = ports(uwd, infect_box) + set_junction!(uwd, ports_infect[1], infectee_junction) # Infectee input + set_junction!(uwd, ports_infect[2], infector_junction) # Infector input + set_junction!(uwd, ports_infect[3], first_infected_junction) # First infected output + return set_junction!(uwd, ports_infect[4], infector_junction) # Infector output (unchanged) +end + +""" +Add disease progression mechanism: from_compartment → to_compartment + +Used for transitions like E→I, I→R, etc. (infected → infected transitions) +""" +function add_disease_progression!( + uwd, from_junction, to_junction, typing::OnePopulationTyping + ) + pop_type = typing.population_type + disease_box = add_box!(uwd, [pop_type, pop_type], name = :disease) + ports_disease = ports(uwd, disease_box) + set_junction!(uwd, ports_disease[1], from_junction) # input + return set_junction!(uwd, ports_disease[2], to_junction) # output +end + +function add_disease_progression!( + uwd, from_junction, to_junction, typing::UninfectedInfectedTyping + ) + pop_type = typing.infected_type + density_box = add_box!(uwd, [pop_type, pop_type], name = :disease) + ports_density = ports(uwd, density_box) + set_junction!(uwd, ports_density[1], from_junction) # input + return set_junction!(uwd, ports_density[2], to_junction) # output +end + +""" +Add uninfected density progression mechanisms. This is useful for transitions like waning partial +immunity (uninfected → uninfected transitions, e.g., V → S). +""" +function add_uninfected_density_progression!( + uwd, from_junction, to_junction, typing::OnePopulationTyping + ) + pop_type = typing.population_type + waning_box = add_box!(uwd, [pop_type, pop_type], name = :waning) + ports_waning = ports(uwd, waning_box) + set_junction!(uwd, ports_waning[1], from_junction) # input + return set_junction!(uwd, ports_waning[2], to_junction) # output +end + +function add_uninfected_density_progression!( + uwd, from_junction, to_junction, typing::UninfectedInfectedTyping + ) + pop_type = typing.uninfected_type + density_box = add_box!(uwd, [pop_type, pop_type], name = :waning) + ports_density = ports(uwd, density_box) + set_junction!(uwd, ports_density[1], from_junction) # input + return set_junction!(uwd, ports_density[2], to_junction) # output +end + +""" +Add reversion progression mechanism: e.g. R → S (infected → uninfected transitions) + +Used for transitions from infected compartments back to uninfected compartments. +""" +function add_reversion_progression!( + uwd, from_junction, to_junction, typing::OnePopulationTyping + ) + pop_type = typing.population_type + reversion_box = add_box!(uwd, [pop_type, pop_type], name = :reversion) + ports_reversion = ports(uwd, reversion_box) + set_junction!(uwd, ports_reversion[1], from_junction) # input + return set_junction!(uwd, ports_reversion[2], to_junction) # output +end + +function add_reversion_progression!( + uwd, from_junction, to_junction, typing::UninfectedInfectedTyping + ) + infected_type = typing.infected_type + uninfected_type = typing.uninfected_type + reversion_box = add_box!(uwd, [infected_type, uninfected_type], name = :reversion) + ports_reversion = ports(uwd, reversion_box) + set_junction!(uwd, ports_reversion[1], from_junction) # input + return set_junction!(uwd, ports_reversion[2], to_junction) # output +end diff --git a/AlgebraicEpiMech/src/typing.jl b/AlgebraicEpiMech/src/typing.jl new file mode 100644 index 0000000..c21b34e --- /dev/null +++ b/AlgebraicEpiMech/src/typing.jl @@ -0,0 +1,147 @@ +""" +Abstract base type for epidemiological typing strategies. + +An `EpidemiologicalTyping` describes how populations and transitions are typed when +constructing epidemiological models. [`type_system`](@ref) materializes the strategy as +the `LabelledPetriNet` used as the codomain of a typed Petri net. + +## Reference + +[Libkind et al. (2023) _An algebraic framework for structured epidemic modelling_](https://royalsocietypublishing.org/rsta/article/380/2233/20210309/112239) +""" +abstract type EpidemiologicalTyping end + +""" + OnePopulationTyping(; population_type = :Population) + +Typing strategy in which every compartment belongs to one population type and therefore +has the same available stratifications. + +# Examples +```julia +typing = OnePopulationTyping() +typing = OnePopulationTyping(population_type = :CityPopulation) +``` +""" +@kwdef struct OnePopulationTyping <: EpidemiologicalTyping + "The symbolic name for the population type." + population_type::Symbol = :Population +end + +""" + UninfectedInfectedTyping(; uninfected_type = :Uninfected, infected_type = :Infected) + +Typing strategy that distinguishes uninfected and infected populations. The two +populations may have different stratifications. + +# Examples +```julia +typing = UninfectedInfectedTyping() +typing = UninfectedInfectedTyping( + uninfected_type = :Susceptible, + infected_type = :Infectious, +) +``` + +See also: [`type_system`](@ref), [`OnePopulationTyping`](@ref) +""" +@kwdef struct UninfectedInfectedTyping <: EpidemiologicalTyping + "The symbolic name for the uninfected population." + uninfected_type::Symbol = :Uninfected + "The symbolic name for the infected population." + infected_type::Symbol = :Infected +end + +""" + type_system(typing::EpidemiologicalTyping, population_transitions...) + +Materialize an epidemiological typing strategy as the `LabelledPetriNet` used as the +codomain of a typed Petri net. Concrete `EpidemiologicalTyping` subtypes must implement +this interface. +""" +function type_system(typing::EpidemiologicalTyping, population_transitions...) + return error( + "`type_system` not implemented for typing $(typeof(typing)). " * + "Implement a method for this typing or use " * + "OnePopulationTyping/UninfectedInfectedTyping." + ) +end + +""" + type_system(typing::OnePopulationTyping, population_transitions...) + +Create the single-population Petri-net type system described by `typing`. Additional +transition specifications are appended to the standard `:transmission`, `:disease`, +`:reversion`, and `:waning` transitions. + +# Examples +```julia +typing = OnePopulationTyping(population_type = :Individual) +codomain = type_system( + typing, + :birth => (:Individual => (:Individual, :Individual)), +) +``` +""" +function type_system(typing::OnePopulationTyping, population_transitions...) + population_type = typing.population_type + return LabelledPetriNet( + [population_type], + :transmission => ( + (population_type, population_type) => (population_type, population_type) + ), + :disease => (population_type => population_type), + :reversion => (population_type => population_type), + :waning => (population_type => population_type), + population_transitions... + ) +end + +""" + type_system(typing::UninfectedInfectedTyping, population_transitions...) + +Create the uninfected/infected Petri-net type system described by `typing`. Additional +transition specifications are appended to the standard `:transmission`, `:disease`, +`:reversion`, and `:waning` transitions. + +# Examples +```julia +typing = UninfectedInfectedTyping( + uninfected_type = :Susceptible, + infected_type = :Infectious, +) +codomain = type_system( + typing, + :vaccination => (:Susceptible => :Susceptible), +) +``` +""" +function type_system(typing::UninfectedInfectedTyping, population_transitions...) + uninfected_type = typing.uninfected_type + infected_type = typing.infected_type + return LabelledPetriNet( + [uninfected_type, infected_type], + :transmission => ( + (uninfected_type, infected_type) => (infected_type, infected_type) + ), + :disease => (infected_type => infected_type), + :reversion => (infected_type => uninfected_type), + :waning => (uninfected_type => uninfected_type), + population_transitions... + ) +end + +""" +Return the type label used for infected compartments (E, I, and R). + +This interface lets model construction remain agnostic to the selected epidemiological +typing strategy. +""" +get_infected_type(typing::OnePopulationTyping) = typing.population_type +get_infected_type(typing::UninfectedInfectedTyping) = typing.infected_type + +""" +Return the type label used for uninfected compartments such as S. +""" +get_uninfected_type(typing::OnePopulationTyping) = typing.population_type +get_uninfected_type(typing::UninfectedInfectedTyping) = typing.uninfected_type diff --git a/AlgebraicEpiMech/src/vectorfield_flat.jl b/AlgebraicEpiMech/src/vectorfield_flat.jl new file mode 100644 index 0000000..286f411 --- /dev/null +++ b/AlgebraicEpiMech/src/vectorfield_flat.jl @@ -0,0 +1,209 @@ +""" +Recursively flatten a tuple of symbols and join them with underscores. + +# Examples +```julia +flatten_symbols((:x, :y, :z)) # returns :x_y_z +flatten_symbols(((:x, :y), :z)) # returns :x_y_z +flatten_symbols((((:a, :b), :c), :d)) # returns :a_b_c_d +``` +""" +function flatten_symbols(name::Tuple) + # Recursively flatten all elements + flattened = _flatten_recursive(name) + # Join all symbols with underscores + symbol_strings = [string(sym) for sym in flattened] + return Symbol(join(symbol_strings, "_")) +end + +function flatten_symbols(name::Symbol) + return name +end + +# Handle recursive flattening of tuples +function _flatten_recursive(name::Tuple) + isempty(name) && throw(ArgumentError("cannot flatten an empty tuple")) + flattened = Symbol[] + for element in name + append!(flattened, _flatten_recursive(element)) + end + return flattened +end + +# Handle single symbol case +_flatten_recursive(name::Symbol) = Symbol[name] + +function _require_unique_flattened_names(names::Vector{Symbol}, part::AbstractString) + seen = Set{Symbol}() + collisions = Set{Symbol}() + for name in names + name in seen && push!(collisions, name) + push!(seen, name) + end + isempty(collisions) || throw( + ArgumentError( + "flattened $part names must be unique; underscore joining produced " * + "duplicate name(s): $(join(sort!(collect(collisions)), ", "))", + ), + ) + return nothing +end + +""" +Generate an ODE vectorfield function from a Petri net using mass action kinetics as per +`AlgebraicPetri.vectorfield`. The only difference is that species and transition names +are flattened using `flatten_symbols`, converting nested tuples into single symbols +joined by underscores. + +# Arguments +- `pn::AbstractPetriNet`: A Petri net representing the reaction network + +# Returns +A function `(du, u, p, t) -> du` where: +- `du`: Array/Dict to store computed derivatives +- `u`: Current state (species concentrations), indexed by flattened species names +- `p`: Parameters (rate constants), indexed by flattened transition names +- `t`: Current time + +This function complies with SciML conventions for in-place ODE vectorfields. + +# Performance + +Everything that can be resolved from the net alone — flattened names, the mass-action input +lists, the nonzero stoichiometry — is resolved once at build time, and the right-hand side is a +sparse loop over arcs rather than a dense loop over `transitions × species`. Name lookup into the +arguments is also done once per concrete argument type: `u`, `du` and `p` may be anything whose +`propertynames` lists the flattened names in the same order as its integer indexing (`LVector`, +`NamedTuple`, ...), and that position map is cached on first sight of the type; an `AbstractDict` +is looked up by name on every call. +""" +function vectorfield_flat(pn::AbstractPetriNet) + S = ns(pn) + T = nt(pn) + species_syms = [flatten_symbols(sname(pn, j)) for j in 1:S] + transition_syms = [flatten_symbols(tname(pn, i)) for i in 1:T] + _require_unique_flattened_names(species_syms, "species") + _require_unique_flattened_names(transition_syms, "transition") + + tm = TransitionMatrices(pn) + # Mass-action inputs as `(species, multiplicity)`, ascending species: the factor is `u_j^k` + # applied in species order — the reference's `prod(u_j^input[i, j] for j)` minus its exact + # `u^0 = 1` factors — so the products agree bit for bit for any multiplicity. + inputs = [Tuple{Int, Int}[] for _ in 1:T] + # Net stoichiometry `output - input`, nonzero entries only, ascending species (a zero entry + # contributes an exact `+0.0` in the dense form). + stoichiometry = [Tuple{Int, Int}[] for _ in 1:T] + for i in 1:T, j in 1:S + k = tm.input[i, j] + k == 0 || push!(inputs[i], (j, k)) + c = tm.output[i, j] - k + c == 0 || push!(stoichiometry[i], (j, c)) + end + + # Caching objects for inplace calculations + u_positions = _NamePositions() + du_positions = _NamePositions() + p_positions = _NamePositions() + + return function vectorfield!(du, u, p, t) + u_m = _values_by_name(u_positions, species_syms, u) + p_m = _values_by_name(p_positions, transition_syms, p) + du_m = _slots_by_name(du_positions, species_syms, du) + + Ty = valtype(du) + acc = _zeroed_accumulator(promote_type(Float64, Ty), Ty, du_m, S) + @inbounds for i in 1:T + factor = one(eltype(u_m)) + for (j, k) in inputs[i] + factor *= u_m[j]^k + end + rate = convert(Ty, AlgebraicPetri.valueat(p_m[i], u, t) * factor) + # only non-zero stoichoimetry looped over + for (j, c) in stoichiometry[i] + acc[j] += rate * c + end + end + _writeback_accumulator!(acc, du_m, S) + return du + end +end + +# Cache for position of each flattened name inside an argument, resolved once per concrete argument type. +# Single-entry and lock-free: the right-hand side is called from threaded ensemble propagation, so +# readers take an acquire load and a miss publishes an immutable `(type, positions)` pair with a +# release store — two threads racing on the same miss compute identical content. +mutable struct _NamePositions + @atomic entry::Union{Nothing, Tuple{DataType, Vector{Int}}} +end +_NamePositions() = _NamePositions(nothing) + +function _positions(cache::_NamePositions, syms::Vector{Symbol}, x) + Tx = typeof(x) + entry = @atomic :acquire cache.entry + if entry !== nothing && entry[1] === Tx + return entry[2] + end + names = collect(propertynames(x)) + positions = Vector{Int}(undef, length(syms)) + for (k, s) in enumerate(syms) + i = findfirst(==(s), names) + i === nothing && throw( + ArgumentError( + "`$s` is not a name of the $(Tx) argument; the vectorfield needs every " * + "flattened species and transition name of the net to be present", + ), + ) + positions[k] = i + end + @atomic :release cache.entry = (Tx, positions) + return positions +end + +# An argument read in net order without copying it: `v[k]` is the value of the k-th net name. +# Materializing the gather would cost one allocation of `length(syms)` per right-hand-side call +struct _ByName{X} + x::X + positions::Vector{Int} +end +Base.@propagate_inbounds Base.getindex(v::_ByName, k::Int) = v.x[v.positions[k]] +Base.@propagate_inbounds Base.setindex!(v::_ByName, val, k::Int) = (v.x[v.positions[k]] = val) +Base.eltype(::Type{_ByName{X}}) where {X} = eltype(X) +Base.length(v::_ByName) = length(v.positions) + +# The same for a dictionary, which indexes only by key. +struct _ByKey{X} + x::X + syms::Vector{Symbol} +end +Base.getindex(v::_ByKey, k::Int) = v.x[v.syms[k]] +Base.setindex!(v::_ByKey, val, k::Int) = (v.x[v.syms[k]] = val) +Base.eltype(::Type{_ByKey{X}}) where {X} = valtype(X) +Base.length(v::_ByKey) = length(v.syms) + +# Fast path: integer indexing through the cached positions (LVector, NamedTuple, ...). +_values_by_name(cache::_NamePositions, syms::Vector{Symbol}, x) = + _ByName(x, _positions(cache, syms, x)) +# By-name path for dictionaries, which index only by key. +_values_by_name(::_NamePositions, syms::Vector{Symbol}, x::AbstractDict) = [x[s] for s in syms] + +# Writable slots of `du` in net order, for the accumulation. +_slots_by_name(cache::_NamePositions, syms::Vector{Symbol}, du) = + _ByName(du, _positions(cache, syms, du)) +_slots_by_name(::_NamePositions, syms::Vector{Symbol}, du::AbstractDict) = _ByKey(du, syms) + +# The accumulator is a cleared `du` when its element type already is the dense form's +# accumulation type, and a zeroed temporary of that type otherwise (both cases resolved by dispatch). +function _zeroed_accumulator(::Type{A}, ::Type{A}, du_m, S) where {A} + @inbounds for k in 1:S + du_m[k] = zero(A) + end + return du_m +end +_zeroed_accumulator(::Type{A}, ::Type{Ty}, du_m, S) where {A, Ty} = zeros(A, S) +_writeback_accumulator!(acc::_ByName, du_m::_ByName, S) = nothing +_writeback_accumulator!(acc::_ByKey, du_m::_ByKey, S) = nothing +_writeback_accumulator!(acc, du_m, S) = ( + @inbounds for k in 1:S + du_m[k] = acc[k] + end +) diff --git a/AlgebraicEpiMech/test/Aqua.jl b/AlgebraicEpiMech/test/Aqua.jl new file mode 100644 index 0000000..e6f60f3 --- /dev/null +++ b/AlgebraicEpiMech/test/Aqua.jl @@ -0,0 +1,5 @@ +@testitem "Aqua.jl" begin + using Aqua + Aqua.test_all(AlgebraicEpiMech, ambiguities = false, persistent_tasks = false) + Aqua.test_ambiguities(AlgebraicEpiMech) +end diff --git a/AlgebraicEpiMech/test/Project.toml b/AlgebraicEpiMech/test/Project.toml new file mode 100644 index 0000000..7397179 --- /dev/null +++ b/AlgebraicEpiMech/test/Project.toml @@ -0,0 +1,10 @@ +[deps] +AlgebraicPetri = "4f99eebe-17bf-4e98-b6a1-2c4f205a959b" +AlgebraicEpiMech = "a6e39d6b-8b2d-49f3-aca0-f4154643f040" +Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +Catlab = "134e5e36-593f-5add-ad60-77f754baafbe" +LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800" +SeeToDee = "1c904df7-48cd-41e7-921b-d889ed9a470c" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a" +TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe" diff --git a/AlgebraicEpiMech/test/runtests.jl b/AlgebraicEpiMech/test/runtests.jl new file mode 100644 index 0000000..b9e874d --- /dev/null +++ b/AlgebraicEpiMech/test/runtests.jl @@ -0,0 +1,3 @@ +using TestItemRunner + +@run_package_tests diff --git a/AlgebraicEpiMech/test/test_age_stratification.jl b/AlgebraicEpiMech/test/test_age_stratification.jl new file mode 100644 index 0000000..297b391 --- /dev/null +++ b/AlgebraicEpiMech/test/test_age_stratification.jl @@ -0,0 +1,82 @@ +# DEPRECATED: This file tests the old AgeStratification API +# New tests should use test_contact_stratification.jl which covers +# ContactStratification and all its convenience constructors (AgeStratification, +# GeographicStratification, etc.) +# +# This file is kept temporarily for backward compatibility but uses the new +# ContactStratification field names (stratum_names instead of age_group_names) + +@testsnippet AgeStratificationSetup begin + using AlgebraicEpiMech + + age_strat = AgeStratification([:adult]) + age_strat_2 = AgeStratification([:child, :adult]) + age_strat_multi = AgeStratification([:child, :adult, :elderly]) +end + +@testitem "AgeStratification constructor (deprecated - see test_contact_stratification.jl)" setup = [AgeStratificationSetup] begin + @testset "Basic construction" begin + # Single age group + @test age_strat.stratum_names == [:adult] + @test age_strat.label == :age + @test length(age_strat.stratum_names) == 1 + + # Two age groups + @test age_strat_2.stratum_names == [:child, :adult] + @test age_strat_2.label == :age + @test length(age_strat_2.stratum_names) == 2 + + # Multiple age groups + @test age_strat_multi.stratum_names == [:child, :adult, :elderly] + @test age_strat_multi.label == :age + @test length(age_strat_multi.stratum_names) == 3 + end + + @testset "Validation" begin + # Empty age group vector should error + @test_throws ErrorException AgeStratification(Symbol[]) + + # Duplicate age group names should error + @test_throws ErrorException AgeStratification([:child, :adult, :child]) + @test_throws ErrorException AgeStratification([:adult, :adult]) + end +end + +@testitem "AgeStratification convenience constructor (deprecated - see test_contact_stratification.jl)" setup = [AgeStratificationSetup] begin + @testset "Splatted construction" begin + # Two age groups via splatting + age_strat = AgeStratification(:child, :adult) + @test age_strat.stratum_names == [:child, :adult] + @test age_strat.label == :age + @test length(age_strat.stratum_names) == 2 + + # Three age groups via splatting + age_strat_3 = AgeStratification(:child, :adult, :elderly) + @test age_strat_3.stratum_names == [:child, :adult, :elderly] + @test age_strat_3.label == :age + @test length(age_strat_3.stratum_names) == 3 + + # Single age group via splatting + age_strat_1 = AgeStratification(:population) + @test age_strat_1.stratum_names == [:population] + @test age_strat_1.label == :age + @test length(age_strat_1.stratum_names) == 1 + end + + @testset "String to Symbol conversion" begin + # Convenience constructor should convert strings to symbols + age_strat_2_splat = AgeStratification("child", "adult") + @test age_strat_2_splat.stratum_names == age_strat_2.stratum_names + @test age_strat_2_splat.label == :age + + # Mixed string and symbol input + age_strat_mixed = AgeStratification("child", :adult, "elderly") + @test age_strat_mixed.stratum_names == age_strat_multi.stratum_names + @test age_strat_mixed.label == :age + end + + @testset "Validation in convenience constructor" begin + # Duplicate names should still error + @test_throws ErrorException AgeStratification(:child, :adult, :child) + end +end diff --git a/AlgebraicEpiMech/test/test_compartmental_models.jl b/AlgebraicEpiMech/test/test_compartmental_models.jl new file mode 100644 index 0000000..7dd1d59 --- /dev/null +++ b/AlgebraicEpiMech/test/test_compartmental_models.jl @@ -0,0 +1,92 @@ +@testitem "Compartmental model constructors" begin + using AlgebraicEpiMech + + @testset "SI model" begin + # Single stage + model = SI() + @test model.number_I_stages == 1 + @test model.number_of_states == 2 # S + I + + # Multi-stage + model_multi = SI(number_I_stages = 3) + @test model_multi.number_I_stages == 3 + @test model_multi.number_of_states == 4 # S + I1 + I2 + I3 + end + + @testset "SEI model" begin + # Single stage + model = SEI() + @test model.number_E_stages == 1 + @test model.number_I_stages == 1 + @test model.number_of_states == 3 # S + E + I + + # Multi-stage + model_multi = SEI(number_E_stages = 2, number_I_stages = 3) + @test model_multi.number_E_stages == 2 + @test model_multi.number_I_stages == 3 + @test model_multi.number_of_states == 6 # S + E1 + E2 + I1 + I2 + I3 + end + + @testset "SIR model" begin + # Single stage + model = SIR() + @test model.number_I_stages == 1 + @test model.number_of_states == 3 # S + I + R + + # Multi-stage + model_multi = SIR(number_I_stages = 2) + @test model_multi.number_I_stages == 2 + @test model_multi.number_of_states == 4 # S + I1 + I2 + R + end + + @testset "SEIR model" begin + model = SEIR(number_E_stages = 2, number_I_stages = 2) + @test model.number_E_stages == 2 + @test model.number_I_stages == 2 + @test model.number_of_states == 6 # S + E1 + E2 + I1 + I2 + R + end + + @testset "SIS model" begin + model = SIS(number_I_stages = 3) + @test model.number_I_stages == 3 + @test model.number_of_states == 4 # S + I1 + I2 + I3 + end + + @testset "SEIS model" begin + model = SEIS(number_E_stages = 1, number_I_stages = 2) + @test model.number_E_stages == 1 + @test model.number_I_stages == 2 + @test model.number_of_states == 4 # S + E + I1 + I2 + end + + @testset "SEIRS model" begin + model = SEIRS() + @test model.number_E_stages == 1 + @test model.number_I_stages == 1 + @test model.number_of_states == 4 # S + E + I + R + end + + @testset "Type hierarchy" begin + @test SI() isa CompartmentalModel + @test SEI() isa CompartmentalModel + @test SIR() isa CompartmentalModel + @test SEIR() isa CompartmentalModel + @test SIS() isa CompartmentalModel + @test SEIS() isa CompartmentalModel + @test SEIRS() isa CompartmentalModel + end + + @testset "Stage counts must be positive" begin + @test_throws ArgumentError SI(number_I_stages = 0) + @test_throws ArgumentError SIR(number_I_stages = -1) + @test_throws ArgumentError SIS(number_I_stages = 0) + @test_throws ArgumentError SEI(number_E_stages = 0) + @test_throws ArgumentError SEI(number_I_stages = -1) + @test_throws ArgumentError SEIR(number_E_stages = -1) + @test_throws ArgumentError SEIR(number_I_stages = 0) + @test_throws ArgumentError SEIS(number_E_stages = 0) + @test_throws ArgumentError SEIS(number_I_stages = -1) + @test_throws ArgumentError SEIRS(number_E_stages = -1) + @test_throws ArgumentError SEIRS(number_I_stages = 0) + end +end diff --git a/AlgebraicEpiMech/test/test_construction_helpers.jl b/AlgebraicEpiMech/test/test_construction_helpers.jl new file mode 100644 index 0000000..ee89c1c --- /dev/null +++ b/AlgebraicEpiMech/test/test_construction_helpers.jl @@ -0,0 +1,199 @@ +@testsnippet HelperSetup begin + using AlgebraicPetri + using Catlab + using AlgebraicEpiMech + import AlgebraicEpiMech as aem + + # Test typings + one_pop_typing = OnePopulationTyping() + ui_typing = UninfectedInfectedTyping() +end + +@testitem "set_S_junction!" setup = [HelperSetup] begin + @testset "OnePopulationTyping" begin + uwd = RelationDiagram(Symbol[]) + + S_junction = aem.set_S_junction!(uwd, one_pop_typing) + + @test S_junction isa Integer + @test length(junctions(uwd)) == 1 # Only S junction added + @test subpart(uwd, S_junction, :variable) == :S + @test isempty(ports(uwd, outer = true)) + end + + @testset "UninfectedInfectedTyping" begin + uwd = RelationDiagram(Symbol[]) + + S_junction = aem.set_S_junction!(uwd, ui_typing) + + @test S_junction isa Integer + @test length(junctions(uwd)) == 1 # Only S junction added + @test subpart(uwd, S_junction, :variable) == :S + @test isempty(ports(uwd, outer = true)) + end +end + +@testitem "variable_name" setup = [HelperSetup] begin + @testset "Single stage returns base name" begin + @test aem.variable_name(:E, 1, 1) == :E + @test aem.variable_name(:I, 1, 1) == :I + end + + @testset "Multiple stages append stage numbers" begin + @test aem.variable_name(:E, 1, 3) == :E1 + @test aem.variable_name(:E, 2, 3) == :E2 + @test aem.variable_name(:E, 3, 3) == :E3 + + @test aem.variable_name(:I, 1, 4) == :I1 + @test aem.variable_name(:I, 2, 4) == :I2 + @test aem.variable_name(:I, 3, 4) == :I3 + @test aem.variable_name(:I, 4, 4) == :I4 + end + + @testset "Works with any symbol" begin + @test aem.variable_name(:R, 1, 2) == :R1 + @test aem.variable_name(:Custom, 5, 10) == :Custom5 + end +end + +@testitem "add_stages!" setup = [HelperSetup] begin + @testset "Single stage" begin + uwd = RelationDiagram(Symbol[]) + aem.set_S_junction!(uwd, one_pop_typing) + + pop_type = aem.get_infected_type(one_pop_typing) + I_junctions = aem.add_stages!(uwd, :I, 1, pop_type, one_pop_typing) + + @test length(I_junctions) == 1 + # Single stage, no progression mechanisms added + @test length(boxes(uwd)) == 0 + end + + @testset "Multiple stages with progression" begin + uwd = RelationDiagram(Symbol[]) + aem.set_S_junction!(uwd, one_pop_typing) + + pop_type = aem.get_infected_type(one_pop_typing) + I_junctions = aem.add_stages!(uwd, :I, 3, pop_type, one_pop_typing) + + @test length(I_junctions) == 3 + # 3 stages → 2 progression mechanisms (I1→I2, I2→I3) + @test length(boxes(uwd)) == 2 + @test isempty(ports(uwd, outer = true)) + end + + @testset "Works with UninfectedInfectedTyping" begin + uwd = RelationDiagram(Symbol[]) + aem.set_S_junction!(uwd, ui_typing) + + pop_type = aem.get_infected_type(ui_typing) + E_junctions = aem.add_stages!(uwd, :E, 2, pop_type, ui_typing) + + @test length(E_junctions) == 2 + # 2 E stages → 1 progression mechanism + @test length(boxes(uwd)) == 1 + + # Add I stages after E + I_junctions = aem.add_stages!(uwd, :I, 3, pop_type, ui_typing) + + @test length(I_junctions) == 3 + # 1 (E progression) + 2 (I progressions) = 3 total + @test length(boxes(uwd)) == 3 + @test isempty(ports(uwd, outer = true)) + end + + @testset "Rejects non-positive stage counts before mutating the UWD" begin + for count in (0, -1) + uwd = RelationDiagram(Symbol[]) + pop_type = aem.get_infected_type(one_pop_typing) + @test_throws ArgumentError aem.add_stages!( + uwd, :I, count, pop_type, one_pop_typing + ) + @test isempty(junctions(uwd)) + @test isempty(boxes(uwd)) + end + end +end + +@testitem "setup_basic! for SI" setup = [HelperSetup] begin + @testset "Single-stage SI" begin + model = SI() + uwd = RelationDiagram(Symbol[]) + + result_uwd, S_junction, + last_I_junction = aem.setup_basic!(uwd, one_pop_typing, model) + + @test result_uwd === uwd # Returns same UWD + @test S_junction isa Integer + @test last_I_junction isa Integer + @test length(junctions(uwd)) == 2 # S and I + # 1 transmission box (S + I → I + I) + @test length(boxes(uwd)) == 1 + end + + @testset "Multi-stage SI" begin + model = SI(number_I_stages = 3) + uwd = RelationDiagram(Symbol[]) + + result_uwd, S_junction, + last_I_junction = aem.setup_basic!(uwd, one_pop_typing, model) + + @test length(junctions(uwd)) == 4 # S + I1 + I2 + I3 + # 3 transmission boxes (one per I stage) + 2 progressions (I1→I2, I2→I3) + @test length(boxes(uwd)) == 5 + end + + @testset "Works with UninfectedInfectedTyping" begin + model = SI(number_I_stages = 2) + uwd = RelationDiagram(Symbol[]) + + result_uwd, S_junction, last_I_junction = aem.setup_basic!(uwd, ui_typing, model) + + @test length(junctions(uwd)) == 3 # S + I1 + I2 + # 2 transmission boxes + 1 progression + @test length(boxes(uwd)) == 3 + end +end + +@testitem "setup_basic! for SEI" setup = [HelperSetup] begin + @testset "Single-stage SEI" begin + model = SEI() + uwd = RelationDiagram(Symbol[]) + + result_uwd, S_junction, + last_I_junction = aem.setup_basic!(uwd, one_pop_typing, model) + + @test result_uwd === uwd + @test S_junction isa Integer + @test last_I_junction isa Integer + @test length(junctions(uwd)) == 3 # S, E, I + # 1 transmission (S + I → E + I) + 1 progression (E → I) + @test length(boxes(uwd)) == 2 + end + + @testset "Multi-stage E and I" begin + model = SEI(number_E_stages = 2, number_I_stages = 3) + uwd = RelationDiagram(Symbol[]) + + result_uwd, S_junction, + last_I_junction = aem.setup_basic!(uwd, one_pop_typing, model) + + @test length(junctions(uwd)) == 6 # S + E1 + E2 + I1 + I2 + I3 + # 3 transmission boxes (one per I stage) + # + 1 E progression (E1→E2) + # + 2 I progressions (I1→I2, I2→I3) + # + 1 E→I progression (E2→I1) + @test length(boxes(uwd)) == 7 + end + + @testset "Different E and I stage counts" begin + model = SEI(number_E_stages = 3, number_I_stages = 1) + uwd = RelationDiagram(Symbol[]) + + result_uwd, S_junction, last_I_junction = aem.setup_basic!(uwd, ui_typing, model) + + @test length(junctions(uwd)) == 5 # S + E1 + E2 + E3 + I + # 1 transmission + 2 E progressions + 1 E→I progression + @test length(boxes(uwd)) == 4 + end +end diff --git a/AlgebraicEpiMech/test/test_contact_stratification.jl b/AlgebraicEpiMech/test/test_contact_stratification.jl new file mode 100644 index 0000000..b9a0949 --- /dev/null +++ b/AlgebraicEpiMech/test/test_contact_stratification.jl @@ -0,0 +1,139 @@ +# Unit tests for ContactStratification and its convenience constructors +# This generalizes the old AgeStratification tests to work with any ContactStratification + +@testsnippet ContactStratificationSetup begin + using AlgebraicEpiMech + + # Various ContactStratification instances + age_single = AgeStratification([:adult]) + age_two = AgeStratification([:child, :adult]) + age_multi = AgeStratification([:child, :adult, :elderly]) + + geo_two = GeographicStratification([:urban, :rural]) + geo_multi = GeographicStratification([:urban, :suburban, :rural]) + + # Direct ContactStratification for testing the base type + risk_strat = ContactStratification([:low, :high], :risk) +end + +@testitem "ContactStratification constructor" setup = [ContactStratificationSetup] begin + @testset "Basic construction" begin + # Single stratum + @test age_single.stratum_names == [:adult] + @test age_single.label == :age + @test length(age_single.stratum_names) == 1 + + # Two strata + @test age_two.stratum_names == [:child, :adult] + @test age_two.label == :age + @test length(age_two.stratum_names) == 2 + + # Multiple strata + @test age_multi.stratum_names == [:child, :adult, :elderly] + @test age_multi.label == :age + @test length(age_multi.stratum_names) == 3 + + # Direct ContactStratification + @test risk_strat.stratum_names == [:low, :high] + @test risk_strat.label == :risk + end + + @testset "Validation" begin + # Empty stratum vector should error + @test_throws ErrorException ContactStratification(Symbol[], :test) + @test_throws ErrorException AgeStratification(Symbol[]) + @test_throws ErrorException GeographicStratification(Symbol[]) + + # Duplicate stratum names should error + @test_throws ErrorException ContactStratification([:a, :b, :a], :test) + @test_throws ErrorException AgeStratification([:child, :adult, :child]) + @test_throws ErrorException GeographicStratification([:urban, :urban]) + end +end + +@testitem "AgeStratification convenience constructor" setup = [ContactStratificationSetup] begin + @testset "Splatted construction" begin + # Two age groups via splatting + age_strat = AgeStratification(:child, :adult) + @test age_strat.stratum_names == [:child, :adult] + @test age_strat.label == :age + @test length(age_strat.stratum_names) == 2 + + # Three age groups via splatting + age_strat_3 = AgeStratification(:child, :adult, :elderly) + @test age_strat_3.stratum_names == [:child, :adult, :elderly] + @test age_strat_3.label == :age + @test length(age_strat_3.stratum_names) == 3 + + # Single age group via splatting + age_strat_1 = AgeStratification(:population) + @test age_strat_1.stratum_names == [:population] + @test age_strat_1.label == :age + @test length(age_strat_1.stratum_names) == 1 + end + + @testset "String to Symbol conversion" begin + # Convenience constructor should convert strings to symbols + age_strat_2_splat = AgeStratification("child", "adult") + @test age_strat_2_splat.stratum_names == age_two.stratum_names + @test age_strat_2_splat.label == :age + + # Mixed string and symbol input + age_strat_mixed = AgeStratification("child", :adult, "elderly") + @test age_strat_mixed.stratum_names == age_multi.stratum_names + @test age_strat_mixed.label == :age + end + + @testset "Validation in convenience constructor" begin + # Duplicate names should still error + @test_throws ErrorException AgeStratification(:child, :adult, :child) + end +end + +@testitem "GeographicStratification convenience constructor" setup = [ContactStratificationSetup] begin + @testset "Splatted construction" begin + # Two geographic regions via splatting + geo_strat = GeographicStratification(:urban, :rural) + @test geo_strat.stratum_names == [:urban, :rural] + @test geo_strat.label == :geography + @test length(geo_strat.stratum_names) == 2 + + # Three geographic regions via splatting + geo_strat_3 = GeographicStratification(:urban, :suburban, :rural) + @test geo_strat_3.stratum_names == [:urban, :suburban, :rural] + @test geo_strat_3.label == :geography + @test length(geo_strat_3.stratum_names) == 3 + end + + @testset "String to Symbol conversion" begin + # Convenience constructor should convert strings to symbols + geo_strat_splat = GeographicStratification("urban", "rural") + @test geo_strat_splat.stratum_names == geo_two.stratum_names + @test geo_strat_splat.label == :geography + end +end + +@testitem "ContactStratification type relationships" setup = [ContactStratificationSetup] begin + @testset "All convenience constructors return ContactStratification" begin + @test age_single isa ContactStratification + @test age_two isa ContactStratification + @test geo_two isa ContactStratification + @test risk_strat isa ContactStratification + end + + @testset "All stratifications are Stratification subtypes" begin + @test age_single isa Stratification + @test geo_two isa Stratification + @test risk_strat isa Stratification + end + + @testset "Labels distinguish stratification types" begin + @test age_two.label == :age + @test geo_two.label == :geography + @test risk_strat.label == :risk + + # Labels can be arbitrary + custom = ContactStratification([:a, :b], :custom_label) + @test custom.label == :custom_label + end +end diff --git a/AlgebraicEpiMech/test/test_create_age_stratification.jl b/AlgebraicEpiMech/test/test_create_age_stratification.jl new file mode 100644 index 0000000..69aa461 --- /dev/null +++ b/AlgebraicEpiMech/test/test_create_age_stratification.jl @@ -0,0 +1,160 @@ +# Unit tests for create_age_stratification.jl +# Tests the create_model function for age stratification typed Petri nets + +@testsnippet AgeStratificationModelSetup begin + using AlgebraicPetri + using AlgebraicPetri.TypedPetri + using Catlab + using AlgebraicEpiMech + + # Create typing instances for testing + one_pop_typing = OnePopulationTyping(population_type = :Individual) + ui_typing = UninfectedInfectedTyping( + uninfected_type = :Susceptible, + infected_type = :Infectious + ) +end + +# Test that create_model returns typed Petri net +@testitem "create_model(AgeStratification) returns ACSetTransformation" setup = [AgeStratificationModelSetup] begin + age_strat = AgeStratification([:child, :adult]) + typed_model = create_model(one_pop_typing, age_strat) + + # Should return an ACSetTransformation (typed Petri net) + @test typed_model isa ACSetTransformation +end + +@testitem "create_model(AgeStratification) typed model has dom as LabelledPetriNet" setup = [AgeStratificationModelSetup] begin + age_strat = AgeStratification([:child, :adult]) + typed_model = create_model(one_pop_typing, age_strat) + + # The domain should be a LabelledPetriNet (usable for ODEs) + pn = dom(typed_model) + @test pn isa LabelledPetriNet +end + +@testitem "create_model(AgeStratification) type system accessible via codom" setup = [AgeStratificationModelSetup] begin + age_strat = AgeStratification([:child, :adult]) + typed_model = create_model(one_pop_typing, age_strat) + + # The type system should be accessible via codom + schema_pn = codom(typed_model) + @test schema_pn isa LabelledPetriNet +end + +# Test with OnePopulationTyping +@testitem "create_model(AgeStratification) with 2 age groups OnePopulationTyping" setup = [AgeStratificationModelSetup] begin + age_strat = AgeStratification([:child, :adult]) + typed_model = create_model(one_pop_typing, age_strat) + + pn = dom(typed_model) + + # Should have 2 species (one per age group) + @test length(AlgebraicPetri.snames(pn)) == 2 + + # Should have n² + 3n = 10 transitions (4 transmission + 6 reflexive) + @test length(AlgebraicPetri.tnames(pn)) == 10 + + # Check transmission names represent contact patterns + tnames = AlgebraicPetri.tnames(pn) + @test :child_child in tnames # within-group child + @test :child_adult in tnames # child infected by adult + @test :adult_child in tnames # adult infected by child + @test :adult_adult in tnames # within-group adult +end + +@testitem "create_model(AgeStratification) with 3 age groups OnePopulationTyping" setup = [AgeStratificationModelSetup] begin + age_strat = AgeStratification([:child, :adult, :elderly]) + typed_model = create_model(one_pop_typing, age_strat) + + pn = dom(typed_model) + + # Should have 3 species + @test length(AlgebraicPetri.snames(pn)) == 3 + + # Should have n² + 3n = 18 transitions + @test length(AlgebraicPetri.tnames(pn)) == 18 +end + +# Test with UninfectedInfectedTyping +@testitem "create_model(AgeStratification) with UninfectedInfectedTyping" setup = [AgeStratificationModelSetup] begin + age_strat = AgeStratification([:child, :adult]) + typed_model = create_model(ui_typing, age_strat) + + pn = dom(typed_model) + + # Should have 4 species (uninfected + infected for each age group) + @test length(AlgebraicPetri.snames(pn)) == 4 + + # Should have n² + 3n = 10 transitions + @test length(AlgebraicPetri.tnames(pn)) == 10 + + # Check species names include _U and _I suffixes + snames = AlgebraicPetri.snames(pn) + @test :child_U in snames + @test :child_I in snames + @test :adult_U in snames + @test :adult_I in snames +end + +# Test generate_transition_names +@testitem "generate_transition_names creates correct names for age groups" setup = [AgeStratificationModelSetup] begin + age_strat = AgeStratification([:young, :old]) + uwd = create_model_uwd(one_pop_typing, age_strat) + + names = generate_transition_names(uwd, age_strat) + + # Should have n² + 3n = 10 names + @test length(names) == 10 + + # Check transmission patterns exist + @test :young_young in names + @test :young_old in names + @test :old_young in names + @test :old_old in names + # Check reflexive stratum names exist + @test :young in names + @test :old in names +end + +# Test single age group edge case +@testitem "create_model(AgeStratification) with single age group" setup = [AgeStratificationModelSetup] begin + age_strat = AgeStratification([:population]) + typed_model = create_model(one_pop_typing, age_strat) + + pn = dom(typed_model) + + # Should have 1 species + @test length(AlgebraicPetri.snames(pn)) == 1 + + # Should have n² + 3n = 4 transitions + @test length(AlgebraicPetri.tnames(pn)) == 4 + + @test :population_population in AlgebraicPetri.tnames(pn) +end + +# Test composition with compartmental models +@testitem "typed_product of SIR and AgeStratification creates age-structured model" setup = [AgeStratificationModelSetup] begin + # Create compartmental model + sir = create_model(one_pop_typing, SIR()) + + # Create age stratification + age_strat = AgeStratification([:child, :adult]) + age_model = create_model(one_pop_typing, age_strat) + + # Compose them + age_sir = typed_product(sir, age_model) + + # Should be a typed model + @test age_sir isa ACSetTransformation + + pn = dom(age_sir) + + # Should have 6 species (S, I, R for each of 2 age groups) + @test length(AlgebraicPetri.snames(pn)) == 6 + + # Species names will be tuples from composition - just verify we have the right count + # The actual naming format depends on typed_product implementation + snames = AlgebraicPetri.snames(pn) + @test length(snames) == 6 +end diff --git a/AlgebraicEpiMech/test/test_create_basic_model.jl b/AlgebraicEpiMech/test/test_create_basic_model.jl new file mode 100644 index 0000000..98d9878 --- /dev/null +++ b/AlgebraicEpiMech/test/test_create_basic_model.jl @@ -0,0 +1,216 @@ +# Unit tests for create_basic_model.jl +# Tests the create_model function that produces typed Petri nets + +@testsnippet BasicModelSetup begin + using AlgebraicPetri + using AlgebraicPetri.TypedPetri + using Catlab + using AlgebraicEpiMech + using LabelledArrays + using SeeToDee: Rk4 + + # Create typing instances for testing + one_pop_typing = OnePopulationTyping(population_type = :Individual) + ui_typing = UninfectedInfectedTyping( + uninfected_type = :Susceptible, + infected_type = :Infectious + ) +end + +# Test that create_model returns typed Petri net +@testitem "create_model returns ACSetTransformation" setup = [BasicModelSetup] begin + model = SEIR() + typed_model = create_model(one_pop_typing, model) + + # Should return an ACSetTransformation (typed Petri net) + @test typed_model isa ACSetTransformation +end + +@testitem "create_model typed model has dom as LabelledPetriNet" setup = [BasicModelSetup] begin + model = SIR() + typed_model = create_model(one_pop_typing, model) + + # The domain should be a LabelledPetriNet (usable for ODEs) + pn = dom(typed_model) + @test pn isa LabelledPetriNet +end + +@testitem "create_model type system accessible via codom" setup = [BasicModelSetup] begin + model = SEIR() + typed_model = create_model(one_pop_typing, model) + + # The type system should be accessible via codom + schema_pn = codom(typed_model) + @test schema_pn isa LabelledPetriNet +end + +# Test SI model +@testitem "create_model(SI) creates valid typed model" setup = [BasicModelSetup] begin + model = SI() + typed_model = create_model(one_pop_typing, model) + + pn = dom(typed_model) + @test length(AlgebraicPetri.snames(pn)) == 2 # S, I + @test length(AlgebraicPetri.tnames(pn)) == 1 # transmission +end + +@testitem "create_model(SI) with multi-stage I" setup = [BasicModelSetup] begin + model = SI(number_I_stages = 3) + typed_model = create_model(one_pop_typing, model) + + pn = dom(typed_model) + @test length(AlgebraicPetri.snames(pn)) == 4 # S, I1, I2, I3 +end + +# Test SEI model +@testitem "create_model(SEI) creates valid typed model" setup = [BasicModelSetup] begin + model = SEI() + typed_model = create_model(one_pop_typing, model) + + pn = dom(typed_model) + @test length(AlgebraicPetri.snames(pn)) == 3 # S, E, I + @test length(AlgebraicPetri.tnames(pn)) == 2 # transmission, density +end + +@testitem "create_model(SEI) with multi-stage E and I" setup = [BasicModelSetup] begin + model = SEI(number_E_stages = 2, number_I_stages = 3) + typed_model = create_model(one_pop_typing, model) + + pn = dom(typed_model) + @test length(AlgebraicPetri.snames(pn)) == 6 # S, E1, E2, I1, I2, I3 +end + +# Test SIR model +@testitem "create_model(SIR) creates valid typed model" setup = [BasicModelSetup] begin + model = SIR() + typed_model = create_model(one_pop_typing, model) + + pn = dom(typed_model) + @test length(AlgebraicPetri.snames(pn)) == 3 # S, I, R + @test length(AlgebraicPetri.tnames(pn)) == 2 # transmission, density +end + +# Test SEIR model +@testitem "create_model(SEIR) creates valid typed model" setup = [BasicModelSetup] begin + model = SEIR() + typed_model = create_model(one_pop_typing, model) + + pn = dom(typed_model) + @test length(AlgebraicPetri.snames(pn)) == 4 # S, E, I, R + @test length(AlgebraicPetri.tnames(pn)) == 3 # transmission, density, density +end + +@testitem "create_model(SEIR) with multi-stage E and I" setup = [BasicModelSetup] begin + model = SEIR(number_E_stages = 2, number_I_stages = 3) + typed_model = create_model(one_pop_typing, model) + + pn = dom(typed_model) + @test length(AlgebraicPetri.snames(pn)) == 7 # S, E1, E2, I1, I2, I3, R +end + +# Test SIS model +@testitem "create_model(SIS) creates valid typed model" setup = [BasicModelSetup] begin + model = SIS() + typed_model = create_model(one_pop_typing, model) + + pn = dom(typed_model) + @test length(AlgebraicPetri.snames(pn)) == 2 # S, I + @test length(AlgebraicPetri.tnames(pn)) == 2 # transmission, reversion +end + +# Test SEIS model +@testitem "create_model(SEIS) creates valid typed model" setup = [BasicModelSetup] begin + model = SEIS() + typed_model = create_model(one_pop_typing, model) + + pn = dom(typed_model) + @test length(AlgebraicPetri.snames(pn)) == 3 # S, E, I + @test length(AlgebraicPetri.tnames(pn)) == 3 # transmission, density, reversion +end + +# Test SEIRS model +@testitem "create_model(SEIRS) creates valid typed model" setup = [BasicModelSetup] begin + model = SEIRS() + typed_model = create_model(one_pop_typing, model) + + pn = dom(typed_model) + @test length(AlgebraicPetri.snames(pn)) == 4 # S, E, I, R + @test length(AlgebraicPetri.tnames(pn)) == 4 # transmission, density, density, reversion +end + +# Test with UninfectedInfectedTyping +@testitem "create_model works with UninfectedInfectedTyping" setup = [BasicModelSetup] begin + model = SEIR() + typed_model = create_model(ui_typing, model) + + pn = dom(typed_model) + @test length(AlgebraicPetri.snames(pn)) == 4 # S, E, I, R + @test typed_model isa ACSetTransformation +end + +# Test that model can be used for ODE simulation +@testitem "create_model output can be used for ODE simulation" setup = [BasicModelSetup] begin + using LabelledArrays + + model = SIR() + typed_model = create_model(one_pop_typing, model) + pn = dom(typed_model) + + # Create vectorfield, wrapped into the out-of-place `(x, u, p, t) -> dx` form SeeToDee + # discretises (`u` is the control input, unused here) + f! = vectorfield_flat(pn) + f(x, u, p, t) = f!(similar(x), x, p, t) + + # Set up initial conditions and parameters using auto-generated transition names + u0 = LVector(S = 990.0, I = 10.0, R = 0.0) + p = LVector(transmission_S_I = 0.0005, I_to_R = 0.25) + + # Discretise to daily RK4 steps and integrate over 40 days + step = Rk4(f, 1.0; supersample = 4) + x = foldl((x, t) -> step(x, nothing, p, t), 0.0:1.0:39.0; init = u0) + + # Verify solution is valid + @test all(isfinite, x) + @test sum(x) ≈ sum(u0) # SIR conserves population + @test x[:R] > 0.0 # Some recovery should occur +end + +# Test model composition capability +@testitem "create_model output can be composed with typed_product" setup = [BasicModelSetup] begin + # Create two simple models + si_model = create_model(one_pop_typing, SI()) + sir_model = create_model(one_pop_typing, SIR()) + + # Should be able to compose them (this tests the typing is compatible) + # typed_product requires both to be ACSetTransformations + @test si_model isa ACSetTransformation + @test sir_model isa ACSetTransformation + + # Verify codomain (type system) is the same type + @test codom(si_model) isa LabelledPetriNet + @test codom(sir_model) isa LabelledPetriNet + + # Verify both codomains are equivalent (same type-system structure) + @test codom(si_model) == codom(sir_model) +end + +# Test that codom matches direct type_system call +@testitem "create_model codom matches type_system" setup = [BasicModelSetup] begin + model = SEIR() + + # Create typed model + typed_model = create_model(one_pop_typing, model) + + # Get type system via codom + codomain_from_typed_model = codom(typed_model) + + # Create the type system directly + direct_type_system = type_system(one_pop_typing) + + # Both should be LabelledPetriNets + @test codomain_from_typed_model isa LabelledPetriNet + @test direct_type_system isa LabelledPetriNet + + # They should be equal in structure + @test codomain_from_typed_model == direct_type_system +end diff --git a/AlgebraicEpiMech/test/test_create_contact_stratification.jl b/AlgebraicEpiMech/test/test_create_contact_stratification.jl new file mode 100644 index 0000000..d40dbaa --- /dev/null +++ b/AlgebraicEpiMech/test/test_create_contact_stratification.jl @@ -0,0 +1,259 @@ +# Unit tests for create_model with ContactStratification and its convenience constructors +# This generalizes create_age_stratification.jl to work with all stratification types + +@testsnippet ContactStratificationModelSetup begin + using AlgebraicPetri + using AlgebraicPetri.TypedPetri + using Catlab + using AlgebraicEpiMech + + # Create typing instances for testing + one_pop_typing = OnePopulationTyping(population_type = :Individual) + ui_typing = UninfectedInfectedTyping( + uninfected_type = :Susceptible, + infected_type = :Infectious + ) +end + +# Test that create_model returns typed Petri net for ContactStratification +@testitem "create_model(ContactStratification) returns ACSetTransformation" setup = [ContactStratificationModelSetup] begin + @testset "AgeStratification" begin + age_strat = AgeStratification([:child, :adult]) + typed_model = create_model(one_pop_typing, age_strat) + @test typed_model isa ACSetTransformation + end + + @testset "GeographicStratification" begin + geo_strat = GeographicStratification([:urban, :rural]) + typed_model = create_model(one_pop_typing, geo_strat) + @test typed_model isa ACSetTransformation + end + + @testset "Direct ContactStratification" begin + risk_strat = ContactStratification([:low, :high], :risk) + typed_model = create_model(one_pop_typing, risk_strat) + @test typed_model isa ACSetTransformation + end +end + +@testitem "create_model(ContactStratification) typed model has dom as LabelledPetriNet" setup = [ContactStratificationModelSetup] begin + strats = [ + AgeStratification([:child, :adult]), + GeographicStratification([:urban, :rural]), + ContactStratification([:low, :high], :risk), + ] + + for strat in strats + typed_model = create_model(one_pop_typing, strat) + pn = dom(typed_model) + @test pn isa LabelledPetriNet + end +end + +@testitem "create_model(ContactStratification) type system accessible via codom" setup = [ContactStratificationModelSetup] begin + strats = [ + AgeStratification([:child, :adult]), + GeographicStratification([:urban, :rural]), + ContactStratification([:low, :high], :risk), + ] + + for strat in strats + typed_model = create_model(one_pop_typing, strat) + schema_pn = codom(typed_model) + @test schema_pn isa LabelledPetriNet + end +end + +# Test with OnePopulationTyping - generalized for any ContactStratification +@testitem "create_model with 2 strata OnePopulationTyping" setup = [ContactStratificationModelSetup] begin + @testset "AgeStratification([:child, :adult])" begin + strat = AgeStratification([:child, :adult]) + typed_model = create_model(one_pop_typing, strat) + pn = dom(typed_model) + + # Should have 2 species (one per stratum) + @test length(AlgebraicPetri.snames(pn)) == 2 + + # Should have n² + 3n = 10 transitions (4 transmission + 6 reflexive) + @test length(AlgebraicPetri.tnames(pn)) == 10 + + # Check transmission names represent contact patterns + tnames = AlgebraicPetri.tnames(pn) + @test :child_child in tnames + @test :child_adult in tnames + @test :adult_child in tnames + @test :adult_adult in tnames + end + + @testset "GeographicStratification([:urban, :rural])" begin + strat = GeographicStratification([:urban, :rural]) + typed_model = create_model(one_pop_typing, strat) + pn = dom(typed_model) + + @test length(AlgebraicPetri.snames(pn)) == 2 + @test length(AlgebraicPetri.tnames(pn)) == 10 + + tnames = AlgebraicPetri.tnames(pn) + @test :urban_urban in tnames + @test :urban_rural in tnames + @test :rural_urban in tnames + @test :rural_rural in tnames + end + + @testset "ContactStratification([:low, :high], :risk)" begin + strat = ContactStratification([:low, :high], :risk) + typed_model = create_model(one_pop_typing, strat) + pn = dom(typed_model) + + @test length(AlgebraicPetri.snames(pn)) == 2 + @test length(AlgebraicPetri.tnames(pn)) == 10 + + tnames = AlgebraicPetri.tnames(pn) + @test :low_low in tnames + @test :low_high in tnames + @test :high_low in tnames + @test :high_high in tnames + end +end + +@testitem "create_model with 3 strata OnePopulationTyping" setup = [ContactStratificationModelSetup] begin + strats_and_names = [ + (AgeStratification([:child, :adult, :elderly]), [:child, :adult, :elderly]), + ( + GeographicStratification([:urban, :suburban, :rural]), + [:urban, :suburban, :rural], + ), + (ContactStratification([:low, :medium, :high], :risk), [:low, :medium, :high]), + ] + + for (strat, expected_names) in strats_and_names + typed_model = create_model(one_pop_typing, strat) + pn = dom(typed_model) + + # Should have 3 species + @test length(AlgebraicPetri.snames(pn)) == 3 + + # Should have n² + 3n = 18 transitions + @test length(AlgebraicPetri.tnames(pn)) == 18 + + # Check some expected transition names + tnames = AlgebraicPetri.tnames(pn) + for name in expected_names + # Within-stratum transmission + within = Symbol(string(name), "_", string(name)) + @test within in tnames + end + end +end + +# Test with UninfectedInfectedTyping +@testitem "create_model(ContactStratification) with UninfectedInfectedTyping" setup = [ContactStratificationModelSetup] begin + @testset "AgeStratification" begin + age_strat = AgeStratification([:child, :adult]) + typed_model = create_model(ui_typing, age_strat) + pn = dom(typed_model) + + # Should have 4 species (uninfected + infected for each age group) + @test length(AlgebraicPetri.snames(pn)) == 4 + + # Should have n² + 3n = 10 transitions + @test length(AlgebraicPetri.tnames(pn)) == 10 + + # Check species names include _U and _I suffixes + snames = AlgebraicPetri.snames(pn) + @test :child_U in snames + @test :child_I in snames + @test :adult_U in snames + @test :adult_I in snames + end + + @testset "GeographicStratification" begin + geo_strat = GeographicStratification([:urban, :rural]) + typed_model = create_model(ui_typing, geo_strat) + pn = dom(typed_model) + + @test length(AlgebraicPetri.snames(pn)) == 4 + @test length(AlgebraicPetri.tnames(pn)) == 10 + + snames = AlgebraicPetri.snames(pn) + @test :urban_U in snames + @test :urban_I in snames + @test :rural_U in snames + @test :rural_I in snames + end +end + +# Test generate_transition_names +@testitem "generate_transition_names creates correct names" setup = [ContactStratificationModelSetup] begin + @testset "AgeStratification" begin + strat = AgeStratification([:young, :old]) + uwd = create_model_uwd(one_pop_typing, strat) + names = generate_transition_names(uwd, strat) + + @test length(names) == 10 # 4 transmission + 8 reflexive + # Transmission names + @test :young_young in names + @test :young_old in names + @test :old_young in names + @test :old_old in names + # Reflexive boxes use stratum name + @test :young in names + @test :old in names + end + + @testset "GeographicStratification" begin + strat = GeographicStratification([:north, :south]) + uwd = create_model_uwd(one_pop_typing, strat) + names = generate_transition_names(uwd, strat) + + @test length(names) == 10 + @test :north_north in names + @test :north_south in names + @test :south_north in names + @test :south_south in names + @test :north in names + @test :south in names + end +end + +# Test single stratum edge case +@testitem "create_model with single stratum" setup = [ContactStratificationModelSetup] begin + strats = [ + AgeStratification([:population]), + GeographicStratification([:region]), + ContactStratification([:group], :custom), + ] + + for strat in strats + typed_model = create_model(one_pop_typing, strat) + pn = dom(typed_model) + + # Should have 1 species + @test length(AlgebraicPetri.snames(pn)) == 1 + + # Should have n² + 3n = 4 transitions + @test length(AlgebraicPetri.tnames(pn)) == 4 + end +end + +# Test composition with compartmental models +@testitem "typed_product of SIR and ContactStratification" setup = [ContactStratificationModelSetup] begin + sir = create_model(one_pop_typing, SIR()) + + strats = [ + AgeStratification([:child, :adult]), + GeographicStratification([:urban, :rural]), + ContactStratification([:low, :high], :risk), + ] + + for strat in strats + strat_model = create_model(one_pop_typing, strat) + composed = typed_product(sir, strat_model) + + @test composed isa ACSetTransformation + pn = dom(composed) + + # Should have 6 species (S, I, R for each of 2 strata) + @test length(AlgebraicPetri.snames(pn)) == 6 + end +end diff --git a/AlgebraicEpiMech/test/test_create_model_uwd.jl b/AlgebraicEpiMech/test/test_create_model_uwd.jl new file mode 100644 index 0000000..c58fb87 --- /dev/null +++ b/AlgebraicEpiMech/test/test_create_model_uwd.jl @@ -0,0 +1,440 @@ +# Unit tests for create_model_uwd.jl +# Tests the exported create_model_uwd methods for all compartmental model types + +@testsnippet ModelUWDSetup begin + using AlgebraicPetri + using Catlab + using AlgebraicEpiMech + + # Create typing instances for testing + one_pop_typing = OnePopulationTyping(population_type = :Individual) + ui_typing = UninfectedInfectedTyping( + uninfected_type = :Susceptible, + infected_type = :Infectious + ) + uninfected_infected_typing = UninfectedInfectedTyping( + uninfected_type = :Susceptible, + infected_type = :Infectious + ) +end + +@testitem "create_model_uwd leaves the outer boundary unspecified" setup = [ModelUWDSetup] begin + models_and_typings = [ + (one_pop_typing, SI()), + (one_pop_typing, SEIR()), + (one_pop_typing, NoCrossImmunity(2)), + (uninfected_infected_typing, CompleteCrossImmunity(2)), + (one_pop_typing, AgeStratification([:child, :adult])), + (uninfected_infected_typing, ImmuneHistory([:strain_1, :strain_2])), + ] + + for (typing, model) in models_and_typings + uwd = create_model_uwd(typing, model) + @test isempty(ports(uwd, outer = true)) + end +end + +# Tests for SI model UWD construction +@testitem "create_model_uwd(SI) creates correct structure" setup = [ModelUWDSetup] begin + model = SI() + uwd = create_model_uwd(one_pop_typing, model) + + # Should have S and I junctions (2 total) + @test length(junctions(uwd)) == 2 + + # Should have 1 box (infection) + @test length(boxes(uwd)) == 1 + + # Should return a valid RelationDiagram + @test uwd isa RelationDiagram +end + +@testitem "create_model_uwd(SI) with multi-stage I" setup = [ModelUWDSetup] begin + model = SI(number_I_stages = 3) + uwd = create_model_uwd(one_pop_typing, model) + + # Should have S and I1, I2, I3 junctions (4 total) + @test length(junctions(uwd)) == 4 + + # Should have 5 boxes (each mechanism creates 2 boxes: infection (2), I1->I2 (2), I2->I3 (1 unidirectional)) + @test length(boxes(uwd)) == 5 +end + +@testitem "create_model_uwd(SI) works with UninfectedInfectedTyping" setup = [ModelUWDSetup] begin + model = SI() + uwd = create_model_uwd(ui_typing, model) + + # Should create valid UWD with typing-specific types + @test length(junctions(uwd)) == 2 + @test length(boxes(uwd)) == 1 + @test uwd isa RelationDiagram +end + +# Tests for SEI model UWD construction +@testitem "create_model_uwd(SEI) creates correct structure" setup = [ModelUWDSetup] begin + model = SEI() + uwd = create_model_uwd(one_pop_typing, model) + + # Should have S, E, and I junctions (3 total) + @test length(junctions(uwd)) == 3 + + # Should have 2 boxes (exposure, progression E->I) + @test length(boxes(uwd)) == 2 + + @test uwd isa RelationDiagram +end + +@testitem "create_model_uwd(SEI) with multi-stage E and I" setup = [ModelUWDSetup] begin + model = SEI(number_E_stages = 2, number_I_stages = 3) + uwd = create_model_uwd(one_pop_typing, model) + + # Should have S, E1, E2, I1, I2, I3 junctions (6 total) + @test length(junctions(uwd)) == 6 + + # Boxes: exposure S->E from each Infected group (3), E1->E2 (1), E2->I1 (1), I1->I2 (1), I2->I3 (1), + @test length(boxes(uwd)) == 7 +end + +@testitem "create_model_uwd(SEI) works with UninfectedInfectedTyping" setup = [ModelUWDSetup] begin + model = SEI() + uwd = create_model_uwd(ui_typing, model) + + # Should create valid UWD with typing-specific types + @test length(junctions(uwd)) == 3 + @test length(boxes(uwd)) == 2 + @test uwd isa RelationDiagram +end + +# Tests for SIR model UWD construction +@testitem "create_model_uwd(SIR) creates correct structure" setup = [ModelUWDSetup] begin + model = SIR() + uwd = create_model_uwd(one_pop_typing, model) + + # Should have S, I, and R junctions (3 total) + @test length(junctions(uwd)) == 3 + + # Should have 2 boxes (infection, recovery) + @test length(boxes(uwd)) == 2 + + @test uwd isa RelationDiagram +end + +@testitem "create_model_uwd(SIR) with multi-stage I" setup = [ModelUWDSetup] begin + model = SIR(number_I_stages = 3) + uwd = create_model_uwd(one_pop_typing, model) + + # Should have S, I1, I2, I3, R junctions (5 total) + @test length(junctions(uwd)) == 5 + + # Boxes: infection from each infectious group (3), I1->I2 (1), I2->I3 (1), I3->R (1) + @test length(boxes(uwd)) == 6 +end + +@testitem "create_model_uwd(SIR) works with UninfectedInfectedTyping" setup = [ModelUWDSetup] begin + model = SIR() + uwd = create_model_uwd(ui_typing, model) + + # Should create valid UWD + @test length(junctions(uwd)) == 3 + @test length(boxes(uwd)) == 2 + @test uwd isa RelationDiagram +end + +# Tests for SEIR model UWD construction +@testitem "create_model_uwd(SEIR) creates correct structure" setup = [ModelUWDSetup] begin + model = SEIR() + uwd = create_model_uwd(one_pop_typing, model) + + # Should have S, E, I, and R junctions (4 total) + @test length(junctions(uwd)) == 4 + + # Should have 3 boxes (exposure, E->I progression, recovery) + @test length(boxes(uwd)) == 3 + + @test uwd isa RelationDiagram +end + +@testitem "create_model_uwd(SEIR) with multi-stage E and I" setup = [ModelUWDSetup] begin + model = SEIR(number_E_stages = 2, number_I_stages = 3) + uwd = create_model_uwd(one_pop_typing, model) + + # Should have S, E1, E2, I1, I2, I3, R junctions (7 total) + @test length(junctions(uwd)) == 7 + + # Boxes: exposure (2), E1->E2 (1), E2->I1 (1), I1->I2 (1), I2->I3 (1), I3->R (1), S->E1 + @test length(boxes(uwd)) == 8 +end + +@testitem "create_model_uwd(SEIR) works with UninfectedInfectedTyping" setup = [ModelUWDSetup] begin + model = SEIR() + uwd = create_model_uwd(ui_typing, model) + + # Should create valid UWD + @test length(junctions(uwd)) == 4 + @test length(boxes(uwd)) == 3 + @test uwd isa RelationDiagram +end + +# Tests for SIS model UWD construction +@testitem "create_model_uwd(SIS) creates correct structure" setup = [ModelUWDSetup] begin + model = SIS() + uwd = create_model_uwd(one_pop_typing, model) + + # Should have S and I junctions (same as SI) + @test length(junctions(uwd)) == 2 + + # Should have 2 boxes (infection, reversion back to S) + @test length(boxes(uwd)) == 2 + + @test uwd isa RelationDiagram +end + +@testitem "create_model_uwd(SIS) with multi-stage I" setup = [ModelUWDSetup] begin + model = SIS(number_I_stages = 3) + uwd = create_model_uwd(one_pop_typing, model) + + # Should have S, I1, I2, I3 junctions (4 total) + @test length(junctions(uwd)) == 4 + + # Boxes: infection (2), I1->I2 (1), I2->I3 (1), I3->S (1), S->I1 + @test length(boxes(uwd)) == 6 +end + +@testitem "create_model_uwd(SIS) works with UninfectedInfectedTyping" setup = [ModelUWDSetup] begin + model = SIS() + uwd = create_model_uwd(ui_typing, model) + + # Should create valid UWD + @test length(junctions(uwd)) == 2 + @test length(boxes(uwd)) == 2 + @test uwd isa RelationDiagram +end + +# Tests for SEIS model UWD construction +@testitem "create_model_uwd(SEIS) creates correct structure" setup = [ModelUWDSetup] begin + model = SEIS() + uwd = create_model_uwd(one_pop_typing, model) + + # Should have S, E, and I junctions (same as SEI) + @test length(junctions(uwd)) == 3 + + # Should have 3 boxes (exposure, E->I progression, reversion to S) + @test length(boxes(uwd)) == 3 + + @test uwd isa RelationDiagram +end + +@testitem "create_model_uwd(SEIS) with multi-stage E and I" setup = [ModelUWDSetup] begin + model = SEIS(number_E_stages = 2, number_I_stages = 3) + uwd = create_model_uwd(one_pop_typing, model) + + # Should have S, E1, E2, I1, I2, I3 junctions (6 total) + @test length(junctions(uwd)) == 6 + + # Boxes: exposure (3), E1->E2 (1), E2->I1 (1), I1->I2 (1), I2->I3 (1), I3->S (1) + @test length(boxes(uwd)) == 8 +end + +@testitem "create_model_uwd(SEIS) works with UninfectedInfectedTyping" setup = [ModelUWDSetup] begin + model = SEIS() + uwd = create_model_uwd(ui_typing, model) + + # Should create valid UWD + @test length(junctions(uwd)) == 3 + @test length(boxes(uwd)) == 3 + @test uwd isa RelationDiagram +end + +# Tests for SEIRS model UWD construction +@testitem "create_model_uwd(SEIRS) creates correct structure" setup = [ModelUWDSetup] begin + model = SEIRS() + uwd = create_model_uwd(one_pop_typing, model) + + # Should have S, E, I, and R junctions (4 total) + @test length(junctions(uwd)) == 4 + + # Should have 4 boxes (exposure, E->I progression, recovery, waning immunity R->S) + @test length(boxes(uwd)) == 4 + + @test uwd isa RelationDiagram +end + +@testitem "create_model_uwd(SEIRS) with multi-stage E and I" setup = [ModelUWDSetup] begin + model = SEIRS(number_E_stages = 2, number_I_stages = 3) + uwd = create_model_uwd(one_pop_typing, model) + + # Should have S, E1, E2, I1, I2, I3, R junctions (7 total) + @test length(junctions(uwd)) == 7 + + # Boxes: exposure (3), E1->E2 (1), E2->I1 (1), I1->I2 (1), I2->I3 (1), I3->R (1), R->S (1) + @test length(boxes(uwd)) == 9 +end + +@testitem "create_model_uwd(SEIRS) works with UninfectedInfectedTyping" setup = [ModelUWDSetup] begin + model = SEIRS() + uwd = create_model_uwd(ui_typing, model) + + # Should create valid UWD + @test length(junctions(uwd)) == 4 + @test length(boxes(uwd)) == 4 + @test uwd isa RelationDiagram +end + +# Test error handling for unimplemented combinations +@testitem "create_model_uwd throws error for unimplemented combinations" setup = [ModelUWDSetup] begin + # Create a custom typing type that has no implementations + struct CustomSchema <: EpidemiologicalTyping end + custom_schema = CustomSchema() + + # Should throw error for any model with unimplemented typing + @test_throws ErrorException create_model_uwd(custom_schema, SI()) + @test_throws ErrorException create_model_uwd(custom_schema, SEIR()) +end + +# Tests for NoCrossImmunity multistrain model UWD construction +@testitem "create_model_uwd(NoCrossImmunity) with auto-generated strain names" setup = [ModelUWDSetup] begin + model = NoCrossImmunity(2) + uwd = create_model_uwd(one_pop_typing, model) + + # Should have 2 strain junctions (no Uninfected in multistrain UWD) + @test length(junctions(uwd)) == 2 + + # Should have 2 * 3 = 6 boxes (per strain: transmission, disease, reversion) + @test length(boxes(uwd)) == 6 + + # Check box names - transmission, disease, reversion for each strain + box_names = [subpart(uwd, b, :name) for b in boxes(uwd)] + @test count(==(:transmission), box_names) == 2 + @test count(==(:disease), box_names) == 2 + @test count(==(:reversion), box_names) == 2 +end + +@testitem "create_model_uwd(NoCrossImmunity) with custom strain names" setup = [ModelUWDSetup] begin + model = NoCrossImmunity([:h1n1, :h3n2, :seasonal]) + uwd = create_model_uwd(one_pop_typing, model) + + # Should have 3 strain junctions (no Uninfected in multistrain UWD) + @test length(junctions(uwd)) == 3 + + # Should have 3 * 3 = 9 boxes + @test length(boxes(uwd)) == 9 + + # Check junction variables are the strain names + junction_vars = [subpart(uwd, j, :variable) for j in junctions(uwd)] + @test :h1n1 in junction_vars + @test :h3n2 in junction_vars + @test :seasonal in junction_vars +end + +# Tests for CompleteCrossImmunity multistrain model UWD construction +@testitem "create_model_uwd(CompleteCrossImmunity) with auto-generated strain names" setup = [ModelUWDSetup] begin + model = CompleteCrossImmunity(2) + uwd = create_model_uwd(uninfected_infected_typing, model) + + # Should have 3 junctions: 1 shared susceptible + 2 strain-specific infected + @test length(junctions(uwd)) == 3 + + # Should have 2 * 3 = 6 boxes (per strain: transmission, disease, reversion) + @test length(boxes(uwd)) == 6 + + # Check junction variables: one :susceptible and two strain names + junction_vars = [subpart(uwd, j, :variable) for j in junctions(uwd)] + @test :susceptible in junction_vars + @test :strain_1 in junction_vars + @test :strain_2 in junction_vars + + # Check box names - transmission, disease, reversion for each strain + box_names = [subpart(uwd, b, :name) for b in boxes(uwd)] + @test count(==(:transmission), box_names) == 2 + @test count(==(:disease), box_names) == 2 + @test count(==(:reversion), box_names) == 2 +end + +@testitem "create_model_uwd(CompleteCrossImmunity) with custom strain names" setup = [ModelUWDSetup] begin + model = CompleteCrossImmunity([:wild_type, :variant]) + uwd = create_model_uwd(uninfected_infected_typing, model) + + # Should have 3 junctions: 1 shared susceptible + 2 strain-specific infected + @test length(junctions(uwd)) == 3 + + # Should have 2 * 3 = 6 boxes + @test length(boxes(uwd)) == 6 + + # Check junction variables + junction_vars = [subpart(uwd, j, :variable) for j in junctions(uwd)] + @test :susceptible in junction_vars + @test :wild_type in junction_vars + @test :variant in junction_vars +end + +# Tests for AgeStratification UWD construction +@testitem "create_model_uwd(AgeStratification) creates correct structure with OnePopulationTyping" setup = [ModelUWDSetup] begin + age_strat = AgeStratification([:child, :adult]) + uwd = create_model_uwd(one_pop_typing, age_strat) + + # Should have 2 age group junctions (single junction per age group for OnePopulationTyping) + @test length(junctions(uwd)) == 2 + + # Should have n² + 3n boxes: 4 transmission + 2 disease + 2 reversion + 2 waning = 10 + @test length(boxes(uwd)) == 10 + + # Check transmission boxes are present + box_names = [subpart(uwd, b, :name) for b in boxes(uwd)] + @test count(==(:transmission), box_names) == 4 + @test count(==(:disease), box_names) == 2 + @test count(==(:reversion), box_names) == 2 + @test count(==(:waning), box_names) == 2 + + # Check junction variables are the age group names (each appears once) + junction_vars = [subpart(uwd, j, :variable) for j in junctions(uwd)] + @test :child in junction_vars + @test :adult in junction_vars + + @test uwd isa RelationDiagram +end + +@testitem "create_model_uwd(AgeStratification) creates correct structure with UninfectedInfectedTyping" setup = [ModelUWDSetup] begin + age_strat = AgeStratification([:child, :adult, :elderly]) + uwd = create_model_uwd(ui_typing, age_strat) + + # Should have 6 age group junctions (uninfected + infected for each of 3 age groups) + @test length(junctions(uwd)) == 6 + + # Should have n² + 3n boxes: 9 transmission + 3 disease + 3 reversion + 3 waning = 18 + @test length(boxes(uwd)) == 18 + + # Check box type counts + box_names = [subpart(uwd, b, :name) for b in boxes(uwd)] + @test count(==(:transmission), box_names) == 9 + @test count(==(:disease), box_names) == 3 + @test count(==(:reversion), box_names) == 3 + + # Check junction variables have unique names with _U and _I suffixes + junction_vars = [subpart(uwd, j, :variable) for j in junctions(uwd)] + @test :child_U in junction_vars + @test :child_I in junction_vars + @test :adult_U in junction_vars + @test :adult_I in junction_vars + @test :elderly_U in junction_vars + @test :elderly_I in junction_vars + + @test uwd isa RelationDiagram +end + +@testitem "create_model_uwd(AgeStratification) with single age group" setup = [ModelUWDSetup] begin + age_strat = AgeStratification([:population]) + uwd = create_model_uwd(one_pop_typing, age_strat) + + # Should have 1 junction for single age group with OnePopulationTyping + @test length(junctions(uwd)) == 1 + + # Should have n² + 3n = 1 + 3 = 4 boxes + @test length(boxes(uwd)) == 4 + + # Check box types + box_names = [subpart(uwd, b, :name) for b in boxes(uwd)] + @test count(==(:transmission), box_names) == 1 + @test count(==(:disease), box_names) == 1 + + @test uwd isa RelationDiagram +end diff --git a/AlgebraicEpiMech/test/test_create_multistrain_model.jl b/AlgebraicEpiMech/test/test_create_multistrain_model.jl new file mode 100644 index 0000000..2b9b033 --- /dev/null +++ b/AlgebraicEpiMech/test/test_create_multistrain_model.jl @@ -0,0 +1,213 @@ +# Unit tests for create_multistrain_model.jl +# Tests the create_model function that produces typed Petri nets for multistrain models + +@testsnippet MultistrainSetup begin + using AlgebraicPetri + using AlgebraicPetri.TypedPetri + using Catlab + using AlgebraicEpiMech + + # Create typing instances for testing + one_pop_typing = OnePopulationTyping(population_type = :Individual) + uninfected_infected_typing = UninfectedInfectedTyping( + uninfected_type = :Susceptible, infected_type = :Infectious + ) +end + +@testitem "create_model - basic NoCrossImmunity" setup = [MultistrainSetup] begin + model = NoCrossImmunity([:h1n1, :h3n2]) + typed_model = create_model(one_pop_typing, model) + + # Should return an ACSetTransformation + @test typed_model isa ACSetTransformation + + # Extract the underlying Petri net + pn = dom(typed_model) + @test pn isa LabelledPetriNet + + # Check transition names are the strain names + transition_names = tnames(pn) + @test :h1n1 in transition_names + @test :h3n2 in transition_names +end + +@testitem "multistrain constructors reject duplicate strain names" begin + using AlgebraicEpiMech + + @test_throws ArgumentError NoCrossImmunity([:a, :a]) + @test_throws ArgumentError CompleteCrossImmunity([:a, :a]) +end + +@testitem "multistrain models require their documented typing" setup = [MultistrainSetup] begin + no_cross = NoCrossImmunity([:a, :b]) + complete = CompleteCrossImmunity([:a, :b]) + + @test applicable(create_model_uwd, one_pop_typing, no_cross) + @test !applicable(create_model_uwd, uninfected_infected_typing, no_cross) + @test applicable(create_model_uwd, uninfected_infected_typing, complete) + @test !applicable(create_model_uwd, one_pop_typing, complete) + + @test_throws MethodError create_model(uninfected_infected_typing, no_cross) + @test_throws MethodError create_model(one_pop_typing, complete) +end + +@testitem "create_model - verify state structure" setup = [MultistrainSetup] begin + model = NoCrossImmunity([:wild, :variant]) + typed_model = create_model(one_pop_typing, model) + + pn = dom(typed_model) + + # Should only have strain states (no Uninfected in pure strain model) + state_names = snames(pn) + @test :wild in state_names + @test :variant in state_names + @test length(state_names) == 2 +end + +@testitem "create_model - auto-generated strain names" setup = [MultistrainSetup] begin + model = NoCrossImmunity(3) + typed_model = create_model(one_pop_typing, model) + + pn = dom(typed_model) + + # Should have auto-generated strain names + state_names = snames(pn) + @test :strain_1 in state_names + @test :strain_2 in state_names + @test :strain_3 in state_names + @test length(state_names) == 3 +end + +@testitem "generate_transition_names - returns strain names" setup = [MultistrainSetup] begin + model = NoCrossImmunity([:flu_a, :flu_b]) + uwd = create_model_uwd(one_pop_typing, model) + + # Generate transition names + names = generate_transition_names(uwd, model) + + # Should return strain names (repeated for each box) + @test all(name in [:flu_a, :flu_b] for name in names) + # Should have 2 strains * 3 boxes each = 6 transition names + @test length(names) == 6 +end + +@testitem "create_model - composition with compartmental model" setup = [MultistrainSetup] begin + # Create both a compartmental and multistrain model + sir_model = create_model(one_pop_typing, SIR()) + strain_model = create_model( + one_pop_typing, NoCrossImmunity( + [ + :alpha, :beta, + ] + ) + ) + + # Compose them + combined = typed_product(sir_model, strain_model) + combined_pn = dom(combined) + + # Should have states for each SIR compartment × strain combination + state_names = snames(combined_pn) + + # All compartments get combined with each strain + @test (:S, :alpha) in state_names + @test (:S, :beta) in state_names + @test (:I, :alpha) in state_names + @test (:I, :beta) in state_names + @test (:R, :alpha) in state_names + @test (:R, :beta) in state_names + + # Total: 3 compartments × 2 strains = 6 states + @test length(state_names) == 6 +end + +# Tests for CompleteCrossImmunity multistrain models + +@testitem "create_model - basic CompleteCrossImmunity" setup = [MultistrainSetup] begin + model = CompleteCrossImmunity([:wild_type, :variant]) + typed_model = create_model(uninfected_infected_typing, model) + + # Should return an ACSetTransformation + @test typed_model isa ACSetTransformation + + # Extract the underlying Petri net + pn = dom(typed_model) + @test pn isa LabelledPetriNet + + # Check transition names are the strain names + transition_names = tnames(pn) + @test :wild_type in transition_names + @test :variant in transition_names +end + +@testitem "create_model - CompleteCrossImmunity state structure" setup = [MultistrainSetup] begin + model = CompleteCrossImmunity([:strain_a, :strain_b]) + typed_model = create_model(uninfected_infected_typing, model) + + pn = dom(typed_model) + + # Should have 1 shared susceptible + 2 strain-specific infected states + state_names = snames(pn) + @test :susceptible in state_names + @test :strain_a in state_names + @test :strain_b in state_names + @test length(state_names) == 3 +end + +@testitem "create_model - CompleteCrossImmunity auto-generated names" setup = [MultistrainSetup] begin + model = CompleteCrossImmunity(3) + typed_model = create_model(uninfected_infected_typing, model) + + pn = dom(typed_model) + + # Should have auto-generated strain names + state_names = snames(pn) + @test :susceptible in state_names + @test :strain_1 in state_names + @test :strain_2 in state_names + @test :strain_3 in state_names + @test length(state_names) == 4 # 1 susceptible + 3 strains +end + +@testitem "generate_transition_names - CompleteCrossImmunity" setup = [MultistrainSetup] begin + model = CompleteCrossImmunity([:h1n1, :h3n2]) + uwd = create_model_uwd(uninfected_infected_typing, model) + + # Generate transition names + names = generate_transition_names(uwd, model) + + # Should return strain names (repeated for each box), excluding :susceptible + @test all(name in [:h1n1, :h3n2] for name in names) + # Should have 2 strains * 3 boxes each = 6 transition names + @test length(names) == 6 +end + +@testitem "create_model - CompleteCrossImmunity composition with SIR" setup = [MultistrainSetup] begin + # Create both a compartmental and multistrain model + sir_model = create_model(uninfected_infected_typing, SIR()) + strain_model = create_model( + uninfected_infected_typing, + CompleteCrossImmunity([:alpha, :beta]) + ) + + # Compose them + combined = typed_product(sir_model, strain_model) + combined_pn = dom(combined) + + # For CompleteCrossImmunity with SIR: + # - S is shared (just one S) + # - I and R are strain-specific (I_alpha, I_beta, R_alpha, R_beta) + state_names = snames(combined_pn) + + # Shared susceptible + @test (:S, :susceptible) in state_names + + # Strain-specific infected and recovered + @test (:I, :alpha) in state_names + @test (:I, :beta) in state_names + @test (:R, :alpha) in state_names + @test (:R, :beta) in state_names + + # Total: 1 shared S + 2*2 strain-specific (I,R) = 5 states + @test length(state_names) == 5 +end diff --git a/AlgebraicEpiMech/test/test_epi_schema.jl b/AlgebraicEpiMech/test/test_epi_schema.jl new file mode 100644 index 0000000..a64746a --- /dev/null +++ b/AlgebraicEpiMech/test/test_epi_schema.jl @@ -0,0 +1,80 @@ +@testitem "epidemiological typing constructors" begin + using AlgebraicEpiMech + + one_population = OnePopulationTyping() + @test one_population isa EpidemiologicalTyping + @test one_population.population_type == :Population + + custom_population = OnePopulationTyping(population_type = :Individual) + @test custom_population.population_type == :Individual + + uninfected_infected = UninfectedInfectedTyping() + @test uninfected_infected isa EpidemiologicalTyping + @test uninfected_infected.uninfected_type == :Uninfected + @test uninfected_infected.infected_type == :Infected + + custom_uninfected_infected = UninfectedInfectedTyping( + uninfected_type = :Susceptible, + infected_type = :Infectious, + ) + @test custom_uninfected_infected.uninfected_type == :Susceptible + @test custom_uninfected_infected.infected_type == :Infectious +end + +@testitem "type_system materializes OnePopulationTyping" begin + using AlgebraicPetri + using AlgebraicPetri.TypedPetri + using AlgebraicEpiMech + + typing = OnePopulationTyping(population_type = :Individual) + codomain = type_system(typing) + + @test codomain isa LabelledPetriNet + @test snames(codomain) == [:Individual] + @test all(in(tnames(codomain)), [:transmission, :disease, :reversion, :waning]) + + system_with_birth = type_system( + typing, + :birth => (:Individual => (:Individual, :Individual)), + ) + @test :birth in tnames(system_with_birth) +end + +@testitem "type_system materializes UninfectedInfectedTyping" begin + using AlgebraicPetri + using AlgebraicPetri.TypedPetri + using AlgebraicEpiMech + + typing = UninfectedInfectedTyping( + uninfected_type = :Susceptible, + infected_type = :Infectious, + ) + codomain = type_system(typing) + + @test codomain isa LabelledPetriNet + @test snames(codomain) == [:Susceptible, :Infectious] + @test all(in(tnames(codomain)), [:transmission, :disease, :reversion, :waning]) + + system_with_recovery = type_system( + typing, + :recovery => (:Infectious => :Susceptible), + ) + @test :recovery in tnames(system_with_recovery) +end + +@testitem "type_system rejects unimplemented typings" begin + using AlgebraicEpiMech + + struct MockTyping <: EpidemiologicalTyping end + + error = try + type_system(MockTyping()) + nothing + catch exception + exception + end + + @test error isa ErrorException + @test occursin("MockTyping", string(error)) + @test occursin("not implemented", string(error)) +end diff --git a/AlgebraicEpiMech/test/test_immune_history.jl b/AlgebraicEpiMech/test/test_immune_history.jl new file mode 100644 index 0000000..ec2f926 --- /dev/null +++ b/AlgebraicEpiMech/test/test_immune_history.jl @@ -0,0 +1,150 @@ +@testsnippet ImmuneHistorySetup begin + using AlgebraicPetri + using AlgebraicPetri.TypedPetri: typed_product + using Catlab + using LabelledArrays + using AlgebraicEpiMech + + typing = UninfectedInfectedTyping() + + # The immune-history factor composed onto a generic SEIRS base. + composed(m) = dom(typed_product(create_model(typing, SEIRS()), create_model(typing, m))) + + fnames(xs) = Set(Symbol.(AlgebraicEpiMech.flatten_symbols.(xs))) + transmissions(pn) = + filter( + t -> startswith(string(t), "transmission_"), + string.(AlgebraicEpiMech.flatten_symbols.(tnames(pn))) + ) +end + +# ---------------------------------------------------------------------------- +# Immune-class logic (internal helpers) +# ---------------------------------------------------------------------------- + +@testitem "immune classes: FullHistory = 2^n, LatestInfection = n+1" setup = [ImmuneHistorySetup] begin + strains = [:current, :invader] + @test length(AlgebraicEpiMech._immune_classes(FullHistory(), strains)) == 4 + @test length(AlgebraicEpiMech._immune_classes(LatestInfection(), strains)) == 3 + @test length(AlgebraicEpiMech._immune_classes(FullHistory(), [:a, :b, :c])) == 8 + @test length(AlgebraicEpiMech._immune_classes(LatestInfection(), [:a, :b, :c])) == 4 +end + +@testitem "susceptibility + reversion target" setup = [ImmuneHistorySetup] begin + strains = [:current, :invader] + @test AlgebraicEpiMech._susceptible_strains((:current,), strains) == [:invader] # escape + @test AlgebraicEpiMech._susceptible_strains((), strains) == [:current, :invader] + @test AlgebraicEpiMech._recover_to(FullHistory(), (:current,), :invader) == (:current, :invader) + @test AlgebraicEpiMech._recover_to(LatestInfection(), (:current,), :invader) == (:invader,) +end + +# ---------------------------------------------------------------------------- +# Composed-model structure — S indexed by {h}, E/I/R by {(h,i)} +# ---------------------------------------------------------------------------- + +@testitem "FullHistory composed with SEIRS (2 strains)" setup = [ImmuneHistorySetup] begin + pn = composed(ImmuneHistory([:current, :invader])) + sp = fnames(snames(pn)) + @test length(sp) == 16 + # S indexed by history h + @test Set([:S_U_naive, :S_U_current, :S_U_invader, :S_U_current_invader]) ⊆ sp + # E/I/R indexed by (h,i) + @test Set([:E_invader_from_current, :I_invader_from_current, :R_invader_from_current]) ⊆ sp + @test length(transmissions(pn)) == 8 +end + +@testitem "LatestInfection composed with SEIRS (2 strains)" setup = [ImmuneHistorySetup] begin + pn = composed(ImmuneHistory([:current, :invader]; mode = LatestInfection())) + sp = fnames(snames(pn)) + @test length(sp) == 15 + @test Set([:S_U_naive, :S_U_current, :S_U_invader]) ⊆ sp + @test !(:S_U_current_invader in sp) # no multi-immunity class under overwrite +end + +@testitem "factor composes with any disease base (SEIR, multi-stage)" setup = [ImmuneHistorySetup] begin + history = create_model(typing, ImmuneHistory([:current, :invader]; mode = LatestInfection())) + # Same factor, different bases — the point of the compositional design. + for base in (SEIR(), SEIRS(number_I_stages = 2)) + pn = dom(typed_product(create_model(typing, base), history)) + @test ns(pn) > 0 && nt(pn) > 0 + @test any(t -> occursin("infect_invader_current", t), transmissions(pn)) # escape survives + end +end + +# ---------------------------------------------------------------------------- +# The escape property — falls out of the pullback +# ---------------------------------------------------------------------------- + +@testitem "escape: invader infects the incumbent-immune S, incumbent does not" setup = [ImmuneHistorySetup] begin + for mode in (FullHistory(), LatestInfection()) + tx = transmissions(composed(ImmuneHistory([:current, :invader]; mode = mode))) + @test any(t -> occursin("infect_invader_current", t), tx) # invader → current-immune S + @test any(t -> occursin("infect_current_naive", t), tx) # naive infectable by both + @test any(t -> occursin("infect_invader_naive", t), tx) + @test !any(t -> occursin("infect_current_current", t), tx) # never re-infected by own strain + @test !any(t -> occursin("infect_invader_invader", t), tx) + end +end + +# ---------------------------------------------------------------------------- +# Guards + validation +# ---------------------------------------------------------------------------- + +@testitem "OnePopulationTyping is rejected" setup = [ImmuneHistorySetup] begin + @test_throws ErrorException create_model_uwd(OnePopulationTyping(), ImmuneHistory([:a, :b])) +end + +@testitem "constructor validation" setup = [ImmuneHistorySetup] begin + @test_throws ArgumentError ImmuneHistory(Symbol[]) + @test_throws ArgumentError ImmuneHistory([:a, :a]) + @test ImmuneHistory([:a, :b]) isa ImmuneHistory{FullHistory} + @test ImmuneHistory([:a, :b]; mode = LatestInfection()) isa ImmuneHistory{LatestInfection} +end + +# ---------------------------------------------------------------------------- +# Behavioral: escape flux, and the history-incrementing reversion +# ---------------------------------------------------------------------------- + +@testitem "escape flux: invader infects the current-immune pool" setup = [ImmuneHistorySetup] begin + pn = composed(ImmuneHistory([:current, :invader]; mode = LatestInfection())) + states = [flatten_symbols(s) for s in snames(pn)] + transitions = [flatten_symbols(t) for t in tnames(pn)] + @test :transmission_S_I_infect_invader_current_by_naive in transitions + + u = LVector(; + ( + s => (s === :S_U_current ? 10.0 : s === :I_invader_from_naive ? 5.0 : 0.0) + for s in states + )... + ) + du = LVector(; (s => 0.0 for s in states)...) + p = LVector(; + ( + t => (t === :transmission_S_I_infect_invader_current_by_naive ? 1.0 : 0.0) + for t in transitions + )... + ) + vectorfield_flat(pn)(du, u, p, 0.0) + + @test du[:E_invader_from_current] > 0.0 # invader infects the current-immune pool + @test du[:S_U_current] < 0.0 +end + +@testitem "reversion increments immune history" setup = [ImmuneHistorySetup] begin + # Recovering from the invader while immune to `current` should land in the class the + # mode dictates: `current_invader` (Full) or `invader` (Latest). + for (mode, target) in ((FullHistory(), :S_U_current_invader), (LatestInfection(), :S_U_invader)) + pn = composed(ImmuneHistory([:current, :invader]; mode = mode)) + states = [flatten_symbols(s) for s in snames(pn)] + transitions = [flatten_symbols(t) for t in tnames(pn)] + @test :R_to_S_invader_from_current in transitions + + u = LVector(; (s => (s === :R_invader_from_current ? 10.0 : 0.0) for s in states)...) + du = LVector(; (s => 0.0 for s in states)...) + p = LVector(; (t => (t === :R_to_S_invader_from_current ? 1.0 : 0.0) for t in transitions)...) + vectorfield_flat(pn)(du, u, p, 0.0) + + @test du[target] > 0.0 # lands in the incremented class + @test du[:R_invader_from_current] < 0.0 + end +end diff --git a/AlgebraicEpiMech/test/test_mechanisms.jl b/AlgebraicEpiMech/test/test_mechanisms.jl new file mode 100644 index 0000000..9700cfd --- /dev/null +++ b/AlgebraicEpiMech/test/test_mechanisms.jl @@ -0,0 +1,195 @@ +@testsnippet MechanismSetup begin + using AlgebraicPetri + using Catlab + using AlgebraicEpiMech + + # Helper function to create a basic UWD for testing mechanisms + function create_test_uwd(typing, num_ports::Int) + port_types = fill( + typing isa OnePopulationTyping ? typing.population_type : typing.infected_type, + num_ports + ) + uwd = RelationDiagram(port_types) + return uwd + end + + # Helper to count boxes in a UWD + count_boxes(uwd) = length(boxes(uwd)) + + # Helper to count junctions in a UWD + count_junctions(uwd) = length(junctions(uwd)) + + # Create test typings + one_pop_typing = OnePopulationTyping() + ui_typing = UninfectedInfectedTyping() +end + +@testitem "Infection mechanism (direct to I)" setup = [MechanismSetup] begin + # OnePopulationTyping: S + I → I + I + @testset "OnePopulationTyping" begin + uwd = create_test_uwd(one_pop_typing, 2) + S_junction = add_junction!(uwd, one_pop_typing.population_type, variable = :S) + I_junction = add_junction!(uwd, one_pop_typing.population_type, variable = :I) + first_I_junction = add_junction!( + uwd, one_pop_typing.population_type, variable = :I1 + ) + + initial_boxes = count_boxes(uwd) + add_infection!(uwd, S_junction, I_junction, first_I_junction, one_pop_typing) + + @test count_boxes(uwd) == initial_boxes + 1 + @test count_junctions(uwd) == 3 # S, I, I1 + # Verify box name + @test subpart(uwd, boxes(uwd)[end], :name) == :transmission + end + + # UninfectedInfectedTyping: S + I → I + I + @testset "UninfectedInfectedTyping" begin + uwd = RelationDiagram([ui_typing.uninfected_type, ui_typing.infected_type]) + S_junction = add_junction!(uwd, ui_typing.uninfected_type, variable = :S) + I_junction = add_junction!(uwd, ui_typing.infected_type, variable = :I) + first_I_junction = add_junction!(uwd, ui_typing.infected_type, variable = :I1) + + initial_boxes = count_boxes(uwd) + add_infection!(uwd, S_junction, I_junction, first_I_junction, ui_typing) + + @test count_boxes(uwd) == initial_boxes + 1 + @test count_junctions(uwd) == 3 # S, I, I1 + # Verify box name + @test subpart(uwd, boxes(uwd)[end], :name) == :transmission + end +end + +@testitem "Infection mechanism (exposure to E)" setup = [MechanismSetup] begin + # OnePopulationTyping: S + I → E + I + @testset "OnePopulationTyping" begin + uwd = create_test_uwd(one_pop_typing, 3) + S_junction = add_junction!(uwd, one_pop_typing.population_type, variable = :S) + I_junction = add_junction!(uwd, one_pop_typing.population_type, variable = :I) + E_junction = add_junction!(uwd, one_pop_typing.population_type, variable = :E) + + initial_boxes = count_boxes(uwd) + add_infection!(uwd, S_junction, I_junction, E_junction, one_pop_typing) + + @test count_boxes(uwd) == initial_boxes + 1 + @test count_junctions(uwd) == 3 # S, I, E + # Verify box name + @test subpart(uwd, boxes(uwd)[end], :name) == :transmission + end + + # UninfectedInfectedTyping: S + I → E + I + @testset "UninfectedInfectedTyping" begin + uwd = RelationDiagram( + [ + ui_typing.uninfected_type, ui_typing.infected_type, ui_typing.infected_type, + ] + ) + S_junction = add_junction!(uwd, ui_typing.uninfected_type, variable = :S) + I_junction = add_junction!(uwd, ui_typing.infected_type, variable = :I) + E_junction = add_junction!(uwd, ui_typing.infected_type, variable = :E) + + initial_boxes = count_boxes(uwd) + add_infection!(uwd, S_junction, I_junction, E_junction, ui_typing) + + @test count_boxes(uwd) == initial_boxes + 1 + @test count_junctions(uwd) == 3 # S, I, E + # Verify box name + @test subpart(uwd, boxes(uwd)[end], :name) == :transmission + end +end + +@testitem "Disease progression mechanism" setup = [MechanismSetup] begin + # OnePopulationTyping: X → Y + @testset "OnePopulationTyping" begin + uwd = create_test_uwd(one_pop_typing, 2) + from_junction = add_junction!(uwd, one_pop_typing.population_type, variable = :E) + to_junction = add_junction!(uwd, one_pop_typing.population_type, variable = :I) + + initial_boxes = count_boxes(uwd) + add_disease_progression!(uwd, from_junction, to_junction, one_pop_typing) + + @test count_boxes(uwd) == initial_boxes + 1 + @test count_junctions(uwd) == 2 # from and to + # Verify box name + @test subpart(uwd, boxes(uwd)[end], :name) == :disease + end + + # UninfectedInfectedTyping: X → Y (infected type) + @testset "UninfectedInfectedTyping" begin + uwd = create_test_uwd(ui_typing, 2) + from_junction = add_junction!(uwd, ui_typing.infected_type, variable = :E) + to_junction = add_junction!(uwd, ui_typing.infected_type, variable = :I) + + initial_boxes = count_boxes(uwd) + add_disease_progression!(uwd, from_junction, to_junction, ui_typing) + + @test count_boxes(uwd) == initial_boxes + 1 + @test count_junctions(uwd) == 2 # from and to + # Verify box name + @test subpart(uwd, boxes(uwd)[end], :name) == :disease + end +end + +@testitem "Uninfected density progression mechanism" setup = [MechanismSetup] begin + # OnePopulationTyping: X → Y + @testset "OnePopulationTyping" begin + uwd = create_test_uwd(one_pop_typing, 2) + from_junction = add_junction!(uwd, one_pop_typing.population_type, variable = :R) + to_junction = add_junction!(uwd, one_pop_typing.population_type, variable = :S) + + initial_boxes = count_boxes(uwd) + add_uninfected_density_progression!(uwd, from_junction, to_junction, one_pop_typing) + + @test count_boxes(uwd) == initial_boxes + 1 + @test count_junctions(uwd) == 2 # from and to + # Verify box name + @test subpart(uwd, boxes(uwd)[end], :name) == :waning + end + + # UninfectedInfectedTyping: X → Y (uninfected type) + @testset "UninfectedInfectedTyping" begin + uwd = RelationDiagram([ui_typing.uninfected_type, ui_typing.uninfected_type]) + from_junction = add_junction!(uwd, ui_typing.uninfected_type, variable = :R) + to_junction = add_junction!(uwd, ui_typing.uninfected_type, variable = :S) + + initial_boxes = count_boxes(uwd) + add_uninfected_density_progression!(uwd, from_junction, to_junction, ui_typing) + + @test count_boxes(uwd) == initial_boxes + 1 + @test count_junctions(uwd) == 2 # from and to + # Verify box name + @test subpart(uwd, boxes(uwd)[end], :name) == :waning + end +end + +@testitem "Reversion progression mechanism" setup = [MechanismSetup] begin + # OnePopulationTyping: infected → uninfected + @testset "OnePopulationTyping" begin + uwd = create_test_uwd(one_pop_typing, 2) + from_junction = add_junction!(uwd, one_pop_typing.population_type, variable = :I) + to_junction = add_junction!(uwd, one_pop_typing.population_type, variable = :S) + + initial_boxes = count_boxes(uwd) + add_reversion_progression!(uwd, from_junction, to_junction, one_pop_typing) + + @test count_boxes(uwd) == initial_boxes + 1 + @test count_junctions(uwd) == 2 # from and to + # Verify box name + @test subpart(uwd, boxes(uwd)[end], :name) == :reversion + end + + # UninfectedInfectedTyping: infected → uninfected (different types) + @testset "UninfectedInfectedTyping" begin + uwd = RelationDiagram([ui_typing.infected_type, ui_typing.uninfected_type]) + from_junction = add_junction!(uwd, ui_typing.infected_type, variable = :I) + to_junction = add_junction!(uwd, ui_typing.uninfected_type, variable = :S) + + initial_boxes = count_boxes(uwd) + add_reversion_progression!(uwd, from_junction, to_junction, ui_typing) + + @test count_boxes(uwd) == initial_boxes + 1 + @test count_junctions(uwd) == 2 # from and to + # Verify box name + @test subpart(uwd, boxes(uwd)[end], :name) == :reversion + end +end diff --git a/AlgebraicEpiMech/test/test_observation_chain.jl b/AlgebraicEpiMech/test/test_observation_chain.jl new file mode 100644 index 0000000..c26b449 --- /dev/null +++ b/AlgebraicEpiMech/test/test_observation_chain.jl @@ -0,0 +1,111 @@ +# `observation_layout` — the READ side of observation chains. +# +# Given a net that already carries chains, recover which species belong to which chain and in what +# order. The construction side it used to accompany (`WithObservation`, `ObservationChainSpec`) +# has been removed: observation is attached to a COMPOSED net by pushout, so these nets are built +# with `attach_observation`. +# +# The tests that used to live here guarding the operadic path — "obs_inflow preserved through +# typed_product with NoCrossImmunity / CompleteCrossImmunity / ImmuneHistory", and the two +# behavioural accumulator-fills checks — are gone rather than ported. The property they asserted +# (that a factor's reflexive `:observation` box survives the pullback) no longer exists, and the +# behaviour they checked is covered directly by `test_observation_rewriting.jl`. + +@testsnippet ObservationLayoutSetup begin + using AlgebraicPetri + using AlgebraicPetri.TypedPetri: typed_product + using Catlab + using AlgebraicEpiMech + + one_pop_typing = OnePopulationTyping(population_type = :Individual) +end + +@testitem "observation_layout returns ordered chain metadata" setup = [ObservationLayoutSetup] begin + pn = attach_observation( + dom(create_model(one_pop_typing, SEIR())), + AtCompartment(:I); n_stages = 2, + ) + + layout = observation_layout(pn) + + @test layout isa ObservationLayout + @test layout.obs_names == (:O_I_1, :O_I_2) + @test length(layout.chains) == 1 + @test layout.chains[1] isa ObservationChainLayout + @test layout.chains[1].source_name == :I + @test layout.chains[1].obs_names == (:O_I_1, :O_I_2) + @test layout.chains[1].cumulative_name == :O_I_2 + @test layout.cumulative_names == (:O_I_2,) +end + +@testitem "observation layout constructors derive cumulative names" begin + using AlgebraicEpiMech + + chain = ObservationChainLayout(:I, (:O_I_1, :O_I_2)) + layout = ObservationLayout(chain.obs_names, (chain,)) + + @test chain.cumulative_name == :O_I_2 + @test layout.cumulative_names == (:O_I_2,) + @test_throws ArgumentError ObservationChainLayout(:I, ()) + @test_throws ArgumentError ObservationLayout((), ((source_name = :I,),)) +end + +@testitem "observation_layout recognises custom prefixes" setup = [ObservationLayoutSetup] begin + pn = attach_observation( + dom(create_model(one_pop_typing, SEIR())), + AtCompartment(:I); n_stages = 2, prefix = :reported_case, + ) + + layout = observation_layout(pn; prefix = :reported_case) + + @test layout.obs_names == (:reported_case_I_1, :reported_case_I_2) + @test only(layout.chains).source_name == :I + @test only(layout.chains).cumulative_name == :reported_case_I_2 + @test isempty(observation_layout(pn).chains) +end + +@testitem "observation_layout separates chains on different sources" setup = [ObservationLayoutSetup] begin + # Two chains, attached one after another. Each is named for the source it samples, which is + # what keeps them apart — a chain naming scheme that dropped the source would collapse both + # onto one label. + pn = attach_observation( + attach_observation( + dom(create_model(one_pop_typing, SEIRS())), + AtCompartment(:I); n_stages = 2, + ), + AtCompartment(:R); n_stages = 1, + ) + + layout = observation_layout(pn) + + @test Set(layout.obs_names) == Set((:O_I_1, :O_I_2, :O_R_1)) + @test Set(chain.source_name for chain in layout.chains) == Set((:I, :R)) + @test Set(layout.cumulative_names) == Set((:O_I_2, :O_R_1)) + + chains = Dict(chain.source_name => chain for chain in layout.chains) + @test chains[:I].obs_names == (:O_I_1, :O_I_2) + @test chains[:R].obs_names == (:O_R_1,) +end + +@testitem "observation_layout resolves stratified chains per stratum" setup = [ObservationLayoutSetup] begin + # Composition FIRST, observation after — so the stratification factor carries no observation + # structure at all and each location's chain is created by the rewrite. + composed = dom( + typed_product( + create_model(one_pop_typing, SEIR()), + create_model(one_pop_typing, GeographicStratification([:ak, :ca])), + ), + ) + pn = attach_observation(composed, AtCompartment(:I); n_stages = 2) + + layout = observation_layout(pn) + + @test length(layout.chains) == 2 + chains = Dict(chain.source_name => chain for chain in layout.chains) + @test Set(keys(chains)) == Set((:I_ak, :I_ca)) + @test Set(layout.cumulative_names) == Set((:O_I_2_ak, :O_I_2_ca)) + for loc in (:ak, :ca) + @test chains[Symbol(:I_, loc)].obs_names == + (Symbol(:O_I_1_, loc), Symbol(:O_I_2_, loc)) + end +end diff --git a/AlgebraicEpiMech/test/test_observation_rewriting.jl b/AlgebraicEpiMech/test/test_observation_rewriting.jl new file mode 100644 index 0000000..6b7ebfb --- /dev/null +++ b/AlgebraicEpiMech/test/test_observation_rewriting.jl @@ -0,0 +1,310 @@ +# Observation attached by pushout, after composition. +# +# The two observation processes are the SAME colimit applied to different rules, so these tests +# assert the two things that distinguish them and would silently be wrong otherwise: +# +# incidence — one observation per event, at the event's own rate. Checked as +# `du[accumulator] == -du[S]` on a net where susceptibles drain only to infection. +# No rate matching is assumed, because there is only one rate. +# prevalence — rate x occupancy, source not consumed. +# +# and the property that motivated the whole approach: attaching observation LEAVES COMPOSITION +# ALONE, so stratification factors never learn that observation exists. + +@testsnippet RewriteSetup begin + using AlgebraicPetri + using AlgebraicPetri: snames, tnames, sname, tname + using AlgebraicPetri.TypedPetri: typed_product + using Catlab + + flat(x) = AlgebraicEpiMech.flatten_symbols(x) + names(pn) = [flat(s) for s in snames(pn)] + tnames_flat(pn) = [flat(t) for t in tnames(pn)] + + "Evaluate the vectorfield with `rates` (by flattened transition name) and `state`." + function derivs(pn; rates, state) + u = Dict(s => get(state, s, 0.0) for s in names(pn)) + p = Dict(t => get(rates, t, 0.0) for t in tnames_flat(pn)) + du = Dict(s => 0.0 for s in names(pn)) + vectorfield_flat(pn)(du, u, p, 0.0) + return du + end + + "All rates whose name starts with `pre` set to `v`." + matching(pn, pre, v) = + Dict(t => v for t in tnames_flat(pn) if startswith(string(t), string(pre))) + + geo_seirs(locs) = dom( + typed_product( + create_model(OnePopulationTyping(), SEIRS()), + create_model(OnePopulationTyping(), GeographicStratification(locs)), + ), + ) +end + +@testitem "incidence: an infection tap counts one per infection, at the event's own rate" setup = [RewriteSetup] begin + # SEIR, not SEIRS: with no waning, S drains only to infection, so `-du[:S]` IS incidence. + base = dom(create_model(OnePopulationTyping(), SEIR())) + obs = attach_observation(base, AtEvent(:transmission); n_stages = 1) + + # The event was AUGMENTED, not shadowed: no new transition. + @test nt(obs) == nt(base) + @test ns(obs) == ns(base) + 1 + @test :O_transmission_1 in names(obs) + + du = derivs( + obs; + rates = merge(matching(obs, :transmission, 0.002), matching(obs, :E_to_I, 0.4)), + state = Dict(:S => 900.0, :E => 40.0, :I => 60.0), + ) + @test du[:O_transmission_1] ≈ -du[:S] # exactly one observation per infection + @test du[:O_transmission_1] ≈ 0.002 * 900.0 * 60.0 + + # Independent of dwell times — the point of tapping the event rather than a compartment. + for (sigma, gamma) in ((0.05, 0.05), (3.0, 2.0)) + alt = derivs( + obs; + rates = merge( + matching(obs, :transmission, 0.002), + matching(obs, :E_to_I, sigma), matching(obs, :I_to_R, gamma), + ), + state = Dict(:S => 900.0, :E => 40.0, :I => 60.0), + ) + @test alt[:O_transmission_1] ≈ du[:O_transmission_1] + @test alt[:O_transmission_1] ≈ -alt[:S] + end +end + +@testitem "prevalence: a compartment tap is rate x occupancy and consumes nobody" setup = [RewriteSetup] begin + base = dom(create_model(OnePopulationTyping(), SEIRS())) + obs = attach_observation(base, AtCompartment(:I); n_stages = 1) + + @test nt(obs) == nt(base) + 1 # a NEW mechanism, unlike the incidence rule + @test :O_I_1 in names(obs) # named for the compartment it samples + + du = derivs( + obs; + rates = matching(obs, :obs_inflow_I, 0.25), + state = Dict(:S => 900.0, :I => 60.0), + ) + @test du[:O_I_1] ≈ 0.25 * 60.0 # occupancy, not flux + @test du[:I] ≈ 0.0 # catalytic +end + +@testitem "chains add Erlang stages and only the last accumulates" setup = [RewriteSetup] begin + base = dom(create_model(OnePopulationTyping(), SEIR())) + obs = attach_observation(base, AtEvent(:transmission); n_stages = 3) + + @test ns(obs) == ns(base) + 3 + @test nt(obs) == nt(base) + 2 # two delay transitions, event still augmented + for s in (:O_transmission_1, :O_transmission_2, :O_transmission_3) + @test s in names(obs) + end + + # The terminal stage has no outflow: it is the accumulator the observation model reads. + terminal = only([s for s in 1:ns(obs) if flat(sname(obs, s)) === :O_transmission_3]) + @test isempty([t for t in 1:nt(obs) if terminal in AlgebraicPetri.inputs(obs, t)]) + + # Flow reaches the end of the chain. + du = derivs( + obs; + rates = merge( + matching(obs, :transmission, 0.002), matching(obs, :O_transmission_1_to_O_transmission_2, 0.5), + ), + state = Dict(:S => 900.0, :I => 60.0, :O_transmission_1 => 10.0), + ) + @test du[:O_transmission_2] > 0.0 +end + +@testitem "composition is untouched: factors never learn observation exists" setup = [RewriteSetup] begin + # The composed net is built with no knowledge of observation whatsoever. If attaching + # observation required factor support, this net would be missing structure BEFORE we start — + # which is the failure mode the pushout approach removes entirely. + base = geo_seirs([:a, :b]) + @test ns(base) == 8 # S,E,I,R x 2 — no accumulators + @test nt(base) == 10 # 4 routes + 3 progressions x 2 + + obs = attach_observation(base, AtEvent(:transmission); n_stages = 1) + @test nt(obs) == nt(base) # still no new transitions + @test ns(obs) == ns(base) + 2 # one accumulator per location, NOT per route + + # Every route into a location records into that location's accumulator. + du = derivs( + obs; + rates = matching(obs, :transmission, 0.002), + state = Dict(:S_a => 900.0, :S_b => 500.0, :I_a => 60.0, :I_b => 40.0), + ) + @test du[:O_transmission_1_a] ≈ -du[:S_a] + @test du[:O_transmission_1_b] ≈ -du[:S_b] + # Both routes into `a` count into the same accumulator — the pushout GLUED rather than added. + @test du[:O_transmission_1_a] ≈ 0.002 * 900.0 * (60.0 + 40.0) +end + +@testitem "escape-selective infection is preserved without any mirroring" setup = [RewriteSetup] begin + # Under ImmuneHistory a class immune to `h` simply has no transmission for `i in h`. Because + # we only rewrite transitions that EXIST, selectivity survives automatically — there is no + # second structure to keep in sync and therefore nothing to get silently wrong. + ui = UninfectedInfectedTyping() + base = dom( + typed_product( + create_model(ui, SEIRS()), + create_model(ui, ImmuneHistory([:current, :invader])), + ), + ) + routes = count(t -> startswith(string(flat(tname(base, t))), "transmission"), 1:nt(base)) + @test routes == 8 + + obs = attach_observation(base, AtEvent(:transmission); n_stages = 1) + @test nt(obs) == nt(base) + @test ns(obs) == ns(base) + 4 # one chain per (h,i) junction, not per route + + accs = sort([s for s in names(obs) if startswith(string(s), "O_transmission_1_")]) + @test length(accs) == 4 + + du = derivs( + obs; + rates = matching(obs, :transmission, 0.002), + state = Dict( + s => ( + startswith(string(s), "S_U") ? 500.0 : + startswith(string(s), "I_") ? 50.0 : 0.0 + ) for s in names(obs) + ), + ) + observed = sum(du[a] for a in accs) + infected = -sum(du[s] for s in names(obs) if startswith(string(s), "S_U")) + @test observed ≈ infected # mass 1 across every escape-permitted route +end + +@testitem "attaching observation is idempotent and order-independent" setup = [RewriteSetup] begin + base = geo_seirs([:a, :b]) + + once = attach_observation(base, AtEvent(:transmission); n_stages = 2) + twice = attach_observation(once, AtEvent(:transmission); n_stages = 2) + @test ns(twice) == ns(once) # already recorded, nothing added + @test nt(twice) == nt(once) + + # Both kinds can coexist: incidence off the event, prevalence off a compartment. + both = attach_observation( + attach_observation(base, AtEvent(:transmission); n_stages = 1), + AtCompartment(:I); n_stages = 1, prefix = :P, + ) + @test :O_transmission_1_a in names(both) # incidence chain + @test :P_I_1_a in names(both) # prevalence chain, distinct prefix +end + +@testitem "reusing an observation namespace requires the same chain shape" setup = [RewriteSetup] begin + base = dom(create_model(OnePopulationTyping(), SEIR())) + + event_one = attach_observation(base, AtEvent(:transmission); n_stages = 1) + event_two = attach_observation(base, AtEvent(:transmission); n_stages = 2) + @test_throws ArgumentError attach_observation( + event_one, AtEvent(:transmission); n_stages = 2 + ) + @test_throws ArgumentError attach_observation( + event_two, AtEvent(:transmission); n_stages = 1 + ) + + compartment_one = attach_observation(base, AtCompartment(:I); n_stages = 1) + compartment_two = attach_observation(base, AtCompartment(:I); n_stages = 2) + @test_throws ArgumentError attach_observation( + compartment_one, AtCompartment(:I); n_stages = 2 + ) + @test_throws ArgumentError attach_observation( + compartment_two, AtCompartment(:I); n_stages = 1 + ) +end + +@testitem "gluing a shared chain identifies every stage, not just the first" setup = [RewriteSetup] begin + # Regression: `L` originally carried only the FIRST stage when gluing, so each additional + # route into a stratum re-created the later stages and delay transitions. The duplicates + # shared a label, so nothing downstream could tell them apart. Needs BOTH a shared chain + # (stratified) and more than one stage to show up — the original tests had one or the other. + base = geo_seirs([:a, :b]) + obs = attach_observation(base, AtEvent(:transmission); n_stages = 3) + + for loc in (:a, :b), stage in 1:3 + label = Symbol("O_transmission_", stage, "_", loc) + @test count(==(label), names(obs)) == 1 + end + for loc in (:a, :b), stage in 1:2 + label = Symbol( + "O_transmission_", stage, "_to_O_transmission_", stage + 1, "_", loc + ) + @test count(==(label), tnames_flat(obs)) == 1 + end + + # 4 routes into 2 strata: 3 stages each, and 2 delay transitions each. + @test ns(obs) == ns(base) + 6 + @test nt(obs) == nt(base) + 4 +end + +@testitem "each stage of a multi-stage compartment gets its own chain" setup = [RewriteSetup] begin + # Regression: chain labels that omitted the source collapsed `I1`, `I2`, `I3` onto one label, + # and every match after the first was skipped as already-tapped. + base = dom(create_model(OnePopulationTyping(), SEIR(number_I_stages = 3))) + obs = attach_observation(base, AtCompartment(:I); n_stages = 1) + + @test nt(obs) == nt(base) + 3 + for stage in ("I1", "I2", "I3") + @test Symbol("O_", stage, "_1") in names(obs) + @test Symbol("obs_inflow_", stage) in tnames_flat(obs) + end +end + +@testitem "rewritten chains are recognised by observation_layout" setup = [RewriteSetup] begin + # The naming contract is not cosmetic: `observation_layout` parses a + # `__` leaf, and `ConfigurableEpi`'s `StateLayout` builds on it. Labels + # that do not parse leave the accumulators looking like core model state. + obs = attach_observation( + dom(create_model(OnePopulationTyping(), SEIR())), + AtEvent(:transmission); n_stages = 2, + ) + layout = observation_layout(obs) + @test Set(layout.obs_names) == Set((:O_transmission_1, :O_transmission_2)) + @test length(layout.chains) == 1 + @test only(layout.chains).obs_names == (:O_transmission_1, :O_transmission_2) + @test only(layout.cumulative_names) == :O_transmission_2 + + # And on a stratified net, one chain per stratum. + geo = attach_observation(geo_seirs([:a, :b]), AtEvent(:transmission); n_stages = 2) + @test length(observation_layout(geo).chains) == 2 +end + +@testitem "direct attachment: structure, input-order stability, idempotency" setup = [RewriteSetup] begin + # The observed net is the base net plus, per matched route, exactly one output arc into the + # infectee location's chain head; base parts keep their order and positions; re-application + # is a no-op. Together with the rate identities above this is the whole contract. + using Catlab: nparts, subpart + + base = geo_seirs([:a, :b, :c]) + obs = attach_observation(base, AtEvent(:transmission); n_stages = 2) + routes = [ + t for t in 1:AlgebraicPetri.nt(base) + if startswith(string(tnames_flat(base)[t]), "transmission") + ] + @test length(routes) == 9 + @test names(obs)[1:AlgebraicPetri.ns(base)] == names(base) + @test tnames_flat(obs)[1:AlgebraicPetri.nt(base)] == tnames_flat(base) + # 3 chains of 2 stages, each with one delay transition + @test AlgebraicPetri.ns(obs) == AlgebraicPetri.ns(base) + 6 + @test AlgebraicPetri.nt(obs) == AlgebraicPetri.nt(base) + 3 + @test nparts(obs, :O) == nparts(base, :O) + length(routes) + 3 + @test nparts(obs, :I) == nparts(base, :I) + 3 + for t in routes + heads = [ + flat(sname(obs, subpart(obs, o, :os))) + for o in 1:nparts(obs, :O) if subpart(obs, o, :ot) == t + ] + @test count(s -> startswith(string(s), "O_transmission_1"), heads) == 1 + end + @test attach_observation(obs, AtEvent(:transmission); n_stages = 2) == obs + + # Prevalence: one catalytic tap transition plus a chain per matched species. + seir = dom(create_model(OnePopulationTyping(), SEIR())) + tap = attach_observation(seir, AtCompartment(:I); n_stages = 2) + @test names(tap)[1:AlgebraicPetri.ns(seir)] == names(seir) + @test AlgebraicPetri.ns(tap) == AlgebraicPetri.ns(seir) + 2 + @test AlgebraicPetri.nt(tap) == AlgebraicPetri.nt(seir) + 2 + @test attach_observation(tap, AtCompartment(:I); n_stages = 2) == tap +end diff --git a/AlgebraicEpiMech/test/test_stratification.jl b/AlgebraicEpiMech/test/test_stratification.jl new file mode 100644 index 0000000..908c152 --- /dev/null +++ b/AlgebraicEpiMech/test/test_stratification.jl @@ -0,0 +1,271 @@ +@testitem "Stratification composition" begin + using AlgebraicEpiMech + + @testset "ContactStratification composition via compose_stratifications" begin + age = ContactStratification([:child, :adult], :age) + geo = ContactStratification([:urban, :rural], :geography) + + # Compose using function + age_geo = compose_stratifications(age, geo) + + # Check stratum names are Cartesian product + @test age_geo.stratum_names == + [:childxurban, :childxrural, :adultxurban, :adultxrural] + + # Check label is combined + @test age_geo.label == :age_x_geography + + # Check it's still a valid ContactStratification + @test age_geo isa ContactStratification + @test length(age_geo.stratum_names) == 4 + end + + @testset "ContactStratification composition via * operator" begin + age = ContactStratification([:child, :adult], :age) + geo = ContactStratification([:urban, :rural], :geography) + + # Compose using * operator + age_geo = age * geo + + # Should give same results as function + @test age_geo.stratum_names == + [:childxurban, :childxrural, :adultxurban, :adultxrural] + @test age_geo.label == :age_x_geography + end + + @testset "Composition with convenience constructors" begin + age = AgeStratification([:child, :adult, :elderly]) + geo = GeographicStratification([:urban, :suburban, :rural]) + + age_geo = age * geo + + # Should create 3×3 = 9 strata + @test length(age_geo.stratum_names) == 9 + + # Check ordering (child × all locations first, then adult × all, etc.) + expected = [ + :childxurban, :childxsuburban, :childxrural, + :adultxurban, :adultxsuburban, :adultxrural, + :elderlyxurban, :elderlyxsuburban, :elderlyxrural, + ] + @test age_geo.stratum_names == expected + + # Check label + @test age_geo.label == :age_x_geography + end + + @testset "Triple composition (associativity)" begin + age = ContactStratification([:child, :adult], :age) + geo = ContactStratification([:urban, :rural], :geography) + risk = ContactStratification([:low, :high], :risk) + + # Compose in different orders + result1 = (age * geo) * risk + result2 = age * (geo * risk) + + # Both should have 2×2×2 = 8 strata + @test length(result1.stratum_names) == 8 + @test length(result2.stratum_names) == 8 + + # Check some expected names + @test :childxurbanxlow in result1.stratum_names + @test :adultxruralxhigh in result1.stratum_names + + # Note: labels will differ due to grouping + @test result1.label == :age_x_geography_x_risk + @test result2.label == :age_x_geography_x_risk # Should still combine nicely + end + + @testset "Composition preserves uniqueness constraint" begin + strat1 = ContactStratification([:a, :b], :test1) + strat2 = ContactStratification([:x, :y], :test2) + + composed = strat1 * strat2 + + # All names should be unique + @test length(composed.stratum_names) == length(unique(composed.stratum_names)) + end + + @testset "Edge case: single stratum compositions" begin + single = ContactStratification([:only], :single) + multi = ContactStratification([:a, :b, :c], :multi) + + result = single * multi + + @test result.stratum_names == [:onlyxa, :onlyxb, :onlyxc] + @test length(result.stratum_names) == 3 + end +end + +@testitem "ContactStratification model integration" begin + using AlgebraicEpiMech + using AlgebraicPetri, AlgebraicPetri.TypedPetri + using Catlab + + @testset "Composed stratification creates valid model" begin + typing = OnePopulationTyping() + + age = AgeStratification([:child, :adult]) + geo = GeographicStratification([:urban, :rural]) + + # Create product stratification + age_geo = age * geo + + # Should be able to create model from composed stratification + age_geo_model = create_model(typing, age_geo) + + # Extract Petri net + pn = dom(age_geo_model) + + # Should have 4 strata (junctions) + @test ns(pn) == 4 + + # Should have n² + 3n = 28 transitions (16 transmission + 12 reflexive) + @test nt(pn) == 28 + end + + @testset "Composed stratification equivalent to sequential composition" begin + typing = OnePopulationTyping() + sir = create_model(typing, SIR()) + age = AgeStratification([:child, :adult]) + geo = GeographicStratification([:urban, :rural]) + + # Approach 1: Sequential typed_product + age_model = create_model(typing, age) + geo_model = create_model(typing, geo) + sir_age = typed_product(sir, age_model) + sir_age_geo_sequential = typed_product(sir_age, geo_model) + + # Approach 2: Pre-composed stratification + age_geo = age * geo + age_geo_model = create_model(typing, age_geo) + sir_age_geo_composed = typed_product(sir, age_geo_model) + + # Both should create same structure + pn_seq = dom(sir_age_geo_sequential) + pn_comp = dom(sir_age_geo_composed) + + # Same number of states (S, I, R) × 4 strata + @test ns(pn_seq) == ns(pn_comp) + @test ns(pn_comp) == 12 # 3 compartments × 4 strata + + # Same number of transitions + @test nt(pn_seq) == nt(pn_comp) + end + + @testset "Transmission transitions preserve infectee stratum (regression test)" begin + # This test catches the bug where S_child + I_adult → I_adult + I_adult + # instead of the correct S_child + I_adult → I_child + I_adult + + typing = OnePopulationTyping() + sir = create_model(typing, SIR()) + age = AgeStratification([:child, :adult]) + + age_model = create_model(typing, age) + age_sir = typed_product(sir, age_model) + pn = dom(age_sir) + + # Find the transmission transition where child is infected by adult + # Transition names format: (:transmission_S_I, :child_adult) + # This means "child infected by adult" + child_by_adult_idx = findfirst( + t -> t == (:transmission_S_I, :child_adult), + tnames(pn) + ) + + @test child_by_adult_idx !== nothing + + # Get the inputs and outputs of this transition + inputs = [pn[edge, :is] for edge in incident(pn, child_by_adult_idx, :it)] + outputs = [pn[edge, :os] for edge in incident(pn, child_by_adult_idx, :ot)] + + input_states = [snames(pn)[i] for i in inputs] + output_states = [snames(pn)[i] for i in outputs] + + # Expected: S_child + I_adult → I_child + I_adult + # The infectee (child) should become I_child, not I_adult + @test Set(input_states) == Set([(:S, :child), (:I, :adult)]) + @test Set(output_states) == Set([(:I, :child), (:I, :adult)]) + + # Critical check: infectee (child) stays in child stratum when infected + @test (:I, :child) in output_states + @test !any( + s -> s == (:I, :adult) && count(==(s), output_states) == 2, [ + ( + :I, :adult, + ), + ] + ) + end + + @testset "Multi-stratification transmission preserves infectee stratum" begin + # Test with composed stratifications (age × geography) + typing = OnePopulationTyping() + sir = create_model(typing, SIR()) + + age = AgeStratification([:child, :adult]) + geo = GeographicStratification([:urban, :rural]) + age_geo = age * geo + + age_geo_model = create_model(typing, age_geo) + age_geo_sir = typed_product(sir, age_geo_model) + pn = dom(age_geo_sir) + + # Test critical case: urban child infected by rural adult + # Transition name format: infectee_infector (with 'x' within each stratum name) + childxurban_by_adultxrural_idx = findfirst( + t -> t == (:transmission_S_I, :childxurban_adultxrural), + tnames(pn) + ) + + @test childxurban_by_adultxrural_idx !== nothing + + inputs = [ + pn[edge, :is] + for edge in incident(pn, childxurban_by_adultxrural_idx, :it) + ] + outputs = [ + pn[edge, :os] + for edge in incident(pn, childxurban_by_adultxrural_idx, :ot) + ] + + input_states = [snames(pn)[i] for i in inputs] + output_states = [snames(pn)[i] for i in outputs] + + # Expected: S_childxurban + I_adultxrural → I_childxurban + I_adultxrural + # The infectee (childxurban) must become I_childxurban, NOT I_adultxrural + @test Set(input_states) == Set([(:S, :childxurban), (:I, :adultxrural)]) + @test Set(output_states) == Set([(:I, :childxurban), (:I, :adultxrural)]) + + # Verify infectee stays in their stratum + @test (:I, :childxurban) in output_states + + # Test several other cross-stratum transmissions + test_cases = [ + ( + (:transmission_S_I, :childxrural_adultxurban), + [(:S, :childxrural), (:I, :adultxurban)], + [(:I, :childxrural), (:I, :adultxurban)], + ), + ( + (:transmission_S_I, :adultxurban_childxrural), + [(:S, :adultxurban), (:I, :childxrural)], + [(:I, :adultxurban), (:I, :childxrural)], + ), + ] + + for (tname, expected_inputs, expected_outputs) in test_cases + tidx = findfirst(t -> t == tname, tnames(pn)) + @test tidx !== nothing + + inputs = [pn[edge, :is] for edge in incident(pn, tidx, :it)] + outputs = [pn[edge, :os] for edge in incident(pn, tidx, :ot)] + + input_states = [snames(pn)[i] for i in inputs] + output_states = [snames(pn)[i] for i in outputs] + + @test Set(input_states) == Set(expected_inputs) + @test Set(output_states) == Set(expected_outputs) + end + end +end diff --git a/AlgebraicEpiMech/test/test_transition_naming.jl b/AlgebraicEpiMech/test/test_transition_naming.jl new file mode 100644 index 0000000..2a335b5 --- /dev/null +++ b/AlgebraicEpiMech/test/test_transition_naming.jl @@ -0,0 +1,252 @@ +# Unit tests for transition naming from UWD structure +# Tests that create_model generates unique transition names for ODE rate assignment + +@testsnippet TransitionNamingSetup begin + using AlgebraicPetri + using Catlab + using AlgebraicEpiMech + + # Create typing for testing + one_pop_typing = OnePopulationTyping(population_type = :Individual) +end + +# Test simple SIR model transition naming +@testitem "SIR model has unique transition names" setup = [TransitionNamingSetup] begin + typed_model = create_model(one_pop_typing, SIR()) + pn = dom(typed_model) + + tnames = AlgebraicPetri.tnames(pn) + + # Should have 2 transitions + @test length(tnames) == 2 + + # All transition names should be unique + @test length(unique(tnames)) == length(tnames) + + # Check expected naming pattern + @test :transmission_S_I in tnames + @test :I_to_R in tnames +end + +# Test multi-stage SI model transition naming +@testitem "Multi-stage SI model has unique transmission names" setup = [TransitionNamingSetup] begin + typed_model = create_model(one_pop_typing, SI(number_I_stages = 3)) + pn = dom(typed_model) + + tnames = AlgebraicPetri.tnames(pn) + + # Should have 5 transitions: 3 transmissions + 2 progressions + @test length(tnames) == 5 + + # All transition names should be unique + @test length(unique(tnames)) == length(tnames) + + # Check expected naming pattern for transmissions + @test :transmission_S_I1 in tnames + @test :transmission_S_I2 in tnames + @test :transmission_S_I3 in tnames + + # Check expected naming pattern for progressions + @test :I1_to_I2 in tnames + @test :I2_to_I3 in tnames +end + +# Test SEIR model transition naming +@testitem "SEIR model has unique transition names" setup = [TransitionNamingSetup] begin + typed_model = create_model(one_pop_typing, SEIR()) + pn = dom(typed_model) + + tnames = AlgebraicPetri.tnames(pn) + + # Should have 3 transitions + @test length(tnames) == 3 + + # All transition names should be unique + @test length(unique(tnames)) == length(tnames) + + # Check expected naming pattern + @test :transmission_S_I in tnames + @test :E_to_I in tnames + @test :I_to_R in tnames +end + +# Test multi-stage SEIR model transition naming +@testitem "Multi-stage SEIR model has unique transition names" setup = [TransitionNamingSetup] begin + typed_model = create_model( + one_pop_typing, + SEIR(number_E_stages = 2, number_I_stages = 3) + ) + pn = dom(typed_model) + + tnames = AlgebraicPetri.tnames(pn) + + # Should have 8 transitions: + # - 3 transmissions (one per I stage) + # - 1 E progression (E1->E2) + # - 2 I progressions (I1->I2, I2->I3) + # - 1 E to I transition (E2->I1) + # - 1 I to R transition (I3->R) + @test length(tnames) == 8 + + # All transition names should be unique + @test length(unique(tnames)) == length(tnames) + + # Check expected naming pattern for transmissions + @test :transmission_S_I1 in tnames + @test :transmission_S_I2 in tnames + @test :transmission_S_I3 in tnames + + # Check expected naming pattern for progressions + @test :E1_to_E2 in tnames + @test :I1_to_I2 in tnames + @test :I2_to_I3 in tnames + @test :E2_to_I1 in tnames + @test :I3_to_R in tnames +end + +# Test SEI model transition naming +@testitem "SEI model has unique transition names" setup = [TransitionNamingSetup] begin + typed_model = create_model(one_pop_typing, SEI()) + pn = dom(typed_model) + + tnames = AlgebraicPetri.tnames(pn) + + # Should have 2 transitions + @test length(tnames) == 2 + + # All transition names should be unique + @test length(unique(tnames)) == length(tnames) + + # Check expected naming pattern + @test :transmission_S_I in tnames + @test :E_to_I in tnames +end + +# Test SIS model transition naming +@testitem "SIS model has unique transition names" setup = [TransitionNamingSetup] begin + typed_model = create_model(one_pop_typing, SIS()) + pn = dom(typed_model) + + tnames = AlgebraicPetri.tnames(pn) + + # Should have 2 transitions + @test length(tnames) == 2 + + # All transition names should be unique + @test length(unique(tnames)) == length(tnames) + + # Check expected naming pattern + @test :transmission_S_I in tnames + @test :I_to_S in tnames +end + +# Test SEIS model transition naming +@testitem "SEIS model has unique transition names" setup = [TransitionNamingSetup] begin + typed_model = create_model(one_pop_typing, SEIS()) + pn = dom(typed_model) + + tnames = AlgebraicPetri.tnames(pn) + + # Should have 3 transitions + @test length(tnames) == 3 + + # All transition names should be unique + @test length(unique(tnames)) == length(tnames) + + # Check expected naming pattern + @test :transmission_S_I in tnames + @test :E_to_I in tnames + @test :I_to_S in tnames +end + +# Test SEIRS model transition naming +@testitem "SEIRS model has unique transition names" setup = [TransitionNamingSetup] begin + typed_model = create_model(one_pop_typing, SEIRS()) + pn = dom(typed_model) + + tnames = AlgebraicPetri.tnames(pn) + + # Should have 4 transitions + @test length(tnames) == 4 + + # All transition names should be unique + @test length(unique(tnames)) == length(tnames) + + # Check expected naming pattern + @test :transmission_S_I in tnames + @test :E_to_I in tnames + @test :I_to_R in tnames + @test :R_to_S in tnames +end + +# Test that transition names enable unambiguous ODE rate assignment +@testitem "Multi-stage model enables unique rate parameter assignment" setup = [TransitionNamingSetup] begin + using LabelledArrays + + typed_model = create_model(one_pop_typing, SI(number_I_stages = 3)) + pn = dom(typed_model) + + tnames = AlgebraicPetri.tnames(pn) + + # Should be able to create a parameter vector with unique names + # This would fail if transition names were not unique + param_dict = Dict(tname => 0.1 for tname in tnames) + + # Verify we can access each transition uniquely + @test haskey(param_dict, :transmission_S_I1) + @test haskey(param_dict, :transmission_S_I2) + @test haskey(param_dict, :transmission_S_I3) + @test haskey(param_dict, :I1_to_I2) + @test haskey(param_dict, :I2_to_I3) + + # Verify all 5 unique parameters can be set + @test length(param_dict) == 5 +end + +# Test that multi-input transitions follow the naming pattern +@testitem "Multi-input transitions follow box_name_input1_input2 pattern" setup = [TransitionNamingSetup] begin + typed_model = create_model(one_pop_typing, SIR()) + pn = dom(typed_model) + + # Find the transmission transition (has 2 inputs) + transmission_idx = findfirst( + t -> AlgebraicPetri.tname(pn, t) == :transmission_S_I, 1:AlgebraicPetri.nt(pn) + ) + @test transmission_idx !== nothing + + # Get its inputs + inputs = AlgebraicPetri.inputs(pn, transmission_idx) + @test length(inputs) == 2 + + # The name should reflect both inputs (sorted in reverse alphabetical order) + input_names = sort([string(AlgebraicPetri.sname(pn, s)) for s in inputs], rev = true) + @test input_names == ["S", "I"] + + # The transition name should follow the pattern + @test AlgebraicPetri.tname(pn, transmission_idx) == :transmission_S_I +end + +# Test that single-input transitions follow input_to_output pattern +@testitem "Single-input transitions follow input_to_output pattern" setup = [TransitionNamingSetup] begin + typed_model = create_model(one_pop_typing, SIR()) + pn = dom(typed_model) + + # Find the recovery transition (I -> R) + recovery_idx = findfirst( + t -> AlgebraicPetri.tname(pn, t) == :I_to_R, 1:AlgebraicPetri.nt(pn) + ) + @test recovery_idx !== nothing + + # Get its inputs and outputs + inputs = AlgebraicPetri.inputs(pn, recovery_idx) + outputs = AlgebraicPetri.outputs(pn, recovery_idx) + + @test length(inputs) == 1 + @test length(outputs) == 1 + + # The name should reflect input_to_output + input_name = AlgebraicPetri.sname(pn, inputs[1]) + output_name = AlgebraicPetri.sname(pn, outputs[1]) + + @test AlgebraicPetri.tname(pn, recovery_idx) == Symbol(input_name, "_to_", output_name) +end diff --git a/AlgebraicEpiMech/test/test_vectorfield_flat.jl b/AlgebraicEpiMech/test/test_vectorfield_flat.jl new file mode 100644 index 0000000..3b981e6 --- /dev/null +++ b/AlgebraicEpiMech/test/test_vectorfield_flat.jl @@ -0,0 +1,213 @@ +@testitem "flatten_symbols - basic tuple flattening" begin + using AlgebraicEpiMech + + result = flatten_symbols((:x, :y, :z)) + @test result == :x_y_z +end + +@testitem "flatten_symbols - nested tuple flattening" begin + using AlgebraicEpiMech + + result = flatten_symbols(((:x, :y), :z)) + @test result == :x_y_z +end + +@testitem "flatten_symbols - nested tuple flattening" begin + using AlgebraicEpiMech + + result = flatten_symbols((((:a, :b), :c), :d)) + @test result == :a_b_c_d +end + +@testitem "flatten_symbols - nested tuple flattening other direction" begin + using AlgebraicEpiMech + + result = flatten_symbols((:a, (:b, (:c, :d)))) + @test result == :a_b_c_d +end + +@testitem "flatten_symbols - single symbol" begin + using AlgebraicEpiMech + + result = flatten_symbols(:single) + @test result == :single +end + +@testitem "flatten_symbols - singleton tuple" begin + using AlgebraicEpiMech + + @test flatten_symbols((:single,)) == :single + @test flatten_symbols((((:single,),),)) == :single +end + +@testitem "vectorfield_flat rejects flattened-name collisions" begin + using AlgebraicPetri + using AlgebraicEpiMech + + TupleNamedPetriNet = AlgebraicPetri.LabelledPetriNetUntyped{Tuple{Symbol, Symbol}} + + species_collision = TupleNamedPetriNet() + add_species!( + species_collision, 2; sname = [(:a_b, :c), (:a, :b_c)] + ) + @test_throws ArgumentError vectorfield_flat(species_collision) + + transition_collision = TupleNamedPetriNet() + add_species!(transition_collision, 1; sname = [(:state, :one)]) + add_transitions!( + transition_collision, 2; tname = [(:a_b, :c), (:a, :b_c)] + ) + @test_throws ArgumentError vectorfield_flat(transition_collision) +end + +@testitem "vectorfield_flat - function signature and type" begin + using AlgebraicPetri, LabelledArrays + using AlgebraicEpiMech + + # Create a simple Petri net + pn = LabelledPetriNet([:S, :I], :infection => ((:S, :I) => (:I, :I))) + + # Generate vectorfield function + f = vectorfield_flat(pn) + + # Test that it returns a function + @test isa(f, Function) + + # Test function can be called with proper arguments + du = LVector(S = 0.0, I = 0.0) + u = LVector(S = 100.0, I = 1.0) + p = LVector(infection = 0.01) + t = 0.0 + + result = f(du, u, p, t) + @test du[:S] == -1.0 # Should return -1.0 change in S for these inputs + @test du[:I] == 1.0 # Should return 1.0 change in I for these inputs +end + +@testitem "vectorfield_flat matches the dense reference bit for bit" begin + # The sparse, name-resolved-once implementation must reproduce the dense reference + # (`dense_vectorfield_flat`, defined below) exactly — same factor and summation order, so `==` rather + # than `≈` — on both the production argument shapes (LVector state, NamedTuple rates whose + # keys are NOT in net order) and the by-name dictionary path the test snippets use. + using AlgebraicPetri, LabelledArrays, Catlab + using AlgebraicPetri: snames, tnames + using AlgebraicPetri.TypedPetri: typed_product + using AlgebraicEpiMech + + # The dense, resolve-everything-per-call form `vectorfield_flat` replaced (kept here, not in + # src, as the executable specification): the new implementation must match it bit for bit. + function dense_vectorfield_flat(pn) + tm = TransitionMatrices(pn) + dt = tm.output - tm.input + return ( + du, + u, + p, + t, + ) -> begin + rates = zeros(valtype(du), nt(pn)) + u_m = [u[flatten_symbols(sname(pn, i))] for i in 1:ns(pn)] + p_m = [p[flatten_symbols(tname(pn, i))] for i in 1:nt(pn)] + for i in 1:nt(pn) + rates[i] = AlgebraicPetri.valueat(p_m[i], u, t) * + prod(u_m[j]^tm.input[i, j] for j in 1:ns(pn)) + end + for j in 1:ns(pn) + du[flatten_symbols(sname(pn, j))] = sum( + rates[i] * dt[i, j] for i in 1:nt(pn); init = 0.0 + ) + end + du + end + end + + # Deterministic, irrational-stride "random" values: no RNG dependency in the test env, and + # every state and rate is distinct and non-round so an ordering slip cannot cancel out. + strided(n, seed) = [mod(seed + k * 0.6180339887, 1.0) for k in 1:n] + + geo = attach_observation( + dom( + typed_product( + create_model(OnePopulationTyping(), SEIRS()), + create_model(OnePopulationTyping(), GeographicStratification([:a, :b, :c])), + ), + ), + AtEvent(:transmission); n_stages = 2, + ) + dimer = LabelledPetriNet([:A, :B], :dimerise => ((:A, :A) => :B), :decay => (:B => :A)) + # Multiplicity 4: `u^4` is a correctly-rounded `pow`, not a repeated product, so only + # applying `^` per (species, multiplicity) as the reference does reproduces it exactly. + quad = LabelledPetriNet( + [:A, :B], :tetramerise => ((:A, :A, :A, :A) => :B), :split => (:B => (:A, :A, :A, :A)) + ) + simple = LabelledPetriNet([:S, :I], :infection => ((:S, :I) => (:I, :I))) + + for pn in (geo, dimer, quad, simple) + f = vectorfield_flat(pn) + g = dense_vectorfield_flat(pn) + ssyms = [flatten_symbols(s) for s in snames(pn)] + tsyms = [flatten_symbols(t) for t in tnames(pn)] + for rep in 1:3 + uvals = 100 .* strided(length(ssyms), 0.1 * rep) + pvals = strided(length(tsyms), 0.3 * rep) + u = LVector(NamedTuple{Tuple(ssyms)}(Tuple(uvals))) + # Rate keys deliberately out of net order: reversed, then rotated by `rep`. + perm = circshift(reverse(1:length(tsyms)), rep) + p = NamedTuple{Tuple(tsyms[perm])}(Tuple(pvals[perm])) + du_fast = LVector(NamedTuple{Tuple(ssyms)}(Tuple(zeros(length(ssyms))))) + du_ref = deepcopy(du_fast) + f(du_fast, u, p, 0.0) + g(du_ref, u, p, 0.0) + @test du_fast == du_ref + + ud = Dict(zip(ssyms, uvals)) + pd = Dict(zip(tsyms, pvals)) + dd_fast = Dict(s => 0.0 for s in ssyms) + dd_ref = copy(dd_fast) + f(dd_fast, ud, pd, 0.0) + g(dd_ref, ud, pd, 0.0) + @test dd_fast == dd_ref + @test all(dd_fast[ssyms[k]] == du_fast[k] for k in eachindex(ssyms)) + end + end + + # The production path (LVector state, NamedTuple rates, Float64) allocates nothing per call. + f_geo = vectorfield_flat(geo) + gsyms = [flatten_symbols(s) for s in snames(geo)] + gtsyms = [flatten_symbols(t) for t in tnames(geo)] + u_geo = LVector(NamedTuple{Tuple(gsyms)}(Tuple(100 .* strided(length(gsyms), 0.7)))) + p_geo = NamedTuple{Tuple(gtsyms)}(Tuple(strided(length(gtsyms), 0.9))) + du_geo = LVector(NamedTuple{Tuple(gsyms)}(Tuple(zeros(length(gsyms))))) + f_geo(du_geo, u_geo, p_geo, 0.0) + @test @allocated(f_geo(du_geo, u_geo, p_geo, 0.0)) == 0 + + # A narrower output type keeps the dense form's numeric path: rates converted to Float32, + # contributions accumulated from a Float64 zero, one final rounding. Six catalytic taps into + # `X` with rates 1e8, 1, 1, 1, 1, 1 give 100000008 that way and 100000000 if accumulated in + # Float32 — so this case discriminates. + taps = LabelledPetriNet( + [:X, :Y], [Symbol("tap", i) => (:Y => (:X, :Y)) for i in 1:6]... + ) + f32 = vectorfield_flat(taps) + g32 = dense_vectorfield_flat(taps) + u32 = LVector(X = 0.0f0, Y = 1.0f0) + p32 = (tap1 = 1.0f8, tap2 = 1.0f0, tap3 = 1.0f0, tap4 = 1.0f0, tap5 = 1.0f0, tap6 = 1.0f0) + du32_fast = LVector(X = 0.0f0, Y = 0.0f0) + du32_ref = LVector(X = 0.0f0, Y = 0.0f0) + f32(du32_fast, u32, p32, 0.0) + g32(du32_ref, u32, p32, 0.0) + @test du32_fast == du32_ref + @test du32_fast[:X] == 100000008.0f0 + + # Function-valued rates still go through `AlgebraicPetri.valueat`. + f = vectorfield_flat(simple) + u = LVector(S = 10.0, I = 2.0) + du = LVector(S = 0.0, I = 0.0) + f(du, u, (infection = (u, t) -> 0.5 * u[:S],), 1.0) + @test du[:S] == -0.5 * 10.0 * 10.0 * 2.0 + @test du[:I] == -du[:S] + + # A rate container missing a transition name is an error naming the transition, not a + # silent zero. + @test_throws ArgumentError f(du, u, (wrong = 1.0,), 0.0) +end