Simulation callback: Error encountered when updating u_s with time

[wpan081]

I want to replicate the swell example with variations in the article by Wagner et al. in 2021, and output the Stokes velocity u_s and the Stokes shear. Both of these two quantities are time-varying. I used the callback method to update u_s in real time, but the simulation always reports an error when I add the callback. In fact, this example can be run on the CPU, but it reports an error when I run it on the GPU (here I have already added GPU() to switch to the GPU framework). Hope someone can help me. Thanks a lot.

using Random
using Printf
using Plots
using JLD2
using Oceananigans
using Oceananigans.Units: minute, minutes, hour, hours
using Oceanostics.FlowDiagnostics: RichardsonNumber
using Oceanostics: TurbulentKineticEnergy, ZShearProductionRate, IsotropicKineticEnergyDissipationRate, KineticEnergyDissipationRate

grid = RectilinearGrid(GPU();
                       size = (32, 32, 20),
                       halo = (3, 3, 3),
                          x = (0, 256),
                          y = (0, 256),
                          z = (-80, 0))

buoyancy = SeawaterBuoyancy(equation_of_state=LinearEquationOfState(thermal_expansion=2e-4, haline_contraction=8e-4))

#Boundary conditions

const k = 2 * 3.14159 / 100
const U = 1.0^2 * k * sqrt(9.8 * k)

@inline ∂z_uˢ(x, y, z, t) = 2k * U * exp(2k * z) * (1 - exp(-t^2 / (2 * (14400)^2)))
@inline ∂t_uˢ(x, y, z, t) =      U * exp(2k * z) * exp(-t^2 / (2 * (14400)^2)) * t / (14400)^2


@inline surface_temperature_flux(x,y,t,p) = - (p.Qʰ) / (p.ρₒ * p.cᴾ)
Qʰ = -50 # W m⁻², surface heat flux
ρₒ = 1026 # kg m⁻³, average density at the surface of the world ocean
cᴾ = 3991 # J K⁻¹ kg⁻¹, typical heat capacity for seawater

dTdz = 0.01 # K m⁻¹

T_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(surface_temperature_flux, parameters=(Qʰ=Qʰ, ρₒ=ρₒ, cᴾ=cᴾ)),bottom = GradientBoundaryCondition(dTdz))
Qᵘ = 0 # m² s⁻²
 
u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵘ))

@inline Qˢ(x, y, t, S, evaporation_rate) = - evaporation_rate * S # [salinity unit] m s⁻¹
nothing # hide

evaporation_rate = 0.0 # m s⁻¹
evaporation_bc = FluxBoundaryCondition(Qˢ, field_dependencies=:S, parameters=evaporation_rate)
S_bcs = FieldBoundaryConditions(top=evaporation_bc)

#initial condition
S₀ = 35.0
T₀ = 20
T0(z) = T₀ + dTdz * z

#Sponge layer

gaussian_mask = GaussianMask{:z}(center=-grid.Lz, width=grid.Lz/10)

u_sponge = v_sponge = w_sponge = Relaxation(rate=1/hour, mask=gaussian_mask)

T_sponge = Relaxation(rate = 1/hour,
                      target = LinearTarget{:z}(intercept=T₀, gradient=dTdz),
                      mask = gaussian_mask)
S_sponge = Relaxation(rate = 1/hour,
                      target = S₀,
                      mask = gaussian_mask)

#Model instantiation
model = NonhydrostaticModel(; grid,
                advection = WENOFifthOrder(),
              timestepper = :RungeKutta3,
                  tracers = (:T, :S),
                 coriolis = FPlane(f=0),
                 buoyancy = buoyancy,
                  closure = AnisotropicMinimumDissipation(),
                  stokes_drift = UniformStokesDrift(∂z_uˢ=∂z_uˢ,∂t_uˢ=∂t_uˢ),
      boundary_conditions = (u=u_bcs, T=T_bcs, S=S_bcs),
                  forcing = (u=u_sponge, v=v_sponge, w=w_sponge, T=T_sponge, S=S_sponge))
                  

