Add ObservationRecipe.reconstruct_vars

docs/src/api.md

@@ -85,6 +85,7 @@ ClimaCalibrate.ObservationRecipe.get_observations_for_nth_iteration
 ClimaCalibrate.ObservationRecipe.get_metadata_for_nth_iteration
 ClimaCalibrate.ObservationRecipe.reconstruct_g_mean_final
 ClimaCalibrate.ObservationRecipe.reconstruct_diag_cov
+ClimaCalibrate.ObservationRecipe.reconstruct_vars
 ClimaCalibrate.ObservationRecipe.seasonally_aligned_yearly_sample_date_ranges
 ClimaCalibrate.ObservationRecipe.change_data_type
 ```

docs/src/observation_recipe.md

@@ -108,9 +108,33 @@ can be used with `ClimaAnalysis.unflatten` to reconstruct the original
 [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).
+## Debugging observational and simulation data
+
+When setting up a calibration, it may be helpful to visualize the
+`EKP.Observation`s or inspect the observational data and metadata together. To
+help with this, `ObservationRecipe` provides the functions
+[`ObservationRecipe.reconstruct_diag_cov`](@ref
+ClimaCalibrate.ObservationRecipe.reconstruct_diag_cov) and
+[`ObservationRecipe.reconstruct_vars`](@ref
+ClimaCalibrate.ObservationRecipe.reconstruct_vars). The function
+`reconstruct_diag_cov` reconstructs the diagonal of the covariance matrix as a
+vector of `OutputVar`s, and the function `reconstruct_vars` reconstructs the
+samples of the `EKP.Observation` as a vector of `OutputVar`s. To visualize the
+data in the `OutputVar`s, see the [visualize
+section](https://clima.github.io/ClimaAnalysis.jl/dev/visualize/) in
+ClimaAnalysis.
+
+```julia
+ObservationRecipe.reconstruct_vars(obs)
+# Reconstructing the diagnonal of a covariance matrix as a OutputVar is only
+# supported for diagonal covariance matrices
+ObservationRecipe.reconstruct_diag_cov(obs)
+```
+
+Finally, `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). This can be helpful
+when debugging whether the G ensemble matrix is formed correctly or not.
 
 ## Frequently asked questions
 
