Merge pull request #426 from SpeedyWeather/mk/forcing

Export forcing!, drag!, etc for extending those with fewer conflicts
SpeedyWeather · Jan 5, 2024 · 56b1fc9 · 56b1fc9
2 parents e5ceca5 + 94b877c
commit 56b1fc9
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 19 deletions.
diff --git a/src/SpeedyWeather.jl b/src/SpeedyWeather.jl
@@ -46,7 +46,8 @@ export  Barotropic,             # abstract
         ShallowWater,
         PrimitiveEquation,
         PrimitiveDry,
-        PrimitiveWet
+        PrimitiveWet,
+        ModelSetup
 
 export  BarotropicModel,        # concrete
         ShallowWaterModel,
@@ -102,7 +103,13 @@ export  NoBoundaryLayerDrag,
         QuadraticDrag
 
 # EXPORT FORCING
-export  JetStreamForcing
+export  forcing!,
+        JetStreamForcing,
+        AbstractForcing
+
+# EXPORT DRAG
+export  drag!,
+        AbstractDrag
 
 # EXPORT VERTICAL DIFFUSION
 export  NoVerticalDiffusion,
@@ -117,7 +124,9 @@ export  DynamicsConstants,
         SpectralTransform,
         Boundaries,
         PrognosticVariables,
+        PrognosticVariablesLayer,
         DiagnosticVariables,
+        DiagnosticVariablesLayer,
         ColumnVariables
 
 # EXPORT SPECTRAL FUNCTIONS

diff --git a/test/extending.jl b/test/extending.jl
@@ -28,7 +28,7 @@
     end
 
     function SpeedyWeather.initialize!( drag::JetDrag,
-                                        model::SpeedyWeather.ModelSetup)
+                                        model::ModelSetup)
 
         (;spectral_grid, geometry) = model
         (;Grid,NF,nlat_half) = spectral_grid
@@ -46,11 +46,11 @@
         return nothing
     end
 
-    function SpeedyWeather.drag!(   diagn::SpeedyWeather.DiagnosticVariablesLayer,
-                                    progn::SpeedyWeather.PrognosticVariablesLayer,
+    function SpeedyWeather.drag!(   diagn::DiagnosticVariablesLayer,
+                                    progn::PrognosticVariablesLayer,
                                     drag::JetDrag,
                                     time::DateTime,
-                                    model::SpeedyWeather.ModelSetup)
+                                    model::ModelSetup)
 
         (;vor) = progn
         (;vor_tend) = diagn.tendencies
@@ -61,10 +61,6 @@
         for lm in eachindex(vor,vor_tend,ζ₀)
             vor_tend[lm] -= r*(vor[lm] - ζ₀[lm])
         end
-
-        SpeedyTransforms.spectral_truncation!(vor_tend)    # set lmax+1 to zero
-
-        return nothing
     end
 
     Base.@kwdef struct StochasticStirring{NF} <: SpeedyWeather.AbstractForcing{NF}
@@ -90,6 +86,17 @@
         "Stirring width [˚]"
         width::Float64 = 24
 
+        "Minimum degree of spherical harmonics to force"
+        lmin::Int = 8
+
+        "Maximum degree of spherical harmonics to force"
+        lmax::Int = 40
+
+        "Minimum order of spherical harmonics to force"
+        mmin::Int = 4
+
+        "Maximum order of spherical harmonics to force"
+        mmax::Int = lmax
 
         # TO BE INITIALISED
         "Stochastic stirring term S"
@@ -112,7 +119,7 @@
     end
 
     function SpeedyWeather.initialize!( forcing::StochasticStirring,
-                                        model::SpeedyWeather.ModelSetup)
+                                        model::ModelSetup)
 
         # precompute forcing strength, scale with radius^2 as is the vorticity equation
         (;radius) = model.spectral_grid
@@ -136,15 +143,15 @@
         return nothing
     end
 
-    function SpeedyWeather.forcing!(diagn::SpeedyWeather.DiagnosticVariablesLayer,
-                                    progn::SpeedyWeather.PrognosticVariablesLayer,
+    function SpeedyWeather.forcing!(diagn::DiagnosticVariablesLayer,
+                                    progn::PrognosticVariablesLayer,
                                     forcing::StochasticStirring,
                                     time::DateTime,
-                                    model::SpeedyWeather.ModelSetup)
+                                    model::ModelSetup)
         SpeedyWeather.forcing!(diagn,forcing,model.spectral_transform)
     end
 
-    function SpeedyWeather.forcing!(diagn::SpeedyWeather.DiagnosticVariablesLayer,
+    function SpeedyWeather.forcing!(diagn::DiagnosticVariablesLayer,
                                     forcing::StochasticStirring{NF},
                                     spectral_transform::SpectralTransform) where NF
 
@@ -153,10 +160,16 @@
         b = forcing.b[]    # = exp(-dt/τ)
 
         (;S) = forcing
-        @inbounds for lm in eachindex(S)
-            # Barnes and Hartmann, 2011 Eq. 2
-            Qi = 2rand(Complex{NF}) - (1 + im)   # ~ [-1,1] in complex
-            S[lm] = a*Qi + b*S[lm]
+        lmax,mmax = size(S)
+        @inbounds for m in 1:mmax
+            for l in m:lmax
+                if (forcing.mmin <= m <= forcing.mmax) &&
+                    (forcing.lmin <= l <= forcing.lmax)
+                    # Barnes and Hartmann, 2011 Eq. 2
+                    Qi = 2rand(Complex{NF}) - (1 + im)   # ~ [-1,1] in complex
+                    S[l,m] = a*Qi + b*S[l,m]
+                end
+            end
         end
 
         # to grid-point space