Merge pull request #197 from CliMA/kp/reconstruct

Add ObservationRecipe.reconstruct_g_mean_final

Author: ph-kev

docs/src/api.md

@@ -80,6 +80,9 @@ ClimaCalibrate.ObservationRecipe.SVDplusDCovariance
 ClimaCalibrate.ObservationRecipe.SVDplusDCovariance(sample_dates)
 ClimaCalibrate.ObservationRecipe.covariance
 ClimaCalibrate.ObservationRecipe.observation
+ClimaCalibrate.ObservationRecipe.short_names
+ClimaCalibrate.ObservationRecipe.get_metadata_for_nth_iteration
+ClimaCalibrate.ObservationRecipe.reconstruct_g_mean_final
 ClimaCalibrate.ObservationRecipe.seasonally_aligned_yearly_sample_date_ranges
 ClimaCalibrate.ObservationRecipe.change_data_type
 ```

docs/src/observation_recipe.md

@@ -88,6 +88,29 @@ obs = ObservationRecipe.observation(
 )
 ```
 
+## Metadata
+
+!!! note
+    Metadata in `EKP.observation` is only added with versions of
+    EnsembleKalmanProcesses later than v2.4.2.
+"""
+
+When creating an observation with [`ObservationRecipe.observation`](@ref
+ClimaCalibrate.ObservationRecipe.observation), metadata is extracted from the
+`OutputVar`s and attached to the observation. The metadata for each observation
+can be accessed with `EKP.get_metadata(obs::EKP.Observation)` and the metadata
+for each iteration can be accessed with
+[`ObservationRecipe.get_metadata_for_nth_iteration`](@ref
+ClimaCalibrate.ObservationRecipe.get_metadata_for_nth_iteration). The metadata
+can be used with `ClimaAnalysis.unflatten` to reconstruct the original
+`OutputVar` before flattening. See the ClimaAnalysis
+[documentation](https://clima.github.io/ClimaAnalysis.jl/dev/api/#FlatVar) about
+`ClimaAnalysis.FlatVar` for more information.
+
+`ObservationRecipe` provides a helper function for reconstructing the mean
+forward map evaluation with [`ObservationRecipe.reconstruct_g_mean_final`](@ref
+ClimaCalibrate.ObservationRecipe.reconstruct_g_mean_final).
+
 ## Frequently asked questions
 
 **Q: I need to compute `g_ensemble` and I do not know how the data of the `OutputVar`s is flattened.**

ext/ClimaAnalysisExt.jl

@@ -506,4 +506,63 @@ function _lat_weights_var(var::OutputVar; min_cosd_lat = 0.1)
     return ClimaAnalysis.remake(var, data = lat_weights)
 end
 
+"""
+    get_metadata_for_nth_iteration(obs_series, N)
+
+For the `N`th iteration, get the metadata of the observation(s) being processed.
+"""
+function ObservationRecipe.get_metadata_for_nth_iteration(obs_series, N)
+    num_epoches = EKP.get_length_epoch(obs_series)
+    # EKP.get_minibatch fails with N > num_epoches, so we use mod1 to go back to
+    # the first epoch which seems consistent with what EKP does
+    minibatch_indices = EKP.get_minibatch(obs_series, mod1(N, num_epoches))
+    minibatch_obs = EKP.get_observations(obs_series)[minibatch_indices]
+    metadata_vec = map(obs -> EKP.get_metadata(obs), minibatch_obs)
+    return vcat(metadata_vec...)
+end
+
+"""
+    reconstruct_g_mean_final(ekp::EKP.EnsembleKalmanProcess,
+                             observation::EKP.Observation)
+
+Reconstruct the mean forward model evaluation at the last iteration as a
+vector of `OutputVar`s.
+
+This function assumes `observation` contains the necessary metadata to reconstruct
+the `OutputVar`s. Note that the metadata comes from the observations.
+"""
+function ObservationRecipe.reconstruct_g_mean_final(
+    ekp::EKP.EnsembleKalmanProcess,
+)
+    obs_series = EKP.get_observation_series(ekp)
+    metadatas = ObservationRecipe.get_metadata_for_nth_iteration(
+        obs_series,
+        EKP.get_N_iterations(ekp),
+    )
+    eltype(metadatas) <: ClimaAnalysis.Var.Metadata || error(
+        "Getting the short names from an observation is only supported with metadata from ClimaAnalysis",
+    )
+    g_mean = EKP.get_g_mean_final(ekp)
+
+    # Check if length of g ensemble is the same as the length of the data in the metadatas
+    total_metadata_length =
+        sum(ClimaAnalysis.Var._data_size(metadata) for metadata in metadatas)
+    length(g_mean) != total_metadata_length && error(
+        "Length of g_mean_final is not the same as the length of all the metadata",
+    )
+
+    # Reconstruct each OutputVar from the metadata
+    index = Ref(1)
+    vars = map(metadatas) do metadata
+        data_size = ClimaAnalysis.Var._data_size(metadata)
+        start_idx = index[]
+        index[] += data_size
+        ClimaAnalysis.unflatten(
+            metadata,
+            g_mean[start_idx:(start_idx + data_size - 1)],
+        )
+    end
+    return vars
+end
+
 end

src/observation_recipe.jl

@@ -7,7 +7,9 @@ export ScalarCovariance,
     observation,
     short_names,
     seasonally_aligned_yearly_sample_date_ranges,
-    change_data_type
+    change_data_type,
+    get_metadata_for_nth_iteration,
+    reconstruct_g_mean_final
 
 import Dates
 
@@ -282,4 +284,8 @@ function seasonally_aligned_yearly_sample_date_ranges end
 
 function change_data_type end
 
+function reconstruct_g_mean_final end
+
+function get_metadata_for_nth_iteration end
+
 end