From 14d5b8b7daf0589a06ee124ed9124da257e9f0c3 Mon Sep 17 00:00:00 2001
From: cbernalz <cbernalz@uci.edu>
Date: Tue, 23 Jul 2024 22:53:09 -0700
Subject: [PATCH] 2024-07-23 update : added agent simulation and edits.

---
 .github/dependabot.yml                        |  7 --
 Project.toml                                  |  1 +
 docs/src/tutorial.md                          | 42 ++++++--
 src/UCIWWEIHR.jl                              | 13 ++-
 src/generalizedtdist.jl                       | 28 ++++++
 src/generate_simulation_data_agent.jl         | 95 +++++++++++++++++++
 ... => generate_simulation_data_uciwweihr.jl} | 11 ++-
 ...tion_functions.jl => negativebinomial2.jl} | 31 +-----
 src/sum_values.jl                             | 11 ---
 9 files changed, 173 insertions(+), 66 deletions(-)
 delete mode 100644 .github/dependabot.yml
 create mode 100644 src/generalizedtdist.jl
 create mode 100644 src/generate_simulation_data_agent.jl
 rename src/{generate_simulation_data_ww_eihr.jl => generate_simulation_data_uciwweihr.jl} (89%)
 rename src/{distribution_functions.jl => negativebinomial2.jl} (55%)
 delete mode 100644 src/sum_values.jl

diff --git a/.github/dependabot.yml b/.github/dependabot.yml
deleted file mode 100644
index 700707c..0000000
--- a/.github/dependabot.yml
+++ /dev/null
@@ -1,7 +0,0 @@
-# https://docs.github.com/github/administering-a-repository/configuration-options-for-dependency-updates
-version: 2
-updates:
-  - package-ecosystem: "github-actions"
-    directory: "/" # Location of package manifests
-    schedule:
-      interval: "weekly"
diff --git a/Project.toml b/Project.toml
index 896136c..71af86f 100644
--- a/Project.toml
+++ b/Project.toml
@@ -20,6 +20,7 @@ Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [compat]
diff --git a/docs/src/tutorial.md b/docs/src/tutorial.md
index ac0885a..6f17c0a 100644
--- a/docs/src/tutorial.md
+++ b/docs/src/tutorial.md
@@ -6,9 +6,9 @@ Plots.reset_defaults()
 
 Welcome to the Tutorials page for the UCIWWEIHR.jl project. This section provides step-by-step guides and examples to help you get started with using the package and understanding its features.
 
-## 1. Getting Started
+## 1. Getting Started.
 
-### 1.1 Installation
+### 1.1 Installation.
 
 To install the UCIWWEIHR.jl package, open the Julia REPL and run the following command:
 
@@ -17,21 +17,23 @@ using Pkg
 Pkg.add("UCIWWEIHR")
 ```
 
-## 2. Generating simulated data
+## 2. Generating simulated data with UCIWWEIHR ODE compartmental based model.
 
-This package provides a way to also simulate data using the model specified in the future paper.  The function called `generate_simulation_data_ww_eihr` can be used to generate synthetic data for a given number of samples and features.  Here we provide a demonstration using the default settings of `generate_simulation_data_ww_eihr`:
+This package provides a way to also simulate data using the UCIWWEIHR ODE compartmental based model specified in the future paper.  The function called `generate_simulation_data_uciwweihr.jl` can be used to generate synthetic data for a given number of samples and features.  Here we provide a demonstration using the default settings of `generate_simulation_data_uciwweihr.jl` :
 
 ``` @example tutorial
 using UCIWWEIHR
 using Plots
 # Running simulation function with defaults
