2-year minibatch w/ 3D fields

Author: nefrathenrici

experiments/calibration/Project.toml

@@ -31,6 +31,7 @@ YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6"
 
 [compat]
 ClimaCalibrate = "0.0.11"
+ClimaLand = "=0.15.9"
 ClimaTimeSteppers = "0.8.2"
 ClimaCore = "0.14.26"
-EnsembleKalmanProcesses = "2.0"
+EnsembleKalmanProcesses = "2.1.2"

experiments/calibration/coarse_amip/generate_observations.jl

@@ -1,16 +1,16 @@
 # Generate and save experiment observations to disk
-using ClimaAnalysis, JLD2
-include("experiments/calibration/coarse_amip/observation_utils.jl")
+using ClimaAnalysis, JLD2, ClimaCoupler
+include(joinpath(pkgdir(ClimaCoupler),"experiments/calibration/coarse_amip/observation_utils.jl"))
 
 const obs_dir = "/home/ext_nefrathe_caltech_edu/calibration_obs"
-const diagnostic_dir = "experiments/calibration/output/old/iteration_000/member_001/model_config/output_0000/clima_atmos/"
+const simdir = SimDir(joinpath(pkgdir(ClimaCoupler),"experiments/calibration/output/iteration_000/member_001/model_config/"))
 
-diagnostic_var2d = OutputVar(joinpath(diagnostic_dir, "rsdt_1M_average.nc"));
-pressure = OutputVar(joinpath(diagnostic_dir, "pfull_1M_average.nc"));
-diagnostic_var3d = OutputVar(joinpath(diagnostic_dir, "ta_1M_average.nc"));
+diagnostic_var2d = get_monthly_averages(simdir, "rsut")
+pressure = get_monthly_averages(simdir, "pfull")
+diagnostic_var3d = get_monthly_averages(simdir, "ta")
 diagnostic_var3d = ClimaAnalysis.Atmos.to_pressure_coordinates(diagnostic_var3d, pressure)
 
 nt = get_all_output_vars(obs_dir, diagnostic_var2d, diagnostic_var3d)
 nyears = 18
 observation_vec = create_observation_vector(nt, nyears)
-JLD2.save_object("experiments/calibration/coarse_amip/observations.jld2", observation_vec)
+JLD2.save_object("experiments/calibration/coarse_amip/observations_3d.jld2", observation_vec)

experiments/calibration/coarse_amip/model_config.yml

@@ -1,5 +1,5 @@
 FLOAT_TYPE: "Float32"
-use_itime: true
+# use_itime: true
 albedo_model: "CouplerAlbedo"
 atmos_config_file: "config/longrun_configs/amip_target_edonly.yml"
 checkpoint_dt: "99999days"

experiments/calibration/coarse_amip/model_interface.jl

@@ -15,12 +15,12 @@ function ClimaCalibrate.forward_model(iter, member)
     output_dir_root = config_dict["coupler_output_dir"]
     eki = JLD2.load_object(ClimaCalibrate.ekp_path(output_dir_root, iter))
     minibatch = EKP.get_current_minibatch(eki)
-    @show minibatch
+    @info "Current minibatch: $minibatch"
     config_dict["start_date"] = minibatch_to_start_date(minibatch)
 
-    spinup_time = 93days
+    spinup_days = 92
     nyears = length(minibatch)
-    t_end_days = spinup_time + 365 * nyears
+    t_end_days = spinup_days + 365 * nyears
     config_dict["t_end"] = "$(t_end_days)days"
 
     # Set member parameter file

experiments/calibration/coarse_amip/observation_map.jl

@@ -10,9 +10,9 @@ function ClimaCalibrate.observation_map(iteration)
     single_member_dims = length(EKP.get_obs(first(observation_vec))) * batch_size
     G_ensemble = Array{Float64}(undef, single_member_dims, ensemble_size)
     for m in 1:ensemble_size
-        @info "Processing member $m"
         member_path = ClimaCalibrate.path_to_ensemble_member(output_dir, iteration, m)
         simdir_path = joinpath(member_path, "model_config/output_active")
+        @info "Processing member $m: $simdir_path"
         try
             G_ensemble[:, m] .= process_member_data(SimDir(simdir_path))
 
@@ -45,7 +45,7 @@ function process_member_data(simdir::SimDir)
     ta = process_outputvar(simdir, "ta")
     hur = process_outputvar(simdir, "hur")
     hus = process_outputvar(simdir, "hus")
-
+    # Map over each year
     year_observations = map(1:4:length(rsut)) do year_start
         year_end = min(year_start + 3, length(rsut))
         yr_ind = year_start:year_end
