Logging events when running a CTMC

Specific Domains Modelling & Simulations
1

At the bottom of this message is a sample piece of code (which should run) that illustrates the reason for my question (slash feature request in disguise): when using DifferentialEquations.jl to run a continuous time Markov chain, is there a way to log which event was triggered, not just how the state changed?

I don’t seem to find it as an option in the solver. In R, in the adaptivetau package, I would add the option reportTransitions=TRUE to a call to ssa.exact or ssa.adaptivetau, which would add a matrix to the solution with events types as columns and each row with a 1 (or more in the adaptive tau case) in the column corresponding to the event having been triggered. There’s a similar option in GillespieSSA2.

If there is no such option here, I am guessing I should be using callbacks, but I am completely baffled with those things in general. How could I go about doing that?

Now note that since infections initiated by both I and A individuals result in the same change in state (S decreases by 1, L increases by 1), I cannot just compare states from one step to the next to figure out what happened (at least not for infection events).

[And here’s my pitch for why that would be a cool feature to add, if it is not already there. My example below is a prototypical “slightly more evolved than the base” epidemiological model, which has two types of infectious individuals. Those in I show symptoms, those in A don’t. In applications, it can be important to be able to say what type of individual initiated an infecting contact, for instance, to understand how much asymptomatic infectious individuals contribute to the overall spread. I am sure there are many other situations where this type of question arises.]

using DifferentialEquations
using Plots

# Define the jumps for each transition
function ctmc_jumps_S_to_L_by_I(u)
    u[1] -= 1  # S decreases
    u[2] += 1  # L increases
end

function ctmc_jumps_S_to_L_by_A(u)
    u[1] -= 1  # S decreases
    u[2] += 1  # L increases
end

function ctmc_jumps_L_to_I(u)
    u[2] -= 1  # L decreases
    u[3] += 1  # I increases
end

function ctmc_jumps_L_to_A(u)
    u[2] -= 1  # L decreases
    u[4] += 1  # A increases
end

function ctmc_jumps_I_to_R(u)
    u[3] -= 1  # I decreases
    u[5] += 1  # R increases
end

function ctmc_jumps_A_to_R(u)
    u[4] -= 1  # A decreases
    u[5] += 1  # R increases
end

# Initial conditions
u0 = [990, 10, 0, 0, 0]  # Initial population: S=990, L=10, I=0, A=0, R=0

# Parameters: β (infection coefficient), p (fraction of infections that are symptomatic), σ (incubation rate), γ (recovery rate), N (total population)
p = (0.3, 0.5, 0.2, 0.1, 1000)

# Time span
tspan = (0.0, 100.0)

# Define the problem
prob = DiscreteProblem(u0, tspan, p)

# Define the rates as functions
rate_S_to_L_by_I = (u, p, t) -> p[1] * u[1] * u[3] / p[4]   # β * S * I / N
rate_S_to_L_by_A = (u, p, t) -> p[1] * u[1] * u[4] / p[4]  # β * S * A / N
rate_L_to_I = (u, p, t) -> p[2] * p[3] * u[2]                      # p * σ * L
rate_L_to_A = (u, p, t) -> (1-p[2]) * p[3] * u[2]               # (1-p) * σ * L
rate_I_to_R = (u, p, t) -> p[4] * u[3]                              # γ * I
rate_A_to_R = (u, p, t) -> p[4] * u[4]                             # γ * A

# Define the jumps
jump_S_to_L_by_I = ConstantRateJump(rate_S_to_L_by_I, ctmc_jumps_S_to_L_by_I)
jump_S_to_L_by_A = ConstantRateJump(rate_S_to_L_by_A, ctmc_jumps_S_to_L_by_A)
jump_L_to_I = ConstantRateJump(rate_L_to_I, ctmc_jumps_L_to_I)
jump_L_to_A = ConstantRateJump(rate_L_to_A, ctmc_jumps_L_to_A)
jump_I_to_R = ConstantRateJump(rate_I_to_R, ctmc_jumps_I_to_R)
jump_A_to_R = ConstantRateJump(rate_A_to_R, ctmc_jumps_A_to_R)

jumps_ordered = [
    jump_S_to_L_by_I,
    jump_S_to_L_by_A,
    jump_L_to_I,
    jump_L_to_A,
    jump_I_to_R,
    jump_A_to_R
]

jump_prob = JumpProblem(prob, Direct(), jumps_ordered...)

sol = solve(jump_prob, SSAStepper())

# Plot the solution
plot(sol, label=["S" "L" "I" "A" "R"], xlabel="Time", ylabel="Population", lw=2)
2

I haven’t tried this, but I think you could add a 6th entry to your state vector (u0=[990, 10, 0, 0, 0,0]) and set it to a different value for each transition function.

