From 6d23b6a43f61bb4588390bc4a4e7347001f8e18d Mon Sep 17 00:00:00 2001
From: Pedro Ribeiro de Almeida <pedroribalmeida@gmail.com>
Date: Fri, 22 Mar 2024 18:36:34 +1100
Subject: [PATCH] Add counterfactual minus guided/unguided diff at location
 level

---
 ext/AvizExt/viz/spatial.jl                | 20 ++++++++++--
 src/metrics/metrics.jl                    |  4 +--
 src/metrics/{site_level.jl => spatial.jl} | 40 +++++++++++++++++++++++
 src/viz/viz.jl                            |  1 +
 4 files changed, 60 insertions(+), 5 deletions(-)
 rename src/metrics/{site_level.jl => spatial.jl} (74%)

diff --git a/ext/AvizExt/viz/spatial.jl b/ext/AvizExt/viz/spatial.jl
index 25e55937d..f070f307a 100644
--- a/ext/AvizExt/viz/spatial.jl
+++ b/ext/AvizExt/viz/spatial.jl
@@ -37,7 +37,7 @@ function create_map!(
     color_map::Union{Symbol, Vector{Symbol}, RGBA{Float32}, Vector{RGBA{Float32}}}=:grayC,
     legend_params::Union{Tuple, Nothing}=nothing,
     axis_opts::Dict=Dict(),
-)
+)::Union{GridLayout,GridPosition}
     axis_opts[:title] = get(axis_opts, :title, "Study Area")
     axis_opts[:xlabel] = get(axis_opts, :xlabel, "Longitude")
     axis_opts[:ylabel] = get(axis_opts, :ylabel, "Latitude")
@@ -58,10 +58,11 @@ function create_map!(
 
     spatial.yticklabelpad = 50
     spatial.ytickalign = 10
+    min_val = @lift(minimum($data))
     max_val = @lift(maximum($data))
 
     # Plot geodata polygons using data as internal color
-    color_range = (0.0, max_val[])
+    color_range = min_val[] < 0 ? (min_val[], max_val[]) : (0, max_val[])
 
     poly!(
         spatial,
@@ -215,6 +216,21 @@ function ADRIA.viz.map!(
     )
 end
 
+function ADRIA.viz.diff_map(
+    rs::ResultSet,
+    diff_outcome::YAXArray{Float64};
+    opts::Dict{Symbol,<:Any}=Dict{Symbol,Any}(),
+    fig_opts::Dict{Symbol,<:Any}=Dict{Symbol,Any}(),
+    axis_opts::Dict{Symbol,<:Any}=Dict{Symbol,Any}()
+)
+    # TODO hande the cases where thes only positive or only negative values
+    min_val, max_res = extrema(diff_outcome)
+    mid_val = -min_val / (max_res - min_val)
+    div_cmap::Vector{RGBA{Float32}} = diverging_palette(10, 200; mid=mid_val)
+    opts[:color_map] = div_cmap
+    return ADRIA.viz.map(rs, diff_outcome; axis_opts=axis_opts, opts=opts, fig_opts=fig_opts)
+end
+
 """
     make_geojson_copy(ds::Union{ResultSet,Domain})::String
 
diff --git a/src/metrics/metrics.jl b/src/metrics/metrics.jl
index cc5c6a04b..863f215eb 100644
--- a/src/metrics/metrics.jl
+++ b/src/metrics/metrics.jl
@@ -15,7 +15,6 @@ using DataFrames
 
 using ADRIA: coral_spec, colony_mean_area, ResultSet, timesteps, site_k_area, site_area
 
-
 abstract type Outcome end
 
 
@@ -608,12 +607,11 @@ function _relative_shelter_volume(rs::ResultSet)::YAXArray
 end
 relative_shelter_volume = Metric(_relative_shelter_volume, (:timesteps, :sites, :scenarios))
 
-
 include("pareto.jl")
 include("ranks.jl")
 include("reef_indices.jl")
 include("scenario.jl")
-include("site_level.jl")
+include("spatial.jl")
 include("temporal.jl")
 include("utils.jl")
 
diff --git a/src/metrics/site_level.jl b/src/metrics/spatial.jl
similarity index 74%
rename from src/metrics/site_level.jl
rename to src/metrics/spatial.jl
index 5c5ba94a1..bb9c436bf 100644
--- a/src/metrics/site_level.jl
+++ b/src/metrics/spatial.jl
@@ -1,3 +1,6 @@
+using Bootstrap
+using Random
+
 """Functions and methods to produce location-level summaries."""
 
 """
@@ -106,3 +109,40 @@ function summarize(
 )::YAXArray where {D,T,N,A}
     return summarize(data[timesteps=timesteps], alongs_axis, metric)
 end
+
+"""
+    cf_difference_map(rs, scens, metric)
+
+Give the mean bootstraped difference from the counterfactual for a given metric.
+"""
+function cf_difference_map(rs, scens, metric)
+    # Extract some metric
+    outcomes = metric(rs)
+    locations = rs.site_data
+
+    # Mean over all timesteps
+    outcomes_mean = dropdims(mean(outcomes, dims=:timesteps), dims=:timesteps)
+
+    cf_outcomes = outcomes_mean[scenarios=scens.guided .== -1]
+    ug_outcomes = outcomes_mean[scenarios=scens.guided .== 0]
+    gd_outcomes = outcomes_mean[scenarios=scens.guided .> 0]
+
+    n_locs = length(outcomes.sites)
+
+    ug_result = ZeroDataCube(; T=Float64, sites=locations.reef_siteid)
+    gd_result = ZeroDataCube(; T=Float64, sites=locations.reef_siteid)
+    n_scens = size(cf_outcomes, :scenarios)
+    for loc in 1:n_locs
+        cf_shuf_set1 = shuffle(1:n_scens)
+        ug_shuf_set = shuffle(1:n_scens)
+        ug_diff = collect(ug_outcomes[loc, ug_shuf_set] .- cf_outcomes[loc, cf_shuf_set1])
+        ug_result[loc] = bootstrap(median, ug_diff, BalancedSampling(100)).t0[1]
+
+        cf_shuf_set2 = shuffle(1:n_scens)
+        gd_shuf_set = shuffle(1:n_scens)
+        gd_diff = collect(gd_outcomes[loc, gd_shuf_set] .- cf_outcomes[loc, cf_shuf_set2])
+        gd_result[loc] = bootstrap(median, gd_diff, BalancedSampling(100)).t0[1]
+    end
+
+    return ug_result, gd_result
+end
diff --git a/src/viz/viz.jl b/src/viz/viz.jl
index 5cdebd5dc..66ce1efee 100644
--- a/src/viz/viz.jl
+++ b/src/viz/viz.jl
@@ -42,5 +42,6 @@ function ranks_to_frequencies!() end
 # Spatial
 function map() end
 function map!() end
+function diff_map() end
 
 end  # module