initial gpu amip work

.buildkite/pipeline.yml

@@ -277,6 +277,13 @@ steps:
       #   command: "julia --color=yes --project=experiments/AMIP/ experiments/AMIP/coupler_driver.jl --coupled true --surface_setup PrescribedSurface --moist equil --vert_diff true --rad allskywithclear --rayleigh_sponge false --energy_check true --mode_name slabplanet --t_end 10days --dt_save_to_sol 3600secs --dt_cpl 21600 --dt 200secs --dt_rad 6hours --idealized_insolation true --mono_surface true --h_elem 4 --precip_model 0M --run_name target_params_in_slab1"
       #   artifact_paths: "experiments/AMIP/output/slabplanet/target_params_in_slab1_artifacts/total_energy*.png"
 
+      - label: "AMIP target: albedo from function"
+        key: "target_amip_albedo_function"
+        command: "julia --color=yes --project=experiments/AMIP/ experiments/AMIP/coupler_driver.jl --config_file $CONFIG_PATH/target_amip_albedo_function.yml"
+        artifact_paths: "experiments/AMIP/output/amip/target_amip_albedo_function_artifacts/*"
+        agents:
+          slurm_mem: 20GB
+
       - label: "AMIP - Float64 + hourly checkpoint"
         key: "amip"
         command: "julia --color=yes --project=experiments/AMIP/ experiments/AMIP/coupler_driver.jl --config_file $CONFIG_PATH/coarse_single_ft64_hourly_checkpoints.yml"
@@ -377,6 +384,17 @@ steps:
           slurm_mem: 20GB
           slurm_gpus: 1
 
+      # GPU RUNS: AMIP
+      - label: "GPU AMIP target: albedo from function"
+        key: "gpu_target_amip_albedo_function"
+        command: "julia --color=yes --project=experiments/AMIP/ experiments/AMIP/coupler_driver.jl --config_file $CONFIG_PATH/gpu_target_amip_albedo_function.yml"
+        artifact_paths: "experiments/AMIP/output/amip/gpu_target_amip_albedo_function_artifacts/*"
+        agents:
+          slurm_mem: 20GB
+          slurm_gpus: 1
+
+      - wait
+
   - wait
 
   # plot job performance history

config/model_configs/default_mono.yml

@@ -1,6 +1,6 @@
 apply_limiter: false
 dt: "200secs"
-dt_cpl 200:
+dt_cpl: 200
 dt_save_to_sol: "3600secs"
 energy_check: true
 h_elem: 4

config/model_configs/gpu_target_amip_albedo_function.yml

@@ -0,0 +1,23 @@
+alpha_rayleigh_uh: 0
+alpha_rayleigh_w: 10
+apply_limiter: false
+dt: "150secs"
+dt_cpl: 150
+dt_rad: "1hours"
+dt_save_to_sol: "1days"
+dz_bottom: 30
+dz_top: 3000
+h_elem: 16
+job_id: "gpu_target_amip_albedo_function"
+kappa_4: 1e16
+land_albedo_type: "function"
+mode_name: "amip"
+moist: "equil"
+precip_model: "0M"
+rad: "gray"
+rayleigh_sponge: true
+run_name: "gpu_target_amip_albedo_function"
+t_end: "300secs"
+vert_diff: "true"
+z_elem: 50
+z_stretch: false

config/model_configs/target_amip_albedo_function.yml

@@ -0,0 +1,23 @@
+alpha_rayleigh_uh: 0
+alpha_rayleigh_w: 10
+apply_limiter: false
+dt: "150secs"
+dt_cpl: 150
+dt_rad: "1hours"
+dt_save_to_sol: "1days"
+dz_bottom: 30
+dz_top: 3000
+h_elem: 16
+job_id: "target_amip_albedo_function"
+kappa_4: 1e16
+land_albedo_type: "function"
+mode_name: "amip"
+moist: "equil"
+precip_model: "0M"
+rad: "gray"
+rayleigh_sponge: true
+run_name: "target_amip_albedo_function"
+t_end: "300secs"
+vert_diff: "true"
+z_elem: 50
+z_stretch: false

