AlgebraicJulia / Structured-Epidemic-Modeling

Results and software for our paper on structured epidemic modeling
BSD 2-Clause "Simplified" License
12 stars 5 forks source link

Problems porting a related epidemiological model #5

Open sdwfrost opened 1 year ago

sdwfrost commented 1 year ago

I'm trying to use the examples here to port an existing high/low risk model to use AlgebraicPetri.jl. I get an error ERROR: Pullbacks of TypeSets that are not products are not supported when I try to compose the model. Can anyone spot where I'm going wrong?

using AlgebraicPetri
using Catlab
using Catlab.Graphics
using Catlab.Graphics.Graphviz
using Catlab.Graphics.Graphviz: Html
using Catlab.WiringDiagrams
using Catlab.CategoricalAlgebra
using Catlab.Programs.RelationalPrograms
using GraphViz
using LabelledArrays
using OrdinaryDiffEq
using Plots

# Functions for graphing typed Petri nets
colors = ["#a08eae","#ffeec6", "#a8dcd9", "#ffeec6", "#a8dcd9"]

function def_trans(typed_petri::ACSetTransformation, colors; labels = true)
  (p, t; pos = "") -> ("t$t", Attributes(
             :label => labels ? Html(flatten(tname(p,t))) : "" ,
             :shape=>"square",
             :color=>colors[typed_petri[:T](t)],
             :pos=>pos))
end

function def_trans(colors = colors; labels = true)
  (p, t; pos = "") -> ("t$t", Attributes(
             :label => labels ? "$(tname(p, t))" : "" ,
             :shape=>"square",
             :color=>colors[t],
             :pos=>pos))
end

flatten(tname::Symbol) = "$tname"

function flatten(tname::Tuple)
    names = split(replace(string(tname), "("=>"", ")"=>"", ":"=>""), ",")
    for i in 1:length(names)
        name = strip(names[i])
        if name[1:2] == "id"
            continue
        end
        return name
    end
    return "id"
end

def_states(p, s; pos="") = ("s$s", Attributes(
        :label => sname(p,s) isa Tuple where T ? Html(replace(string(sname(p,s)), ":"=>"", "," => "", "("=>"", ")"=>"")) : "$(sname(p,s))",
        :shape=>"circle",
        :color=>"#6C9AC3",
        :pos=>pos
))

Graph_typed(typed_petri::ACSetTransformation, colors = colors; labels = true) = Graph(dom(typed_petri),
    make_trans = def_trans(typed_petri, colors; labels = labels),
    make_states = def_states
)

### Define models

SIR = LabelledPetriNet([:S, :I, :R],
  :β => ((:S, :I)=>(:I, :I)),
  :γ => (:I=>:R),
  :∅ => (:S => :S),
  :∅ => (:I => :I),
  :∅ => (:R => :R)
)

infectious_type = LabelledPetriNet([:Pop],
  :interact=>((:Pop, :Pop)=>(:Pop, :Pop)), 
  :t_disease=>(:Pop=>:Pop),
  :t_strata=>(:Pop=>:Pop)
)

## Pull out indexing from model

s, = parts(infectious_type, :S)
t_interact, t_disease, t_strata = parts(infectious_type, :T)
i_interact1, i_interact2, i_disease, i_strata = parts(infectious_type, :I)
o_interact1, o_interact2, o_disease, o_strata = parts(infectious_type, :O)

typed_SIR = ACSetTransformation(SIR, infectious_type,
  S = [s, s, s],
  T = [t_interact, t_disease, t_strata, t_strata, t_strata],
  I = [i_interact1, i_interact2, i_disease, i_strata, i_strata, i_strata],
  O = [o_interact1, o_interact2, o_disease, o_strata, o_strata, o_strata],
  Name = name -> nothing # specify the mapping for the loose ACSet transform
)

Graph_typed(typed_SIR)

risk = LabelledPetriNet([:H, :L],
    :∅ => (:H => :H),
    :∅ => (:L => :L),
    :interact => ((:H, :H) => (:H, :H)),
    :interact => ((:L, :L) => (:L, :L)),
    :interact => ((:H, :L) => (:H, :L)),
    :interact => ((:L, :H) => (:L, :H)),
)

