Adding computed fields to a model

[ammtD]

Following up on a thread on the #oceananigans channel on the Julia Slack:

I am trying to add computed fields to a model output; fields derived from simulated tracers and velocities, e.g. N2, Ri
This is related to the closed issue #2337 I believe, but the solution there doesn't yet seem to work for me.
Here's my code, which is roughly based on the Langmuir Turbulence example.

using Pkg
using Oceananigans
using Oceananigans.Units: seconds, minute, minutes, hour, hours
using Oceananigans.BuoyancyModels: g_Earth
using Printf
using Plots

using DelimitedFiles
using Polynomials
using Dierckx

grid = RegularRectilinearGrid(size=(30, 32, 50), x=(0.0, 60.0), y=(0.0, 64.0), z=(-50.0, 0.0),halo=(3,3,3)) 
    #extent=(128, 128, 64),

coriolis = FPlane(f=1e-4) # s⁻
buoyancy = SeawaterBuoyancy(equation_of_state=LinearEquationOfState(α=2e-4, β=8e-4))

# ---- BC's   -------------------------------
# Constant heat flux (warming)
Qh = 400  # W m⁻², surface _heat_ flux
ρ0 = 1026
cp =3991
QT = Qh / (ρ0 * cp)   # K m s⁻¹, surface _temperature_ flux

dTdz = 0.01  # K/m

T_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(QT), bottom = GradientBoundaryCondition(dTdz))

@inline Qs(x, y, t, S, evaporation_rate) = - evaporation_rate * S # [salinity unit] m s⁻¹
evaporation_rate = 1e-3 / hour # m s⁻¹

S_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qs, field_dependencies=:S, parameters=evaporation_rate))

 #  constant wind forcing at surface: 
wspd= 5;
u10= wspd;  # m s⁻¹, average wind velocity 10 meters above the ocean
cd = 2.5e-3  # dimensionless drag coefficient
ρa = 1.225   # kg m⁻³, average density of air at sea-level
Qu = - ρa / ρ0 * cd * u10 * abs(u10) # water side friction velocity, m² s⁻²

u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qu))

#---- Stokes forcing-----------------------------------------------------------------
amplitude = 0.8 # m
wavelength = 60 # m
wavenumber = 2π / wavelength # m⁻¹
frequency = sqrt(g_Earth * wavenumber) # s⁻¹

const vertical_scale = wavelength / 4π
const Us = amplitude^2 * wavenumber * frequency # m s⁻¹
us(z)= Us * exp(z / vertical_scale)

@inline dus_dz(z,t) = 1 / vertical_scale * Us * exp(z/vertical_scale)


#----  Model definition --------LES type with rotation affecting resolved velocity field... not quite what we want but struggling
# to get background fields to function ---------

model = NonhydrostaticModel(
           architecture = CPU(),
              advection = WENO5(),
            timestepper = :RungeKutta3,
                   grid = grid,
                tracers = (:T, :S), 
               buoyancy = buoyancy,
               coriolis = coriolis, 
                closure = AnisotropicMinimumDissipation(),
           stokes_drift = UniformStokesDrift(∂z_uˢ=dus_dz),
    boundary_conditions = (u=u_bcs, T=T_bcs, S=S_bcs))

    
# -------Initializations---------------------------------------------------------------------

R(z) = randn() * exp(z/4)
usr = sqrt(abs(Qu))
ui(x,y,z) = us(z) + usr* 1e-1 * R(z)
wi(x,y,z) = usr * 1e-1*R(z)


Nz= grid.Nz;  Nx = grid.Nx;  Ny = grid.Ny
zc = znodes(Center,grid)
yc = ynodes(Center,grid)
xc= xnodes(Center,grid);


tsdata = readdlm("CTD_dyerelease.txt", Float64) 
zdata= -1*reverse(tsdata[:,1], dims=1); Tdata= reverse(tsdata[:,2], dims=1); 
Sdata = reverse(tsdata[:,3], dims=1);
# interpolate onto zc
x=zdata; y1=Tdata; ; y2=Sdata; xq= zc; 
Tj = Spline1D(x, y1; k=1)(xq);
Sj = Spline1D(x, y2; k=1)(xq);

Ti=  repeat(reshape(Tj,1,1,Nz), Nx, Ny, 1);
Si = repeat(reshape(Sj,1,1,Nz), Nx,Ny, 1); 

set!(model, u=ui, T=Ti, S=Si, w=wi)  

# --- plot initialization ----------------------------
u1= model.velocities.u[1,1,:] ;  w1=model.velocities.w[1,1,:]
T1=model.tracers.T[1,1,:];  S1=model.tracers.S[1,1,:];
ind= 1:Nz;

uplot= plot([u1[ind], w1[ind]],zc,label = ["uᵢ" "wᵢ"], xlabel="m/s",ylabel = "zC [m]",linewidth = 2,legend=:bottomright)     # , v1[ind]
Tplot= plot(T1[ind],zc,label = ["Tᵢ"], xlabel="deg C",ylabel = "zC [m]",linewidth = 2,legend=:bottomright)
Splot = plot(S1[ind],zc,label= "Sᵢ", xlabel = "g/kg", ylabel = "zC [m]",linewidth=2)
plot(Tplot, Splot, uplot)

