From 1368b9096531bf357f4b269f47d5d7ee2ecff059 Mon Sep 17 00:00:00 2001
From: akshaysridhar <asridhar@caltech.edu>
Date: Wed, 6 Mar 2024 18:20:44 -0800
Subject: [PATCH] renamed:    ../../ekp_config.yml -> ekp_config.yml modified: 
  ../../generate_observations.sbatch modified:   ../../observation_map.jl

---
 .../amip_coupled}/ekp_config.yml              |  0
 calibration/generate_observations.sbatch      |  3 +-
 calibration/observation_map.jl                | 34 +++++++++++++++++++
 3 files changed, 36 insertions(+), 1 deletion(-)
 rename calibration/{ => experiments/amip_coupled}/ekp_config.yml (100%)

diff --git a/calibration/ekp_config.yml b/calibration/experiments/amip_coupled/ekp_config.yml
similarity index 100%
rename from calibration/ekp_config.yml
rename to calibration/experiments/amip_coupled/ekp_config.yml
diff --git a/calibration/generate_observations.sbatch b/calibration/generate_observations.sbatch
index 4ef003a605..068be41bba 100644
--- a/calibration/generate_observations.sbatch
+++ b/calibration/generate_observations.sbatch
@@ -27,7 +27,8 @@ include("coupler_component_init.jl")
 include("coupler_driver_calibration.jl");
 solve_coupler!(cs); # Integrate the coupled model
 
-ta = ncread(joinpath(output_dir, "", "ta_150s_average.nc"), "ta")
+testdir = "/Users/akshaysridhar/Research/Codes/ClimaCoupler.jl/calibration/experiments/amip_coupled/truth_simulation/"
+wa = NCDataset(joinpath(testdir, "", "wa_inst.nc"))["wa"]
 include(joinpath(experiment_dir, "observation_map.jl"))
 (; observation, variance) = process_member_data(ta; output_variance = true)
 JLD2.save_object(joinpath(experiment_dir, "obs_mean.jld2"), observation)
diff --git a/calibration/observation_map.jl b/calibration/observation_map.jl
index 71318c8d57..5c2a236a3b 100644
--- a/calibration/observation_map.jl
+++ b/calibration/observation_map.jl
@@ -1,2 +1,36 @@
 ### Place holder for NCEP data from the ClimaCoupler outputs
 #
+
+function observation_map(iteration)
+    experiment_id = "amip_coupled"
+    config = YAML.load_file(joinpath("experiments", experiment_id, "ekp_config.yml"))
+    output_dir = config["output_dir"]
+    ensemble_size = config["ensemble_size"]
+    model_output = "wa_inst.nc"
+    dims = 1
+    G_ensemble = Array{Float64}(undef, dims..., ensemble_size)
+    for m in 1:ensemble_size
+        member_path =
+            TOMLInterface.path_to_ensemble_member(output_dir, iteration, m)
+        ta = ncread(joinpath(member_path, model_output), "wa")
+        G_ensemble[:, m] = process_member_data(ta)
+    end
+    return G_ensemble
+end
+
+function process_member_data(wa; output_variance = false)
+    # Cut off first 120 days to get equilibrium, take second level slice
+    level_slice = 2
+    wa_second_height = wa[3:size(wa)[1], :, :, level_slice]
+    # Average over long and latitude
+    area_avg_wa_second_height =
+        longitudinal_avg(latitudinal_avg(wa_second_height))
+    observation = Float64[area_avg_wa_second_height[3]]
+    if !(output_variance)
+        return observation
+    else
+        variance = Matrix{Float64}(undef, 1, 1)
+        variance[1] = var(area_avg_wa_second_height)
+        return (; observation, variance)
+    end
+end
\ No newline at end of file