From d407c36e2f4dee55ce09e9e1a8c78dc9d72ee4a2 Mon Sep 17 00:00:00 2001
From: Pedro Ribeiro de Almeida <pedroribalmeida@gmail.com>
Date: Fri, 13 Sep 2024 12:41:33 +1000
Subject: [PATCH] Add cf_difference_loc

---
 src/metrics/metrics.jl                    |  2 +-
 src/metrics/{site_level.jl => spatial.jl} | 63 +++++++++++++++++++++++
 test/metrics/spatial.jl                   | 27 ++++++++++
 test/runtests.jl                          |  1 +
 4 files changed, 92 insertions(+), 1 deletion(-)
 rename src/metrics/{site_level.jl => spatial.jl} (61%)
 create mode 100644 test/metrics/spatial.jl

diff --git a/src/metrics/metrics.jl b/src/metrics/metrics.jl
index 87777972e..672d791b3 100644
--- a/src/metrics/metrics.jl
+++ b/src/metrics/metrics.jl
@@ -696,7 +696,7 @@ 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 61%
rename from src/metrics/site_level.jl
rename to src/metrics/spatial.jl
index 6d32e9995..2b132cfdf 100644
--- a/src/metrics/site_level.jl
+++ b/src/metrics/spatial.jl
@@ -1,5 +1,8 @@
 """Functions and methods to produce location-level summaries."""
 
+using Bootstrap
+using Random
+
 """
     per_loc(metric, data::YAXArray{D,T,N,A})::YAXArray where {D,T,N,A}
 
@@ -107,3 +110,63 @@ function summarize(
 )::YAXArray where {D,T,N,A}
     return summarize(data[timesteps=timesteps], alongs_axis, metric)
 end
+
+"""
+    cf_difference_loc(rs, scens, metric)
+
+Mean bootstrapped differences (counterfactual - guided) and (counterfactual - unguided) for
+each location.
+
+# Arguments
+- `outcome` : Metric outcome with dimensions (:timesteps, :locations, :scenarios)
+- `scens` : Scenarios DataFrame
+
+# Returns
+Two elements tuple with mean bootstrapped difference (counterfactual - guided) and
+(counterfactual - unguided) for each location.
+"""
+function cf_difference_loc(
+    outcome::YAXArray{T,3}, scens::DataFrame; conf::Float64=0.95
+)::Tuple where {T}
+    # Mean over all timesteps
+    outcomes_agg = dropdims(mean(outcome; dims=:timesteps); dims=:timesteps)
+
+    # Counterfactual, guided and unguided outcomes
+    cf_outcomes = outcomes_agg[scenarios=scens.guided .== -1]
+    gd_outcomes = outcomes_agg[scenarios=scens.guided .> 0]
+    ug_outcomes = outcomes_agg[scenarios=scens.guided .== 0]
+
+    # Find the smallest set of outcomes because they might not have the same size
+    n_cf_outcomes = size(cf_outcomes, :scenarios)
+    n_gd_outcomes = size(gd_outcomes, :scenarios)
+    n_ug_outcomes = size(ug_outcomes, :scenarios)
+    min_n_outcomes = min(n_cf_outcomes, n_gd_outcomes, n_ug_outcomes)
+
+    # Build (counterfactual-guided) and (counterfactual-unguided) result DataCubes
+    _locations = axis_labels(outcome, :locations)
+    gd_result = ZeroDataCube(;
+        T=T, value=[:lower_bound, :value, :upper_bound], locations=_locations
+    )
+    ug_result = ZeroDataCube(;
+        T=T, value=[:lower_bound, :value, :upper_bound], locations=_locations
+    )
+
+    n_locs = length(_locations)
+    for loc in 1:n_locs
+        cf_shuf_set::Vector{Int64} = shuffle(1:n_cf_outcomes)[1:min_n_outcomes]
+        gd_shuf_set::Vector{Int64} = shuffle(1:n_gd_outcomes)[1:min_n_outcomes]
+        ug_shuf_set::Vector{Int64} = shuffle(1:n_ug_outcomes)[1:min_n_outcomes]
+
+        # Counterfactual - Guided
+        gd_diff = collect(gd_outcomes[loc, gd_shuf_set] .- cf_outcomes[loc, cf_shuf_set])
+        cf_gd_bootstrap = bootstrap(median, gd_diff, BalancedSampling(100))
+        gd_result[[2, 1, 3], loc] .= confint(cf_gd_bootstrap, PercentileConfInt(conf))[1]
+
+        # Counterfactual - Unguided
+        ug_diff = collect(ug_outcomes[loc, ug_shuf_set] .- cf_outcomes[loc, cf_shuf_set])
+        cf_ug_bootstrap = bootstrap(median, ug_diff, BalancedSampling(100))
+        ug_result[[2, 1, 3], loc] .= confint(cf_ug_bootstrap, PercentileConfInt(conf))[1]
+    end
+
+    return gd_result, ug_result
+end
diff --git a/test/metrics/spatial.jl b/test/metrics/spatial.jl
new file mode 100644
index 000000000..817bafdad
--- /dev/null
+++ b/test/metrics/spatial.jl
@@ -0,0 +1,27 @@
+using ADRIA: metrics, metrics.cf_difference_loc
+using ADRIA: YAXArray
+
+if !@isdefined(TEST_RS)
+    const TEST_DOM, TEST_N_SAMPLES, TEST_SCENS, TEST_RS = test_rs()
+end
+
+@testset "cf_difference_loc" begin
+    test_metrics = [
+        metrics.relative_cover,
+        metrics.juvenile_indicator,
+        metrics.absolute_shelter_volume,
+        metrics.relative_juveniles,
+        metrics.coral_evenness,
+        metrics.relative_shelter_volume
+    ]
+
+    for metric in test_metrics
+        gd_diff::YAXArray{<:Real,2}, ug_diff::YAXArray{<:Real,2} = cf_difference_loc(
+            metric(TEST_RS), TEST_SCENS
+        )
+
+        for diff in [gd_diff, ug_diff]
+            @test all(diff[1, :] .<= diff[2, :] .<= diff[3, :])
+        end
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 9fa5c8692..85a9caf74 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -76,6 +76,7 @@ include("seeding.jl")
 include("spec.jl")
 include("mcda.jl")
 include("metrics/metrics.jl")
+include("metrics/spatial.jl")
 include("utils/scale.jl")
 include("utils/text_display.jl")
 include("viz.jl")