#Notes:
#* To change the `architecture` to `GPU`, replace `architecture = CPU()` with
#`architecture = GPU()`.

#Initial conditions
Ξ(z) = randn() * z / model.grid.Lz * (1 + z / model.grid.Lz) # noise
Tᵢ(x, y, z) = T0(z) + dTdz * model.grid.Lz * 1e-6 * Ξ(z)
uᵢ(x, y, z) = sqrt(abs(Qᵘ)) * 1e-3 * Ξ(z)
set!(model, u=uᵢ, w=uᵢ, T=Tᵢ, S=S₀)

#Setting up a simulation
simulation = Simulation(model, Δt=10.0, stop_time=30hours)
wizard = TimeStepWizard(cfl=1.0, max_change=1.1, max_Δt=1minute)
simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))

#Nice progress messaging is helpful:
progress_message(sim) = @printf("Iteration: %04d, time: %s, Δt: %s, max(|w|) = %.1e ms⁻¹, wall time: %s\n",
                                iteration(sim),
                                prettytime(sim),
                                prettytime(sim.Δt),
                                maximum(abs, sim.model.velocities.w),
                                prettytime(sim.run_wall_time))

simulation.callbacks[:progress] = Callback(progress_message, IterationInterval(100))

#We then set up the simulation:

#Define a function for the Stokes drift that includes the time variable
us = CenterField(grid)
usz = CenterField(grid)
function update_us!(sim)
    t = sim.model.clock.time
    z = znodes(sim.model.grid, Center()) 
    zC = reshape(z, (1, 1, length(z)))
    us_z = U * exp(2k * z) * (1 - exp(-t^2 / (2 * (14400)^2)))
    us_zz =2k * U * exp(2k * z) * (1 - exp(-t^2 / (2 * (14400)^2)))
    us .= us_z 
    usz .=us_zz
    return nothing
end

simulation.callbacks[:update_us] = Callback(update_us!, IterationInterval(10))
u_eulerian = @at (Center, Center, Center) model.velocities.u - us

us_average = Average(us, dims=(1, 2))
usz_average = Average(usz,dims=(1, 2))
u_lagrangian = @at (Center, Center, Center) model.velocities.u
w = @at (Center, Center, Center) model.velocities.w
U_lagrangian = Field(Average(u_lagrangian, dims=(1, 2)))
W = Field(Average(w, dims=(1, 2)))

#Calculate stokes shear production
up =   Field(u_lagrangian - U_lagrangian)
wp =   Field(w - W)
wu_p = Average(up*wp, dims=(1, 2))
wpup = Field(Average(wp*up, dims = (1, 2)))
usz = Field(Average(usz, dims = (1, 2)))
stokes_shear = Average(- wpup * usz, dims = (1, 2))

simulation.output_writers[:averages] =
     NetCDFOutputWriter(model, (; us_average,usz_average,stokes_shear),
                         schedule = AveragedTimeInterval(3minute),
                         filename = "test.nc",
             overwrite_existing = true) 

run!(simulation)

[glwagner]

@wpan081 your Stokes shear are functions:

@inline ∂z_uˢ(x, y, z, t) = 2k * U * exp(2k * z) * (1 - exp(-t^2 / (2 * (14400)^2)))
@inline ∂t_uˢ(x, y, z, t) =      U * exp(2k * z) * exp(-t^2 / (2 * (14400)^2)) * t / (14400)^2

which cannot be updated. You can only update arrays.

For the use case you describe, a function will work. The function doesn't need to be "updated", because it depends on time and is evaluated on the fly, each time step. Your functions, however, do not use the right function signature. Here's the docstring:

