From 0c64aa32939b4262f5d84210c9e21b26f444e415 Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Wed, 9 Dec 2020 19:48:04 -0800
Subject: [PATCH] Fix second order rain flux

Move precipitation flux def up before ref
---
 src/Atmos/Model/multiphysics_types.jl       | 11 ++++
 src/Atmos/Model/precipitation.jl            | 63 +++------------------
 src/Atmos/Model/tendencies_precipitation.jl | 39 +++++++++++++
 3 files changed, 57 insertions(+), 56 deletions(-)

diff --git a/src/Atmos/Model/multiphysics_types.jl b/src/Atmos/Model/multiphysics_types.jl
index c85c64a226e..92d8ccab0ee 100644
--- a/src/Atmos/Model/multiphysics_types.jl
+++ b/src/Atmos/Model/multiphysics_types.jl
@@ -48,6 +48,17 @@ RemovePrecipitation(b::Bool) = (
     RemovePrecipitation{TotalMoisture}(b),
 )
 
+"""
+    PrecipitationFlux{PV <: Union{Rain, Snow}} <: TendencyDef{Flux{FirstOrder}, PV}
+
+Computes the precipitation flux as a sum of air velocity and terminal velocity
+multiplied by the advected variable.
+"""
+struct PrecipitationFlux{PV <: Union{Rain, Snow}} <:
+       TendencyDef{Flux{FirstOrder}, PV} end
+
+PrecipitationFlux() = (PrecipitationFlux{Rain}(), PrecipitationFlux{Snow}())
+
 function remove_precipitation_sources(
     s::RemovePrecipitation{PV},
     atmos,
diff --git a/src/Atmos/Model/precipitation.jl b/src/Atmos/Model/precipitation.jl
index f3a28d4e6af..0c02a5d0797 100644
--- a/src/Atmos/Model/precipitation.jl
+++ b/src/Atmos/Model/precipitation.jl
@@ -56,54 +56,6 @@ function compute_gradient_argument!(
 
 source!(::PrecipitationModel, args...) = nothing
 
-"""
-    PrecipitationFlux{PV <: Union{Rain, Snow}} <: TendencyDef{Flux{FirstOrder}, PV}
-
-Computes the precipitation flux as a sum of air velocity and terminal velocity
-multiplied by the advected variable.
-"""
-struct PrecipitationFlux{PV <: Union{Rain, Snow}} <:
-       TendencyDef{Flux{FirstOrder}, PV} end
-
-PrecipitationFlux() = (PrecipitationFlux{Rain()}, PrecipitationFlux{Snow()})
-
-function flux(::PrecipitationFlux{Rain}, m, state, aux, t, ts, direction)
-    FT = eltype(state)
-    u = state.ρu / state.ρ
-    q_rai = state.precipitation.ρq_rai / state.ρ
-
-    v_term_rai::FT = FT(0)
-    if q_rai > FT(0)
-        v_term_rai = terminal_velocity(
-            m.param_set,
-            m.param_set.microphys.rai,
-            state.ρ,
-            q_rai,
-        )
-    end
-
-    k̂ = vertical_unit_vector(m, aux)
-    return state.precipitation.ρq_rai * (u - k̂ * v_term_rai)
-end
-function flux(::PrecipitationFlux{Snow}, m, state, aux, t, ts, direction)
-    FT = eltype(state)
-    u = state.ρu / state.ρ
-    q_sno = state.precipitation.ρq_sno / state.ρ
-
-    v_term_sno::FT = FT(0)
-    if q_sno > FT(0)
-        v_term_sno = terminal_velocity(
-            m.param_set,
-            m.param_set.microphys.sno,
-            state.ρ,
-            q_sno,
-        )
-    end
-
-    k̂ = vertical_unit_vector(m, aux)
-    return state.precipitation.ρq_sno * (u - k̂ * v_term_sno)
-end
-
 """
     NoPrecipitation <: PrecipitationModel
 
@@ -170,18 +122,17 @@ end
 function flux_second_order!(
     precip::RainModel,
     flux::Grad,
+    atmos::AtmosModel,
     state::Vars,
-    diffusive::Vars,
     aux::Vars,
     t::Real,
-    D_t,
+    ts,
+    diffusive::Vars,
+    hyperdiffusive::Vars,
 )
-    d_q_rai = (-D_t) .* diffusive.precipitation.∇q_rai
-
-    flux_second_order!(precip, flux, state, d_q_rai)
-end
-function flux_second_order!(precip::RainModel, flux::Grad, state::Vars, d_q_rai)
-    flux.precipitation.ρq_rai += d_q_rai * state.ρ
+    tend = Flux{SecondOrder}()
+    args = (atmos, state, aux, t, ts, diffusive, hyperdiffusive)
+    flux.precipitation.ρq_rai = Σfluxes(eq_tends(Rain(), atmos, tend), args...)
 end
 
 function source!(
diff --git a/src/Atmos/Model/tendencies_precipitation.jl b/src/Atmos/Model/tendencies_precipitation.jl
index 20eb536af03..41433af63f2 100644
--- a/src/Atmos/Model/tendencies_precipitation.jl
+++ b/src/Atmos/Model/tendencies_precipitation.jl
@@ -4,6 +4,45 @@
 ##### First order fluxes
 #####
 
+function flux(::PrecipitationFlux{Rain}, m, state, aux, t, ts, direction)
+    FT = eltype(state)
+    u = state.ρu / state.ρ
+    q_rai = state.precipitation.ρq_rai / state.ρ
+
+    v_term_rai::FT = FT(0)
+    if q_rai > FT(0)
+        v_term_rai = terminal_velocity(
+            m.param_set,
+            m.param_set.microphys.rai,
+            state.ρ,
+            q_rai,
+        )
+    end
+
+    k̂ = vertical_unit_vector(m, aux)
+    return state.precipitation.ρq_rai * (u - k̂ * v_term_rai)
+end
+
+function flux(::PrecipitationFlux{Snow}, m, state, aux, t, ts, direction)
+    FT = eltype(state)
+    u = state.ρu / state.ρ
+    q_sno = state.precipitation.ρq_sno / state.ρ
+
+    v_term_sno::FT = FT(0)
+    if q_sno > FT(0)
+        v_term_sno = terminal_velocity(
+            m.param_set,
+            m.param_set.microphys.sno,
+            state.ρ,
+            q_sno,
+        )
+    end
+
+    k̂ = vertical_unit_vector(m, aux)
+    return state.precipitation.ρq_sno * (u - k̂ * v_term_sno)
+end
+
+
 #####
 ##### Second order fluxes
 #####