@@ -62,7 +62,7 @@ function process_member_data(simdir::SimDir)
         hur_yr = downsample_and_vectorize(hur[yr_ind])
         hus_yr = downsample_and_vectorize(hus[yr_ind])
         # ql, qi
-        vcat(net_rad_yr, rsut_yr, rlut_yr, cre_yr, pr_yr, ts_yr)#, ta, hur, hus)
+        vcat(net_rad_yr, rsut_yr, rlut_yr, cre_yr, pr_yr, ts_yr, ta_yr, hur_yr, hus_yr)
     end
     return vcat(year_observations...)
 end
@@ -72,6 +72,16 @@ function downsample_and_vectorize(seasonal_avgs)
     return vcat(vec.(downsampled_seasonal_avg_arrays)...)
 end
 
+# Process an outputvar into a vector of seasonal averages
+function process_outputvar(simdir, name)
+    monthly_avgs = preprocess_monthly_averages(simdir, name)
+    seasons = split_monthly_averages_into_seasons(monthly_avgs)
+    # Ensure each season has three months
+    @assert all(map(x -> length(times(x)) == 3, seasons))
+    seasonal_avgs = average_time.(seasons)
+    return seasonal_avgs
+end
+
 # Preprocess monthly averages to the right dimensions and dates, remove NaNs
 function preprocess_monthly_averages(simdir, name)
     monthly_avgs = get_monthly_averages(simdir, name)
@@ -81,26 +91,23 @@ function preprocess_monthly_averages(simdir, name)
         monthly_avgs = ClimaAnalysis.Atmos.to_pressure_coordinates(monthly_avgs, pressure)
         monthly_avgs = limit_pressure_dim_to_era5_range(monthly_avgs)
     end
-    # TODO: Ask ollie how to replace nans in a way that makes sense
-    monthly_avgs = ClimaAnalysis.replace(monthly_avgs, missing => 0.0, NaN => 0.0)
+    # TODO: Replace NaNs with global mean
+    monthly_avgs = ClimaAnalysis.replace(monthly_avgs, NaN => 0.0)
     monthly_avgs = ClimaAnalysis.shift_to_start_of_previous_month(monthly_avgs)
     # Remove spinup time
     monthly_avgs = window(monthly_avgs, "time"; left = spinup_time)
     return monthly_avgs
 end
 
-# Process an outputvar into a vector of seasonal averages
-function process_outputvar(simdir, name)
-    monthly_avgs = preprocess_monthly_averages(simdir, name)
-    seasons = split_by_season_across_time(monthly_avgs)
-    # Ensure each season has three months
-    if !all(map(x -> length(times(x)) == 3, seasons))
-        @info "Uneven months per season, rebalancing..."
-        rebalance_months_per_season!(seasons)
-    end
-    seasonal_avgs = average_time.(seasons)
+function split_monthly_averages_into_seasons(monthly_avgs)
+    all_times = times(monthly_avgs)
+    @assert length(all_times) % 3 == 0
 
-    return seasonal_avgs
+    # Window over 3 months at a time to create seasonal outputvars
+    split_by_seasons = map(all_times[1:3:end]) do t
+        window(monthly_avgs, "time"; left = t, right = t + 2months)
+    end
+    return split_by_seasons
 end
 
 function rebalance_months_per_season!(seasons)

experiments/calibration/coarse_amip/observation_utils.jl

@@ -75,19 +75,19 @@ function get_all_output_vars(obs_dir, diagnostic_var2d, diagnostic_var3d)
     # specific humidity
     hus = resample(era5_outputvar(joinpath(obs_dir, "era5_monthly_avg_pressure_level_q.nc")))
 
-    # Cloud specific liquid water content
-    ql = era5_outputvar(joinpath(obs_dir, "era5_specific_cloud_liquid_water_content_1deg.nc"))
-    # Cloud specific ice water content
-    qi = era5_outputvar(joinpath(obs_dir, "era5_specific_cloud_ice_water_content_1deg.nc"))
-    foreach((ql, qi)) do var
-        # Convert from hPa to Pa in-place so we don't create more huge OutputVars
-        @assert var.dim_attributes[pressure_name(var)]["units"] == "hPa"
-        var.dims[pressure_name(var)] .*= 100.0
-        set_dim_units!(var, pressure_name(var), "Pa")
-    end
+    # # Cloud specific liquid water content
+    # ql = era5_outputvar(joinpath(obs_dir, "era5_specific_cloud_liquid_water_content_1deg.nc"))
+    # # Cloud specific ice water content
+    # qi = era5_outputvar(joinpath(obs_dir, "era5_specific_cloud_ice_water_content_1deg.nc"))
+    # foreach((ql, qi)) do var
+    #     # Convert from hPa to Pa in-place so we don't create more huge OutputVars
+    #     @assert var.dim_attributes[pressure_name(var)]["units"] == "hPa"
+    #     var.dims[pressure_name(var)] .*= 100.0
+    #     set_dim_units!(var, pressure_name(var), "Pa")
+    # end
     # TODO: determine where time is spent here
