Merge branch 'main' into mg/dummy-gpu-ci

CHANGELOG.md

@@ -3,6 +3,11 @@
 ## Unreleased
 
 - Buildkite CI with dummy pipeline [#646](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/646)
+- ConvectiveHeating implemented [#639](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/639)
+- Number format flexibility with set! [#634](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/634)
+- Forcing/drag for primitive models [#635](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/635)
+- RingGrids indexing with leading Colon should now always return another RingGrid instance [#637](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/637)
+- 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)
 
 ## v0.13

Project.toml

@@ -56,12 +56,12 @@ Enzyme = "0.13"
 FFTW = "1"
 FastGaussQuadrature = "0.4, 0.5, 1"
 FiniteDifferences = "0.12"
-GPUArrays = "11"
+GPUArrays = "10"
 GenericFFT = "0.1"
 GeoMakie = "0.7.6"
-JLArrays = "0.1.4, 0.2"
+JLArrays = "0.1"
 JLD2 = "0.4, 0.5"
-KernelAbstractions = "0.9"
+KernelAbstractions = "0.9.0 - 0.9.29"
 LinearAlgebra = "1.10"
 Makie = "0.20, 0.21"
 NCDatasets = "0.12, 0.13, 0.14"

docs/src/forcing_drag.md

@@ -197,8 +197,8 @@ function SpeedyWeather.forcing!(
     diagn::DiagnosticVariables,
     progn::PrognosticVariables,
     forcing::StochasticStirring,
-    model::AbstractModel,
     lf::Integer,
+    model::AbstractModel,
 )
     # function barrier only
     forcing!(diagn, forcing, model.spectral_transform)

src/LowerTriangularMatrices/lower_triangular_array.jl

@@ -630,10 +630,10 @@ function Base.similar(
     return LowerTriangularArray{T, N, ArrayType{T,N}}(undef, size(L; as=Matrix))
 end
 
-function KernelAbstractions.get_backend(
-    a::LowerTriangularArray{T, N, ArrayType}
+function GPUArrays.backend(
+    ::Type{LowerTriangularArray{T, N, ArrayType}}
 ) where {T, N, ArrayType <: GPUArrays.AbstractGPUArray}
-    return KernelAbstractions.get_backend(a.data)
+    return GPUArrays.backend(ArrayType)
 end
 
 Adapt.adapt_structure(to, L::LowerTriangularArray) =

src/RingGrids/general.jl

@@ -76,7 +76,7 @@ Base.@propagate_inbounds Base.getindex(G::AbstractGridArray, ijk...) = getindex(
 @inline function Base.getindex(
     G::GridArray,
     col::Colon,
-    k::Integer...,
+    k...,
 ) where {GridArray<:AbstractGridArray}
     GridArray_ = nonparametric_type(GridArray)  # obtain parameters from G.data
     return GridArray_(getindex(G.data, col, k...), G.nlat_half, G.rings)
@@ -481,11 +481,11 @@ AbstractGPUGridArrayStyle{2, ArrayType, Grid}(::Val{3}) where {ArrayType, Grid}
 AbstractGPUGridArrayStyle{2, ArrayType, Grid}(::Val{1}) where {ArrayType, Grid} = AbstractGPUGridArrayStyle{2, ArrayType, Grid}()
 AbstractGPUGridArrayStyle{3, ArrayType, Grid}(::Val{4}) where {ArrayType, Grid} = AbstractGPUGridArrayStyle{4, ArrayType, Grid}()
 AbstractGPUGridArrayStyle{3, ArrayType, Grid}(::Val{2}) where {ArrayType, Grid} = AbstractGPUGridArrayStyle{3, ArrayType, Grid}()
-    
-function KernelAbstractions.get_backend(
-    g::Grid
+
+function GPUArrays.backend(
+    ::Type{Grid}
 ) where {Grid <: AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType <: GPUArrays.AbstractGPUArray}
-    return KernelAbstractions.get_backend(g.data)
+    return GPUArrays.backend(ArrayType)
 end
 
 function Base.similar(

src/SpeedyTransforms/fourier.jl

@@ -33,6 +33,7 @@ function _fourier_batched!(                 # GRID TO SPECTRAL
     (; rfft_plans) = S                      # pre-planned transforms
     nlayers = size(grids, 2)                # number of vertical layers
 
+    @assert eltype(grids) == eltype(S) "Number format of grid $(eltype(grids)) and SpectralTransform $(eltype(S)) need too match."
     @boundscheck ismatching(S, grids) || throw(DimensionMismatch(S, grids))
     @boundscheck nlayers == S.nlayers || throw(DimensionMismatch(S, grids))
     @boundscheck size(f_north) == size(f_south) == (S.nfreq_max, S.nlayers, nlat_half) || throw(DimensionMismatch(S, grids))
@@ -82,6 +83,7 @@ function _fourier_serial!(                  # GRID TO SPECTRAL
     rfft_plans = S.rfft_plans_1D            # pre-planned transforms
     nlayers = size(grids, 2)                # number of vertical layers
 
+    @assert eltype(grids) == eltype(S) "Number format of grid $(eltype(grids)) and SpectralTransform $(eltype(S)) need too match."
     @boundscheck ismatching(S, grids) || throw(DimensionMismatch(S, grids))
     @boundscheck nlayers <= S.nlayers || throw(DimensionMismatch(S, grids))
     @boundscheck size(f_north) == size(f_south) == (S.nfreq_max, S.nlayers, nlat_half) || throw(DimensionMismatch(S, grids))

src/SpeedyTransforms/spectral_transform.jl

@@ -85,6 +85,9 @@ struct SpectralTransform{
     eigenvalues⁻¹::VectorType           # = -1/(l*(l+1))
 end
 
+# eltype of a transform is the number format used within
+Base.eltype(S::SpectralTransform{NF}) where NF = NF
+
 """
 $(TYPEDSIGNATURES)
 Generator function for a SpectralTransform struct. With `NF` the number format,

src/dynamics/drag.jl

@@ -1,5 +1,15 @@
 abstract type AbstractDrag <: AbstractModelComponent end
 
+# function barrier for all drags to unpack model.drag
+function drag!(
+    diagn::DiagnosticVariables,
+    progn::PrognosticVariables,
+    lf::Integer,
+    model::AbstractModel,
+)
+    drag!(diagn, progn, model.drag, lf, model)
+end
+
 ## NO DRAG
 export NoDrag
 struct NoDrag <: AbstractDrag end
@@ -9,8 +19,7 @@ initialize!(::NoDrag, ::AbstractModel) = nothing
 function drag!(     diagn::DiagnosticVariables,
                     progn::PrognosticVariables,
                     drag::NoDrag,
-                    model::AbstractModel,
-                    lf::Integer)
+                    args...)
     return nothing
 end
 
@@ -37,8 +46,8 @@ function drag!(
     diagn::DiagnosticVariables,
     progn::PrognosticVariables,
     drag::QuadraticDrag,
-    model::AbstractModel,
     lf::Integer,
+    model::AbstractModel,
 )
     drag!(diagn, drag)
 end
@@ -71,4 +80,70 @@ function drag!(
         Fu[ij, k] -= c*speed*u[ij, k]     # -= as the tendencies already contain forcing
         Fv[ij, k] -= c*speed*v[ij, k]
     end
+end
+
+export JetDrag
+@kwdef struct JetDrag{NF, SpectralVariable2D} <: SpeedyWeather.AbstractDrag
+
+    # DIMENSIONS from SpectralGrid
+    "Spectral resolution as max degree of spherical harmonics"
+    trunc::Int
+
+    "[OPTION] Relaxation time scale τ"
+    time_scale::Second = Day(6)
+
+    "[OPTION] Jet strength [m/s]"
+    u₀::NF = 20
+
+    "[OPTION] latitude of Gaussian jet [˚N]"
+    latitude::NF = 30
+
+    "[OPTION] Width of Gaussian jet [˚]"
+    width::NF = 6
+
+    # TO BE INITIALISED
+    "Relaxation back to reference vorticity"
+    ζ₀::SpectralVariable2D = zeros(LowerTriangularMatrix{Complex{NF}}, trunc+2, trunc+1)
+end
+
+function JetDrag(SG::SpectralGrid; kwargs...)
+    return JetDrag{SG.NF, SG.SpectralVariable2D}(; SG.trunc, kwargs...)
+end
+
+function initialize!(drag::JetDrag, model::AbstractModel)
+    (; 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
+
+    û = transform(u, model.spectral_transform)
+    v̂ = zero(û)
+    curl!(drag.ζ₀, û, v̂, model.spectral_transform)
+    return nothing
+end
+
+function drag!(
+    diagn::DiagnosticVariables,
+    progn::PrognosticVariables,
+    drag::JetDrag,
+    lf::Integer,
+    model::AbstractModel,
+)
+    vor = progn.vor[lf]
+    (; vor_tend) = diagn.tendencies
+    (; ζ₀) = drag
+
+    # scale by radius as is vorticity
+    (; radius) = model.spectral_grid
+    r = radius/drag.time_scale.value
+
+    k = diagn.nlayers   # drag only on surface layer
+    for lm in eachharmonic(vor_tend)
+        vor_tend[lm, k] -= r*(vor[lm, k] - ζ₀[lm])
+    end
 end

src/dynamics/forcing.jl

@@ -1,16 +1,27 @@
 abstract type AbstractForcing <: AbstractModelComponent end
 
+# function barrier for all forcings to unpack model.forcing
+function forcing!(
+    diagn::DiagnosticVariables,
+    progn::PrognosticVariables,
+    lf::Integer,
+    model::AbstractModel,
+)
+    forcing!(diagn, progn, model.forcing, lf, model)
+end
+
 ## NO FORCING = dummy forcing
 export NoForcing
 struct NoForcing <: AbstractForcing end
 NoForcing(SG::SpectralGrid) = NoForcing()
 initialize!(::NoForcing, ::AbstractModel) = nothing
 
-function forcing!(  diagn::DiagnosticVariables,
-                    progn::PrognosticVariables,
-                    forcing::NoForcing,
-                    model::AbstractModel,
-                    lf::Integer)
+function forcing!(  
+    diagn::DiagnosticVariables,
+    progn::PrognosticVariables,
+    forcing::NoForcing,
+    args...,
+)
     return nothing
 end
 
@@ -98,8 +109,8 @@ function forcing!(
     diagn::DiagnosticVariables,
     progn::PrognosticVariables,
     forcing::JetStreamForcing,
-    model::AbstractModel,
     lf::Integer,
+    model::AbstractModel,
 )
     forcing!(diagn, forcing)
 end
@@ -119,8 +130,84 @@ function forcing!(
         for (j, ring) in enumerate(eachring(Fu))
             F = amplitude[j]
             for ij in ring
-                Fu[ij] = tapering[k]*F
+                # += to accumulate, not overwrite previous parameterizations/terms
+                Fu[ij] += tapering[k]*F
             end
         end
     end
+end
+
+export StochasticStirring
+@kwdef struct StochasticStirring{NF, VectorType} <: AbstractForcing
+
+    "Number of latitude rings, used for latitudinal mask"
+    nlat::Int
+
+    "[OPTION] Stirring strength A [1/s²]"
+    strength::NF = 1e-9
+
+    "[OPTION] Stirring latitude [˚N]"
+    latitude::NF = 45
+
+    "[OPTION] Stirring width [˚]"
+    width::NF = 24
+
+    # TO BE INITIALISED        
+    "Latitudinal mask, confined to mid-latitude storm track by default [1]"
+    lat_mask::VectorType = zeros(NF, nlat)
+end
+
+function StochasticStirring(SG::SpectralGrid; kwargs...)
+    return StochasticStirring{SG.NF, SG.VectorType}(; nlat=SG.nlat, kwargs...)
+end
+
+function initialize!(
+    forcing::StochasticStirring,
+    model::AbstractModel)
+
+    model.random_process isa NoRandomProcess &&
+        @warn "StochasticStirring needs a random process. model.random_process is a NoRandomProcess."
+
+    # precompute the latitudinal mask
+    (; latd) = model.geometry
+    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
+
+function forcing!(
+    diagn::DiagnosticVariables,
+    progn::PrognosticVariables,
+    forcing::StochasticStirring,
+    lf::Integer,
+    model::AbstractModel,
+)
+    forcing!(diagn, forcing, model.spectral_transform)
+end
+
+function forcing!(
+    diagn::DiagnosticVariables,
+    forcing::StochasticStirring,
+    spectral_transform::SpectralTransform,
+)
+    # get random values from random process
+    S_grid = diagn.grid.random_pattern
+
+    # mask everything but mid-latitudes
+    RingGrids._scale_lat!(S_grid, forcing.lat_mask)
+
+    # back to spectral space
+    S_masked = diagn.dynamics.a_2D
+    transform!(S_masked, S_grid, spectral_transform)
+
+    # scale by radius^2 as is the vorticity equation, and scale to forcing strength
+    S_masked .*= (diagn.scale[]^2 * forcing.strength)
+
+    # force every layer
+    (; vor_tend) = diagn.tendencies
+
+    for k in eachmatrix(vor_tend)
+        vor_tend[:, k] .+= S_masked
+    end
 end