-df = generate_simulation_data_ww_eihr()
+df = generate_simulation_data_uciwweihr()
 first(df, 5)
 ```
 
+### 2.1. Visualizing UCIWWEIHR model results.
+
 Here we can make simple plots to visualize the data generated using the [Plots.jl](https://github.com/JuliaPlots/Plots.jl) package.
 
-### 2.1. Concentration of Pathogen Genome in WW
+### 2.1.1. Concentration of pathogen genome in wastewater(WW).
 ```@example tutorial
 plot(df.obstimes, df.log_ww_conc,
     label=nothing,
@@ -40,7 +42,7 @@ plot(df.obstimes, df.log_ww_conc,
     title="Plot of Conc. of Pathogen Genome in WW Over Time")
 ```
 
-### 2.2. Hospitalizations
+### 2.1.2. Hospitalizations.
 ```@example tutorial
 plot(df.obstimes, df.hosp, 
     label=nothing,
@@ -49,7 +51,7 @@ plot(df.obstimes, df.hosp,
     title="Plot of Hosp Over Time")
 ```
 
-### 2.3. Reproductive Number
+### 2.1.3. Reproductive number.
 ```@example tutorial
 plot(df.obstimes, df.rt, 
     label=nothing,
@@ -57,3 +59,27 @@ plot(df.obstimes, df.rt,
     ylabel="RT", 
     title="Plot of RT Over Time")
 ```
+
+## 3. Generating simulated data with an agent based model.
+
+This package provides a way to also simulate data using the agent based model in the future paper.  The function called `generate_simulation_data_agent.jl` can be used to generate synthetic data for a given population size and features.  Here we provide a demonstration using the default settings of `generate_simulation_data_agent.jl` :
+
+``` @example tutorial
+using UCIWWEIHR
+using Plots
+# Running simulation function with defaults
+df = generate_simulation_data_agent()
+first(df, 5)
+```
+
+### 3.1. Visualizing SEIHR compartments.
+
+We can also use the [Plots.jl](https://github.com/JuliaPlots/Plots.jl) package to visualize the data generated.
+
+```@example tutorial
+plot(df.Time, df.S, label = "Suseptible", xlabel = "Time", ylabel = "Number of Individuals", title = "Agent Based Model Simulation Results")
+plot!(df.Time, df.E, label = "Exposed")
+plot!(df.Time, df.I, label = "Infected")
+plot!(df.Time, df.H, label = "Hospitalized")
+plot!(df.Time, df.R, label = "Recovered")
+```
diff --git a/src/UCIWWEIHR.jl b/src/UCIWWEIHR.jl
index 26c6c3d..3bf48b1 100644
--- a/src/UCIWWEIHR.jl
+++ b/src/UCIWWEIHR.jl
@@ -18,15 +18,18 @@ using CSV
 using DataFrames
 using DifferentialEquations
 using Plots
+using StatsBase
 
-include("sum_values.jl")
-include("generate_simulation_data_ww_eihr.jl")
+include("generate_simulation_data_uciwweihr.jl")
+include("generate_simulation_data_agent.jl")
 include("bayes_eihr_model.jl")
-include("distribution_functions.jl")
 include("eihr_ode.jl")
+include("negativebinomial2.jl")
+include("generalizedtdist.jl")
+
 export eihr_ode
-export sum_values
-export generate_simulation_data_ww_eihr
+export generate_simulation_data_uciwweihr
+export generate_simulation_data_agent
 export bayes_eihr_model
 export NegativeBinomial2
 export GeneralizedTDist
diff --git a/src/generalizedtdist.jl b/src/generalizedtdist.jl
new file mode 100644
index 0000000..7b91784
--- /dev/null
+++ b/src/generalizedtdist.jl
@@ -0,0 +1,28 @@
+# credit to Damon Bayer for all of these helper functions
+"""
+Create a generalized t distribution with location and scale.
+
+# Arguments
+- `μ`: Location parameter.
+- `σ`: Scale parameter.
+- `ν`: Degrees of freedom.
+
+# Returns
+A generalized t distribution object.
+
+"""
+
+function GeneralizedTDist(μ, σ, ν)
+  if ν <= zero(ν)
+    ν = nextfloat(zero(ν))
+  end
+
+  if σ <= zero(σ)
+    σ = nextfloat(zero(σ))
+  end
+
+  Generalized_T = μ + σ*Distributions.TDist(ν)
+
+  return(Generalized_T)
+
+end
diff --git a/src/generate_simulation_data_agent.jl b/src/generate_simulation_data_agent.jl
new file mode 100644
index 0000000..d85705d
--- /dev/null
+++ b/src/generate_simulation_data_agent.jl
@@ -0,0 +1,95 @@
+"""
+## Generating Simulation Data for Agent Based Model
+
+To generate simulation data using the agent based model, you can use the `generate_simulation_data_agent` function defined in the `UCIWWEIHR.jl` package. This function allows you to customize various parameters for the simulation.
+NOT FINISHED, STILL NEEDS WW AND RT
+
+### Function Signature
+
+# Arguments
+- seed::Int64: Seed for random number generation. Default value is 1.
+- pop_size::Int64: Size of the population. Default value is 1000.
+- I_init::Int64: Initial number of infected individuals. Default value is 200.
+- H_init::Int64: Initial number of hospitalized individuals. Default value is 20.
+- beta::Float64: Transmission rate. Default value is 0.2.
+- gamma::Float64: Rate of exposed individuals becoming infectious. Default value is 1/4.
+- nu::Float64: Rate of infected individuals recovering or getting hospitalized. Default value is 1/7.
+- epsilon::Float64: Rate of hospitalized individuals recovering. Default value is 1/5.
+- w_init::Float64: Probability of an infected individual becoming hospitalized. Default value is 0.35.
+
+# Returns
+- df::DataFrame: A DataFrame containing the simulation data with columns `Time`, `S`, `E`, `I`, `H`, and `R`.
+"""
+function generate_simulation_data_agent(
+    seed::Int64=1, pop_size::Int64=1000, 
+    I_init::Int64=200, H_init::Int64=20, 
+    beta::Float64=0.2, gamma::Float64=1/4, 
+    nu::Float64=1/7, epsilon::Float64=1/5, 
+    w_init::Float64=0.35
+)
+    
+    Random.seed!(seed)
+
+    pop_vec = repeat(["S"], pop_size)
+    pop_vec[1:I_init] = repeat(["I"], I_init)
+    pop_vec[I_init+1:I_init+H_init] = repeat(["H"], H_init)
+
+    rate_vec = zeros(pop_size)
+    id_vec = 1:pop_size
+
+    t = 0.0
+
+    rate_frame = DataFrame(id = id_vec, state = pop_vec, rate = rate_vec)
+
+    n_suseptible = sum(rate_frame.state .== "S"); n_exposed = sum(rate_frame.state .== "E")
+    n_infected = sum(rate_frame.state .== "I"); n_hospitalized = sum(rate_frame.state .== "H")
+    n_recovered = sum(rate_frame.state .== "R")
+
+    state_frame = DataFrame(T = Float64[], S = Int[], E = Int[], I = Int[], H = Int[], R = Int[])
+    push!(state_frame, (t, n_suseptible, n_exposed, n_infected, n_hospitalized, n_recovered))
+
+    rate_frame.rate .= ifelse.(rate_frame.state .== "S", beta * n_infected,
+                            ifelse.(rate_frame.state .== "E", gamma,
+                                ifelse.(rate_frame.state .== "I", nu,
+                                    ifelse.(rate_frame.state .== "H", epsilon, 0))))
+
+    while n_infected > 0 || n_hospitalized > 0
+        #global t; 
+        #global n_suseptible; global n_exposed; global n_infected; global n_hospitalized; global n_recovered
+        #global state_frame = copy(state_frame); global rate_frame = copy(rate_frame)
+        next_event = rand(Exponential(1/sum(rate_frame.rate)),1)
+
+        which_id = sample(rate_frame.id, Weights(rate_frame.rate), 1)
+
+        rate_frame.state[which_id] .= ifelse(rate_frame.state[which_id] == ["S"], "E",
+                                        ifelse(rate_frame.state[which_id] == ["E"], "I",
+                                            ifelse(rate_frame.state[which_id] == ["I"], ifelse.(rand() < w_init, "H", "R"), "R")))
+
+        n_infected_old = n_infected
+        n_suseptible = sum(rate_frame.state .== "S"); n_exposed = sum(rate_frame.state .== "E")
+        n_infected = sum(rate_frame.state .== "I"); n_hospitalized = sum(rate_frame.state .== "H")
+        n_recovered = sum(rate_frame.state .== "R")
+
+        #new_infected = n_infected - n_infected_old
+        #Rt = ?
+    
+        rate_frame.rate .= ifelse.(rate_frame.state .== "S", beta * n_infected,
+                                ifelse.(rate_frame.state .== "E", gamma,
+                                    ifelse.(rate_frame.state .== "I", nu,
+                                        ifelse.(rate_frame.state .== "H", epsilon, 0))))
+
+        t += next_event[1]
+        current_state = (t, n_suseptible, n_exposed, n_infected, n_hospitalized, n_recovered)
+        state_frame = push!(state_frame, current_state)
+    
+    end
+
+    state_frame.Time = ceil.(Int, state_frame.T)
+    state_frame.Timediff = state_frame.Time - state_frame.T
+    grouped_state_frame = groupby(state_frame, :Time)
+    day_data = combine(sdf -> sdf[argmin(sdf.Timediff), :], grouped_state_frame)
+    select!(day_data, Not(:Timediff,:T))
+
+    return day_data
+
+end
diff --git a/src/generate_simulation_data_ww_eihr.jl b/src/generate_simulation_data_uciwweihr.jl
similarity index 89%
rename from src/generate_simulation_data_ww_eihr.jl
rename to src/generate_simulation_data_uciwweihr.jl
index 0b9a19a..b5bb2d1 100644
--- a/src/generate_simulation_data_ww_eihr.jl
+++ b/src/generate_simulation_data_uciwweihr.jl
@@ -1,12 +1,10 @@
 """
-## Generating Simulation Data for Wastewater EIHR
+## Generating Simulation Data for UCIWWEIHR ODE Compartmental Based Model
 
-To generate simulation data for wastewater EIHR, you can use the `generate_simulation_data_ww_eihr` function defined in the `UCIWWEIHR.jl` package. This function allows you to customize various parameters for the simulation.
+To generate simulation data using the UCIWWEIHR ODE compartmental based model, you can use the `generate_simulation_data_uciwweihr` function defined in the `UCIWWEIHR.jl` package. This function allows you to customize various parameters for the simulation.
 
 ### Function Signature
 
-```julia
-
 # Arguments
 - time_points::Int64: Number of time points wanted for simulation. Default value is 150.
 - seed::Int64: Seed for random number generation. Default value is 1.
@@ -24,8 +22,11 @@ To generate simulation data for wastewater EIHR, you can use the `generate_simul
 - sigma_Rt::Float64: Standard deviation for random walk of time-varying reproduction number. Default value is sqrt(0.02).
 - w_init::Float64: Initial value of the time-varying hospitalization rate. Default value is 0.35.
 - sigma_w::Float64: Standard deviation for random walk of time-varying hospitalization rate. Default value is sqrt(0.02).
+
+# Returns
+- df::DataFrame: A DataFrame containing the simulation data with columns `obstimes`, `log_ww_conc`, `hosp`, and `rt`.
 """
-function generate_simulation_data_ww_eihr(
+function generate_simulation_data_uciwweihr(
     time_points::Int64=150, seed::Int64=1,
     E_init::Int64=200, I_init::Int64=100, H_init::Int64=20,
     gamma::Float64=1/4, nu::Float64=1/7, epsilon::Float64=1/5,
diff --git a/src/distribution_functions.jl b/src/negativebinomial2.jl
similarity index 55%
rename from src/distribution_functions.jl
rename to src/negativebinomial2.jl
index bb1388a..eaeb09d 100644
--- a/src/distribution_functions.jl
+++ b/src/negativebinomial2.jl
@@ -1,4 +1,3 @@
-
 # credit to Damon Bayer for all of these helper functions
 """
 Create a re-parametrized negative binomial distribution in terms of mean and overdispersion.
@@ -29,32 +28,4 @@ function NegativeBinomial2(μ, ϕ)
   end
 
   Distributions.NegativeBinomial(r, p)
-end
-
-"""
-Create a generalized t distribution with location and scale.
-
-# Arguments
-- `μ`: Location parameter.
-- `σ`: Scale parameter.
-- `ν`: Degrees of freedom.
-
-# Returns
-A generalized t distribution object.
-
-"""
-
-function GeneralizedTDist(μ, σ, ν)
-  if ν <= zero(ν)
-    ν = nextfloat(zero(ν))
-  end
-
-  if σ <= zero(σ)
-    σ = nextfloat(zero(σ))
-  end
-
-  Generalized_T = μ + σ*Distributions.TDist(ν)
-
-  return(Generalized_T)
-
-end
+end
\ No newline at end of file
diff --git a/src/sum_values.jl b/src/sum_values.jl
deleted file mode 100644
index 73747e9..0000000
--- a/src/sum_values.jl
+++ /dev/null
@@ -1,11 +0,0 @@
-"""
-    sum_values(x,y)
-This is an example of Docstring. This function receives two 
-numbers x and y and returns the sum of the squares.
-```math
-x^1 + y^1 
-```
-"""
-function sum_values(x,y)
-  return x+y
-end
\ No newline at end of file