Compositional epidemiological modeling using structured cospans

Structured cospans

Structured cospans are a category-theoretic machine for turning closed systems into open ones, which interact with other systems along a designated boundary. Open systems can be composed at their boundaries to form larger systems. Structured cospans were invented by John Baez and Kenny Courser (Baez and Courser 2019; Courser 2020), building on Brendan Fong’s decorated cospans (Fong 2015). Baez has also written a blog post introduction to structured cospans at the n-Category Cafe.

The setting for structured cospans involves two categories, a category \mathsf{X} inhabited by the original, closed systems and a category \mathsf{A} describing the boundary of the open systems. The category \mathsf{A} is usually interpreted as a substructure of \mathsf{X}. For example, \mathsf{X} could be the category of graphs and \mathsf{A} the category \mathsf{FinSet} of finite sets, interpreted as sets of vertices. The connection between two categories is formally established by a functor L: \mathsf{A} \to \mathsf{X}. In the case of graphs, L: \mathsf{FinSet}\to \mathsf{Graph} is the discrete graph functor, assigning to any finite set V the graph having vertices V and no edges.

For a given functor L: \mathsf{A} \to \mathsf{X}, a structured cospan consists of a pair of objects a,b in \mathsf{A} together with a cospan in X of the form:

The structured cospan represents an open system. The apex x of the cospan is the system itself, while the left leg gives the “inputs” and the right leg gives the “outputs.”¹ When L is the discrete graph functor, the left foot L(a) is a discrete graph on the input vertices a and the right foot is a discrete graph L(b) on the output vertices b.

In practice, the functor L: \mathsf{A} \to \mathsf{X} usually has a right adjoint R: \mathsf{X} \to \mathsf{A}. Although not strictly necessary, it is technically convenient to assume that the right adjoint exists and some authors do so (Cicala 2020). In our implementation, we also assume that the adjoint exists. Recall that L \dashv R being adjoint functors means that there is a natural bijection between morphisms L(a) \to x in \mathsf{X} and morphisms a \to R(x) in \mathsf{A}. The right adjoints that show up in structured cospans are usually forgetful functors. For example, the right adjoint to the discrete graph functor L: \mathsf{FinSet}\to \mathsf{Graph} is the functor R: \mathsf{Graph}\to \mathsf{FinSet} that sends a graph to its vertex set. Indeed, graph homomorphisms L(a) \to x out of a discrete graph naturally correspond to functions a \to R(x) into a graph’s vertex set.

Adjointness allows us to pass freely between two different forms of a structured cospan. Let us call the defining data, namely objects a,b in \mathsf{A} and a cospan L(a) \rightarrow x \leftarrow L(b) in \mathsf{X}, the L-form of the structured cospan. By adjointness, this data is equivalent to an object x in \mathsf{X} together with a cospan

in \mathsf{A}. Call this new data a structured cospan in R-form. Readers familiar with decorated cospans will recognize the R-form of a structured cospan as being very similar to a decorated cospan: a cospan in the base category \mathsf{A}, often just \mathsf{FinSet}, together with a decoration x of its apex in a richer category \mathsf{X}.

Both forms of a structured cospan are useful. The L-form is important because composition happens inside the category \mathsf{X}. Specifically, the composite of structured cospans L(a) \rightarrow x \leftarrow L(b) and L(b) \rightarrow y \leftarrow L(c) is given by the following pushout in \mathsf{X}, which is assumed to exist.

As we will see, the R-form is useful for specifying structured cospans, because it is more convenient to write down a cospan in a category like \mathsf{FinSet} than a category like \mathsf{Graph}.

Structured cospans in Catlab

We recently implemented structured cospans in Catlab.jl, comprising a generic interface plus special support for \mathsf{C}-sets. Given any schema \mathsf{C}, structured cospans in the category \mathsf{X} := [\mathsf{C},\mathsf{Set}] of \mathsf{C}-sets can be constructed and then composed in just a few lines of code. The category \mathsf{A}, containing the boundaries of the \mathsf{C}-sets, is defined by choosing an object of \mathsf{C}.² Let’s see how this works for open graphs.

As we have seen previously, graphs are \mathsf{C}-sets on the category \mathsf{C} = \mathsf{Sch}(\mathsf{Graph}) = \{E \rightrightarrows V\} with two parallel arrows. Choosing the object V \in \mathsf{Sch}(\mathsf{Graph}) representing the vertices is enough to specify the left adjoint L: \mathsf{FinSet}\to \mathsf{Graph} discussed above. In code, given the Graph type defined in the Catlab standard library, we can define new types OpenGraph, which are structured cospans of graphs, and OpenGraphOb, which are the objects of this category.

