Issues when running stratified model

[sdwfrost]

Hi @mehalter!

I'm trying to implement this simple multigroup model using AlgebraicPetri, and (probably due to my lack of understanding), have run into issues.

1. Firstly, when I use typed_product, my recovery rate disappears. How do I keep it in my stratified model?
2. Secondly, how can I set/access the rates and concentrations in a combined model? As is, concentrations and rates throws an error, as the keys are tuples rather than symbols.

using AlgebraicPetri
using AlgebraicPetri.TypedPetri
using Catlab.Programs
using Catlab.Graphics
using Catlab.CategoricalAlgebra
using Catlab.WiringDiagrams
using GraphViz
using LabelledArrays
using OrdinaryDiffEq
using Plots

display_uwd(ex) = to_graphviz(ex, box_labels=:name, junction_labels=:variable, edge_attrs=Dict(:len=>".75"));

infectious_ontology = LabelledPetriNet(
  [:Pop],
  :infect => ((:Pop, :Pop)=>(:Pop, :Pop)), 
  :disease => (:Pop => :Pop),
  :strata => (:Pop => :Pop)
)

sir_relation = @relation () where {(S::Pop, I::Pop, R::Pop)} begin
    infect(S, I, I, I)
    disease(I, R)
end
display_uwd(sir_relation)

sir_model = oapply_typed(infectious_ontology, sir_relation, [:infection, :recovery])
sir_model = add_params(sir_model,
                       Dict(:S => 0.99, :I => 0.01, :R => 0.0),
                       Dict(:infection => 0.5, :recovery => 0.25))
sir_pn = dom(sir_model)
sir_prob = ODEProblem(vectorfield(sir_pn), concentrations(sir_pn), (0.0, 40.0), rates(sir_pn))
sir_sol = solve(sir_prob, Rosenbrock32())
plot(sir_sol)

risk_relation = @relation () where {(H::Pop, L::Pop)} begin
  infect(H, H, H, H)
  infect(L, L, L, L)
  infect(H, L, H, L)
  infect(L, H, L, H)
end
risk_model = oapply_typed(infectious_ontology, risk_relation, [:infect_hh, :infect_ll, :infect_hl, :infect_lh])
risk_model = add_params(risk_model,
                        Dict(:H => 500.0, :L => 500.0),
                        Dict(:infect_hh => 1.0, :infect_ll => 1.0, :infect_hl => 1.0, :infect_lh => 1.0))

sir_model = add_reflexives(sir_model, [[:strata], [:strata], [:strata]], infectious_ontology)
risk_model = add_reflexives(risk_model, [[:strata], [:strata]], infectious_ontology)
sir_risk = typed_product(sir_model, risk_model)
sir_risk_pn = dom(sir_risk)
sir_risk_pn

[slwu89]

@sdwfrost is the closer to what you are expecting? I haven't tried to do the stratification with parameters yet though.

infectious_ontology = LabelledPetriNet(
  [:Pop],
  :infect => ((:Pop, :Pop)=>(:Pop, :Pop)), 
  :disease => (:Pop => :Pop),
  :strata => (:Pop => :Pop)
)

sir_relation = @relation () where {(S::Pop, I::Pop, R::Pop)} begin
    infect(S, I, I, I)
    disease(I, R)
end

sir_model = oapply_typed(infectious_ontology, sir_relation, [:infection, :recovery])

typed_risk = pairwise_id_typed_petri(infectious_ontology, :Pop, :infect,
                                    [:L, :H])

typed_risk_aug = add_reflexives(
  typed_risk,
  repeat([[:disease]], 2),
  infectious_ontology
)

stratified = typed_product(typed_risk_aug, sir_model)
dom(stratified)

It produces the following PN:

julia> dom(stratified)
AlgebraicPetri.LabelledPetriNetUntyped{Tuple{Symbol, Symbol}} with elements T = 1:6, S = 1:6, I = 1:10, O = 1:10
┌───┬───────────────────────┐
│ T │                 tname │
├───┼───────────────────────┤
│ 1 │    (:L_L, :infection) │
│ 2 │    (:L_H, :infection) │
│ 3 │    (:H_L, :infection) │
│ 4 │    (:H_H, :infection) │
│ 5 │ (:disease, :recovery) │
│ 6 │ (:disease, :recovery) │
└───┴───────────────────────┘
┌───┬──────────┐
│ S │    sname │
├───┼──────────┤
│ 1 │ (:L, :S) │
│ 2 │ (:H, :S) │
│ 3 │ (:L, :I) │
│ 4 │ (:H, :I) │
│ 5 │ (:L, :R) │
│ 6 │ (:H, :R) │
└───┴──────────┘

It looks like the stratification API is new, I'm not 100% clear on what all the arguments for pairwise_id_typed_petri are doing. I'd also be interested to hear what @mehalter suggests is the best approach.