Merge pull request #163 from CliMA/kp/postprocess

Add `analyze_iteration` and `postprocess_g_ensemble`

Author: ph-kev

docs/src/api.md

@@ -5,6 +5,8 @@
 ```@docs
 ClimaCalibrate.forward_model
 ClimaCalibrate.observation_map
+ClimaCalibrate.analyze_iteration
+ClimaCalibrate.postprocess_g_ensemble
 ```
 
 ## Worker Interface

docs/src/quickstart.md

@@ -46,6 +46,41 @@ function observation_map(iteration)
     return G_ensemble
 end
 ```
+
+### Optional postprocessing
+
+It may be the case that `observation_map` is insufficient as you need to more information,
+such as information from the `ekp` object to compute `G_ensemble`. Further postprocessing of the
+`G_ensemble` object can be done by implementing the `postprocess_g_ensemble` as shown
+below.
+
+```julia
+function postprocess_g_ensemble(ekp, g_ensemble, prior, output_dir, iteration)
+    return g_ensemble
+end
+```
+
+After each evaluation of the observation map and before updating the ensemble, it may be
+helpful to print the errors from the `ekp` object or plot `G_ensemble`. This can be done
+by implementing the `analyze_iteration` as shown below.
+
+```julia
+function ClimaCalibrate.analyze_iteration(
+    ekp,
+    g_ensemble,
+    prior,
+    output_dir,
+    iteration,
+)
+    @info "Analyzing iteration"
+    @info "Iteration $iteration"
+    @info "Current mean parameter: $(EnsembleKalmanProcesses.get_ϕ_mean_final(prior, ekp))"
+    @info "g_ensemble: $g_ensemble"
+    @info "output_dir: $output_dir"
+    return nothing
+end
+```
+
 ### Parameters
 
 Every parameter that is being calibrated requires a prior distribution to sample from.

src/ekp_interface.jl

@@ -377,8 +377,11 @@ Compute the observation map and update the given EKP object.
 """
 function observation_map_and_update!(ekp, output_dir, iteration, prior)
     g_ensemble = observation_map(iteration)
+    g_ensemble =
+        postprocess_g_ensemble(ekp, g_ensemble, prior, output_dir, iteration)
     save_G_ensemble(output_dir, iteration, g_ensemble)
     terminate = update_ensemble!(ekp, g_ensemble, output_dir, iteration, prior)
+    analyze_iteration(ekp, g_ensemble, prior, output_dir, iteration)
     return terminate
 end
 

src/model_interface.jl

@@ -1,7 +1,7 @@
 import EnsembleKalmanProcesses as EKP
 import YAML
 
-export forward_model, observation_map
+export forward_model, observation_map, analyze_iteration
 
 """
     forward_model(iteration, member)
@@ -24,3 +24,30 @@ This function must be implemented for each calibration experiment.
 function observation_map(iteration)
     error("observation_map not implemented")
 end
+
+"""
+    analyze_iteration(ekp, g_ensemble, prior, output_dir, iteration)
+
+After each evaluation of the observation map and before updating the ensemble,
+`analyze_iteration` is evaluated.
+
+This function is optional to implement.
+
+For example, one may want to print information from the `eki` object or plot
+`g_ensemble`.
+"""
+function analyze_iteration(ekp, g_ensemble, prior, output_dir, iteration)
+    @info "Mean constrained parameter(s): $(EKP.get_ϕ_mean_final(prior, ekp))"
+    @info "Covariance-weighted error: $(last(EKP.get_error(ekp)))"
+    return nothing
+end
+
+"""
+    postprocess_g_ensemble(ekp, g_ensemble, prior, output_dir, iteration)
+
+Postprocess `g_ensemble` after evaluating the observation map and before
+updating the ensemble.
+"""
+function postprocess_g_ensemble(ekp, g_ensemble, prior, output_dir, iteration)
+    return g_ensemble
+end

test/julia_backend.jl

@@ -50,6 +50,15 @@ function CAL.observation_map(iteration)
     return G_ensemble
 end
 
+function CAL.analyze_iteration(ekp, g_ensemble, prior, output_dir, iteration)
+    @info "Analyzing iteration"
+    @info "Iteration $iteration"
+    @info "Current mean constrained parameter: $(EKP.get_ϕ_mean_final(prior, ekp))"
+    @info "g_ensemble: $g_ensemble"
+    @info "output_dir: $output_dir"
+    return nothing
+end
+
 # Test!
 ekp = CAL.calibrate(CAL.JuliaBackend, experiment_config)
 

test/model_interface.jl

@@ -3,8 +3,10 @@ import ClimaCalibrate
 using EnsembleKalmanProcesses.ParameterDistributions
 using Test
 
-# Tests for ensuring ClimaCalibrate has protected interfaces. The API tested below must be defined for each model,
-# otherwise ClimaCalibrate will throw an error.
+# Tests for ensuring ClimaCalibrate has protected interfaces. The API tested
+# below must be defined for each model, otherwise ClimaCalibrate will throw an
+# error. However, `analyze_iteration` and `postprocess_g_ensemble` are
+# optional to implement.
 
 @testset "Model Interface stubs" begin
     @test_throws ErrorException("forward_model not implemented") ClimaCalibrate.forward_model(
@@ -14,4 +16,6 @@ using Test
     @test_throws ErrorException("observation_map not implemented") ClimaCalibrate.observation_map(
         1,
     )
+    @test isnothing(ClimaCalibrate.analyze_iteration(1, 1, 1, 1, 1))
+    @test 2 == ClimaCalibrate.postprocess_g_ensemble(1, 2, 3, 4, 5)
 end