for example

function ctmc_jumps_S_to_L_by_I(u)
    u[1] -= 1  # S decreases
    u[2] += 1  # L increases
    u[6] = 1
end

function ctmc_jumps_S_to_L_by_A(u)
    u[1] -= 1  # S decreases
    u[2] += 1  # L increases
    u[6] = 2
end

Not as nice as something built-in, but worth a try.

3

Ha, cool idea. I will try and report back.

4

Works fine! So I’ll write your suggestion as the solution, even though I agree, if anybody developing this were to stumble here, a proper mechanism for doing this would be great.

Including my revised code (sans plot) in case someone needs it.

using DifferentialEquations
using Plots

# Define the jumps for each transition
function ctmc_jumps_S_to_L_by_I(u)
    u[1] -= 1  # S decreases
    u[2] += 1  # L increases
    u[6] = 1   # Record jump index
end

function ctmc_jumps_S_to_L_by_A(u)
    u[1] -= 1  # S decreases
    u[2] += 1  # L increases
    u[6] = 2   # Record jump index
end

function ctmc_jumps_L_to_I(u)
    u[2] -= 1  # L decreases
    u[3] += 1  # I increases
    u[6] = 3   # Record jump index
end

function ctmc_jumps_L_to_A(u)
    u[2] -= 1  # L decreases
    u[4] += 1  # A increases
    u[6] = 4   # Record jump index
end

function ctmc_jumps_I_to_R(u)
    u[3] -= 1  # I decreases
    u[5] += 1  # R increases
    u[6] = 5   # Record jump index
end

function ctmc_jumps_A_to_R(u)
    u[4] -= 1  # A decreases
    u[5] += 1  # R increases
    u[6] = 6   # Record jump index
end

# Initial conditions
u0 = [990, 10, 0, 0, 0, 0]  # Initial population: S=990, L=10, I=0, A=0, R=0, jump index=0

# Parameters: β (infection coefficient), p (fraction of infections that are symptomatic), σ (incubation rate), γ (recovery rate), N (total population)
p = (0.3, 0.5, 0.2, 0.1, 1000)

# Time span
tspan = (0.0, 100.0)

# Define the problem
prob = DiscreteProblem(u0, tspan, p)

# Define the rates as functions
rate_S_to_L_by_I = (u, p, t) -> p[1] * u[1] * u[3] / p[4]  # β * S * I / N
rate_S_to_L_by_A = (u, p, t) -> p[1] * u[1] * u[4] / p[4]  # β * S * A / N
rate_L_to_I = (u, p, t) -> p[2] * p[3] * u[2]               # p * σ * L
rate_L_to_A = (u, p, t) -> (1-p[2]) * p[3] * u[2]           # (1-p) * σ * L
rate_I_to_R = (u, p, t) -> p[4] * u[3]                      # γ * I
rate_A_to_R = (u, p, t) -> p[4] * u[4]                      # γ * A

# Define the jumps
jump_S_to_L_by_I = ConstantRateJump(rate_S_to_L_by_I, ctmc_jumps_S_to_L_by_I)
jump_S_to_L_by_A = ConstantRateJump(rate_S_to_L_by_A, ctmc_jumps_S_to_L_by_A)
jump_L_to_I = ConstantRateJump(rate_L_to_I, ctmc_jumps_L_to_I)
jump_L_to_A = ConstantRateJump(rate_L_to_A, ctmc_jumps_L_to_A)
jump_I_to_R = ConstantRateJump(rate_I_to_R, ctmc_jumps_I_to_R)
jump_A_to_R = ConstantRateJump(rate_A_to_R, ctmc_jumps_A_to_R)

jumps_ordered = [
    jump_S_to_L_by_I,
    jump_S_to_L_by_A,
    jump_L_to_I,
    jump_L_to_A,
    jump_I_to_R,
    jump_A_to_R
]

jump_prob = JumpProblem(prob, Direct(), jumps_ordered...)

sol = solve(jump_prob, SSAStepper())

# Prepare event type names in order
event_types = [
    "Infection by I (S → L)",
    "Infection by A (S → L)",
    "Incubation (L → I)",
    "Incubation (L → A)",
    "Recovery (I → R)",
    "Recovery (A → R)"
]

# Extract jump indices for each time point
jump_indices = [u[6] for u in sol.u]

# Build the matrix: time, event number, event type
event_matrix = [ (t, idx, idx > 0 && idx <= length(event_types) ? event_types[Int(idx)] : "None")
                 for (t, idx) in zip(sol.t, jump_indices) ]

# Convert to a matrix of Any for display or further processing
event_matrix = reduce(vcat, [event_matrix...])

# Example: print the first 10 events
for row in event_matrix[1:10, :]
    println(row)
end