typed_risk = ACSetTransformation(risk, infectious_type,
  S = [s, s],
  T = [t_disease, t_disease, t_interact, t_interact, t_interact, t_interact],
  I = [i_disease, i_disease, i_interact1, i_interact2, i_interact1, i_interact2, i_interact2, i_interact1, i_interact2, i_interact1],
  O = [o_disease, o_disease, o_interact1, o_interact2, o_interact1, o_interact2, o_interact2, o_interact1, o_interact2, o_interact1],
  Name = name -> nothing # specify the mapping for the loose ACSet transform
)
[i_disease, i_disease, i_interact1, i_interact2, i_interact1, i_interact2, i_interact2, i_interact1, i_interact2, i_interact1]

Graph_typed(typed_risk)

typed_stratify(typed_model1, typed_model2) =
  compose(proj1(pullback(typed_model1, typed_model2)), typed_model1)

SIR_risk = typed_stratify(typed_SIR, typed_risk)
jpfairbanks commented 1 year ago

I think the problem here is the labels. We've made some changes in AlgebraicPetri and Catlab to better support attributes. You can cast a LabelledPetriNet to a regular PetriNet with copy_parts! and then compute the corresponding pullbacks.

@kris-brown has had some new features land in Catlab v0.15 to provide better support for computing with Attributes. He can speak to a better way to do stratification with labels now.

kris-brown commented 1 year ago

At this instant in time, I agree with James that currently the best option is to perform stratification on unlabeled Petri Nets. When v0.15 is released it will be much more natural to stratify with attributes present!

sdwfrost commented 1 year ago

Thanks @jpfairbanks @kris-brown,

I see how I can cast to an unlabelled Petri net; how do I go about manually specifying the stratification by pullbacks?

SIR = LabelledPetriNet([:S, :I, :R],
  :β => ((:S, :I)=>(:I, :I)),
  :γ => (:I=>:R),
  :∅ => (:S => :S),
  :∅ => (:I => :I),
  :∅ => (:R => :R)
)
SIR_unlabelled = PetriNet()
copy_parts!(SIR, SIR_unlabelled)

risk = LabelledPetriNet([:H, :L],
    :∅ => (:H => :H),
    :∅ => (:L => :L),
    :interact => ((:H, :H) => (:H, :H)),
    :interact => ((:L, :L) => (:L, :L)),
    :interact => ((:H, :L) => (:H, :L)),
    :interact => ((:L, :H) => (:L, :H)),
)
risk_unlabelled = PetriNet()
copy_parts!(risk, risk_unlabelled)
kris-brown commented 1 year ago

So I see now that AlgebraicPetri has added some specific code to help with this. Importantly, it only makes sense to stratify typed Petri net models, which are Petri Nets with a map into some fixed Petri Net of 'types'. E.g. infectious_ontology below says there is one type of species and three types of transitions. I attempted to adapt your example below in order to show some of the special syntax for constructing these things, but keep in mind sir, sir_refl, and res are all just LabeledPetriNet morphisms that you could have constructed manually.

using AlgebraicPetri,AlgebraicPetri.TypedPetri
using Catlab, Catlab.CategoricalAlgebra, Catlab.Programs

infectious_ontology = LabelledPetriNet(
  [:Pop],
  :infect=>((:Pop, :Pop)=>(:Pop, :Pop)),
  :disease=>(:Pop=>:Pop),
  :strata=>(:Pop=>:Pop)
)
sir_uwd = @relation () where (S::Pop, I::Pop, R::Pop) begin
    infect(S, I, I, I); disease(I, R)
end
sir = oapply_typed(infectious_ontology, sir_uwd, [:beta, :gamma]) # SIR -> Inf.Ont.
sir_refl = add_reflexives(sir, fill([:strata],3), infectious_ontology)
Graph(sir_refl |> dom) # visualize the added reflexive transitions in graphviz

risk_uwd = @relation () where (H::Pop, L::Pop) begin
    disease(H,H); disease(L,L); 
    infect(H,H,H,H); infect(H,L,H,L); infect(L,H,L,H); infect(L,L,L,L)
end
risk = oapply_typed(infectious_ontology, risk_uwd, [:dh, :dl, :hh, :hl, :lh,:ll]) # HL -> Inf.Ont.
Graph(risk |> dom)

res = typed_product(sir_refl, risk)
Graph(res |> dom)

Does this help?

sdwfrost commented 1 year ago

Hi Kris,

