2024-08-23 update : added all result visuals & updated docs.

docs/make.jl

@@ -22,7 +22,8 @@ makedocs(;
       "GETTING STARTED" => "tutorials/getting_started.md",
       "UCIWWEIHR SIMULATION DATA" => "tutorials/uciwweihr_simulation_data.md",
       "AGENT-BASED SIMULATION DATA" => "tutorials/agent_based_simulation_data.md",
-      "UCIWWEIHR FITTING MODEL" => "tutorials/uciwweihr_model_fitting.md",
+      "UCIWWEIHR FITTING MODEL W/OUT FORECASTING" => "tutorials/uciwwiehr_model_fitting_no_forecast.md",
+      "UCIWWEIHR FITTING MODEL W/ FORECASTING" => "tutorials/uciwwiehr_model_fitting_forecast.md",
     ]
     ,
     "NEWS" => "news.md",

docs/src/tutorial_index.md

@@ -7,5 +7,7 @@ Future Description.
 - [Getting Started](@ref getting_started)
 - [Generating simulated data with UCIWWEIHR ODE compartmental based model.](@ref uciwweihr_simulation_data)
 - [Generating simulated data with an agent based model.](@ref agent_based_simulation_data)
+- [Generating posterior distribution samples with UCIWWEIHR ODE compartmental based model without forecasting.](@ref uciwwiehr_model_fitting_no_forecast)
+- [Generating posterior distribution samples with UCIWWEIHR ODE compartmental based model with forecasting.](@ref uciwwiehr_model_fitting_with_forecast)
 
 

docs/src/tutorials/uciwwiehr_model_fitting_forecast.md

@@ -0,0 +1,110 @@
+```@setup tutorial
+using Plots, StatsPlots; gr()
+Plots.reset_defaults()
+
+```
+
+# [Generating Posterior Distribution Samples with UCIWWEIHR ODE Compartmental Based Model with Forecasting.](@id uciwwiehr_model_fitting_with_forecast)
+
+Here we extend the [previous tutorial](@ref uciwwiehr_model_fitting_no_forecast)  to include forecasting capabilities.  We start with generating out data using `generate_simulation_data_uciwweihr`'s alternate parameterization where we do not prespecify the effective reproduction number and hospitalization probability but instead preform a log-normal random walk and a logit-normal random walk respectively.  We then sample from the posterior distribution using the `uciwweihr_fit.jl` function.  We then generate desired quantities and forecast for a given time period with the posterior predictive distribution, using `uciwweihr_gq_pp.jl`.
+
+
+## 1. Data Generation.
+
+Here we generate two datasets, one with 150 time points and one with 178 time points.  We will use the 150 time point dataset for fitting and the 178 time point dataset for forecast evaluation.
+
+``` @example tutorial
+using UCIWWEIHR
+# Running simulation function with presets
+params = create_uciwweihr_params(
+    time_points = 150
+)
+df = generate_simulation_data_uciwweihr(params)
+
+params_ext = create_uciwweihr_params(
+    time_points = 178
+)
+df_ext = generate_simulation_data_uciwweihr(params_ext)
+first(df, 5)
+first(df_ext, 5)
+```
+
+## 2. Sampling from the Posterior Distribution and Posterior Predictive Distribution.
+
+Here we sample from the posterior distribution using the `uciwweihr_fit.jl` function.  First, we setup some presets, then have an array where index 1 contains the posterior/prior predictive samples, index 2 contains the posterior/prior generated quantities samples, and index 3 contains the original sampled parameters for the model.  The diference here is that we set `forecast = true` and `forecast_weeks = 4` to forecast 4 weeks into the future.
+
+``` @example tutorial
+data_hosp = df.hosp
+data_wastewater = df.log_ww_conc
+obstimes = df.obstimes
+param_change_times = 1:7:length(obstimes) # Change every week
+priors_only = false
+n_samples = 200
+forecast = true
+forecast_weeks = 4
+
+samples = uciwweihr_fit(
+    data_hosp,
+    data_wastewater,
+    obstimes,
+    param_change_times,
+    priors_only,
+    n_samples
+)
+model_output = uciwweihr_gq_pp(
+    samples = samples,
+    data_hosp = data_hosp,
+    data_wastewater = data_wastewater,
+    obstimes = obstimes,
+    param_change_times = param_change_times,
+    forecast = forecast,
+    forecast_weeks = forecast_weeks
+)
+
+first(model_output[1][:,1:5], 5)
+```
+
+``` @example tutorial
+first(model_output[2][:,1:5], 5)
+```
+
+``` @example tutorial
+first(model_output[3][:,1:5], 5)
+```
+
+## 3. MCMC Diagnostic Plots/Results Along with Posterior Predictive Distribution.
+
+We can again look at model diagnostics, posterior distribution of time or non-time varying parameters, and the posterior predictive distribution extended for forecasting.
+
+```@example tutorial
+uciwweihr_visualizer(
+    pp_samples = model_output[1],
+    gq_samples = model_output[2],
+    data_hosp = df_ext.hosp,
+    data_wastewater = df_ext.log_ww_conc, 
+    actual_rt_vals = df_ext.rt, 
+    actual_w_t = df_ext.wt, 
+    actual_non_time_varying_vals = params,
+    forecast_weeks = forecast_weeks,
+    bayes_dist_type = "Posterior",
+    save_plots = false
+)
+```
+
+### 3.1. MCMC Diagnostic Plots.
+
+![Plot 1](plots/mcmc_diagnosis_plots.png)
+
+### 3.2. Time Varying Parameter Results Plot.
+
+![Plot 2](plots/mcmc_time_varying_parameter_plots.png)
+
+### 3.3. Non-Time Varying Parameter Results Plot.
+![Plot 3](plots/mcmc_nontime_varying_parameter_plots.png)
+
+### 3.4. Posterior Predictive Distribution Plot.
+
+![Plot 4](plots/mcmc_pred_parameter_plots.png)
+
+
+### [Tutorial Contents](@ref tutorial_home)

