Merge pull request #724 from open-AIMS/iv-labels

Save right intervention label at ResultSet ranks YAXArray

src/interventions/Interventions.jl

@@ -154,3 +154,7 @@ Base.@kwdef struct Intervention <: EcoModel
         description="Start of fogging deployments after this number of years has elapsed.",
     )
 end
+
+function interventions()
+    return [:seed, :fog]
+end

src/io/ResultSet.jl

@@ -73,14 +73,25 @@ function ResultSet(
         input_set.attrs["sim_constants"],
         model_spec,
         outcomes,
-        DataCube(log_set["rankings"], Symbol.(Tuple(log_set["rankings"].attrs["structure"]))),
+        _rankings_data(log_set["rankings"]),
         DataCube(log_set["seed"], Symbol.(Tuple(log_set["seed"].attrs["structure"]))),
         DataCube(log_set["fog"], Symbol.(Tuple(log_set["fog"].attrs["structure"]))),
         DataCube(log_set["shade"], Symbol.(Tuple(log_set["shade"].attrs["structure"]))),
         DataCube(log_set["coral_dhw_log"], Symbol.(Tuple(log_set["coral_dhw_log"].attrs["structure"]))),
     )
 end
 
+function _rankings_data(rankings_set::ZArray{T})::YAXArray{T} where {T}
+    ax_names = Symbol.(Tuple(rankings_set.attrs["structure"]))
+    ax_labels::Vector{Union{UnitRange{Int64},Vector{Symbol}}} = range.([1], size(rankings_set))
+
+    # Replace intervention
+    intervention_idx = findfirst(x -> x == :intervention, ax_names)
+    ax_labels[intervention_idx] = interventions()
+
+    return DataCube(rankings_set; NamedTuple{ax_names}(ax_labels)...)
+end
+
 """
     _copy_env_data(src::String, dst::String, subdir::String)::Nothing
 

src/io/result_io.jl

@@ -197,7 +197,8 @@ function setup_logs(z_store, unique_sites, n_scens, tf, n_sites)
     log_fn::String = joinpath(z_store.folder, LOG_GRP)
 
     # Store ranked sites
-    rank_dims::Tuple{Int64,Int64,Int64,Int64} = (tf, n_sites, 2, n_scens)  # sites, site id and rank, no. scenarios
+    n_interventions = length(interventions())
+    rank_dims::Tuple{Int64,Int64,Int64,Int64} = (tf, n_sites, n_interventions, n_scens)  # sites, site id and rank, no. scenarios
     fog_dims::Tuple{Int64,Int64,Int64} = (tf, n_sites, n_scens)  # timeframe, sites, no. scenarios
 
     # tf, no. species to seed, site id and rank, no. scenarios

src/scenario.jl

@@ -196,9 +196,9 @@ Run individual scenarios for a given domain, saving results to a Zarr data store
 Results are stored in Zarr format at a pre-configured location.
 Sets up a new `cache` if not provided.
 
-# Notes
-Logs of site ranks only store the mean site rankings over all environmental scenarios.
-This is to reduce the volume of data stored.
+# Arguments
+- `domain` : Domain
+- `idx` : Scenario index
 
 # Returns
 Nothing
@@ -320,9 +320,6 @@ end
 
 Core scenario running function.
 
-# Notes
-Only the mean site rankings are kept
-
 # Returns
 NamedTuple of collated results
 """
@@ -417,11 +414,11 @@ function run_model(domain::Domain, param_set::YAXArray)::NamedTuple
     # Avoid placing importance on sites that were not considered
     # Lower values are higher importance/ranks.
     # Values of n_locs+1 indicate locations that were not considered in rankings.
-    log_location_ranks = DataCube(
-        fill(0.0, tf, n_locs, 2);  # log seeding/fogging ranks
+    log_location_ranks = ZeroDataCube(;     # log seeding/fogging ranks
+        T=Float64,
         timesteps=1:tf,
         locations=domain.site_ids,
-        intervention=[:seed, :fog]
+        intervention=interventions()
     )
 
     Yshade = SparseArray(spzeros(tf, n_locs))