modify target_levels

Author: costachris

config/model_configs/diagnostic_edmfx_era5_initial_condition.yml

@@ -4,9 +4,7 @@ z_elem: 63
 dz_bottom: 30.0
 dt: 120secs
 t_end: 12hours
-# start_date: "20250627-0000"
-start_date: "20250831"
-era5_initial_condition_dir: "/net/sampo/data1/wxquest_data/initial_conditions"
+start_date: "20250627-0000"
 initial_condition: "WeatherModel"
 discrete_hydrostatic_balance: false # TODO: change to true once we fix discrete hydrostatic balance to work with orography
 topo_smoothing: true

src/initial_conditions/initial_conditions.jl

@@ -380,8 +380,10 @@ function overwrite_initial_conditions!(
     extrapolation_bc = (Intp.Periodic(), Intp.Flat(), Intp.Flat())
 
     # Extract face coordinates and compute center midpoints
-
-    target_levels = Fields.field2array(Fields.coordinate_field(Y.c).z)[:, argmin(Fields.field2array(Fields.coordinate_field(Y.c).z))[2]]
+    # Compute target levels on CPU to avoid GPU reductions
+    z_arr_cpu = Array(Fields.field2array(Fields.coordinate_field(Y.c).z))
+    icol = argmin(z_arr_cpu[1, :])
+    target_levels = z_arr_cpu[:, icol]
 
     file_path = weather_model_data_path(
         initial_condition.start_date,

src/utils/weather_model.jl

@@ -144,7 +144,8 @@ function to_z_levels(source_file, target_file, target_levels, FT)
     # See https://agupubs.onlinelibrary.wiley.com/doi/full/10.1002/2017MS001059
     req3d = ["u", "v", "t", "q", "w"]
     for var_name in req3d
-        var_obj = defVar(ncout, var_name, FT, ("lon", "lat", "z"), attrib = ncin[var_name].attrib)
+        var_obj =
+            defVar(ncout, var_name, FT, ("lon", "lat", "z"), attrib = ncin[var_name].attrib)
         if var_name == "w"
             var_obj[:, :, :] = zeros(FT, length(lon), length(lat), length(target_levels))
         else
@@ -160,7 +161,8 @@ function to_z_levels(source_file, target_file, target_levels, FT)
     # Simply repeat the surface values for all levels
     surf_map = Dict("skt" => "skt", "sp" => "p")
     for (src_name, dst_name) in surf_map
-        var_obj = defVar(ncout, dst_name, FT, ("lon", "lat", "z"), attrib = ncin[src_name].attrib)
+        var_obj =
+            defVar(ncout, dst_name, FT, ("lon", "lat", "z"), attrib = ncin[src_name].attrib)
         for k in 1:length(target_levels)
             var_obj[:, :, k] = FT.(ncin[src_name][:, :, 1])
         end