using Catlab.CategoricalAlgebra, Catlab.Graphs, Catlab.Graphics

const OpenGraphOb, OpenGraph = OpenCSetTypes(Graph, :V)

Now let’s construct an open graph.

x = Graph()
add_vertices!(x, 4)
add_edges!(x, [1,1,2,2], [2,2,3,4])
f = OpenGraph(x, FinFunction([1],4), FinFunction([3,4],4))

to_graphviz(apex(f), node_labels=true)

The feet of this structured cospan, viewed as objects in \mathsf{A} = \mathsf{FinSet}, are the finite sets [1] := \{1\} and [2] := \{1,2\}.

feet(f)

2-element StaticArraysCore.SVector{2, Catlab.CategoricalAlgebra.FinSets.FinSetInt} with indices SOneTo(2):
 FinSet(1)
 FinSet(2)

The legs of the cospan, which are internally stored in L-form, are graph homomorphisms out of the discrete graphs on 1 and 2 vertices.

map(components, legs(f))

2-element StaticArraysCore.SVector{2, @NamedTuple{V::Catlab.CategoricalAlgebra.FinSets.FinDomFunctionVector{Int64, Vector{Int64}, Catlab.CategoricalAlgebra.FinSets.FinSetInt}, E::Catlab.CategoricalAlgebra.FinSets.FinDomFunctionVector{Int64, UnitRange{Int64}, Catlab.CategoricalAlgebra.FinSets.FinSetInt}}} with indices SOneTo(2):
 (V = FinFunction([1], 1, 4), E = FinFunction(1:0, 0, 4))
 (V = FinFunction([3, 4], 2, 4), E = FinFunction(1:0, 0, 4))

As the output above shows, the left leg has vertex map [1] \to [4],\ 1 \mapsto 1, identifying vertex 1 as the input, while the right leg has vertex map [2] \to [4],\ (1,2) \mapsto (3,4), identifying vertices 3 and 4 as outputs.

If we construct another open graph

y = Graph()
add_vertices!(y, 4)
add_edges!(y, [1,2,3], [3,3,4])
g = OpenGraph(y, FinFunction([1,2],4), FinFunction([4],4))

to_graphviz(apex(g), node_labels=true)

we can compose it with the first one to construct a larger graph.

using Catlab.Theories

h = compose(f, g)
to_graphviz(apex(h), node_labels=true)

We can also place the graphs in parallel by taking their monoidal product.

h = f ⊗ g
to_graphviz(apex(h), node_labels=true)

The implementation takes advantage of Catlab’s support for limits and colimits of \mathsf{C}-sets. Composites of structured cospans are computed by pushouts and monoidal products by coproducts.

Petri nets and epidemiological models

Structured cospans of \mathsf{C}-sets are general machinery that can be used to build compositional modeling tools. Specifically, they provide a simple interface for building domain-specific modeling frameworks that support hierarchically defined models with multiple levels of abtraction. To illustrate this paradigm, we build a Petri net modeling framework and use it to construct complex epidemiology models.

A Petri net is a compartmental model defined by a bipartite graph with two types of nodes. There are state nodes, drawn as circles, that contain some concentration of objects, and transition nodes, drawn as boxes, that decrease the concentration of incoming state nodes and increase the concentration of outgoing state nodes.

To fit in with \mathsf{C}-set formalism, we work with whole-grain Petri nets, a variant of Petri nets introduced by Joachim Kock (Kock 2020). Since we need data attributes like initial concentrations and transition rates, we will actually use attributed \mathsf{C}-sets, as described in a previous blog post. The schema \mathsf{C} is shown in the following diagram, where S and T represent the sets of states and transitions, I is the set of input edges from some state s \in S to some transition t \in T, and O is the set of output edges from some transition t \in T to some state s \in S. The attributes on S to \texttt{Str} and \mathbb{N} represent labels and initial concentrations, respectively, and similarly the attributes on T to \texttt{Str} and \mathbb{R}_+ represent labels and transition rates.

With this setup, we can extend the (whole-grain) Petri nets to open Petri nets (Baez and Master 2020), where we will glue along states of the Petri nets:

const OpenLabelledReactionNetOb, OpenLabelledReactionNet =
  OpenACSetTypes(LabelledReactionNet, :S)

These definitions are contained in the package AlgebraicPetri.jl, which we will use in the remainder of the post.