docs/src/tutorials/uciwweihr_model_fitting.md → docs/src/tutorials/uciwwiehr_model_fitting_no_forecast.md

@@ -4,7 +4,7 @@ Plots.reset_defaults()
 
 ```
 
-# [Generating Posterior Distribution Samples with UCIWWEIHR ODE compartmental based model.](@id uciwwiehr_model_fitting)
+# [Generating Posterior Distribution Samples with UCIWWEIHR ODE Compartmental Based Model without Forecasting.](@id uciwwiehr_model_fitting_no_forecast)
 
 This package has a way to sample from a posterior or prior that is defined in the future paper using the `uciwweihr_fit.jl` and `uciwweihr_model.jl`.  We can then generate desired quantities and forecast for a given time period with the posterior predictive distribution, using `uciwweihr_gq_pp.jl`.  We first generate data using the `generate_simulation_data_uciwweihr` function which is a non-mispecified version of the model, we will also be using prespecified effective reporduction curves and prespecified hospitalization probability curves.
 
@@ -49,7 +49,7 @@ data_wastewater = df.log_ww_conc
 obstimes = df.obstimes
 param_change_times = 1:7:length(obstimes) # Change every week
 priors_only = false
-n_samples = 50
+n_samples = 200
 
 samples = uciwweihr_fit(
     data_hosp,
@@ -83,10 +83,17 @@ first(model_output[3][:,1:5], 5)
 We also provide a very basic way to visualize some MCMC diagnostics along with effective sample sizes of desired generated quantities(does not include functionality for time-varying quantities).  Along with this, we can also visualize the posterior predictive distribution with actual observed values, which can be used to examine forecasts generated by the model.
 
 ```@example tutorial