# ---------------------------------------------------
simulation = Simulation(model, Δt=30.0, stop_time=2hours)
wizard = TimeStepWizard(cfl=0.7, max_change=1.1, max_Δt=40.0seconds)
simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))

# ---- Computed Fields, for e.g. N2 and Ri------------------# following issue 2337...where does this go?? 

N2= CenterField(grid)

α=2e-4; β=8e-4
g=g_Earth

T= model.tracers.T;    S = model.tracers.S; 

N2_op = @at(Center, Center, Center) g * α * ∂z(T) - g* β* ∂z(S)

function compute_N2!(sim)
    N2 .= N2_op
    
    fill_halo_regions!(N2, sim.model.architecture)
    return nothing
end
simulation.callbacks[:compute_N2]= Callback(compute_N2!)

# -----------------------------------------------------------------------------------

outputs = (u=model.velocities.u, v=model.velocities.v, w=model.velocities.w, T=model.tracers.T, S=model.tracers.S, 
      νₑ = model.diffusivity_fields.νₑ, kT=model.diffusivity_fields.κₑ.T)

output_interval = 1minutes

using NetCDF

simulation.output_writers[:field_output] = 
    NetCDFOutputWriter(model, outputs,
                     schedule = TimeInterval(output_interval),
                     filepath = joinpath(@__DIR__,"junk_LT_TSUini_2h.nc"),
                     mode = "c")

run!(simulation)

Other than the ----Computed field --- section this code runs.

[tomchor]

I think you just need to pass N2_op to the outputs list.

I haven't tested it, but your code would look something like this:

α=2e-4; β=8e-4
g=g_Earth

T= model.tracers.T;    S = model.tracers.S; 

N2 = @at(Center, Center, Center) g * α * ∂z(T) - g* β* ∂z(S)

outputs = (u=model.velocities.u, v=model.velocities.v, w=model.velocities.w, T=model.tracers.T, S=model.tracers.S, 
      νₑ = model.diffusivity_fields.νₑ, kT=model.diffusivity_fields.κₑ.T, N2=N2,)

output_interval = 1minutes

simulation.output_writers[:field_output] = 
    NetCDFOutputWriter(model, outputs,
                     schedule = TimeInterval(output_interval),
                     filepath = joinpath(@__DIR__,"junk_LT_TSUini_2h.nc"),
                     mode = "c")

run!(simulation)

Can you test this and see if this works?

Also I don't think you need the using NetCDF call.

Also(2) in the latest version of Oceananigans you'll have to slighly change your call to NetCDFOutputWriter:

• mode = "c" becomes overwrite_existing=true
• filepath keywords becomes filename

[ammtD]

In your suggestion here you're omitting the callback, is that right?

you're right, the using NetCDF call seems unnecessary. I have been holding off with updating to the latest stable version, because so many things have changed, and I'm on a tight schedule to produce something with the model. I know it's a bit risky and not sustainable for much longer....

[tomchor]

In your suggestion here you're omitting the callback, is that right?

Yes.

I definitely suggest updating it though. It'll be much worse now since you waited for so long and changes have accumulated, but it's doable and we can help you out either here or on slack.

[ammtD]

I hadn't gotten to the outputs yet, since this line doesn't work for me:
N2 = @at(Center, Center, Center) g * α * ∂z(T) - g* β* ∂z(S)

It gives error:

syntax: extra token "g" after end of expression

Stacktrace:
 [1] top-level scope
   @ In[3]:10
 [2] eval
   @ .\boot.jl:360 [inlined]
 [3] include_string(mapexpr::typeof(REPL.softscope), mod::Module, code::String, filename::String)
   @ Base .\loading.jl:1116

Could it be something that wasn't working yet in v. 64.1?   

[ammtD]

...because even with the N2=@at( .... line now working, the simulation still crashes.

[iuryt]

Which version are you using?

using Pkg; Pkg.status("Oceananigans")

[ammtD]

N2_op = @at((Center, Center, Center), g * α * ∂z(T) - g* β* ∂z(S))

function compute_N2!(sim)
    cf_N2 .= N2_op
    
   fill_halo_regions!(N2, sim.model.architecture)
    return nothing
end
simulation.callbacks[:compute_N2]= Callback(compute_N2!)

run!(simulation) crashes on initialization, because it doesn't know fill_halo_regions!
If I comment out that line, and re-run the code, the simulation initialization starts ... and completes successfully,

So what is fill_halo_regions! ? Should I call this as well when initializing the T/S tracerfields and/or background fields (...and if so, is this another good reason to update to the most recent version?) Getting bg-fields to work is another issue I am working on right now...

[ammtD]

Which version are you using?

using Pkg; Pkg.status("Oceananigans")

I am on v0.64.1 still.... I know that's getting a bit old.

[iuryt]

That's pretty old and things are moving fast here. I suggest you update Oceananigans.

using Pkg; Pkg.update("Oceananigans")