Merge pull request #12 from SpeedyWeather/mk/drag_curl

mask=true/false for StochasticStirring, JetDrag initialization with model.spectral_transform

Author: milankl

src/jet_drag.jl

@@ -56,8 +56,8 @@ function SpeedyWeather.initialize!( drag::JetDrag,
         u[ij] = drag.u₀ * exp(-(lat[ij]-drag.latitude)^2/(2*drag.width^2))
     end
 
-    # to spectral space of size lmax+1 × mmax as required by the curl 
-    û = SpeedyTransforms.spectral(u, one_more_degree=true)
+    # to spectral space, reusing the precomputed spectral transform from model
+    û = SpeedyTransforms.spectral(u, model.spectral_transform)
     v̂ = zero(û)
     SpeedyTransforms.curl!(drag.ζ₀, û, v̂, model.spectral_transform)
     return nothing

src/stochastic_stirring.jl

@@ -23,6 +23,9 @@ Base.@kwdef struct StochasticStirring{NF} <: SpeedyWeather.AbstractForcing
     "Stirring latitude [˚N]"
     latitude::NF = 45
 
+    "En/disable stirring mask"
+    mask::Bool = true
+
     "Stirring width [˚]"
     width::NF = 24
 
@@ -81,12 +84,13 @@ function SpeedyWeather.initialize!( forcing::StochasticStirring,
     forcing.a[] = A*sqrt(1 - exp(-2dt/τ))
     forcing.b[] = exp(-dt/τ)
 
-    # precompute the Gaussian latitudinal mask
-    (;latd) = model.geometry                 # in ˚N on every latitude ring
-    
-    for j in eachindex(forcing.lat_mask)
-        # Gaussian centred at forcing.latitude of width forcing.width
-        forcing.lat_mask[j] = exp(-(forcing.latitude-latd[j])^2/forcing.width^2*2)
+    if forcing.mask
+        # precompute the Gaussian latitudinal mask
+        (;latd) = model.geometry                 # in ˚N on every latitude ring
+        for j in eachindex(forcing.lat_mask)
+            # Gaussian centred at forcing.latitude of width forcing.width
+            forcing.lat_mask[j] = exp(-(forcing.latitude-latd[j])^2/forcing.width^2*2)
+        end
     end
 
     return nothing
@@ -134,16 +138,22 @@ function SpeedyWeather.forcing!(
         end
     end
 
-    # to grid-point space
-    S_grid = diagn.dynamics_variables.a_grid        # reuse general work array
-    SpeedyTransforms.gridded!(S_grid, S, spectral_transform)
-    
-    # mask everything but mid-latitudes
-    RingGrids._scale_lat!(S_grid, forcing.lat_mask)
-    
-    # back to spectral space, write directly into vorticity tendency
+    # add stochastic stirring term S to vorticity tendency
+    # with masked or without (=skip additional transform)
     (; vor_tend) = diagn.tendencies
-    SpeedyTransforms.spectral!(vor_tend, S_grid, spectral_transform)
-    
+    if forcing.mask
+        # to grid-point space
+        S_grid = diagn.dynamics_variables.a_grid        # reuse general work array
+        SpeedyTransforms.gridded!(S_grid, S, spectral_transform)
+        
+        # mask everything but mid-latitudes
+        RingGrids._scale_lat!(S_grid, forcing.lat_mask)
+        
+        # back to spectral space, write directly into vorticity tendency
+        SpeedyTransforms.spectral!(vor_tend, S_grid, spectral_transform)
+    else
+        vor_tend .= S   # copy forcing S over into vor_tend
+    end
+        
     return nothing
 end

test/runtests.jl

@@ -23,3 +23,29 @@ using Test
     run!(simulation, period=Day(5))
     @test simulation.model.feedback.nars_detected == false
 end
+
+@testset "Initialize at various resolutions" begin
+    for trunc in (31, 42, 63, 85, 127, 170, 255, 340)
+        spectral_grid = SpectralGrid(trunc=trunc, nlev=1)
+        drag = JetDrag(spectral_grid, time_scale=Day(6))
+        forcing = StochasticStirring(spectral_grid)
+        initial_conditions = StartFromRest()
+    
+        # with barotropic model
+        model = BarotropicModel(;spectral_grid, initial_conditions, forcing, drag)
+        simulation = initialize!(model)
+    end
+end
+
+@testset "Mask off" begin
+    spectral_grid = SpectralGrid(trunc=31, nlev=1)
+    drag = JetDrag(spectral_grid, time_scale=Day(6))
+    forcing = StochasticStirring(spectral_grid, mask=false)
+    initial_conditions = StartFromRest()
+
+    model = BarotropicModel(;spectral_grid, initial_conditions, forcing, drag)
+    simulation = initialize!(model)
+
+    run!(simulation, period=Day(5))
+    @test simulation.model.feedback.nars_detected == false
+end