From cbe761e433c7297e73c948c782024e29b1f3f044 Mon Sep 17 00:00:00 2001 From: Milan Date: Thu, 14 Dec 2023 00:26:23 +0000 Subject: [PATCH 1/5] Make prognostics available for forcing, drag in Barotropic/SWM --- docs/src/extending.md | 6 +++++- src/dynamics/drag.jl | 4 +++- src/dynamics/forcing.jl | 2 ++ src/dynamics/tendencies.jl | 12 +++++++----- 4 files changed, 17 insertions(+), 7 deletions(-) diff --git a/docs/src/extending.md b/docs/src/extending.md index 6a9986a70..b44972249 100644 --- a/docs/src/extending.md +++ b/docs/src/extending.md @@ -82,6 +82,7 @@ that is called on _every_ step of the time integration. This new method for `forcing!` needs to have the following function signature ```julia function forcing!( diagn::DiagnosticVariablesLayer, + progn::PrognosticVariablesLayer, forcing::MyForcing, time::DateTime, model::ModelSetup) @@ -339,6 +340,7 @@ then this will be called automatically with multiple dispatch. ```@example extend function SpeedyWeather.forcing!(diagn::SpeedyWeather.DiagnosticVariablesLayer, + progn::SpeedyWeather.PrognosticVariablesLayer, forcing::StochasticStirring, time::DateTime, model::SpeedyWeather.ModelSetup) @@ -349,7 +351,9 @@ end The function has to be as outlined above. The first argument has to be of type ``DiagnosticVariablesLayer`` as it acts on a layer of the diagnostic variables, where the current model state in grid-point space and the tendencies (in spectral space) -are defined. The second argument has to be of the type of our new forcing, +are defined. The second argument has to be a `PrognosticVariablesLayer` because, +in general, the forcing may use the prognostic variables in spectral space. +The third argument has to be of the type of our new forcing, the third argument is time which you may use or skip, the last element is a `ModelSetup`, but as before you can be more restrictive to define a forcing only for the `BarotropicModel` for example, use ``model::Barotropic`` in that case. diff --git a/src/dynamics/drag.jl b/src/dynamics/drag.jl index 50589c68f..629f7812d 100644 --- a/src/dynamics/drag.jl +++ b/src/dynamics/drag.jl @@ -14,6 +14,7 @@ function initialize!( drag::NoDrag, end function drag!( diagn::DiagnosticVariablesLayer, + progn::PrognosticVariablesLayer, drag::NoDrag, time::DateTime, model::ModelSetup) @@ -23,7 +24,7 @@ end # Quadratic drag Base.@kwdef struct QuadraticDrag{NF} <: AbstractDrag{NF} "drag coefficient [1]" - c_D::Float64 = 1e-5 + c_D::NF = 1e-5 end QuadraticDrag(SG::SpectralGrid;kwargs...) = QuadraticDrag{SG.NF}(;kwargs...) @@ -35,6 +36,7 @@ end # function barrier function drag!( diagn::DiagnosticVariablesLayer, + progn::PrognosticVariablesLayer, drag::QuadraticDrag, time::DateTime, model::ModelSetup) diff --git a/src/dynamics/forcing.jl b/src/dynamics/forcing.jl index a087c72e3..9f4d194d8 100644 --- a/src/dynamics/forcing.jl +++ b/src/dynamics/forcing.jl @@ -14,6 +14,7 @@ function initialize!( forcing::NoForcing, end function forcing!( diagn::DiagnosticVariablesLayer, + progn::PrognosticVariablesLayer, forcing::NoForcing, time::DateTime, model::ModelSetup) @@ -83,6 +84,7 @@ end # function barrier function forcing!( diagn::DiagnosticVariablesLayer, + progn::PrognosticVariablesLayer, forcing::JetStreamForcing, time::DateTime, model::ModelSetup) diff --git a/src/dynamics/tendencies.jl b/src/dynamics/tendencies.jl index 5bbc365f1..4de28b857 100644 --- a/src/dynamics/tendencies.jl +++ b/src/dynamics/tendencies.jl @@ -2,17 +2,19 @@ $(TYPEDSIGNATURES) Calculate all tendencies for the BarotropicModel.""" function dynamics_tendencies!( diagn::DiagnosticVariablesLayer, + progn::PrognosticVariablesLayer, time::DateTime, model::Barotropic) - forcing!(diagn,model.forcing,time,model) # = (Fᵤ, Fᵥ) forcing for u,v - drag!(diagn,model.drag,time,model) # drag term for u,v - vorticity_flux!(diagn,model) # = ∇×(v(ζ+f) + Fᵤ,-u(ζ+f) + Fᵥ) + forcing!(diagn,progn,model.forcing,time,model) # = (Fᵤ, Fᵥ) forcing for u,v + drag!(diagn,progn,model.drag,time,model) # drag term for u,v + vorticity_flux!(diagn,model) # = ∇×(v(ζ+f) + Fᵤ,-u(ζ+f) + Fᵥ) end """ $(TYPEDSIGNATURES) Calculate all tendencies for the ShallowWaterModel.""" function dynamics_tendencies!( diagn::DiagnosticVariablesLayer, + progn::PrognosticVariablesLayer, surface::SurfaceVariables, pres::LowerTriangularMatrix, # spectral pressure/η for geopotential time::DateTime, # time to evaluate the tendencies at @@ -22,8 +24,8 @@ function dynamics_tendencies!( diagn::DiagnosticVariablesLayer, F,D = model.forcing, model.drag # for compatibility with other ModelSetups pressure pres = interface displacement η here - forcing!(diagn,F,time,model) # = (Fᵤ, Fᵥ, Fₙ) forcing for u,v,η - drag!(diagn,D,time,model) # drag term for momentum u,v + forcing!(diagn,progn,F,time,model) # = (Fᵤ, Fᵥ, Fₙ) forcing for u,v,η + drag!(diagn,progn,D,time,model) # drag term for momentum u,v vorticity_flux!(diagn,model) # = ∇×(v(ζ+f) + Fᵤ,-u(ζ+f) + Fᵥ), tendency for vorticity # = ∇⋅(v(ζ+f) + Fᵤ,-u(ζ+f) + Fᵥ), tendency for divergence From 3ffc815c6df5f93566a8d5af1357358d99b30ad8 Mon Sep 17 00:00:00 2001 From: Milan Date: Thu, 14 Dec 2023 00:27:28 +0000 Subject: [PATCH 2/5] Move lf outside of gridded! Move spectral_truncation! into time_integration --- src/dynamics/tendencies_dynamics.jl | 52 ++++++++++--------------- src/dynamics/time_integration.jl | 59 +++++++++++++++++------------ 2 files changed, 55 insertions(+), 56 deletions(-) diff --git a/src/dynamics/tendencies_dynamics.jl b/src/dynamics/tendencies_dynamics.jl index 4e3487aa5..f21132d86 100644 --- a/src/dynamics/tendencies_dynamics.jl +++ b/src/dynamics/tendencies_dynamics.jl @@ -131,7 +131,6 @@ function surface_pressure_tendency!( @. pres_tend = -pres_tend - D̄ # the -D̄ term in spectral and swap sign pres_tend[1] = 0 # for mass conservation - spectral_truncation!(pres_tend) # remove lmax+1 row, only vectors use it return nothing end @@ -292,8 +291,6 @@ function temperature_tendency!( # now add the -∇⋅((u,v)*T') term flux_divergence!(temp_tend,temp_grid,diagn,G,S,add=true,flipsign=true) - # only vectors make use of the lmax+1 row, set to zero for scalars - spectral_truncation!(temp_tend) return nothing end @@ -505,7 +502,6 @@ function bernoulli_potential!( diagn::DiagnosticVariablesLayer{NF}, spectral!(bernoulli,bernoulli_grid,S) # to spectral space bernoulli .+= geopot # add geopotential Φ ∇²!(div_tend,bernoulli,S,add=true,flipsign=true) # add -∇²(½(u² + v²) + ϕ) - spectral_truncation!(div_tend) # set lmax+1 row to zero end """ @@ -574,7 +570,8 @@ function SpeedyTransforms.gridded!( # all variables on layers for (progn_layer,diagn_layer) in zip(progn.layers,diagn.layers) - gridded!(diagn_layer,progn_layer,lf,model) + progn_layer_lf = progn_layer.timesteps[lf] + gridded!(diagn_layer,progn_layer_lf,model) end # surface only for ShallowWaterModel or PrimitiveEquation @@ -592,22 +589,21 @@ Propagate the spectral state of the prognostic variables `progn` to the diagnostic variables in `diagn` for the barotropic vorticity model. Updates grid vorticity, spectral stream function and spectral and grid velocities u,v.""" function SpeedyTransforms.gridded!( diagn::DiagnosticVariablesLayer, - progn::PrognosticLayerTimesteps, - lf::Int, # leapfrog index + progn::PrognosticVariablesLayer, model::Barotropic) (; vor_grid, u_grid, v_grid ) = diagn.grid_variables + (; vor ) = progn # relative vorticity U = diagn.dynamics_variables.a # reuse work arrays for velocities in spectral V = diagn.dynamics_variables.b # U = u*coslat, V=v*coslat S = model.spectral_transform - vor_lf = progn.timesteps[lf].vor # relative vorticity at leapfrog step lf - gridded!(vor_grid,vor_lf,S) # get vorticity on grid from spectral vor + gridded!(vor_grid,vor,S) # get vorticity on grid from spectral vor # get spectral U,V from spectral vorticity via stream function Ψ # U = u*coslat = -coslat*∂Ψ/∂lat # V = v*coslat = ∂Ψ/∂lon, radius omitted in both cases - UV_from_vor!(U,V,vor_lf,S) + UV_from_vor!(U,V,vor,S) # transform from U,V in spectral to u,v on grid (U,V = u,v*coslat) gridded!(u_grid,U,S,unscale_coslat=true) @@ -623,26 +619,23 @@ diagnostic variables in `diagn` for the shallow water model. Updates grid vortic grid divergence, grid interface displacement (`pres_grid`) and the velocities u,v.""" function SpeedyTransforms.gridded!( diagn::DiagnosticVariablesLayer, - progn::PrognosticLayerTimesteps, - lf::Int, # leapfrog index + progn::PrognosticVariablesLayer, model::ShallowWater, # everything that's constant ) (; vor_grid, div_grid, u_grid, v_grid ) = diagn.grid_variables + (; vor, div) = progn U = diagn.dynamics_variables.a # reuse work arrays for velocities spectral V = diagn.dynamics_variables.b # U = u*coslat, V=v*coslat S = model.spectral_transform - vor_lf = progn.timesteps[lf].vor # pick leapfrog index without memory allocation - div_lf = progn.timesteps[lf].div - # get spectral U,V from vorticity and divergence via stream function Ψ and vel potential ϕ # U = u*coslat = -coslat*∂Ψ/∂lat + ∂ϕ/dlon # V = v*coslat = coslat*∂ϕ/∂lat + ∂Ψ/dlon - UV_from_vordiv!(U,V,vor_lf,div_lf,S) + UV_from_vordiv!(U,V,vor,div,S) - gridded!(vor_grid,vor_lf,S) # get vorticity on grid from spectral vor - gridded!(div_grid,div_lf,S) # get divergence on grid from spectral div + gridded!(vor_grid,vor,S) # get vorticity on grid from spectral vor + gridded!(div_grid,div,S) # get divergence on grid from spectral div # transform from U,V in spectral to u,v on grid (U,V = u,v*coslat) gridded!(u_grid,U,S,unscale_coslat=true) @@ -658,13 +651,13 @@ diagnostic variables in `diagn` for primitive equation models. Updates grid vort grid divergence, grid temperature, pressure (`pres_grid`) and the velocities u,v.""" function SpeedyTransforms.gridded!( diagn::DiagnosticVariablesLayer, - progn::PrognosticLayerTimesteps, - lf::Int, # leapfrog index + progn::PrognosticVariablesLayer, model::PrimitiveEquation) # everything that's constant (; vor_grid, div_grid, u_grid, v_grid ) = diagn.grid_variables (; temp_grid, humid_grid ) = diagn.grid_variables (; temp_grid_prev, u_grid_prev, v_grid_prev) = diagn.grid_variables + (; vor, div, temp, humid) = progn U = diagn.dynamics_variables.a # reuse work arrays for velocities spectral V = diagn.dynamics_variables.b # U = u*coslat, V=v*coslat @@ -675,24 +668,19 @@ function SpeedyTransforms.gridded!( diagn::DiagnosticVariablesLayer, S = model.spectral_transform wet_core = model isa PrimitiveWet - vor_lf = progn.timesteps[lf].vor # pick leapfrog index without memory allocation - div_lf = progn.timesteps[lf].div - temp_lf = progn.timesteps[lf].temp - wet_core && (humid_lf = progn.timesteps[lf].humid) - # get spectral U,V from vorticity and divergence via stream function Ψ and vel potential ϕ # U = u*coslat = -coslat*∂Ψ/∂lat + ∂ϕ/dlon # V = v*coslat = coslat*∂ϕ/∂lat + ∂Ψ/dlon - UV_from_vordiv!(U,V,vor_lf,div_lf,S) + UV_from_vordiv!(U,V,vor,div,S) - gridded!(vor_grid,vor_lf,S) # get vorticity on grid from spectral vor - gridded!(div_grid,div_lf,S) # get divergence on grid from spectral div - gridded!(temp_grid,temp_lf,S) # (absolute) temperature - wet_core && gridded!(humid_grid,humid_lf,S) # specific humidity (wet core only) + gridded!(vor_grid,vor,S) # get vorticity on grid from spectral vor + gridded!(div_grid,div,S) # get divergence on grid from spectral div + gridded!(temp_grid,temp,S) # (absolute) temperature + wet_core && gridded!(humid_grid,humid,S)# specific humidity (wet core only) # include humidity effect into temp for everything stability-related - temperature_average!(diagn,temp_lf,S) # TODO: do at frequency of reinitialize implicit? - virtual_temperature!(diagn,temp_lf,model) # temp = virt temp for dry core + temperature_average!(diagn,temp,S) # TODO: do at frequency of reinitialize implicit? + virtual_temperature!(diagn,temp,model) # temp = virt temp for dry core # transform from U,V in spectral to u,v on grid (U,V = u,v*coslat) gridded!(u_grid,U,S,unscale_coslat=true) diff --git a/src/dynamics/time_integration.jl b/src/dynamics/time_integration.jl index a7b21f5de..a838e5fb7 100644 --- a/src/dynamics/time_integration.jl +++ b/src/dynamics/time_integration.jl @@ -132,7 +132,6 @@ function leapfrog!( A_old::LowerTriangularMatrix{Complex{NF}}, # prognostic A_lf = lf == 1 ? A_old : A_new # view on either t or t+dt to dis/enable Williams filter (;robert_filter, williams_filter) = L # coefficients for the Robert and Williams filter - two = convert(NF,2) # 2 in number format NF dt_NF = convert(NF,dt) # time step dt in number format NF # LEAP FROG time step with or without Robert+Williams filter @@ -140,13 +139,13 @@ function leapfrog!( A_old::LowerTriangularMatrix{Complex{NF}}, # prognostic # see Williams (2009), Eq. 7-9 # for lf == 1 (initial time step) no filter applied (w1=w2=0) # for lf == 2 (later steps) Robert+Williams filter is applied - w1 = lf == 1 ? zero(NF) : robert_filter*williams_filter/two # = ν*α/2 in Williams (2009, Eq. 8) - w2 = lf == 1 ? zero(NF) : robert_filter*(1-williams_filter)/two # = ν(1-α)/2 in Williams (2009, Eq. 9) + w1 = lf == 1 ? zero(NF) : robert_filter*williams_filter/2 # = ν*α/2 in Williams (2009, Eq. 8) + w2 = lf == 1 ? zero(NF) : robert_filter*(1-williams_filter)/2 # = ν(1-α)/2 in Williams (2009, Eq. 9) @inbounds for lm in eachharmonic(A_old,A_new,A_lf,tendency) a_old = A_old[lm] # double filtered value from previous time step (t-Δt) a_new = a_old + dt_NF*tendency[lm] # Leapfrog/Euler step depending on dt=Δt,2Δt (unfiltered at t+Δt) - a_update = a_old - two*A_lf[lm] + a_new # Eq. 8&9 in Williams (2009), calculate only once + a_update = a_old - 2A_lf[lm] + a_new # Eq. 8&9 in Williams (2009), calculate only once A_old[lm] = A_lf[lm] + w1*a_update # Robert's filter: A_old[lm] becomes 2xfiltered value at t A_new[lm] = a_new - w2*a_update # Williams filter: A_new[lm] becomes 1xfiltered value at t+Δt end @@ -168,10 +167,24 @@ function leapfrog!( progn::PrognosticLayerTimesteps, var_old = getproperty(progn.timesteps[1],var) var_new = getproperty(progn.timesteps[2],var) var_tend = getproperty(diagn.tendencies,Symbol(var,:_tend)) + spectral_truncation!(var_tend) # set lmax+1 mode to zero leapfrog!(var_old,var_new,var_tend,dt,lf,model.time_stepping) end end +function leapfrog!( progn::PrognosticSurfaceTimesteps, + diagn::SurfaceVariables, + dt::Real, # time step (mostly =2Δt, but for init steps =Δt,Δt/2) + lf::Int, # leapfrog index to dis/enable Williams filter + model::ModelSetup) + + (;pres_tend) = diagn + pres_old = progn.timesteps[1].pres + pres_new = progn.timesteps[2].pres + spectral_truncation!(pres_tend) # set lmax+1 mode to zero + leapfrog!(pres_old,pres_new,pres_tend,dt,lf,model.time_stepping) +end + """ $(TYPEDSIGNATURES) Performs the first two initial time steps (Euler forward, unfiltered leapfrog) to populate the @@ -184,7 +197,7 @@ function first_timesteps!( ) (;clock) = progn - clock.n_timesteps == 0 && return time # exit immediately for no time steps + clock.n_timesteps == 0 && return nothing # exit immediately for no time steps (;implicit) = model (;Δt, Δt_millisec) = model.time_stepping @@ -211,7 +224,7 @@ function first_timesteps!( # from now on precomputed implicit terms with 2Δt initialize!(implicit,2Δt,diagn,model) - return time + return nothing end """ @@ -230,10 +243,11 @@ function timestep!( progn::PrognosticVariables, # all prognostic variables # LOOP OVER LAYERS FOR TENDENCIES, DIFFUSION, LEAPFROGGING AND PROPAGATE STATE TO GRID for (progn_layer,diagn_layer) in zip(progn.layers,diagn.layers) - dynamics_tendencies!(diagn_layer,time,model) + progn_lf = progn_layer.timesteps[lf2] # pick the leapfrog time step lf2 for tendencies + dynamics_tendencies!(diagn_layer,progn_lf,time,model) horizontal_diffusion!(diagn_layer,progn_layer,model) leapfrog!(progn_layer,diagn_layer,dt,lf1,model) - gridded!(diagn_layer,progn_layer,lf2,model) + gridded!(diagn_layer,progn_lf,model) end end @@ -252,28 +266,26 @@ function timestep!( progn::PrognosticVariables{NF}, # all prognostic variables (;time) = progn.clock # current time progn_layer = progn.layers[1] # only calculate tendencies for the first layer - diagn_layer = diagn.layers[1] - diagn_surface = diagn.surface - progn_surface = progn.surface + diagn_layer = diagn.layers[1] # multi-layer shallow water not supported + + progn_lf = progn_layer.timesteps[lf2] # pick the leapfrog time step lf2 for tendencies (;pres) = progn.surface.timesteps[lf2] - (;implicit, time_stepping, spectral_transform) = model + (;implicit, spectral_transform) = model zero_tendencies!(diagn) # GET TENDENCIES, CORRECT THEM FOR SEMI-IMPLICIT INTEGRATION - dynamics_tendencies!(diagn_layer,diagn_surface,pres,time,model) - implicit_correction!(diagn_layer,progn_layer,diagn_surface,progn_surface,implicit) + dynamics_tendencies!(diagn_layer,progn_lf,diagn.surface,pres,time,model) + implicit_correction!(diagn_layer,progn_layer,diagn.surface,progn.surface,implicit) # APPLY DIFFUSION, STEP FORWARD IN TIME, AND TRANSFORM NEW TIME STEP TO GRID horizontal_diffusion!(progn_layer,diagn_layer,model) leapfrog!(progn_layer,diagn_layer,dt,lf1,model) - gridded!(diagn_layer,progn_layer,lf2,model) + gridded!(diagn_layer,progn_lf,model) # SURFACE LAYER (pressure), no diffusion though - (;pres_grid,pres_tend) = diagn.surface - pres_old = progn.surface.timesteps[1].pres - pres_new = progn.surface.timesteps[2].pres - leapfrog!(pres_old,pres_new,pres_tend,dt,lf1,time_stepping) + (;pres_grid) = diagn.surface + leapfrog!(progn.surface,diagn.surface,dt,lf1,model) gridded!(pres_grid,pres,spectral_transform) end @@ -312,16 +324,15 @@ function timestep!( progn::PrognosticVariables{NF}, # all prognostic variables if k <= diagn.nlev # model levels diagn_layer = diagn.layers[k] progn_layer = progn.layers[k] + progn_layer_lf = progn_layer.timesteps[lf2] horizontal_diffusion!(progn_layer,diagn_layer,model) # implicit diffusion of vor, div, temp leapfrog!(progn_layer,diagn_layer,dt,lf1,model) # time step forward for vor, div, temp - gridded!(diagn_layer,progn_layer,lf2,model) # propagate spectral state to grid + gridded!(diagn_layer,progn_layer_lf,model) # propagate spectral state to grid else # surface level - (;pres_grid,pres_tend) = diagn.surface - pres_old = progn.surface.timesteps[1].pres - pres_new = progn.surface.timesteps[2].pres + leapfrog!(progn.surface,diagn.surface,dt,lf1,model) + (;pres_grid) = diagn.surface pres_lf = progn.surface.timesteps[lf2].pres - leapfrog!(pres_old,pres_new,pres_tend,dt,lf1,model.time_stepping) gridded!(pres_grid,pres_lf,model.spectral_transform) end end From 675c384f22d15df8e8150765663525adb766915b Mon Sep 17 00:00:00 2001 From: Milan Date: Thu, 14 Dec 2023 00:27:55 +0000 Subject: [PATCH 3/5] spectral_truncation! for all rows below the triangle --- src/SpeedyTransforms/spectral_truncation.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/SpeedyTransforms/spectral_truncation.jl b/src/SpeedyTransforms/spectral_truncation.jl index a0b396d36..2b79ac5f1 100644 --- a/src/SpeedyTransforms/spectral_truncation.jl +++ b/src/SpeedyTransforms/spectral_truncation.jl @@ -71,7 +71,8 @@ to enforce that all coefficients for which the degree l is larger than order m a spectral_truncation!(alms::AbstractMatrix) = spectral_truncation!(alms,size(alms)...) function spectral_truncation!(alms::LowerTriangularMatrix{NF}) where NF - alms[end,:] .= zero(NF) # set the lmax+1 mode to zero + lmax,mmax = size(alms) + alms[mmax+1:lmax,:] .= 0 # set everything to zero below the triangle return alms end From 88abb307b60ec3710612035acdf4def771c81f33 Mon Sep 17 00:00:00 2001 From: Milan Date: Thu, 14 Dec 2023 00:28:06 +0000 Subject: [PATCH 4/5] adapt tests --- test/geopotential.jl | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/test/geopotential.jl b/test/geopotential.jl index 860d671a3..037a02857 100644 --- a/test/geopotential.jl +++ b/test/geopotential.jl @@ -11,9 +11,10 @@ temp = 280 # in Kelvin lf = 1 for (progn_layer,diagn_layer) in zip(p.layers,d.layers) - progn_layer.timesteps[lf].temp[1] = temp*model.spectral_transform.norm_sphere - fill!(progn_layer.timesteps[lf].humid,0) # dry core - SpeedyWeather.gridded!(diagn_layer,progn_layer,lf,model) # propagate spectral state to grid + progn_layer_lf = progn_layer.timesteps[lf] + progn_layer_lf.temp[1] = temp*model.spectral_transform.norm_sphere + fill!(progn_layer_lf.humid,0) # dry core + SpeedyWeather.gridded!(diagn_layer,progn_layer_lf,model) # propagate spectral state to grid SpeedyWeather.linear_virtual_temperature!(diagn_layer,progn_layer,model,lf) end @@ -49,7 +50,8 @@ end lf = 1 for (progn_layer,diagn_layer) in zip(p.layers,d.layers) - SpeedyWeather.gridded!(diagn_layer,progn_layer,lf,m) # propagate spectral state to grid + progn_layer_lf = progn_layer.timesteps[lf] + SpeedyWeather.gridded!(diagn_layer,progn_layer_lf,m) # propagate spectral state to grid SpeedyWeather.bernoulli_potential!(diagn_layer,m.spectral_transform) end end @@ -68,10 +70,11 @@ end temp = 280 # in Kelvin lf = 1 for (progn_layer,diagn_layer) in zip(p.layers,d.layers) - progn_layer.timesteps[lf].temp[1] = temp*m.spectral_transform.norm_sphere - fill!(progn_layer.timesteps[lf].humid,0) - progn_layer.timesteps[lf].humid[1] = rand(NF) - SpeedyWeather.gridded!(diagn_layer,progn_layer,lf,m) # propagate spectral state to grid + progn_layer_lf = progn_layer.timesteps[lf] + progn_layer_lf.temp[1] = temp*m.spectral_transform.norm_sphere + fill!(progn_layer_lf.humid,0) + progn_layer_lf.humid[1] = rand(NF) + SpeedyWeather.gridded!(diagn_layer,progn_layer_lf,m) # propagate spectral state to grid temp_lf = progn_layer.timesteps[lf].temp # always passed on for compat with DryCore SpeedyWeather.virtual_temperature!(diagn_layer,temp_lf,m) From 3c8259d456e16ff12e25868358e1d0d2e27886b1 Mon Sep 17 00:00:00 2001 From: Milan Date: Sat, 16 Dec 2023 12:07:22 +0100 Subject: [PATCH 5/5] Add forcing and drag extensions to tests --- test/extending.jl | 196 ++++++++++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 3 + 2 files changed, 199 insertions(+) create mode 100644 test/extending.jl diff --git a/test/extending.jl b/test/extending.jl new file mode 100644 index 000000000..495127303 --- /dev/null +++ b/test/extending.jl @@ -0,0 +1,196 @@ +@testset "Extending forcing and drag" begin + Base.@kwdef struct JetDrag{NF} <: SpeedyWeather.AbstractDrag{NF} + + # DIMENSIONS from SpectralGrid + "Spectral resolution as max degree of spherical harmonics" + trunc::Int + + # OPTIONS + "Relaxation time scale τ" + time_scale::Second = Day(6) + + "Jet strength [m/s]" + u₀::Float64 = 20 + + "latitude of Gaussian jet [˚N]" + latitude::Float64 = 30 + + "Width of Gaussian jet [˚]" + width::Float64 = 6 + + # TO BE INITIALISED + "Relaxation back to reference vorticity" + ζ₀::LowerTriangularMatrix{Complex{NF}} = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) + end + + function JetDrag(SG::SpectralGrid;kwargs...) + return JetDrag{SG.NF}(;SG.trunc,kwargs...) + end + + function SpeedyWeather.initialize!( drag::JetDrag, + model::SpeedyWeather.ModelSetup) + + (;spectral_grid, geometry) = model + (;Grid,NF,nlat_half) = spectral_grid + u = zeros(Grid{NF},nlat_half) + + lat = geometry.latds + + for ij in eachindex(u) + u[ij] = drag.u₀ * exp(-(lat[ij]-drag.latitude)^2/(2*drag.width^2)) + end + + û = SpeedyTransforms.spectral(u,one_more_degree=true) + v̂ = zero(û) + SpeedyTransforms.curl!(drag.ζ₀,û,v̂,model.spectral_transform) + return nothing + end + + function SpeedyWeather.drag!( diagn::SpeedyWeather.DiagnosticVariablesLayer, + progn::SpeedyWeather.PrognosticVariablesLayer, + drag::JetDrag, + time::DateTime, + model::SpeedyWeather.ModelSetup) + + (;vor) = progn + (;vor_tend) = diagn.tendencies + (;ζ₀) = drag + + (;radius) = model.spectral_grid + r = radius/drag.time_scale.value + 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} + + # DIMENSIONS from SpectralGrid + "Spectral resolution as max degree of spherical harmonics" + trunc::Int + + "Number of latitude rings, used for latitudinal mask" + nlat::Int + + + # OPTIONS + "Decorrelation time scale τ [days]" + decorrelation_time::Second = Day(2) + + "Stirring strength A [1/s²]" + strength::Float64 = 1e-11 + + "Stirring latitude [˚N]" + latitude::Float64 = 45 + + "Stirring width [˚]" + width::Float64 = 24 + + + # TO BE INITIALISED + "Stochastic stirring term S" + S::LowerTriangularMatrix{Complex{NF}} = zeros(LowerTriangularMatrix{Complex{NF}},trunc+2,trunc+1) + + "a = A*sqrt(1 - exp(-2dt/τ)), the noise factor times the stirring strength [1/s²]" + a::Base.RefValue{NF} = Ref(zero(NF)) + + "b = exp(-dt/τ), the auto-regressive factor [1]" + b::Base.RefValue{NF} = Ref(zero(NF)) + + "Latitudinal mask, confined to mid-latitude storm track by default [1]" + lat_mask::Vector{NF} = zeros(NF,nlat) + end + + function StochasticStirring(SG::SpectralGrid;kwargs...) + (;trunc,Grid,nlat_half) = SG + nlat = RingGrids.get_nlat(Grid,nlat_half) + return StochasticStirring{SG.NF}(;trunc,nlat,kwargs...) + end + + function SpeedyWeather.initialize!( forcing::StochasticStirring, + model::SpeedyWeather.ModelSetup) + + # precompute forcing strength, scale with radius^2 as is the vorticity equation + (;radius) = model.spectral_grid + A = radius^2 * forcing.strength + + # precompute noise and auto-regressive factor, packed in RefValue for mutability + dt = model.time_stepping.Δt_sec + τ = forcing.decorrelation_time.value # in seconds + forcing.a[] = A*sqrt(1 - exp(-2dt/τ)) + forcing.b[] = exp(-dt/τ) + + # precompute the latitudinal mask + (;Grid,nlat_half) = model.spectral_grid + latd = RingGrids.get_latd(Grid,nlat_half) + + 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 + + return nothing + end + + function SpeedyWeather.forcing!(diagn::SpeedyWeather.DiagnosticVariablesLayer, + progn::SpeedyWeather.PrognosticVariablesLayer, + forcing::StochasticStirring, + time::DateTime, + model::SpeedyWeather.ModelSetup) + SpeedyWeather.forcing!(diagn,forcing,model.spectral_transform) + end + + function SpeedyWeather.forcing!(diagn::SpeedyWeather.DiagnosticVariablesLayer, + forcing::StochasticStirring{NF}, + spectral_transform::SpectralTransform) where NF + + # noise and auto-regressive factors + a = forcing.a[] # = sqrt(1 - exp(-2dt/τ)) + 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] + end + + # to grid-point space + S_grid = diagn.dynamics_variables.a_grid + SpeedyTransforms.gridded!(S_grid,S,spectral_transform) + + # mask everything but mid-latitudes + RingGrids._scale_lat!(S_grid,forcing.lat_mask) + + # back to spectral space + (;vor_tend) = diagn.tendencies + SpeedyTransforms.spectral!(vor_tend,S_grid,spectral_transform) + SpeedyTransforms.spectral_truncation!(vor_tend) # set lmax+1 to zero + + return nothing + end + + spectral_grid = SpectralGrid(trunc=31,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) + + run!(simulation,period=Day(5)) + @test simulation.model.feedback.nars_detected == false + + # with shallow water model + model = ShallowWaterModel(;spectral_grid,initial_conditions,forcing,drag) + simulation = initialize!(model) + + run!(simulation,period=Day(5)) + @test simulation.model.feedback.nars_detected == false +end \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index 5cc276a07..69dd0c752 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -36,5 +36,8 @@ include("column_variables.jl") # INITIALIZATION AND INTEGRATION include("run_speedy.jl") +# EXTENSION +include("extending.jl") + # OUTPUT include("netcdf_output.jl") \ No newline at end of file