From b91d61cf9070c217286aaf9eabf70f51bc53a172 Mon Sep 17 00:00:00 2001
From: Zhaoyi Shen <11598433+szy21@users.noreply.github.com>
Date: Wed, 2 Oct 2024 13:45:09 -0700
Subject: [PATCH] add rainfall and snowfall flux  diagnostics

---
 src/diagnostics/core_diagnostics.jl    | 64 ++++++++++++++++++++++++++
 src/diagnostics/default_diagnostics.jl |  2 +
 2 files changed, 66 insertions(+)

diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
index ab76d8a4a8..c6aa345f3a 100644
--- a/src/diagnostics/core_diagnostics.jl
+++ b/src/diagnostics/core_diagnostics.jl
@@ -711,6 +711,70 @@ add_diagnostic_variable!(
     compute! = compute_pr!,
 )
 
+compute_prra!(out, state, cache, time) =
+    compute_prra!(out, state, cache, time, cache.atmos.precip_model)
+compute_prra!(_, _, _, _, precip_model::T) where {T} =
+    error_diagnostic_variable("prra", precip_model)
+
+function compute_prra!(
+    out,
+    state,
+    cache,
+    time,
+    precip_model::Union{
+        NoPrecipitation,
+        Microphysics0Moment,
+        Microphysics1Moment,
+    },
+)
+    if isnothing(out)
+        return cache.precipitation.surface_rain_flux
+    else
+        out .= cache.precipitation.surface_rain_flux
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "prra",
+    long_name = "Rainfall Flux",
+    standard_name = "rainfall_flux",
+    units = "kg m^-2 s^-1",
+    comments = "Precipitation including all forms of water in the liquid phase",
+    compute! = compute_prra!,
+)
+
+compute_prsn!(out, state, cache, time) =
+    compute_prsn!(out, state, cache, time, cache.atmos.precip_model)
+compute_prsn!(_, _, _, _, precip_model::T) where {T} =
+    error_diagnostic_variable("prsn", precip_model)
+
+function compute_prsn!(
+    out,
+    state,
+    cache,
+    time,
+    precip_model::Union{
+        NoPrecipitation,
+        Microphysics0Moment,
+        Microphysics1Moment,
+    },
+)
+    if isnothing(out)
+        return cache.precipitation.surface_snow_flux
+    else
+        out .= cache.precipitation.surface_snow_flux
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "prsn",
+    long_name = "Snowfall Flux",
+    standard_name = "snowfall_flux",
+    units = "kg m^-2 s^-1",
+    comments = "Precipitation including all forms of water in the solid phase",
+    compute! = compute_prsn!,
+)
+
 ###
 # Precipitation (3d)
 ###
diff --git a/src/diagnostics/default_diagnostics.jl b/src/diagnostics/default_diagnostics.jl
index 477d497237..03702fed82 100644
--- a/src/diagnostics/default_diagnostics.jl
+++ b/src/diagnostics/default_diagnostics.jl
@@ -200,6 +200,8 @@ function default_diagnostics(
         "hussfc",
         "evspsbl",
         "pr",
+        "prra",
+        "prsn",
         "prw",
         "lwp",
         "clwvi",