From fcb2501c0fdbc691206348f927eff637098c4276 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Wed, 6 Nov 2024 00:36:29 +0000
Subject: [PATCH] update ODEProcess benchmark

---
 benchmark/Project.toml                        |  1 +
 .../bench/EpiInfModels/InfectionODEProcess.jl | 24 -------------------
 benchmark/bench/EpiInfModels/ODEProcess.jl    | 20 ++++++++++++++++
 3 files changed, 21 insertions(+), 24 deletions(-)
 delete mode 100644 benchmark/bench/EpiInfModels/InfectionODEProcess.jl
 create mode 100644 benchmark/bench/EpiInfModels/ODEProcess.jl

diff --git a/benchmark/Project.toml b/benchmark/Project.toml
index 13be17133..0b550876f 100644
--- a/benchmark/Project.toml
+++ b/benchmark/Project.toml
@@ -3,6 +3,7 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
 EpiAware = "b2eeebe4-5992-4301-9193-7ebc9f62c855"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 PkgBenchmark = "32113eaa-f34f-5b0d-bd6c-c81e245fc73d"
 Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 TuringBenchmarking = "0db1332d-5c25-4deb-809f-459bc696f94f"
diff --git a/benchmark/bench/EpiInfModels/InfectionODEProcess.jl b/benchmark/bench/EpiInfModels/InfectionODEProcess.jl
deleted file mode 100644
index bc0c4b4b1..000000000
--- a/benchmark/bench/EpiInfModels/InfectionODEProcess.jl
+++ /dev/null
@@ -1,24 +0,0 @@
-let
-    r = log(2) / 7 # Growth rate corresponding to 7 day doubling time
-    u0 = [1.0]
-    p = [r]
-    params = ODEParams(u0 = u0, p = p)
-
-    # Define the ODE problem using SciML
-    # We use a simple exponential growth model
-
-    function expgrowth(du, u, p, t)
-        du[1] = p[1] * u[1]
-    end
-    prob = ODEProblem(expgrowth, u0, (0.0, 10.0), p)
-
-    # Define the ODEProcess
-
-    expgrowth_model = ODEProcess(prob::ODEProblem; ts = 0:1:10,
-        solver = Tsit5(),
-        sol2infs = sol -> sol[1, :])
-
-    # Generate the latent infections
-    I_t = generate_latent_infs(expgrowth_model, params)()
-    suite["ODEProcess"] = make_epiaware_suite(mdl)
-end
diff --git a/benchmark/bench/EpiInfModels/ODEProcess.jl b/benchmark/bench/EpiInfModels/ODEProcess.jl
new file mode 100644
index 000000000..58997725a
--- /dev/null
+++ b/benchmark/bench/EpiInfModels/ODEProcess.jl
@@ -0,0 +1,20 @@
+let
+    # Create an instance of SIRParams
+    sirparams = SIRParams(
+        tspan = (0.0, 30.0),
+        infectiousness = LogNormal(log(0.3), 0.05),
+        recovery_rate = LogNormal(log(0.1), 0.05),
+        initial_prop_infected = Beta(1, 99)
+    )
+
+    # Define the ODEProcess
+    sir_process = ODEProcess(
+        params = sirparams,
+        sol2infs = sol -> sol[2, :],
+        solver_options = Dict(:verbose => false, :saveat => 1.0)
+    )
+
+    # Generate the latent infections
+    mdl = generate_latent_infs(sir_process, nothing)
+    suite["ODEProcess"] = make_epiaware_suite(mdl)
+end