Add functions to format and use ResultStores more effectively (particularly counterfactual-only runs)

Project.toml

@@ -7,6 +7,7 @@ version = "1.4.1"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
+DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
 NetCDF = "30363a11-5582-574a-97bb-aa9a979735b9"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
@@ -17,6 +18,7 @@ YAXArrays = "c21b50f5-aa40-41ea-b809-c0f5e47bfa5c"
 CSV = "0.10"
 DataFrames = "1"
 Dates = "1.11.0"
+DimensionalData = "0.29.1"
 NetCDF = "0.11.8"
 Random = "1"
 Statistics = "1.10.0"

src/ReefModEngine.jl

@@ -54,7 +54,12 @@ export
 
 # IO
 export
-    ResultStore, concat_results!, save_result_store
+    ResultStore,
+    concat_results!,
+    save_result_store,
+    load_result_store,
+    remove_duplicate_reps,
+    concat_separate_reps
 
 # Run reps
 export run_rme

src/ResultStore.jl

@@ -1,4 +1,4 @@
-using CSV, Dates, DataFrames, NetCDF, YAXArrays
+using CSV, Dates, DataFrames, NetCDF, YAXArrays, DimensionalData
 
 using Base: num_bit_chunks
 mutable struct ResultStore
@@ -36,7 +36,7 @@ function save_result_store(dir_name::String, result_store::ResultStore)::Nothing
     mkpath(dir_name)
 
     result_path = joinpath(dir_name, "results.nc")
-    savedataset(result_store.results; path=result_path, driver=:netcdf)
+    savedataset(result_store.results; path=result_path, driver=:netcdf, overwrite=true)
 
     scenario_path = joinpath(dir_name, "scenarios.csv")
     CSV.write(scenario_path, result_store.scenario)
@@ -60,49 +60,49 @@ function create_dataset(start_year::Int, end_year::Int, n_reefs::Int, reps::Int)
         zeros(arr_size...);
         timesteps=start_year:end_year,
         locations=1:n_reefs,
-        scenarios=1:(2 * reps)
+        scenarios=1:(2*reps)
     )
     # DHW [degree heating weeks]
     dhw = DataCube(
         zeros(arr_size...);
         timesteps=start_year:end_year,
         locations=1:n_reefs,
-        scenarios=1:(2 * reps)
+        scenarios=1:(2*reps)
     )
     # DHW mortality [% of population (to be confirmed)]
     dhw_mortality = DataCube(
         zeros(arr_size...);
         timesteps=start_year:end_year,
         locations=1:n_reefs,
-        scenarios=1:(2 * reps)
+        scenarios=1:(2*reps)
     )
     # Cyclone mortality [% of population (to be confirmed)]
     cyc_mortality = DataCube(
         zeros(arr_size...);
         timesteps=start_year:end_year,
         locations=1:n_reefs,
-        scenarios=1:(2 * reps)
+        scenarios=1:(2*reps)
     )
     # Cyclone categories [0 to 5]
     cyc_cat = DataCube(
         zeros(arr_size...);
         timesteps=start_year:end_year,
         locations=1:n_reefs,
-        scenarios=1:(2 * reps)
+        scenarios=1:(2*reps)
     )
     # Crown-of-Thorn Starfish population [per ha]
     cots = DataCube(
         zeros(arr_size...);
         timesteps=start_year:end_year,
         locations=1:n_reefs,
-        scenarios=1:(2 * reps)
+        scenarios=1:(2*reps)
     )
     # Mortality caused by Crown-of-Thorn Starfish [% of population (to be confirmed)]
     cots_mortality = DataCube(
         zeros(arr_size...);
         timesteps=start_year:end_year,
         locations=1:n_reefs,
-        scenarios=1:(2 * reps)
+        scenarios=1:(2*reps)
     )
     # Total Species cover [% of total reef area]
     n_species = 6
@@ -112,7 +112,7 @@ function create_dataset(start_year::Int, end_year::Int, n_reefs::Int, reps::Int)
         timesteps=start_year:end_year,
         locations=1:n_reefs,
         taxa=1:n_species,
