Kernels For SSAStepper

[sivasathyaseeelan]

using JumpProcesses, DiffEqBase, SciMLBase, Plots, CUDA
using Test, LinearAlgebra
using StableRNGs
rng = StableRNG(12345)

# SEIR-VR model jumps
function seir_vr_jumps()
    # 1. Infection: S + I → E + I
    infection_rate = (u, p, t) -> p[1] * u[1] * u[3]  # β * S * I
    infection_affect! = function (integrator)
        integrator.u[1] -= 1  # S decreases
        integrator.u[2] += 1  # E increases
        integrator.u[6] = sum(integrator.u[1:5])  # Update N
    end

    # 2. Progression: E → I
    progression_rate = (u, p, t) -> p[2] * u[2]  # σ * E
    progression_affect! = function (integrator)
        integrator.u[2] -= 1  # E decreases
        integrator.u[3] += 1  # I increases
        integrator.u[6] = sum(integrator.u[1:5])  # Update N
    end

    # 3. Recovery: I → R
    recovery_rate = (u, p, t) -> p[3] * u[3]  # γ * I
    recovery_affect! = function (integrator)
        integrator.u[3] -= 1  # I decreases
        integrator.u[4] += 1  # R increases
        integrator.u[6] = sum(integrator.u[1:5])  # Update N
    end

    # 4. Vaccination: S → V
    vaccination_rate = (u, p, t) -> p[4] * u[1]  # ν * S
    vaccination_affect! = function (integrator)
        integrator.u[1] -= 1  # S decreases
        integrator.u[5] += 1  # V increases
        integrator.u[6] = sum(integrator.u[1:5])  # Update N
    end

    # 5. Vaccine waning: V → S
    waning_rate = (u, p, t) -> p[5] * u[5]  # ω * V
    waning_affect! = function (integrator)
        integrator.u[5] -= 1  # V decreases
        integrator.u[1] += 1  # S increases
        integrator.u[6] = sum(integrator.u[1:5])  # Update N
    end

    # 6. Reinfection: V + I → E + I
    reinfection_rate = (u, p, t) -> p[6] * u[5] * u[3]  # α * V * I
    reinfection_affect! = function (integrator)
        integrator.u[5] -= 1  # V decreases
        integrator.u[2] += 1  # E increases
        integrator.u[6] = sum(integrator.u[1:5])  # Update N
    end

    # 7. Natural immunity waning: R → S
    immunity_waning_rate = (u, p, t) -> p[7] * u[4]  # μ * R
    immunity_waning_affect! = function (integrator)
        integrator.u[4] -= 1  # R decreases
        integrator.u[1] += 1  # S increases
        integrator.u[6] = sum(integrator.u[1:5])  # Update N
    end

    # 8. Natural infection: S → E
    natural_infection_rate = (u, p, t) -> p[8]  # λ (constant)
    natural_infection_affect! = function (integrator)
        integrator.u[1] -= 1  # S decreases
        integrator.u[2] += 1  # E increases
        integrator.u[6] = sum(integrator.u[1:5])  # Update N
    end

    return [
        ConstantRateJump(infection_rate, infection_affect!),
        ConstantRateJump(progression_rate, progression_affect!),
        ConstantRateJump(recovery_rate, recovery_affect!),
        ConstantRateJump(vaccination_rate, vaccination_affect!),
        ConstantRateJump(waning_rate, waning_affect!),
        ConstantRateJump(reinfection_rate, reinfection_affect!),
        ConstantRateJump(immunity_waning_rate, immunity_waning_affect!),
        ConstantRateJump(natural_infection_rate, natural_infection_affect!)
    ]
end

# Main execution
u0 = [990.0, 5.0, 5.0, 0.0, 0.0, 1000.0]  # S, E, I, R, V, N
tspan = (0.0, 50.0)
p = [0.002, 0.2, 0.25, 0.01, 0.05, 0.001, 0.02, 0.001]  # β, σ, γ, ν, ω, α, μ, λ
rng = StableRNG(12345)

prob = DiscreteProblem(u0, tspan, p)
jumps = seir_vr_jumps()
jump_prob = JumpProblem(prob, Direct(), jumps...; rng=rng)

# Solve on GPU
ensemble_prob = EnsembleProblem(jump_prob)
sol = solve(ensemble_prob, SSAStepper(), EnsembleGPUKernel(CUDABackend()), trajectories = 10)
plot(sol)

Checklist

• Appropriate tests were added
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

Add any other context about the problem here.

[ChrisRackauckas]

Make the same test in the CPU tests just using the CPU backend to kernel abstractions.

[ChrisRackauckas]

Chop this out with the new tooling.

[ChrisRackauckas]

why is this not using SSAStepper?

[isaacsas]

@sivasathyaseeelan do you plan to update per @ChrisRackauckas comments?

[sivasathyaseeelan]

@sivasathyaseeelan do you plan to update per @ChrisRackauckas comments?

I will get back after merging adaptive and implicit tau leaping