help?> UniformStokesDrift
search: UniformStokesDrift StokesDrift

  UniformStokesDrift(; ∂z_uˢ=zerofunction, ∂z_vˢ=zerofunction, ∂t_uˢ=zerofunction, ∂t_vˢ=zerofunction, parameters=nothing)

  Construct a set of functions for a Stokes drift velocity field corresponding to a horizontally-uniform surface gravity wave field, with optional
  parameters.

  If parameters=nothing, then the functions ∂z_uˢ, ∂z_vˢ, ∂t_uˢ, ∂t_vˢ must be callable with signature (z, t). If !isnothing(parameters), then
  functions must be callable with the signature (z, t, parameters).

  To resolve the evolution of the Lagrangian-mean momentum, we require vertical-derivatives and time-derivatives of the horizontal components of the
  Stokes drift, uˢ and vˢ.

  Examples
  ≡≡≡≡≡≡≡≡

  Exponentially decaying Stokes drift corresponding to a surface Stokes drift of uˢ(z=0) = 0.005 and decay scale h = 20:

  using Oceananigans

  @inline uniform_stokes_shear(z, t) = 0.005 * exp(z / 20)

  stokes_drift = UniformStokesDrift(∂z_uˢ=uniform_stokes_shear)

  # output

  UniformStokesDrift{Nothing}:
  ├── ∂z_uˢ: uniform_stokes_shear
  ├── ∂z_vˢ: zerofunction
  ├── ∂t_uˢ: zerofunction
  └── ∂t_vˢ: zerofunction

  Exponentially-decaying Stokes drift corresponding to a surface Stokes drift of uˢ = 0.005 and decay scale h = 20, using parameters:

  using Oceananigans

  @inline uniform_stokes_shear(z, t, p) = p.uˢ * exp(z / p.h)

  stokes_drift_parameters = (uˢ = 0.005, h = 20)
  stokes_drift = UniformStokesDrift(∂z_uˢ=uniform_stokes_shear, parameters=stokes_drift_parameters)

  # output

  UniformStokesDrift with parameters (uˢ=0.005, h=20):
  ├── ∂z_uˢ: uniform_stokes_shear
  ├── ∂z_vˢ: zerofunction
  ├── ∂t_uˢ: zerofunction
  └── ∂t_vˢ: zerofunction

[wpan081]

Thanks for your reply! Maybe I'm not skilled enough, so I can't understand how this can help me output the changing u_s. Maybe my problem description wasn't clear enough. I want to output u_s and its derivative in the z - direction to calculate the Stokes shear production in the TKE budget. So, I used the following code to update the fields of u_s and u_sz:

us = CenterField(grid)
usz = CenterField(grid)

function update_us!(sim)
    t = sim.model.clock.time
    z = znodes(sim.model.grid, Center()) 
    zC = reshape(z, (1, 1, length(z)))
    us_z = U * exp(2k * z) * (1 - exp(-t^2 / (2 * (14400)^2)))
    us_zz = 2k * U * exp(2k * z) * (1 - exp(-t^2 / (2 * (14400)^2)))
    us .= us_z 
    usz .= us_zz
    return nothing
end

simulation.callbacks[:update_us] = Callback(update_us!, IterationInterval(10))
u_eulerian = @at (Center, Center, Center) model.velocities.u - us
us_average = Average(us, dims=(1, 2))
usz_average = Average(usz, dims=(1, 2))

In this way, I can obtain a time varying u_s that can be output and used to calculate the Stokes shear production? In CPU, this code works well, but GPU not.

[glwagner]

You may want to use FunctionField for that (eg form a FunctionField from your Stokes shear). Note that it's technically more correct to compute the components of shear production at the u, v locations respectively, and then reconstruct afterwards to obtain the total contribution to the TKE budget. https://github.com/tomchor/Oceanostics.jl has some utilities for computing shear production that might be useful to take a look at.

Shear production is problematic as a diagnostic more generally, but you may know this...

[wpan081]

Thanks a lot! I'll try it.