-    ql = resample(reverse_dim(reverse_dim(ql, latitude_name(ql)), pressure_name(ql)))
-    qi = resample(reverse_dim(reverse_dim(qi, latitude_name(qi)), pressure_name(qi)))
+    # ql = resample(reverse_dim(reverse_dim(ql, latitude_name(ql)), pressure_name(ql)))
+    # qi = resample(reverse_dim(reverse_dim(qi, latitude_name(qi)), pressure_name(qi)))
 
     return (; rlut, rsut, rsutcs, rlutcs, pr, net_rad, cre, shf, ts, ta, hur, hus)#, ql, qi)
 end
@@ -111,7 +111,10 @@ end
 #####
 # Processing to create EKP.ObservationSeries
 #####
-
+function to_datetime(var::OutputVar)
+    start_date = DateTime(var.attributes["start_date"], dateformat"yyyy-mm-ddTHH:MM:SS")
+    return [start_date + Second(t) for t in times(var)]
+end
 to_datetime(start_date, time) = DateTime(start_date) + Second(time)
 to_datetime(time) = DateTime(start_date) + Second(time)
 
@@ -140,14 +143,17 @@ function get_yearly_averages(var)
     return year_averaged_matrices
 end
 
-# Given an outputvar, compute the standard deviation at each point for each season.
+# TODO: compute seasonal stdev properly
 function get_seasonal_stdev(output_var)
     all_seasonal_averages = get_seasonal_averages(output_var)
     all_seasonal_averages = downsample.(all_seasonal_averages, 3)
-    seasonal_average_matrix = cat(all_seasonal_averages...; dims = 3)
-    interannual_stdev = std(seasonal_average_matrix, dims = 3)
-    # TODO: Add spatial variance
-    return dropdims(interannual_stdev; dims = 3)
+    
+    # Determine dimensionality of the data
+    dims = ndims(all_seasonal_averages[1]) + 1
+    
+    seasonal_average_matrix = cat(all_seasonal_averages...; dims)
+    interannual_stdev = std(seasonal_average_matrix; dims)
+    return dropdims(interannual_stdev; dims)
 end
 
 # Given an outputvar, compute the covariance for each season.
@@ -226,15 +232,15 @@ function create_observation_vector(nt, yrs = 19)
         # shf_obs = make_single_year_of_seasonal_observations(shf, yr)
         ts_obs = make_single_year_of_seasonal_observations(ts, yr)
 
-        # ta_obs = make_single_year_of_seasonal_observations(ta, yr)
-        # hur_obs = make_single_year_of_seasonal_observations(hur, yr)
-        # hus_obs = make_single_year_of_seasonal_observations(hus, yr)
-        EKP.combine_observations([net_rad_obs, rsut_obs, rlut_obs, cre_obs, pr_obs, ts_obs])#, ta_obs, hur_obs, hus_obs])
+        ta_obs = make_single_year_of_seasonal_observations(ta, yr)
+        hur_obs = make_single_year_of_seasonal_observations(hur, yr)
+        hus_obs = make_single_year_of_seasonal_observations(hus, yr)
+        EKP.combine_observations([net_rad_obs, rsut_obs, rlut_obs, cre_obs, pr_obs, ts_obs, ta_obs, hur_obs, hus_obs])
     end
 
     return all_observations # NOT an EKP.ObservationSeries
 end
-# TODO: Ask  kevin to implement in 
+# TODO: Ask  kevin to implement in ClimaAnalysis
 downsample(var::ClimaAnalysis.OutputVar, n) = downsample(var.data, n)
 
 function downsample(arr::AbstractArray, n)

experiments/calibration/coarse_amip/run_calibration.jl

@@ -29,7 +29,7 @@ priors = [
     constrained_gaussian("precipitation_timescale", 600, 400, 0, 1200),
 ]
 prior = combine_distributions(priors)
-observation_path = joinpath(experiment_dir, "observations.jld2")
+observation_path = joinpath(experiment_dir, "observations_3d.jld2")
 observation_vec = JLD2.load_object(observation_path)
 
 # Create the EKP.ObservationSeries
@@ -49,6 +49,7 @@ eki = EKP.EnsembleKalmanProcess(
     EKP.construct_initial_ensemble(prior, ensemble_size),
     observation_series,
     EKP.TransformInversion(),
+    verbose=true
 )
 
 eki = CAL.calibrate(CAL.WorkerBackend, eki, ensemble_size, n_iterations, prior, output_dir)