-uciwweihr_visualizer(gq_samples = model_output[2], 
-                        actual_rt_vals = df.rt, 
-                        actual_w_t = df.wt, 
-                        save_plots = true)
+uciwweihr_visualizer(
+    pp_samples = model_output[1],
+    gq_samples = model_output[2],
+    data_hosp = data_hosp,
+    data_wastewater = data_wastewater, 
+    actual_rt_vals = df.rt, 
+    actual_w_t = df.wt, 
+    actual_non_time_varying_vals = params,
+    bayes_dist_type = "Posterior",
+    save_plots = true
+)
 ```
 
 ### 3.1. MCMC Diagnostic Plots.
@@ -97,5 +104,12 @@ uciwweihr_visualizer(gq_samples = model_output[2],
 
 ![Plot 2](plots/mcmc_time_varying_parameter_plots.png)
 
+### 3.3. Non-Time Varying Parameter Results Plot.
+![Plot 3](plots/mcmc_nontime_varying_parameter_plots.png)
+
+### 3.4. Posterior Predictive Distribution Plot.
+
+![Plot 4](plots/mcmc_pred_parameter_plots.png)
+
 
 ### [Tutorial Contents](@ref tutorial_home)

src/UCIWWEIHR.jl

@@ -30,15 +30,18 @@ include("generalizedtdist.jl")
 include("uciwweihr_model.jl")
 include("uciwweihr_fit.jl")
 include("uciwweihr_gq_pp.jl")
-include("uciwweihr_visualizer.jl")
 include("helper_functions.jl")
 include("mcmcdiags_vis.jl")
 include("time_varying_param_vis.jl")
+include("non_time_varying_param_vis.jl")
+include("predictive_param_vis.jl")
+include("uciwweihr_visualizer.jl")
 
 export eihr_ode
 export uciwweihr_sim_params
 export create_uciwweihr_params
 export generate_random_walk
+export generate_logit_normal_random_walk
 export generate_simulation_data_uciwweihr
 export generate_simulation_data_agent
 export NegativeBinomial2
@@ -50,8 +53,12 @@ export uciwweihr_visualizer
 export ChainsCustomIndexs
 export save_plots_to_docs
 export startswith_any
+export generate_colors
 export calculate_quantiles
+export repeat_last_n_elements
 export mcmcdiags_vis
 export time_varying_param_vis
+export non_time_varying_param_vis
+export predictive_param_vis
 
 end

src/generate_simulation_data_uciwweihr.jl

@@ -1,6 +1,7 @@
 """
     uciwweihr_sim_params
-Struct for holding parameters used in the UCIWWEIHR ODE compartmental model simulation.
+
+Struct for holding parameters used in the UCIWWEIHR ODE compartmental model simulation.  Use `create_uciwweihr_params` to create an instance of this struct.
 
 # Fields
 - `time_points::Int64`: Number of time points for the simulation.
@@ -19,6 +20,8 @@ Struct for holding parameters used in the UCIWWEIHR ODE compartmental model simu
 - `sigma_Rt::Float64`: Standard deviation for random walk of time-varying reproduction number.
 - `w::Union{Float64, Vector{Float64}}`: Initial value or time series of the time-varying hospitalization rate.
 - `sigma_w::Float64`: Standard deviation for random walk of time-varying hospitalization rate.
+- `rt_init::Float64`: Initial value of the time-varying reproduction number, NOT USER SPECIFIED `create_uciwweihr_params` TAKES CARE OF THIS.
+- `w_init::Float64`: Initial value of the time-varying hospitalization rate, NOT USER SPECIFIED `create_uciwweihr_params` TAKES CARE OF THIS.
 """
 struct uciwweihr_sim_params
     time_points::Int64
@@ -37,10 +40,13 @@ struct uciwweihr_sim_params
     sigma_Rt::Float64
     w::Union{Float64, Vector{Float64}}
     sigma_w::Float64
+    rt_init::Float64
+    w_init::Float64
 end
 
 """
     create_uciwweihr_params(; kwargs...)
+
 Creates a `uciwweihr_sim_params` struct with the option to either use a predetermined `Rt` and `w` or generate them as random walks.
 
 # Arguments
@@ -60,16 +66,19 @@ function create_uciwweihr_params(; time_points::Int64=150, seed::Int64=1,
                                   sigma_w::Float64=sqrt(0.001))
 
     Random.seed!(seed)
+    rt_init = isa(Rt, Float64) ? Rt : Rt[1]
+    w_init = isa(w, Float64) ? w : w[1]
 
     Rt_t = isa(Rt, Float64) ? generate_random_walk(time_points, sigma_Rt, Rt) : Rt
-    w_t = isa(w, Float64) ? generate_random_walk(time_points, sigma_w, w) : w
+    w_t = isa(w, Float64) ? generate_logit_normal_random_walk(time_points, sigma_w, w) : w
 
     return uciwweihr_sim_params(time_points, seed, E_init, I_init, H_init, gamma, nu, epsilon,
-                                rho_gene, tau, df, sigma_hosp, Rt_t, sigma_Rt, w_t, sigma_w)
+                                rho_gene, tau, df, sigma_hosp, Rt_t, sigma_Rt, w_t, sigma_w, rt_init, w_init)
 end
 
 """
     generate_random_walk(time_points::Int64, sigma::Float64, init_val::Float64)
+
 Generates a random walk time series.
 
 # Arguments
@@ -90,6 +99,29 @@ function generate_random_walk(time_points::Int64, sigma::Float64, init_val::Floa
     return walk
 end
 
+"""
+    generate_logit_normal_random_walk(time_points::Int64, sigma::Float64, init_val::Float64)
+
+Generates a logit-normal random walk time series.
+
+# Arguments
+- `time_points::Int64`: Number of time points.
+- `sigma::Float64`: Standard deviation of the random walk in logit space.
+- `init_val::Float64`: Initial value of the random walk on the probability scale.
+
+# Returns
+- `walk::Vector{Float64}`: Generated random walk on the probability scale.
+"""
+function generate_logit_normal_random_walk(time_points::Int64, sigma::Float64, init_val::Float64)
+    walk = Float64[]
+    logit_val = logit(init_val)
+    for _ in 1:time_points
+        logit_val = rand(Normal(0, sigma)) + logit_val
+        push!(walk, logistic(logit_val))
+    end
+    return walk
+end
+
 """
     generate_simulation_data(params::UCIWWEIHRParams)
 

src/helper_functions.jl

@@ -110,7 +110,6 @@ end
 
 Saves plots to docs/plots directory.
 
-Function created by Christian Bernal Zelaya. 
 """
 function save_plots_to_docs(plot, filename; format = "png")
     doc_loc = "plots"
@@ -130,7 +129,6 @@ end
 
 Checks if the name of time varying paramter starts with any of the patterns.
 
-Function created by Christian Bernal Zelaya. 
 """
 function startswith_any(name, patterns)
     for pattern in patterns
@@ -147,7 +145,6 @@ end
 
 Calculate quantiles for a given chain and variable prefix.  Quantiles can be any user desired quantile.
 
-Function created by Christian Bernal Zelaya. 
 """
 function calculate_quantiles(df, chain, var_prefix, quantiles)
     df_chain = filter(row -> row.chain == chain, df)
@@ -167,9 +164,28 @@ end
 
 Generates a vector with colors for ribbons in plots.
 
-Function created by Christian Bernal Zelaya. 
 """
 function generate_colors(number_of_colors)
     alpha_values = range(0.1, stop=0.7, length=number_of_colors)
     return [RGBA(colorant"blue", alpha) for alpha in alpha_values]
+end
+
+
+"""
+    repeat_last_n_elements(x::Vector{T}, n::Int, w::Int) where T
+
+Modifies a given array so that the last n elements are repeated w times.
+    
+"""
+function repeat_last_n_elements(x::Vector{T}, n::Int, w::Int) where T
+    if n == 0
+        return x
+    else
+        n = min(n, length(x))
+        last_n_elements = x[end-n+1:end]
+        repeated_elements = [elem for elem in last_n_elements for _ in 1:w]
+        x_new = vcat(x, repeated_elements)
+
+        return x_new
+    end
 end