2024-07-21 update : added main model and added sim func.

cbernalz · Jul 21, 2024 · 784e107 · 784e107
1 parent c04006f
commit 784e107
Show file tree

Hide file tree

Showing 11 changed files with 482 additions and 7 deletions.
diff --git a/Project.toml b/Project.toml
@@ -3,6 +3,26 @@ uuid = "176850df-a79a-485c-b76d-f2c16e00fafb"
 authors = ["Christian O. Bernal Zelaya"]
 version = "1.0.0-DEV"
 
+[deps]
+AxisArrays = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
+CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
+Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+
+
 [compat]
 julia = "1"
 

diff --git a/docs/Project.toml b/docs/Project.toml
@@ -1,3 +1,4 @@
 [deps]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 UCIWWEIHR = "176850df-a79a-485c-b76d-f2c16e00fafb"
diff --git a/docs/make.jl b/docs/make.jl
@@ -11,7 +11,7 @@ makedocs(;
   sitename = "UCIWWEIHR.jl",
   format = Documenter.HTML(;
     prettyurls = get(ENV, "CI", "false") == "true",
-    canonical = "https://cbernalz.github.io/UCIWWEIHR.jl",
+    repolink = "https://cbernalz.github.io/UCIWWEIHR.jl",
   ),
   pages = [
     "HOME" => "index.md", 

diff --git a/docs/src/license.md b/docs/src/license.md
@@ -1,6 +1,6 @@
 # [UCIWWEIHR.jl Package License](@id license)
 
-The UCIWWEIHR.jl package is licensed under the [MIT License](https://opensource.org/licenses/MIT).
+The UCIWWEIHR.jl package is licensed under the MIT License.
 
 ## MIT License
 

diff --git a/docs/src/tutorial.md b/docs/src/tutorial.md
@@ -1,14 +1,59 @@
+```@setup tutorial
+using Plots; gr()
+Plots.reset_defaults()
+```
 # [Tutorials](@id tutorial)
 
 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.
 
-## Getting Started
+## 1. Getting Started
 
-### Installation
+### 1.1 Installation
 
 To install the UCIWWEIHR.jl package, open the Julia REPL and run the following command:
 
-```julia
+``` julia
 using Pkg
 Pkg.add("UCIWWEIHR")
+```
+
+## 2. Generating simulated data
+
+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 way to generate synthetic data for the default settings of `generate_simulation_data_ww_eihr`:
+
+``` @example tutorial
+using UCIWWEIHR
+using Plots
+# Running simulation function with defaults
+df = generate_simulation_data_ww_eihr()
+first(df, 5)
+```
+
+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
+```@example tutorial
+plot(df.obstimes, df.log_ww_conc,
+    label=nothing,
+    xlabel="Obstimes", 
+    ylabel="Conc. of Pathogen Genome in WW", 
+    title="Plot of Conc. of Pathogen Genome in WW Over Time")
+```
+
+### 2.2. Hospitalizations
+```@example tutorial
+plot(df.obstimes, df.hosp, 
+    label=nothing,
+    xlabel="Obstimes", 
+    ylabel="Hosp", 
+    title="Plot of Hosp Over Time")
+```
+
+### 2.3. Reproductive Number
+```@example tutorial
+plot(df.obstimes, df.rt, 
+    label=nothing,
+    xlabel="Obstimes", 
+    ylabel="RT", 
+    title="Plot of RT Over Time")
 ```
diff --git a/src/UCIWWEIHR.jl b/src/UCIWWEIHR.jl
@@ -1,8 +1,34 @@
 module UCIWWEIHR
 
-# Write your package code here.
+using AxisArrays
+using MCMCChains
+using Optim
+using LineSearches
+using Random
+using LinearAlgebra
+using NaNMath
+using LogExpFunctions
+using PreallocationTools
+using Distributions
+using Turing
+using Random
+using ForwardDiff
+using Logging
+using CSV
+using DataFrames
+using DifferentialEquations
+using Plots
 
 include("sum_values.jl")
+include("generate_simulation_data_ww_eihr.jl")
+include("bayes_eihr_model.jl")
+include("distribution_functions.jl")
+include("eihr_ode.jl")
+export eihr_ode
 export sum_values
+export generate_simulation_data_ww_eihr
+export bayes_eihr_model
+export NegativeBinomial2
+export GeneralizedTDist
 
 end
diff --git a/src/bayes_eihr_model.jl b/src/bayes_eihr_model.jl
@@ -0,0 +1,190 @@
+# bayesian non-constant Rt and hosp rate model eihr
+# -------------------------------------------------
+"""
+    bayes_eihr_model(...)
+This is the bayesian semi-parametric model for the wastewater EIHR compartmental model.  
+The defaults for this fuction will follow those of the default simulation in generate_simulation_data_ww_eihr.jl function.
+
+# Arguments
+- `data_hosp::Array{Float64}`: An array of hospital data.
+- `data_wastewater::Array{Float64}`: An array of pathogen genome concentration in localized wastewater data.
+- `obstimes::Array{Float64}`: An array of timepoints for observed hosp/wastewater.
+- `E_init_sd::Float64`: Standard deviation for the initial number of exposed individuals.
+- `E_init_mean::Int64`: Mean for the initial number of exposed individuals.
+- `I_init_sd::Float64`: Standard deviation for the initial number of infected individuals.
+- `I_init_mean::Int64`: Mean for the initial number of infected individuals.
+- `H_init_sd::Float64`: Standard deviation for the initial number of hospitalized individuals.
+- `H_init_mean::Int64`: Mean for the initial number of hospitalized individuals.
+- `gamma_sd::Float64`: Standard deviation for the rate of incubation.
+- `log_gamma_mean::Float64`: Mean for the rate of incubation on log scale.
+- `nu_sd::Float64`: Standard deviation for the rate of laeving the infected compartment.
+- `log_nu_mean::Float64`: Mean for the rate of leaving the infected compartment on the log scale.
+- `epsilon_sd::Float64`: Standard deviation for the rate of hospitalization recovery.
+- `log_epsilon_mean::Float64`: Mean for the rate of hospitalization recovery on the log scale.
+- `rho_gene_sd::Float64`: Standard deviation for the rho prior.
+- `log_rho_gene_mean::Float64`: Mean for the row prior on log scale.
+- `tau_sd::Float64`: Standard deviation for the scale/variation of the log scale data.
+- `log_tau_mean::Float64`: Mean for the scale/variation of the log scale data on log scale itself.
+- `df_shape::Float64`: Shape parameter for the gamma distribution.
+- `df_scale::Float64`: Scale parameter for the gamma distribution.
+- `sigma_hosp_sd::Float64`: Standard deviation for the negative binomial distribution for hospital data.
+- `sigma_hosp_mean::Float64`: Mean for the negative binomial distribution for hospital data.
+- `Rt_init_sd::Float64`: Standard deviation for the initial value of the time-varying reproduction number.
+- `Rt_init_mean::Float64`: Mean for the initial value of the time-varying reproduction number.
+- `sigma_Rt_sd::Float64`: Standard deviation for normal prior of log time-varying reproduction number standard deviation.
+- `sigma_Rt_mean::Float64`: Mean for normal prior of log time-varying reproduction number standard deviation.
+- `w_init_sd::Float64`: Standard deviation for the initial value of the time-varying hospitalization rate.
+- `w_init_mean::Float64`: Mean for the initial value of the time-varying hospitalization rate.
+- `sigma_w_sd::Float64`: Standard deviation for normal prior of log time-varying hospitalization rate standard devaiation.
+- `sigma_w_mean::Float64`: Mean for normal prior of time-varying hospitalization rate standard devaiation.
+- `param_change_times::Array{Float64}`: An array of timepoints where the parameters change.
+
+"""
+
+@model function bayes_eihr_model(
+    data_hosp::Array{Float64},
+    data_wastewater::Array{Float64},
+    obstimes::Array{Float64},
+    E_init_sd::Float64, E_init_mean::Int64,
+    I_init_sd::Float64, I_init_mean::Int64,
+    H_init_sd::Float64, H_init_mean::Int64,
+    gamma_sd::Float64, log_gamma_mean::Float64,
+    nu_sd::Float64, log_nu_mean::Float64,
+    epsilon_sd::Float64, log_epsilon_mean::Float64,
+    rho_gene_sd::Float64, log_rho_gene_mean::Float64,
+    tau_sd::Float64, log_tau_mean::Float64,
+    df_shape::Float64, df_scale::Float64,
+    sigma_hosp_sd::Float64, sigma_hosp_mean::Float64,
+    Rt_init_sd::Float64, Rt_init_mean::Float64,
+    sigma_Rt_sd::Float64, sigma_Rt_mean::Float64,
+    w_init_sd::Float64, w_init_mean::Float64,
+    sigma_w_sd::Float64, sigma_w_mean::Float64,
+    param_change_times::Array{Float64}
+    )
+
+
+        # Prelims
+        max_neg_bin_sigma = 1e10
+        min_neg_bin_sigma = 1e-10
+
+
+        # Calculate number of observed datapoints timepoints
+        l_obs = length(obstimes)
+        l_param_change_times = length(param_change_times)
+
+
+        # PRIORS-----------------------------
+        # Compartments
+        E_init_non_centered ~ Normal()
+        I_init_non_centered ~ Normal()
+        H_init_non_centered ~ Normal()
+        # Parameters for compartments
+        gamma_non_centered ~ Normal()
+        nu_non_centered ~ Normal()
+        epsilon_non_centered ~ Normal()
+        # Parameters for wastewater
+        rho_gene_non_centered ~ Normal() # gene detection rate
+        tau_non_centered ~ Normal() # standard deviation for log scale data
+        df ~ Gamma(df_shape, df_scale)
+        # Parameters for hospital
+        sigma_hosp_non_centered ~ Normal()
+        # Non-constant Rt
+        Rt_params_non_centered ~ MvNormal(zeros(l_param_change_times + 2), I) # +2 for sigma and init
+        sigma_Rt_non_centered = Rt_params_non_centered[1]
+        Rt_init_non_centered = Rt_params_non_centered[2]
+        log_Rt_steps_non_centered = Rt_params_non_centered[3:end]
+        # Non-constant Hosp Rate w
+        w_params_non_centered ~ MvNormal(zeros(l_param_change_times + 2), I) # +2 for sigma and init
+        sigma_w_non_centered = w_params_non_centered[1]
+        w_init_non_centered = w_params_non_centered[2]
+        log_w_steps_non_centered = w_params_non_centered[3:end]
+
+
+        # TRANSFORMATIONS--------------------
+        # Compartments
+        E_init = E_init_non_centered * E_init_sd + E_init_mean
+        I_init = I_init_non_centered * I_init_sd + I_init_mean
+        H_init = H_init_non_centered * H_init_sd + H_init_mean
+        # Parameters for compartments
+        gamma = exp(gamma_non_centered * gamma_sd + log_gamma_mean)
+        nu = exp(nu_non_centered * nu_sd + log_nu_mean)
+        epsilon = exp(epsilon_non_centered * epsilon_sd + log_epsilon_mean)
+        # Parameters for wastewater
+        rho_gene = exp(rho_gene_non_centered * rho_gene_sd + log_rho_gene_mean)
+        tau = exp(tau_non_centered * tau_sd + log_tau_mean)
+        # Parameters for hospital
+        sigma_hosp = clamp.(sigma_hosp_non_centered * sigma_hosp_sd + sigma_hosp_mean, min_neg_bin_sigma, max_neg_bin_sigma)
+        # Non-constant Rt
+        Rt_init = exp(Rt_init_non_centered * Rt_init_sd + Rt_init_mean)
+        sigma_Rt = exp(sigma_Rt_non_centered * sigma_Rt_sd + sigma_Rt_mean)
+        alpha_t_no_init = exp.(log(Rt_init) .+ cumsum(log_Rt_steps_non_centered) * sigma_Rt) * nu
+        alpha_init = Rt_init * nu
+        alpha_t = vcat(alpha_init, alpha_t_no_init)
+        # Non-constant Hosp Rate w
+        w_init = exp(w_init_non_centered * w_init_sd + w_init_mean)
+        sigma_w = exp(sigma_w_non_centered * sigma_w_sd + sigma_w_mean)
+        w_no_init = exp.(log(w_init) .+ cumsum(log_w_steps_non_centered) * sigma_w)
+        w_t = vcat(w_init, w_no_init)
+
+
+        # ODE SETUP--------------------------
+        prob = ODEProblem{true}(eihr_ode!, zeros(3), (0.0, obstimes[end]), ones(5))
+        function param_affect_beta!(integrator)
+            ind_t = searchsortedfirst(param_change_times, integrator.t) # Find the index of param_change_times that contains the current timestep
+            integrator.p[1] = alpha_t_no_init[ind_t] # Replace alpha with a new value from alpha_t_no_init
+            integrator.p[4] = w_no_init[ind_t] # Replace w with a new value from w_no_init
+        end
+        param_callback = PresetTimeCallback(param_change_times, param_affect_beta!, save_positions=(false, false))
+        u0 = [E_init, I_init, H_init]
+        p0 = [alpha_init, gamma, nu, w_init, epsilon]
+        extra_ode_precision = false
+        abstol = extra_ode_precision ? 1e-11 : 1e-9
+        reltol = extra_ode_precision ? 1e-8 : 1e-6
+        sol = solve(prob, Tsit5(); callback=param_callback, saveat=obstimes, save_start=true, 
+                    verbose=false, abstol=abstol, reltol=reltol, u0=u0, p=p0, tspan=(0.0, obstimes[end]))
+        # If the ODE solver fails, reject the sample by adding -Inf to the likelihood
+        if sol.retcode != :Success
+            println("An error occurred during ODE solution!!!")
+            Turing.@addlogprob! -Inf
+            return
+        end
+        sol_array = Array(sol)
+        I_comp_sol = clamp.(sol_array[2,2:end],1, 1e10)
+
+
+        # W-W means--------------------------
+        # E - 1 // I - 2 // H - 3 // R - 4
+        log_genes_mean = log.(I_comp_sol) .+ log(rho_gene) # first entry is the initial conditions, we want 2:end
+        # Likelihood calculations------------
+        sol_hosp = clamp.(sol_array[3,2:end], 1, 1e10)
+        for i in 1:l_obs
+            data_wastewater[i] ~ GeneralizedTDist(log_genes_mean[i], tau, df)
+            data_hosp[i] ~ NegativeBinomial2(sol_hosp[i], sigma_hosp)
+        end
+
+
+        # Generated quantities
+        H_comp = sol_array[3, :]
+        rt_vals = alpha_t / nu
+        w_t = w_t
+
+        return (
+            E_init,
+            I_init,
+            H_init,
+            alpha_t = alpha_t,
+            gamma = gamma,
+            nu = nu,
+            w_t = w_t,
+            epsilon = epsilon,
+            rt_vals = rt_vals,
+            rho_gene = rho_gene,
+            tau = tau,
+            df = df,
+            sigma_hosp = sigma_hosp,
+            H = H_comp,
+            log_genes_mean = log_genes_mean
+        )
+
+
+    end