We are going to build epidemiology models where the states represent stages of a disease, the tokens in the states represent people, and the transitions represent people moving between stages of the disease. First, we import AlgebraicPetri and define a few type constants that enforce the constraint that rates are real numbers. The state populations must be whole numbers since it does not make sense to have fractions of people.

using AlgebraicPetri

const EpiRxnNet = LabelledReactionNet{Number,Int}
const OpenEpiRxnNet = OpenLabelledReactionNet{Number,Int}
const OpenEpiRxnNetOb = OpenLabelledReactionNetOb{Number,Int}

AlgebraicPetri provides a straightforward API for defining Petri nets. Here, the first argument is a tuple states and their initial populations, and each following argument specifies a transition in the network with a given label and rate. In this way, we define the classic SIR model of infectious diseases:

sir=EpiRxnNet((:S=>100, :I=>1, :R=>0),
              (:inf,.03)=>((:S,:I)=>(:I,:I)), (:rec,.25)=>(:I=>:R))

To provide a more abstract, hierarchical approach to model building, we can think of epidemiology models as being made up of two basic operations. First, there can be “spontaneous” changes, such as when an infected person recovers from a disease. Second, there are “exposure” transitions where one state causes another state to change, such as when an infected person exposes a susceptible person. We define helper functions to create Petri nets that express these two patterns of interaction.

using Catlab

ob(x::Symbol,xn::Int) = codom(Open([x], EpiRxnNet(x=>xn), [x]))
function spontaneous_petri(transition::Symbol, rate::Number,
                           s::Symbol, s₀::Int,
                           t::Symbol, t₀::Int)
  Open([s], EpiRxnNet((s=>s₀,t=>t₀), (transition,rate)=>(s=>t)), [t])
end
function exposure_petri(transition::Symbol, rate::Number,
                        s::Symbol, s₀::Int,
                        e::Symbol, e₀::Int,
                        t::Symbol, t₀::Int)
  Open([s, e], EpiRxnNet((s=>s₀,e=>e₀,t=>t₀), (transition,rate)=>((s,e)=>(t,e))), [t])
end

An example of spontaneous change is recovery, where an infected person recovers from an illness and moves to the recovered state.

recover = spontaneous_petri(:rec, .25, :I, 1, :R, 0)

An example of the exposure phenomena is when an infected person exposes a susceptible person.

expose = exposure_petri(:exp, .1, :S, 100, :I, 1, :E, 1)

We will focus on building part of the COEXIST model that the UK has been using to make policy decisions about COVID-19. To begin, we need to make a simple presentation with the appropriate objects and morphisms. Each of the morphims in the presentation specifies either a spontaneous transition or an exposure transition as described above. The presentation defines a “model space” in which we can build models that contain only the provided interactions.

using Catlab
using Catlab.Theories
using AlgebraicPetri

@present Epidemiology(FreeBiproductCategory) begin
  (S, E, A, I, I2, R, R2, D)::Ob

  exposure_e::Hom(S⊗E,E)
  exposure_a::Hom(S⊗A,E)
  exposure_i::Hom(S⊗I,E)
  exposure_i2::Hom(S⊗I2,E)
  illness::Hom(E,I)
  illness_progression::Hom(I,I2)
  asymptomatic_illness::Hom(E,A)
  asymptomatic_recovery::Hom(A,R)
  illness_recovery::Hom(I2,R)
  recovery_progression::Hom(R,R2)
  death::Hom(I2,D)
end

Having defined this presentation, we can use Catlab’s @program macro to easily build up a complex model using a Julia-esque syntax. Here, the input variables are objects in our presentation and the “function calls” represent applying morphisms to those objects.

using Catlab.Programs

coexist = @program Epidemiology (s::S, e::E, i::I, i2::I2, a::A, r::R, r2::R2, d::D) begin
    e_2 = exposure_e(s, e)
    e_3 = exposure_a(s, a)
    e_4 = exposure_i(s, i)
    e_5 = exposure_i2(s, i2)
    e_all = [e, e_2, e_3, e_4, e_5]
    a_2 = asymptomatic_illness(e_all)
    a_all = [a, a_2]
    r_2 = asymptomatic_recovery(a_all)
    i_2 = illness(e_all)
    i_all = [i, i_2]
    i2_2 = illness_progression(i)
    i2_all = [i2, i2_2]
    d_2 = death(i2_all)
    r_3 = illness_recovery(i2_all)
    r_all = [r, r_2, r_3]
    r2_2 = recovery_progression(r_all)
    r2_all = [r2, r2_2]
    d_all = [d, d_2]
    return s, e_all, i_all, i2_all, a_all, r_all, r2_all, d_all