experiments/AMIP/user_io/amip_visualizer.jl

@@ -78,8 +78,9 @@ function read_latest_model_data(name::Symbol, filedir::String, root::String)
     varfile_root = @sprintf "%s%s" string(name) root
     filename = glob("*" * varfile_root * "*", filedir)[end]
 
-    comms_ctx = ClimaComms.SingletonCommsContext()
-    hdfreader = InputOutput.HDF5Reader(filename)
+    # Ensure file gets read onto CPU for postprocessing
+    cpu_singleton_context = ClimaComms.SingletonCommsContext(ClimaComms.CPUSingleThreaded())
+    hdfreader = InputOutput.HDF5Reader(filename, cpu_singleton_context)
     var = InputOutput.read_field(hdfreader, string(name))
     close(hdfreader)
     return var

experiments/AMIP/user_io/user_diagnostics.jl

@@ -56,25 +56,23 @@ Top of the atmosphere radiation fluxes (W m⁻²).
 function get_var(cs::CoupledSimulation, ::Val{:toa_fluxes})
     atmos_sim = cs.model_sims.atmos_sim
     face_space = axes(atmos_sim.integrator.u.f)
-    z = parent(Fields.coordinate_field(face_space).z)
-    Δz_top = round(FT(0.5) * (z[end, 1, 1, 1, 1] - z[end - 1, 1, 1, 1, 1]))
-    n_faces = length(z[:, 1, 1, 1, 1])
+    nz_faces = length(face_space.grid.vertical_grid.topology.mesh.faces)
 
     LWd_TOA = Fields.level(
         CA.RRTMGPI.array2field(FT.(atmos_sim.integrator.p.radiation.radiation_model.face_lw_flux_dn), face_space),
-        n_faces - half,
+        nz_faces - half,
     )
     LWu_TOA = Fields.level(
         CA.RRTMGPI.array2field(FT.(atmos_sim.integrator.p.radiation.radiation_model.face_lw_flux_up), face_space),
-        n_faces - half,
+        nz_faces - half,
     )
     SWd_TOA = Fields.level(
         CA.RRTMGPI.array2field(FT.(atmos_sim.integrator.p.radiation.radiation_model.face_sw_flux_dn), face_space),
-        n_faces - half,
+        nz_faces - half,
     )
     SWu_TOA = Fields.level(
         CA.RRTMGPI.array2field(FT.(atmos_sim.integrator.p.radiation.radiation_model.face_sw_flux_up), face_space),
-        n_faces - half,
+        nz_faces - half,
     )
 
     radiation_sources = @. -(LWd_TOA + SWd_TOA - LWu_TOA - SWu_TOA)

src/Regridder.jl

@@ -163,10 +163,10 @@ function hdwrite_regridfile_rll_to_cgll(
     end
 
     # If doesn't make sense to regrid with GPUs/MPI processes
-    cpu_context = ClimaComms.SingletonCommsContext(ClimaComms.CPUSingleThreaded())
+    cpu_singleton_context = ClimaComms.SingletonCommsContext(ClimaComms.CPUSingleThreaded())
 
     topology = Topologies.Topology2D(
-        cpu_context,
+        cpu_singleton_context,
         Spaces.topology(space2d).mesh,
         Topologies.spacefillingcurve(Spaces.topology(space2d).mesh),
     )
@@ -245,7 +245,7 @@ function hdwrite_regridfile_rll_to_cgll(
     )
 
     map(
-        x -> write_to_hdf5(REGRID_DIR, hd_outfile_root, times[x], offline_fields[x], varname, cpu_context),
+        x -> write_to_hdf5(REGRID_DIR, hd_outfile_root, times[x], offline_fields[x], varname, cpu_singleton_context),
         1:length(times),
     )
     jldsave(joinpath(REGRID_DIR, hd_outfile_root * "_times.jld2"); times = times)