I ran into problems with using the AlgebraicPetri helper functions too - see https://github.com/AlgebraicJulia/AlgebraicPetri.jl/issues/152

I see you added disease(H,H) and disease(L,L) to the definition of risk_uwd compared to my code; I thought that the line on adding reflexives would do that in my example. I'm still stuck on running the vector field though, as I don't know how to label initial conditions and parameter values when the indices are of the type (:S, :H) etc.

sdwfrost commented 1 year ago

I see where I'm going wrong in the other issue; firstly, typed_product followed by add_reflexives was giving me the full product in states (which I don't want - I was getting six transitions, and not two). Secondly, there's a function that I hadn't spotted (flatten_labels) which is needed to run the model in OrdinaryDiffEq.

jpfairbanks commented 1 year ago

Thanks for being a guinea pig on this for us. You are one of the first people trying to use stratification who didn't help develop it. We can use this example in the docs once we get it working!

sdwfrost commented 1 year ago

Happy to play 'red team' @jpfairbanks! This seems to work, and mirror my multigroup ODE model. It would be nice if the risk model didn't refer to any terms in the infectious_ontology, and that these could be added later, and still be able to give a typed product with the right names. @slwu89 suggestion here does this, but the labels are off (two gamma_disease).

using AlgebraicPetri,AlgebraicPetri.TypedPetri
using Catlab, Catlab.CategoricalAlgebra, Catlab.Programs
using OrdinaryDiffEq
using LabelledArrays
using Plots

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

sir_uwd = @relation () where (S::Pop, I::Pop, R::Pop) begin
    infect(S, I, I, I)
    disease(I, R)
end
sir = oapply_typed(infectious_ontology, sir_uwd, [:beta, :gamma])
Graph(sir |> dom)

sir_refl = add_reflexives(sir, fill([:strata],3), infectious_ontology)
Graph(sir_refl |> dom)

risk_uwd = @relation () where (H::Pop, L::Pop) begin
    infect(H,H,H,H)
    infect(H,L,H,L)
    infect(L,H,L,H)
    infect(L,L,L,L)
    disease(H,H)
    disease(L,L)
end
risk = oapply_typed(infectious_ontology, risk_uwd, [:hh, :hl, :lh, :ll, :h, :l])
Graph(risk |> dom)

res = typed_product(sir_refl, risk)
Graph(res |> dom)

relabelled_res = flatten_labels(res |> dom)
Graph(relabelled_res)

K = 2
S = [495.0, 495.0]
I = [5.0, 5.0]
R = [0.0, 0.0]
N = [S[i]+I[i]+R[i] for i in 1:K]

β = 0.05
c = [20.0, 5.0]
pij = hcat([[c[j]*N[j]/sum([c[k]*N[k] for k in 1:K]) for j in 1:K] for i in 1:K]...)'
betas = β .* (c .* pij) ./ N
γ = 0.25

u0 = @LArray vec([S I R]')[:,1] Tuple(snames(relabelled_res))
p = @LArray [vec(betas); γ; γ] Tuple(tnames(relabelled_res))
tspan = (0.0, 40.0)

vf = vectorfield(relabelled_res)
prob = ODEProblem(vf, u0, tspan, p)
sol = solve(prob, Rosenbrock32())
plot(sol)
jpfairbanks commented 1 year ago

Nice! This is a great example.

It would be nice if the risk model didn't refer to any terms in the infectious_ontology

Yeah, it is very appealing to avoid that, but the nature of the stratification with pullbacks is that these types are important for figuring out which processes can interact. If compute a bigger model with product of untyped petri nets, it will have too many terms in it. Some interaction terms that are logically 0, will be free parameters in the product model. The typing information is the additional feature that makes this work correctly.

sdwfrost commented 1 year ago

I was thinking more about adding the disease(H,H) and disease(L,L) terms subsequently in the UWD. When making stratifications, you don't want to think about what other transitions from other models until it actually comes to doing a pullback. Now to try to work out how to make SDEs and jump problems - Petri.jl did the dispatch for you...

jpfairbanks commented 1 year ago

I was thinking more about adding the disease(H,H) and disease(L,L) terms subsequently in the UWD. When making stratifications, you don't want to think about what other transitions from other models until it actually comes to doing a pullback.

Yes this is a fundamental problem. The terms in the stratification model depend on an understanding of the disease model that you are going to stratify. We haven't found a good way to decouple them any more than the way you do it in this example.