From 0aed291cf8259e3502cce7f46b708b73e0af7ffc Mon Sep 17 00:00:00 2001
From: imreddyTeja <treddy0205@gmail.com>
Date: Mon, 24 Feb 2025 15:21:45 -0800
Subject: [PATCH] Add more get_fields to ClimaAtmosSimulation

Add diffuse_fraction, SW_d, LW_d, and cos of the zenith_angle
to get_field on AtmosModelSimulation in Interfacer.jl.

Add get_field methods for the above values in
ClimaAtmosSimulation
---
 .../components/atmosphere/climaatmos.jl       | 24 +++++++++++++++++++
 src/Interfacer.jl                             |  6 +++++
 src/Utilities.jl                              |  7 ++++++
 3 files changed, 37 insertions(+)

diff --git a/experiments/ClimaEarth/components/atmosphere/climaatmos.jl b/experiments/ClimaEarth/components/atmosphere/climaatmos.jl
index 58a1a46626..2558c389d7 100644
--- a/experiments/ClimaEarth/components/atmosphere/climaatmos.jl
+++ b/experiments/ClimaEarth/components/atmosphere/climaatmos.jl
@@ -171,12 +171,36 @@ moisture_flux(::Union{CA.EquilMoistModel, CA.NonEquilMoistModel}, integrator) =
 ρq_tot(::Union{CA.EquilMoistModel, CA.NonEquilMoistModel}, integrator) = integrator.u.c.ρq_tot
 
 # extensions required by the Interfacer
+function Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:diffuse_fraction})
+    radiation_model = sim.integrator.p.radiation.rrtmgp_model
+    total_flux_dn = radiation_model.face_sw_flux_dn
+    face_space = axes(sim.integrator.u.f)
+    if radiation_model.radiation_mode isa CA.RRTMGPInterface.GrayRadiation
+        diffuse_fractions = zero(total_flux_dn)
+    else
+        direct_flux_dn = radiation_model.face_sw_direct_flux_dn
+        diffuse_fractions = Utilities.safe_div.(total_flux_dn .- direct_flux_dn, total_flux_dn)
+    end
+    return CC.Fields.level(CC.Fields.array2field(diffuse_fractions, face_space), CC.Utilities.half)
+end
+Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{Symbol("cos(zenith_angle)")}) = CC.Fields.array2field(
+    sim.integrator.p.radiation.rrtmgp_model.cos_zenith,
+    CC.Fields.level(axes(sim.integrator.u.c), 1),
+)
 Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:liquid_precipitation}) =
     surface_rain_flux(sim.integrator.p.atmos.moisture_model, sim.integrator)
+Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:LW_d}) = CC.Fields.level(
+    CC.Fields.array2field(sim.integrator.p.radiation.rrtmgp_model.face_lw_flux_dn, axes(sim.integrator.u.f)),
+    CC.Utilities.half,
+)
 Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:radiative_energy_flux_sfc}) =
     surface_radiation_flux(sim.integrator.p.atmos.radiation_mode, sim.integrator)
 Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:snow_precipitation}) =
     surface_snow_flux(sim.integrator.p.atmos.moisture_model, sim.integrator)
+Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:SW_d}) = CC.Fields.level(
+    CC.Fields.array2field(sim.integrator.p.radiation.rrtmgp_model.face_sw_flux_dn, axes(sim.integrator.u.f)),
+    CC.Utilities.half,
+)
 Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:turbulent_energy_flux}) =
     CC.Geometry.WVector.(sim.integrator.p.precomputed.sfc_conditions.ρ_flux_h_tot)
 Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:turbulent_moisture_flux}) =
diff --git a/src/Interfacer.jl b/src/Interfacer.jl
index abdb092836..7912c2f346 100644
--- a/src/Interfacer.jl
+++ b/src/Interfacer.jl
@@ -119,13 +119,19 @@ an atmosphere component model.
 get_field(
     sim::AtmosModelSimulation,
     val::Union{
+        Val{Symbol("cos(zenith_angle)")},
+        Val{:diffuse_fraction},
         Val{:energy},
         Val{:height_int},
         Val{:height_sfc},
         Val{:liquid_precipitation},
+        Val{:LHF},
+        Val{:LW_d},
         Val{:radiative_energy_flux_sfc},
         Val{:radiative_energy_flux_toa},
         Val{:snow_precipitation},
+        Val{:SHF},
+        Val{:SW_d},
         Val{:turblent_energy_flux},
         Val{:turbulent_moisture_flux},
         Val{:thermo_state_int},
diff --git a/src/Utilities.jl b/src/Utilities.jl
index 75b359a729..0ede1185ac 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -36,6 +36,13 @@ or 0 otherwise, keeping the same type.
 """
 binary_mask(var) = binary_mask(var, eps(eltype(var)))
 
+"""
+    safe_div(x, y)
+
+Divides `x` by `y` if `y` is not zero, otherwise returns zero.
+"""
+safe_div(x, y) = y == zero(y) ? zero(x) : x / y
+
 """
     swap_space!(space_out::CC.Spaces.AbstractSpace, field_in::CC.Fields.Field)