Merge branch 'main' into mk/summarysize

docs/src/callbacks.md

@@ -184,13 +184,17 @@ add!(callbacks, :key1 => NoCallback(), :key2 => NoCallback())   # keys provided
 ```
 Meaning that callbacks can be added before and after model construction
 
-
 ```@example callbacks
 spectral_grid = SpectralGrid()
 callbacks = CallbackDict(:callback_added_before => NoCallback())
 model = PrimitiveWetModel(; spectral_grid, callbacks)
 add!(model.callbacks, :callback_added_afterwards => NoCallback())
+add!(model, :callback_added_afterwards2 => NoCallback())
 ```
+
+Note how the first argument can be `model.callbacks` as outlined in the sections above
+because this is the callbacks dictionary, but also simply
+`model`, which will add the callback to `model.callbacks`. It's equivalent.
 Let us add two more meaningful callbacks
 
 ```@example callbacks

docs/src/speedytransforms.md

@@ -379,14 +379,16 @@ is obtained, and we do the following for ``f\zeta/g``.
 
 ```@example speedytransforms
 vor_grid = gridded(vor, Grid=spectral_grid.Grid)
-f = SpeedyWeather.coriolis(vor_grid)
-fζ_g = spectral_grid.Grid(vor_grid .* f ./ model.planet.gravity)
+f = coriolis(vor_grid)      # create Coriolis parameter f on same grid with default rotation
+g = model.planet.gravity
+fζ_g = zero(f)              # preallocate on same grid
+@. fζ_g = vor_grid * f / g  # in-place and element-wise
 nothing # hide
 ```
