From a6c06f05693f830c301fc97d9c972862bc6f791d Mon Sep 17 00:00:00 2001
From: milankl <milankloewer@gmx.de>
Date: Sat, 14 Dec 2024 00:02:03 +0000
Subject: [PATCH 1/4] ConvectiveHeating implemented

---
 src/physics/convection.jl | 59 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 59 insertions(+)

diff --git a/src/physics/convection.jl b/src/physics/convection.jl
index 44c57421b..a0a0bc6dc 100644
--- a/src/physics/convection.jl
+++ b/src/physics/convection.jl
@@ -430,4 +430,63 @@ function dry_adiabat!(
     # level of zero buoyancy is reached when the loop stops, but in case it's at the top it's still buoyant
     level_zero_buoyancy = k + (1-buoyant)
     return level_zero_buoyancy
+end
+
+export ConvectiveHeating
+@kwdef struct ConvectiveHeating{NF} <: AbstractConvection
+    # DIMENSION
+    nlat::Int
+    
+    "[OPTION] Q_max heating strength as 1K/time_scale"
+    time_scale::Second = Hour(12)
+
+    "[OPTION] Pressure of maximum heating [hPa]"
+    p₀::NF = 525
+
+    "[OPTION] Vertical extent of heating [hPa]"
+    σₚ::NF = 200
+
+    "[OPTION] Latitude of heating [˚N]"
+    θ₀::NF = 0
+
+    "[OPTION] Latitudinal width of heating [˚]"
+    σθ::NF = 20
+
+    # precomputed latitude mask
+    lat_mask::Vector{NF} = zeros(NF, nlat)
+end
+
+# generator
+ConvectiveHeating(SG::SpectralGrid; kwargs...) = ConvectiveHeating{SG.NF}(nlat=SG.nlat; kwargs...)
+
+function initialize!(C::ConvectiveHeating, model::PrimitiveEquation)
+    
+    (; latd) = model.geometry
+    (; θ₀, σθ) = C
+    
+    for (j, θ) in enumerate(latd)
+        # Lee and Kim, 2003, eq. 2
+        C.lat_mask[j] = cosd((θ-θ₀)/σθ)^2
+    end
+end
+
+function convection!(
+    column::ColumnVariables,
+    scheme::ConvectiveHeating,
+    model::PrimitiveEquation,
+)
+    # escape immediately if not in the tropics
+    abs(column.latd) > scheme.σθ && return nothing
+
+    p₀ = scheme.p₀*100      # hPa -> Pa
+    σₚ = scheme.σₚ*100      # hPa -> Pa
+    cos²θ_term = scheme.lat_mask[column.jring]
+    Qmax = 1/Second(scheme.time_scale).value
+
+    for k in eachindex(column)
+        p = column.pres[k]      # Pressure in Pa
+
+        # Lee and Kim, 2003, eq. 2, 
+        column.temp_tend[k] += Qmax*exp(-((p-p₀)/σₚ)^2 / 2)*cos²θ_term
+    end
 end
\ No newline at end of file

From 02f5238b24e5ec57a1dd7aa9a5e990ee2a397a9b Mon Sep 17 00:00:00 2001
From: milankl <milankloewer@gmx.de>
Date: Sat, 14 Dec 2024 00:04:28 +0000
Subject: [PATCH 2/4] update changelog

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9f113bbda..73a98bd4b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- ConvectiveHeating implemented [#639](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/639)
 - Roll back GPUArrays upgrade to ensure CUDA compatibility [#636](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/636)
 - Change default timestep to 40min at T31 [#623](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/623)
 

From 54a0b0eb762c6b2100648ed78219de1e164015cf Mon Sep 17 00:00:00 2001
From: milankl <milankloewer@gmx.de>
Date: Sat, 14 Dec 2024 00:08:28 +0000
Subject: [PATCH 3/4] docstring added

---
 src/physics/convection.jl | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/physics/convection.jl b/src/physics/convection.jl
index a0a0bc6dc..0aed75ad2 100644
--- a/src/physics/convection.jl
+++ b/src/physics/convection.jl
@@ -433,6 +433,10 @@ function dry_adiabat!(
 end
 
 export ConvectiveHeating
+
+"""Convective heating as defined by Lee and Kim, 2003, JAS
+implemented as convection parameterization. Fields are
+$(TYPEDFIELDS)"""
 @kwdef struct ConvectiveHeating{NF} <: AbstractConvection
     # DIMENSION
     nlat::Int
@@ -459,6 +463,7 @@ end
 # generator
 ConvectiveHeating(SG::SpectralGrid; kwargs...) = ConvectiveHeating{SG.NF}(nlat=SG.nlat; kwargs...)
 
+# precompute latitudinal mask
 function initialize!(C::ConvectiveHeating, model::PrimitiveEquation)
     
     (; latd) = model.geometry
@@ -476,7 +481,7 @@ function convection!(
     model::PrimitiveEquation,
 )
     # escape immediately if not in the tropics
-    abs(column.latd) > scheme.σθ && return nothing
+    abs(column.latd) >= scheme.σθ && return nothing
 
     p₀ = scheme.p₀*100      # hPa -> Pa
     σₚ = scheme.σₚ*100      # hPa -> Pa
@@ -486,7 +491,7 @@ function convection!(
     for k in eachindex(column)
         p = column.pres[k]      # Pressure in Pa
 
-        # Lee and Kim, 2003, eq. 2, 
+        # Lee and Kim, 2003, eq. 2
         column.temp_tend[k] += Qmax*exp(-((p-p₀)/σₚ)^2 / 2)*cos²θ_term
     end
 end
\ No newline at end of file

From cfa9708e25c84a7958d53125e9e61ef94410bca7 Mon Sep 17 00:00:00 2001
From: milankl <milankloewer@gmx.de>
Date: Sat, 14 Dec 2024 00:17:06 +0000
Subject: [PATCH 4/4] tests for ConvectiveHeating

---
 test/convection.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/test/convection.jl b/test/convection.jl
index a833264d9..407a0d291 100644
--- a/test/convection.jl
+++ b/test/convection.jl
@@ -4,7 +4,8 @@
 
     for Convection in ( NoConvection,
                         SimplifiedBettsMiller,
-                        DryBettsMiller)
+                        DryBettsMiller,
+                        ConvectiveHeating)
         for Model in (  PrimitiveDryModel,
                         PrimitiveWetModel)