fix and refactor figure 1

pipeline/scripts/create_figure1.jl

@@ -1,8 +1,10 @@
 using EpiAwarePipeline, EpiAware, AlgebraOfGraphics, JLD2, DrWatson, DataFramesMeta,
       Statistics, Distributions, CSV, CairoMakie
 
-## Define scenarios
+## Define scenarios and targets
 scenarios = ["measures_outbreak", "smooth_outbreak", "smooth_endemic", "rough_endemic"]
+targets = ["log_I_t", "rt", "Rt"]
+gi_means = [2.0, 10.0, 20.0]
 
 ## load some data and create a dataframe for the plot
 truth_data_df = CSV.File(plotsdir("plotting_data/truthdata.csv")) |> DataFrame
@@ -17,9 +19,9 @@ scenario_dict = Dict(
 )
 
 target_dict = Dict(
-    "log_I_t" => (title = "log(Incidence)", ylims = (3.5, 12)),
-    "rt" => (title = "Exp. growth rate", ylims = (-0.5, 0.5)),
-    "Rt" => (title = "Reproductive number", ylims = (-0.1, 2.5))
+    "log_I_t" => (title = "log(Incidence)",),
+    "rt" => (title = "Exp. growth rate",),
+    "Rt" => (title = "Reproductive number",)
 )
 
 latent_model_dict = Dict(
@@ -28,16 +30,16 @@ latent_model_dict = Dict(
     "diff_ar" => (title = "Diff. AR(1)",)
 )
 
-##
-targets = ["log_I_t", "rt", "Rt"]
-plt_truth_mat = [EpiAwarePipeline._figure_one_truth_data_panel(
-                     truth_data_df, scenario, target; true_gi_choice = 2.0)
-                 for scenario in keys(scenario_dict), target in targets]
-
 ## `ar` is the default latent model which we show as figure 1, others are for SI
 
-_ = map(latent_model_dict |> keys |> collect) do latent_model
-    fig = figureone(
-        truth_data_df, prediction_df, latent_model, scenario_dict, target_dict, latent_model_dict)
-    save(plotsdir("figure1_$(latent_model).png"), fig)
+_ = mapreduce(vcat, latent_model_dict |> keys |> collect) do latent_model
+    map(Iterators.product(gi_means, gi_means)) do (true_gi_choice, used_gi_choice)
+        fig = figureone(
+            prediction_df, truth_data_df, scenarios, targets; scenario_dict, target_dict,
+            latent_model_dict, latent_model, true_gi_choice, used_gi_choice)
+        # save(plotsdir("figure1_$(latent_model).png"), fig)
+        save(
+            plotsdir("figure1_$(latent_model)_trueGI_$(true_gi_choice)_usedGI_$(used_gi_choice).png"),
+            fig)
+    end
 end