-The last line is a bit awkward for now, as the element-wise multiplication between
-two grids escapes the grid and returns a vector that we wrap again into a grid.
-We will fix that in future releases. Now we need to apply the inverse
-Laplace operator to ``f\zeta/g`` which we do as follows
+The last two lines are a bit awkward for now, as the element-wise multiplication between
+two grids would escapes the grid and returns a vector that we wrap again into a Grid.
+However, by preallocating the new variable with `zero` we guarantee to stay on that grid.
+Now we need to apply the inverse Laplace operator to ``f\zeta/g`` which we do as follows
 
 ```@example speedytransforms
 fζ_g_spectral = spectral(fζ_g, one_more_degree=true)

src/dynamics/coriolis.jl

@@ -20,18 +20,20 @@ function initialize!(coriolis::Coriolis, model::ModelSetup)
     return nothing
 end
 
+export coriolis
+
 """
 $(TYPEDSIGNATURES)
 Return the Coriolis parameter `f` on the grid `Grid` of resolution `nlat_half`
 on a planet of `ratation` [1/s]. Default rotation of Earth."""
 function coriolis(
     ::Type{Grid},
     nlat_half::Integer;
-    rotation=DEFAULT_ROTATION
+    rotation = DEFAULT_ROTATION
 ) where {Grid<:AbstractGrid}
 
-    f = zeros(Grid, nlat_half)
-    lat = get_lat(Grid, nlat_half)
+    f = zeros(Grid, nlat_half)      # preallocate
+    lat = get_lat(Grid, nlat_half)  # in radians [-π/2, π/2]
 
     for (j, ring) in enumerate(eachring(f))
         fⱼ = 2rotation*sin(lat[j])
@@ -48,7 +50,7 @@ Return the Coriolis parameter `f` on a grid like `grid`
 on a planet of `ratation` [1/s]. Default rotation of Earth."""
 function coriolis(
     grid::Grid;
-    rotation=DEFAULT_ROTATION
+    kwargs...
 ) where {Grid<:AbstractGrid}
-    return coriolis(Grid, grid.nlat_half; rotation)
+    return coriolis(Grid, grid.nlat_half; kwargs...)
 end

src/dynamics/scaling.jl

@@ -21,13 +21,28 @@ function scale!(
     end
 end
 
+"""
+$(TYPEDSIGNATURES)
+Scale the variable `var` inside `diagn` with scalar `scale`.
+"""
+function scale!(
+    diagn::DiagnosticVariables,
+    var::Symbol,
+    scale::Real,
+)
+    for layer in diagn.layers
+        variable = getfield(layer.grid_variables, var)
+        variable .*= scale
+    end
+end
+
 """
 $(TYPEDSIGNATURES)
 Scales the prognostic variables vorticity and divergence with
 the Earth's radius which is used in the dynamical core."""
 function scale!(progn::PrognosticVariables,
                 scale::Real)
-    new_scale = scale/progn.scale[]     # undo previous scale and new scale in one go
+    new_scale = scale/progn.scale[]     # undo previous scale and new scale in one go
     scale!(progn, :vor, new_scale)
     scale!(progn, :div, new_scale)
     progn.scale[] = scale               # store scaling information
@@ -43,6 +58,16 @@ function unscale!(progn::PrognosticVariables)
     progn.scale[] = 1                   # set scale back to 1=unscaled
 end
 
+"""
+$(TYPEDSIGNATURES)
+Undo the radius-scaling of vorticity and divergence from scale!(diagn, scale::Real)."""
+function unscale!(diagn::DiagnosticVariables)
+    inv_scale = inv(diagn.scale[])
+    scale!(diagn, :vor_grid, inv_scale)
+    scale!(diagn, :div_grid, inv_scale)
+    diagn.scale[] = 1                   # set scale back to 1=unscaled
+end
+
 """
 $(TYPEDSIGNATURES)
 Undo the radius-scaling for any variable. Method used for netcdf output."""

src/dynamics/time_integration.jl

@@ -52,13 +52,13 @@ function get_Δt_millisec(
     adjust_with_output::Bool,
     output_dt::Dates.TimePeriod = DEFAULT_OUTPUT_DT,
 )
-    # linearly scale Δt with trunc+1 (which are often powers of two)
+    # linearly scale Δt with trunc+1 (which are often powers of two)
     resolution_factor = (DEFAULT_TRUNC+1)/(trunc+1)
 
     # radius also affects grid spacing, scale proportionally
     radius_factor = radius/DEFAULT_RADIUS
 
-    # maybe rename to _at_trunc_and_radius?
+    # maybe rename to _at_trunc_and_radius?
     Δt_at_trunc = Second(Δt_at_T31).value * resolution_factor * radius_factor
 
     if adjust_with_output && (output_dt > Millisecond(0))
@@ -73,7 +73,7 @@ function get_Δt_millisec(
         Δt_millisec_unadjusted = round(Int, 1000*Δt_at_trunc)
         Δt_ratio = Δt_millisec.value/Δt_millisec_unadjusted
 
-        if abs(Δt_ratio - 1) > 0.05     # only when +-5% changes
+        if abs(Δt_ratio - 1) > 0.05     # only when +-5% changes
             p = round(Int, (Δt_ratio - 1)*100)
             ps = p > 0 ? "+" : ""
             @info "Time step changed from $Δt_millisec_unadjusted to $Δt_millisec ($ps$p%) to match output frequency."
@@ -114,7 +114,7 @@ function initialize!(L::Leapfrog, model::ModelSetup)
     n = round(Int, Millisecond(output_dt).value/L.Δt_millisec.value)
     nΔt = n*L.Δt_millisec
     if nΔt != output_dt
-        @info "$n steps of Δt = $(L.Δt_millisec.value)ms yield output every $(nΔt.value)ms (=$(nΔt.value/1000)s), but output_dt = $(output_dt.value)s"
+        @warn "$n steps of Δt = $(L.Δt_millisec.value)ms yield output every $(nΔt.value)ms (=$(nΔt.value/1000)s), but output_dt = $(output_dt.value)s"
     end
 end
 
@@ -133,7 +133,7 @@ function leapfrog!( A_old::LowerTriangularMatrix{Complex{NF}},      # prognostic
     @boundscheck lf == 1 || lf == 2 || throw(BoundsError())         # index lf picks leapfrog dim
 
     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
+    (; robert_filter, williams_filter) = L              # coefficients for the Robert and Williams filter
     dt_NF = convert(NF, dt)                             # time step dt in number format NF
 
     # LEAP FROG time step with or without Robert+Williams filter
@@ -183,7 +183,7 @@ function leapfrog!( progn::PrognosticSurfaceTimesteps,
     (; 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
+    spectral_truncation!(pres_tend)         # set lmax+1 mode to zero
     leapfrog!(pres_old, pres_new, pres_tend, dt, lf, model.time_stepping)
 end
 
@@ -214,13 +214,13 @@ function first_timesteps!(
     timestep!(clock, Δt_millisec_half, increase_counter=false)      
 
     # output, callbacks not called after the first Euler step as it's a half-step (i=0 to i=1/2)
-    # populating the second leapfrog index to perform the second time step
+    # populating the second leapfrog index to perform the second time step
 
     # SECOND TIME STEP (UNFILTERED LEAPFROG with dt=Δt, leapfrogging from t=0 over t=Δt/2 to t=Δt)
     initialize!(implicit, Δt, diagn, model)    # update precomputed implicit terms with time step Δt
     lf1 = 1                             # without Robert+Williams filter
     lf2 = 2                             # evaluate all tendencies at t=dt/2,
-                                        # the 2nd leapfrog index (=>Leapfrog)
+                                        # the 2nd leapfrog index (=>Leapfrog)
     timestep!(progn, diagn, Δt, model, lf1, lf2)
     # remove prev Δt/2 in case not even milliseconds, otherwise time is off by 1ms
     timestep!(clock, -Δt_millisec_half, increase_counter=false) 
@@ -290,11 +290,11 @@ function timestep!( progn::PrognosticVariables{NF}, # all prognostic variables
     (; pres) = progn.surface.timesteps[lf2]
     (; implicit, spectral_transform) = model
 
-    # GET TENDENCIES, CORRECT THEM FOR SEMI-IMPLICIT INTEGRATION
+    # GET TENDENCIES, CORRECT THEM FOR SEMI-IMPLICIT INTEGRATION
     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
+    # 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_lf, model)
@@ -397,7 +397,7 @@ function time_stepping!(
     # considered part of the model initialisation
     first_timesteps!(progn,diagn, model, output)
 
-    # only now initialise feedback for benchmark accuracy
+    # only now initialise feedback for benchmark accuracy
     initialize!(feedback, clock, model)
 
     # MAIN LOOP
@@ -413,6 +413,7 @@ function time_stepping!(
     # UNSCALE, CLOSE, FINISH
     finish!(feedback)                       # finish the progress meter, do first for benchmark accuracy
     unscale!(progn)                         # undo radius-scaling for vor, div from the dynamical core
+    unscale!(diagn)                         # undo radius-scaling for vor, div from the dynamical core
     close(output)                           # close netCDF file
     write_restart_file(progn, output)       # as JLD2 
     finish!(model.callbacks, progn, diagn, model)

src/output/callbacks.jl

@@ -30,7 +30,7 @@ export NoCallback
 struct NoCallback <: AbstractCallback end
 initialize!(::NoCallback, args...) = nothing     # executed once before the main time loop
 callback!(::NoCallback, args...) = nothing       # executed after every time step
-finish!(::NoCallback, args...) = nothing         # executed after main time loop finishes
+finish!(::NoCallback, args...) = nothing         # executed after main time loop finishes
 
 # simply loop over dict of callbacks
 for func in (:initialize!, :callback!, :finish!)
@@ -59,12 +59,24 @@ function add!(D::CALLBACK_DICT, key_callbacks::Pair{Symbol, <:AbstractCallback}.
     end
 end
 
+"""
+$(TYPEDSIGNATURES)
+Add a or several callbacks to a model::ModelSetup. To be used like
+
+    add!(model, :my_callback => callback)
+    add!(model, :my_callback1 => callback, :my_callback2 => other_callback)
+"""
+add!(model::ModelSetup, key_callbacks::Pair{Symbol, <:AbstractCallback}...) =
+    add!(model.callbacks, key_callbacks...)
 add!(D::CALLBACK_DICT, key::Symbol, callback::AbstractCallback) = add!(D, Pair(key, callback))
+add!(model::ModelSetup, key::Symbol, callback::AbstractCallback) =
+    add!(model.callbacks, Pair(key, callback))
+
 
-# also with string but flag conversion
+# also with string but flag conversion
 function add!(D::CALLBACK_DICT, key::String, callback::AbstractCallback)
     key_symbol = Symbol(key)
-    @info "Callback keys are Symbols. String \"$key\" converted to Symbol :$key_symbol."
+    @warn "Callback keys are Symbols. String \"$key\" converted to Symbol :$key_symbol."
     add!(D, key_symbol, callback)
 end
 
@@ -75,14 +87,24 @@ key which is randomly created like callback_????. To be used like
 
     add!(model.callbacks, callback)
     add!(model.callbacks, callback1, callback2)."""
-function add!(D::CALLBACK_DICT, callbacks::AbstractCallback...)
+function add!(D::CALLBACK_DICT, callbacks::AbstractCallback...; verbose=true)
     for callback in callbacks
         key = Symbol("callback_"*Random.randstring(4))
-        @info "$(typeof(callback)) callback added with key $key"
+        verbose && @info "$(typeof(callback)) callback added with key $key"
         add!(D, key => callback)
     end
 end
 
+"""
+$(TYPEDSIGNATURES)
+Add a or several callbacks to a mdoel without specifying the
+key which is randomly created like callback_????. To be used like
+
+    add!(model.callbacks, callback)
+    add!(model.callbacks, callback1, callback2)."""
+add!(model::ModelSetup, callbacks::AbstractCallback...) =
+    add!(model.callbacks, callbacks..., verbose = model.feedback.verbose)
+
 # delete!(dict, key) already defined in Base
 
 export GlobalSurfaceTemperatureCallback
@@ -109,8 +131,8 @@ function initialize!(
     model::ModelSetup,
 ) where NF
     callback.temp = Vector{NF}(undef, progn.clock.n_timesteps+1) # replace with vector of correct length
-    callback.temp[1] = diagn.layers[diagn.nlev].temp_average[]   # set initial conditions
-    callback.timestep_counter = 1                                # (re)set counter to 1
+    callback.temp[1] = diagn.layers[diagn.nlev].temp_average[]   # set initial conditions
+    callback.timestep_counter = 1                                # (re)set counter to 1
 end
 
 """

test/callbacks.jl

@@ -38,13 +38,13 @@ end
         # allocate recorder: number of time steps (incl initial conditions) in simulation  
         callback.maximum_surface_wind_speed = zeros(progn.clock.n_timesteps + 1)
 
-        # where surface (=lowermost model layer) u, v on the grid are stored
+        # where surface (=lowermost model layer) u, v on the grid are stored
         (; u_grid, v_grid) = diagn.layers[diagn.nlev].grid_variables
 
         # maximum wind speed of initial conditions
         callback.maximum_surface_wind_speed[1] = max_2norm(u_grid, v_grid)
 
-        # (re)set counter to 1
+        # (re)set counter to 1
         callback.timestep_counter = 1
     end
 
@@ -85,6 +85,7 @@ end
     key = :storm_chaser
     add!(model.callbacks, key => storm_chaser)      # with :storm_chaser key
     add!(model.callbacks, NoCallback())             # add dummy too 
+    add!(model, NoCallback())                       # add dummy with ::ModelSetup interface
 
     simulation = initialize!(model)
     run!(simulation, period=Day(1))