diff --git a/pipeline/scripts/create_figure1.jl b/pipeline/scripts/create_figure1.jl index 81c0c29b0..a4dd6ed12 100644 --- a/pipeline/scripts/create_figure1.jl +++ b/pipeline/scripts/create_figure1.jl @@ -1,4 +1,4 @@ -## Script to make figure 1 +## Script to make figure 1 and alternate latent models for SI using Pkg Pkg.activate(joinpath(@__DIR__(), "..")) @@ -19,8 +19,6 @@ truth_df = mapreduce(vcat, truth_data_files) do filename make_truthdata_dataframe(filename, D, pipelines) end -## Make mainfigure plots - # Define scenario titles and reference times for figure 1 scenario_dict = Dict( "measures_outbreak" => (title = "Outbreak with measures", T = 28), @@ -29,7 +27,22 @@ scenario_dict = Dict( "rough_endemic" => (title = "Rough endemic", T = 35) ) -fig1 = figureone(truth_df, analysis_df, scenario_dict) +target_dict = Dict( + "log_I_t" => (title = "log(Incidence)", ylims = (3.5, 6)), + "rt" => (title = "Exp. growth rate", ylims = (-0.1, 0.1)), + "Rt" => (title = "Reproductive number", ylims = (-0.1, 3)) +) + +latent_model_dict = Dict( + "wkly_rw" => (title = "Random walk",), + "wkly_ar" => (title = "AR(1)",), + "wkly_diff_ar" => (title = "Diff. AR(1)",) +) + +## `wkly_ar` is the default latent model which we show as figure 1, others are for SI -## Save the figure -save(plotsdir("figure1.png"), fig1) +_ = map(latent_model_dict |> keys |> collect) do latent_model + fig = figureone( + truth_df, analysis_df, latent_model, scenario_dict, target_dict, latent_model_dict) + save(plotsdir("figure1_$(latent_model).png"), fig) +end diff --git a/pipeline/src/mainplots/figureone.jl b/pipeline/src/mainplots/figureone.jl index 13410afe0..72e34d380 100644 --- a/pipeline/src/mainplots/figureone.jl +++ b/pipeline/src/mainplots/figureone.jl @@ -110,7 +110,8 @@ function _figure_one_scenario_truth_data(truth_df, scenario; true_gi_choice) end """ -Generate figure 1 showing the analysis and truth data for different scenarios. +Generate a version figure 1 showing the analysis and truth data for different scenarios _and_ +different latent process models. ## Arguments - `truth_df`: DataFrame containing the truth data. @@ -127,7 +128,7 @@ Generate figure 1 showing the analysis and truth data for different scenarios. - `fig`: Figure object containing the generated figure. """ -function figureone( +function figureone_with_latent_model( truth_df, analysis_df, scenario_dict; fig_kws = (; size = (1000, 2000)), true_gi_choice = 10.0, used_gi_choice = 10.0, legend_title = "Process type") # Perform checks on the dataframes @@ -163,3 +164,181 @@ function figureone( return fig end + +""" +Internal method for creating a model panel plot for Figure One. + +This function takes in various parameters to filter the `analysis_df` DataFrame and create a model panel plot for Figure One. +The filtered DataFrame is used to generate the plot using the `model_plotting_data` variable. +The plot includes process values, color-coded by the infection generating process, +and credible intervals defined by `lower_sym` and `upper_sym`. + +## Arguments +- `analysis_df`: The DataFrame containing the analysis data. +- `scenario`: The scenario to filter the DataFrame. +- `target`: The target to filter the DataFrame. +- `latentmodel`: The latent model to filter the DataFrame. +- `reference_time`: The reference time to filter the DataFrame. +- `true_gi_choice`: The true GI mean value to filter the DataFrame. +- `used_gi_choice`: The used GI mean value to filter the DataFrame. +- `lower_sym`: The symbol representing the lower confidence interval (default: `:q_025`). +- `upper_sym`: The symbol representing the upper confidence interval (default: `:q_975`). + +## Returns +- `plt_model`: The model panel plot. + +""" +function _figure_one_model_panel( + analysis_df, scenario, target, latentmodel; reference_time, true_gi_choice, + used_gi_choice, lower_sym = :q_025, upper_sym = :q_975) + model_plotting_data = analysis_df |> + df -> @subset(df, :True_GI_Mean.==true_gi_choice) |> + df -> @subset(df, :Used_GI_Mean.==used_gi_choice) |> + df -> @subset(df, :Reference_Time.==reference_time) |> + df -> @subset(df, :Scenario.==scenario) |> + df -> @subset(df, :Target.==target) |> + df -> @subset(df, + :Latent_Model.==latentmodel) |> + data + + plt_model = model_plotting_data * + mapping(:target_times => "T", :q_5 => "Process values", + color = :IGP_Model => "Infection generating process") * + mapping(lower = lower_sym, upper = upper_sym) * visual(LinesFill) + + return plt_model +end + +""" +Internal method for creating a truth data panel plot for a given scenario and + target using the provided truth data. + +## Arguments +- `truth_df`: DataFrame containing the truth data. +- `scenario`: Scenario to plot. +- `target`: Target to plot. +- `true_gi_choice`: True GI choice to filter the data. + +## Returns +- `plt_truth`: Plot object representing the truth data panel. + +""" +function _figure_one_truth_data_panel(truth_df, scenario, target; true_gi_choice) + truth_plotting_data = truth_df |> + df -> @subset(df, :True_GI_Mean.==true_gi_choice) |> + df -> @subset(df, :Scenario.==scenario) |> + df -> @subset(df, :Target.==target) |> data + plt_truth = truth_plotting_data * + mapping( + :target_times => "T", :target_values => "values", color = :IGP_Model) * + visual(Scatter) + return plt_truth +end + +""" +Create figure one with multiple panels showing the analysis results and truth data for different scenarios and targets. + +# Arguments +- `truth_df::DataFrame`: The truth data as a DataFrame. +- `analysis_df::DataFrame`: The analysis data as a DataFrame. +- `latent_model::AbstractString`: The latent model to use for the infection generating process. +- `scenario_dict::Dict{AbstractString, Scenario}`: A dictionary mapping scenario names to Scenario objects. +- `target_dict::Dict{AbstractString, Target}`: A dictionary mapping target names to Target objects. +- `latent_model_dict::Dict{AbstractString, LatentModel}`: A dictionary mapping latent model names to LatentModel objects. + +# Optional Arguments +- `fig_kws::NamedTuple`: Keyword arguments for the Figure object. +- `true_gi_choice::Float64`: The true value of the infection generating process. +- `used_gi_choice::Float64`: The value of the infection generating process used in the analysis. +- `legend_title::AbstractString`: The title of the legend. +- `targets::Vector{AbstractString}`: The names of the targets to include in the figure. +- `scenarios::Vector{AbstractString}`: The names of the scenarios to include in the figure. + +# Returns +- `fig::Figure`: The figure object containing the panels. + +# Example +This assumes that the user already has the necessary dataframes `truth_df` and `analysis_df` loaded. + +```julia +using EpiAwarePipeline +# Define scenario titles and reference times for figure 1 +scenario_dict = Dict( + "measures_outbreak" => (title = "Outbreak with measures", T = 28), + "smooth_outbreak" => (title = "Outbreak no measures", T = 35), + "smooth_endemic" => (title = "Smooth endemic", T = 35), + "rough_endemic" => (title = "Rough endemic", T = 35) +) + +target_dict = Dict( + "log_I_t" => (title = "log(Incidence)", ylims = (3.5, 6)), + "rt" => (title = "Exp. growth rate", ylims = (-0.1, 0.1)), + "Rt" => (title = "Reproductive number", ylims = (-0.1, 3)) +) + +latent_model_dict = Dict( + "wkly_rw" => (title = "Random walk",), + "wkly_ar" => (title = "AR(1)",), + "wkly_diff_ar" => (title = "Diff. AR(1)",) +) + +fig1 = figureone( + truth_df, analysis_df, "wkly_ar", scenario_dict, target_dict, latent_model_dict) +``` + +""" +function figureone( + truth_df, analysis_df, latent_model, scenario_dict, target_dict, + latent_model_dict; fig_kws = (; size = (1000, 1000)), + true_gi_choice = 10.0, used_gi_choice = 10.0, + legend_title = "Infection generating\n process", + targets = ["log_I_t", "rt", "Rt"], + scenarios = [ + "measures_outbreak", "smooth_outbreak", "smooth_endemic", "rough_endemic"]) + # Perform checks on the dataframes + EpiAwarePipeline._figure_one_dataframe_checks(truth_df, analysis_df, scenario_dict) + latent_models = analysis_df.Latent_Model |> unique + @assert latent_model in latent_models "The latent model is not in the analysis data" + @assert latent_model in keys(latent_model_dict) "The latent model is not in the latent_model_dict dictionary" + @assert all([target in keys(target_dict) for target in targets]) "Not all targets are in the target dictionary" + @assert all([scenario in keys(scenario_dict) for scenario in scenarios]) "Not all scenarios are in the scenario dictionary" + + # Treat the truth data as a Latent model option + truth_df[!, "IGP_Model"] .= "Truth data" + + plt_truth_mat = [_figure_one_truth_data_panel( + truth_df, scenario, target; true_gi_choice) + for scenario in keys(scenario_dict), target in targets] + + plt_analysis_mat = [_figure_one_model_panel( + analysis_df, scenario, target, latent_model; + reference_time = scenario_dict[scenario].T, + true_gi_choice, used_gi_choice) + for scenario in keys(scenario_dict), target in targets] + + fig = Figure(; fig_kws...) + leg = nothing + for (i, scenario) in enumerate(scenarios) + for (j, target) in enumerate(targets) + sf = fig[i, j] + ag = draw!( + sf, plt_analysis_mat[i, j] + plt_truth_mat[i, j], + axis = (; limits = (nothing, target_dict[target].ylims))) + leg = AlgebraOfGraphics.compute_legend(ag) + i == 1 && + Label(sf[0, 1], target_dict[target].title, fontsize = 22, font = :bold) + j == 3 && Label(sf[1, 2], scenario_dict[scenario].title, + fontsize = 18, font = :bold, rotation = -pi / 2) + end + end + + Label(fig[:, 0], "Process values", fontsize = 28, font = :bold, rotation = pi / 2) + Label(fig[3, 4], + "Latent model\n for infection\n generating\n process:\n$(latent_model_dict[latent_model].title)", + fontsize = 18, + font = :bold) + _leg = (leg[1], leg[2], [legend_title]) + Legend(fig[2, 4], _leg...) + + return fig +end