@@ -121,7 +145,10 @@ ClimaCalibrate.ObservationRecipe.reconstruct_g_mean_final).
 [`ClimaAnalysis.flatten`](https://clima.github.io/ClimaAnalysis.jl/dev/flat/#Flatten)
 in the ClimaAnalysis documentation for more information. The order of the
 variables in the observation is the same as the order of the `OutputVar`s when
-creating the `EKP.Observation` using `ObservationRecipe.observation`.
+creating the `EKP.Observation` using `ObservationRecipe.observation`. If you are
+using `ObservationRecipe`, it is recommended that you also use
+[`GEnsembleBuilder`](ensemble_builder.md) which simplifies building the
+G ensemble matrix.
 
 **Q: How do I handle `NaN`s in the `OutputVar`s so that there are no `NaN`s in the sample and covariance matrix?**
 

ext/observation_recipe.jl

@@ -650,6 +650,30 @@ function ObservationRecipe.reconstruct_diag_cov(obs::EKP.Observation)
     return vars
 end
 
+"""
+    reconstruct_vars(obs::EKP.Observation)
+
+Reconstruct the `OutputVar`s from the `samples` in `obs`.
+"""
+function ObservationRecipe.reconstruct_vars(obs::EKP.Observation)
+    all_metadata = EKP.get_metadata(obs)
+    samples = EKP.get_samples(obs)
+    stacked_sample = reduce(vcat, samples)
+
+    start_index = 1
+    vars = OutputVar[]
+    for metadata in all_metadata
+        data_size = ClimaAnalysis.flattened_length(metadata)
+        var = ClimaAnalysis.unflatten(
+            metadata,
+            view(stacked_sample, start_index:(start_index + data_size - 1)),
+        )
+        push!(vars, var)
+        start_index += data_size
+    end
+    return vars
+end
+
 """
     _get_minibatch_indices_for_nth_iteration(ekp::EKP.EnsembleKalmanProcess, N)
 

src/observation_recipe.jl

@@ -10,7 +10,9 @@ export ScalarCovariance,
     change_data_type,
     get_observations_for_nth_iteration,
     get_metadata_for_nth_iteration,
-    reconstruct_g_mean_final
+    reconstruct_g_mean_final,
+    reconstruct_diag_cov,
+    reconstruct_vars
 
 import Dates
 
@@ -315,6 +317,8 @@ function reconstruct_g_mean_final end
 
 function reconstruct_diag_cov end
 
+function reconstruct_vars end
+
 function get_observations_for_nth_iteration end
 
 function get_metadata_for_nth_iteration end

test/observation_recipe.jl

@@ -1531,7 +1531,7 @@ end
     end
 end
 
-@testset "Reconstruct diagonal of covariance" begin
+@testset "Reconstruct diagonal of covariance and OutputVars from observations" begin
     time =
         ClimaAnalysis.Utils.date_to_time.(
             Dates.DateTime(2007, 12),
@@ -1589,6 +1589,38 @@ end
     @test isequal(time_var_from_covs1.data, time_var_from_covs2.data)
     @test isequal(lat_var_from_covs2.data, lat_var_from_covs3.data)
 
+    # Reconstruct OutputVars from observations
+    vars1 = ObservationRecipe.reconstruct_vars(obs1)
+    vars2 = ObservationRecipe.reconstruct_vars(obs2)
+    vars3 = ObservationRecipe.reconstruct_vars(obs3)
+
+    @test length(vars1) == 2
+    @test length(vars2) == 1
+    @test length(vars3) == 1
+
+    # Check dimensions
+    @test vars1[1].dims["time"] == time[1:8]
+    @test vars1[2].dims["lon"] == lon
+    @test vars1[2].dims["time"] == time[1:8]
+    @test vars2[1].dims["time"] == time[1:8]
+    @test vars3[1].dims["time"] == time[1:8]
+    @test vars3[1].dims["lon"] == lon
+
+    # Check attributes and dimension attributes
+    for (reconstructed_var, original_var) in zip(
+        (vars1[1], vars1[2], vars2[1], vars3[1]),
+        (time_var, lon_var, time_var, lon_var),
+    )
+        @test reconstructed_var.attributes == original_var.attributes
+        @test reconstructed_var.dim_attributes == original_var.dim_attributes
+    end
+
+    # Check data
+    @test vars1[1].data == time_var.data[1:8]
+    @test vars1[2].data == lon_var.data[:, 1:8]
+    @test vars2[1].data == time_var.data[1:8]
+    @test vars3[1].data == permutedims(lon_var.data[:, 1:8], (2, 1))
+
     var3d =
         TemplateVar() |>
         add_dim("time", time, units = "s") |>
@@ -1622,4 +1654,14 @@ end
     @test !isequal(var_from_covs4.data[:, 1, :], var_from_covs4.data[:, 2, :])
     @test !isequal(var_from_covs4.data[:, 2, :], var_from_covs4.data[:, 3, :])
     @test !isequal(var_from_covs4.data[:, 1, :], var_from_covs4.data[:, 3, :])
+
+    # Test reconstruction of obs4
+    vars4 = ObservationRecipe.reconstruct_vars(obs4)
+    @test length(vars4) == 1
+    @test vars4[1].attributes == var3d.attributes
+    @test vars4[1].dim_attributes == var3d.dim_attributes
+    @test vars4[1].dims["time"] == [0.0]
+    @test vars4[1].dims["lat"] == var3d.dims["lat"]
+    @test vars4[1].dims["lon"] == var3d.dims["lon"]
+    @test vars4[1].data == var3d.data[[1], :, :]
 end