-        scenarios=1:(2 * reps)
+        scenarios=1:(2*reps)
     )
 
     return Dataset(
@@ -313,7 +313,7 @@ function append_scenarios!(rs::ResultStore, reps::Int)::Nothing
     )
 
     if size(rs.scenario) == (0, 0)
-        rs.scenario = vcat(df_cf, df_iv);
+        rs.scenario = vcat(df_cf, df_iv)
     else
         rs.scenario = vcat(rs.scenario, df_cf, df_iv)
     end
@@ -414,7 +414,7 @@ function concat_results!(
                 tmp::Ref{Cdouble},
                 n_reefs::Cint
             )::Cint
-            rs.results.dhw_mortality[timesteps=At(yr), scenarios=rep_offset +reps + r] = tmp
+            rs.results.dhw_mortality[timesteps=At(yr), scenarios=rep_offset + reps + r] = tmp
 
             @RME runGetData(
                 "cyclone_loss_pct"::Cstring,
@@ -539,3 +539,201 @@ function concat_results!(
 
     return nothing
 end
+
+"""
+    load_result_store(dir_name::String, n_reps::Int64)::ResultStore
+
+Save ResultStore from saved results.nc and scenarios.csv files to allow modification.
+
+# Arguments
+- `dir_name` : Directory where result store files are held.
+- `n_reps` : The number of reps held in resultstore (should not include duplicate reps for counterfactual-only runs).
+"""
+function load_result_store(dir_name::String, n_reps::Int64)::ResultStore
+    result_path = joinpath(dir_name, "results.nc")
+    results = open_dataset(result_path, driver=:netcdf)
+    start_year = first(results.timesteps)
+    end_year = last(results.timesteps)
+    n_reefs = length(results.locations)
+
+    if n_reps ∉ [length(results.scenarios), length(results.scenarios) / 2]
+        throw("Input n_reps does not match the number of scenarios stored in data file.")
+    end
+
+    scenario_path = joinpath(dir_name, "scenarios.csv")
+    scenario = CSV.read(scenario_path, DataFrame)
+
+    return ResultStore(
+        results,
+        scenario,
+        start_year,
+        end_year,
+        (end_year - start_year) + 1,
+        n_reefs,
+        n_reps
+    )
+end
+
+"""
+    remove_duplicate_reps(result_store::ResultStore, n_reps::Int64)
+
+Find the indices of unique scenarios when there are duplicated scenarios and rebuild
+the scenarios axis in `rebuild_RME_dataset()` to contain only a single copy of unique scenarios.
+"""
+function remove_duplicate_reps(result_store::ResultStore, n_reps::Int64)
+    cover = result_store.results.total_cover
+
+    for year_reef1 in cover.timesteps
+        cover_scen = cover[At(year_reef1), 1, :]
+        if size(unique(cover_scen.data), 1) == n_reps
+            global unique_indices = unique(i -> cover_scen.data[i], 1:length(cover_scen.data))
+            break
+        end
+    end
+
+    result_store.results = rebuild_RME_dataset(
+        result_store.results,
+        first(result_store.results.timesteps),
+        last(result_store.results.timesteps),
+        length(result_store.results.locations),
+        n_reps,
+        unique_indices
+    )
+
+    result_store.scenario = result_store.scenario[unique_indices, :]
+    result_store.reps = n_reps
+
+    return result_store
+end
+
+"""
+    rebuild_RME_dataset(
+        rs_dataset::Dataset,
+        start_year::Int64,
+        end_year::Int64,
+        n_reefs::Int64,
+        n_reps::Int64,
+        unique_indices::Vector{Int64}
+    )
+
+Rebuild a RME dataset that has duplicated scenarios. For example, when RME outputs counterfactual runs with duplicate scenario data.
+
+# Arguments
+- `rs_dataset` : The RME dataset with duplicated scenarios.
+- `start_year` : Start year of timesteps dimension.
+- `end_year` : End year of timesteps dimension.
+- `location_ids` : Location IDs to be held in sites dimension.
+- `n_reps` : The intended number of scenarios that should be in the returned dataset (after removing duplicate scenarios).
+- `unique_indices` : The first index of each unique scenario to keep (excludes indices of duplicate scenarios).
+"""
+function rebuild_RME_dataset(
+    rs_dataset::Dataset,
+    start_year::Int64,
+    end_year::Int64,
+    n_reefs::Int64,
+    n_reps::Int64,
+    unique_indices::Vector{Int64}
+)
+    variable_keys = keys(rs_dataset.cubes)
+
+    arrays = Dict()
+    for variable in variable_keys
+        if variable == :total_taxa_cover
+            axlist = (
+                Dim{:timesteps}(start_year:end_year),
+                Dim{:locations}(1:n_reefs),
+                Dim{:taxa}(1:6),
+                Dim{:scenarios}(1:n_reps)
+            )
+        else
+            axlist = (
+                Dim{:timesteps}(start_year:end_year),
+                Dim{:locations}(1:n_reefs),
+                Dim{:scenarios}(1:n_reps)
+            )
+        end
+
+        # Remove duplicated scenarios
+        yarray = rs_dataset[variable][scenarios=unique_indices]
+        # Rebuild to ensure correct scenario lookup axis.
+        yarray = DimensionalData.rebuild(yarray, dims=axlist)
+        push!(arrays, variable => yarray)
+    end
+
+    return Dataset(; arrays...)
+end
+
+"""
+    concat_separate_reps(results_store_1::ResultStore, result_store_s::ResultStore...)
+
+Concatenate ResultStores that have been saved separately along the `scenarios` axis.
+Intended use: When additional scenarios have been run after saving an initial scenario set.
+All variables and factors such as start_year, end_year, n_reefs must be identical across
+ResultStores.
+"""
+function concat_separate_reps(results_store_1::ResultStore, result_store_s::ResultStore...)
+    stores = [results_store_1, result_store_s...]
+    datasets = [store.results for store in stores]
+    scenarios = [store.scenario for store in stores]
+
+    start_year = results_store_1.start_year
+    end_year = results_store_1.end_year
+    year_range = results_store_1.year_range
+    n_reefs = results_store_1.n_reefs
+
+
+    results = concat_RME_datasets(datasets)
+    scenarios = vcat(scenarios...)
+    reps = size(scenarios, 1)
+
+    return ResultStore(results, scenarios, start_year, end_year, year_range, n_reefs, reps)
+end
+
+"""
+    concat_RME_datasets(datasets::Vector{Dataset})
+
+Combine RME result datasets along the `scenarios` dimension to
+combine scenarios that have been run separately into a single dataset.
+
+# Example
+results_dataset_300scens = concat_RME_netcdfs(
+    results_dataset_200scens,
+    results_dataset_50scens,
+    results_dataset_50scens
+)
+"""
+function concat_RME_datasets(datasets::Vector{Dataset})
+    variable_keys = keys(datasets[1].cubes)
+    arrays = Dict()
+
+    for variable in variable_keys
+        if variable == :total_taxa_cover
+            yarrays = [x[variable] for x in datasets]
+            yarray = YAXArrays.cat(yarrays...; dims=4) # In RME YAXArrays with taxa the 4th dimension is scenarios
+
+            # For some reason after concattenating you need to rebuild the scenario axis
+            axlist = (
+                yarray.axes[1],
+                yarray.axes[2],
+                yarray.axes[3],
+                Dim{:scenarios}(1:size(yarray, 4))
+            )
+            yarray = rebuild(yarray, dims=axlist)
+        else
+            yarrays = [x[variable] for x in datasets]
+            yarray = YAXArrays.cat(yarrays...; dims=3) # In RME YAXArrays without taxa the 3rd dimension is scenarios
+
+            # For some reason after concattenating you need to rebuild the scenario axis
+            axlist = (
+                yarray.axes[1],
+                yarray.axes[2],
+                Dim{:scenarios}(1:size(yarray, 3))
+            )
+            yarray = rebuild(yarray, dims=axlist)
+        end
+
+        push!(arrays, variable => yarray)
+    end
+
+    return Dataset(; arrays...)
+end