add back state comparisons

CliMA · Mar 21, 2024 · ff4e8ee · ff4e8ee
1 parent ac1035b
commit ff4e8ee
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 38 deletions.
diff --git a/experiments/AMIP/compare_cpu_gpu.jl b/experiments/AMIP/compare_cpu_gpu.jl
@@ -90,8 +90,8 @@ gpu_land_state_2d = DLM.readdlm(joinpath(gpu_artifacts_dir, "land_state_2d_tend_
 
 # Read in ocean state variables (if not AMIP)
 if !(mode_name == "amip")
-    cpu_land_state = DLM.readdlm(joinpath(cpu_artifacts_dir, "land_state_tend_cpu.txt"), ',')
-    gpu_land_state = DLM.readdlm(joinpath(gpu_artifacts_dir, "land_state_tend_gpu.txt"), ',')
+    cpu_ocean_state = DLM.readdlm(joinpath(cpu_artifacts_dir, "ocean_state_tend_cpu.txt"), ',')
+    gpu_ocean_state = DLM.readdlm(joinpath(gpu_artifacts_dir, "ocean_state_tend_gpu.txt"), ',')
 
     @show abs(maximum(cpu_ocean_state .- gpu_ocean_state))
     @show abs(median(cpu_ocean_state .- gpu_ocean_state))

diff --git a/experiments/AMIP/coupler_driver.jl b/experiments/AMIP/coupler_driver.jl
@@ -771,42 +771,49 @@ atmos_ρ = cs.model_sims.atmos_sim.integrator.u.c.ρ
 atmos_uₕ = cs.model_sims.atmos_sim.integrator.u.c.uₕ
 atmos_u₃ = cs.model_sims.atmos_sim.integrator.u.f.u₃
 
-# # Write to text files
-# open(joinpath(COUPLER_ARTIFACTS_DIR, "atmos_state_center_tend_$device_suffix.txt"), "w") do io
-#     DLM.writedlm(
-#         io,
-#         hcat(parent(atmos_ρe_tot)[:], parent(atmos_ρq_tot)[:], parent(atmos_ρ)[:], parent(atmos_uₕ)[:]),
-#         ',',
-#     )
-# end;
-
-# open(joinpath(COUPLER_ARTIFACTS_DIR, "atmos_state_face_tend_$device_suffix.txt"), "w") do io
-#     DLM.writedlm(io, parent(atmos_u₃)[:], ',')
-# end;
-
-# # Extract land state variables
-# land_T = cs.model_sims.land_sim.integrator.u.bucket.T
-# land_W = cs.model_sims.land_sim.integrator.u.bucket.W
-# land_Ws = cs.model_sims.land_sim.integrator.u.bucket.Ws
-# land_σS = cs.model_sims.land_sim.integrator.u.bucket.σS
-
-# # Write to text files
-# open(joinpath(COUPLER_ARTIFACTS_DIR, "land_state_3d_tend_$device_suffix.txt"), "w") do io
-#     DLM.writedlm(io, parent(land_T)[:], ',')
-# end;
-# open(joinpath(COUPLER_ARTIFACTS_DIR, "land_state_2d_tend_$device_suffix.txt"), "w") do io
-#     DLM.writedlm(io, hcat(parent(land_W)[:], parent(land_Ws)[:], parent(land_σS)[:]), ',')
-# end;
-
-# # ocean state
-# if cs.mode.name != "amip"
-#     ocean_T_sfc = cs.model_sims.ocean_sim.integrator.u.T_sfc
-
-#     # Write to text files
-#     open(joinpath(COUPLER_ARTIFACTS_DIR, "ocean_state_tend_$device_suffix.txt"), "w") do io
-#         DLM.writedlm(io, parent(ocean_T_sfc)[:], ',')
-#     end
-# end;
+# Write to text files
+open(joinpath(COUPLER_ARTIFACTS_DIR, "atmos_state_center_tend_$device_suffix.txt"), "w") do io
+    # note: `atmos_uₕ` has two components that we have to index into
+    DLM.writedlm(
+        io,
+        hcat(
+            Array(parent(atmos_ρe_tot))[:],
+            Array(parent(atmos_ρq_tot))[:],
+            Array(parent(atmos_ρ))[:],
+            Array(parent(atmos_uₕ.components))[:, :, :, 1, :][:],
+            Array(parent(atmos_uₕ.components))[:, :, :, 2, :][:],
+        ),
+        ',',
+    )
+end;
+
+open(joinpath(COUPLER_ARTIFACTS_DIR, "atmos_state_face_tend_$device_suffix.txt"), "w") do io
+    DLM.writedlm(io, Array(parent(atmos_u₃))[:], ',')
+end;
+
+# Extract land state variables
+land_T = cs.model_sims.land_sim.integrator.u.bucket.T
+land_W = cs.model_sims.land_sim.integrator.u.bucket.W
+land_Ws = cs.model_sims.land_sim.integrator.u.bucket.Ws
+land_σS = cs.model_sims.land_sim.integrator.u.bucket.σS
+
+# Write to text files
+open(joinpath(COUPLER_ARTIFACTS_DIR, "land_state_3d_tend_$device_suffix.txt"), "w") do io
+    DLM.writedlm(io, Array(parent(land_T))[:], ',')
+end;
+open(joinpath(COUPLER_ARTIFACTS_DIR, "land_state_2d_tend_$device_suffix.txt"), "w") do io
+    DLM.writedlm(io, hcat(Array(parent(land_W))[:], Array(parent(land_Ws))[:], Array(parent(land_σS))[:]), ',')
+end;
+
+# ocean state
+if cs.mode.name != "amip"
+    ocean_T_sfc = cs.model_sims.ocean_sim.integrator.u.T_sfc
+
+    # Write to text files
+    open(joinpath(COUPLER_ARTIFACTS_DIR, "ocean_state_tend_$device_suffix.txt"), "w") do io
+        DLM.writedlm(io, Array(parent(ocean_T_sfc))[:], ',')
+    end
+end;
 
 #=
 ## Postprocessing