end

By defining the morphisms as open \mathsf{C}-sets, we can use structured cospans to functorially construct the Petri net that corresponds to the model defined in the @program macro.

So far the presentation assumes that all the people involved in the epidemic are well-mixed. It cannot handle geographically or demographically heterogeneous populations. What if we want to extend the model to two different populations that interact through cross exposure? For example, suppose we have two age groups with different rates of infection and recovery, where the two generations interact through cross exposure.

We can address this problem by defining a new presentation that represents the cross exposure of two separate sets of populations. Then, using the @program macro, we define a new model for cross exposure.

@present EpiCrossExposure(FreeBiproductCategory) begin
    (S, E, A, I, I2, R, R2, D)::Ob
    (S′, E′, A′, I′, I2′, R′, R2′, D′)::Ob

    exposure_i::Hom(S⊗I′,E)
    exposure_e::Hom(S⊗E′,E)
    exposure_a::Hom(S⊗A′,E)
    exposure_i2::Hom(S⊗I2′,E)
    exposure_i′::Hom(S′⊗I,E′)
    exposure_e′::Hom(S′⊗E,E′)
    exposure_a′::Hom(S′⊗A,E′)
    exposure_i2′::Hom(S′⊗I2,E′)
end;

crossexposure = @program EpiCrossExposure (s::S, e::E, i::I, i2::I2, a::A, r::R, r2::R2, d::D,
                                           s′::S′, e′::E′, i′::I′, i2′::I2′, a′::A′, r′::R′, r2′::R2′, d′::D′) begin
    e_2 = exposure_i(s, i′)
    e_3 = exposure_i2(s, i2′)
    e_4 = exposure_a(s, a′)
    e_5 = exposure_e(s, e′)
    e_all = [e, e_2, e_3, e_4, e_5]
    e′_2 = exposure_i′(s′, i)
    e′_3 = exposure_i2′(s′, i2)
    e′_4 = exposure_a′(s′, a)
    e′_5 = exposure_e′(s′, e_all)
    e′_all = [e′, e′_2, e′_3, e′_4, e′_5]
    return s, e_all, i, i2, a, r, r2, d,
           s′, e′_all, i′, i2′, a′, r′, r2′, d′
end

Having defined COEXIST and cross exposure models, we can use both as new Petri net building blocks and construct a larger model of two populations interacting through cross exposure, where each goes to their own COEXIST model with unique rates and populations. This method allows us to quickly create complex models, while the hierarchical construction also makes them easier to understand.

@present DualCoexist(FreeBiproductCategory) begin
    (Pop1,Pop2)::Ob

    crossexp::Hom(Pop1⊗Pop2,Pop1⊗Pop2)
    coex1::Hom(Pop1,Pop1)
    coex2::Hom(Pop2,Pop2)
end

dualCoexist = @program DualCoexist (pop1::Pop1, pop2::Pop2) begin
    pop1′, pop2′ = crossexp(pop1, pop2)
    return coex1(pop1′), coex2(pop2′)
end

In this wiring diagram, the boxes no longer represent our basic spontaneous or exposure Petri nets, but the full Petri nets shown above for cross exposure and the COEXIST epidemiology model. Thanks to structured cospans, we can convert this wiring diagram into its corresponding Petri net. We can see that even with just two populations, the number of states and transitions grows rapidly. Without the ability to define our models hierarchically and with user-defined abstractions, a Petri net of this complexity would be difficult to reliably encode by hand.

Since we kept track of all the initial populations and rates during the construction of the model, we can now directly generate the appropriate ordinary differential equations as input to DifferentialEquations.jl, run the simulation, and plot the results.

WARNING: using PlotMeasures.h in module Notebook conflicts with an existing identifier.

The plot forecasts the expected rate of COVID-19 infection over time. Looking at the plot, we can see that in the absence of any interventions, there is an extremely high rate of infection. Our construction of the model makes it easy to adjust the rates and parameters in response to changes in intervention policy. For instance, we can decrease the rates of exposure to simulate putting in place a stay-at-home order and promoting social distancing. We can then generate a new simulation, compute a solution, and compare the results.

In this new system, we see that the simulated interventions decrease the rates of both infection and transmission and also shift the time of peak infection.

This style of rapid modeling can enable policy makers to quickly build new models during an emerging pandemic, simulate several different intervention policy scenarios, and make policy decisions informed by the results of these simulations.

The complete code reproducing these simulations is available in the AlgebraicPetri.jl documentation.