From 13cd63209aba02a211a6a35d14697f74b93db213 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Mon, 24 Jun 2024 05:08:43 -0400 Subject: [PATCH 001/114] working version on CPU --- NN_2D_channel.jl | 310 ++++++++++++++++++++++++++ NN_closure.jl | 170 ++++++++++++++ feature_scaling.jl | 84 +++++++ xin_kai_vertical_diffusivity_local.jl | 230 +++++++++++++++++++ 4 files changed, 794 insertions(+) create mode 100644 NN_2D_channel.jl create mode 100644 NN_closure.jl create mode 100644 feature_scaling.jl create mode 100644 xin_kai_vertical_diffusivity_local.jl diff --git a/NN_2D_channel.jl b/NN_2D_channel.jl new file mode 100644 index 0000000000..56e2204153 --- /dev/null +++ b/NN_2D_channel.jl @@ -0,0 +1,310 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("NN_closure.jl") +include("xin_kai_vertical_diffusivity_local.jl") +include("feature_scaling.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 + +const Ly = 2000kilometers # meridional domain length [m] + +# Architecture +model_architecture = CPU() + +# number of grid points +Ny = 192 +Nz = 128 + +const Lz = 1024 + +grid = RectilinearGrid(model_architecture, + topology = (Flat, Bounded, Bounded), + size = (Ny, Nz), + halo = (3, 3), + y = (0, Ly), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +T_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(1e-4), bottom=FluxBoundaryCondition(0.0)) + +##### +##### Coriolis +##### + +const f₀ = 8e-5 +const β = 1e-11 +coriolis = BetaPlane(f₀=f₀, β = β) + +##### +##### Forcing and initial condition +##### +const dTdz = 0.014 +const dSdz = 0.0021 + +const T_surface = 20.0 +const S_surface = 36.6 + +T_initial(y, z) = dTdz * z + T_surface +S_initial(y, z) = dSdz * z + S_surface + +# closure +κh = 0.5e-5 # [m²/s] horizontal diffusivity +νh = 30.0 # [m²/s] horizontal viscocity +κz = 0.5e-5 # [m²/s] vertical diffusivity +νz = 3e-4 # [m²/s] vertical viscocity + +horizontal_closure = HorizontalScalarDiffusivity(ν = νh, κ = κh) +vertical_closure = VerticalScalarDiffusivity(ν = νz, κ = κz) + +convective_adjustment = ConvectiveAdjustmentVerticalDiffusivity(convective_κz = 1.0, + convective_νz = 0.0) + +nn_closure = NNFluxClosure(model_architecture) +base_closure = XinKaiLocalVerticalDiffusivity() + +##### +##### Model building +##### + +@info "Building a model..." + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = ImplicitFreeSurface(), + momentum_advection = WENO(grid = grid), + tracer_advection = WENO(grid = grid), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = (nn_closure, base_closure), + # closure = (horizontal_closure, vertical_closure, convective_adjustment), + tracers = (:T, :S), + boundary_conditions = (; T = T_bcs), + # forcing = (; b = Fb) +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(y, z) = rand() * exp(z / 8) + +T_initial_noisy(y, z) = T_initial(y, z) + 1e-6 * noise(y, z) +S_initial_noisy(y, z) = S_initial(y, z) + 1e-6 * noise(y, z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 1days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.1, max_change=1.1, max_Δt=20minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(20)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): (%6.3e, %6.3e, %6.3e) m/s, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.velocities.w), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(20)) + + +##### +##### Diagnostics +##### + +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S + +ζ = Field(∂x(v) - ∂y(u)) + +Tbar = Field(Average(T, dims = 1)) +Sbar = Field(Average(S, dims = 1)) +V = Field(Average(v, dims = 1)) +W = Field(Average(w, dims = 1)) + +T′ = T - Tbar +S′ = S - Sbar +v′ = v - V +w′ = w - W + +v′T′ = Field(Average(v′ * T′, dims = 1)) +w′T′ = Field(Average(w′ * T′, dims = 1)) +v′S′ = Field(Average(v′ * S′, dims = 1)) +w′S′ = Field(Average(w′ * S′, dims = 1)) + +outputs = (; T, S, ζ, w) + +averaged_outputs = (; v′T′, w′T′, v′S′, w′S′, Tbar, Sbar) + +##### +##### Build checkpointer and output writer +##### + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(100days), + prefix = "NN_channel", + overwrite_existing = true) + +simulation.output_writers[:fields] = JLD2OutputWriter(model, outputs, + schedule = TimeInterval(5days), + filename = "NN_channel", + # field_slicer = nothing, + verbose = true, + overwrite_existing = true) + +simulation.output_writers[:averages] = JLD2OutputWriter(model, averaged_outputs, + schedule = AveragedTimeInterval(1days, window = 1days, stride = 1), + filename = "NN_channel_averages", + verbose = true, + overwrite_existing = true) + +@info "Running the simulation..." + +try + run!(simulation, pickup = false) +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#= +# ##### +# ##### Visualization +# ##### + +using Plots + +grid = RectilinearGrid(architecture = CPU(), + topology = (Periodic, Bounded, Bounded), + size = (grid.Nx, grid.Ny, grid.Nz), + halo = (3, 3, 3), + x = (0, grid.Lx), + y = (0, grid.Ly), + z = z_faces) + +xζ, yζ, zζ = nodes((Face, Face, Center), grid) +xc, yc, zc = nodes((Center, Center, Center), grid) +xw, yw, zw = nodes((Center, Center, Face), grid) + +j′ = round(Int, grid.Ny / 2) +y′ = yζ[j′] + +b_timeseries = FieldTimeSeries("abernathey_channel.jld2", "b", grid = grid) +ζ_timeseries = FieldTimeSeries("abernathey_channel.jld2", "ζ", grid = grid) +w_timeseries = FieldTimeSeries("abernathey_channel.jld2", "w", grid = grid) + +@show b_timeseries + +anim = @animate for i in 1:length(b_timeseries.times) + b = b_timeseries[i] + ζ = ζ_timeseries[i] + w = w_timeseries[i] + + b′ = interior(b) .- mean(b) + b_xy = b′[:, :, grid.Nz] + ζ_xy = interior(ζ)[:, :, grid.Nz] + ζ_xz = interior(ζ)[:, j′, :] + w_xz = interior(w)[:, j′, :] + + @show bmax = max(1e-9, maximum(abs, b_xy)) + @show ζmax = max(1e-9, maximum(abs, ζ_xy)) + @show wmax = max(1e-9, maximum(abs, w_xz)) + + blims = (-bmax, bmax) .* 0.8 + ζlims = (-ζmax, ζmax) .* 0.8 + wlims = (-wmax, wmax) .* 0.8 + + blevels = vcat([-bmax], range(blims[1], blims[2], length = 31), [bmax]) + ζlevels = vcat([-ζmax], range(ζlims[1], ζlims[2], length = 31), [ζmax]) + wlevels = vcat([-wmax], range(wlims[1], wlims[2], length = 31), [wmax]) + + xlims = (-grid.Lx / 2, grid.Lx / 2) .* 1e-3 + ylims = (0, grid.Ly) .* 1e-3 + zlims = (-grid.Lz, 0) + + w_xz_plot = contourf(xw * 1e-3, zw, w_xz', + xlabel = "x (km)", + ylabel = "z (m)", + aspectratio = 0.05, + linewidth = 0, + levels = wlevels, + clims = wlims, + xlims = xlims, + ylims = zlims, + color = :balance) + + ζ_xy_plot = contourf(xζ * 1e-3, yζ * 1e-3, ζ_xy', + xlabel = "x (km)", + ylabel = "y (km)", + aspectratio = :equal, + linewidth = 0, + levels = ζlevels, + clims = ζlims, + xlims = xlims, + ylims = ylims, + color = :balance) + + b_xy_plot = contourf(xc * 1e-3, yc * 1e-3, b_xy', + xlabel = "x (km)", + ylabel = "y (km)", + aspectratio = :equal, + linewidth = 0, + levels = blevels, + clims = blims, + xlims = xlims, + ylims = ylims, + color = :balance) + + w_xz_title = @sprintf("w(x, z) at t = %s", prettytime(ζ_timeseries.times[i])) + ζ_xz_title = @sprintf("ζ(x, z) at t = %s", prettytime(ζ_timeseries.times[i])) + ζ_xy_title = "ζ(x, y)" + b_xy_title = "b(x, y)" + + layout = @layout [upper_slice_plot{0.2h} + Plots.grid(1, 2)] + + plot(w_xz_plot, ζ_xy_plot, b_xy_plot, layout = layout, size = (1200, 1200), title = [w_xz_title ζ_xy_title b_xy_title]) +end + +mp4(anim, "abernathey_channel.mp4", fps = 8) #hide +=# \ No newline at end of file diff --git a/NN_closure.jl b/NN_closure.jl new file mode 100644 index 0000000000..8475b3ea5c --- /dev/null +++ b/NN_closure.jl @@ -0,0 +1,170 @@ +import Oceananigans.TurbulenceClosures: + compute_diffusivities!, + DiffusivityFields, + viscosity, + diffusivity, + diffusive_flux_x, + diffusive_flux_y, + diffusive_flux_z, + viscous_flux_ux, + viscous_flux_uy, + viscous_flux_uz, + viscous_flux_vx, + viscous_flux_vy, + viscous_flux_vz, + viscous_flux_wx, + viscous_flux_wy, + viscous_flux_wz + +using Oceananigans.BuoyancyModels: ∂x_b, ∂y_b, ∂z_b +using Oceananigans.Coriolis +using Oceananigans.Grids: φnode +using Lux, LuxCUDA +using JLD2 +using ComponentArrays + +using KernelAbstractions: @index, @kernel, @private + +import Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure, ExplicitTimeDiscretization + +using Adapt + +include("./feature_scaling.jl") + +@inline hack_sind(φ) = sin(φ * π / 180) + +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::HydrostaticSphericalCoriolis) = 2 * coriolis.rotation_rate * hack_sind(φnode(i, j, k, grid, Center(), Center(), Center())) +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::FPlane) = coriolis.f +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::BetaPlane) = coriolis.f₀ + coriolis.β * ynode(i, j, k, grid, Center(), Center(), Center()) + +struct NN{M, P, S} + model :: M + ps :: P + st :: S +end + +# @inline (neural_network::NN)(∂Tᵢ₋₁, ∂Tᵢ, ∂Tᵢ₊₁, ∂Sᵢ₋₁, ∂Sᵢ, ∂Sᵢ₊₁, ∂σᵢ₋₁, ∂σᵢ, ∂σᵢ₊₁, Jᵇ, fᶜᶜ, dev) = first(LuxCore.apply(neural_network.model, dev([Jᵇ, fᶜᶜ, ∂Tᵢ, ∂Tᵢ₋₁, ∂Tᵢ₊₁, ∂Sᵢ, ∂Sᵢ₋₁, ∂Sᵢ₊₁, ∂σᵢ, ∂σᵢ₋₁, ∂σᵢ₊₁]), neural_network.ps, neural_network.st)) +@inline (neural_network::NN)(input) = first(first(neural_network.model(input, neural_network.ps, neural_network.st))) + +struct NNFluxClosure{A <: NN, S} <: AbstractTurbulenceClosure{ExplicitTimeDiscretization, 3} + wT :: A + wS :: A + scaling :: S + # dev :: D +end + +Adapt.adapt_structure(to, nn :: NNFluxClosure) = + NNFluxClosure(Adapt.adapt(to, nn.wT), + Adapt.adapt(to, nn.wS), + Adapt.adapt(to, nn.scaling)) + +Adapt.adapt_structure(to, nn :: NN) = + NN(Adapt.adapt(to, nn.model), + Adapt.adapt(to, nn.ps), + Adapt.adapt(to, nn.st)) + +function NNFluxClosure(arch) + # arch = model.architecture + dev = ifelse(arch == GPU(), gpu_device(), cpu_device()) + nn_model = Chain(Dense(11, 128, relu), Dense(128, 128, relu), Dense(128, 1)) + + # scaling = jldopen("./NN_model2.jld2")["scaling"] + scaling = (; ∂T∂z = ZeroMeanUnitVarianceScaling(-0.0006850967567052092, 0.019041912105983983), + ∂S∂z = ZeroMeanUnitVarianceScaling(-0.00042981832021978374, 0.0028927446724707905), + ∂ρ∂z = ZeroMeanUnitVarianceScaling(-0.0011311157767216616, 0.0008333035237211424), + f = ZeroMeanUnitVarianceScaling(-1.5e-5, 8.73212459828649e-5), + wb = ZeroMeanUnitVarianceScaling(6.539366623323223e-8, 1.827377562065243e-7), + wT = ZeroMeanUnitVarianceScaling(1.8169228278423015e-5, 0.00010721779595955453), + wS = ZeroMeanUnitVarianceScaling(-5.8185988680682135e-6, 1.7691239104281005e-5)) + + # NNs = jldopen("./NN_model2.jld2")["NNs"] + ps = jldopen("./NN_model2.jld2")["u"] |> dev + sts = jldopen("./NN_model2.jld2")["sts"] |> dev + + wT_NN = NN(nn_model, ps.wT, sts.wT) + wS_NN = NN(nn_model, ps.wS, sts.wS) + + return NNFluxClosure(wT_NN, wS_NN, scaling) +end + +DiffusivityFields(grid, tracer_names, bcs, ::NNFluxClosure) = + (; wT = ZFaceField(grid), + wS = ZFaceField(grid)) + +function compute_diffusivities!(diffusivities, closure::NNFluxClosure, model; parameters = :xyz) + arch = model.architecture + grid = model.grid + velocities = model.velocities + tracers = model.tracers + buoyancy = model.buoyancy + coriolis = model.coriolis + clock = model.clock + top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) + + launch!(arch, grid, parameters, + _compute_residual_fluxes!, diffusivities, grid, closure, tracers, velocities, buoyancy, coriolis, top_tracer_bcs, clock) + + + return nothing +end + +@kernel function _compute_residual_fluxes!(diffusivities, grid, closure, tracers, velocities, buoyancy, coriolis, top_tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + + # Find a way to extract the type FT + nn_input = @private eltype(grid) 11 + + scaling = closure.scaling + + nn_input[10] = Jᵇ = scaling.wb(top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, tracers)) + nn_input[11] = fᶜᶜ = scaling.f(fᶜᶜᵃ(i, j, k, grid, coriolis)) + + nn_input[1] = ∂Tᵢ₋₁ = scaling.∂T∂z(∂zᶜᶜᶠ(i, j, k-1, grid, tracers.T)) + nn_input[2] = ∂Tᵢ = scaling.∂T∂z(∂zᶜᶜᶠ(i, j, k, grid, tracers.T)) + nn_input[3] = ∂Tᵢ₊₁ = scaling.∂T∂z(∂zᶜᶜᶠ(i, j, k+1, grid, tracers.T)) + + nn_input[4] = ∂Sᵢ₋₁ = scaling.∂S∂z(∂zᶜᶜᶠ(i, j, k-1, grid, tracers.S)) + nn_input[5] = ∂Sᵢ = scaling.∂S∂z(∂zᶜᶜᶠ(i, j, k, grid, tracers.S)) + nn_input[6] = ∂Sᵢ₊₁ = scaling.∂S∂z(∂zᶜᶜᶠ(i, j, k+1, grid, tracers.S)) + + ρ₀ = buoyancy.model.equation_of_state.reference_density + g = buoyancy.model.gravitational_acceleration + + nn_input[7] = ∂σᵢ = scaling.∂ρ∂z(ρ₀ * ∂z_b(i, j, k, grid, buoyancy, tracers) / g) + nn_input[8] = ∂σᵢ₋₁ = scaling.∂ρ∂z(ρ₀ * ∂z_b(i, j, k, grid, buoyancy, tracers) / g) + nn_input[9] = ∂σᵢ₊₁ = scaling.∂ρ∂z(ρ₀ * ∂z_b(i, j, k, grid, buoyancy, tracers) / g) + + @inbounds wT = inv(scaling.wT)(closure.wT(nn_input)) + @inbounds wS = inv(scaling.wS)(closure.wS(nn_input)) + + @inbounds diffusivities.wT[i, j, k] = ifelse(k > grid.Nz - 2, 0, wT) + @inbounds diffusivities.wS[i, j, k] = ifelse(k > grid.Nz - 2, 0, wS) +end + +# Write here your constructor +# NNFluxClosure() = ... insert NN here ... (make sure it is on GPU if you need it on GPU!) + +const NNC = NNFluxClosure + +##### +##### Abstract Smagorinsky functionality +##### + +# Horizontal fluxes are zero! +@inline viscous_flux_wz( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_ux( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_uy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline diffusive_flux_x(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) +@inline diffusive_flux_y(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) + +# Viscous fluxes are zero (for now) +@inline viscous_flux_uz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) + +# The only function extended by NNFluxClosure +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{1}, c, clock, fields, buoyancy) = @inbounds K.wT[i, j, k] +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{2}, c, clock, fields, buoyancy) = @inbounds K.wS[i, j, k] \ No newline at end of file diff --git a/feature_scaling.jl b/feature_scaling.jl new file mode 100644 index 0000000000..6f141e9ae2 --- /dev/null +++ b/feature_scaling.jl @@ -0,0 +1,84 @@ +using Statistics + +abstract type AbstractFeatureScaling end + +##### +##### Zero-mean unit-variance feature scaling +##### + +struct ZeroMeanUnitVarianceScaling{T} <: AbstractFeatureScaling + μ :: T + σ :: T +end + +""" + ZeroMeanUnitVarianceScaling(data) + +Returns a feature scaler for `data` with zero mean and unit variance. +""" +function ZeroMeanUnitVarianceScaling(data) + μ, σ = mean(data), std(data) + return ZeroMeanUnitVarianceScaling(μ, σ) +end + +scale(x, s::ZeroMeanUnitVarianceScaling) = (x .- s.μ) / s.σ +unscale(y, s::ZeroMeanUnitVarianceScaling) = s.σ * y .+ s.μ + +##### +##### Min-max feature scaling +##### + +struct MinMaxScaling{T} <: AbstractFeatureScaling + a :: T + b :: T + data_min :: T + data_max :: T +end + +""" + MinMaxScaling(data; a=0, b=1) + +Returns a feature scaler for `data` with minimum `a` and `maximum `b`. +""" +function MinMaxScaling(data; a=0, b=1) + data_min, data_max = extrema(data) + return MinMaxScaling{typeof(data_min)}(a, b, data_min, data_max) +end + +scale(x, s::MinMaxScaling) = s.a + (x - s.data_min) * (s.b - s.a) / (s.data_max - s.data_min) +unscale(y, s::MinMaxScaling) = s.data_min .+ (y .- s.a) * (s.data_max - s.data_min) / (s.b - s.a) + +##### +##### Convenience functions +##### + +(s::AbstractFeatureScaling)(x) = scale(x, s) +Base.inv(s::AbstractFeatureScaling) = y -> unscale(y, s) + +struct DiffusivityScaling{T} <: AbstractFeatureScaling + ν₀ :: T + κ₀ :: T + ν₁ :: T + κ₁ :: T +end + +function DiffusivityScaling(ν₀=1e-5, κ₀=1e-5, ν₁=0.1, κ₁=0.1) + return DiffusivityScaling(ν₀, κ₀, ν₁, κ₁) +end + +function scale(x, s::DiffusivityScaling) + ν, κ = x + ν₀, κ₀, ν₁, κ₁ = s.ν₀, s.κ₀, s.ν₁, s.κ₁ + return ν₀ + ν * ν₁, κ₀ + κ * κ₁ +end + +function unscale(y, s::DiffusivityScaling) + ν, κ = y + ν₀, κ₀, ν₁, κ₁ = s.ν₀, s.κ₀, s.ν₁, s.κ₁ + return (ν - ν₀) / ν₁, (κ - κ₀) / κ₁ +end + +(s::DiffusivityScaling)(x) = scale(x, s) +Base.inv(s::DiffusivityScaling) = y -> unscale(y, s) + + diff --git a/xin_kai_vertical_diffusivity_local.jl b/xin_kai_vertical_diffusivity_local.jl new file mode 100644 index 0000000000..240487542c --- /dev/null +++ b/xin_kai_vertical_diffusivity_local.jl @@ -0,0 +1,230 @@ +using Oceananigans +using Oceananigans.Architectures: architecture +using Oceananigans.BuoyancyModels: ∂z_b +using Oceananigans.Operators +using Oceananigans.Grids: inactive_node +using Oceananigans.Operators: ℑzᵃᵃᶜ, ℑxyᶠᶠᵃ, ℑxyᶜᶜᵃ + +using Adapt + +using KernelAbstractions: @index, @kernel +using KernelAbstractions.Extras.LoopInfo: @unroll + +using Oceananigans.TurbulenceClosures: + tapering_factorᶠᶜᶜ, + tapering_factorᶜᶠᶜ, + tapering_factorᶜᶜᶠ, + tapering_factor, + SmallSlopeIsopycnalTensor, + AbstractScalarDiffusivity, + ExplicitTimeDiscretization, + FluxTapering, + isopycnal_rotation_tensor_xz_ccf, + isopycnal_rotation_tensor_yz_ccf, + isopycnal_rotation_tensor_zz_ccf + +import Oceananigans.TurbulenceClosures: + compute_diffusivities!, + DiffusivityFields, + viscosity, + diffusivity, + getclosure, + top_buoyancy_flux, + diffusive_flux_x, + diffusive_flux_y, + diffusive_flux_z, + viscous_flux_ux, + viscous_flux_vx, + viscous_flux_uy, + viscous_flux_vy + +using Oceananigans.Utils: launch! +using Oceananigans.Coriolis: fᶠᶠᵃ +using Oceananigans.Operators +using Oceananigans.BuoyancyModels: ∂x_b, ∂y_b, ∂z_b + +using Oceananigans.TurbulenceClosures +using Oceananigans.TurbulenceClosures: HorizontalFormulation, VerticalFormulation, AbstractScalarDiffusivity +using Oceananigans.TurbulenceClosures: AbstractScalarBiharmonicDiffusivity +using Oceananigans.Operators +using Oceananigans.Operators: Δxᶜᶜᶜ, Δyᶜᶜᶜ, ℑxyᶜᶜᵃ, ζ₃ᶠᶠᶜ, div_xyᶜᶜᶜ +using Oceananigans.Operators: Δx, Δy +using Oceananigans.Operators: ℑxyz + +using Oceananigans.Operators: ℑxyzᶜᶜᶠ, ℑyzᵃᶜᶠ, ℑxzᶜᵃᶠ, Δxᶜᶜᶜ, Δyᶜᶜᶜ + +using Oceananigans.BoundaryConditions + +struct XinKaiLocalVerticalDiffusivity{TD, FT} <: AbstractScalarDiffusivity{TD, VerticalFormulation, 2} + ν₀ :: FT + νˢʰ :: FT + νᶜⁿ :: FT + Prₜ :: FT + Riᶜ :: FT + δRi :: FT +end + +function XinKaiLocalVerticalDiffusivity{TD}(ν₀ :: FT, + νˢʰ :: FT, + νᶜⁿ :: FT, + Prₜ :: FT, + Riᶜ :: FT, + δRi :: FT) where {TD, FT} + + return XinKaiLocalVerticalDiffusivity{TD, FT}(ν₀, νˢʰ, νᶜⁿ, Prₜ, Riᶜ, δRi) +end + +function XinKaiLocalVerticalDiffusivity(time_discretization = VerticallyImplicitTimeDiscretization(), + FT = Float64; + ν₀ = 1e-5, + νˢʰ = 0.0885, + νᶜⁿ = 4.3668, + Prₜ = 1.207, + Riᶜ = - 0.21982, + δRi = 8.342e-4) + + TD = typeof(time_discretization) + + return XinKaiLocalVerticalDiffusivity{TD}(convert(FT, ν₀), + convert(FT, νˢʰ), + convert(FT, νᶜⁿ), + convert(FT, Prₜ), + convert(FT, Riᶜ), + convert(FT, δRi)) +end + +XinKaiLocalVerticalDiffusivity(FT::DataType; kw...) = + XinKaiLocalVerticalDiffusivity(VerticallyImplicitTimeDiscretization(), FT; kw...) + +Adapt.adapt_structure(to, clo::XinKaiLocalVerticalDiffusivity{TD, FT}) where {TD, FT} = + XinKaiLocalVerticalDiffusivity{TD, FT}(clo.ν₀, clo.νˢʰ, clo.νᶜⁿ, clo.Prₜ, clo.Riᶜ, clo.δRi) + +##### +##### Diffusivity field utilities +##### + +const RBVD = XinKaiLocalVerticalDiffusivity +const RBVDArray = AbstractArray{<:RBVD} +const FlavorOfXKVD = Union{RBVD, RBVDArray} +const c = Center() +const f = Face() + +@inline viscosity_location(::FlavorOfXKVD) = (c, c, f) +@inline diffusivity_location(::FlavorOfXKVD) = (c, c, f) + +@inline viscosity(::FlavorOfXKVD, diffusivities) = diffusivities.κᵘ +@inline diffusivity(::FlavorOfXKVD, diffusivities, id) = diffusivities.κᶜ + +with_tracers(tracers, closure::FlavorOfXKVD) = closure + +# Note: computing diffusivities at cell centers for now. +function DiffusivityFields(grid, tracer_names, bcs, closure::FlavorOfXKVD) + κᶜ = Field((Center, Center, Face), grid) + κᵘ = Field((Center, Center, Face), grid) + Ri = Field((Center, Center, Face), grid) + return (; κᶜ, κᵘ, Ri) +end + +function compute_diffusivities!(diffusivities, closure::FlavorOfXKVD, model; parameters = :xyz) + arch = model.architecture + grid = model.grid + clock = model.clock + tracers = model.tracers + buoyancy = model.buoyancy + velocities = model.velocities + top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) + + launch!(arch, grid, parameters, + compute_ri_number!, + diffusivities, + grid, + closure, + velocities, + tracers, + buoyancy, + top_tracer_bcs, + clock) + + # Use `only_local_halos` to ensure that no communication occurs during + # this call to fill_halo_regions! + fill_halo_regions!(diffusivities.Ri; only_local_halos=true) + + launch!(arch, grid, parameters, + compute_xinkai_diffusivities!, + diffusivities, + grid, + closure, + velocities, + tracers, + buoyancy, + top_tracer_bcs, + clock) + + return nothing +end + +@inline ϕ²(i, j, k, grid, ϕ, args...) = ϕ(i, j, k, grid, args...)^2 + +@inline function shear_squaredᶜᶜᶠ(i, j, k, grid, velocities) + ∂z_u² = ℑxᶜᵃᵃ(i, j, k, grid, ϕ², ∂zᶠᶜᶠ, velocities.u) + ∂z_v² = ℑyᵃᶜᵃ(i, j, k, grid, ϕ², ∂zᶜᶠᶠ, velocities.v) + return ∂z_u² + ∂z_v² +end + +@inline function Riᶜᶜᶠ(i, j, k, grid, velocities, buoyancy, tracers) + S² = shear_squaredᶜᶜᶠ(i, j, k, grid, velocities) + N² = ∂z_b(i, j, k, grid, buoyancy, tracers) + Ri = N² / S² + + # Clip N² and avoid NaN + return ifelse(N² <= 0, zero(grid), Ri) +end + +const c = Center() +const f = Face() + +@kernel function compute_ri_number!(diffusivities, grid, closure::FlavorOfXKVD, + velocities, tracers, buoyancy, tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + @inbounds diffusivities.Ri[i, j, k] = Riᶜᶜᶠ(i, j, k, grid, velocities, buoyancy, tracers) +end + +@kernel function compute_xinkai_diffusivities!(diffusivities, grid, closure::FlavorOfXKVD, + velocities, tracers, buoyancy, tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + _compute_xinkai_diffusivities!(i, j, k, diffusivities, grid, closure, + velocities, tracers, buoyancy, tracer_bcs, clock) +end + + +@inline function _compute_xinkai_diffusivities!(i, j, k, diffusivities, grid, closure, + velocities, tracers, buoyancy, tracer_bcs, clock) + + # Ensure this works with "ensembles" of closures, in addition to ordinary single closures + closure_ij = getclosure(i, j, closure) + + ν₀ = closure_ij.ν₀ + νˢʰ = closure_ij.νˢʰ + νᶜⁿ = closure_ij.νᶜⁿ + Prₜ = closure_ij.Prₜ + Riᶜ = closure_ij.Riᶜ + δRi = closure_ij.δRi + + # Convection and entrainment + N² = ∂z_b(i, j, k, grid, buoyancy, tracers) + + # Conditions + convecting = N² < 0 # applies regardless of Qᵇ + + # (Potentially) apply a horizontal filter to the Richardson number + Ri = ℑxyᶜᶜᵃ(i, j, k, grid, ℑxyᶠᶠᵃ, diffusivities.Ri) + + # Convective adjustment diffusivity + ν_local = ifelse(convecting, - (νᶜⁿ - νˢʰ) / 2 * tanh(Ri / δRi) + νˢʰ, clamp(Riᶜ * Ri + νˢʰ + ν₀, ν₀, νˢʰ)) + + # Update by averaging in time + @inbounds diffusivities.κᵘ[i, j, k] = ν_local + @inbounds diffusivities.κᶜ[i, j, k] = ν_local / Prₜ + + return nothing +end From bb70e5c2eb048c63e6c933f61c91af15626a37d9 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Mon, 24 Jun 2024 10:28:06 +0100 Subject: [PATCH 002/114] a working model for GPU! --- NN_closure.jl | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/NN_closure.jl b/NN_closure.jl index 8475b3ea5c..e728c63679 100644 --- a/NN_closure.jl +++ b/NN_closure.jl @@ -22,6 +22,7 @@ using Oceananigans.Grids: φnode using Lux, LuxCUDA using JLD2 using ComponentArrays +using StaticArrays using KernelAbstractions: @index, @kernel, @private @@ -43,14 +44,13 @@ struct NN{M, P, S} st :: S end -# @inline (neural_network::NN)(∂Tᵢ₋₁, ∂Tᵢ, ∂Tᵢ₊₁, ∂Sᵢ₋₁, ∂Sᵢ, ∂Sᵢ₊₁, ∂σᵢ₋₁, ∂σᵢ, ∂σᵢ₊₁, Jᵇ, fᶜᶜ, dev) = first(LuxCore.apply(neural_network.model, dev([Jᵇ, fᶜᶜ, ∂Tᵢ, ∂Tᵢ₋₁, ∂Tᵢ₊₁, ∂Sᵢ, ∂Sᵢ₋₁, ∂Sᵢ₊₁, ∂σᵢ, ∂σᵢ₋₁, ∂σᵢ₊₁]), neural_network.ps, neural_network.st)) @inline (neural_network::NN)(input) = first(first(neural_network.model(input, neural_network.ps, neural_network.st))) +@inline tosarray(x::AbstractArray) = SArray{Tuple{size(x)...}}(x) struct NNFluxClosure{A <: NN, S} <: AbstractTurbulenceClosure{ExplicitTimeDiscretization, 3} wT :: A wS :: A scaling :: S - # dev :: D end Adapt.adapt_structure(to, nn :: NNFluxClosure) = @@ -64,7 +64,6 @@ Adapt.adapt_structure(to, nn :: NN) = Adapt.adapt(to, nn.st)) function NNFluxClosure(arch) - # arch = model.architecture dev = ifelse(arch == GPU(), gpu_device(), cpu_device()) nn_model = Chain(Dense(11, 128, relu), Dense(128, 128, relu), Dense(128, 1)) @@ -78,8 +77,11 @@ function NNFluxClosure(arch) wS = ZeroMeanUnitVarianceScaling(-5.8185988680682135e-6, 1.7691239104281005e-5)) # NNs = jldopen("./NN_model2.jld2")["NNs"] - ps = jldopen("./NN_model2.jld2")["u"] |> dev - sts = jldopen("./NN_model2.jld2")["sts"] |> dev + ps = jldopen("./NN_model2.jld2")["u"] + sts = jldopen("./NN_model2.jld2")["sts"] + + ps_static = Lux.recursive_map(tosarray, ps) + sts_static = Lux.recursive_map(tosarray, sts) wT_NN = NN(nn_model, ps.wT, sts.wT) wS_NN = NN(nn_model, ps.wS, sts.wS) @@ -104,7 +106,6 @@ function compute_diffusivities!(diffusivities, closure::NNFluxClosure, model; pa launch!(arch, grid, parameters, _compute_residual_fluxes!, diffusivities, grid, closure, tracers, velocities, buoyancy, coriolis, top_tracer_bcs, clock) - return nothing end From e845dd5d8edb38041185b5195c5746c7cb377136 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 25 Jun 2024 18:23:23 +0100 Subject: [PATCH 003/114] it actually works on the gpu now! --- NN_1D_model.jl | 195 +++++++++++++++++++++++++++++++++++++++++++ NN_closure_global.jl | 188 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 383 insertions(+) create mode 100644 NN_1D_model.jl create mode 100644 NN_closure_global.jl diff --git a/NN_1D_model.jl b/NN_1D_model.jl new file mode 100644 index 0000000000..a1f1b118fe --- /dev/null +++ b/NN_1D_model.jl @@ -0,0 +1,195 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("NN_closure_global.jl") +include("xin_kai_vertical_diffusivity_local.jl") +include("feature_scaling.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 + +# Architecture +model_architecture = GPU() + +# number of grid points +Nz = 32 + +const Lz = 256 + +grid = RectilinearGrid(model_architecture, + topology = (Flat, Flat, Bounded), + size = Nz, + halo = 3, + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +T_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(1e-4), bottom=FluxBoundaryCondition(0.0)) + +##### +##### Coriolis +##### + +const f₀ = 8e-5 +const β = 1e-11 +# coriolis = BetaPlane(f₀=f₀, β = β) +coriolis = FPlane(f=f₀) + +##### +##### Forcing and initial condition +##### +const dTdz = 0.014 +const dSdz = 0.0021 + +const T_surface = 20.0 +const S_surface = 36.6 + +T_initial(z) = dTdz * z + T_surface +S_initial(z) = dSdz * z + S_surface + +nn_closure = NNFluxClosure(model_architecture) +base_closure = XinKaiLocalVerticalDiffusivity() + +##### +##### Model building +##### + +@info "Building a model..." + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = ImplicitFreeSurface(), + momentum_advection = WENO(grid = grid), + tracer_advection = WENO(grid = grid), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = (nn_closure, base_closure), + tracers = (:T, :S), + boundary_conditions = (; T = T_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(z) = T_initial(z) + 1e-6 * noise(z) +S_initial_noisy(z) = S_initial(z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 2days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.1, max_change=1.1, max_Δt=20minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(20)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(20)) + +##### +##### Diagnostics +##### + +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S + +Tbar = Field(Average(T, dims = (1,2))) +Sbar = Field(Average(S, dims = (1,2))) + +averaged_outputs = (; Tbar, Sbar) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:jld2] = JLD2OutputWriter(model, averaged_outputs, + filename = "NN_1D_channel_averages", + schedule = TimeInterval(10minutes), + overwrite_existing = true) + +@info "Running the simulation..." + +try + run!(simulation, pickup = false) +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +# ##### +# ##### Visualization +# ##### + +using CairoMakie + +Tbar_data = FieldTimeSeries("./NN_1D_channel_averages.jld2", "Tbar") +Sbar_data = FieldTimeSeries("./NN_1D_channel_averages.jld2", "Sbar") + +zC = znodes(Tbar_data.grid, Center()) +zF = znodes(Tbar_data.grid, Face()) + +Nt = length(Tbar_data.times) +#%% +fig = Figure(size = (900, 600)) +axT = CairoMakie.Axis(fig[1, 1], xlabel = "T (°C)", ylabel = "z (m)") +axS = CairoMakie.Axis(fig[1, 2], xlabel = "S (g kg⁻¹)", ylabel = "z (m)") +n = Observable(1) + +Tbarₙ = @lift interior(Tbar_data[$n], 1, 1, :) +Sbarₙ = @lift interior(Sbar_data[$n], 1, 1, :) + +title_str = @lift "Time: $(round(Tbar_data.times[$n] / 86400, digits=3)) days" + +lines!(axT, Tbarₙ, zC) +lines!(axS, Sbarₙ, zC) + +Label(fig[0, :], title_str, tellwidth = false) + +CairoMakie.record(fig, "./NN_1D_fields.mp4", 1:Nt, framerate=10) do nn + n[] = nn +end + +display(fig) +#%% \ No newline at end of file diff --git a/NN_closure_global.jl b/NN_closure_global.jl new file mode 100644 index 0000000000..622be5372c --- /dev/null +++ b/NN_closure_global.jl @@ -0,0 +1,188 @@ +import Oceananigans.TurbulenceClosures: + compute_diffusivities!, + DiffusivityFields, + viscosity, + diffusivity, + diffusive_flux_x, + diffusive_flux_y, + diffusive_flux_z, + viscous_flux_ux, + viscous_flux_uy, + viscous_flux_uz, + viscous_flux_vx, + viscous_flux_vy, + viscous_flux_vz, + viscous_flux_wx, + viscous_flux_wy, + viscous_flux_wz + +using Oceananigans.BuoyancyModels: ∂x_b, ∂y_b, ∂z_b +using Oceananigans.Coriolis +using Oceananigans.Grids: φnode +using Oceananigans.Grids: total_size +using Oceananigans: architecture, on_architecture +using Lux, LuxCUDA +using JLD2 +using ComponentArrays +using StaticArrays +using OffsetArrays + +using KernelAbstractions: @index, @kernel, @private + +import Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure, ExplicitTimeDiscretization + +using Adapt + +include("./feature_scaling.jl") + +@inline hack_sind(φ) = sin(φ * π / 180) + +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::HydrostaticSphericalCoriolis) = 2 * coriolis.rotation_rate * hack_sind(φnode(i, j, k, grid, Center(), Center(), Center())) +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::FPlane) = coriolis.f +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::BetaPlane) = coriolis.f₀ + coriolis.β * ynode(i, j, k, grid, Center(), Center(), Center()) + +struct NN{M, P, S} + model :: M + ps :: P + st :: S +end + +@inline (neural_network::NN)(input) = first(neural_network.model(input, neural_network.ps, neural_network.st)) +@inline tosarray(x::AbstractArray) = SArray{Tuple{size(x)...}}(x) + +struct NNFluxClosure{A <: NN, S} <: AbstractTurbulenceClosure{ExplicitTimeDiscretization, 3} + wT :: A + wS :: A + scaling :: S +end + +Adapt.adapt_structure(to, nn :: NNFluxClosure) = + NNFluxClosure(Adapt.adapt(to, nn.wT), + Adapt.adapt(to, nn.wS), + Adapt.adapt(to, nn.scaling)) + +Adapt.adapt_structure(to, nn :: NN) = + NN(Adapt.adapt(to, nn.model), + Adapt.adapt(to, nn.ps), + Adapt.adapt(to, nn.st)) + +function NNFluxClosure(arch) + dev = ifelse(arch == GPU(), gpu_device(), cpu_device()) + nn_model = Chain(Dense(11, 128, relu), Dense(128, 128, relu), Dense(128, 1)) + + # scaling = jldopen("./NN_model2.jld2")["scaling"] + scaling = (; ∂T∂z = ZeroMeanUnitVarianceScaling(-0.0006850967567052092, 0.019041912105983983), + ∂S∂z = ZeroMeanUnitVarianceScaling(-0.00042981832021978374, 0.0028927446724707905), + ∂ρ∂z = ZeroMeanUnitVarianceScaling(-0.0011311157767216616, 0.0008333035237211424), + f = ZeroMeanUnitVarianceScaling(-1.5e-5, 8.73212459828649e-5), + wb = ZeroMeanUnitVarianceScaling(6.539366623323223e-8, 1.827377562065243e-7), + wT = ZeroMeanUnitVarianceScaling(1.8169228278423015e-5, 0.00010721779595955453), + wS = ZeroMeanUnitVarianceScaling(-5.8185988680682135e-6, 1.7691239104281005e-5)) + + # NNs = jldopen("./NN_model2.jld2")["NNs"] + ps = jldopen("./NN_model2.jld2")["u"] |> dev + sts = jldopen("./NN_model2.jld2")["sts"] |> dev + + ps = ps .= 0 + + wT_NN = NN(nn_model, ps.wT, sts.wT) + wS_NN = NN(nn_model, ps.wS, sts.wS) + + return NNFluxClosure(wT_NN, wS_NN, scaling) +end + +function DiffusivityFields(grid, tracer_names, bcs, ::NNFluxClosure) + arch = architecture(grid) + wT = ZFaceField(grid) + wS = ZFaceField(grid) + + Nx_in, Ny_in, Nz_in = total_size(wT) + ox_in, oy_in, oz_in = wT.data.offsets + + wrk_in = OffsetArray(zeros(11, Nx_in, Ny_in, Nz_in), 0, ox_in, oy_in, oz_in) + wrk_in = on_architecture(arch, wrk_in) + + return (; wrk_in, wT, wS) +end + +function compute_diffusivities!(diffusivities, closure::NNFluxClosure, model; parameters = :xyz) + arch = model.architecture + grid = model.grid + velocities = model.velocities + tracers = model.tracers + buoyancy = model.buoyancy + coriolis = model.coriolis + clock = model.clock + top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) + input = diffusivities.wrk_in + wT = diffusivities.wT + wS = diffusivities.wS + + launch!(arch, grid, parameters, + _populate_input!, input, grid, closure, tracers, velocities, buoyancy, coriolis, top_tracer_bcs, clock) + + wT.data.parent .= dropdims(closure.wT(input.parent), dims=1) + wS.data.parent .= dropdims(closure.wS(input.parent), dims=1) + + launch!(arch, grid, parameters, _rescale_nn_fluxes!, diffusivities, grid, closure) + return nothing +end + +@kernel function _populate_input!(input, grid, closure, tracers, velocities, buoyancy, coriolis, top_tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + + scaling = closure.scaling + + ρ₀ = buoyancy.model.equation_of_state.reference_density + g = buoyancy.model.gravitational_acceleration + + @inbounds input[1, i, j, k] = ∂Tᵢ₋₁ = scaling.∂T∂z(∂zᶜᶜᶠ(i, j, k-1, grid, tracers.T)) + @inbounds input[2, i, j, k] = ∂Tᵢ = scaling.∂T∂z(∂zᶜᶜᶠ(i, j, k, grid, tracers.T)) + @inbounds input[3, i, j, k] = ∂Tᵢ₊₁ = scaling.∂T∂z(∂zᶜᶜᶠ(i, j, k+1, grid, tracers.T)) + + @inbounds input[4, i, j, k] = ∂Sᵢ₋₁ = scaling.∂S∂z(∂zᶜᶜᶠ(i, j, k-1, grid, tracers.S)) + @inbounds input[5, i, j, k] = ∂Sᵢ = scaling.∂S∂z(∂zᶜᶜᶠ(i, j, k, grid, tracers.S)) + @inbounds input[6, i, j, k] = ∂Sᵢ₊₁ = scaling.∂S∂z(∂zᶜᶜᶠ(i, j, k+1, grid, tracers.S)) + + @inbounds input[7, i, j, k] = ∂σᵢ₋₁ = scaling.∂ρ∂z(ρ₀ * ∂z_b(i, j, k-1, grid, buoyancy, tracers) / g) + @inbounds input[8, i, j, k] = ∂σᵢ = scaling.∂ρ∂z(ρ₀ * ∂z_b(i, j, k, grid, buoyancy, tracers) / g) + @inbounds input[9, i, j, k] = ∂σᵢ₊₁ = scaling.∂ρ∂z(ρ₀ * ∂z_b(i, j, k+1, grid, buoyancy, tracers) / g) + + @inbounds input[10, i, j, k] = Jᵇ = scaling.wb(top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, tracers)) + @inbounds input[11, i, j, k] = fᶜᶜ = scaling.f(fᶜᶜᵃ(i, j, k, grid, coriolis)) +end + +@kernel function _rescale_nn_fluxes!(diffusivities, grid, closure) + i, j, k = @index(Global, NTuple) + scaling = closure.scaling + @inbounds diffusivities.wT[i, j, k] = ifelse(k >= grid.Nz - 1 || k <= 2, 0, inv(scaling.wT)(diffusivities.wT[i, j, k])) + @inbounds diffusivities.wS[i, j, k] = ifelse(k >= grid.Nz - 1 || k <= 2, 0, inv(scaling.wS)(diffusivities.wS[i, j, k])) +end + +# Write here your constructor +# NNFluxClosure() = ... insert NN here ... (make sure it is on GPU if you need it on GPU!) + +const NNC = NNFluxClosure + +##### +##### Abstract Smagorinsky functionality +##### + +# Horizontal fluxes are zero! +@inline viscous_flux_wz( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_ux( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_uy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline diffusive_flux_x(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) +@inline diffusive_flux_y(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) + +# Viscous fluxes are zero (for now) +@inline viscous_flux_uz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) + +# The only function extended by NNFluxClosure +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{1}, c, clock, fields, buoyancy) = @inbounds K.wT[i, j, k] +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{2}, c, clock, fields, buoyancy) = @inbounds K.wS[i, j, k] \ No newline at end of file From b6552abdb083b478d50d343ea81d1d552b7b5f30 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 25 Jun 2024 18:23:37 +0100 Subject: [PATCH 004/114] build scaling from tuple values --- feature_scaling.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/feature_scaling.jl b/feature_scaling.jl index 6f141e9ae2..0de273e982 100644 --- a/feature_scaling.jl +++ b/feature_scaling.jl @@ -81,4 +81,8 @@ end (s::DiffusivityScaling)(x) = scale(x, s) Base.inv(s::DiffusivityScaling) = y -> unscale(y, s) +function construct_scaling(scaling_params, scaling_type::ZeroMeanUnitVarianceScaling) + return NamedTuple(key=>ZeroMeanUnitVarianceScaling(scaling_params[key].μ, scaling_params[key].σ) for key in keys(scaling_params)) +end + From c9a9c656fae9c826d526bca33c8b012b95312b13 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Wed, 26 Jun 2024 11:30:08 +0100 Subject: [PATCH 005/114] Update XinKaiLocalVerticalDiffusivity with 2 Pr values --- xin_kai_vertical_diffusivity_local.jl | 42 ++++++++++++++++----------- 1 file changed, 25 insertions(+), 17 deletions(-) diff --git a/xin_kai_vertical_diffusivity_local.jl b/xin_kai_vertical_diffusivity_local.jl index 240487542c..533288f213 100644 --- a/xin_kai_vertical_diffusivity_local.jl +++ b/xin_kai_vertical_diffusivity_local.jl @@ -59,36 +59,40 @@ struct XinKaiLocalVerticalDiffusivity{TD, FT} <: AbstractScalarDiffusivity{TD, V ν₀ :: FT νˢʰ :: FT νᶜⁿ :: FT - Prₜ :: FT + Pr_convₜ :: FT + Pr_shearₜ :: FT Riᶜ :: FT δRi :: FT end function XinKaiLocalVerticalDiffusivity{TD}(ν₀ :: FT, - νˢʰ :: FT, - νᶜⁿ :: FT, - Prₜ :: FT, - Riᶜ :: FT, - δRi :: FT) where {TD, FT} + νˢʰ :: FT, + νᶜⁿ :: FT, + Pr_convₜ :: FT, + Pr_shearₜ :: FT, + Riᶜ :: FT, + δRi :: FT) where {TD, FT} - return XinKaiLocalVerticalDiffusivity{TD, FT}(ν₀, νˢʰ, νᶜⁿ, Prₜ, Riᶜ, δRi) + return XinKaiLocalVerticalDiffusivity{TD, FT}(ν₀, νˢʰ, νᶜⁿ, Pr_convₜ, Pr_shearₜ, Riᶜ, δRi) end function XinKaiLocalVerticalDiffusivity(time_discretization = VerticallyImplicitTimeDiscretization(), FT = Float64; ν₀ = 1e-5, - νˢʰ = 0.0885, - νᶜⁿ = 4.3668, - Prₜ = 1.207, - Riᶜ = - 0.21982, - δRi = 8.342e-4) + νˢʰ = 0.04569735882746968, + νᶜⁿ = 0.47887785611155065, + Pr_convₜ = 0.1261854430705509, + Pr_shearₜ = 1.594794053970444, + Riᶜ = 0.9964350402840053, + δRi = 0.05635304878092709) TD = typeof(time_discretization) return XinKaiLocalVerticalDiffusivity{TD}(convert(FT, ν₀), convert(FT, νˢʰ), convert(FT, νᶜⁿ), - convert(FT, Prₜ), + convert(FT, Pr_convₜ), + convert(FT, Pr_shearₜ), convert(FT, Riᶜ), convert(FT, δRi)) end @@ -97,7 +101,7 @@ XinKaiLocalVerticalDiffusivity(FT::DataType; kw...) = XinKaiLocalVerticalDiffusivity(VerticallyImplicitTimeDiscretization(), FT; kw...) Adapt.adapt_structure(to, clo::XinKaiLocalVerticalDiffusivity{TD, FT}) where {TD, FT} = - XinKaiLocalVerticalDiffusivity{TD, FT}(clo.ν₀, clo.νˢʰ, clo.νᶜⁿ, clo.Prₜ, clo.Riᶜ, clo.δRi) + XinKaiLocalVerticalDiffusivity{TD, FT}(clo.ν₀, clo.νˢʰ, clo.νᶜⁿ, clo.Pr_convₜ, clo.Pr_shearₜ, clo.Riᶜ, clo.δRi) ##### ##### Diffusivity field utilities @@ -206,7 +210,8 @@ end ν₀ = closure_ij.ν₀ νˢʰ = closure_ij.νˢʰ νᶜⁿ = closure_ij.νᶜⁿ - Prₜ = closure_ij.Prₜ + Pr_convₜ = closure_ij.Pr_convₜ + Pr_shearₜ = closure_ij.Pr_shearₜ Riᶜ = closure_ij.Riᶜ δRi = closure_ij.δRi @@ -220,11 +225,14 @@ end Ri = ℑxyᶜᶜᵃ(i, j, k, grid, ℑxyᶠᶠᵃ, diffusivities.Ri) # Convective adjustment diffusivity - ν_local = ifelse(convecting, - (νᶜⁿ - νˢʰ) / 2 * tanh(Ri / δRi) + νˢʰ, clamp(Riᶜ * Ri + νˢʰ + ν₀, ν₀, νˢʰ)) + ν_local = ifelse(convecting, (νˢʰ - νᶜⁿ) * tanh(Ri / δRi) + νˢʰ, clamp((ν₀ - νˢʰ) * Ri / Riᶜ + νˢʰ, ν₀, νˢʰ)) # Update by averaging in time + # @inbounds diffusivities.κᵘ[i, j, k] = ifelse(k == 1, 0, ν_local) + # @inbounds diffusivities.κᶜ[i, j, k] = ifelse(k == 1, 0, ifelse(convecting, ν_local / Pr_convₜ, ν_local / Pr_shearₜ)) + @inbounds diffusivities.κᵘ[i, j, k] = ν_local - @inbounds diffusivities.κᶜ[i, j, k] = ν_local / Prₜ + @inbounds diffusivities.κᶜ[i, j, k] = ifelse(convecting, ν_local / Pr_convₜ, ν_local / Pr_shearₜ) return nothing end From 9a07b82ba19f1bda8a74037491c688fa633204d2 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Wed, 26 Jun 2024 14:51:42 +0100 Subject: [PATCH 006/114] fix function construct_scaling to construct_zeromeanunitvariance_scaling --- feature_scaling.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feature_scaling.jl b/feature_scaling.jl index 0de273e982..4d3d8592a6 100644 --- a/feature_scaling.jl +++ b/feature_scaling.jl @@ -81,7 +81,7 @@ end (s::DiffusivityScaling)(x) = scale(x, s) Base.inv(s::DiffusivityScaling) = y -> unscale(y, s) -function construct_scaling(scaling_params, scaling_type::ZeroMeanUnitVarianceScaling) +function construct_zeromeanunitvariance_scaling(scaling_params) return NamedTuple(key=>ZeroMeanUnitVarianceScaling(scaling_params[key].μ, scaling_params[key].σ) for key in keys(scaling_params)) end From 74c8b611f6652aa4a3146071ba4f93f124931d2b Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Wed, 26 Jun 2024 15:43:28 +0100 Subject: [PATCH 007/114] update using KernelParameters --- xin_kai_vertical_diffusivity_local.jl | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/xin_kai_vertical_diffusivity_local.jl b/xin_kai_vertical_diffusivity_local.jl index 533288f213..01c3aa6d4a 100644 --- a/xin_kai_vertical_diffusivity_local.jl +++ b/xin_kai_vertical_diffusivity_local.jl @@ -138,7 +138,12 @@ function compute_diffusivities!(diffusivities, closure::FlavorOfXKVD, model; par velocities = model.velocities top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) - launch!(arch, grid, parameters, + Nx_in, Ny_in, Nz_in = total_size(diffusivities.κᶜ) + ox_in, oy_in, oz_in = diffusivities.κᶜ.data.offsets + + kp = KernelParameters((Nx_in, Ny_in, Nz_in), (ox_in, oy_in, oz_in)) + + launch!(arch, grid, kp, compute_ri_number!, diffusivities, grid, @@ -153,7 +158,7 @@ function compute_diffusivities!(diffusivities, closure::FlavorOfXKVD, model; par # this call to fill_halo_regions! fill_halo_regions!(diffusivities.Ri; only_local_halos=true) - launch!(arch, grid, parameters, + launch!(arch, grid, kp, compute_xinkai_diffusivities!, diffusivities, grid, @@ -228,11 +233,8 @@ end ν_local = ifelse(convecting, (νˢʰ - νᶜⁿ) * tanh(Ri / δRi) + νˢʰ, clamp((ν₀ - νˢʰ) * Ri / Riᶜ + νˢʰ, ν₀, νˢʰ)) # Update by averaging in time - # @inbounds diffusivities.κᵘ[i, j, k] = ifelse(k == 1, 0, ν_local) - # @inbounds diffusivities.κᶜ[i, j, k] = ifelse(k == 1, 0, ifelse(convecting, ν_local / Pr_convₜ, ν_local / Pr_shearₜ)) - - @inbounds diffusivities.κᵘ[i, j, k] = ν_local - @inbounds diffusivities.κᶜ[i, j, k] = ifelse(convecting, ν_local / Pr_convₜ, ν_local / Pr_shearₜ) + @inbounds diffusivities.κᵘ[i, j, k] = ifelse(k <= 1 || k >= grid.Nz+1, 0, ν_local) + @inbounds diffusivities.κᶜ[i, j, k] = ifelse(k <= 1 || k >= grid.Nz+1, 0, ifelse(convecting, ν_local / Pr_convₜ, ν_local / Pr_shearₜ)) return nothing end From 6f97c3fe242d4788ec11eeca077818869fa58fb8 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Wed, 26 Jun 2024 15:43:49 +0100 Subject: [PATCH 008/114] fix nn closure bug --- NN_closure_global.jl | 47 +++++++++++++++++++++++--------------------- 1 file changed, 25 insertions(+), 22 deletions(-) diff --git a/NN_closure_global.jl b/NN_closure_global.jl index 622be5372c..ddebe03cc0 100644 --- a/NN_closure_global.jl +++ b/NN_closure_global.jl @@ -20,6 +20,7 @@ using Oceananigans.BuoyancyModels: ∂x_b, ∂y_b, ∂z_b using Oceananigans.Coriolis using Oceananigans.Grids: φnode using Oceananigans.Grids: total_size +using Oceananigans.Utils: KernelParameters using Oceananigans: architecture, on_architecture using Lux, LuxCUDA using JLD2 @@ -68,25 +69,23 @@ Adapt.adapt_structure(to, nn :: NN) = function NNFluxClosure(arch) dev = ifelse(arch == GPU(), gpu_device(), cpu_device()) - nn_model = Chain(Dense(11, 128, relu), Dense(128, 128, relu), Dense(128, 1)) + nn_path = "./NDE_FC_Qb_15simnew_2layer_64_relu_2Pr_model.jld2" - # scaling = jldopen("./NN_model2.jld2")["scaling"] - scaling = (; ∂T∂z = ZeroMeanUnitVarianceScaling(-0.0006850967567052092, 0.019041912105983983), - ∂S∂z = ZeroMeanUnitVarianceScaling(-0.00042981832021978374, 0.0028927446724707905), - ∂ρ∂z = ZeroMeanUnitVarianceScaling(-0.0011311157767216616, 0.0008333035237211424), - f = ZeroMeanUnitVarianceScaling(-1.5e-5, 8.73212459828649e-5), - wb = ZeroMeanUnitVarianceScaling(6.539366623323223e-8, 1.827377562065243e-7), - wT = ZeroMeanUnitVarianceScaling(1.8169228278423015e-5, 0.00010721779595955453), - wS = ZeroMeanUnitVarianceScaling(-5.8185988680682135e-6, 1.7691239104281005e-5)) + ps, sts, scaling_params, wT_model, wS_model = jldopen(nn_path, "r") do file + ps = file["u"] |> dev + sts = file["sts"] |> dev + scaling_params = file["scaling"] + wT_model = file["model"].wT + wS_model = file["model"].wS + return ps, sts, scaling_params, wT_model, wS_model + end - # NNs = jldopen("./NN_model2.jld2")["NNs"] - ps = jldopen("./NN_model2.jld2")["u"] |> dev - sts = jldopen("./NN_model2.jld2")["sts"] |> dev + scaling = construct_zeromeanunitvariance_scaling(scaling_params) - ps = ps .= 0 + # ps = ps .= 0 - wT_NN = NN(nn_model, ps.wT, sts.wT) - wS_NN = NN(nn_model, ps.wS, sts.wS) + wT_NN = NN(wT_model, ps.wT, sts.wT) + wS_NN = NN(wS_model, ps.wS, sts.wS) return NNFluxClosure(wT_NN, wS_NN, scaling) end @@ -118,13 +117,17 @@ function compute_diffusivities!(diffusivities, closure::NNFluxClosure, model; pa wT = diffusivities.wT wS = diffusivities.wS - launch!(arch, grid, parameters, + Nx_in, Ny_in, Nz_in = total_size(wT) + ox_in, oy_in, oz_in = wT.data.offsets + kp = KernelParameters((Nx_in, Ny_in, Nz_in), (ox_in, oy_in, oz_in)) + + launch!(arch, grid, kp, _populate_input!, input, grid, closure, tracers, velocities, buoyancy, coriolis, top_tracer_bcs, clock) wT.data.parent .= dropdims(closure.wT(input.parent), dims=1) wS.data.parent .= dropdims(closure.wS(input.parent), dims=1) - launch!(arch, grid, parameters, _rescale_nn_fluxes!, diffusivities, grid, closure) + launch!(arch, grid, kp, _rescale_nn_fluxes!, diffusivities, grid, closure) return nothing end @@ -144,9 +147,9 @@ end @inbounds input[5, i, j, k] = ∂Sᵢ = scaling.∂S∂z(∂zᶜᶜᶠ(i, j, k, grid, tracers.S)) @inbounds input[6, i, j, k] = ∂Sᵢ₊₁ = scaling.∂S∂z(∂zᶜᶜᶠ(i, j, k+1, grid, tracers.S)) - @inbounds input[7, i, j, k] = ∂σᵢ₋₁ = scaling.∂ρ∂z(ρ₀ * ∂z_b(i, j, k-1, grid, buoyancy, tracers) / g) - @inbounds input[8, i, j, k] = ∂σᵢ = scaling.∂ρ∂z(ρ₀ * ∂z_b(i, j, k, grid, buoyancy, tracers) / g) - @inbounds input[9, i, j, k] = ∂σᵢ₊₁ = scaling.∂ρ∂z(ρ₀ * ∂z_b(i, j, k+1, grid, buoyancy, tracers) / g) + @inbounds input[7, i, j, k] = ∂σᵢ₋₁ = scaling.∂ρ∂z(-ρ₀ * ∂z_b(i, j, k-1, grid, buoyancy, tracers) / g) + @inbounds input[8, i, j, k] = ∂σᵢ = scaling.∂ρ∂z(-ρ₀ * ∂z_b(i, j, k, grid, buoyancy, tracers) / g) + @inbounds input[9, i, j, k] = ∂σᵢ₊₁ = scaling.∂ρ∂z(-ρ₀ * ∂z_b(i, j, k+1, grid, buoyancy, tracers) / g) @inbounds input[10, i, j, k] = Jᵇ = scaling.wb(top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, tracers)) @inbounds input[11, i, j, k] = fᶜᶜ = scaling.f(fᶜᶜᵃ(i, j, k, grid, coriolis)) @@ -155,8 +158,8 @@ end @kernel function _rescale_nn_fluxes!(diffusivities, grid, closure) i, j, k = @index(Global, NTuple) scaling = closure.scaling - @inbounds diffusivities.wT[i, j, k] = ifelse(k >= grid.Nz - 1 || k <= 2, 0, inv(scaling.wT)(diffusivities.wT[i, j, k])) - @inbounds diffusivities.wS[i, j, k] = ifelse(k >= grid.Nz - 1 || k <= 2, 0, inv(scaling.wS)(diffusivities.wS[i, j, k])) + @inbounds diffusivities.wT[i, j, k] = ifelse(k >= grid.Nz - 1 || k <= 3, inv(scaling.wT)(0), inv(scaling.wT)(diffusivities.wT[i, j, k])) + @inbounds diffusivities.wS[i, j, k] = ifelse(k >= grid.Nz - 1 || k <= 3, inv(scaling.wS)(0), inv(scaling.wS)(diffusivities.wS[i, j, k])) end # Write here your constructor From 65f15e4922768b35bfd372791138b0cb12309f40 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Wed, 26 Jun 2024 15:43:58 +0100 Subject: [PATCH 009/114] test script for nn closure --- NN_1D_model.jl | 35 ++++++++++++++++++++--------------- 1 file changed, 20 insertions(+), 15 deletions(-) diff --git a/NN_1D_model.jl b/NN_1D_model.jl index a1f1b118fe..fc8615247b 100644 --- a/NN_1D_model.jl +++ b/NN_1D_model.jl @@ -18,14 +18,16 @@ using Oceananigans.OutputReaders: FieldTimeSeries using Oceananigans.Grids: xnode, ynode, znode using SeawaterPolynomials using SeawaterPolynomials:TEOS10 +import SeawaterPolynomials.TEOS10: s, ΔS, Sₐᵤ +s(Sᴬ::Number) = Sᴬ + ΔS >= 0 ? √((Sᴬ + ΔS) / Sₐᵤ) : NaN # Architecture -model_architecture = GPU() +model_architecture = CPU() # number of grid points -Nz = 32 +Nz = 64 -const Lz = 256 +const Lz = 512 grid = RectilinearGrid(model_architecture, topology = (Flat, Flat, Bounded), @@ -38,7 +40,13 @@ grid = RectilinearGrid(model_architecture, ##### ##### Boundary conditions ##### -T_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(1e-4), bottom=FluxBoundaryCondition(0.0)) +const dTdz = 0.014 +const dSdz = 0.0021 + +const T_surface = 20.0 +const S_surface = 36.6 + +T_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(1e-4), bottom=GradientBoundaryCondition(dTdz)) ##### ##### Coriolis @@ -52,12 +60,6 @@ coriolis = FPlane(f=f₀) ##### ##### Forcing and initial condition ##### -const dTdz = 0.014 -const dSdz = 0.0021 - -const T_surface = 20.0 -const S_surface = 36.6 - T_initial(z) = dTdz * z + T_surface S_initial(z) = dSdz * z + S_surface @@ -78,6 +80,7 @@ model = HydrostaticFreeSurfaceModel( buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), coriolis = coriolis, closure = (nn_closure, base_closure), + # closure = base_closure, tracers = (:T, :S), boundary_conditions = (; T = T_bcs), ) @@ -101,7 +104,7 @@ update_state!(model) ##### Simulation building ##### Δt₀ = 5minutes -stop_time = 2days +stop_time = 10days simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) @@ -175,7 +178,7 @@ Nt = length(Tbar_data.times) fig = Figure(size = (900, 600)) axT = CairoMakie.Axis(fig[1, 1], xlabel = "T (°C)", ylabel = "z (m)") axS = CairoMakie.Axis(fig[1, 2], xlabel = "S (g kg⁻¹)", ylabel = "z (m)") -n = Observable(1) +n = Observable(Nt) Tbarₙ = @lift interior(Tbar_data[$n], 1, 1, :) Sbarₙ = @lift interior(Sbar_data[$n], 1, 1, :) @@ -187,9 +190,11 @@ lines!(axS, Sbarₙ, zC) Label(fig[0, :], title_str, tellwidth = false) -CairoMakie.record(fig, "./NN_1D_fields.mp4", 1:Nt, framerate=10) do nn - n[] = nn -end +# CairoMakie.record(fig, "./NN_1D_fields.mp4", 1:Nt, framerate=10) do nn +# n[] = nn +# xlims!(axT, nothing, nothing) +# xlims!(axS, nothing, nothing) +# end display(fig) #%% \ No newline at end of file From cb601ce3e8e6b46d6c07c618d58e939538645869 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 27 Jun 2024 19:52:53 +0100 Subject: [PATCH 010/114] fix N2 average, tracer diffusivity expression for local diffusivity closure with 2 Pr --- xin_kai_vertical_diffusivity_local.jl | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/xin_kai_vertical_diffusivity_local.jl b/xin_kai_vertical_diffusivity_local.jl index 01c3aa6d4a..ad85509c00 100644 --- a/xin_kai_vertical_diffusivity_local.jl +++ b/xin_kai_vertical_diffusivity_local.jl @@ -186,7 +186,7 @@ end Ri = N² / S² # Clip N² and avoid NaN - return ifelse(N² <= 0, zero(grid), Ri) + return ifelse(N² == 0, zero(grid), Ri) end const c = Center() @@ -205,7 +205,6 @@ end velocities, tracers, buoyancy, tracer_bcs, clock) end - @inline function _compute_xinkai_diffusivities!(i, j, k, diffusivities, grid, closure, velocities, tracers, buoyancy, tracer_bcs, clock) @@ -220,21 +219,23 @@ end Riᶜ = closure_ij.Riᶜ δRi = closure_ij.δRi - # Convection and entrainment - N² = ∂z_b(i, j, k, grid, buoyancy, tracers) - - # Conditions - convecting = N² < 0 # applies regardless of Qᵇ + κ₀ = ν₀ / Pr_shearₜ + κˢʰ = νˢʰ / Pr_shearₜ + κᶜⁿ = νᶜⁿ / Pr_convₜ # (Potentially) apply a horizontal filter to the Richardson number Ri = ℑxyᶜᶜᵃ(i, j, k, grid, ℑxyᶠᶠᵃ, diffusivities.Ri) + # Conditions + convecting = Ri < 0 # applies regardless of Qᵇ + # Convective adjustment diffusivity ν_local = ifelse(convecting, (νˢʰ - νᶜⁿ) * tanh(Ri / δRi) + νˢʰ, clamp((ν₀ - νˢʰ) * Ri / Riᶜ + νˢʰ, ν₀, νˢʰ)) + κ_local = ifelse(convecting, (κˢʰ - κᶜⁿ) * tanh(Ri / δRi) + κˢʰ, clamp((κ₀ - κˢʰ) * Ri / Riᶜ + κˢʰ, κ₀, κˢʰ)) # Update by averaging in time @inbounds diffusivities.κᵘ[i, j, k] = ifelse(k <= 1 || k >= grid.Nz+1, 0, ν_local) - @inbounds diffusivities.κᶜ[i, j, k] = ifelse(k <= 1 || k >= grid.Nz+1, 0, ifelse(convecting, ν_local / Pr_convₜ, ν_local / Pr_shearₜ)) + @inbounds diffusivities.κᶜ[i, j, k] = ifelse(k <= 1 || k >= grid.Nz+1, 0, κ_local) return nothing end From dbe575942583738fb0d69c192551624bd863be2b Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Fri, 28 Jun 2024 11:30:44 +0100 Subject: [PATCH 011/114] 2Pr version of local physical closure --- xin_kai_vertical_diffusivity_local.jl | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/xin_kai_vertical_diffusivity_local.jl b/xin_kai_vertical_diffusivity_local.jl index ad85509c00..6388d67a98 100644 --- a/xin_kai_vertical_diffusivity_local.jl +++ b/xin_kai_vertical_diffusivity_local.jl @@ -77,14 +77,14 @@ function XinKaiLocalVerticalDiffusivity{TD}(ν₀ :: FT, end function XinKaiLocalVerticalDiffusivity(time_discretization = VerticallyImplicitTimeDiscretization(), - FT = Float64; - ν₀ = 1e-5, - νˢʰ = 0.04569735882746968, - νᶜⁿ = 0.47887785611155065, - Pr_convₜ = 0.1261854430705509, - Pr_shearₜ = 1.594794053970444, - Riᶜ = 0.9964350402840053, - δRi = 0.05635304878092709) + FT = Float64; + ν₀ = 1e-5, + νˢʰ = 0.04569735882746968, + νᶜⁿ = 0.47887785611155065, + Pr_convₜ = 0.1261854430705509, + Pr_shearₜ = 1.594794053970444, + Riᶜ = 0.9964350402840053, + δRi = 0.05635304878092709) TD = typeof(time_discretization) From e9ac435e3dc8a06620e14bcd50b29e4e501fe026 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Fri, 28 Jun 2024 11:31:06 +0100 Subject: [PATCH 012/114] nonlocal physical closure of vertical diffusivity --- xin_kai_vertical_diffusivity.jl | 255 ++++++++++++++++++++++++++++++++ 1 file changed, 255 insertions(+) create mode 100644 xin_kai_vertical_diffusivity.jl diff --git a/xin_kai_vertical_diffusivity.jl b/xin_kai_vertical_diffusivity.jl new file mode 100644 index 0000000000..fed133a780 --- /dev/null +++ b/xin_kai_vertical_diffusivity.jl @@ -0,0 +1,255 @@ +using Oceananigans +using Oceananigans.Architectures: architecture +using Oceananigans.BuoyancyModels: ∂z_b +using Oceananigans.Operators +using Oceananigans.BoundaryConditions +using Oceananigans.Grids: inactive_node +using Oceananigans.Operators: ℑzᵃᵃᶜ, ℑxyᶠᶠᵃ, ℑxyᶜᶜᵃ + +using Adapt + +using KernelAbstractions: @index, @kernel +using KernelAbstractions.Extras.LoopInfo: @unroll + +using Oceananigans.TurbulenceClosures: + tapering_factorᶠᶜᶜ, + tapering_factorᶜᶠᶜ, + tapering_factorᶜᶜᶠ, + tapering_factor, + SmallSlopeIsopycnalTensor, + AbstractScalarDiffusivity, + ExplicitTimeDiscretization, + FluxTapering, + isopycnal_rotation_tensor_xz_ccf, + isopycnal_rotation_tensor_yz_ccf, + isopycnal_rotation_tensor_zz_ccf + +import Oceananigans.TurbulenceClosures: + compute_diffusivities!, + DiffusivityFields, + viscosity, + diffusivity, + getclosure, + top_buoyancy_flux, + diffusive_flux_x, + diffusive_flux_y, + diffusive_flux_z, + viscous_flux_ux, + viscous_flux_vx, + viscous_flux_uy, + viscous_flux_vy + +using Oceananigans.Utils: launch! +using Oceananigans.Coriolis: fᶠᶠᵃ +using Oceananigans.Operators +using Oceananigans.BuoyancyModels: ∂x_b, ∂y_b, ∂z_b + +using Oceananigans.TurbulenceClosures +using Oceananigans.TurbulenceClosures: HorizontalFormulation, VerticalFormulation, AbstractScalarDiffusivity +using Oceananigans.TurbulenceClosures: AbstractScalarBiharmonicDiffusivity +using Oceananigans.Operators +using Oceananigans.Operators: Δxᶜᶜᶜ, Δyᶜᶜᶜ, ℑxyᶜᶜᵃ, ζ₃ᶠᶠᶜ, div_xyᶜᶜᶜ +using Oceananigans.Operators: Δx, Δy +using Oceananigans.Operators: ℑxyz + +using Oceananigans.Operators: ℑxyzᶜᶜᶠ, ℑyzᵃᶜᶠ, ℑxzᶜᵃᶠ, Δxᶜᶜᶜ, Δyᶜᶜᶜ + +struct XinKaiVerticalDiffusivity{TD, FT} <: AbstractScalarDiffusivity{TD, VerticalFormulation, 2} + ν₀ :: FT + νˢʰ :: FT + νᶜⁿ :: FT + Cᵉⁿ :: FT + Prₜ :: FT + Riᶜ :: FT + δRi :: FT + Q₀ :: FT + δQ :: FT +end + +function XinKaiVerticalDiffusivity{TD}(ν₀ :: FT, + νˢʰ :: FT, + νᶜⁿ :: FT, + Cᵉⁿ :: FT, + Prₜ :: FT, + Riᶜ :: FT, + δRi :: FT, + Q₀ :: FT, + δQ :: FT) where {TD, FT} + + return XinKaiVerticalDiffusivity{TD, FT}(ν₀, νˢʰ, νᶜⁿ, Cᵉⁿ, Prₜ, Riᶜ, δRi, Q₀, δQ) +end + +function XinKaiVerticalDiffusivity(time_discretization = VerticallyImplicitTimeDiscretization(), + FT = Float64; + ν₀ = 1e-5, + νˢʰ = 0.0885, + νᶜⁿ = 4.3668, + Cᵉⁿ = 0.2071, + Prₜ = 1.207, + Riᶜ = -0.21982, + δRi = 8.342e-4, + Q₀ = 0.08116, + δQ = 0.02622) + + TD = typeof(time_discretization) + + return XinKaiVerticalDiffusivity{TD}(convert(FT, ν₀), + convert(FT, νˢʰ), + convert(FT, νᶜⁿ), + convert(FT, Cᵉⁿ), + convert(FT, Prₜ), + convert(FT, Riᶜ), + convert(FT, δRi), + convert(FT, Q₀), + convert(FT, δQ)) +end + +XinKaiVerticalDiffusivity(FT::DataType; kw...) = + XinKaiVerticalDiffusivity(VerticallyImplicitTimeDiscretization(), FT; kw...) + +Adapt.adapt_structure(to, clo::XinKaiVerticalDiffusivity{TD, FT}) where {TD, FT} = + XinKaiVerticalDiffusivity{TD, FT}(clo.ν₀, clo.νˢʰ, clo.νᶜⁿ, clo.Cᵉⁿ, clo.Prₜ, clo.Riᶜ, clo.δRi, clo.Q₀, clo.δQ) + +##### +##### Diffusivity field utilities +##### + +const RBVD = XinKaiVerticalDiffusivity +const RBVDArray = AbstractArray{<:RBVD} +const FlavorOfXKVD = Union{RBVD, RBVDArray} +const c = Center() +const f = Face() + +@inline viscosity_location(::FlavorOfXKVD) = (c, c, f) +@inline diffusivity_location(::FlavorOfXKVD) = (c, c, f) + +@inline viscosity(::FlavorOfXKVD, diffusivities) = diffusivities.κᵘ +@inline diffusivity(::FlavorOfXKVD, diffusivities, id) = diffusivities.κᶜ + +with_tracers(tracers, closure::FlavorOfXKVD) = closure + +# Note: computing diffusivities at cell centers for now. +function DiffusivityFields(grid, tracer_names, bcs, closure::FlavorOfXKVD) + κᶜ = Field((Center, Center, Face), grid) + κᵘ = Field((Center, Center, Face), grid) + N² = Field((Center, Center, Face), grid) + Ri = Field((Center, Center, Face), grid) + return (; κᶜ, κᵘ, Ri, N²) +end + +function compute_diffusivities!(diffusivities, closure::FlavorOfXKVD, model; parameters = :xyz) + arch = model.architecture + grid = model.grid + clock = model.clock + tracers = model.tracers + buoyancy = model.buoyancy + velocities = model.velocities + top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) + + Nx_in, Ny_in, Nz_in = total_size(diffusivities.κᶜ) + ox_in, oy_in, oz_in = diffusivities.κᶜ.data.offsets + + kp = KernelParameters((Nx_in, Ny_in, Nz_in), (ox_in, oy_in, oz_in)) + + launch!(arch, grid, kp, compute_N²!, diffusivities, grid, closure, tracers, buoyancy) + launch!(arch, grid, kp, compute_ri_number!, diffusivities, grid, closure, velocities) + + # Use `only_local_halos` to ensure that no communication occurs during + # this call to fill_halo_regions! + fill_halo_regions!(diffusivities.Ri; only_local_halos=true) + + launch!(arch, grid, kp, + compute_xinkai_diffusivities!, + diffusivities, + grid, + closure, + velocities, + tracers, + buoyancy, + top_tracer_bcs, + clock) + + return nothing +end + +@inline ϕ²(i, j, k, grid, ϕ, args...) = ϕ(i, j, k, grid, args...)^2 + +@inline function shear_squaredᶜᶜᶠ(i, j, k, grid, velocities) + ∂z_u² = ℑxᶜᵃᵃ(i, j, k, grid, ϕ², ∂zᶠᶜᶠ, velocities.u) + ∂z_v² = ℑyᵃᶜᵃ(i, j, k, grid, ϕ², ∂zᶜᶠᶠ, velocities.v) + return ∂z_u² + ∂z_v² +end + +@inline function N²ᶜᶜᶠ(i, j, k, grid, buoyancy, tracers) + return ∂z_b(i, j, k, grid, buoyancy, tracers) +end + +@inline function Riᶜᶜᶠ(i, j, k, grid, velocities, diffusivities) + S² = shear_squaredᶜᶜᶠ(i, j, k, grid, velocities) + N² = diffusivities.N²[i, j, k] + Ri = N² / S² + + # Clip N² and avoid NaN + return ifelse(N² == 0, zero(grid), Ri) +end + +const c = Center() +const f = Face() + +@kernel function compute_N²!(diffusivities, grid, closure::FlavorOfXKVD, tracers, buoyancy) + i, j, k = @index(Global, NTuple) + @inbounds diffusivities.N²[i, j, k] = N²ᶜᶜᶠ(i, j, k, grid, buoyancy, tracers) +end + +@kernel function compute_ri_number!(diffusivities, grid, closure::FlavorOfXKVD, velocities) + i, j, k = @index(Global, NTuple) + @inbounds diffusivities.Ri[i, j, k] = Riᶜᶜᶠ(i, j, k, grid, velocities, diffusivities) +end + +@kernel function compute_xinkai_diffusivities!(diffusivities, grid, closure::FlavorOfXKVD, + velocities, tracers, buoyancy, tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + _compute_xinkai_diffusivities!(i, j, k, diffusivities, grid, closure, + velocities, tracers, buoyancy, tracer_bcs, clock) +end + +@inline function _compute_xinkai_diffusivities!(i, j, k, diffusivities, grid, closure, + velocities, tracers, buoyancy, tracer_bcs, clock) + + # Ensure this works with "ensembles" of closures, in addition to ordinary single closures + closure_ij = getclosure(i, j, closure) + + ν₀ = closure_ij.ν₀ + νˢʰ = closure_ij.νˢʰ + νᶜⁿ = closure_ij.νᶜⁿ + Cᵉⁿ = closure_ij.Cᵉⁿ + Prₜ = closure_ij.Prₜ + Riᶜ = closure_ij.Riᶜ + δRi = closure_ij.δRi + Q₀ = closure_ij.Q₀ + δQ = closure_ij.δQ + + Qᵇ = top_buoyancy_flux(i, j, grid, buoyancy, tracer_bcs, clock, merge(velocities, tracers)) + + # (Potentially) apply a horizontal filter to the Richardson number + Ri = ℑxyᶜᶜᵃ(i, j, k, grid, ℑxyᶠᶠᵃ, diffusivities.Ri) + Ri_above = ℑxyᶜᶜᵃ(i, j, k + 1, grid, ℑxyᶠᶠᵃ, diffusivities.Ri) + N² = ℑxyᶜᶜᵃ(i, j, k, grid, ℑxyᶠᶠᵃ, diffusivities.N²) + + # Conditions + convecting = Ri < 0 # applies regardless of Qᵇ + entraining = (Ri > 0) & (Ri_above < 0) & (Qᵇ > 0) + + # Convective adjustment diffusivity + ν_local = ifelse(convecting, - (νᶜⁿ - νˢʰ) / 2 * tanh(Ri / δRi) + νˢʰ, clamp(Riᶜ * Ri + νˢʰ + ν₀, ν₀, νˢʰ)) + + # Entrainment diffusivity + x = Qᵇ / (N² + 1e-11) + ν_nonlocal = ifelse(entraining, Cᵉⁿ * νᶜⁿ * 0.5 * (tanh((x - Q₀) / δQ) + 1), 0) + + # Update by averaging in time + @inbounds diffusivities.κᵘ[i, j, k] = ifelse(k <= 1 || k >= grid.Nz+1, 0, ν_local + ν_nonlocal) + @inbounds diffusivities.κᶜ[i, j, k] = ifelse(k <= 1 || k >= grid.Nz+1, 0, (ν_local + ν_nonlocal) / Prₜ) + + return nothing +end From 5e548da91b4bfa8b8b90e03c0e901f00afceeb32 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Fri, 28 Jun 2024 11:31:14 +0100 Subject: [PATCH 013/114] 2Pr nonlocal physical closure --- xin_kai_vertical_diffusivity_2Pr.jl | 266 ++++++++++++++++++++++++++++ 1 file changed, 266 insertions(+) create mode 100644 xin_kai_vertical_diffusivity_2Pr.jl diff --git a/xin_kai_vertical_diffusivity_2Pr.jl b/xin_kai_vertical_diffusivity_2Pr.jl new file mode 100644 index 0000000000..79a2765de9 --- /dev/null +++ b/xin_kai_vertical_diffusivity_2Pr.jl @@ -0,0 +1,266 @@ +using Oceananigans +using Oceananigans.Architectures: architecture +using Oceananigans.BuoyancyModels: ∂z_b +using Oceananigans.Operators +using Oceananigans.BoundaryConditions +using Oceananigans.Grids: inactive_node +using Oceananigans.Operators: ℑzᵃᵃᶜ, ℑxyᶠᶠᵃ, ℑxyᶜᶜᵃ + +using Adapt + +using KernelAbstractions: @index, @kernel +using KernelAbstractions.Extras.LoopInfo: @unroll + +using Oceananigans.TurbulenceClosures: + tapering_factorᶠᶜᶜ, + tapering_factorᶜᶠᶜ, + tapering_factorᶜᶜᶠ, + tapering_factor, + SmallSlopeIsopycnalTensor, + AbstractScalarDiffusivity, + ExplicitTimeDiscretization, + FluxTapering, + isopycnal_rotation_tensor_xz_ccf, + isopycnal_rotation_tensor_yz_ccf, + isopycnal_rotation_tensor_zz_ccf + +import Oceananigans.TurbulenceClosures: + compute_diffusivities!, + DiffusivityFields, + viscosity, + diffusivity, + getclosure, + top_buoyancy_flux, + diffusive_flux_x, + diffusive_flux_y, + diffusive_flux_z, + viscous_flux_ux, + viscous_flux_vx, + viscous_flux_uy, + viscous_flux_vy + +using Oceananigans.Utils: launch! +using Oceananigans.Coriolis: fᶠᶠᵃ +using Oceananigans.Operators +using Oceananigans.BuoyancyModels: ∂x_b, ∂y_b, ∂z_b + +using Oceananigans.TurbulenceClosures +using Oceananigans.TurbulenceClosures: HorizontalFormulation, VerticalFormulation, AbstractScalarDiffusivity +using Oceananigans.TurbulenceClosures: AbstractScalarBiharmonicDiffusivity +using Oceananigans.Operators +using Oceananigans.Operators: Δxᶜᶜᶜ, Δyᶜᶜᶜ, ℑxyᶜᶜᵃ, ζ₃ᶠᶠᶜ, div_xyᶜᶜᶜ +using Oceananigans.Operators: Δx, Δy +using Oceananigans.Operators: ℑxyz + +using Oceananigans.Operators: ℑxyzᶜᶜᶠ, ℑyzᵃᶜᶠ, ℑxzᶜᵃᶠ, Δxᶜᶜᶜ, Δyᶜᶜᶜ + +struct XinKaiVerticalDiffusivity{TD, FT} <: AbstractScalarDiffusivity{TD, VerticalFormulation, 2} + ν₀ :: FT + νˢʰ :: FT + νᶜⁿ :: FT + Cᵉⁿ :: FT + Pr_convₜ :: FT + Pr_shearₜ :: FT + Riᶜ :: FT + δRi :: FT + Q₀ :: FT + δQ :: FT +end + +function XinKaiVerticalDiffusivity{TD}(ν₀ :: FT, + νˢʰ :: FT, + νᶜⁿ :: FT, + Cᵉⁿ :: FT, + Pr_convₜ :: FT, + Pr_shearₜ :: FT, + Riᶜ :: FT, + δRi :: FT, + Q₀ :: FT, + δQ :: FT) where {TD, FT} + + return XinKaiVerticalDiffusivity{TD, FT}(ν₀, νˢʰ, νᶜⁿ, Cᵉⁿ, Pr_convₜ, Pr_shearₜ, Riᶜ, δRi, Q₀, δQ) +end + +function XinKaiVerticalDiffusivity(time_discretization = VerticallyImplicitTimeDiscretization(), + FT = Float64; + ν₀ = 1e-5, + νˢʰ = 0.07738088203341657, + νᶜⁿ = 0.533741914196933, + Cᵉⁿ = 0.5196272898085122, + Pr_convₜ = 0.01632117727992826, + Pr_shearₜ = 1.8499159986192901, + Riᶜ = 0.4923581673007292, + δRi = 0.00012455519496760374, + Q₀ = 0.048232078296680234, + δQ = 0.01884938627051353) + + TD = typeof(time_discretization) + + return XinKaiVerticalDiffusivity{TD}(convert(FT, ν₀), + convert(FT, νˢʰ), + convert(FT, νᶜⁿ), + convert(FT, Cᵉⁿ), + convert(FT, Pr_convₜ), + convert(FT, Pr_shearₜ), + convert(FT, Riᶜ), + convert(FT, δRi), + convert(FT, Q₀), + convert(FT, δQ)) +end + +XinKaiVerticalDiffusivity(FT::DataType; kw...) = + XinKaiVerticalDiffusivity(VerticallyImplicitTimeDiscretization(), FT; kw...) + +Adapt.adapt_structure(to, clo::XinKaiVerticalDiffusivity{TD, FT}) where {TD, FT} = + XinKaiVerticalDiffusivity{TD, FT}(clo.ν₀, clo.νˢʰ, clo.νᶜⁿ, clo.Cᵉⁿ, clo.Pr_convₜ, clo.Pr_shearₜ, clo.Riᶜ, clo.δRi, clo.Q₀, clo.δQ) + +##### +##### Diffusivity field utilities +##### + +const RBVD = XinKaiVerticalDiffusivity +const RBVDArray = AbstractArray{<:RBVD} +const FlavorOfXKVD = Union{RBVD, RBVDArray} +const c = Center() +const f = Face() + +@inline viscosity_location(::FlavorOfXKVD) = (c, c, f) +@inline diffusivity_location(::FlavorOfXKVD) = (c, c, f) + +@inline viscosity(::FlavorOfXKVD, diffusivities) = diffusivities.κᵘ +@inline diffusivity(::FlavorOfXKVD, diffusivities, id) = diffusivities.κᶜ + +with_tracers(tracers, closure::FlavorOfXKVD) = closure + +# Note: computing diffusivities at cell centers for now. +function DiffusivityFields(grid, tracer_names, bcs, closure::FlavorOfXKVD) + κᶜ = Field((Center, Center, Face), grid) + κᵘ = Field((Center, Center, Face), grid) + N² = Field((Center, Center, Face), grid) + Ri = Field((Center, Center, Face), grid) + return (; κᶜ, κᵘ, Ri, N²) +end + +function compute_diffusivities!(diffusivities, closure::FlavorOfXKVD, model; parameters = :xyz) + arch = model.architecture + grid = model.grid + clock = model.clock + tracers = model.tracers + buoyancy = model.buoyancy + velocities = model.velocities + top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) + + Nx_in, Ny_in, Nz_in = total_size(diffusivities.κᶜ) + ox_in, oy_in, oz_in = diffusivities.κᶜ.data.offsets + + kp = KernelParameters((Nx_in, Ny_in, Nz_in), (ox_in, oy_in, oz_in)) + + launch!(arch, grid, kp, compute_N²!, diffusivities, grid, closure, tracers, buoyancy) + launch!(arch, grid, kp, compute_ri_number!, diffusivities, grid, closure, velocities) + + # Use `only_local_halos` to ensure that no communication occurs during + # this call to fill_halo_regions! + fill_halo_regions!(diffusivities.Ri; only_local_halos=true) + + launch!(arch, grid, kp, + compute_xinkai_diffusivities!, + diffusivities, + grid, + closure, + velocities, + tracers, + buoyancy, + top_tracer_bcs, + clock) + + return nothing +end + +@inline ϕ²(i, j, k, grid, ϕ, args...) = ϕ(i, j, k, grid, args...)^2 + +@inline function shear_squaredᶜᶜᶠ(i, j, k, grid, velocities) + ∂z_u² = ℑxᶜᵃᵃ(i, j, k, grid, ϕ², ∂zᶠᶜᶠ, velocities.u) + ∂z_v² = ℑyᵃᶜᵃ(i, j, k, grid, ϕ², ∂zᶜᶠᶠ, velocities.v) + return ∂z_u² + ∂z_v² +end + +@inline function N²ᶜᶜᶠ(i, j, k, grid, buoyancy, tracers) + return ∂z_b(i, j, k, grid, buoyancy, tracers) +end + +@inline function Riᶜᶜᶠ(i, j, k, grid, velocities, diffusivities) + S² = shear_squaredᶜᶜᶠ(i, j, k, grid, velocities) + N² = diffusivities.N²[i, j, k] + Ri = N² / S² + + # Clip N² and avoid NaN + return ifelse(N² == 0, zero(grid), Ri) +end + +const c = Center() +const f = Face() + +@kernel function compute_N²!(diffusivities, grid, closure::FlavorOfXKVD, tracers, buoyancy) + i, j, k = @index(Global, NTuple) + @inbounds diffusivities.N²[i, j, k] = N²ᶜᶜᶠ(i, j, k, grid, buoyancy, tracers) +end + +@kernel function compute_ri_number!(diffusivities, grid, closure::FlavorOfXKVD, velocities) + i, j, k = @index(Global, NTuple) + @inbounds diffusivities.Ri[i, j, k] = Riᶜᶜᶠ(i, j, k, grid, velocities, diffusivities) +end + +@kernel function compute_xinkai_diffusivities!(diffusivities, grid, closure::FlavorOfXKVD, + velocities, tracers, buoyancy, tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + _compute_xinkai_diffusivities!(i, j, k, diffusivities, grid, closure, + velocities, tracers, buoyancy, tracer_bcs, clock) +end + +@inline function _compute_xinkai_diffusivities!(i, j, k, diffusivities, grid, closure, + velocities, tracers, buoyancy, tracer_bcs, clock) + + # Ensure this works with "ensembles" of closures, in addition to ordinary single closures + closure_ij = getclosure(i, j, closure) + + ν₀ = closure_ij.ν₀ + νˢʰ = closure_ij.νˢʰ + νᶜⁿ = closure_ij.νᶜⁿ + Cᵉⁿ = closure_ij.Cᵉⁿ + Pr_convₜ = closure_ij.Pr_convₜ + Pr_shearₜ = closure_ij.Pr_shearₜ + Riᶜ = closure_ij.Riᶜ + δRi = closure_ij.δRi + Q₀ = closure_ij.Q₀ + δQ = closure_ij.δQ + + κ₀ = ν₀ / Pr_shearₜ + κˢʰ = νˢʰ / Pr_shearₜ + κᶜⁿ = νᶜⁿ / Pr_convₜ + + Qᵇ = top_buoyancy_flux(i, j, grid, buoyancy, tracer_bcs, clock, merge(velocities, tracers)) + + # (Potentially) apply a horizontal filter to the Richardson number + Ri = ℑxyᶜᶜᵃ(i, j, k, grid, ℑxyᶠᶠᵃ, diffusivities.Ri) + Ri_above = ℑxyᶜᶜᵃ(i, j, k + 1, grid, ℑxyᶠᶠᵃ, diffusivities.Ri) + N² = ℑxyᶜᶜᵃ(i, j, k, grid, ℑxyᶠᶠᵃ, diffusivities.N²) + + # Conditions + convecting = Ri < 0 # applies regardless of Qᵇ + entraining = (Ri > 0) & (Ri_above < 0) & (Qᵇ > 0) + + # Convective adjustment diffusivity + ν_local = ifelse(convecting, (νˢʰ - νᶜⁿ) * tanh(Ri / δRi) + νˢʰ, clamp((ν₀ - νˢʰ) * Ri / Riᶜ + νˢʰ, ν₀, νˢʰ)) + κ_local = ifelse(convecting, (κˢʰ - κᶜⁿ) * tanh(Ri / δRi) + κˢʰ, clamp((κ₀ - κˢʰ) * Ri / Riᶜ + κˢʰ, κ₀, κˢʰ)) + + # Entrainment diffusivity + x = Qᵇ / (N² + 1e-11) + ν_nonlocal = ifelse(entraining, Cᵉⁿ * νᶜⁿ * 0.5 * (tanh((x - Q₀) / δQ) + 1), 0) + κ_nonlocal = ifelse(entraining, ν_nonlocal / Pr_shearₜ, 0) + + # Update by averaging in time + @inbounds diffusivities.κᵘ[i, j, k] = ifelse(k <= 1 || k >= grid.Nz+1, 0, ν_local + ν_nonlocal) + @inbounds diffusivities.κᶜ[i, j, k] = ifelse(k <= 1 || k >= grid.Nz+1, 0, κ_local + κ_nonlocal) + + return nothing +end From 772c2358fe491d66047a1d956aeee743dca0ad4c Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Fri, 28 Jun 2024 14:41:54 +0100 Subject: [PATCH 014/114] working GPU version of NN closure with scaling and correction --- NN_closure_global.jl | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/NN_closure_global.jl b/NN_closure_global.jl index ddebe03cc0..001dc5700f 100644 --- a/NN_closure_global.jl +++ b/NN_closure_global.jl @@ -72,8 +72,8 @@ function NNFluxClosure(arch) nn_path = "./NDE_FC_Qb_15simnew_2layer_64_relu_2Pr_model.jld2" ps, sts, scaling_params, wT_model, wS_model = jldopen(nn_path, "r") do file - ps = file["u"] |> dev - sts = file["sts"] |> dev + ps = file["u"] |> dev |> f64 + sts = file["sts"] |> dev |> f64 scaling_params = file["scaling"] wT_model = file["model"].wT wS_model = file["model"].wS @@ -82,8 +82,6 @@ function NNFluxClosure(arch) scaling = construct_zeromeanunitvariance_scaling(scaling_params) - # ps = ps .= 0 - wT_NN = NN(wT_model, ps.wT, sts.wT) wS_NN = NN(wS_model, ps.wS, sts.wS) @@ -127,11 +125,12 @@ function compute_diffusivities!(diffusivities, closure::NNFluxClosure, model; pa wT.data.parent .= dropdims(closure.wT(input.parent), dims=1) wS.data.parent .= dropdims(closure.wS(input.parent), dims=1) - launch!(arch, grid, kp, _rescale_nn_fluxes!, diffusivities, grid, closure) + launch!(arch, grid, kp, _rescale_nn_fluxes!, diffusivities, grid, closure, tracers, buoyancy, top_tracer_bcs, clock) + launch!(arch, grid, kp, _adjust_nn_bottom_fluxes!, diffusivities, grid, closure) return nothing end -@kernel function _populate_input!(input, grid, closure, tracers, velocities, buoyancy, coriolis, top_tracer_bcs, clock) +@kernel function _populate_input!(input, grid, closure::NNFluxClosure, tracers, velocities, buoyancy, coriolis, top_tracer_bcs, clock) i, j, k = @index(Global, NTuple) scaling = closure.scaling @@ -155,11 +154,20 @@ end @inbounds input[11, i, j, k] = fᶜᶜ = scaling.f(fᶜᶜᵃ(i, j, k, grid, coriolis)) end -@kernel function _rescale_nn_fluxes!(diffusivities, grid, closure) +@kernel function _rescale_nn_fluxes!(diffusivities, grid, closure::NNFluxClosure, tracers, buoyancy, top_tracer_bcs, clock) i, j, k = @index(Global, NTuple) scaling = closure.scaling - @inbounds diffusivities.wT[i, j, k] = ifelse(k >= grid.Nz - 1 || k <= 3, inv(scaling.wT)(0), inv(scaling.wT)(diffusivities.wT[i, j, k])) - @inbounds diffusivities.wS[i, j, k] = ifelse(k >= grid.Nz - 1 || k <= 3, inv(scaling.wS)(0), inv(scaling.wS)(diffusivities.wS[i, j, k])) + convecting = top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, tracers) > 0 + interior_point = k <= grid.Nz - 1 & k >= 2 + + @inbounds diffusivities.wT[i, j, k] = ifelse(convecting & interior_point, inv(scaling.wT)(diffusivities.wT[i, j, k]) - inv(scaling.wT)(0), 0) + @inbounds diffusivities.wS[i, j, k] = ifelse(convecting & interior_point, inv(scaling.wS)(diffusivities.wS[i, j, k]) - inv(scaling.wS)(0), 0) +end + +@kernel function _adjust_nn_bottom_fluxes!(diffusivities, grid, closure::NNFluxClosure) + i, j, k = @index(Global, NTuple) + @inbounds diffusivities.wT[i, j, k] = ifelse(k <= 3, diffusivities.wT[i, j, 4], diffusivities.wT[i, j, k]) + @inbounds diffusivities.wS[i, j, k] = ifelse(k <= 3, diffusivities.wS[i, j, 4], diffusivities.wS[i, j, k]) end # Write here your constructor From 9c6873fccb914ba88b4c734946c43b1ccc46a746 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Mon, 1 Jul 2024 21:28:06 +0100 Subject: [PATCH 015/114] validation script for oceananigans NN implementation --- validate_NN_1D_model.jl | 204 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 204 insertions(+) create mode 100644 validate_NN_1D_model.jl diff --git a/validate_NN_1D_model.jl b/validate_NN_1D_model.jl new file mode 100644 index 0000000000..115342f5cc --- /dev/null +++ b/validate_NN_1D_model.jl @@ -0,0 +1,204 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("NN_closure_global.jl") +include("xin_kai_vertical_diffusivity_local.jl") +include("feature_scaling.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +import SeawaterPolynomials.TEOS10: s, ΔS, Sₐᵤ +s(Sᴬ::Number) = Sᴬ + ΔS >= 0 ? √((Sᴬ + ΔS) / Sₐᵤ) : NaN + +# Architecture +model_architecture = CPU() + +file = jldopen("model_inference_run.jld2", "r") + +# number of grid points +const Nz = file["Nz"] +const Lz = file["Lz"] + +grid = RectilinearGrid(model_architecture, + topology = (Flat, Flat, Bounded), + size = Nz, + halo = 3, + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const dTdz = file["dTdz"] +const dSdz = file["dSdz"] + +const T_surface = file["T_surface"] +const S_surface = file["S_surface"] + +T_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(file["wT_top"])) +S_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(file["wS_top"])) + +##### +##### Coriolis +##### + +const f₀ = file["f₀"] +coriolis = FPlane(f=f₀) + +##### +##### Forcing and initial condition +##### +T_initial(z) = dTdz * z + T_surface +S_initial(z) = dSdz * z + S_surface + +nn_closure = NNFluxClosure(model_architecture) +base_closure = XinKaiLocalVerticalDiffusivity() + +##### +##### Model building +##### + +@info "Building a model..." + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = ImplicitFreeSurface(), + momentum_advection = WENO(grid = grid), + tracer_advection = WENO(grid = grid), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = (nn_closure, base_closure), + # closure = base_closure, + tracers = (:T, :S), + boundary_conditions = (; T = T_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(z) = T_initial(z) + 1e-6 * noise(z) +S_initial_noisy(z) = S_initial(z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = file["Δt"] +stop_time = file["τ"] + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(20)) + +##### +##### Diagnostics +##### + +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S + +Tbar = Field(Average(T, dims = (1,2))) +Sbar = Field(Average(S, dims = (1,2))) + +averaged_outputs = (; Tbar, Sbar) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:jld2] = JLD2OutputWriter(model, averaged_outputs, + filename = "NN_1D_channel_averages", + schedule = TimeInterval(Δt₀), + overwrite_existing = true) + +@info "Running the simulation..." + +try + run!(simulation, pickup = false) +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +# ##### +# ##### Visualization +# ##### +#%% +using GLMakie + +Tbar_data = FieldTimeSeries("./NN_1D_channel_averages.jld2", "Tbar") +Sbar_data = FieldTimeSeries("./NN_1D_channel_averages.jld2", "Sbar") + +zC = znodes(Tbar_data.grid, Center()) +zF = znodes(Tbar_data.grid, Face()) + +Nt = length(Tbar_data.times) + +fig = Figure(size = (900, 600)) +axT = GLMakie.Axis(fig[1, 1], xlabel = "T (°C)", ylabel = "z (m)") +axS = GLMakie.Axis(fig[1, 2], xlabel = "S (g kg⁻¹)", ylabel = "z (m)") +slider = Slider(fig[2, :], range=1:Nt) +n = slider.value + +Tbarₙ = @lift interior(Tbar_data[$n], 1, 1, :) +Sbarₙ = @lift interior(Sbar_data[$n], 1, 1, :) + +Tbar_truthₙ = @lift file["sol_T"][:, $n] +Sbar_truthₙ = @lift file["sol_S"][:, $n] + +title_str = @lift "Time: $(round(Tbar_data.times[$n] / 86400, digits=3)) days" + +lines!(axT, Tbarₙ, zC, label="Oceananigans") +lines!(axS, Sbarₙ, zC, label="Oceananigans") + +lines!(axT, Tbar_truthₙ, zC, label="Truth") +lines!(axS, Sbar_truthₙ, zC, label="Truth") + +axislegend(axT, position = :lb) +Label(fig[0, :], title_str, tellwidth = false) + +# CairoMakie.record(fig, "./NN_1D_fields.mp4", 1:Nt, framerate=10) do nn +# n[] = nn +# xlims!(axT, nothing, nothing) +# xlims!(axS, nothing, nothing) +# end + +display(fig) +#%% \ No newline at end of file From f6d508ad7caecbf672f44dd19a722c1d7f093d74 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Mon, 1 Jul 2024 21:28:24 +0100 Subject: [PATCH 016/114] close file and record video of validation --- validate_NN_1D_model.jl | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/validate_NN_1D_model.jl b/validate_NN_1D_model.jl index 115342f5cc..f2df7292b8 100644 --- a/validate_NN_1D_model.jl +++ b/validate_NN_1D_model.jl @@ -194,11 +194,11 @@ lines!(axS, Sbar_truthₙ, zC, label="Truth") axislegend(axT, position = :lb) Label(fig[0, :], title_str, tellwidth = false) -# CairoMakie.record(fig, "./NN_1D_fields.mp4", 1:Nt, framerate=10) do nn -# n[] = nn -# xlims!(axT, nothing, nothing) -# xlims!(axS, nothing, nothing) -# end +GLMakie.record(fig, "./NN_1D_validation.mp4", 1:Nt, framerate=60, px_per_unit=4) do nn + @info nn + n[] = nn +end display(fig) -#%% \ No newline at end of file +#%% +close(file) \ No newline at end of file From 1faa5a69d803668973c4dd8b365f63a77c8780fe Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 2 Jul 2024 12:36:16 +0100 Subject: [PATCH 017/114] Update Project.toml with new package dependencies --- Manifest.toml | 411 ++++++++++++++++++++++++++++++++++++++++++-------- Project.toml | 7 +- 2 files changed, 353 insertions(+), 65 deletions(-) diff --git a/Manifest.toml b/Manifest.toml index 961740127e..6da492bf6a 100644 --- a/Manifest.toml +++ b/Manifest.toml @@ -1,23 +1,30 @@ # This file is machine-generated - editing it directly is not advised -julia_version = "1.10.3" +julia_version = "1.10.1" manifest_format = "2.0" -project_hash = "04d395caf937b0921325a77873167e8baa293a99" +project_hash = "f43673c2cec178a7eac57c04f7c0f97bcf64e09b" + +[[deps.ADTypes]] +git-tree-sha1 = "fa0822e5baee6e23081c2685ae27265dabee23d8" +uuid = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +version = "1.4.0" +weakdeps = ["ChainRulesCore", "EnzymeCore"] + + [deps.ADTypes.extensions] + ADTypesChainRulesCoreExt = "ChainRulesCore" + ADTypesEnzymeCoreExt = "EnzymeCore" [[deps.AbstractFFTs]] deps = ["LinearAlgebra"] git-tree-sha1 = "d92ad398961a3ed262d8bf04a1a2b8340f915fef" uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c" version = "1.5.0" +weakdeps = ["ChainRulesCore", "Test"] [deps.AbstractFFTs.extensions] AbstractFFTsChainRulesCoreExt = "ChainRulesCore" AbstractFFTsTestExt = "Test" - [deps.AbstractFFTs.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" - [[deps.Adapt]] deps = ["LinearAlgebra", "Requires"] git-tree-sha1 = "6a55b747d1812e699320963ffde36f1ebdda4099" @@ -28,15 +35,20 @@ weakdeps = ["StaticArrays"] [deps.Adapt.extensions] AdaptStaticArraysExt = "StaticArrays" +[[deps.ArgCheck]] +git-tree-sha1 = "a3a402a35a2f7e0b87828ccabbd5ebfbebe356b4" +uuid = "dce04be8-c92d-5529-be00-80e4d2c0e197" +version = "2.3.0" + [[deps.ArgTools]] uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f" version = "1.1.1" [[deps.ArrayInterface]] deps = ["Adapt", "LinearAlgebra", "SparseArrays", "SuiteSparse"] -git-tree-sha1 = "133a240faec6e074e07c31ee75619c90544179cf" +git-tree-sha1 = "ed2ec3c9b483842ae59cd273834e5b46206d6dda" uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" -version = "7.10.0" +version = "7.11.0" [deps.ArrayInterface.extensions] ArrayInterfaceBandedMatricesExt = "BandedMatrices" @@ -78,6 +90,12 @@ version = "0.5.0" [[deps.Base64]] uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f" +[[deps.BitTwiddlingConvenienceFunctions]] +deps = ["Static"] +git-tree-sha1 = "0c5f81f47bbbcf4aea7b2959135713459170798b" +uuid = "62783981-4cbd-42fc-bca8-16325de8dc4b" +version = "0.1.5" + [[deps.Blosc_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Lz4_jll", "Zlib_jll", "Zstd_jll"] git-tree-sha1 = "19b98ee7e3db3b4eff74c5c9c72bf32144e24f10" @@ -101,39 +119,63 @@ git-tree-sha1 = "5afb5c5ba2688ca43a9ad2e5a91cbb93921ccfa1" uuid = "179af706-886a-5703-950a-314cd64e0468" version = "0.1.3" +[[deps.CPUSummary]] +deps = ["CpuId", "IfElse", "PrecompileTools", "Static"] +git-tree-sha1 = "585a387a490f1c4bd88be67eea15b93da5e85db7" +uuid = "2a0fbf3d-bb9c-48f3-b0a9-814d99fd7ab9" +version = "0.2.5" + [[deps.CUDA]] deps = ["AbstractFFTs", "Adapt", "BFloat16s", "CEnum", "CUDA_Driver_jll", "CUDA_Runtime_Discovery", "CUDA_Runtime_jll", "Crayons", "DataFrames", "ExprTools", "GPUArrays", "GPUCompiler", "KernelAbstractions", "LLVM", "LLVMLoopInfo", "LazyArtifacts", "Libdl", "LinearAlgebra", "Logging", "NVTX", "Preferences", "PrettyTables", "Printf", "Random", "Random123", "RandomNumbers", "Reexport", "Requires", "SparseArrays", "StaticArrays", "Statistics"] -git-tree-sha1 = "abc3c845165c2d5c03ab61754a90c7f7cff0f6a4" +git-tree-sha1 = "b8c28cb78014f7ae81a652ce1524cba7667dea5c" uuid = "052768ef-5323-5732-b1bb-66c8b64840ba" -version = "5.4.1" +version = "5.3.5" +weakdeps = ["ChainRulesCore", "EnzymeCore", "SpecialFunctions"] [deps.CUDA.extensions] ChainRulesCoreExt = "ChainRulesCore" EnzymeCoreExt = "EnzymeCore" SpecialFunctionsExt = "SpecialFunctions" - [deps.CUDA.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" - SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" - [[deps.CUDA_Driver_jll]] deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "Pkg"] -git-tree-sha1 = "c48f9da18efd43b6b7adb7ee1f93fe5f2926c339" +git-tree-sha1 = "dc172b558adbf17952001e15cf0d6364e6d78c2f" uuid = "4ee394cb-3365-5eb0-8335-949819d2adfc" -version = "0.9.0+0" +version = "0.8.1+0" [[deps.CUDA_Runtime_Discovery]] deps = ["Libdl"] -git-tree-sha1 = "5db9da5fdeaa708c22ba86b82c49528f402497f2" +git-tree-sha1 = "38f830504358e9972d2a0c3e5d51cb865e0733df" uuid = "1af6417a-86b4-443c-805f-a4643ffb695f" -version = "0.3.3" +version = "0.2.4" [[deps.CUDA_Runtime_jll]] deps = ["Artifacts", "CUDA_Driver_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"] -git-tree-sha1 = "bcba305388e16aa5c879e896726db9e71b4942c6" +git-tree-sha1 = "4ca7d6d92075906c2ce871ea8bba971fff20d00c" uuid = "76a88914-d11a-5bdc-97e0-2f5a05c973a2" -version = "0.14.0+1" +version = "0.12.1+0" + +[[deps.CUDNN_jll]] +deps = ["Artifacts", "CUDA_Runtime_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"] +git-tree-sha1 = "cbf7d75f8c58b147bdf6acea2e5bc96cececa6d4" +uuid = "62b44479-cb7b-5706-934f-f13b2eb2e645" +version = "9.0.0+1" + +[[deps.ChainRulesCore]] +deps = ["Compat", "LinearAlgebra"] +git-tree-sha1 = "71acdbf594aab5bbb2cec89b208c41b4c411e49f" +uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +version = "1.24.0" +weakdeps = ["SparseArrays"] + + [deps.ChainRulesCore.extensions] + ChainRulesCoreSparseArraysExt = "SparseArrays" + +[[deps.CloseOpenIntervals]] +deps = ["Static", "StaticArrayInterface"] +git-tree-sha1 = "70232f82ffaab9dc52585e0dd043b5e0c6b714f1" +uuid = "fb6a15b2-703c-40df-9091-08a04967cfa9" +version = "0.1.12" [[deps.ColorTypes]] deps = ["FixedPointNumbers", "Random"] @@ -166,7 +208,12 @@ weakdeps = ["Dates", "LinearAlgebra"] [[deps.CompilerSupportLibraries_jll]] deps = ["Artifacts", "Libdl"] uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae" -version = "1.1.1+0" +version = "1.1.0+0" + +[[deps.ConcreteStructs]] +git-tree-sha1 = "f749037478283d372048690eb3b5f92a79432b34" +uuid = "2569d6c7-a4a2-43d3-a901-331e8e4be471" +version = "0.2.3" [[deps.ConstructionBase]] deps = ["LinearAlgebra"] @@ -182,6 +229,12 @@ version = "1.5.5" IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +[[deps.CpuId]] +deps = ["Markdown"] +git-tree-sha1 = "fcbb72b032692610bfbdb15018ac16a36cf2e406" +uuid = "adafc99b-e345-5852-983c-f28acb93d879" +version = "0.3.1" + [[deps.Crayons]] git-tree-sha1 = "249fe38abf76d48563e2f4556bebd215aa317e15" uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f" @@ -230,15 +283,12 @@ deps = ["LinearAlgebra", "Statistics", "StatsAPI"] git-tree-sha1 = "66c4c81f259586e8f002eacebc177e1fb06363b0" uuid = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7" version = "0.10.11" +weakdeps = ["ChainRulesCore", "SparseArrays"] [deps.Distances.extensions] DistancesChainRulesCoreExt = "ChainRulesCore" DistancesSparseArraysExt = "SparseArrays" - [deps.Distances.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" - [[deps.Distributed]] deps = ["Random", "Serialization", "Sockets"] uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b" @@ -259,6 +309,15 @@ git-tree-sha1 = "71c79e77221ab3a29918aaf6db4f217b89138608" uuid = "b305315f-e792-5b7a-8f41-49f472929428" version = "1.0.1" +[[deps.EnzymeCore]] +git-tree-sha1 = "88bc63137eb033acc3afe1b9875717889c718c46" +uuid = "f151be2c-9106-41f4-ab19-57ee4f262869" +version = "0.7.5" +weakdeps = ["Adapt"] + + [deps.EnzymeCore.extensions] + AdaptExt = "Adapt" + [[deps.ExprTools]] git-tree-sha1 = "27415f162e6028e81c72b82ef756bf321213b6ec" uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04" @@ -276,6 +335,17 @@ git-tree-sha1 = "c6033cc3892d0ef5bb9cd29b7f2f0331ea5184ea" uuid = "f5851436-0d7a-5f13-b9de-f02708fd171a" version = "3.3.10+0" +[[deps.FastBroadcast]] +deps = ["ArrayInterface", "LinearAlgebra", "Polyester", "Static", "StaticArrayInterface", "StrideArraysCore"] +git-tree-sha1 = "2be93e36303143c6fffd07e2222bbade35194d9e" +uuid = "7034ab61-46d4-4ed7-9d0f-46aef9175898" +version = "0.3.3" + +[[deps.FastClosures]] +git-tree-sha1 = "acebe244d53ee1b461970f8910c235b259e772ef" +uuid = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" +version = "0.3.2" + [[deps.FileIO]] deps = ["Pkg", "Requires", "UUIDs"] git-tree-sha1 = "82d8afa92ecf4b52d78d869f038ebfb881267322" @@ -291,6 +361,12 @@ git-tree-sha1 = "05882d6995ae5c12bb5f36dd2ed3f61c98cbb172" uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93" version = "0.8.5" +[[deps.Functors]] +deps = ["LinearAlgebra"] +git-tree-sha1 = "8a66c07630d6428eaab3506a0eabfcf4a9edea05" +uuid = "d9f16b24-f501-4c13-a1f2-28368ffc5196" +version = "0.4.11" + [[deps.Future]] deps = ["Random"] uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820" @@ -302,9 +378,9 @@ version = "6.2.1+6" [[deps.GPUArrays]] deps = ["Adapt", "GPUArraysCore", "LLVM", "LinearAlgebra", "Printf", "Random", "Reexport", "Serialization", "Statistics"] -git-tree-sha1 = "38cb19b8a3e600e509dc36a6396ac74266d108c1" +git-tree-sha1 = "c154546e322a9c73364e8a60430b0f79b812d320" uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" -version = "10.1.1" +version = "10.2.0" [[deps.GPUArraysCore]] deps = ["Adapt"] @@ -314,9 +390,9 @@ version = "0.1.6" [[deps.GPUCompiler]] deps = ["ExprTools", "InteractiveUtils", "LLVM", "Libdl", "Logging", "Scratch", "TimerOutputs", "UUIDs"] -git-tree-sha1 = "518ebd058c9895de468a8c255797b0c53fdb44dd" +git-tree-sha1 = "1600477fba37c9fc067b9be21f5e8101f24a8865" uuid = "61eb1bfa-7361-4325-ad38-22787b887f55" -version = "0.26.5" +version = "0.26.4" [[deps.Glob]] git-tree-sha1 = "97285bbd5230dd766e9ef6749b80fc617126d496" @@ -354,9 +430,15 @@ version = "0.2.1" [[deps.InlineStrings]] deps = ["Parsers"] -git-tree-sha1 = "9cc2baf75c6d09f9da536ddf58eb2f29dedaf461" +git-tree-sha1 = "86356004f30f8e737eff143d57d41bd580e437aa" uuid = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48" -version = "1.4.0" +version = "1.4.1" + + [deps.InlineStrings.extensions] + ArrowTypesExt = "ArrowTypes" + + [deps.InlineStrings.weakdeps] + ArrowTypes = "31f734f8-188a-4ce0-8406-c8a06bd891cd" [[deps.IntelOpenMP_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] @@ -421,21 +503,19 @@ version = "0.2.1+0" [[deps.KernelAbstractions]] deps = ["Adapt", "Atomix", "InteractiveUtils", "LinearAlgebra", "MacroTools", "PrecompileTools", "Requires", "SparseArrays", "StaticArrays", "UUIDs", "UnsafeAtomics", "UnsafeAtomicsLLVM"] -git-tree-sha1 = "db02395e4c374030c53dc28f3c1d33dec35f7272" +git-tree-sha1 = "ed7167240f40e62d97c1f5f7735dea6de3cc5c49" uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c" -version = "0.9.19" +version = "0.9.18" +weakdeps = ["EnzymeCore"] [deps.KernelAbstractions.extensions] EnzymeExt = "EnzymeCore" - [deps.KernelAbstractions.weakdeps] - EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" - [[deps.LLVM]] deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Preferences", "Printf", "Requires", "Unicode"] -git-tree-sha1 = "065c36f95709dd4a676dc6839a35d6fa6f192f24" +git-tree-sha1 = "839c82932db86740ae729779e610f07a1640be9a" uuid = "929cbde3-209d-540e-8aea-75f648917ca0" -version = "7.1.0" +version = "6.6.3" weakdeps = ["BFloat16s"] [deps.LLVM.extensions] @@ -466,6 +546,12 @@ git-tree-sha1 = "50901ebc375ed41dbf8058da26f9de442febbbec" uuid = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f" version = "1.3.1" +[[deps.LayoutPointers]] +deps = ["ArrayInterface", "LinearAlgebra", "ManualMemory", "SIMDTypes", "Static", "StaticArrayInterface"] +git-tree-sha1 = "62edfee3211981241b57ff1cedf4d74d79519277" +uuid = "10f19ff3-798f-405d-979b-55457f8fc047" +version = "0.1.15" + [[deps.LazyArtifacts]] deps = ["Artifacts", "Pkg"] uuid = "4af54fe1-eca0-43a8-85a7-787d91b784e3" @@ -509,9 +595,9 @@ uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" [[deps.LogExpFunctions]] deps = ["DocStringExtensions", "IrrationalConstants", "LinearAlgebra"] -git-tree-sha1 = "18144f3e9cbe9b15b070288eef858f71b291ce37" +git-tree-sha1 = "a2d09619db4e765091ee5c6ffe8872849de0feea" uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688" -version = "0.3.27" +version = "0.3.28" [deps.LogExpFunctions.extensions] LogExpFunctionsChainRulesCoreExt = "ChainRulesCore" @@ -526,6 +612,130 @@ version = "0.3.27" [[deps.Logging]] uuid = "56ddb016-857b-54e1-b83d-db4d58db5568" +[[deps.LossFunctions]] +deps = ["Markdown", "Requires", "Statistics"] +git-tree-sha1 = "df9da07efb9b05ca7ef701acec891ee8f73c99e2" +uuid = "30fc2ffe-d236-52d8-8643-a9d8f7c094a7" +version = "0.11.1" + + [deps.LossFunctions.extensions] + LossFunctionsCategoricalArraysExt = "CategoricalArrays" + + [deps.LossFunctions.weakdeps] + CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" + +[[deps.Lux]] +deps = ["ADTypes", "Adapt", "ArgCheck", "ArrayInterface", "ChainRulesCore", "Compat", "ConcreteStructs", "ConstructionBase", "EnzymeCore", "FastClosures", "Functors", "GPUArraysCore", "LinearAlgebra", "LossFunctions", "LuxCore", "LuxDeviceUtils", "LuxLib", "MacroTools", "Markdown", "PrecompileTools", "Preferences", "Random", "Reexport", "Setfield", "Statistics", "WeightInitializers"] +git-tree-sha1 = "c889341b5ab4bbeb59ea5d693eba43fb44dc40d6" +repo-rev = "ap/cuda_kernel" +repo-url = "https://github.com/LuxDL/Lux.jl.git" +uuid = "b2108857-7c20-44ae-9111-449ecde12c47" +version = "0.5.57" + + [deps.Lux.extensions] + LuxAMDGPUExt = "AMDGPU" + LuxComponentArraysExt = "ComponentArrays" + LuxDynamicExpressionsExt = "DynamicExpressions" + LuxDynamicExpressionsForwardDiffExt = ["DynamicExpressions", "ForwardDiff"] + LuxEnzymeExt = "Enzyme" + LuxFluxExt = "Flux" + LuxForwardDiffExt = "ForwardDiff" + LuxMLUtilsExt = "MLUtils" + LuxMPIExt = "MPI" + LuxMPINCCLExt = ["CUDA", "MPI", "NCCL"] + LuxOptimisersExt = "Optimisers" + LuxReverseDiffExt = "ReverseDiff" + LuxSimpleChainsExt = "SimpleChains" + LuxTrackerExt = "Tracker" + LuxZygoteExt = "Zygote" + + [deps.Lux.weakdeps] + AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" + CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" + ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" + DynamicExpressions = "a40a106e-89c9-4ca8-8020-a735e8728b6b" + Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" + Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" + ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" + MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" + MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" + NCCL = "3fe64909-d7a1-4096-9b7d-7a0f12cf0f6b" + Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" + ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" + SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5" + Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" + Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[[deps.LuxCUDA]] +deps = ["CUDA", "Reexport", "cuDNN"] +git-tree-sha1 = "7cf839d5eecb41e50dd430a9dcf9398727b71e8a" +uuid = "d0bbae9a-e099-4d5b-a835-1c6931763bda" +version = "0.3.2" + +[[deps.LuxCore]] +deps = ["Functors", "Random", "Setfield"] +git-tree-sha1 = "c96985555a9fe41d7ec2bd5625d6c2077e05e33e" +uuid = "bb33d45b-7691-41d6-9220-0943567d0623" +version = "0.1.15" + +[[deps.LuxDeviceUtils]] +deps = ["Adapt", "ChainRulesCore", "Functors", "LuxCore", "PrecompileTools", "Preferences", "Random"] +git-tree-sha1 = "d3a5cb86d3f4a5ba0ee4d2ca501055e400960c4c" +uuid = "34f89e08-e1d5-43b4-8944-0b49ac560553" +version = "0.1.23" + + [deps.LuxDeviceUtils.extensions] + LuxDeviceUtilsAMDGPUExt = "AMDGPU" + LuxDeviceUtilsCUDAExt = "CUDA" + LuxDeviceUtilsFillArraysExt = "FillArrays" + LuxDeviceUtilsGPUArraysExt = "GPUArrays" + LuxDeviceUtilsLuxCUDAExt = "LuxCUDA" + LuxDeviceUtilsMetalExt = ["GPUArrays", "Metal"] + LuxDeviceUtilsRecursiveArrayToolsExt = "RecursiveArrayTools" + LuxDeviceUtilsReverseDiffExt = "ReverseDiff" + LuxDeviceUtilsSparseArraysExt = "SparseArrays" + LuxDeviceUtilsTrackerExt = "Tracker" + LuxDeviceUtilsZygoteExt = "Zygote" + LuxDeviceUtilsoneAPIExt = ["GPUArrays", "oneAPI"] + + [deps.LuxDeviceUtils.weakdeps] + AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" + CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" + FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" + GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" + LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda" + Metal = "dde4c033-4e86-420c-a63e-0dd931031962" + RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" + ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" + SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" + Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" + Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" + +[[deps.LuxLib]] +deps = ["ArrayInterface", "ChainRulesCore", "EnzymeCore", "FastBroadcast", "FastClosures", "GPUArraysCore", "LinearAlgebra", "LuxCore", "Markdown", "NNlib", "PrecompileTools", "Random", "Reexport", "Statistics"] +git-tree-sha1 = "7b203688117c3e21f074d5a551618ec3126c4e50" +uuid = "82251201-b29d-42c6-8e01-566dec8acb11" +version = "0.3.27" + + [deps.LuxLib.extensions] + LuxLibAMDGPUExt = "AMDGPU" + LuxLibCUDAExt = "CUDA" + LuxLibForwardDiffExt = "ForwardDiff" + LuxLibReverseDiffExt = "ReverseDiff" + LuxLibTrackerAMDGPUExt = ["AMDGPU", "Tracker"] + LuxLibTrackerExt = "Tracker" + LuxLibTrackercuDNNExt = ["CUDA", "Tracker", "cuDNN"] + LuxLibcuDNNExt = ["CUDA", "cuDNN"] + + [deps.LuxLib.weakdeps] + AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" + CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" + ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" + ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" + Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" + cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" + [[deps.Lz4_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] git-tree-sha1 = "6c26c5e8a4203d43b5497be3ec5d4e0c3cde240a" @@ -576,6 +786,11 @@ git-tree-sha1 = "2fa9ee3e63fd3a4f7a9a4f4744a52f4856de82df" uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" version = "0.5.13" +[[deps.ManualMemory]] +git-tree-sha1 = "bcaef4fc7a0cfe2cba636d84cda54b5e4e4ca3cd" +uuid = "d125e4d3-2237-4719-b19c-fa641b8a4667" +version = "0.1.8" + [[deps.Markdown]] deps = ["Base64"] uuid = "d6f4376e-aef5-505a-96c1-9c027394607a" @@ -610,6 +825,24 @@ git-tree-sha1 = "a640912695952b074672edb5f9aaee2f7f9fd59a" uuid = "85f8d34a-cbdd-5861-8df4-14fed0d494ab" version = "0.14.4" +[[deps.NNlib]] +deps = ["Adapt", "Atomix", "ChainRulesCore", "GPUArraysCore", "KernelAbstractions", "LinearAlgebra", "Pkg", "Random", "Requires", "Statistics"] +git-tree-sha1 = "78de319bce99d1d8c1d4fe5401f7cfc2627df396" +uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" +version = "0.9.18" + + [deps.NNlib.extensions] + NNlibAMDGPUExt = "AMDGPU" + NNlibCUDACUDNNExt = ["CUDA", "cuDNN"] + NNlibCUDAExt = "CUDA" + NNlibEnzymeCoreExt = "EnzymeCore" + + [deps.NNlib.weakdeps] + AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" + CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" + EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" + cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" + [[deps.NVTX]] deps = ["Colors", "JuliaNVTXCallbacks_jll", "Libdl", "NVTX_jll"] git-tree-sha1 = "53046f0483375e3ed78e49190f1154fa0a4083a1" @@ -665,9 +898,9 @@ version = "4.1.6+0" [[deps.OpenSSL_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "3da7367955dcc5c54c1ba4d402ccdc09a1a3e046" +git-tree-sha1 = "a028ee3cb5641cccc4c24e90c36b0a4f7707bdf5" uuid = "458c3c95-2e84-50aa-8efc-19380b2a3a95" -version = "3.0.13+1" +version = "3.0.14+0" [[deps.OpenSpecFun_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"] @@ -698,17 +931,25 @@ git-tree-sha1 = "8489905bcdbcfac64d1daa51ca07c0d8f0283821" uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0" version = "2.8.1" +[[deps.PartialFunctions]] +deps = ["MacroTools"] +git-tree-sha1 = "47b49a4dbc23b76682205c646252c0f9e1eb75af" +uuid = "570af359-4316-4cb7-8c74-252c00c2016b" +version = "1.2.0" + [[deps.PencilArrays]] deps = ["Adapt", "JSON3", "LinearAlgebra", "MPI", "OffsetArrays", "Random", "Reexport", "StaticArrayInterface", "StaticArrays", "StaticPermutations", "Strided", "TimerOutputs", "VersionParsing"] -git-tree-sha1 = "6510e851700a851944f7ffa5cd990cced4802ad2" +git-tree-sha1 = "fa85ac32172d96cfdb91dbc53e8e57007e5a2b5a" uuid = "0e08944d-e94e-41b1-9406-dcf66b6a9d2e" -version = "0.19.3" +version = "0.19.5" [deps.PencilArrays.extensions] + PencilArraysAMDGPUExt = ["AMDGPU"] PencilArraysDiffEqExt = ["DiffEqBase"] PencilArraysHDF5Ext = ["HDF5"] [deps.PencilArrays.weakdeps] + AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" @@ -729,6 +970,18 @@ git-tree-sha1 = "f9501cc0430a26bc3d156ae1b5b0c1b47af4d6da" uuid = "eebad327-c553-4316-9ea0-9fa01ccd7688" version = "0.3.3" +[[deps.Polyester]] +deps = ["ArrayInterface", "BitTwiddlingConvenienceFunctions", "CPUSummary", "IfElse", "ManualMemory", "PolyesterWeave", "Requires", "Static", "StaticArrayInterface", "StrideArraysCore", "ThreadingUtilities"] +git-tree-sha1 = "b3e2bae88cf07baf0a051fe09666b8ef97aefe93" +uuid = "f517fe37-dbe3-4b94-8317-1923a5111588" +version = "0.7.14" + +[[deps.PolyesterWeave]] +deps = ["BitTwiddlingConvenienceFunctions", "CPUSummary", "IfElse", "Static", "ThreadingUtilities"] +git-tree-sha1 = "240d7170f5ffdb285f9427b92333c3463bf65bf6" +uuid = "1d0040c9-8b98-4ee7-8388-3f51789ca0ad" +version = "0.2.1" + [[deps.PooledArrays]] deps = ["DataAPI", "Future"] git-tree-sha1 = "36d8b4b899628fb92c2749eb488d884a926614d3" @@ -826,6 +1079,11 @@ weakdeps = ["RecipesBase"] uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce" version = "0.7.0" +[[deps.SIMDTypes]] +git-tree-sha1 = "330289636fb8107c5f32088d2741e9fd7a061a5c" +uuid = "94e857df-77ce-4151-89e5-788b33177be4" +version = "0.1.0" + [[deps.Scratch]] deps = ["Dates"] git-tree-sha1 = "3bac05bc7e74a75fd9cba4295cde4045d9fe2386" @@ -846,6 +1104,12 @@ version = "1.4.3" [[deps.Serialization]] uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b" +[[deps.Setfield]] +deps = ["ConstructionBase", "Future", "MacroTools", "StaticArraysCore"] +git-tree-sha1 = "e2cc6d8c88613c05e1defb55170bf5ff211fbeac" +uuid = "efcf1570-3423-57d1-acb7-fd33fddbac46" +version = "1.1.1" + [[deps.Sockets]] uuid = "6462fe0b-24de-5631-8697-dd941f90decc" @@ -865,13 +1129,11 @@ deps = ["IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_j git-tree-sha1 = "2f5d4697f21388cbe1ff299430dd169ef97d7e14" uuid = "276daf66-3868-5448-9aa4-cd146d93841b" version = "2.4.0" +weakdeps = ["ChainRulesCore"] [deps.SpecialFunctions.extensions] SpecialFunctionsChainRulesCoreExt = "ChainRulesCore" - [deps.SpecialFunctions.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - [[deps.Static]] deps = ["IfElse"] git-tree-sha1 = "d2fdac9ff3906e27f7a618d47b676941baa6c80c" @@ -891,22 +1153,19 @@ weakdeps = ["OffsetArrays", "StaticArrays"] [[deps.StaticArrays]] deps = ["LinearAlgebra", "PrecompileTools", "Random", "StaticArraysCore"] -git-tree-sha1 = "9ae599cd7529cfce7fea36cf00a62cfc56f0f37c" +git-tree-sha1 = "6e00379a24597be4ae1ee6b2d882e15392040132" uuid = "90137ffa-7385-5640-81b9-e52037218182" -version = "1.9.4" +version = "1.9.5" +weakdeps = ["ChainRulesCore", "Statistics"] [deps.StaticArrays.extensions] StaticArraysChainRulesCoreExt = "ChainRulesCore" StaticArraysStatisticsExt = "Statistics" - [deps.StaticArrays.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" - [[deps.StaticArraysCore]] -git-tree-sha1 = "36b3d696ce6366023a0ea192b4cd442268995a0d" +git-tree-sha1 = "192954ef1208c7019899fbf8049e717f92959682" uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" -version = "1.4.2" +version = "1.4.3" [[deps.StaticPermutations]] git-tree-sha1 = "193c3daa18ff3e55c1dae66acb6a762c4a3bdb0b" @@ -924,17 +1183,23 @@ git-tree-sha1 = "1ff449ad350c9c4cbc756624d6f8a8c3ef56d3ed" uuid = "82ae8749-77ed-4fe6-ae5f-f523153014b0" version = "1.7.0" +[[deps.StrideArraysCore]] +deps = ["ArrayInterface", "CloseOpenIntervals", "IfElse", "LayoutPointers", "LinearAlgebra", "ManualMemory", "SIMDTypes", "Static", "StaticArrayInterface", "ThreadingUtilities"] +git-tree-sha1 = "25349bf8f63aa36acbff5e3550a86e9f5b0ef682" +uuid = "7792a7ef-975c-4747-a70f-980b88e8d1da" +version = "0.5.6" + [[deps.Strided]] deps = ["LinearAlgebra", "StridedViews", "TupleTools"] -git-tree-sha1 = "40c69be0e1b72ee2f42923b7d1ff13e0b04e675c" +git-tree-sha1 = "bd9bd1c70cfc115cc3a30213fc725125a6b43652" uuid = "5e0ebb24-38b0-5f93-81fe-25c709ecae67" -version = "2.0.4" +version = "2.1.0" [[deps.StridedViews]] deps = ["LinearAlgebra", "PackageExtensionCompat"] -git-tree-sha1 = "5b765c4e401693ab08981989f74a36a010aa1d8e" +git-tree-sha1 = "2917996ce0fa6b8a3a85240a5e9ff930e2aeaa43" uuid = "4db3bf67-4bd7-4b4e-b153-31dc3fb37143" -version = "0.2.2" +version = "0.3.1" weakdeps = ["CUDA"] [deps.StridedViews.extensions] @@ -1012,6 +1277,12 @@ version = "0.16.0" deps = ["InteractiveUtils", "Logging", "Random", "Serialization"] uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +[[deps.ThreadingUtilities]] +deps = ["ManualMemory"] +git-tree-sha1 = "eda08f7e9818eb53661b3deb74e3159460dfbc27" +uuid = "8290d209-cae3-49c0-8002-c8c24d57dab5" +version = "0.5.2" + [[deps.TimerOutputs]] deps = ["ExprTools", "Printf"] git-tree-sha1 = "5a13ae8a41237cff5ecf34f73eb1b8f42fff6531" @@ -1019,9 +1290,9 @@ uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" version = "0.5.24" [[deps.TranscodingStreams]] -git-tree-sha1 = "5d54d076465da49d6746c647022f3b3674e64156" +git-tree-sha1 = "a947ea21087caba0a798c5e494d0bb78e3a1a3a0" uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa" -version = "0.10.8" +version = "0.10.9" weakdeps = ["Random", "Test"] [deps.TranscodingStreams.extensions] @@ -1055,6 +1326,16 @@ git-tree-sha1 = "58d6e80b4ee071f5efd07fda82cb9fbe17200868" uuid = "81def892-9a0e-5fdd-b105-ffc91e053289" version = "1.3.0" +[[deps.WeightInitializers]] +deps = ["ChainRulesCore", "LinearAlgebra", "PartialFunctions", "PrecompileTools", "Random", "SpecialFunctions", "Statistics"] +git-tree-sha1 = "f0e6760ef9d22f043710289ddf29e4a4048c4822" +uuid = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d" +version = "0.1.7" +weakdeps = ["CUDA"] + + [deps.WeightInitializers.extensions] + WeightInitializersCUDAExt = "CUDA" + [[deps.XML2_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Libiconv_jll", "Zlib_jll"] git-tree-sha1 = "52ff2af32e591541550bd753c0da8b9bc92bb9d9" @@ -1078,6 +1359,12 @@ git-tree-sha1 = "e678132f07ddb5bfa46857f0d7620fb9be675d3b" uuid = "3161d3a3-bdf6-5164-811a-617609db77b4" version = "1.5.6+0" +[[deps.cuDNN]] +deps = ["CEnum", "CUDA", "CUDA_Runtime_Discovery", "CUDNN_jll"] +git-tree-sha1 = "1f6a185a8da9bbbc20134b7b935981f70c9b26ad" +uuid = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" +version = "1.3.1" + [[deps.libaec_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] git-tree-sha1 = "46bf7be2917b59b761247be3f317ddf75e50e997" diff --git a/Project.toml b/Project.toml index 5d2eb1405e..049d83a8b4 100644 --- a/Project.toml +++ b/Project.toml @@ -20,6 +20,8 @@ JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Logging = "56ddb016-857b-54e1-b83d-db4d58db5568" +Lux = "b2108857-7c20-44ae-9111-449ecde12c47" +LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab" OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" @@ -32,6 +34,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Rotations = "6038ab10-8711-5258-84ad-4b1120ba62dc" SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" @@ -90,6 +93,4 @@ TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" TimesDates = "bdfc003b-8df8-5c39-adcd-3a9087f5df4a" [targets] -test = ["BenchmarkTools", "Coverage", "CUDA_Runtime_jll", "DataDeps", "Enzyme", - "InteractiveUtils", "MPIPreferences", "OpenMPI_jll", "Test", "TimerOutputs", - "TimesDates", "SafeTestsets"] +test = ["BenchmarkTools", "Coverage", "CUDA_Runtime_jll", "DataDeps", "Enzyme", "InteractiveUtils", "MPIPreferences", "OpenMPI_jll", "Test", "TimerOutputs", "TimesDates", "SafeTestsets"] From 930cc7ee6ab93e8442b3c3dac7769673ff1a960e Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 2 Jul 2024 12:36:32 +0100 Subject: [PATCH 018/114] Update NN closure model to use a larger neural network --- NN_closure_global.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/NN_closure_global.jl b/NN_closure_global.jl index 001dc5700f..4df648f355 100644 --- a/NN_closure_global.jl +++ b/NN_closure_global.jl @@ -69,7 +69,7 @@ Adapt.adapt_structure(to, nn :: NN) = function NNFluxClosure(arch) dev = ifelse(arch == GPU(), gpu_device(), cpu_device()) - nn_path = "./NDE_FC_Qb_15simnew_2layer_64_relu_2Pr_model.jld2" + nn_path = "./NDE_FC_Qb_18simnew_2layer_128_relu_2Pr_model.jld2" ps, sts, scaling_params, wT_model, wS_model = jldopen(nn_path, "r") do file ps = file["u"] |> dev |> f64 @@ -158,7 +158,7 @@ end i, j, k = @index(Global, NTuple) scaling = closure.scaling convecting = top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, tracers) > 0 - interior_point = k <= grid.Nz - 1 & k >= 2 + interior_point = k <= grid.Nz - 1 & k >= 3 @inbounds diffusivities.wT[i, j, k] = ifelse(convecting & interior_point, inv(scaling.wT)(diffusivities.wT[i, j, k]) - inv(scaling.wT)(0), 0) @inbounds diffusivities.wS[i, j, k] = ifelse(convecting & interior_point, inv(scaling.wS)(diffusivities.wS[i, j, k]) - inv(scaling.wS)(0), 0) From fd95fc6ae57af0ee120594fbd58f25206b586b6d Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 2 Jul 2024 14:35:53 +0100 Subject: [PATCH 019/114] Remove unused import of StaticArrays --- NN_closure_global.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/NN_closure_global.jl b/NN_closure_global.jl index 4df648f355..3f9dcade77 100644 --- a/NN_closure_global.jl +++ b/NN_closure_global.jl @@ -25,7 +25,6 @@ using Oceananigans: architecture, on_architecture using Lux, LuxCUDA using JLD2 using ComponentArrays -using StaticArrays using OffsetArrays using KernelAbstractions: @index, @kernel, @private From 72aebc1535edfda270f462ac9aceba3a656c00d1 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 2 Jul 2024 15:12:30 +0100 Subject: [PATCH 020/114] Add ComponentArrays dependency --- Manifest.toml | 66 ++++++++++++++++++++++++++++++++++++++++++++++++++- Project.toml | 1 + 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/Manifest.toml b/Manifest.toml index 6da492bf6a..4e69b75d27 100644 --- a/Manifest.toml +++ b/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.10.1" manifest_format = "2.0" -project_hash = "f43673c2cec178a7eac57c04f7c0f97bcf64e09b" +project_hash = "2d1be40f4d74ba6922f076fdb530a846bb95a6f4" [[deps.ADTypes]] git-tree-sha1 = "fa0822e5baee6e23081c2685ae27265dabee23d8" @@ -195,6 +195,12 @@ git-tree-sha1 = "d6fb5bf939a2753c74984b11434ea25d6c397a58" uuid = "1fbeeb36-5f17-413c-809b-666fb144f157" version = "0.3.6" +[[deps.CommonSubexpressions]] +deps = ["MacroTools", "Test"] +git-tree-sha1 = "7b8a93dba8af7e3b42fecabf646260105ac373f7" +uuid = "bbf7d656-a473-5ed7-a52c-81e309532950" +version = "0.3.0" + [[deps.Compat]] deps = ["TOML", "UUIDs"] git-tree-sha1 = "b1c55339b7c6c350ee89f2c1604299660525b248" @@ -210,6 +216,36 @@ deps = ["Artifacts", "Libdl"] uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae" version = "1.1.0+0" +[[deps.ComponentArrays]] +deps = ["ArrayInterface", "ChainRulesCore", "ForwardDiff", "Functors", "LinearAlgebra", "PackageExtensionCompat", "StaticArrayInterface", "StaticArraysCore"] +git-tree-sha1 = "c2663c30580894680c793d6b8043567b5f4d4878" +uuid = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" +version = "0.15.14" + + [deps.ComponentArrays.extensions] + ComponentArraysAdaptExt = "Adapt" + ComponentArraysConstructionBaseExt = "ConstructionBase" + ComponentArraysGPUArraysExt = "GPUArrays" + ComponentArraysOptimisersExt = "Optimisers" + ComponentArraysRecursiveArrayToolsExt = "RecursiveArrayTools" + ComponentArraysReverseDiffExt = "ReverseDiff" + ComponentArraysSciMLBaseExt = "SciMLBase" + ComponentArraysTrackerExt = "Tracker" + ComponentArraysTruncatedStacktracesExt = "TruncatedStacktraces" + ComponentArraysZygoteExt = "Zygote" + + [deps.ComponentArrays.weakdeps] + Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" + ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9" + GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" + Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" + RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" + ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" + SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" + Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" + TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77" + Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + [[deps.ConcreteStructs]] git-tree-sha1 = "f749037478283d372048690eb3b5f92a79432b34" uuid = "2569d6c7-a4a2-43d3-a901-331e8e4be471" @@ -272,6 +308,18 @@ version = "1.0.0" deps = ["Printf"] uuid = "ade2ca70-3891-5945-98fb-dc099432e06a" +[[deps.DiffResults]] +deps = ["StaticArraysCore"] +git-tree-sha1 = "782dd5f4561f5d267313f23853baaaa4c52ea621" +uuid = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" +version = "1.1.0" + +[[deps.DiffRules]] +deps = ["IrrationalConstants", "LogExpFunctions", "NaNMath", "Random", "SpecialFunctions"] +git-tree-sha1 = "23163d55f885173722d1e4cf0f6110cdbaf7e272" +uuid = "b552c78f-8df3-52c6-915a-8e097449b14b" +version = "1.15.1" + [[deps.DiskArrays]] deps = ["LRUCache", "OffsetArrays"] git-tree-sha1 = "ef25c513cad08d7ebbed158c91768ae32f308336" @@ -361,6 +409,16 @@ git-tree-sha1 = "05882d6995ae5c12bb5f36dd2ed3f61c98cbb172" uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93" version = "0.8.5" +[[deps.ForwardDiff]] +deps = ["CommonSubexpressions", "DiffResults", "DiffRules", "LinearAlgebra", "LogExpFunctions", "NaNMath", "Preferences", "Printf", "Random", "SpecialFunctions"] +git-tree-sha1 = "cf0fe81336da9fb90944683b8c41984b08793dad" +uuid = "f6369f11-7733-5829-9624-2563aa707210" +version = "0.10.36" +weakdeps = ["StaticArrays"] + + [deps.ForwardDiff.extensions] + ForwardDiffStaticArraysExt = "StaticArrays" + [[deps.Functors]] deps = ["LinearAlgebra"] git-tree-sha1 = "8a66c07630d6428eaab3506a0eabfcf4a9edea05" @@ -855,6 +913,12 @@ git-tree-sha1 = "ce3269ed42816bf18d500c9f63418d4b0d9f5a3b" uuid = "e98f9f5b-d649-5603-91fd-7774390e6439" version = "3.1.0+2" +[[deps.NaNMath]] +deps = ["OpenLibm_jll"] +git-tree-sha1 = "0877504529a3e5c3343c6f8b4c0381e57e4387e4" +uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3" +version = "1.0.2" + [[deps.NetCDF_jll]] deps = ["Artifacts", "Blosc_jll", "Bzip2_jll", "HDF5_jll", "JLLWrappers", "LazyArtifacts", "LibCURL_jll", "Libdl", "MPICH_jll", "MPIPreferences", "MPItrampoline_jll", "MicrosoftMPI_jll", "OpenMPI_jll", "TOML", "XML2_jll", "Zlib_jll", "Zstd_jll", "libzip_jll"] git-tree-sha1 = "4686378c4ae1d1948cfbe46c002a11a4265dcb07" diff --git a/Project.toml b/Project.toml index 049d83a8b4..87947544a0 100644 --- a/Project.toml +++ b/Project.toml @@ -6,6 +6,7 @@ version = "0.91.2" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" Crayons = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f" CubedSphere = "7445602f-e544-4518-8976-18f8e8ae6cdb" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" From c2d336e7a48de58ce7c7352f06bf06ecd416e4e8 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 2 Jul 2024 22:20:13 +0100 Subject: [PATCH 021/114] add total_size and KernelParameters dependency --- xin_kai_vertical_diffusivity_2Pr.jl | 2 ++ 1 file changed, 2 insertions(+) diff --git a/xin_kai_vertical_diffusivity_2Pr.jl b/xin_kai_vertical_diffusivity_2Pr.jl index 79a2765de9..ca28a5fdfd 100644 --- a/xin_kai_vertical_diffusivity_2Pr.jl +++ b/xin_kai_vertical_diffusivity_2Pr.jl @@ -4,6 +4,8 @@ using Oceananigans.BuoyancyModels: ∂z_b using Oceananigans.Operators using Oceananigans.BoundaryConditions using Oceananigans.Grids: inactive_node +using Oceananigans.Grids: total_size +using Oceananigans.Utils: KernelParameters using Oceananigans.Operators: ℑzᵃᵃᶜ, ℑxyᶠᶠᵃ, ℑxyᶜᶜᵃ using Adapt From d4a1156cc246c40c0670c0241a6288777ca155fb Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Wed, 17 Jul 2024 02:43:16 +0100 Subject: [PATCH 022/114] Coarsen LES data for NN closure model --- coarsen_LES_NN_closure.jl | 111 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 111 insertions(+) create mode 100644 coarsen_LES_NN_closure.jl diff --git a/coarsen_LES_NN_closure.jl b/coarsen_LES_NN_closure.jl new file mode 100644 index 0000000000..2a7d2dcbd8 --- /dev/null +++ b/coarsen_LES_NN_closure.jl @@ -0,0 +1,111 @@ +using Oceananigans +using Statistics +using JLD2 + +LES_FILE_DIR = "./NN_2D_channel_horizontal_convection_0.0003_LES.jld2" +MODEL_FILE_DIR = "./NN_closure_2D_channel.jld2" + +u_data_LES = FieldTimeSeries(LES_FILE_DIR, "u", backend=OnDisk()) +v_data_LES = FieldTimeSeries(LES_FILE_DIR, "v", backend=OnDisk()) +T_data_LES = FieldTimeSeries(LES_FILE_DIR, "T", backend=OnDisk()) +S_data_LES = FieldTimeSeries(LES_FILE_DIR, "S", backend=OnDisk()) + +u_data_model = FieldTimeSeries(MODEL_FILE_DIR, "u") +v_data_model = FieldTimeSeries(MODEL_FILE_DIR, "v") +T_data_model = FieldTimeSeries(MODEL_FILE_DIR, "T") +S_data_model = FieldTimeSeries(MODEL_FILE_DIR, "S") + +end_time = 23 * 60^2 * 24 + +yC_LES = ynodes(T_data_LES.grid, Center()) +yF_LES = ynodes(T_data_LES.grid, Face()) + +zC_LES = znodes(T_data_LES.grid, Center()) +zF_LES = znodes(T_data_LES.grid, Face()) + +Nt_LES = findfirst(x -> x ≈ end_time, T_data_LES.times) + +Δy_LES = T_data_LES.grid.Ly / T_data_LES.grid.Ny +Δz_LES = T_data_LES.grid.Lz / T_data_LES.grid.Nz + +Ny_model = T_data_model.grid.Ny +Nz_model = T_data_model.grid.Nz + +Δy_model = T_data_model.grid.Ly / Ny_model +Δz_model = T_data_model.grid.Lz / Nz_model + +coarse_ratio_y = Int(Δy_model / Δy_LES) +coarse_ratio_z = Int(Δz_model / Δz_LES) + +function coarsen_dataᵃᶜᵃ(data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt_LES) + Ny_LES = data_LES.grid.Ny + Nz_LES = data_LES.grid.Nz + data_LES_coarse = zeros(1, Ny_model, Nz_model, Nt_LES) + LES_temp = zeros(Ny_LES, Nz_LES) + + for nt in 1:Nt_LES + LES_temp .= interior(data_LES[nt], 1, :, :) + Threads.@threads for j in axes(data_LES_coarse, 2) + @info "nt = $nt, Processing j = $j" + for k in axes(data_LES_coarse, 3) + data_LES_coarse[1, j, k, nt] = mean(LES_temp[(j-1)*coarse_ratio_y+1:j*coarse_ratio_y, (k-1)*coarse_ratio_z+1:k*coarse_ratio_z]) + end + end + end + + return data_LES_coarse +end + +T_data_LES_coarse = coarsen_dataᵃᶜᵃ(T_data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt_LES) +S_data_LES_coarse = coarsen_dataᵃᶜᵃ(S_data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt_LES) +u_data_LES_coarse = coarsen_dataᵃᶜᵃ(u_data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt_LES) + +function coarsen_dataᵃᶠᵃ(data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt) + Ny_LES = data_LES.grid.Ny + Nz_LES = data_LES.grid.Nz + + dataᵃᶠᵃ_temp = zeros(Ny_LES+1, Nz_LES) + dataᵃᶜᵃ_temp = zeros(Ny_LES, Nz_LES) + + data_coarse = zeros(1, Ny_model, Nz_model, Nt) + + for nt in 1:Nt + @info "nt = $nt, Interpolating for LES data" + dataᵃᶠᵃ_temp .= interior(data_LES[nt], 1, :, :) + Threads.@threads for j in 1:Ny_LES + # for j in 1:Ny_LES + for k in 1:Nz_LES + dataᵃᶜᵃ_temp[j, k] = mean(dataᵃᶠᵃ_temp[j:j+1, k]) + end + end + + @info "nt = $nt, Coarsening data" + Threads.@threads for j in axes(data_coarse, 2) + # for j in axes(data_coarse, 2) + for k in axes(data_coarse, 3) + data_coarse[1, j, k, nt] = mean(dataᵃᶜᵃ_temp[(j-1)*coarse_ratio_y+1:j*coarse_ratio_y, (k-1)*coarse_ratio_z+1:k*coarse_ratio_z]) + end + end + end + + return data_coarse +end + +v_data_LES_coarse = coarsen_dataᵃᶠᵃ(v_data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt_LES) +v_data_model_coarse = coarsen_dataᵃᶠᵃ(v_data_model, 1, 1, Ny_model, Nz_model, Nt_LES) + +u_data_model_coarse = interior(u_data_model) +T_data_model_coarse = interior(T_data_model) +S_data_model_coarse = interior(S_data_model) + +jldopen("./LES_NDE_FC_Qb_absf_24simnew_2layer_128_relu_2Pr_coarsened.jld2", "w") do file + file["u_LES"] = u_data_LES_coarse + file["v_LES"] = v_data_LES_coarse + file["T_LES"] = T_data_LES_coarse + file["S_LES"] = S_data_LES_coarse + file["u_NN_model"] = u_data_model_coarse + file["v_NN_model"] = v_data_model_coarse + file["T_NN_model"] = T_data_model_coarse + file["S_NN_model"] = S_data_model_coarse +end + From ae7b563472cfa0b677e82153cbd401cdc86f109a Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 16 Jul 2024 22:19:21 -0400 Subject: [PATCH 023/114] rename file and compare LES with NN closure --- coarsen_LES_NN_closure.jl | 111 -------------------- compare_LES_NN_closure.jl | 208 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 208 insertions(+), 111 deletions(-) delete mode 100644 coarsen_LES_NN_closure.jl create mode 100644 compare_LES_NN_closure.jl diff --git a/coarsen_LES_NN_closure.jl b/coarsen_LES_NN_closure.jl deleted file mode 100644 index 2a7d2dcbd8..0000000000 --- a/coarsen_LES_NN_closure.jl +++ /dev/null @@ -1,111 +0,0 @@ -using Oceananigans -using Statistics -using JLD2 - -LES_FILE_DIR = "./NN_2D_channel_horizontal_convection_0.0003_LES.jld2" -MODEL_FILE_DIR = "./NN_closure_2D_channel.jld2" - -u_data_LES = FieldTimeSeries(LES_FILE_DIR, "u", backend=OnDisk()) -v_data_LES = FieldTimeSeries(LES_FILE_DIR, "v", backend=OnDisk()) -T_data_LES = FieldTimeSeries(LES_FILE_DIR, "T", backend=OnDisk()) -S_data_LES = FieldTimeSeries(LES_FILE_DIR, "S", backend=OnDisk()) - -u_data_model = FieldTimeSeries(MODEL_FILE_DIR, "u") -v_data_model = FieldTimeSeries(MODEL_FILE_DIR, "v") -T_data_model = FieldTimeSeries(MODEL_FILE_DIR, "T") -S_data_model = FieldTimeSeries(MODEL_FILE_DIR, "S") - -end_time = 23 * 60^2 * 24 - -yC_LES = ynodes(T_data_LES.grid, Center()) -yF_LES = ynodes(T_data_LES.grid, Face()) - -zC_LES = znodes(T_data_LES.grid, Center()) -zF_LES = znodes(T_data_LES.grid, Face()) - -Nt_LES = findfirst(x -> x ≈ end_time, T_data_LES.times) - -Δy_LES = T_data_LES.grid.Ly / T_data_LES.grid.Ny -Δz_LES = T_data_LES.grid.Lz / T_data_LES.grid.Nz - -Ny_model = T_data_model.grid.Ny -Nz_model = T_data_model.grid.Nz - -Δy_model = T_data_model.grid.Ly / Ny_model -Δz_model = T_data_model.grid.Lz / Nz_model - -coarse_ratio_y = Int(Δy_model / Δy_LES) -coarse_ratio_z = Int(Δz_model / Δz_LES) - -function coarsen_dataᵃᶜᵃ(data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt_LES) - Ny_LES = data_LES.grid.Ny - Nz_LES = data_LES.grid.Nz - data_LES_coarse = zeros(1, Ny_model, Nz_model, Nt_LES) - LES_temp = zeros(Ny_LES, Nz_LES) - - for nt in 1:Nt_LES - LES_temp .= interior(data_LES[nt], 1, :, :) - Threads.@threads for j in axes(data_LES_coarse, 2) - @info "nt = $nt, Processing j = $j" - for k in axes(data_LES_coarse, 3) - data_LES_coarse[1, j, k, nt] = mean(LES_temp[(j-1)*coarse_ratio_y+1:j*coarse_ratio_y, (k-1)*coarse_ratio_z+1:k*coarse_ratio_z]) - end - end - end - - return data_LES_coarse -end - -T_data_LES_coarse = coarsen_dataᵃᶜᵃ(T_data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt_LES) -S_data_LES_coarse = coarsen_dataᵃᶜᵃ(S_data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt_LES) -u_data_LES_coarse = coarsen_dataᵃᶜᵃ(u_data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt_LES) - -function coarsen_dataᵃᶠᵃ(data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt) - Ny_LES = data_LES.grid.Ny - Nz_LES = data_LES.grid.Nz - - dataᵃᶠᵃ_temp = zeros(Ny_LES+1, Nz_LES) - dataᵃᶜᵃ_temp = zeros(Ny_LES, Nz_LES) - - data_coarse = zeros(1, Ny_model, Nz_model, Nt) - - for nt in 1:Nt - @info "nt = $nt, Interpolating for LES data" - dataᵃᶠᵃ_temp .= interior(data_LES[nt], 1, :, :) - Threads.@threads for j in 1:Ny_LES - # for j in 1:Ny_LES - for k in 1:Nz_LES - dataᵃᶜᵃ_temp[j, k] = mean(dataᵃᶠᵃ_temp[j:j+1, k]) - end - end - - @info "nt = $nt, Coarsening data" - Threads.@threads for j in axes(data_coarse, 2) - # for j in axes(data_coarse, 2) - for k in axes(data_coarse, 3) - data_coarse[1, j, k, nt] = mean(dataᵃᶜᵃ_temp[(j-1)*coarse_ratio_y+1:j*coarse_ratio_y, (k-1)*coarse_ratio_z+1:k*coarse_ratio_z]) - end - end - end - - return data_coarse -end - -v_data_LES_coarse = coarsen_dataᵃᶠᵃ(v_data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt_LES) -v_data_model_coarse = coarsen_dataᵃᶠᵃ(v_data_model, 1, 1, Ny_model, Nz_model, Nt_LES) - -u_data_model_coarse = interior(u_data_model) -T_data_model_coarse = interior(T_data_model) -S_data_model_coarse = interior(S_data_model) - -jldopen("./LES_NDE_FC_Qb_absf_24simnew_2layer_128_relu_2Pr_coarsened.jld2", "w") do file - file["u_LES"] = u_data_LES_coarse - file["v_LES"] = v_data_LES_coarse - file["T_LES"] = T_data_LES_coarse - file["S_LES"] = S_data_LES_coarse - file["u_NN_model"] = u_data_model_coarse - file["v_NN_model"] = v_data_model_coarse - file["T_NN_model"] = T_data_model_coarse - file["S_NN_model"] = S_data_model_coarse -end - diff --git a/compare_LES_NN_closure.jl b/compare_LES_NN_closure.jl new file mode 100644 index 0000000000..a18fbd5b9d --- /dev/null +++ b/compare_LES_NN_closure.jl @@ -0,0 +1,208 @@ +using Oceananigans +using Statistics +using JLD2 +using CairoMakie + +# LES_FILE_DIR = "./NN_2D_channel_horizontal_convection_0.0003_LES.jld2" + +# u_data_LES = FieldTimeSeries(LES_FILE_DIR, "u", backend=OnDisk()) +# v_data_LES = FieldTimeSeries(LES_FILE_DIR, "v", backend=OnDisk()) +# T_data_LES = FieldTimeSeries(LES_FILE_DIR, "T", backend=OnDisk()) +# S_data_LES = FieldTimeSeries(LES_FILE_DIR, "S", backend=OnDisk()) + +# yC_LES = ynodes(T_data_LES.grid, Center()) +# yF_LES = ynodes(T_data_LES.grid, Face()) + +# zC_LES = znodes(T_data_LES.grid, Center()) +# zF_LES = znodes(T_data_LES.grid, Face()) + +# Nt_LES = findfirst(x -> x ≈ end_time, T_data_LES.times) + +# Δy_LES = T_data_LES.grid.Ly / T_data_LES.grid.Ny +# Δz_LES = T_data_LES.grid.Lz / T_data_LES.grid.Nz + +MODEL_FILE_DIR = "./NN_closure_2D_channel.jld2" + +u_data_model = FieldTimeSeries(MODEL_FILE_DIR, "u") +v_data_model = FieldTimeSeries(MODEL_FILE_DIR, "v") +T_data_model = FieldTimeSeries(MODEL_FILE_DIR, "T") +S_data_model = FieldTimeSeries(MODEL_FILE_DIR, "S") + +end_time = 23 * 60^2 * 24 + +Ny_model = T_data_model.grid.Ny +Nz_model = T_data_model.grid.Nz + +Δy_model = T_data_model.grid.Ly / Ny_model +Δz_model = T_data_model.grid.Lz / Nz_model + +yC_model = ynodes(T_data_model.grid, Center()) +zC_model = znodes(T_data_model.grid, Center()) + +# coarse_ratio_y = Int(Δy_model / Δy_LES) +# coarse_ratio_z = Int(Δz_model / Δz_LES) + +# function coarsen_dataᵃᶜᵃ(data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt_LES) +# Ny_LES = data_LES.grid.Ny +# Nz_LES = data_LES.grid.Nz +# data_LES_coarse = zeros(1, Ny_model, Nz_model, Nt_LES) +# LES_temp = zeros(Ny_LES, Nz_LES) + +# for nt in 1:Nt_LES +# LES_temp .= interior(data_LES[nt], 1, :, :) +# Threads.@threads for j in axes(data_LES_coarse, 2) +# @info "nt = $nt, Processing j = $j" +# for k in axes(data_LES_coarse, 3) +# data_LES_coarse[1, j, k, nt] = mean(LES_temp[(j-1)*coarse_ratio_y+1:j*coarse_ratio_y, (k-1)*coarse_ratio_z+1:k*coarse_ratio_z]) +# end +# end +# end + +# return data_LES_coarse +# end + +# T_data_LES_coarse = coarsen_dataᵃᶜᵃ(T_data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt_LES) +# S_data_LES_coarse = coarsen_dataᵃᶜᵃ(S_data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt_LES) +# u_data_LES_coarse = coarsen_dataᵃᶜᵃ(u_data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt_LES) + +# function coarsen_dataᵃᶠᵃ(data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt) +# Ny_LES = data_LES.grid.Ny +# Nz_LES = data_LES.grid.Nz + +# dataᵃᶠᵃ_temp = zeros(Ny_LES+1, Nz_LES) +# dataᵃᶜᵃ_temp = zeros(Ny_LES, Nz_LES) + +# data_coarse = zeros(1, Ny_model, Nz_model, Nt) + +# for nt in 1:Nt +# @info "nt = $nt, Interpolating for LES data" +# dataᵃᶠᵃ_temp .= interior(data_LES[nt], 1, :, :) +# Threads.@threads for j in 1:Ny_LES +# # for j in 1:Ny_LES +# for k in 1:Nz_LES +# dataᵃᶜᵃ_temp[j, k] = mean(dataᵃᶠᵃ_temp[j:j+1, k]) +# end +# end + +# @info "nt = $nt, Coarsening data" +# Threads.@threads for j in axes(data_coarse, 2) +# # for j in axes(data_coarse, 2) +# for k in axes(data_coarse, 3) +# data_coarse[1, j, k, nt] = mean(dataᵃᶜᵃ_temp[(j-1)*coarse_ratio_y+1:j*coarse_ratio_y, (k-1)*coarse_ratio_z+1:k*coarse_ratio_z]) +# end +# end +# end + +# return data_coarse +# end + +# v_data_LES_coarse = coarsen_dataᵃᶠᵃ(v_data_LES, coarse_ratio_y, coarse_ratio_z, Ny_model, Nz_model, Nt_LES) +# v_data_model_coarse = coarsen_dataᵃᶠᵃ(v_data_model, 1, 1, Ny_model, Nz_model, Nt_LES) + +# u_data_model_coarse = interior(u_data_model) +# T_data_model_coarse = interior(T_data_model) +# S_data_model_coarse = interior(S_data_model) + +# jldopen("./LES_NDE_FC_Qb_absf_24simnew_2layer_128_relu_2Pr_coarsened.jld2", "w") do file +# file["u_LES"] = u_data_LES_coarse +# file["v_LES"] = v_data_LES_coarse +# file["T_LES"] = T_data_LES_coarse +# file["S_LES"] = S_data_LES_coarse +# file["u_NN_model"] = u_data_model_coarse +# file["v_NN_model"] = v_data_model_coarse +# file["T_NN_model"] = T_data_model_coarse +# file["S_NN_model"] = S_data_model_coarse +# end + +FILE_DIR = "./LES_NDE_FC_Qb_absf_24simnew_2layer_128_relu_2Pr_coarsened.jld2" + +u_data_LES_coarse, v_data_LES_coarse, T_data_LES_coarse, S_data_LES_coarse, u_data_model_coarse, v_data_model_coarse, T_data_model_coarse, S_data_model_coarse = jldopen(FILE_DIR, "r") do file + u_data_LES_coarse = file["u_LES"] + v_data_LES_coarse = file["v_LES"] + T_data_LES_coarse = file["T_LES"] + S_data_LES_coarse = file["S_LES"] + u_data_model_coarse = file["u_NN_model"] + v_data_model_coarse = file["v_NN_model"] + T_data_model_coarse = file["T_NN_model"] + S_data_model_coarse = file["S_NN_model"] + return u_data_LES_coarse, v_data_LES_coarse, T_data_LES_coarse, S_data_LES_coarse, u_data_model_coarse, v_data_model_coarse, T_data_model_coarse, S_data_model_coarse +end + +Nt = size(u_data_LES_coarse, 4) +#%% +fig = CairoMakie.Figure(size = (2000, 900)) +axu_LES = CairoMakie.Axis(fig[1, 1], xlabel = "y (m)", ylabel = "z (m)", title = "u (LES) m/s") +axv_LES = CairoMakie.Axis(fig[1, 3], xlabel = "y (m)", ylabel = "z (m)", title = "v (LES) m/s") +axT_LES = CairoMakie.Axis(fig[1, 5], xlabel = "y (m)", ylabel = "z (m)", title = "Temperature (LES) °C") +axS_LES = CairoMakie.Axis(fig[1, 7], xlabel = "y (m)", ylabel = "z (m)", title = "Salinity (LES) psu") +axu_model = CairoMakie.Axis(fig[2, 1], xlabel = "y (m)", ylabel = "z (m)", title = "u (NN closure) (m/s)") +axv_model = CairoMakie.Axis(fig[2, 3], xlabel = "y (m)", ylabel = "z (m)", title = "v (NN closure) (m/s)") +axT_model = CairoMakie.Axis(fig[2, 5], xlabel = "y (m)", ylabel = "z (m)", title = "Temperature (NN closure) °C") +axS_model = CairoMakie.Axis(fig[2, 7], xlabel = "y (m)", ylabel = "z (m)", title = "Salinity (NN closure) psu") +axΔu = CairoMakie.Axis(fig[3, 1], xlabel = "y (m)", ylabel = "z (m)", title = "u (LES) - u(NN closure) (m/s)") +axΔv = CairoMakie.Axis(fig[3, 3], xlabel = "y (m)", ylabel = "z (m)", title = "v (LES) - v(NN closure) (m/s)") +axΔT = CairoMakie.Axis(fig[3, 5], xlabel = "y (m)", ylabel = "z (m)", title = "Temperature (LES) - Temperature (NN closure) (°C)") +axΔS = CairoMakie.Axis(fig[3, 7], xlabel = "y (m)", ylabel = "z (m)", title = "Salinity (LES) - Salinity (NN closure) (psu)") + +n = Observable(1) +u_LESₙ = @lift u_data_LES_coarse[1, :, :, $n] +v_LESₙ = @lift v_data_LES_coarse[1, :, :, $n] +T_LESₙ = @lift T_data_LES_coarse[1, :, :, $n] +S_LESₙ = @lift S_data_LES_coarse[1, :, :, $n] + +u_modelₙ = @lift u_data_model_coarse[1, :, :, $n] +v_modelₙ = @lift v_data_model_coarse[1, :, :, $n] +T_modelₙ = @lift T_data_model_coarse[1, :, :, $n] +S_modelₙ = @lift S_data_model_coarse[1, :, :, $n] + +Δuₙ = @lift $u_LESₙ .- $u_modelₙ +Δvₙ = @lift $v_LESₙ .- $v_modelₙ +ΔTₙ = @lift $T_LESₙ .- $T_modelₙ +ΔSₙ = @lift $S_LESₙ .- $S_modelₙ + +ulim = @lift (-maximum([maximum(abs, $u_LESₙ), 1e-16, maximum(abs, $u_modelₙ)]), + maximum([maximum(abs, $u_LESₙ), 1e-16, maximum(abs, $u_modelₙ)])) +vlim = @lift (-maximum([maximum(abs, $v_LESₙ), 1e-16, maximum(abs, $v_modelₙ)]), + maximum([maximum(abs, $v_LESₙ), 1e-16, maximum(abs, $v_modelₙ)])) +Tlim = (minimum(T_data_LES_coarse[1, :, :, 1]), maximum(T_data_LES_coarse[1, :, :, 1])) +Slim = (minimum(S_data_LES_coarse[1, :, :, 1]), maximum(S_data_LES_coarse[1, :, :, 1])) + +Δulim = @lift (-maximum([maximum(abs, $Δuₙ), 1e-16]), maximum([maximum(abs, $Δuₙ), 1e-16])) +Δvlim = @lift (-maximum([maximum(abs, $Δvₙ), 1e-16]), maximum([maximum(abs, $Δvₙ), 1e-16])) +ΔTlim = @lift (-maximum([maximum(abs, $ΔTₙ), 1e-16]), maximum([maximum(abs, $ΔTₙ), 1e-16])) +ΔSlim = @lift (-maximum([maximum(abs, $ΔSₙ), 1e-16]), maximum([maximum(abs, $ΔSₙ), 1e-16])) + +hu = heatmap!(axu_LES, yC_model, zC_model, u_LESₙ, colormap = :RdBu_9, colorrange = ulim) +hv = heatmap!(axv_LES, yC_model, zC_model, v_LESₙ, colormap = :RdBu_9, colorrange = vlim) +hT = heatmap!(axT_LES, yC_model, zC_model, T_LESₙ, colorrange = Tlim) +hS = heatmap!(axS_LES, yC_model, zC_model, S_LESₙ, colorrange = Slim) + +hu_model = heatmap!(axu_model, yC_model, zC_model, u_modelₙ, colormap = :RdBu_9, colorrange = ulim) +hv_model = heatmap!(axv_model, yC_model, zC_model, v_modelₙ, colormap = :RdBu_9, colorrange = vlim) +hT_model = heatmap!(axT_model, yC_model, zC_model, T_modelₙ, colorrange = Tlim) +hS_model = heatmap!(axS_model, yC_model, zC_model, S_modelₙ, colorrange = Slim) + +hΔu = heatmap!(axΔu, yC_model, zC_model, Δuₙ, colormap = :RdBu_9, colorrange = Δulim) +hΔv = heatmap!(axΔv, yC_model, zC_model, Δvₙ, colormap = :RdBu_9, colorrange = Δvlim) +hΔT = heatmap!(axΔT, yC_model, zC_model, ΔTₙ, colormap = :RdBu_9, colorrange = ΔTlim) +hΔS = heatmap!(axΔS, yC_model, zC_model, ΔSₙ, colormap = :RdBu_9, colorrange = ΔSlim) + +Colorbar(fig[1:2, 2], hu, label = "u (m/s)") +Colorbar(fig[1:2, 4], hv, label = "v (m/s)") +Colorbar(fig[1:2, 6], hT, label = "T (°C)") +Colorbar(fig[1:2, 8], hS, label = "S (psu)") + +Colorbar(fig[3, 2], hΔu, label = "u (m/s)") +Colorbar(fig[3, 4], hΔv, label = "v (m/s)") +Colorbar(fig[3, 6], hΔT, label = "T (°C)") +Colorbar(fig[3, 8], hΔS, label = "S (psu)") + +# display(fig) + +CairoMakie.record(fig, "./LES_NN_2D_sin_cooling_heating_3e-4_23_days_comparison.mp4", 1:Nt, framerate=30) do nn + @info nn + n[] = nn +end + + +#%% \ No newline at end of file From 4d118c20bc0369de914dea1a0284d41ae2d5138a Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 18 Jul 2024 16:57:33 -0400 Subject: [PATCH 024/114] run LES for hald sinusoid cooling --- 2D_model_LES_sin.jl | 225 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 225 insertions(+) create mode 100644 2D_model_LES_sin.jl diff --git a/2D_model_LES_sin.jl b/2D_model_LES_sin.jl new file mode 100644 index 0000000000..3fda63741c --- /dev/null +++ b/2D_model_LES_sin.jl @@ -0,0 +1,225 @@ +#using Pkg +using Oceananigans +using Printf +using Statistics + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +import SeawaterPolynomials.TEOS10: s, ΔS, Sₐᵤ +s(Sᴬ::Number) = Sᴬ + ΔS >= 0 ? √((Sᴬ + ΔS) / Sₐᵤ) : NaN +using Glob + +# Architecture +model_architecture = GPU() + +# number of grid points +Ny = 20000 +Nz = 256 + +const Ly = 40kilometers +const Lz = 512meters + +grid = RectilinearGrid(model_architecture, + topology = (Flat, Bounded, Bounded), + size = (Ny, Nz), + halo = (5, 5), + y = (0, Ly), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const dTdz = 0.014 +const dSdz = 0.0021 + +const T_surface = 20.0 +const S_surface = 36.6 +const max_temperature_flux = 3e-4 + +FILE_DIR = "./LES/NN_2D_channel_sin_cooling_$(max_temperature_flux)_LES" +mkpath(FILE_DIR) + +@inline function temperature_flux(y, t) + return max_temperature_flux * sin(π * y / Ly) +end + +T_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(temperature_flux)) + +##### +##### Coriolis +##### + +const f₀ = 8e-5 +coriolis = FPlane(f=f₀) + +##### +##### Forcing and initial condition +##### +T_initial(y, z) = dTdz * z + T_surface +S_initial(y, z) = dSdz * z + S_surface + +##### +##### Model building +##### + +@info "Building a model..." + +model = NonhydrostaticModel(; grid = grid, + advection = WENO(order=9), + coriolis = coriolis, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + tracers = (:T, :S), + timestepper = :RungeKutta3, + closure = nothing, + boundary_conditions = (; T=T_bcs)) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(y, z) = T_initial(y, z) + 1e-6 * noise(z) +S_initial_noisy(y, z) = S_initial(y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +##### +##### Simulation building +##### +simulation = Simulation(model, Δt = 0.1, stop_time = 30days) + +# add timestep wizard callback +wizard = TimeStepWizard(cfl=0.6, max_change=1.05, max_Δt=20minutes) +simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(sim.model.velocities.u), + maximum(sim.model.velocities.v), + maximum(sim.model.tracers.T), + maximum(sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(1000)) + +##### +##### Diagnostics +##### + +u, w = model.velocities.u, model.velocities.w +v = @at (Center, Center, Center) model.velocities.v +T, S = model.tracers.T, model.tracers.S + +outputs = (; u, v, w, T, S) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:jld2] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields.jld2", + schedule = TimeInterval(1hour)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(1day), + prefix = "$(FILE_DIR)/checkpointer", + overwrite_existing = true) + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +checkpointers = glob("$(FILE_DIR)/checkpointer_iteration*.jld2") +if !isempty(checkpointers) + rm.(checkpointers) +end + +# ##### +# ##### Visualization +# ##### +#%% +using CairoMakie + + +u_data = FieldTimeSeries("./NN_2D_channel_sin_cooling_$(max_temperature_flux)_LES.jld2", "u", backend=OnDisk()) +v_data = FieldTimeSeries("./NN_2D_channel_sin_cooling_$(max_temperature_flux)_LES.jld2", "v", backend=OnDisk()) +T_data = FieldTimeSeries("./NN_2D_channel_sin_cooling_$(max_temperature_flux)_LES.jld2", "T", backend=OnDisk()) +S_data = FieldTimeSeries("./NN_2D_channel_sin_cooling_$(max_temperature_flux)_LES.jld2", "S", backend=OnDisk()) + +yC = ynodes(T_data.grid, Center()) +yF = ynodes(T_data.grid, Face()) + +zC = znodes(T_data.grid, Center()) +zF = znodes(T_data.grid, Face()) + +Nt = length(T_data.times) +#%% +fig = Figure(size = (1500, 900)) +axu = CairoMakie.Axis(fig[1, 1], xlabel = "y (m)", ylabel = "z (m)", title = "u") +axv = CairoMakie.Axis(fig[1, 3], xlabel = "y (m)", ylabel = "z (m)", title = "v") +axT = CairoMakie.Axis(fig[2, 1], xlabel = "y (m)", ylabel = "z (m)", title = "Temperature") +axS = CairoMakie.Axis(fig[2, 3], xlabel = "y (m)", ylabel = "z (m)", title = "Salinity") +n = Obeservable(1) + +uₙ = @lift interior(u_data[$n], 1, :, :) +vₙ = @lift interior(v_data[$n], 1, :, :) +Tₙ = @lift interior(T_data[$n], 1, :, :) +Sₙ = @lift interior(S_data[$n], 1, :, :) + +ulim = @lift (-maximum([maximum(abs, $uₙ), 1e-16]), maximum([maximum(abs, $uₙ), 1e-16])) +vlim = @lift (-maximum([maximum(abs, $vₙ), 1e-16]), maximum([maximum(abs, $vₙ), 1e-16])) +Tlim = (minimum(interior(T_data[1])), maximum(interior(T_data[1]))) +Slim = (minimum(interior(S_data[1])), maximum(interior(S_data[1]))) + +title_str = @lift "Time: $(round(T_data.times[$n] / 86400, digits=2)) days" +Label(fig[0, :], title_str, tellwidth = false) + +hu = heatmap!(axu, yC, zC, uₙ, colormap=:RdBu_9, colorrange=ulim) +hv = heatmap!(axv, yC, zC, vₙ, colormap=:RdBu_9, colorrange=vlim) +hT = heatmap!(axT, yC, zC, Tₙ, colorrange=Tlim) +hS = heatmap!(axS, yC, zC, Sₙ, colorrange=Slim) + +Colorbar(fig[1, 2], hu, label = "(m/s)") +Colorbar(fig[1, 4], hv, label = "(m/s)") +Colorbar(fig[2, 2], hT, label = "(°C)") +Colorbar(fig[2, 4], hS, label = "(psu)") + +CairoMakie.record(fig, "$(FILE_DIR)/2D_sin_cooling_$(max_temperature_flux)_30days.mp4", 1:Nt, framerate=15) do nn + n[] = nn +end + +# display(fig) +#%% \ No newline at end of file From fc081cccd06edfc11b714896a17f8c47f065e804 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 22 Aug 2024 00:50:08 +0800 Subject: [PATCH 025/114] run 3D simulation with limited extent --- 3D_model_LES_sin.jl | 228 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 228 insertions(+) create mode 100644 3D_model_LES_sin.jl diff --git a/3D_model_LES_sin.jl b/3D_model_LES_sin.jl new file mode 100644 index 0000000000..6308740dda --- /dev/null +++ b/3D_model_LES_sin.jl @@ -0,0 +1,228 @@ +#using Pkg +using Oceananigans +using Printf +using Statistics + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +import SeawaterPolynomials.TEOS10: s, ΔS, Sₐᵤ +s(Sᴬ::Number) = Sᴬ + ΔS >= 0 ? √((Sᴬ + ΔS) / Sₐᵤ) : NaN +using Glob + +# Architecture +model_architecture = GPU() + +# number of grid points +Nx = 125 +Ny = 4000 +Nz = 250 + +const Lx = 250metres +const Ly = 8kilometers +const Lz = 500meters + +grid = RectilinearGrid(model_architecture, + topology = (Periodic, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (5, 5, 5), + x = (0, Lx), + y = (0, Ly), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const dTdz = 0.014 +const dSdz = 0.0021 + +const T_surface = 20.0 +const S_surface = 36.6 +const max_temperature_flux = 3e-4 + +FILE_DIR = "./LES/NN_3D_channel_sin_cooling_$(max_temperature_flux)_LES_Lx_$(Lx)_Ly_$(Ly)_Lz_$(Lz)_Nx_$(Nx)_Ny_$(Ny)_Nz_$(Nz)" +mkpath(FILE_DIR) + +@inline function temperature_flux(y, t) + return max_temperature_flux * sin(π * y / Ly) +end + +T_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(temperature_flux)) + +##### +##### Coriolis +##### + +const f₀ = 8e-5 +coriolis = FPlane(f=f₀) + +##### +##### Forcing and initial condition +##### +T_initial(x, y, z) = dTdz * z + T_surface +S_initial(x, y, z) = dSdz * z + S_surface + +##### +##### Model building +##### + +@info "Building a model..." + +model = NonhydrostaticModel(; grid = grid, + advection = WENO(order=9), + coriolis = coriolis, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + tracers = (:T, :S), + timestepper = :RungeKutta3, + closure = nothing, + boundary_conditions = (; T=T_bcs)) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +##### +##### Simulation building +##### +simulation = Simulation(model, Δt = 0.1, stop_time = 10days) + +# add timestep wizard callback +wizard = TimeStepWizard(cfl=0.6, max_change=1.05, max_Δt=20minutes) +simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(sim.model.velocities.u), + maximum(sim.model.velocities.v), + maximum(sim.model.tracers.T), + maximum(sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(1000)) + +##### +##### Diagnostics +##### + +u, w = model.velocities.u, model.velocities.w +v = @at (Center, Center, Center) model.velocities.v +T, S = model.tracers.T, model.tracers.S + +outputs = (; u, v, w, T, S) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:jld2] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields.jld2", + schedule = TimeInterval(1hour)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(0.5day), + prefix = "$(FILE_DIR)/checkpointer", + overwrite_existing = true) + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +checkpointers = glob("$(FILE_DIR)/checkpointer_iteration*.jld2") +if !isempty(checkpointers) + rm.(checkpointers) +end + +# ##### +# ##### Visualization +# ##### +#%% +using CairoMakie + + +u_data = FieldTimeSeries("./NN_2D_channel_sin_cooling_$(max_temperature_flux)_LES.jld2", "u", backend=OnDisk()) +v_data = FieldTimeSeries("./NN_2D_channel_sin_cooling_$(max_temperature_flux)_LES.jld2", "v", backend=OnDisk()) +T_data = FieldTimeSeries("./NN_2D_channel_sin_cooling_$(max_temperature_flux)_LES.jld2", "T", backend=OnDisk()) +S_data = FieldTimeSeries("./NN_2D_channel_sin_cooling_$(max_temperature_flux)_LES.jld2", "S", backend=OnDisk()) + +yC = ynodes(T_data.grid, Center()) +yF = ynodes(T_data.grid, Face()) + +zC = znodes(T_data.grid, Center()) +zF = znodes(T_data.grid, Face()) + +Nt = length(T_data.times) +#%% +fig = Figure(size = (1500, 900)) +axu = CairoMakie.Axis(fig[1, 1], xlabel = "y (m)", ylabel = "z (m)", title = "u") +axv = CairoMakie.Axis(fig[1, 3], xlabel = "y (m)", ylabel = "z (m)", title = "v") +axT = CairoMakie.Axis(fig[2, 1], xlabel = "y (m)", ylabel = "z (m)", title = "Temperature") +axS = CairoMakie.Axis(fig[2, 3], xlabel = "y (m)", ylabel = "z (m)", title = "Salinity") +n = Obeservable(1) + +uₙ = @lift interior(u_data[$n], 1, :, :) +vₙ = @lift interior(v_data[$n], 1, :, :) +Tₙ = @lift interior(T_data[$n], 1, :, :) +Sₙ = @lift interior(S_data[$n], 1, :, :) + +ulim = @lift (-maximum([maximum(abs, $uₙ), 1e-16]), maximum([maximum(abs, $uₙ), 1e-16])) +vlim = @lift (-maximum([maximum(abs, $vₙ), 1e-16]), maximum([maximum(abs, $vₙ), 1e-16])) +Tlim = (minimum(interior(T_data[1])), maximum(interior(T_data[1]))) +Slim = (minimum(interior(S_data[1])), maximum(interior(S_data[1]))) + +title_str = @lift "Time: $(round(T_data.times[$n] / 86400, digits=2)) days" +Label(fig[0, :], title_str, tellwidth = false) + +hu = heatmap!(axu, yC, zC, uₙ, colormap=:RdBu_9, colorrange=ulim) +hv = heatmap!(axv, yC, zC, vₙ, colormap=:RdBu_9, colorrange=vlim) +hT = heatmap!(axT, yC, zC, Tₙ, colorrange=Tlim) +hS = heatmap!(axS, yC, zC, Sₙ, colorrange=Slim) + +Colorbar(fig[1, 2], hu, label = "(m/s)") +Colorbar(fig[1, 4], hv, label = "(m/s)") +Colorbar(fig[2, 2], hT, label = "(°C)") +Colorbar(fig[2, 4], hS, label = "(psu)") + +CairoMakie.record(fig, "$(FILE_DIR)/3D_sin_cooling_$(max_temperature_flux)_Lx_$(Lx)_Ly_$(Ly)_Lz_$(Lz)_Nx_$(Nx)_Ny_$(Ny)_Nz_$(Nz)_10days.mp4", 1:Nt, framerate=10) do nn + n[] = nn +end + +# display(fig) +#%% \ No newline at end of file From a78109e4d48c4675ecacfa77bf3b9219526c6dc6 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 22 Aug 2024 00:55:26 +0800 Subject: [PATCH 026/114] add sponge at bottom --- 3D_model_LES_sin.jl | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/3D_model_LES_sin.jl b/3D_model_LES_sin.jl index 6308740dda..66a2b4a46f 100644 --- a/3D_model_LES_sin.jl +++ b/3D_model_LES_sin.jl @@ -67,6 +67,17 @@ coriolis = FPlane(f=f₀) T_initial(x, y, z) = dTdz * z + T_surface S_initial(x, y, z) = dSdz * z + S_surface +damping_rate = 1/15minute + +T_target(x, y, z, t) = T_initial(x, y, z) +S_target(x, y, z, t) = S_initial(x, y, z) + +bottom_mask = GaussianMask{:z}(center=-grid.Lz, width=grid.Lz/10) + +uvw_sponge = Relaxation(rate=damping_rate, mask=bottom_mask) +T_sponge = Relaxation(rate=damping_rate, mask=bottom_mask, target=T_target) +S_sponge = Relaxation(rate=damping_rate, mask=bottom_mask, target=S_target) + ##### ##### Model building ##### @@ -80,7 +91,8 @@ model = NonhydrostaticModel(; grid = grid, tracers = (:T, :S), timestepper = :RungeKutta3, closure = nothing, - boundary_conditions = (; T=T_bcs)) + boundary_conditions = (; T=T_bcs), + forcing = (u=uvw_sponge, v=uvw_sponge, w=uvw_sponge, T=T_sponge, S=S_sponge)) @info "Built $model." @@ -176,7 +188,6 @@ end #%% using CairoMakie - u_data = FieldTimeSeries("./NN_2D_channel_sin_cooling_$(max_temperature_flux)_LES.jld2", "u", backend=OnDisk()) v_data = FieldTimeSeries("./NN_2D_channel_sin_cooling_$(max_temperature_flux)_LES.jld2", "v", backend=OnDisk()) T_data = FieldTimeSeries("./NN_2D_channel_sin_cooling_$(max_temperature_flux)_LES.jld2", "T", backend=OnDisk()) From ca8d78c49d4faf0216d293489ed0cac291b4584b Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 22 Aug 2024 01:02:51 +0800 Subject: [PATCH 027/114] fix metres to meters --- 3D_model_LES_sin.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3D_model_LES_sin.jl b/3D_model_LES_sin.jl index 66a2b4a46f..0837df3b0f 100644 --- a/3D_model_LES_sin.jl +++ b/3D_model_LES_sin.jl @@ -21,7 +21,7 @@ Nx = 125 Ny = 4000 Nz = 250 -const Lx = 250metres +const Lx = 250meters const Ly = 8kilometers const Lz = 500meters From 78e7a673a72975c2832e707957c9c88b5e828f92 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 22 Aug 2024 01:09:40 +0800 Subject: [PATCH 028/114] fix temperature flux --- 3D_model_LES_sin.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3D_model_LES_sin.jl b/3D_model_LES_sin.jl index 0837df3b0f..e10b28340e 100644 --- a/3D_model_LES_sin.jl +++ b/3D_model_LES_sin.jl @@ -48,7 +48,7 @@ const max_temperature_flux = 3e-4 FILE_DIR = "./LES/NN_3D_channel_sin_cooling_$(max_temperature_flux)_LES_Lx_$(Lx)_Ly_$(Ly)_Lz_$(Lz)_Nx_$(Nx)_Ny_$(Ny)_Nz_$(Nz)" mkpath(FILE_DIR) -@inline function temperature_flux(y, t) +@inline function temperature_flux(x, y, t) return max_temperature_flux * sin(π * y / Ly) end From eaf2183b8dab953be5e8875d3e7ab2fcc5fb2a92 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 22 Aug 2024 01:17:34 +0800 Subject: [PATCH 029/114] Calculate average velocities and tracers in 3D model LES simulation --- 3D_model_LES_sin.jl | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/3D_model_LES_sin.jl b/3D_model_LES_sin.jl index e10b28340e..6915d540fc 100644 --- a/3D_model_LES_sin.jl +++ b/3D_model_LES_sin.jl @@ -146,7 +146,14 @@ u, w = model.velocities.u, model.velocities.w v = @at (Center, Center, Center) model.velocities.v T, S = model.tracers.T, model.tracers.S -outputs = (; u, v, w, T, S) +ubar = Average(u, dims=1) +vbar = Average(v, dims=1) +wbar = Average(w, dims=1) +Tbar = Average(T, dims=1) +Sbar = Average(S, dims=1) + +# outputs = (; u, v, w, T, S) +outputs = (; ubar, vbar, wbar, Tbar, Sbar) ##### ##### Build checkpointer and output writer From 7e144fffeed2d6d506c416dfe892d12e42141ba9 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 22 Aug 2024 01:34:51 +0800 Subject: [PATCH 030/114] reduce size of model --- 3D_model_LES_sin.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/3D_model_LES_sin.jl b/3D_model_LES_sin.jl index 6915d540fc..c886878ceb 100644 --- a/3D_model_LES_sin.jl +++ b/3D_model_LES_sin.jl @@ -18,11 +18,11 @@ model_architecture = GPU() # number of grid points Nx = 125 -Ny = 4000 +Ny = 3000 Nz = 250 const Lx = 250meters -const Ly = 8kilometers +const Ly = 6kilometers const Lz = 500meters grid = RectilinearGrid(model_architecture, From ba7f58dc4147c0d87822f972ef432fb64b7748da Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 29 Aug 2024 10:58:32 +0800 Subject: [PATCH 031/114] run double gyre with physical closure --- NN_doublegyre_model.jl | 536 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 536 insertions(+) create mode 100644 NN_doublegyre_model.jl diff --git a/NN_doublegyre_model.jl b/NN_doublegyre_model.jl new file mode 100644 index 0000000000..a154e2fae2 --- /dev/null +++ b/NN_doublegyre_model.jl @@ -0,0 +1,536 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +# include("NN_closure_global.jl") +# include("xin_kai_vertical_diffusivity_local.jl") +include("xin_kai_vertical_diffusivity_2Pr.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes +import SeawaterPolynomials.TEOS10: s, ΔS, Sₐᵤ +s(Sᴬ::Number) = Sᴬ + ΔS >= 0 ? √((Sᴬ + ΔS) / Sₐᵤ) : NaN + +#%% +# Architecture +model_architecture = GPU() + +# nn_closure = NNFluxClosure(model_architecture) +# base_closure = XinKaiLocalVerticalDiffusivity() +# closure = (nn_closure, base_closure) + +vertical_base_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +convection_closure = XinKaiVerticalDiffusivity() +closure = (vertical_base_closure, convection_closure) +# closure = vertical_base_closure + +# number of grid points +const Nx = 96 +const Ny = 96 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +# const dSdz = 0.0021 +const dSdz = 0 +const S_surface = 36.6 + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/8days + +##### +##### Forcing and initial condition +##### + +@inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y +@inline surface_T_flux(x, y, t, T) = -μ_T * (T - T_ref(y)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_south - S_north) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = -μ_T * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + # closure = (nn_closure, base_closure), + closure = closure, + # closure = RiBasedVerticalDiffusivity(), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 730days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(20)) + +##### +##### Diagnostics +##### + +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +ν, κ = model.diffusivity_fields[2].κᵘ, model.diffusivity_fields[2].κᶜ +Ri = model.diffusivity_fields[2].Ri +# wT, wS = model.diffusivity_fields[2].wT, model.diffusivity_fields[2].wS + +# outputs = (; u, v, w, T, S, ν, κ, Ri, wT, wS) +# outputs = (; u, v, w, T, S, ν, κ, Ri) +outputs = (; u, v, w, T, S) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + # filename = "NN_closure_2D_channel_NDE_FC_Qb_18simnew_2layer_128_relu_2Pr", + filename = "doublegyre_xin_kai_vertical_diffusivity_2Pr_xy", + indices = (:, :, Nz), + schedule = TimeInterval(1day), + overwrite_existing = true) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + # filename = "NN_closure_2D_channel_NDE_FC_Qb_18simnew_2layer_128_relu_2Pr", + filename = "doublegyre_xin_kai_vertical_diffusivity_2Pr_yz", + indices = (1, :, :), + schedule = TimeInterval(1day), + overwrite_existing = true) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + # filename = "NN_closure_2D_channel_NDE_FC_Qb_18simnew_2layer_128_relu_2Pr", + filename = "doublegyre_xin_kai_vertical_diffusivity_2Pr_xz", + indices = (:, 1, :), + schedule = TimeInterval(1day), + overwrite_existing = true) + +@info "Running the simulation..." + +try + run!(simulation, pickup = false) +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end +#%% +T_xy_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_xy.jld2", "T") +T_xz_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_xz.jld2", "T") +T_yz_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_xy.jld2", "S") +S_xz_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_xz.jld2", "S") +S_yz_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_xy.jld2", "u") +u_xz_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_xz.jld2", "u") +u_yz_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_xy.jld2", "v") +v_xz_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_xz.jld2", "v") +v_yz_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +CairoMakie.record(fig, "./$(dataname)_test.mp4", 1:Nt, framerate=5, px_per_unit=2) do nn + @info nn + n[] = nn +end + +# display(fig) +#%% + +# # ##### +# # ##### Visualization +# # ##### +# using CairoMakie + +# dataname = "NN_closure_doublegyre_NDE_FC_Qb_absf_24simnew_2layer_128_relu_2Pr" +# DATA_DIR = "./$(dataname).jld2" + +# u_data = FieldTimeSeries("$(DATA_DIR)", "u") +# v_data = FieldTimeSeries("$(DATA_DIR)", "v") +# T_data = FieldTimeSeries("$(DATA_DIR)", "T") +# S_data = FieldTimeSeries("$(DATA_DIR)", "S") +# # ν_data = FieldTimeSeries("$(DATA_DIR)", "ν") +# # κ_data = FieldTimeSeries("$(DATA_DIR)", "κ") +# # Ri_data = FieldTimeSeries("$(DATA_DIR)", "Ri") +# # wT_data = FieldTimeSeries("$(DATA_DIR)", "wT") +# # wS_data = FieldTimeSeries("$(DATA_DIR)", "wS") + +# yC = ynodes(T_data.grid, Center()) +# yF = ynodes(T_data.grid, Face()) + +# zC = znodes(T_data.grid, Center()) +# zF = znodes(T_data.grid, Face()) + +# Nt = length(T_data.times) +# #%% +# fig = Figure(size = (1500, 900)) +# axu = CairoMakie.Axis(fig[1, 1], xlabel = "y (m)", ylabel = "z (m)", title = "u (m/s)") +# axv = CairoMakie.Axis(fig[1, 3], xlabel = "y (m)", ylabel = "z (m)", title = "v (m/s)") +# axT = CairoMakie.Axis(fig[2, 1], xlabel = "y (m)", ylabel = "z (m)", title = "Temperature (°C)") +# axS = CairoMakie.Axis(fig[2, 3], xlabel = "y (m)", ylabel = "z (m)", title = "Salinity (psu)") +# n = Observable(1) + +# uₙ = @lift interior(u_data[$n], 45, :, :) +# vₙ = @lift interior(v_data[$n], 45, :, :) +# Tₙ = @lift interior(T_data[$n], 45, :, :) +# Sₙ = @lift interior(S_data[$n], 45, :, :) + +# ulim = @lift (-maximum([maximum(abs, $uₙ), 1e-16]), maximum([maximum(abs, $uₙ), 1e-16])) +# vlim = @lift (-maximum([maximum(abs, $vₙ), 1e-16]), maximum([maximum(abs, $vₙ), 1e-16])) +# Tlim = (minimum(interior(T_data[1])), maximum(interior(T_data[1]))) +# Slim = (minimum(interior(S_data[1])), maximum(interior(S_data[1]))) + +# title_str = @lift "Time: $(round(T_data.times[$n] / 86400, digits=2)) days" +# Label(fig[0, :], title_str, tellwidth = false) + +# hu = heatmap!(axu, yC, zC, uₙ, colormap=:RdBu_9, colorrange=ulim) +# hv = heatmap!(axv, yF, zC, vₙ, colormap=:RdBu_9, colorrange=vlim) +# hT = heatmap!(axT, yC, zC, Tₙ, colorrange=Tlim) +# hS = heatmap!(axS, yC, zC, Sₙ, colorrange=Slim) + +# Colorbar(fig[1, 2], hu, label = "u (m/s)") +# Colorbar(fig[1, 4], hv, label = "v (m/s)") +# Colorbar(fig[2, 2], hT, label = "T (°C)") +# Colorbar(fig[2, 4], hS, label = "S (psu)") + +# CairoMakie.record(fig, "./$(dataname)_test.mp4", 1:Nt, framerate=10) do nn +# n[] = nn +# end + +# display(fig) +# #%% +# fig = Figure(size = (1920, 1080)) +# axu = CairoMakie.Axis(fig[1, 1], xlabel = "y (m)", ylabel = "z (m)", title = "u") +# axv = CairoMakie.Axis(fig[1, 3], xlabel = "y (m)", ylabel = "z (m)", title = "v") +# axT = CairoMakie.Axis(fig[2, 1], xlabel = "y (m)", ylabel = "z (m)", title = "Temperature") +# axS = CairoMakie.Axis(fig[2, 3], xlabel = "y (m)", ylabel = "z (m)", title = "Salinity") +# axν = CairoMakie.Axis(fig[1, 5], xlabel = "y (m)", ylabel = "z (m)", title = "Viscosity (log10 scale)") +# axκ = CairoMakie.Axis(fig[2, 5], xlabel = "y (m)", ylabel = "z (m)", title = "Diffusivity (log10 scale)") +# axRi = CairoMakie.Axis(fig[3, 5], xlabel = "y (m)", ylabel = "z (m)", title = "Richardson number") +# axwT = CairoMakie.Axis(fig[3, 1], xlabel = "y (m)", ylabel = "z (m)", title = "wT(NN)") +# axwS = CairoMakie.Axis(fig[3, 3], xlabel = "y (m)", ylabel = "z (m)", title = "wS(NN)") + +# n = Observable(1) + +# uₙ = @lift interior(u_data[$n], 1, :, :) +# vₙ = @lift interior(v_data[$n], 1, :, :) +# Tₙ = @lift interior(T_data[$n], 1, :, :) +# Sₙ = @lift interior(S_data[$n], 1, :, :) +# νₙ = @lift log10.(interior(ν_data[$n], 1, :, :)) +# κₙ = @lift log10.(interior(κ_data[$n], 1, :, :)) +# Riₙ = @lift clamp.(interior(Ri_data[$n], 1, :, :), -20, 20) +# wTₙ = @lift interior(wT_data[$n], 1, :, :) +# wSₙ = @lift interior(wS_data[$n], 1, :, :) + +# ulim = @lift (-maximum([maximum(abs, $uₙ), 1e-7]), maximum([maximum(abs, $uₙ), 1e-7])) +# vlim = @lift (-maximum([maximum(abs, $vₙ), 1e-7]), maximum([maximum(abs, $vₙ), 1e-7])) +# Tlim = (minimum(interior(T_data[1])), maximum(interior(T_data[1]))) +# Slim = (minimum(interior(S_data[1])), maximum(interior(S_data[1]))) +# νlim = (-6, 2) +# κlim = (-6, 2) +# wTlim = @lift (-maximum([maximum(abs, $wTₙ), 1e-7]), maximum([maximum(abs, $wTₙ), 1e-7])) +# wSlim = @lift (-maximum([maximum(abs, $wSₙ), 1e-7]), maximum([maximum(abs, $wSₙ), 1e-7])) + +# title_str = @lift "Time: $(round(T_data.times[$n] / 86400, digits=2)) days" +# Label(fig[0, :], title_str, tellwidth = false) + +# hu = heatmap!(axu, yC, zC, uₙ, colormap=:RdBu_9, colorrange=ulim) +# hv = heatmap!(axv, yF, zC, vₙ, colormap=:RdBu_9, colorrange=vlim) +# hT = heatmap!(axT, yC, zC, Tₙ, colorrange=Tlim) +# hS = heatmap!(axS, yC, zC, Sₙ, colorrange=Slim) +# hν = heatmap!(axν, yC, zC, νₙ, colorrange=νlim) +# hκ = heatmap!(axκ, yC, zC, κₙ, colorrange=κlim) +# hRi = heatmap!(axRi, yC, zF, Riₙ, colormap=:RdBu_9, colorrange=(-20, 20)) +# hwT = heatmap!(axwT, yC, zF, wTₙ, colormap=:RdBu_9, colorrange=wTlim) +# hwS = heatmap!(axwS, yC, zF, wSₙ, colormap=:RdBu_9, colorrange=wSlim) + +# cbu = Colorbar(fig[1, 2], hu, label = "(m/s)") +# cbv = Colorbar(fig[1, 4], hv, label = "(m/s)") +# cbT = Colorbar(fig[2, 2], hT, label = "(°C)") +# cbS = Colorbar(fig[2, 4], hS, label = "(psu)") +# cbν = Colorbar(fig[1, 6], hν, label = "(m²/s)") +# cbκ = Colorbar(fig[2, 6], hκ, label = "(m²/s)") +# cbRi = Colorbar(fig[3, 6], hRi) +# cbwT = Colorbar(fig[3, 2], hwT, label = "(m/s °C)") +# cbwS = Colorbar(fig[3, 4], hwS, label = "(m/s psu)") + +# tight_ticklabel_spacing!(cbu) +# tight_ticklabel_spacing!(cbv) +# tight_ticklabel_spacing!(cbT) +# tight_ticklabel_spacing!(cbS) +# tight_ticklabel_spacing!(cbν) +# tight_ticklabel_spacing!(cbκ) +# tight_ticklabel_spacing!(cbRi) +# tight_ticklabel_spacing!(cbwT) +# tight_ticklabel_spacing!(cbwS) + +# CairoMakie.record(fig, "./$(dataname)_2D_sin_cooling_heating_23days_fluxes.mp4", 1:Nt, framerate=30) do nn +# n[] = nn +# end +# #%% \ No newline at end of file From fcb466a9259875872c458eb175beb733928302f0 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 29 Aug 2024 10:59:57 +0800 Subject: [PATCH 032/114] rename file --- NN_doublegyre_model.jl => physicalclosure_doublegyre_model.jl | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename NN_doublegyre_model.jl => physicalclosure_doublegyre_model.jl (100%) diff --git a/NN_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl similarity index 100% rename from NN_doublegyre_model.jl rename to physicalclosure_doublegyre_model.jl From b63e7d6d90118fe64de002db86e34713813ed0d3 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 29 Aug 2024 00:01:43 -0400 Subject: [PATCH 033/114] fix S forcing --- physicalclosure_doublegyre_model.jl | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index a154e2fae2..00cee83fda 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -64,10 +64,6 @@ const T_south = 30 const T_mid = (T_north + T_south) / 2 const ΔT = T_south - T_north -# const dSdz = 0.0021 -const dSdz = 0 -const S_surface = 36.6 - const S_north = 34 const S_south = 37 const S_mid = (S_north + S_south) / 2 @@ -108,7 +104,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) -@inline S_ref(y) = (S_south - S_north) / Ly * y + S_mid +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid @inline S_initial(x, y, z) = S_ref(y) @inline surface_S_flux(x, y, t, S) = -μ_T * (S - S_ref(y)) surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) From 67e011c7f2aa7ff5dfe43541df5d9495b097a27f Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 29 Aug 2024 01:05:05 -0400 Subject: [PATCH 034/114] use RiBasedVerticalDiffusivity as closure --- physicalclosure_doublegyre_model.jl | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index 00cee83fda..72436f1340 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -3,7 +3,7 @@ using Oceananigans # include("NN_closure_global.jl") # include("xin_kai_vertical_diffusivity_local.jl") -include("xin_kai_vertical_diffusivity_2Pr.jl") +# include("xin_kai_vertical_diffusivity_2Pr.jl") ENV["GKSwstype"] = "100" @@ -32,7 +32,8 @@ model_architecture = GPU() # closure = (nn_closure, base_closure) vertical_base_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) -convection_closure = XinKaiVerticalDiffusivity() +# convection_closure = XinKaiVerticalDiffusivity() +convection_closure = RiBasedVerticalDiffusivity() closure = (vertical_base_closure, convection_closure) # closure = vertical_base_closure @@ -190,8 +191,8 @@ simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterv u, v, w = model.velocities T, S = model.tracers.T, model.tracers.S -ν, κ = model.diffusivity_fields[2].κᵘ, model.diffusivity_fields[2].κᶜ -Ri = model.diffusivity_fields[2].Ri +# ν, κ = model.diffusivity_fields[2].κᵘ, model.diffusivity_fields[2].κᶜ +# Ri = model.diffusivity_fields[2].Ri # wT, wS = model.diffusivity_fields[2].wT, model.diffusivity_fields[2].wS # outputs = (; u, v, w, T, S, ν, κ, Ri, wT, wS) From 55d635a176366bfd77d2f772d24761c80c1d40df Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 29 Aug 2024 13:00:09 -0400 Subject: [PATCH 035/114] fix sign errors in tracer fluxes --- physicalclosure_doublegyre_model.jl | 38 ++++++++++++++--------------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index 72436f1340..08db4035b3 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -101,13 +101,13 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), west = ValueBoundaryCondition(0)) @inline T_ref(y) = T_mid - ΔT / Ly * y -@inline surface_T_flux(x, y, t, T) = -μ_T * (T - T_ref(y)) +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) @inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid @inline S_initial(x, y, z) = S_ref(y) -@inline surface_S_flux(x, y, t, S) = -μ_T * (S - S_ref(y)) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) @@ -183,7 +183,7 @@ function print_progress(sim) return nothing end -simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(20)) +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) ##### ##### Diagnostics @@ -204,21 +204,21 @@ outputs = (; u, v, w, T, S) ##### simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, # filename = "NN_closure_2D_channel_NDE_FC_Qb_18simnew_2layer_128_relu_2Pr", - filename = "doublegyre_xin_kai_vertical_diffusivity_2Pr_xy", + filename = "doublegyre_Ri_based_vertical_diffusivity_2Pr_xy", indices = (:, :, Nz), schedule = TimeInterval(1day), overwrite_existing = true) simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, # filename = "NN_closure_2D_channel_NDE_FC_Qb_18simnew_2layer_128_relu_2Pr", - filename = "doublegyre_xin_kai_vertical_diffusivity_2Pr_yz", + filename = "doublegyre_Ri_based_vertical_diffusivity_2Pr_yz", indices = (1, :, :), schedule = TimeInterval(1day), overwrite_existing = true) simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, # filename = "NN_closure_2D_channel_NDE_FC_Qb_18simnew_2layer_128_relu_2Pr", - filename = "doublegyre_xin_kai_vertical_diffusivity_2Pr_xz", + filename = "doublegyre_Ri_based_vertical_diffusivity_2Pr_xz", indices = (:, 1, :), schedule = TimeInterval(1day), overwrite_existing = true) @@ -232,21 +232,21 @@ catch err showerror(stdout, err) end #%% -T_xy_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_xy.jld2", "T") -T_xz_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_xz.jld2", "T") -T_yz_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_yz.jld2", "T") +T_xy_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_xy.jld2", "T") +T_xz_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_xz.jld2", "T") +T_yz_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_yz.jld2", "T") -S_xy_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_xy.jld2", "S") -S_xz_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_xz.jld2", "S") -S_yz_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_yz.jld2", "S") +S_xy_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_xy.jld2", "S") +S_xz_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_xz.jld2", "S") +S_yz_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_yz.jld2", "S") -u_xy_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_xy.jld2", "u") -u_xz_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_xz.jld2", "u") -u_yz_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_yz.jld2", "u") +u_xy_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_xy.jld2", "u") +u_xz_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_xz.jld2", "u") +u_yz_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_yz.jld2", "u") -v_xy_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_xy.jld2", "v") -v_xz_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_xz.jld2", "v") -v_yz_data = FieldTimeSeries("doublegyre_xin_kai_vertical_diffusivity_2Pr_yz.jld2", "v") +v_xy_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_xy.jld2", "v") +v_xz_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_xz.jld2", "v") +v_yz_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_yz.jld2", "v") times = T_xy_data.times ./ 24 ./ 60^2 Nt = length(times) @@ -392,7 +392,7 @@ zlims!(axS, (-Lz, 0)) zlims!(axu, (-Lz, 0)) zlims!(axv, (-Lz, 0)) -CairoMakie.record(fig, "./$(dataname)_test.mp4", 1:Nt, framerate=5, px_per_unit=2) do nn +CairoMakie.record(fig, "./doublegyre_Ri_based_vertical_diffusivity_2Pr.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn @info nn n[] = nn end From 8319ba8002487c360159da4df8d3620d101c206e Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 29 Aug 2024 15:21:20 -0400 Subject: [PATCH 036/114] fix boundary conditions, initialize trivial model first --- physicalclosure_doublegyre_model.jl | 26 +++++++++++++++++++++++--- 1 file changed, 23 insertions(+), 3 deletions(-) diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index 08db4035b3..f926aa0e62 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -38,8 +38,8 @@ closure = (vertical_base_closure, convection_closure) # closure = vertical_base_closure # number of grid points -const Nx = 96 -const Ny = 96 +const Nx = 100 +const Ny = 100 const Nz = 200 const Δz = 8meters @@ -93,10 +93,14 @@ v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, bottom = u_drag_bc, north = ValueBoundaryCondition(0), - south = ValueBoundaryCondition(0)) + south = ValueBoundaryCondition(0), + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), bottom = v_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0), east = ValueBoundaryCondition(0), west = ValueBoundaryCondition(0)) @@ -122,6 +126,22 @@ coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) @info "Building a model..." +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + model = HydrostaticFreeSurfaceModel( grid = grid, free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), From 229feeb399d170a50457d4c0704fdee9e93e480a Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 29 Aug 2024 15:38:22 -0400 Subject: [PATCH 037/114] update boundary conditions, initial state --- physicalclosure_doublegyre_model.jl | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index f926aa0e62..01951076cd 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -78,7 +78,7 @@ const μ_T = 1/8days ##### Forcing and initial condition ##### -@inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) +@inline T_initial(x, y, z) = T_south @inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) @@ -93,14 +93,10 @@ v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, bottom = u_drag_bc, north = ValueBoundaryCondition(0), - south = ValueBoundaryCondition(0), - east = ValueBoundaryCondition(0), - west = ValueBoundaryCondition(0)) + south = ValueBoundaryCondition(0)) v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), bottom = v_drag_bc, - north = ValueBoundaryCondition(0), - south = ValueBoundaryCondition(0), east = ValueBoundaryCondition(0), west = ValueBoundaryCondition(0)) From a85327597047abb20d995afe1340c56f5ff44b4f Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 29 Aug 2024 17:40:01 -0400 Subject: [PATCH 038/114] fix tyop_buoyancy_flux bug --- NN_closure_global.jl | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/NN_closure_global.jl b/NN_closure_global.jl index 3f9dcade77..63021bcd52 100644 --- a/NN_closure_global.jl +++ b/NN_closure_global.jl @@ -68,6 +68,7 @@ Adapt.adapt_structure(to, nn :: NN) = function NNFluxClosure(arch) dev = ifelse(arch == GPU(), gpu_device(), cpu_device()) + # nn_path = "./NDE_FC_Qb_absf_24simnew_2layer_128_relu_2Pr_model.jld2" nn_path = "./NDE_FC_Qb_18simnew_2layer_128_relu_2Pr_model.jld2" ps, sts, scaling_params, wT_model, wS_model = jldopen(nn_path, "r") do file @@ -124,7 +125,7 @@ function compute_diffusivities!(diffusivities, closure::NNFluxClosure, model; pa wT.data.parent .= dropdims(closure.wT(input.parent), dims=1) wS.data.parent .= dropdims(closure.wS(input.parent), dims=1) - launch!(arch, grid, kp, _rescale_nn_fluxes!, diffusivities, grid, closure, tracers, buoyancy, top_tracer_bcs, clock) + launch!(arch, grid, kp, _rescale_nn_fluxes!, diffusivities, grid, closure, tracers, velocities, buoyancy, top_tracer_bcs, clock) launch!(arch, grid, kp, _adjust_nn_bottom_fluxes!, diffusivities, grid, closure) return nothing end @@ -149,14 +150,15 @@ end @inbounds input[8, i, j, k] = ∂σᵢ = scaling.∂ρ∂z(-ρ₀ * ∂z_b(i, j, k, grid, buoyancy, tracers) / g) @inbounds input[9, i, j, k] = ∂σᵢ₊₁ = scaling.∂ρ∂z(-ρ₀ * ∂z_b(i, j, k+1, grid, buoyancy, tracers) / g) - @inbounds input[10, i, j, k] = Jᵇ = scaling.wb(top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, tracers)) + @inbounds input[10, i, j, k] = Jᵇ = scaling.wb(top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers))) + # @inbounds input[11, i, j, k] = fᶜᶜ = scaling.f(abs(fᶜᶜᵃ(i, j, k, grid, coriolis))) @inbounds input[11, i, j, k] = fᶜᶜ = scaling.f(fᶜᶜᵃ(i, j, k, grid, coriolis)) end -@kernel function _rescale_nn_fluxes!(diffusivities, grid, closure::NNFluxClosure, tracers, buoyancy, top_tracer_bcs, clock) +@kernel function _rescale_nn_fluxes!(diffusivities, grid, closure::NNFluxClosure, tracers, velocities, buoyancy, top_tracer_bcs, clock) i, j, k = @index(Global, NTuple) scaling = closure.scaling - convecting = top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, tracers) > 0 + convecting = top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers)) > 0 interior_point = k <= grid.Nz - 1 & k >= 3 @inbounds diffusivities.wT[i, j, k] = ifelse(convecting & interior_point, inv(scaling.wT)(diffusivities.wT[i, j, k]) - inv(scaling.wT)(0), 0) From 477c6be17ce3f5b0af7e438707030c61ad0e752f Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 29 Aug 2024 18:37:01 -0400 Subject: [PATCH 039/114] fix initial conditions, run for 10 years, set up checkpointing --- physicalclosure_doublegyre_model.jl | 68 +++++++++++++++++------------ 1 file changed, 41 insertions(+), 27 deletions(-) diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index 01951076cd..33bb4fba43 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -20,10 +20,14 @@ using Oceananigans.Grids: xnode, ynode, znode using SeawaterPolynomials using SeawaterPolynomials:TEOS10 using ColorSchemes +using Glob import SeawaterPolynomials.TEOS10: s, ΔS, Sₐᵤ s(Sᴬ::Number) = Sᴬ + ΔS >= 0 ? √((Sᴬ + ΔS) / Sₐᵤ) : NaN #%% +FILE_DIR = "./Output/doublegyre_RiBasedVerticalDiffusivity" +mkpath(FILE_DIR) + # Architecture model_architecture = GPU() @@ -78,7 +82,7 @@ const μ_T = 1/8days ##### Forcing and initial condition ##### -@inline T_initial(x, y, z) = T_south +@inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) @inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) @@ -171,7 +175,7 @@ update_state!(model) ##### Simulation building ##### Δt₀ = 5minutes -stop_time = 730days +stop_time = 3650days simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) @@ -220,49 +224,59 @@ outputs = (; u, v, w, T, S) ##### simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, # filename = "NN_closure_2D_channel_NDE_FC_Qb_18simnew_2layer_128_relu_2Pr", - filename = "doublegyre_Ri_based_vertical_diffusivity_2Pr_xy", + filename = "$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xy", indices = (:, :, Nz), - schedule = TimeInterval(1day), - overwrite_existing = true) + schedule = TimeInterval(5days)) simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, # filename = "NN_closure_2D_channel_NDE_FC_Qb_18simnew_2layer_128_relu_2Pr", - filename = "doublegyre_Ri_based_vertical_diffusivity_2Pr_yz", + filename = "$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_yz", indices = (1, :, :), - schedule = TimeInterval(1day), - overwrite_existing = true) + schedule = TimeInterval(5days)) simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, # filename = "NN_closure_2D_channel_NDE_FC_Qb_18simnew_2layer_128_relu_2Pr", - filename = "doublegyre_Ri_based_vertical_diffusivity_2Pr_xz", + filename = "$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xz", indices = (:, 1, :), - schedule = TimeInterval(1day), - overwrite_existing = true) + schedule = TimeInterval(5days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(365days), + prefix = "$(FILE_DIR)/checkpointer") @info "Running the simulation..." try - run!(simulation, pickup = false) + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end catch err - @info "run! threw an error! The error message is" - showerror(stdout, err) + @info "run! threw an error! The error message is" + showerror(stdout, err) end + #%% -T_xy_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_xy.jld2", "T") -T_xz_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_xz.jld2", "T") -T_yz_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_yz.jld2", "T") +T_xy_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_yz.jld2", "T") -S_xy_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_xy.jld2", "S") -S_xz_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_xz.jld2", "S") -S_yz_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_yz.jld2", "S") +S_xy_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_yz.jld2", "S") -u_xy_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_xy.jld2", "u") -u_xz_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_xz.jld2", "u") -u_yz_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_yz.jld2", "u") +u_xy_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_yz.jld2", "u") -v_xy_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_xy.jld2", "v") -v_xz_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_xz.jld2", "v") -v_yz_data = FieldTimeSeries("doublegyre_Ri_based_vertical_diffusivity_2Pr_yz.jld2", "v") +v_xy_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_yz.jld2", "v") times = T_xy_data.times ./ 24 ./ 60^2 Nt = length(times) @@ -408,7 +422,7 @@ zlims!(axS, (-Lz, 0)) zlims!(axu, (-Lz, 0)) zlims!(axv, (-Lz, 0)) -CairoMakie.record(fig, "./doublegyre_Ri_based_vertical_diffusivity_2Pr.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn +CairoMakie.record(fig, "$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn @info nn n[] = nn end From 232b67dd6f87cb089f5aad2d556d2002cf86c3ff Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 29 Aug 2024 23:58:04 -0400 Subject: [PATCH 040/114] run double gyre with new physical closure --- xk_physicalclosure_doublegyre_model.jl | 560 +++++++++++++++++++++++++ 1 file changed, 560 insertions(+) create mode 100644 xk_physicalclosure_doublegyre_model.jl diff --git a/xk_physicalclosure_doublegyre_model.jl b/xk_physicalclosure_doublegyre_model.jl new file mode 100644 index 0000000000..488c109cd0 --- /dev/null +++ b/xk_physicalclosure_doublegyre_model.jl @@ -0,0 +1,560 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +# include("NN_closure_global.jl") +# include("xin_kai_vertical_diffusivity_local.jl") +include("xin_kai_vertical_diffusivity_2Pr.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes +using Glob +import SeawaterPolynomials.TEOS10: s, ΔS, Sₐᵤ +s(Sᴬ::Number) = Sᴬ + ΔS >= 0 ? √((Sᴬ + ΔS) / Sₐᵤ) : NaN + +#%% +FILE_DIR = "./Output/doublegyre_XinKaiVerticalDiffusivity" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +# nn_closure = NNFluxClosure(model_architecture) +# base_closure = XinKaiLocalVerticalDiffusivity() +# closure = (nn_closure, base_closure) + +vertical_base_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +convection_closure = XinKaiVerticalDiffusivity() +# convection_closure = RiBasedVerticalDiffusivity() +closure = (vertical_base_closure, convection_closure) +# closure = vertical_base_closure + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/8days + +##### +##### Forcing and initial condition +##### + +@inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + # closure = (nn_closure, base_closure), + closure = closure, + # closure = RiBasedVerticalDiffusivity(), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 3650days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### + +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +# ν, κ = model.diffusivity_fields[2].κᵘ, model.diffusivity_fields[2].κᶜ +# Ri = model.diffusivity_fields[2].Ri +# wT, wS = model.diffusivity_fields[2].wT, model.diffusivity_fields[2].wS + +# outputs = (; u, v, w, T, S, ν, κ, Ri, wT, wS) +# outputs = (; u, v, w, T, S, ν, κ, Ri) +outputs = (; u, v, w, T, S) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(5days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(5days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(5days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(365days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +CairoMakie.record(fig, "$(FILE_DIR)/3D_instantaneous_fields.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn + @info nn + n[] = nn +end + +# display(fig) +#%% + +# # ##### +# # ##### Visualization +# # ##### +# using CairoMakie + +# dataname = "NN_closure_doublegyre_NDE_FC_Qb_absf_24simnew_2layer_128_relu_2Pr" +# DATA_DIR = "./$(dataname).jld2" + +# u_data = FieldTimeSeries("$(DATA_DIR)", "u") +# v_data = FieldTimeSeries("$(DATA_DIR)", "v") +# T_data = FieldTimeSeries("$(DATA_DIR)", "T") +# S_data = FieldTimeSeries("$(DATA_DIR)", "S") +# # ν_data = FieldTimeSeries("$(DATA_DIR)", "ν") +# # κ_data = FieldTimeSeries("$(DATA_DIR)", "κ") +# # Ri_data = FieldTimeSeries("$(DATA_DIR)", "Ri") +# # wT_data = FieldTimeSeries("$(DATA_DIR)", "wT") +# # wS_data = FieldTimeSeries("$(DATA_DIR)", "wS") + +# yC = ynodes(T_data.grid, Center()) +# yF = ynodes(T_data.grid, Face()) + +# zC = znodes(T_data.grid, Center()) +# zF = znodes(T_data.grid, Face()) + +# Nt = length(T_data.times) +# #%% +# fig = Figure(size = (1500, 900)) +# axu = CairoMakie.Axis(fig[1, 1], xlabel = "y (m)", ylabel = "z (m)", title = "u (m/s)") +# axv = CairoMakie.Axis(fig[1, 3], xlabel = "y (m)", ylabel = "z (m)", title = "v (m/s)") +# axT = CairoMakie.Axis(fig[2, 1], xlabel = "y (m)", ylabel = "z (m)", title = "Temperature (°C)") +# axS = CairoMakie.Axis(fig[2, 3], xlabel = "y (m)", ylabel = "z (m)", title = "Salinity (psu)") +# n = Observable(1) + +# uₙ = @lift interior(u_data[$n], 45, :, :) +# vₙ = @lift interior(v_data[$n], 45, :, :) +# Tₙ = @lift interior(T_data[$n], 45, :, :) +# Sₙ = @lift interior(S_data[$n], 45, :, :) + +# ulim = @lift (-maximum([maximum(abs, $uₙ), 1e-16]), maximum([maximum(abs, $uₙ), 1e-16])) +# vlim = @lift (-maximum([maximum(abs, $vₙ), 1e-16]), maximum([maximum(abs, $vₙ), 1e-16])) +# Tlim = (minimum(interior(T_data[1])), maximum(interior(T_data[1]))) +# Slim = (minimum(interior(S_data[1])), maximum(interior(S_data[1]))) + +# title_str = @lift "Time: $(round(T_data.times[$n] / 86400, digits=2)) days" +# Label(fig[0, :], title_str, tellwidth = false) + +# hu = heatmap!(axu, yC, zC, uₙ, colormap=:RdBu_9, colorrange=ulim) +# hv = heatmap!(axv, yF, zC, vₙ, colormap=:RdBu_9, colorrange=vlim) +# hT = heatmap!(axT, yC, zC, Tₙ, colorrange=Tlim) +# hS = heatmap!(axS, yC, zC, Sₙ, colorrange=Slim) + +# Colorbar(fig[1, 2], hu, label = "u (m/s)") +# Colorbar(fig[1, 4], hv, label = "v (m/s)") +# Colorbar(fig[2, 2], hT, label = "T (°C)") +# Colorbar(fig[2, 4], hS, label = "S (psu)") + +# CairoMakie.record(fig, "./$(dataname)_test.mp4", 1:Nt, framerate=10) do nn +# n[] = nn +# end + +# display(fig) +# #%% +# fig = Figure(size = (1920, 1080)) +# axu = CairoMakie.Axis(fig[1, 1], xlabel = "y (m)", ylabel = "z (m)", title = "u") +# axv = CairoMakie.Axis(fig[1, 3], xlabel = "y (m)", ylabel = "z (m)", title = "v") +# axT = CairoMakie.Axis(fig[2, 1], xlabel = "y (m)", ylabel = "z (m)", title = "Temperature") +# axS = CairoMakie.Axis(fig[2, 3], xlabel = "y (m)", ylabel = "z (m)", title = "Salinity") +# axν = CairoMakie.Axis(fig[1, 5], xlabel = "y (m)", ylabel = "z (m)", title = "Viscosity (log10 scale)") +# axκ = CairoMakie.Axis(fig[2, 5], xlabel = "y (m)", ylabel = "z (m)", title = "Diffusivity (log10 scale)") +# axRi = CairoMakie.Axis(fig[3, 5], xlabel = "y (m)", ylabel = "z (m)", title = "Richardson number") +# axwT = CairoMakie.Axis(fig[3, 1], xlabel = "y (m)", ylabel = "z (m)", title = "wT(NN)") +# axwS = CairoMakie.Axis(fig[3, 3], xlabel = "y (m)", ylabel = "z (m)", title = "wS(NN)") + +# n = Observable(1) + +# uₙ = @lift interior(u_data[$n], 1, :, :) +# vₙ = @lift interior(v_data[$n], 1, :, :) +# Tₙ = @lift interior(T_data[$n], 1, :, :) +# Sₙ = @lift interior(S_data[$n], 1, :, :) +# νₙ = @lift log10.(interior(ν_data[$n], 1, :, :)) +# κₙ = @lift log10.(interior(κ_data[$n], 1, :, :)) +# Riₙ = @lift clamp.(interior(Ri_data[$n], 1, :, :), -20, 20) +# wTₙ = @lift interior(wT_data[$n], 1, :, :) +# wSₙ = @lift interior(wS_data[$n], 1, :, :) + +# ulim = @lift (-maximum([maximum(abs, $uₙ), 1e-7]), maximum([maximum(abs, $uₙ), 1e-7])) +# vlim = @lift (-maximum([maximum(abs, $vₙ), 1e-7]), maximum([maximum(abs, $vₙ), 1e-7])) +# Tlim = (minimum(interior(T_data[1])), maximum(interior(T_data[1]))) +# Slim = (minimum(interior(S_data[1])), maximum(interior(S_data[1]))) +# νlim = (-6, 2) +# κlim = (-6, 2) +# wTlim = @lift (-maximum([maximum(abs, $wTₙ), 1e-7]), maximum([maximum(abs, $wTₙ), 1e-7])) +# wSlim = @lift (-maximum([maximum(abs, $wSₙ), 1e-7]), maximum([maximum(abs, $wSₙ), 1e-7])) + +# title_str = @lift "Time: $(round(T_data.times[$n] / 86400, digits=2)) days" +# Label(fig[0, :], title_str, tellwidth = false) + +# hu = heatmap!(axu, yC, zC, uₙ, colormap=:RdBu_9, colorrange=ulim) +# hv = heatmap!(axv, yF, zC, vₙ, colormap=:RdBu_9, colorrange=vlim) +# hT = heatmap!(axT, yC, zC, Tₙ, colorrange=Tlim) +# hS = heatmap!(axS, yC, zC, Sₙ, colorrange=Slim) +# hν = heatmap!(axν, yC, zC, νₙ, colorrange=νlim) +# hκ = heatmap!(axκ, yC, zC, κₙ, colorrange=κlim) +# hRi = heatmap!(axRi, yC, zF, Riₙ, colormap=:RdBu_9, colorrange=(-20, 20)) +# hwT = heatmap!(axwT, yC, zF, wTₙ, colormap=:RdBu_9, colorrange=wTlim) +# hwS = heatmap!(axwS, yC, zF, wSₙ, colormap=:RdBu_9, colorrange=wSlim) + +# cbu = Colorbar(fig[1, 2], hu, label = "(m/s)") +# cbv = Colorbar(fig[1, 4], hv, label = "(m/s)") +# cbT = Colorbar(fig[2, 2], hT, label = "(°C)") +# cbS = Colorbar(fig[2, 4], hS, label = "(psu)") +# cbν = Colorbar(fig[1, 6], hν, label = "(m²/s)") +# cbκ = Colorbar(fig[2, 6], hκ, label = "(m²/s)") +# cbRi = Colorbar(fig[3, 6], hRi) +# cbwT = Colorbar(fig[3, 2], hwT, label = "(m/s °C)") +# cbwS = Colorbar(fig[3, 4], hwS, label = "(m/s psu)") + +# tight_ticklabel_spacing!(cbu) +# tight_ticklabel_spacing!(cbv) +# tight_ticklabel_spacing!(cbT) +# tight_ticklabel_spacing!(cbS) +# tight_ticklabel_spacing!(cbν) +# tight_ticklabel_spacing!(cbκ) +# tight_ticklabel_spacing!(cbRi) +# tight_ticklabel_spacing!(cbwT) +# tight_ticklabel_spacing!(cbwS) + +# CairoMakie.record(fig, "./$(dataname)_2D_sin_cooling_heating_23days_fluxes.mp4", 1:Nt, framerate=30) do nn +# n[] = nn +# end +# #%% \ No newline at end of file From 148a5827425ea213142368eaae599eeedf368039 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Wed, 4 Sep 2024 16:06:49 -0400 Subject: [PATCH 041/114] plotting barotropic streamfunction --- plot_barotropic_streamfunction.jl | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) create mode 100644 plot_barotropic_streamfunction.jl diff --git a/plot_barotropic_streamfunction.jl b/plot_barotropic_streamfunction.jl new file mode 100644 index 0000000000..b7cf736dca --- /dev/null +++ b/plot_barotropic_streamfunction.jl @@ -0,0 +1,30 @@ +using CairoMakie +using Oceananigans + +filename = "doublegyre_RiBasedVerticalDiffusivity_streamfunction" +FILE_DIR = "./Output/$(filename)/" + +Ψ_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_streamfunction.jld2", "Ψ") + +Nx = Ψ_data.grid.Nx +Ny = Ψ_data.grid.Ny + +xF = Ψ_data.grid.xᶠᵃᵃ[1:Nx+1] +yC = Ψ_data.grid.yᵃᶜᵃ[1:Ny] + +Nt = length(Ψ_data) +times = Ψ_data.times / 24 / 60^2 / 365 +#%% +timeframe = 31 +Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +clim = maximum(abs, Ψ_frame) +N_levels = 16 +levels = range(-clim, stop=clim, length=N_levels) +fig = Figure(size=(800, 800)) +ax = Axis(fig[1, 1], xlabel="x (km)", ylabel="y (km)", title="Ri-based Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +tightlimits!(ax) +save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +display(fig) +#%% \ No newline at end of file From e9f2f0e4742f78914668229284c2ec799b61383d Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Wed, 4 Sep 2024 16:18:17 -0400 Subject: [PATCH 042/114] fix plot units --- plot_barotropic_streamfunction.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plot_barotropic_streamfunction.jl b/plot_barotropic_streamfunction.jl index b7cf736dca..6bc5454513 100644 --- a/plot_barotropic_streamfunction.jl +++ b/plot_barotropic_streamfunction.jl @@ -21,7 +21,7 @@ clim = maximum(abs, Ψ_frame) N_levels = 16 levels = range(-clim, stop=clim, length=N_levels) fig = Figure(size=(800, 800)) -ax = Axis(fig[1, 1], xlabel="x (km)", ylabel="y (km)", title="Ri-based Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +ax = Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="Ri-based Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) Colorbar(fig[1, 2], cf, label="Ψ (Sv)") tightlimits!(ax) From ed8ec7d6eb77eeaa5d4355200ac592f59cbc9a1c Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Mon, 9 Sep 2024 15:52:57 -0400 Subject: [PATCH 043/114] run double gyre with CATKE --- physicalclosure_doublegyre_model.jl | 291 ++++++++++++---------------- 1 file changed, 124 insertions(+), 167 deletions(-) diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index 33bb4fba43..6808a492b0 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -1,6 +1,8 @@ #using Pkg # pkg"add Oceananigans CairoMakie" using Oceananigans +using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity +using Oceananigans.BuoyancyModels: ∂z_b # include("NN_closure_global.jl") # include("xin_kai_vertical_diffusivity_local.jl") # include("xin_kai_vertical_diffusivity_2Pr.jl") @@ -25,7 +27,8 @@ import SeawaterPolynomials.TEOS10: s, ΔS, Sₐᵤ s(Sᴬ::Number) = Sᴬ + ΔS >= 0 ? √((Sᴬ + ΔS) / Sₐᵤ) : NaN #%% -FILE_DIR = "./Output/doublegyre_RiBasedVerticalDiffusivity" +filename = "doublegyre_CATKEVerticalDiffusivity_streamfunction" +FILE_DIR = "./Output/$(filename)" mkpath(FILE_DIR) # Architecture @@ -37,7 +40,7 @@ model_architecture = GPU() vertical_base_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) # convection_closure = XinKaiVerticalDiffusivity() -convection_closure = RiBasedVerticalDiffusivity() +convection_closure = CATKEVerticalDiffusivity() closure = (vertical_base_closure, convection_closure) # closure = vertical_base_closure @@ -149,9 +152,7 @@ model = HydrostaticFreeSurfaceModel( tracer_advection = WENO(order=5), buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), coriolis = coriolis, - # closure = (nn_closure, base_closure), closure = closure, - # closure = RiBasedVerticalDiffusivity(), tracers = (:T, :S), boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), ) @@ -175,7 +176,7 @@ update_state!(model) ##### Simulation building ##### Δt₀ = 5minutes -stop_time = 3650days +stop_time = 10950days simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) @@ -211,37 +212,52 @@ simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterv u, v, w = model.velocities T, S = model.tracers.T, model.tracers.S -# ν, κ = model.diffusivity_fields[2].κᵘ, model.diffusivity_fields[2].κᶜ -# Ri = model.diffusivity_fields[2].Ri -# wT, wS = model.diffusivity_fields[2].wT, model.diffusivity_fields[2].wS +U_bt = Field(Integral(u, dims=3)) +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)) -# outputs = (; u, v, w, T, S, ν, κ, Ri, wT, wS) -# outputs = (; u, v, w, T, S, ν, κ, Ri) -outputs = (; u, v, w, T, S) +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +outputs = (; u, v, w, T, S, N²) ##### ##### Build checkpointer and output writer ##### simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, - # filename = "NN_closure_2D_channel_NDE_FC_Qb_18simnew_2layer_128_relu_2Pr", - filename = "$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xy", + filename = "$(FILE_DIR)/instantaneous_fields_xy", indices = (:, :, Nz), - schedule = TimeInterval(5days)) + schedule = TimeInterval(10days)) simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, - # filename = "NN_closure_2D_channel_NDE_FC_Qb_18simnew_2layer_128_relu_2Pr", - filename = "$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_yz", + filename = "$(FILE_DIR)/instantaneous_fields_yz", indices = (1, :, :), - schedule = TimeInterval(5days)) + schedule = TimeInterval(10days)) simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, - # filename = "NN_closure_2D_channel_NDE_FC_Qb_18simnew_2layer_128_relu_2Pr", - filename = "$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xz", + filename = "$(FILE_DIR)/instantaneous_fields_xz", indices = (:, 1, :), - schedule = TimeInterval(5days)) + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_south] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_south", + indices = (:, 25, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_north] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_north", + indices = (:, 75, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) simulation.output_writers[:checkpointer] = Checkpointer(model, - schedule = TimeInterval(365days), + schedule = TimeInterval(730days), prefix = "$(FILE_DIR)/checkpointer") @info "Running the simulation..." @@ -262,21 +278,21 @@ catch err end #%% -T_xy_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xy.jld2", "T") -T_xz_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xz.jld2", "T") -T_yz_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_yz.jld2", "T") +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") -S_xy_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xy.jld2", "S") -S_xz_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xz.jld2", "S") -S_yz_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_yz.jld2", "S") +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") -u_xy_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xy.jld2", "u") -u_xz_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xz.jld2", "u") -u_yz_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_yz.jld2", "u") +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") -v_xy_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xy.jld2", "v") -v_xz_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_xz.jld2", "v") -v_yz_data = FieldTimeSeries("$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr_yz.jld2", "v") +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") times = T_xy_data.times ./ 24 ./ 60^2 Nt = length(times) @@ -422,142 +438,83 @@ zlims!(axS, (-Lz, 0)) zlims!(axu, (-Lz, 0)) zlims!(axv, (-Lz, 0)) -CairoMakie.record(fig, "$(FILE_DIR)/doublegyre_Ri_based_vertical_diffusivity_2Pr.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn - @info nn +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn n[] = nn end -# display(fig) +@info "Done!" +#%% +Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +Nt = length(Ψ_data) +times = Ψ_data.times / 24 / 60^2 / 365 +#%% +timeframe = Nt +Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +clim = maximum(abs, Ψ_frame) + 1e-13 +N_levels = 16 +levels = range(-clim, stop=clim, length=N_levels) +fig = Figure(size=(800, 800)) +ax = Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +tightlimits!(ax) +save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +display(fig) #%% +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +N²_xz_north_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_north.jld2", "N²") +N²_xz_south_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_south.jld2", "N²") + +xC = N²_xz_north_data.grid.xᶜᵃᵃ[1:N²_xz_north_data.grid.Nx] +zf = N²_xz_north_data.grid.zᵃᵃᶠ[1:N²_xz_north_data.grid.Nz+1] + +yloc_north = N²_xz_north_data.grid.yᵃᶠᵃ[N²_xz_north_data.indices[2][1]] +yloc_south = N²_xz_south_data.grid.yᵃᶠᵃ[N²_xz_south_data.indices[2][1]] + +Nt = length(N²_xz_north_data) +times = N²_xz_north_data.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +N²_lim = (find_min(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) - 1e-13, + find_max(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) + 1e-13) +#%% +fig = Figure(size=(800, 800)) +ax_north = Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_south / 1e3)) km") +ax_south = Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_north / 1e3)) km") + +n = Observable(2) + +N²_north = @lift interior(N²_xz_north_data[$n], :, 1, :) +N²_south = @lift interior(N²_xz_south_data[$n], :, 1, :) + +colorscheme = colorschemes[:jet] + +N²_north_surface = heatmap!(ax_north, xC, zf, N²_north, colormap=colorscheme, colorrange=N²_lim) +N²_south_surface = heatmap!(ax_south, xC, zf, N²_south, colormap=colorscheme, colorrange=N²_lim) + +Colorbar(fig[1:2, 2], N²_north_surface, label="N² (s⁻²)") + +title_str = @lift "Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +@info "Recording buoyancy frequency xz slice" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_buoyancy_frequency_xz_slice.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn + n[] = nn +end -# # ##### -# # ##### Visualization -# # ##### -# using CairoMakie - -# dataname = "NN_closure_doublegyre_NDE_FC_Qb_absf_24simnew_2layer_128_relu_2Pr" -# DATA_DIR = "./$(dataname).jld2" - -# u_data = FieldTimeSeries("$(DATA_DIR)", "u") -# v_data = FieldTimeSeries("$(DATA_DIR)", "v") -# T_data = FieldTimeSeries("$(DATA_DIR)", "T") -# S_data = FieldTimeSeries("$(DATA_DIR)", "S") -# # ν_data = FieldTimeSeries("$(DATA_DIR)", "ν") -# # κ_data = FieldTimeSeries("$(DATA_DIR)", "κ") -# # Ri_data = FieldTimeSeries("$(DATA_DIR)", "Ri") -# # wT_data = FieldTimeSeries("$(DATA_DIR)", "wT") -# # wS_data = FieldTimeSeries("$(DATA_DIR)", "wS") - -# yC = ynodes(T_data.grid, Center()) -# yF = ynodes(T_data.grid, Face()) - -# zC = znodes(T_data.grid, Center()) -# zF = znodes(T_data.grid, Face()) - -# Nt = length(T_data.times) -# #%% -# fig = Figure(size = (1500, 900)) -# axu = CairoMakie.Axis(fig[1, 1], xlabel = "y (m)", ylabel = "z (m)", title = "u (m/s)") -# axv = CairoMakie.Axis(fig[1, 3], xlabel = "y (m)", ylabel = "z (m)", title = "v (m/s)") -# axT = CairoMakie.Axis(fig[2, 1], xlabel = "y (m)", ylabel = "z (m)", title = "Temperature (°C)") -# axS = CairoMakie.Axis(fig[2, 3], xlabel = "y (m)", ylabel = "z (m)", title = "Salinity (psu)") -# n = Observable(1) - -# uₙ = @lift interior(u_data[$n], 45, :, :) -# vₙ = @lift interior(v_data[$n], 45, :, :) -# Tₙ = @lift interior(T_data[$n], 45, :, :) -# Sₙ = @lift interior(S_data[$n], 45, :, :) - -# ulim = @lift (-maximum([maximum(abs, $uₙ), 1e-16]), maximum([maximum(abs, $uₙ), 1e-16])) -# vlim = @lift (-maximum([maximum(abs, $vₙ), 1e-16]), maximum([maximum(abs, $vₙ), 1e-16])) -# Tlim = (minimum(interior(T_data[1])), maximum(interior(T_data[1]))) -# Slim = (minimum(interior(S_data[1])), maximum(interior(S_data[1]))) - -# title_str = @lift "Time: $(round(T_data.times[$n] / 86400, digits=2)) days" -# Label(fig[0, :], title_str, tellwidth = false) - -# hu = heatmap!(axu, yC, zC, uₙ, colormap=:RdBu_9, colorrange=ulim) -# hv = heatmap!(axv, yF, zC, vₙ, colormap=:RdBu_9, colorrange=vlim) -# hT = heatmap!(axT, yC, zC, Tₙ, colorrange=Tlim) -# hS = heatmap!(axS, yC, zC, Sₙ, colorrange=Slim) - -# Colorbar(fig[1, 2], hu, label = "u (m/s)") -# Colorbar(fig[1, 4], hv, label = "v (m/s)") -# Colorbar(fig[2, 2], hT, label = "T (°C)") -# Colorbar(fig[2, 4], hS, label = "S (psu)") - -# CairoMakie.record(fig, "./$(dataname)_test.mp4", 1:Nt, framerate=10) do nn -# n[] = nn -# end - -# display(fig) -# #%% -# fig = Figure(size = (1920, 1080)) -# axu = CairoMakie.Axis(fig[1, 1], xlabel = "y (m)", ylabel = "z (m)", title = "u") -# axv = CairoMakie.Axis(fig[1, 3], xlabel = "y (m)", ylabel = "z (m)", title = "v") -# axT = CairoMakie.Axis(fig[2, 1], xlabel = "y (m)", ylabel = "z (m)", title = "Temperature") -# axS = CairoMakie.Axis(fig[2, 3], xlabel = "y (m)", ylabel = "z (m)", title = "Salinity") -# axν = CairoMakie.Axis(fig[1, 5], xlabel = "y (m)", ylabel = "z (m)", title = "Viscosity (log10 scale)") -# axκ = CairoMakie.Axis(fig[2, 5], xlabel = "y (m)", ylabel = "z (m)", title = "Diffusivity (log10 scale)") -# axRi = CairoMakie.Axis(fig[3, 5], xlabel = "y (m)", ylabel = "z (m)", title = "Richardson number") -# axwT = CairoMakie.Axis(fig[3, 1], xlabel = "y (m)", ylabel = "z (m)", title = "wT(NN)") -# axwS = CairoMakie.Axis(fig[3, 3], xlabel = "y (m)", ylabel = "z (m)", title = "wS(NN)") - -# n = Observable(1) - -# uₙ = @lift interior(u_data[$n], 1, :, :) -# vₙ = @lift interior(v_data[$n], 1, :, :) -# Tₙ = @lift interior(T_data[$n], 1, :, :) -# Sₙ = @lift interior(S_data[$n], 1, :, :) -# νₙ = @lift log10.(interior(ν_data[$n], 1, :, :)) -# κₙ = @lift log10.(interior(κ_data[$n], 1, :, :)) -# Riₙ = @lift clamp.(interior(Ri_data[$n], 1, :, :), -20, 20) -# wTₙ = @lift interior(wT_data[$n], 1, :, :) -# wSₙ = @lift interior(wS_data[$n], 1, :, :) - -# ulim = @lift (-maximum([maximum(abs, $uₙ), 1e-7]), maximum([maximum(abs, $uₙ), 1e-7])) -# vlim = @lift (-maximum([maximum(abs, $vₙ), 1e-7]), maximum([maximum(abs, $vₙ), 1e-7])) -# Tlim = (minimum(interior(T_data[1])), maximum(interior(T_data[1]))) -# Slim = (minimum(interior(S_data[1])), maximum(interior(S_data[1]))) -# νlim = (-6, 2) -# κlim = (-6, 2) -# wTlim = @lift (-maximum([maximum(abs, $wTₙ), 1e-7]), maximum([maximum(abs, $wTₙ), 1e-7])) -# wSlim = @lift (-maximum([maximum(abs, $wSₙ), 1e-7]), maximum([maximum(abs, $wSₙ), 1e-7])) - -# title_str = @lift "Time: $(round(T_data.times[$n] / 86400, digits=2)) days" -# Label(fig[0, :], title_str, tellwidth = false) - -# hu = heatmap!(axu, yC, zC, uₙ, colormap=:RdBu_9, colorrange=ulim) -# hv = heatmap!(axv, yF, zC, vₙ, colormap=:RdBu_9, colorrange=vlim) -# hT = heatmap!(axT, yC, zC, Tₙ, colorrange=Tlim) -# hS = heatmap!(axS, yC, zC, Sₙ, colorrange=Slim) -# hν = heatmap!(axν, yC, zC, νₙ, colorrange=νlim) -# hκ = heatmap!(axκ, yC, zC, κₙ, colorrange=κlim) -# hRi = heatmap!(axRi, yC, zF, Riₙ, colormap=:RdBu_9, colorrange=(-20, 20)) -# hwT = heatmap!(axwT, yC, zF, wTₙ, colormap=:RdBu_9, colorrange=wTlim) -# hwS = heatmap!(axwS, yC, zF, wSₙ, colormap=:RdBu_9, colorrange=wSlim) - -# cbu = Colorbar(fig[1, 2], hu, label = "(m/s)") -# cbv = Colorbar(fig[1, 4], hv, label = "(m/s)") -# cbT = Colorbar(fig[2, 2], hT, label = "(°C)") -# cbS = Colorbar(fig[2, 4], hS, label = "(psu)") -# cbν = Colorbar(fig[1, 6], hν, label = "(m²/s)") -# cbκ = Colorbar(fig[2, 6], hκ, label = "(m²/s)") -# cbRi = Colorbar(fig[3, 6], hRi) -# cbwT = Colorbar(fig[3, 2], hwT, label = "(m/s °C)") -# cbwS = Colorbar(fig[3, 4], hwS, label = "(m/s psu)") - -# tight_ticklabel_spacing!(cbu) -# tight_ticklabel_spacing!(cbv) -# tight_ticklabel_spacing!(cbT) -# tight_ticklabel_spacing!(cbS) -# tight_ticklabel_spacing!(cbν) -# tight_ticklabel_spacing!(cbκ) -# tight_ticklabel_spacing!(cbRi) -# tight_ticklabel_spacing!(cbwT) -# tight_ticklabel_spacing!(cbwS) - -# CairoMakie.record(fig, "./$(dataname)_2D_sin_cooling_heating_23days_fluxes.mp4", 1:Nt, framerate=30) do nn -# n[] = nn -# end -# #%% \ No newline at end of file +@info "Done!" +#%% \ No newline at end of file From 893e5295d5ee7506ede1fe2c7c0b2c1cfff8223a Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Mon, 9 Sep 2024 16:49:00 -0400 Subject: [PATCH 044/114] add TKE tracer for CATKE --- physicalclosure_doublegyre_model.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index 6808a492b0..f372d07f9a 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -141,7 +141,7 @@ model = HydrostaticFreeSurfaceModel( buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), coriolis = coriolis, closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), - tracers = (:T, :S), + tracers = (:T, :S, :e), boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), ) @@ -153,7 +153,7 @@ model = HydrostaticFreeSurfaceModel( buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), coriolis = coriolis, closure = closure, - tracers = (:T, :S), + tracers = (:T, :S, :e), boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), ) From ba5491d725de7f742e3078fcd33b2f5eafa9f601 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 10 Sep 2024 01:09:42 -0400 Subject: [PATCH 045/114] fix closure to use only CATKE --- physicalclosure_doublegyre_model.jl | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index f372d07f9a..4a65a488e6 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -34,14 +34,10 @@ mkpath(FILE_DIR) # Architecture model_architecture = GPU() -# nn_closure = NNFluxClosure(model_architecture) -# base_closure = XinKaiLocalVerticalDiffusivity() -# closure = (nn_closure, base_closure) - -vertical_base_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +# vertical_base_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) # convection_closure = XinKaiVerticalDiffusivity() convection_closure = CATKEVerticalDiffusivity() -closure = (vertical_base_closure, convection_closure) +closure = convection_closure # closure = vertical_base_closure # number of grid points From 21a83560aa090559a51113c8c95b357aaa9f8bd4 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 17 Sep 2024 11:45:05 -0400 Subject: [PATCH 046/114] update CATKE configuration --- physicalclosure_doublegyre_model.jl | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index 4a65a488e6..631f337e80 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -2,6 +2,8 @@ # pkg"add Oceananigans CairoMakie" using Oceananigans using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity +using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: CATKEVerticalDiffusivity, CATKEMixingLength, CATKEEquation + using Oceananigans.BuoyancyModels: ∂z_b # include("NN_closure_global.jl") # include("xin_kai_vertical_diffusivity_local.jl") @@ -36,7 +38,12 @@ model_architecture = GPU() # vertical_base_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) # convection_closure = XinKaiVerticalDiffusivity() -convection_closure = CATKEVerticalDiffusivity() +function CATKE_ocean_closure() + mixing_length = CATKEMixingLength(Cᵇ=0.01) + turbulent_kinetic_energy_equation = CATKEEquation(Cᵂϵ=1.0) + return CATKEVerticalDiffusivity(; mixing_length, turbulent_kinetic_energy_equation) +end +convection_closure = CATKE_ocean_closure() closure = convection_closure # closure = vertical_base_closure From de190d2a023df5f2ac0c1a3906a0d53bb43c6ee5 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 17 Sep 2024 12:05:18 -0400 Subject: [PATCH 047/114] use older dependencies for compatibility with neural networks --- Manifest.toml | 414 ++++++++++++++++++++++++++++++++++++++++---------- Project.toml | 22 ++- 2 files changed, 346 insertions(+), 90 deletions(-) diff --git a/Manifest.toml b/Manifest.toml index 1e3ebd4d2c..4e69b75d27 100644 --- a/Manifest.toml +++ b/Manifest.toml @@ -1,8 +1,18 @@ # This file is machine-generated - editing it directly is not advised -julia_version = "1.10.5" +julia_version = "1.10.1" manifest_format = "2.0" -project_hash = "029baf2a08759d3f1940dcb219270c1432addb60" +project_hash = "2d1be40f4d74ba6922f076fdb530a846bb95a6f4" + +[[deps.ADTypes]] +git-tree-sha1 = "fa0822e5baee6e23081c2685ae27265dabee23d8" +uuid = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +version = "1.4.0" +weakdeps = ["ChainRulesCore", "EnzymeCore"] + + [deps.ADTypes.extensions] + ADTypesChainRulesCoreExt = "ChainRulesCore" + ADTypesEnzymeCoreExt = "EnzymeCore" [[deps.AbstractFFTs]] deps = ["LinearAlgebra"] @@ -34,6 +44,34 @@ version = "2.3.0" uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f" version = "1.1.1" +[[deps.ArrayInterface]] +deps = ["Adapt", "LinearAlgebra", "SparseArrays", "SuiteSparse"] +git-tree-sha1 = "ed2ec3c9b483842ae59cd273834e5b46206d6dda" +uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" +version = "7.11.0" + + [deps.ArrayInterface.extensions] + ArrayInterfaceBandedMatricesExt = "BandedMatrices" + ArrayInterfaceBlockBandedMatricesExt = "BlockBandedMatrices" + ArrayInterfaceCUDAExt = "CUDA" + ArrayInterfaceCUDSSExt = "CUDSS" + ArrayInterfaceChainRulesExt = "ChainRules" + ArrayInterfaceGPUArraysCoreExt = "GPUArraysCore" + ArrayInterfaceReverseDiffExt = "ReverseDiff" + ArrayInterfaceStaticArraysCoreExt = "StaticArraysCore" + ArrayInterfaceTrackerExt = "Tracker" + + [deps.ArrayInterface.weakdeps] + BandedMatrices = "aae01518-5342-5314-be14-df237901396f" + BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" + CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" + CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e" + ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2" + GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" + ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" + StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" + Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" + [[deps.Artifacts]] uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33" @@ -89,9 +127,10 @@ version = "0.2.5" [[deps.CUDA]] deps = ["AbstractFFTs", "Adapt", "BFloat16s", "CEnum", "CUDA_Driver_jll", "CUDA_Runtime_Discovery", "CUDA_Runtime_jll", "Crayons", "DataFrames", "ExprTools", "GPUArrays", "GPUCompiler", "KernelAbstractions", "LLVM", "LLVMLoopInfo", "LazyArtifacts", "Libdl", "LinearAlgebra", "Logging", "NVTX", "Preferences", "PrettyTables", "Printf", "Random", "Random123", "RandomNumbers", "Reexport", "Requires", "SparseArrays", "StaticArrays", "Statistics"] -git-tree-sha1 = "fdd9dfb67dfefd548f51000cc400bb51003de247" +git-tree-sha1 = "b8c28cb78014f7ae81a652ce1524cba7667dea5c" uuid = "052768ef-5323-5732-b1bb-66c8b64840ba" -version = "5.4.3" +version = "5.3.5" +weakdeps = ["ChainRulesCore", "EnzymeCore", "SpecialFunctions"] [deps.CUDA.extensions] ChainRulesCoreExt = "ChainRulesCore" @@ -99,22 +138,44 @@ version = "5.4.3" SpecialFunctionsExt = "SpecialFunctions" [[deps.CUDA_Driver_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] -git-tree-sha1 = "325058b426c2b421e3d2df3d5fa646d72d2e3e7e" +deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "Pkg"] +git-tree-sha1 = "dc172b558adbf17952001e15cf0d6364e6d78c2f" uuid = "4ee394cb-3365-5eb0-8335-949819d2adfc" -version = "0.9.2+0" +version = "0.8.1+0" [[deps.CUDA_Runtime_Discovery]] deps = ["Libdl"] -git-tree-sha1 = "33576c7c1b2500f8e7e6baa082e04563203b3a45" +git-tree-sha1 = "38f830504358e9972d2a0c3e5d51cb865e0733df" uuid = "1af6417a-86b4-443c-805f-a4643ffb695f" -version = "0.3.5" +version = "0.2.4" [[deps.CUDA_Runtime_jll]] deps = ["Artifacts", "CUDA_Driver_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"] -git-tree-sha1 = "afea94249b821dc754a8ca6695d3daed851e1f5a" +git-tree-sha1 = "4ca7d6d92075906c2ce871ea8bba971fff20d00c" uuid = "76a88914-d11a-5bdc-97e0-2f5a05c973a2" -version = "0.14.1+0" +version = "0.12.1+0" + +[[deps.CUDNN_jll]] +deps = ["Artifacts", "CUDA_Runtime_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"] +git-tree-sha1 = "cbf7d75f8c58b147bdf6acea2e5bc96cececa6d4" +uuid = "62b44479-cb7b-5706-934f-f13b2eb2e645" +version = "9.0.0+1" + +[[deps.ChainRulesCore]] +deps = ["Compat", "LinearAlgebra"] +git-tree-sha1 = "71acdbf594aab5bbb2cec89b208c41b4c411e49f" +uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +version = "1.24.0" +weakdeps = ["SparseArrays"] + + [deps.ChainRulesCore.extensions] + ChainRulesCoreSparseArraysExt = "SparseArrays" + +[[deps.CloseOpenIntervals]] +deps = ["Static", "StaticArrayInterface"] +git-tree-sha1 = "70232f82ffaab9dc52585e0dd043b5e0c6b714f1" +uuid = "fb6a15b2-703c-40df-9091-08a04967cfa9" +version = "0.1.12" [[deps.ColorTypes]] deps = ["FixedPointNumbers", "Random"] @@ -142,9 +203,9 @@ version = "0.3.0" [[deps.Compat]] deps = ["TOML", "UUIDs"] -git-tree-sha1 = "8ae8d32e09f0dcf42a36b90d4e17f5dd2e4c4215" +git-tree-sha1 = "b1c55339b7c6c350ee89f2c1604299660525b248" uuid = "34da2185-b29b-5c13-b0c7-acf172513d20" -version = "4.16.0" +version = "4.15.0" weakdeps = ["Dates", "LinearAlgebra"] [deps.Compat.extensions] @@ -191,18 +252,17 @@ uuid = "2569d6c7-a4a2-43d3-a901-331e8e4be471" version = "0.2.3" [[deps.ConstructionBase]] -git-tree-sha1 = "76219f1ed5771adbb096743bff43fb5fdd4c1157" +deps = ["LinearAlgebra"] +git-tree-sha1 = "260fd2400ed2dab602a7c15cf10c1933c59930a2" uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9" -version = "1.5.8" +version = "1.5.5" [deps.ConstructionBase.extensions] ConstructionBaseIntervalSetsExt = "IntervalSets" - ConstructionBaseLinearAlgebraExt = "LinearAlgebra" ConstructionBaseStaticArraysExt = "StaticArrays" [deps.ConstructionBase.weakdeps] IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953" - LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [[deps.CpuId]] @@ -217,10 +277,10 @@ uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f" version = "4.1.1" [[deps.CubedSphere]] -deps = ["TaylorSeries"] -git-tree-sha1 = "51bb25de518b4c62b7cdf26e5fbb84601bb27a60" +deps = ["Elliptic", "FFTW", "Printf", "ProgressBars", "SpecialFunctions", "TaylorSeries", "Test"] +git-tree-sha1 = "10134667d7d3569b191a65801514271b8a93b292" uuid = "7445602f-e544-4518-8976-18f8e8ae6cdb" -version = "0.3.0" +version = "0.2.5" [[deps.DataAPI]] git-tree-sha1 = "abe83f3a2f1b857aac70ef8b269080af17764bbe" @@ -292,6 +352,20 @@ deps = ["ArgTools", "FileWatching", "LibCURL", "NetworkOptions"] uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6" version = "1.6.0" +[[deps.Elliptic]] +git-tree-sha1 = "71c79e77221ab3a29918aaf6db4f217b89138608" +uuid = "b305315f-e792-5b7a-8f41-49f472929428" +version = "1.0.1" + +[[deps.EnzymeCore]] +git-tree-sha1 = "88bc63137eb033acc3afe1b9875717889c718c46" +uuid = "f151be2c-9106-41f4-ab19-57ee4f262869" +version = "0.7.5" +weakdeps = ["Adapt"] + + [deps.EnzymeCore.extensions] + AdaptExt = "Adapt" + [[deps.ExprTools]] git-tree-sha1 = "27415f162e6028e81c72b82ef756bf321213b6ec" uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04" @@ -362,9 +436,9 @@ version = "6.2.1+6" [[deps.GPUArrays]] deps = ["Adapt", "GPUArraysCore", "LLVM", "LinearAlgebra", "Printf", "Random", "Reexport", "Serialization", "Statistics"] -git-tree-sha1 = "62ee71528cca49be797076a76bdc654a170a523e" +git-tree-sha1 = "c154546e322a9c73364e8a60430b0f79b812d320" uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" -version = "10.3.1" +version = "10.2.0" [[deps.GPUArraysCore]] deps = ["Adapt"] @@ -373,10 +447,10 @@ uuid = "46192b85-c4d5-4398-a991-12ede77f4527" version = "0.1.6" [[deps.GPUCompiler]] -deps = ["ExprTools", "InteractiveUtils", "LLVM", "Libdl", "Logging", "Preferences", "Scratch", "Serialization", "TOML", "TimerOutputs", "UUIDs"] -git-tree-sha1 = "ab29216184312f99ff957b32cd63c2fe9c928b91" +deps = ["ExprTools", "InteractiveUtils", "LLVM", "Libdl", "Logging", "Scratch", "TimerOutputs", "UUIDs"] +git-tree-sha1 = "1600477fba37c9fc067b9be21f5e8101f24a8865" uuid = "61eb1bfa-7361-4325-ad38-22787b887f55" -version = "0.26.7" +version = "0.26.4" [[deps.Glob]] git-tree-sha1 = "97285bbd5230dd766e9ef6749b80fc617126d496" @@ -397,9 +471,14 @@ version = "1.14.3+3" [[deps.Hwloc_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "5e19e1e4fa3e71b774ce746274364aef0234634e" +git-tree-sha1 = "ca0f6bf568b4bfc807e7537f081c81e35ceca114" uuid = "e33a78d0-f292-5ffc-b300-72abe9b543c8" -version = "2.11.1+0" +version = "2.10.0+0" + +[[deps.IfElse]] +git-tree-sha1 = "debdd00ffef04665ccbb3e150747a77560e8fad1" +uuid = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173" +version = "0.1.1" [[deps.IncompleteLU]] deps = ["LinearAlgebra", "SparseArrays"] @@ -408,23 +487,22 @@ uuid = "40713840-3770-5561-ab4c-a76e7d0d7895" version = "0.2.1" [[deps.InlineStrings]] -git-tree-sha1 = "45521d31238e87ee9f9732561bfee12d4eebd52d" +deps = ["Parsers"] +git-tree-sha1 = "86356004f30f8e737eff143d57d41bd580e437aa" uuid = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48" -version = "1.4.2" +version = "1.4.1" [deps.InlineStrings.extensions] ArrowTypesExt = "ArrowTypes" - ParsersExt = "Parsers" [deps.InlineStrings.weakdeps] ArrowTypes = "31f734f8-188a-4ce0-8406-c8a06bd891cd" - Parsers = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0" [[deps.IntelOpenMP_jll]] -deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl"] -git-tree-sha1 = "10bd689145d2c3b2a9844005d01087cc1194e79e" +deps = ["Artifacts", "JLLWrappers", "Libdl"] +git-tree-sha1 = "be50fe8df3acbffa0274a744f1a99d29c45a57f4" uuid = "1d5cc7b8-4909-519e-a0f8-d0f5ad9712d0" -version = "2024.2.1+0" +version = "2024.1.0+0" [[deps.InteractiveUtils]] deps = ["Markdown"] @@ -435,6 +513,11 @@ git-tree-sha1 = "0dc7b50b8d436461be01300fd8cd45aa0274b038" uuid = "41ab1584-1d38-5bbf-9106-f11c6c58b48f" version = "1.3.0" +[[deps.IrrationalConstants]] +git-tree-sha1 = "630b497eafcc20001bba38a4651b327dcfc491d2" +uuid = "92d709cd-6900-40b7-9082-c6be49f344b6" +version = "0.2.2" + [[deps.IterativeSolvers]] deps = ["LinearAlgebra", "Printf", "Random", "RecipesBase", "SparseArrays"] git-tree-sha1 = "59545b0a2b27208b0650df0a46b8e3019f85055b" @@ -448,15 +531,27 @@ version = "1.0.0" [[deps.JLD2]] deps = ["FileIO", "MacroTools", "Mmap", "OrderedCollections", "Pkg", "PrecompileTools", "Reexport", "Requires", "TranscodingStreams", "UUIDs", "Unicode"] -git-tree-sha1 = "5fe858cb863e211c6dedc8cce2dc0752d4ab6e2b" +git-tree-sha1 = "bdbe8222d2f5703ad6a7019277d149ec6d78c301" uuid = "033835bb-8acc-5ee8-8aae-3f567f8a3819" -version = "0.4.50" +version = "0.4.48" [[deps.JLLWrappers]] deps = ["Artifacts", "Preferences"] -git-tree-sha1 = "f389674c99bfcde17dc57454011aa44d5a260a40" +git-tree-sha1 = "7e5d6779a1e09a36db2a7b6cff50942a0a7d0fca" uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210" -version = "1.6.0" +version = "1.5.0" + +[[deps.JSON3]] +deps = ["Dates", "Mmap", "Parsers", "PrecompileTools", "StructTypes", "UUIDs"] +git-tree-sha1 = "eb3edce0ed4fa32f75a0a11217433c31d56bd48b" +uuid = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" +version = "1.14.0" + + [deps.JSON3.extensions] + JSON3ArrowExt = ["ArrowTypes"] + + [deps.JSON3.weakdeps] + ArrowTypes = "31f734f8-188a-4ce0-8406-c8a06bd891cd" [[deps.JuliaNVTXCallbacks_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] @@ -466,18 +561,19 @@ version = "0.2.1+0" [[deps.KernelAbstractions]] deps = ["Adapt", "Atomix", "InteractiveUtils", "LinearAlgebra", "MacroTools", "PrecompileTools", "Requires", "SparseArrays", "StaticArrays", "UUIDs", "UnsafeAtomics", "UnsafeAtomicsLLVM"] -git-tree-sha1 = "d0448cebd5919e06ca5edc7a264631790de810ec" +git-tree-sha1 = "ed7167240f40e62d97c1f5f7735dea6de3cc5c49" uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c" -version = "0.9.22" +version = "0.9.18" +weakdeps = ["EnzymeCore"] [deps.KernelAbstractions.extensions] EnzymeExt = "EnzymeCore" [[deps.LLVM]] deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Preferences", "Printf", "Requires", "Unicode"] -git-tree-sha1 = "2470e69781ddd70b8878491233cd09bc1bd7fc96" +git-tree-sha1 = "839c82932db86740ae729779e610f07a1640be9a" uuid = "929cbde3-209d-540e-8aea-75f648917ca0" -version = "8.1.0" +version = "6.6.3" weakdeps = ["BFloat16s"] [deps.LLVM.extensions] @@ -485,9 +581,9 @@ weakdeps = ["BFloat16s"] [[deps.LLVMExtra_jll]] deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"] -git-tree-sha1 = "597d1c758c9ae5d985ba4202386a607c675ee700" +git-tree-sha1 = "88b916503aac4fb7f701bb625cd84ca5dd1677bc" uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab" -version = "0.0.31+0" +version = "0.0.29+0" [[deps.LLVMLoopInfo]] git-tree-sha1 = "2e5c102cfc41f48ae4740c7eca7743cc7e7b75ea" @@ -555,6 +651,22 @@ version = "1.17.0+0" deps = ["Libdl", "OpenBLAS_jll", "libblastrampoline_jll"] uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +[[deps.LogExpFunctions]] +deps = ["DocStringExtensions", "IrrationalConstants", "LinearAlgebra"] +git-tree-sha1 = "a2d09619db4e765091ee5c6ffe8872849de0feea" +uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688" +version = "0.3.28" + + [deps.LogExpFunctions.extensions] + LogExpFunctionsChainRulesCoreExt = "ChainRulesCore" + LogExpFunctionsChangesOfVariablesExt = "ChangesOfVariables" + LogExpFunctionsInverseFunctionsExt = "InverseFunctions" + + [deps.LogExpFunctions.weakdeps] + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" + InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112" + [[deps.Logging]] uuid = "56ddb016-857b-54e1-b83d-db4d58db5568" @@ -684,21 +796,21 @@ version = "0.3.27" [[deps.Lz4_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "7f26c8fc5229e68484e0b3447312c98e16207d11" +git-tree-sha1 = "6c26c5e8a4203d43b5497be3ec5d4e0c3cde240a" uuid = "5ced341a-0733-55b8-9ab6-a4889d929147" -version = "1.10.0+0" +version = "1.9.4+0" [[deps.MKL_jll]] deps = ["Artifacts", "IntelOpenMP_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "oneTBB_jll"] -git-tree-sha1 = "f046ccd0c6db2832a9f639e2c669c6fe867e5f4f" +git-tree-sha1 = "80b2833b56d466b3858d565adcd16a4a05f2089b" uuid = "856f044c-d86e-5d09-b602-aeab76dc8ba7" -version = "2024.2.0+0" +version = "2024.1.0+0" [[deps.MPI]] deps = ["Distributed", "DocStringExtensions", "Libdl", "MPICH_jll", "MPIPreferences", "MPItrampoline_jll", "MicrosoftMPI_jll", "OpenMPI_jll", "PkgVersion", "PrecompileTools", "Requires", "Serialization", "Sockets"] -git-tree-sha1 = "14cef41baf5b675b192b02a22c710f725ab333a7" +git-tree-sha1 = "4e3136db3735924f96632a5b40a5979f1f53fa07" uuid = "da04e1cc-30fd-572f-bb4f-1f8673147195" -version = "0.20.20" +version = "0.20.19" [deps.MPI.extensions] AMDGPUExt = "AMDGPU" @@ -710,9 +822,9 @@ version = "0.20.20" [[deps.MPICH_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "Hwloc_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "MPIPreferences", "TOML"] -git-tree-sha1 = "19d4bd098928a3263693991500d05d74dbdc2004" +git-tree-sha1 = "4099bb6809ac109bfc17d521dad33763bcf026b7" uuid = "7cb0a576-ebde-5e09-9194-50597f1243b4" -version = "4.2.2+0" +version = "4.2.1+1" [[deps.MPIPreferences]] deps = ["Libdl", "Preferences"] @@ -824,9 +936,9 @@ uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908" version = "1.2.0" [[deps.OffsetArrays]] -git-tree-sha1 = "1a27764e945a152f7ca7efa04de513d473e9542e" +git-tree-sha1 = "e64b4f5ea6b7389f6f046d13d4896a8f9c1ba71e" uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" -version = "1.14.1" +version = "1.14.0" weakdeps = ["Adapt"] [deps.OffsetArrays.extensions] @@ -837,6 +949,11 @@ deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"] uuid = "4536629a-c528-5b80-bd46-f80d51c5b363" version = "0.3.23+4" +[[deps.OpenLibm_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "05823500-19ac-5b8b-9628-191a04bc5112" +version = "0.8.1+2" + [[deps.OpenMPI_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "MPIPreferences", "TOML"] git-tree-sha1 = "e25c1778a98e34219a00455d6e4384e017ea9762" @@ -845,9 +962,15 @@ version = "4.1.6+0" [[deps.OpenSSL_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "1b35263570443fdd9e76c76b7062116e2f374ab8" +git-tree-sha1 = "a028ee3cb5641cccc4c24e90c36b0a4f7707bdf5" uuid = "458c3c95-2e84-50aa-8efc-19380b2a3a95" -version = "3.0.15+0" +version = "3.0.14+0" + +[[deps.OpenSpecFun_jll]] +deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"] +git-tree-sha1 = "13652491f6856acfd2db29360e1bbcd4565d04f1" +uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e" +version = "0.5.5+0" [[deps.OrderedCollections]] git-tree-sha1 = "dfdf5519f235516220579f949664f1bf44e741c5" @@ -860,6 +983,46 @@ git-tree-sha1 = "2cd396108e178f3ae8dedbd8e938a18726ab2fbf" uuid = "c2071276-7c44-58a7-b746-946036e04d0a" version = "0.24.1+0" +[[deps.PackageExtensionCompat]] +git-tree-sha1 = "fb28e33b8a95c4cee25ce296c817d89cc2e53518" +uuid = "65ce6f38-6b18-4e1d-a461-8949797d7930" +version = "1.0.2" +weakdeps = ["Requires", "TOML"] + +[[deps.Parsers]] +deps = ["Dates", "PrecompileTools", "UUIDs"] +git-tree-sha1 = "8489905bcdbcfac64d1daa51ca07c0d8f0283821" +uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0" +version = "2.8.1" + +[[deps.PartialFunctions]] +deps = ["MacroTools"] +git-tree-sha1 = "47b49a4dbc23b76682205c646252c0f9e1eb75af" +uuid = "570af359-4316-4cb7-8c74-252c00c2016b" +version = "1.2.0" + +[[deps.PencilArrays]] +deps = ["Adapt", "JSON3", "LinearAlgebra", "MPI", "OffsetArrays", "Random", "Reexport", "StaticArrayInterface", "StaticArrays", "StaticPermutations", "Strided", "TimerOutputs", "VersionParsing"] +git-tree-sha1 = "fa85ac32172d96cfdb91dbc53e8e57007e5a2b5a" +uuid = "0e08944d-e94e-41b1-9406-dcf66b6a9d2e" +version = "0.19.5" + + [deps.PencilArrays.extensions] + PencilArraysAMDGPUExt = ["AMDGPU"] + PencilArraysDiffEqExt = ["DiffEqBase"] + PencilArraysHDF5Ext = ["HDF5"] + + [deps.PencilArrays.weakdeps] + AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" + DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" + HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" + +[[deps.PencilFFTs]] +deps = ["AbstractFFTs", "FFTW", "LinearAlgebra", "MPI", "PencilArrays", "Reexport", "TimerOutputs"] +git-tree-sha1 = "bd69f3f0ee248cfb4241800aefb705b5ded592ff" +uuid = "4a48f351-57a6-4416-9ec4-c37015456aae" +version = "0.15.1" + [[deps.Pkg]] deps = ["Artifacts", "Dates", "Downloads", "FileWatching", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"] uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" @@ -911,6 +1074,12 @@ version = "2.3.2" deps = ["Unicode"] uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7" +[[deps.ProgressBars]] +deps = ["Printf"] +git-tree-sha1 = "b437cdb0385ed38312d91d9c00c20f3798b30256" +uuid = "49802e3a-d2f1-5c88-81d8-b72133a6f568" +version = "1.5.1" + [[deps.Quaternions]] deps = ["LinearAlgebra", "Random", "RealDot"] git-tree-sha1 = "994cc27cdacca10e68feb291673ec3a76aa2fae9" @@ -932,10 +1101,10 @@ uuid = "74087812-796a-5b5d-8853-05524746bad3" version = "1.7.0" [[deps.RandomNumbers]] -deps = ["Random"] -git-tree-sha1 = "c6ec94d2aaba1ab2ff983052cf6a606ca5985902" +deps = ["Random", "Requires"] +git-tree-sha1 = "043da614cc7e95c703498a491e2c21f58a2b8111" uuid = "e6cf234a-135c-5ec9-84dd-332b85af5143" -version = "1.6.0" +version = "1.5.3" [[deps.RealDot]] deps = ["LinearAlgebra"] @@ -986,15 +1155,15 @@ uuid = "6c6a2e73-6563-6170-7368-637461726353" version = "1.2.1" [[deps.SeawaterPolynomials]] -git-tree-sha1 = "78f965a2f0cd5250a20c9aba9979346dd2b35734" +git-tree-sha1 = "6d85acd6de472f8e6da81c61c7c5b6280a55e0bc" uuid = "d496a93d-167e-4197-9f49-d3af4ff8fe40" -version = "0.3.5" +version = "0.3.4" [[deps.SentinelArrays]] deps = ["Dates", "Random"] -git-tree-sha1 = "ff11acffdb082493657550959d4feb4b6149e73a" +git-tree-sha1 = "90b4f68892337554d31cdcdbe19e48989f26c7e6" uuid = "91c51154-3ec4-41a3-a24f-3f23e20d615c" -version = "1.4.5" +version = "1.4.3" [[deps.Serialization]] uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b" @@ -1019,11 +1188,39 @@ deps = ["Libdl", "LinearAlgebra", "Random", "Serialization", "SuiteSparse_jll"] uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" version = "1.10.0" +[[deps.SpecialFunctions]] +deps = ["IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"] +git-tree-sha1 = "2f5d4697f21388cbe1ff299430dd169ef97d7e14" +uuid = "276daf66-3868-5448-9aa4-cd146d93841b" +version = "2.4.0" +weakdeps = ["ChainRulesCore"] + + [deps.SpecialFunctions.extensions] + SpecialFunctionsChainRulesCoreExt = "ChainRulesCore" + +[[deps.Static]] +deps = ["IfElse"] +git-tree-sha1 = "d2fdac9ff3906e27f7a618d47b676941baa6c80c" +uuid = "aedffcd0-7271-4cad-89d0-dc628f76c6d3" +version = "0.8.10" + +[[deps.StaticArrayInterface]] +deps = ["ArrayInterface", "Compat", "IfElse", "LinearAlgebra", "PrecompileTools", "Requires", "SparseArrays", "Static", "SuiteSparse"] +git-tree-sha1 = "5d66818a39bb04bf328e92bc933ec5b4ee88e436" +uuid = "0d7ed370-da01-4f52-bd93-41d350b8b718" +version = "1.5.0" +weakdeps = ["OffsetArrays", "StaticArrays"] + + [deps.StaticArrayInterface.extensions] + StaticArrayInterfaceOffsetArraysExt = "OffsetArrays" + StaticArrayInterfaceStaticArraysExt = "StaticArrays" + [[deps.StaticArrays]] deps = ["LinearAlgebra", "PrecompileTools", "Random", "StaticArraysCore"] -git-tree-sha1 = "eeafab08ae20c62c44c8399ccb9354a04b80db50" +git-tree-sha1 = "6e00379a24597be4ae1ee6b2d882e15392040132" uuid = "90137ffa-7385-5640-81b9-e52037218182" -version = "1.9.7" +version = "1.9.5" +weakdeps = ["ChainRulesCore", "Statistics"] [deps.StaticArrays.extensions] StaticArraysChainRulesCoreExt = "ChainRulesCore" @@ -1034,6 +1231,11 @@ git-tree-sha1 = "192954ef1208c7019899fbf8049e717f92959682" uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" version = "1.4.3" +[[deps.StaticPermutations]] +git-tree-sha1 = "193c3daa18ff3e55c1dae66acb6a762c4a3bdb0b" +uuid = "15972242-4b8f-49a0-b8a1-9ac0e7a1a45d" +version = "0.3.0" + [[deps.Statistics]] deps = ["LinearAlgebra", "SparseArrays"] uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" @@ -1051,6 +1253,22 @@ git-tree-sha1 = "25349bf8f63aa36acbff5e3550a86e9f5b0ef682" uuid = "7792a7ef-975c-4747-a70f-980b88e8d1da" version = "0.5.6" +[[deps.Strided]] +deps = ["LinearAlgebra", "StridedViews", "TupleTools"] +git-tree-sha1 = "bd9bd1c70cfc115cc3a30213fc725125a6b43652" +uuid = "5e0ebb24-38b0-5f93-81fe-25c709ecae67" +version = "2.1.0" + +[[deps.StridedViews]] +deps = ["LinearAlgebra", "PackageExtensionCompat"] +git-tree-sha1 = "2917996ce0fa6b8a3a85240a5e9ff930e2aeaa43" +uuid = "4db3bf67-4bd7-4b4e-b153-31dc3fb37143" +version = "0.3.1" +weakdeps = ["CUDA"] + + [deps.StridedViews.extensions] + StridedViewsCUDAExt = "CUDA" + [[deps.StringManipulation]] deps = ["PrecompileTools"] git-tree-sha1 = "a04cabe79c5f01f4d723cc6704070ada0b9d46d5" @@ -1070,6 +1288,16 @@ weakdeps = ["Adapt", "GPUArraysCore", "SparseArrays", "StaticArrays"] StructArraysSparseArraysExt = "SparseArrays" StructArraysStaticArraysExt = "StaticArrays" +[[deps.StructTypes]] +deps = ["Dates", "UUIDs"] +git-tree-sha1 = "ca4bccb03acf9faaf4137a9abc1881ed1841aa70" +uuid = "856f2bd8-1eba-4b0a-8007-ebc267875bd4" +version = "1.10.0" + +[[deps.SuiteSparse]] +deps = ["Libdl", "LinearAlgebra", "Serialization", "SparseArrays"] +uuid = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9" + [[deps.SuiteSparse_jll]] deps = ["Artifacts", "Libdl", "libblastrampoline_jll"] uuid = "bea87d4a-7f5b-5778-9afe-8cc45184846c" @@ -1087,10 +1315,10 @@ uuid = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c" version = "1.0.1" [[deps.Tables]] -deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "OrderedCollections", "TableTraits"] -git-tree-sha1 = "598cd7c1f68d1e205689b1c2fe65a9f85846f297" +deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "LinearAlgebra", "OrderedCollections", "TableTraits"] +git-tree-sha1 = "cb76cf677714c095e535e3501ac7954732aeea2d" uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" -version = "1.12.0" +version = "1.11.1" [[deps.Tar]] deps = ["ArgTools", "SHA"] @@ -1098,10 +1326,16 @@ uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e" version = "1.10.0" [[deps.TaylorSeries]] -deps = ["InteractiveUtils", "LinearAlgebra", "Markdown", "Requires", "SparseArrays"] -git-tree-sha1 = "66f4d1993bae49eeba21a1634b5f65782585a42c" +deps = ["LinearAlgebra", "Markdown", "Requires", "SparseArrays"] +git-tree-sha1 = "1c7170668366821b0c4c4fe03ee78f8d6cf36e2c" uuid = "6aa5eb33-94cf-58f4-a9d0-e4b2c4fc25ea" -version = "0.10.13" +version = "0.16.0" + + [deps.TaylorSeries.extensions] + TaylorSeriesIAExt = "IntervalArithmetic" + + [deps.TaylorSeries.weakdeps] + IntervalArithmetic = "d1acc4aa-44c8-5952-acd4-ba5d80a2a253" [[deps.Test]] deps = ["InteractiveUtils", "Logging", "Random", "Serialization"] @@ -1120,9 +1354,18 @@ uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" version = "0.5.24" [[deps.TranscodingStreams]] -git-tree-sha1 = "e84b3a11b9bece70d14cce63406bbc79ed3464d2" +git-tree-sha1 = "a947ea21087caba0a798c5e494d0bb78e3a1a3a0" uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa" -version = "0.11.2" +version = "0.10.9" +weakdeps = ["Random", "Test"] + + [deps.TranscodingStreams.extensions] + TestExt = ["Test", "Random"] + +[[deps.TupleTools]] +git-tree-sha1 = "41d61b1c545b06279871ef1a4b5fcb2cac2191cd" +uuid = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6" +version = "1.5.0" [[deps.UUIDs]] deps = ["Random", "SHA"] @@ -1138,15 +1381,30 @@ version = "0.2.1" [[deps.UnsafeAtomicsLLVM]] deps = ["LLVM", "UnsafeAtomics"] -git-tree-sha1 = "bf2c553f25e954a9b38c9c0593a59bb13113f9e5" +git-tree-sha1 = "d9f5962fecd5ccece07db1ff006fb0b5271bdfdd" uuid = "d80eeb9a-aca5-4d75-85e5-170c8b632249" -version = "0.1.5" +version = "0.1.4" + +[[deps.VersionParsing]] +git-tree-sha1 = "58d6e80b4ee071f5efd07fda82cb9fbe17200868" +uuid = "81def892-9a0e-5fdd-b105-ffc91e053289" +version = "1.3.0" + +[[deps.WeightInitializers]] +deps = ["ChainRulesCore", "LinearAlgebra", "PartialFunctions", "PrecompileTools", "Random", "SpecialFunctions", "Statistics"] +git-tree-sha1 = "f0e6760ef9d22f043710289ddf29e4a4048c4822" +uuid = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d" +version = "0.1.7" +weakdeps = ["CUDA"] + + [deps.WeightInitializers.extensions] + WeightInitializersCUDAExt = "CUDA" [[deps.XML2_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Libiconv_jll", "Zlib_jll"] -git-tree-sha1 = "1165b0443d0eca63ac1e32b8c0eb69ed2f4f8127" +git-tree-sha1 = "52ff2af32e591541550bd753c0da8b9bc92bb9d9" uuid = "02c8fc9c-b97f-50b9-bbe4-9be30ff0a78a" -version = "2.13.3+0" +version = "2.12.7+0" [[deps.XZ_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] @@ -1180,7 +1438,7 @@ version = "1.1.2+0" [[deps.libblastrampoline_jll]] deps = ["Artifacts", "Libdl"] uuid = "8e850b90-86db-534c-a0d3-1478176c7d93" -version = "5.11.0+0" +version = "5.8.0+1" [[deps.libzip_jll]] deps = ["Artifacts", "Bzip2_jll", "GnuTLS_jll", "JLLWrappers", "Libdl", "XZ_jll", "Zlib_jll", "Zstd_jll"] diff --git a/Project.toml b/Project.toml index c94cb55b15..87947544a0 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Oceananigans" uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09" authors = ["Climate Modeling Alliance and contributors"] -version = "0.91.13" +version = "0.91.2" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" @@ -27,6 +27,8 @@ MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab" OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" +PencilArrays = "0e08944d-e94e-41b1-9406-dcf66b6a9d2e" +PencilFFTs = "4a48f351-57a6-4416-9ec4-c37015456aae" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" @@ -39,42 +41,39 @@ StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" [weakdeps] Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" -Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" -MakieCore = "20f20a25-4f0e-4fdf-b5d1-57303727442b" [extensions] OceananigansEnzymeExt = "Enzyme" -OceananigansMakieExt = ["MakieCore", "Makie"] [compat] Adapt = "3, 4" CUDA = "4.1.1, 5" Crayons = "4" -CubedSphere = "0.2, 0.3" +CubedSphere = "0.1, 0.2" Dates = "1.9" Distances = "0.10" DocStringExtensions = "0.8, 0.9" -Enzyme = "0.12.20" +Enzyme = "0.11.14" FFTW = "1" Glob = "1.3" IncompleteLU = "0.2" InteractiveUtils = "1.9" IterativeSolvers = "0.9" JLD2 = "0.4" -KernelAbstractions = "0.9.21" +KernelAbstractions = "0.9" LinearAlgebra = "1.9" Logging = "1.9" MPI = "0.16, 0.17, 0.18, 0.19, 0.20" -Makie = "0.21" -MakieCore = "0.7, 0.8" NCDatasets = "0.12.10, 0.13.1, 0.14" OffsetArrays = "1.4" OrderedCollections = "1.1" +PencilArrays = "0.16, 0.17, 0.18, 0.19" +PencilFFTs = "0.13.5, 0.14, 0.15" Pkg = "1.9" Printf = "1.9" Random = "1.9" Rotations = "1.0" -SeawaterPolynomials = "0.3.5" +SeawaterPolynomials = "0.3.4" SparseArrays = "1.9" Statistics = "1.9" StructArrays = "0.4, 0.5, 0.6" @@ -88,7 +87,6 @@ DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267" -OrthogonalSphericalShellGrids = "c2be9673-fb75-4747-82dc-aa2bb9f4aed0" OpenMPI_jll = "fe0851c0-eecd-5654-98d4-656369965a5c" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" @@ -96,4 +94,4 @@ TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" TimesDates = "bdfc003b-8df8-5c39-adcd-3a9087f5df4a" [targets] -test = ["BenchmarkTools", "Coverage", "CUDA_Runtime_jll", "DataDeps", "Enzyme", "InteractiveUtils", "MPIPreferences", "OrthogonalSphericalShellGrids", "OpenMPI_jll", "SafeTestsets", "Test", "TimerOutputs", "TimesDates"] +test = ["BenchmarkTools", "Coverage", "CUDA_Runtime_jll", "DataDeps", "Enzyme", "InteractiveUtils", "MPIPreferences", "OpenMPI_jll", "Test", "TimerOutputs", "TimesDates", "SafeTestsets"] From 7e7a04fe29305aa66776679c9881175755529c8e Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 19 Sep 2024 01:47:59 -0400 Subject: [PATCH 048/114] add fields to be calculated for CATKE --- physicalclosure_doublegyre_model.jl | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index 631f337e80..4583664fa8 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -29,7 +29,7 @@ import SeawaterPolynomials.TEOS10: s, ΔS, Sₐᵤ s(Sᴬ::Number) = Sᴬ + ΔS >= 0 ? √((Sᴬ + ΔS) / Sₐᵤ) : NaN #%% -filename = "doublegyre_CATKEVerticalDiffusivity_streamfunction" +filename = "doublegyre_CATKEVerticalDiffusivity" FILE_DIR = "./Output/$(filename)" mkpath(FILE_DIR) @@ -225,7 +225,25 @@ end N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) N² = Field(N²_op) -outputs = (; u, v, w, T, S, N²) +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) +compute!(ρ) + +ρᶠ = @at (Center, Center, Face) ρ +∂ρ∂z = ∂z(ρ) +∂²ρ∂z² = ∂z(∂ρ∂z) + +κc = model.diffusivity_fields.κc +wT = κc * ∂z(T) +wS = κc * ∂z(S) + +outputs = (; u, v, w, T, S, ρ, ρᶠ, ∂ρ∂z, ∂²ρ∂z², N², wT, wS) ##### ##### Build checkpointer and output writer From 1c3186872df0eddf53f78e1e07236a18cb749bbe Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 19 Sep 2024 18:03:20 -0400 Subject: [PATCH 049/114] local diffusivity for 2step calibration --- xin_kai_vertical_diffusivity_local_2step.jl | 241 ++++++++++++++++++++ 1 file changed, 241 insertions(+) create mode 100644 xin_kai_vertical_diffusivity_local_2step.jl diff --git a/xin_kai_vertical_diffusivity_local_2step.jl b/xin_kai_vertical_diffusivity_local_2step.jl new file mode 100644 index 0000000000..e1f995bc00 --- /dev/null +++ b/xin_kai_vertical_diffusivity_local_2step.jl @@ -0,0 +1,241 @@ +using Oceananigans +using Oceananigans.Architectures: architecture +using Oceananigans.BuoyancyModels: ∂z_b +using Oceananigans.Operators +using Oceananigans.Grids: inactive_node +using Oceananigans.Operators: ℑzᵃᵃᶜ, ℑxyᶠᶠᵃ, ℑxyᶜᶜᵃ + +using Adapt + +using KernelAbstractions: @index, @kernel +using KernelAbstractions.Extras.LoopInfo: @unroll + +using Oceananigans.TurbulenceClosures: + tapering_factorᶠᶜᶜ, + tapering_factorᶜᶠᶜ, + tapering_factorᶜᶜᶠ, + tapering_factor, + SmallSlopeIsopycnalTensor, + AbstractScalarDiffusivity, + ExplicitTimeDiscretization, + FluxTapering, + isopycnal_rotation_tensor_xz_ccf, + isopycnal_rotation_tensor_yz_ccf, + isopycnal_rotation_tensor_zz_ccf + +import Oceananigans.TurbulenceClosures: + compute_diffusivities!, + DiffusivityFields, + viscosity, + diffusivity, + getclosure, + top_buoyancy_flux, + diffusive_flux_x, + diffusive_flux_y, + diffusive_flux_z, + viscous_flux_ux, + viscous_flux_vx, + viscous_flux_uy, + viscous_flux_vy + +using Oceananigans.Utils: launch! +using Oceananigans.Coriolis: fᶠᶠᵃ +using Oceananigans.Operators +using Oceananigans.BuoyancyModels: ∂x_b, ∂y_b, ∂z_b + +using Oceananigans.TurbulenceClosures +using Oceananigans.TurbulenceClosures: HorizontalFormulation, VerticalFormulation, AbstractScalarDiffusivity +using Oceananigans.TurbulenceClosures: AbstractScalarBiharmonicDiffusivity +using Oceananigans.Operators +using Oceananigans.Operators: Δxᶜᶜᶜ, Δyᶜᶜᶜ, ℑxyᶜᶜᵃ, ζ₃ᶠᶠᶜ, div_xyᶜᶜᶜ +using Oceananigans.Operators: Δx, Δy +using Oceananigans.Operators: ℑxyz + +using Oceananigans.Operators: ℑxyzᶜᶜᶠ, ℑyzᵃᶜᶠ, ℑxzᶜᵃᶠ, Δxᶜᶜᶜ, Δyᶜᶜᶜ + +using Oceananigans.BoundaryConditions + +struct XinKaiLocalVerticalDiffusivity{TD, FT} <: AbstractScalarDiffusivity{TD, VerticalFormulation, 2} + ν₀ :: FT + νˢʰ :: FT + νᶜⁿ :: FT + Pr_convₜ :: FT + Pr_shearₜ :: FT + Riᶜ :: FT + δRi :: FT +end + +function XinKaiLocalVerticalDiffusivity{TD}(ν₀ :: FT, + νˢʰ :: FT, + νᶜⁿ :: FT, + Pr_convₜ :: FT, + Pr_shearₜ :: FT, + Riᶜ :: FT, + δRi :: FT) where {TD, FT} + + return XinKaiLocalVerticalDiffusivity{TD, FT}(ν₀, νˢʰ, νᶜⁿ, Pr_convₜ, Pr_shearₜ, Riᶜ, δRi) +end + +function XinKaiLocalVerticalDiffusivity(time_discretization = VerticallyImplicitTimeDiscretization(), + FT = Float64; + ν₀ = 1e-5, + νˢʰ = 0.0615914063656973, + νᶜⁿ = 1.5364711416895118, + Pr_convₜ = 0.18711389733455402, + Pr_shearₜ = 1.0842017486284887, + Riᶜ = 0.4366901962987793, + δRi = 0.0009691362773690692) + + TD = typeof(time_discretization) + + return XinKaiLocalVerticalDiffusivity{TD}(convert(FT, ν₀), + convert(FT, νˢʰ), + convert(FT, νᶜⁿ), + convert(FT, Pr_convₜ), + convert(FT, Pr_shearₜ), + convert(FT, Riᶜ), + convert(FT, δRi)) +end + +XinKaiLocalVerticalDiffusivity(FT::DataType; kw...) = + XinKaiLocalVerticalDiffusivity(VerticallyImplicitTimeDiscretization(), FT; kw...) + +Adapt.adapt_structure(to, clo::XinKaiLocalVerticalDiffusivity{TD, FT}) where {TD, FT} = + XinKaiLocalVerticalDiffusivity{TD, FT}(clo.ν₀, clo.νˢʰ, clo.νᶜⁿ, clo.Pr_convₜ, clo.Pr_shearₜ, clo.Riᶜ, clo.δRi) + +##### +##### Diffusivity field utilities +##### + +const RBVD = XinKaiLocalVerticalDiffusivity +const RBVDArray = AbstractArray{<:RBVD} +const FlavorOfXKVD = Union{RBVD, RBVDArray} +const c = Center() +const f = Face() + +@inline viscosity_location(::FlavorOfXKVD) = (c, c, f) +@inline diffusivity_location(::FlavorOfXKVD) = (c, c, f) + +@inline viscosity(::FlavorOfXKVD, diffusivities) = diffusivities.κᵘ +@inline diffusivity(::FlavorOfXKVD, diffusivities, id) = diffusivities.κᶜ + +with_tracers(tracers, closure::FlavorOfXKVD) = closure + +# Note: computing diffusivities at cell centers for now. +function DiffusivityFields(grid, tracer_names, bcs, closure::FlavorOfXKVD) + κᶜ = Field((Center, Center, Face), grid) + κᵘ = Field((Center, Center, Face), grid) + Ri = Field((Center, Center, Face), grid) + return (; κᶜ, κᵘ, Ri) +end + +function compute_diffusivities!(diffusivities, closure::FlavorOfXKVD, model; parameters = :xyz) + arch = model.architecture + grid = model.grid + clock = model.clock + tracers = model.tracers + buoyancy = model.buoyancy + velocities = model.velocities + top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) + + Nx_in, Ny_in, Nz_in = total_size(diffusivities.κᶜ) + ox_in, oy_in, oz_in = diffusivities.κᶜ.data.offsets + + kp = KernelParameters((Nx_in, Ny_in, Nz_in), (ox_in, oy_in, oz_in)) + + launch!(arch, grid, kp, + compute_ri_number!, + diffusivities, + grid, + closure, + velocities, + tracers, + buoyancy, + top_tracer_bcs, + clock) + + # Use `only_local_halos` to ensure that no communication occurs during + # this call to fill_halo_regions! + fill_halo_regions!(diffusivities.Ri; only_local_halos=true) + + launch!(arch, grid, kp, + compute_xinkai_diffusivities!, + diffusivities, + grid, + closure, + velocities, + tracers, + buoyancy, + top_tracer_bcs, + clock) + + return nothing +end + +@inline ϕ²(i, j, k, grid, ϕ, args...) = ϕ(i, j, k, grid, args...)^2 + +@inline function shear_squaredᶜᶜᶠ(i, j, k, grid, velocities) + ∂z_u² = ℑxᶜᵃᵃ(i, j, k, grid, ϕ², ∂zᶠᶜᶠ, velocities.u) + ∂z_v² = ℑyᵃᶜᵃ(i, j, k, grid, ϕ², ∂zᶜᶠᶠ, velocities.v) + return ∂z_u² + ∂z_v² +end + +@inline function Riᶜᶜᶠ(i, j, k, grid, velocities, buoyancy, tracers) + S² = shear_squaredᶜᶜᶠ(i, j, k, grid, velocities) + N² = ∂z_b(i, j, k, grid, buoyancy, tracers) + Ri = N² / S² + + # Clip N² and avoid NaN + return ifelse(N² == 0, zero(grid), Ri) +end + +const c = Center() +const f = Face() + +@kernel function compute_ri_number!(diffusivities, grid, closure::FlavorOfXKVD, + velocities, tracers, buoyancy, tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + @inbounds diffusivities.Ri[i, j, k] = Riᶜᶜᶠ(i, j, k, grid, velocities, buoyancy, tracers) +end + +@kernel function compute_xinkai_diffusivities!(diffusivities, grid, closure::FlavorOfXKVD, + velocities, tracers, buoyancy, tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + _compute_xinkai_diffusivities!(i, j, k, diffusivities, grid, closure, + velocities, tracers, buoyancy, tracer_bcs, clock) +end + +@inline function _compute_xinkai_diffusivities!(i, j, k, diffusivities, grid, closure, + velocities, tracers, buoyancy, tracer_bcs, clock) + + # Ensure this works with "ensembles" of closures, in addition to ordinary single closures + closure_ij = getclosure(i, j, closure) + + ν₀ = closure_ij.ν₀ + νˢʰ = closure_ij.νˢʰ + νᶜⁿ = closure_ij.νᶜⁿ + Pr_convₜ = closure_ij.Pr_convₜ + Pr_shearₜ = closure_ij.Pr_shearₜ + Riᶜ = closure_ij.Riᶜ + δRi = closure_ij.δRi + + κ₀ = ν₀ / Pr_shearₜ + κˢʰ = νˢʰ / Pr_shearₜ + κᶜⁿ = νᶜⁿ / Pr_convₜ + + # (Potentially) apply a horizontal filter to the Richardson number + Ri = ℑxyᶜᶜᵃ(i, j, k, grid, ℑxyᶠᶠᵃ, diffusivities.Ri) + + # Conditions + convecting = Ri < 0 # applies regardless of Qᵇ + + # Convective adjustment diffusivity + ν_local = ifelse(convecting, (νˢʰ - νᶜⁿ) * tanh(Ri / δRi) + νˢʰ, clamp((ν₀ - νˢʰ) * Ri / Riᶜ + νˢʰ, ν₀, νˢʰ)) + κ_local = ifelse(convecting, (κˢʰ - κᶜⁿ) * tanh(Ri / δRi) + κˢʰ, clamp((κ₀ - κˢʰ) * Ri / Riᶜ + κˢʰ, κ₀, κˢʰ)) + + # Update by averaging in time + @inbounds diffusivities.κᵘ[i, j, k] = ifelse(k <= 1 || k >= grid.Nz+1, 0, ν_local) + @inbounds diffusivities.κᶜ[i, j, k] = ifelse(k <= 1 || k >= grid.Nz+1, 0, κ_local) + + return nothing +end From 3609d6c84a3c146e0a6f83dfe0f6285fd668c64b Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 19 Sep 2024 18:03:44 -0400 Subject: [PATCH 050/114] new NN closure with nof and base boundary layer criteria --- NN_closure_global_nof_BBL.jl | 218 +++++++++++++++++++++++++++++++++++ 1 file changed, 218 insertions(+) create mode 100644 NN_closure_global_nof_BBL.jl diff --git a/NN_closure_global_nof_BBL.jl b/NN_closure_global_nof_BBL.jl new file mode 100644 index 0000000000..9695cc8dfd --- /dev/null +++ b/NN_closure_global_nof_BBL.jl @@ -0,0 +1,218 @@ +import Oceananigans.TurbulenceClosures: + compute_diffusivities!, + DiffusivityFields, + viscosity, + diffusivity, + diffusive_flux_x, + diffusive_flux_y, + diffusive_flux_z, + viscous_flux_ux, + viscous_flux_uy, + viscous_flux_uz, + viscous_flux_vx, + viscous_flux_vy, + viscous_flux_vz, + viscous_flux_wx, + viscous_flux_wy, + viscous_flux_wz + +using Oceananigans.BuoyancyModels: ∂x_b, ∂y_b, ∂z_b +using Oceananigans.Coriolis +using Oceananigans.Grids: φnode +using Oceananigans.Grids: total_size +using Oceananigans.Utils: KernelParameters +using Oceananigans: architecture, on_architecture +using Lux, LuxCUDA +using JLD2 +using ComponentArrays +using OffsetArrays + +using KernelAbstractions: @index, @kernel, @private + +import Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure, ExplicitTimeDiscretization + +using Adapt + +include("./feature_scaling.jl") + +@inline hack_sind(φ) = sin(φ * π / 180) + +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::HydrostaticSphericalCoriolis) = 2 * coriolis.rotation_rate * hack_sind(φnode(i, j, k, grid, Center(), Center(), Center())) +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::FPlane) = coriolis.f +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::BetaPlane) = coriolis.f₀ + coriolis.β * ynode(i, j, k, grid, Center(), Center(), Center()) + +struct NN{M, P, S} + model :: M + ps :: P + st :: S +end + +@inline (neural_network::NN)(input) = first(neural_network.model(input, neural_network.ps, neural_network.st)) +@inline tosarray(x::AbstractArray) = SArray{Tuple{size(x)...}}(x) + +struct NNFluxClosure{A <: NN, S} <: AbstractTurbulenceClosure{ExplicitTimeDiscretization, 3} + wT :: A + wS :: A + scaling :: S +end + +Adapt.adapt_structure(to, nn :: NNFluxClosure) = + NNFluxClosure(Adapt.adapt(to, nn.wT), + Adapt.adapt(to, nn.wS), + Adapt.adapt(to, nn.scaling)) + +Adapt.adapt_structure(to, nn :: NN) = + NN(Adapt.adapt(to, nn.model), + Adapt.adapt(to, nn.ps), + Adapt.adapt(to, nn.st)) + +function NNFluxClosure(arch) + dev = ifelse(arch == GPU(), gpu_device(), cpu_device()) + nn_path = "./NDE_FC_Qb_nof_BBL_trainFC24new_scalingtrain54new_2layer_64_relu_2Pr_model_temp.jld2" + + ps, sts, scaling_params, wT_model, wS_model = jldopen(nn_path, "r") do file + ps = file["u"] |> dev |> f64 + sts = file["sts"] |> dev |> f64 + scaling_params = file["scaling"] + wT_model = file["model"].wT + wS_model = file["model"].wS + return ps, sts, scaling_params, wT_model, wS_model + end + + scaling = construct_zeromeanunitvariance_scaling(scaling_params) + + wT_NN = NN(wT_model, ps.wT, sts.wT) + wS_NN = NN(wS_model, ps.wS, sts.wS) + + return NNFluxClosure(wT_NN, wS_NN, scaling) +end + +function DiffusivityFields(grid, tracer_names, bcs, ::NNFluxClosure) + arch = architecture(grid) + wT = ZFaceField(grid) + wS = ZFaceField(grid) + ∂ρ²∂z² = ZFaceField(grid) + BBL_index = Field((Center, Center, Nothing), grid, Int32) + + Nx_in, Ny_in, Nz_in = total_size(wT) + ox_in, oy_in, oz_in = wT.data.offsets + + wrk_in = OffsetArray(zeros(10, Nx_in, Ny_in, Nz_in), 0, ox_in, oy_in, oz_in) + wrk_in = on_architecture(arch, wrk_in) + + return (; wrk_in, wT, wS, ∂ρ²∂z², BBL_index) +end + +function compute_diffusivities!(diffusivities, closure::NNFluxClosure, model; parameters = :xyz) + arch = model.architecture + grid = model.grid + velocities = model.velocities + tracers = model.tracers + buoyancy = model.buoyancy + coriolis = model.coriolis + clock = model.clock + top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) + input = diffusivities.wrk_in + wT = diffusivities.wT + wS = diffusivities.wS + ∂ρ²∂z² = diffusivities.∂ρ²∂z² + BBL_index = diffusivities.BBL_index + + Nx_in, Ny_in, Nz_in = total_size(wT) + ox_in, oy_in, oz_in = wT.data.offsets + kp = KernelParameters((Nx_in, Ny_in, Nz_in), (ox_in, oy_in, oz_in)) + kp_2D = KernelParameters((Nx_in, Ny_in), (ox_in, oy_in)) + + launch!(arch, grid, kp, + _populate_input!, input, ∂ρ²∂z², grid, closure, tracers, velocities, buoyancy, coriolis, top_tracer_bcs, clock) + + launch!(arch, grid, kp_2D, _find_base_boundary_layer!, ∂ρ²∂z², grid, BBL_index) + + wT.data.parent .= dropdims(closure.wT(input.parent), dims=1) + wS.data.parent .= dropdims(closure.wS(input.parent), dims=1) + + launch!(arch, grid, kp, _adjust_nn_fluxes!, diffusivities, grid, closure, tracers, velocities, buoyancy, top_tracer_bcs, clock) + return nothing +end + +@kernel function _populate_input!(input, ∂ρ²∂z², grid, closure::NNFluxClosure, tracers, velocities, buoyancy, coriolis, top_tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + + scaling = closure.scaling + + ρ₀ = buoyancy.model.equation_of_state.reference_density + g = buoyancy.model.gravitational_acceleration + + @inbounds input[1, i, j, k] = ∂Tᵢ₋₁ = scaling.∂T∂z(∂zᶜᶜᶠ(i, j, k-1, grid, tracers.T)) + @inbounds input[2, i, j, k] = ∂Tᵢ = scaling.∂T∂z(∂zᶜᶜᶠ(i, j, k, grid, tracers.T)) + @inbounds input[3, i, j, k] = ∂Tᵢ₊₁ = scaling.∂T∂z(∂zᶜᶜᶠ(i, j, k+1, grid, tracers.T)) + + @inbounds input[4, i, j, k] = ∂Sᵢ₋₁ = scaling.∂S∂z(∂zᶜᶜᶠ(i, j, k-1, grid, tracers.S)) + @inbounds input[5, i, j, k] = ∂Sᵢ = scaling.∂S∂z(∂zᶜᶜᶠ(i, j, k, grid, tracers.S)) + @inbounds input[6, i, j, k] = ∂Sᵢ₊₁ = scaling.∂S∂z(∂zᶜᶜᶠ(i, j, k+1, grid, tracers.S)) + + @inbounds input[7, i, j, k] = ∂σᵢ₋₁ = scaling.∂ρ∂z(-ρ₀ * ∂z_b(i, j, k-1, grid, buoyancy, tracers) / g) + @inbounds input[8, i, j, k] = ∂σᵢ = scaling.∂ρ∂z(-ρ₀ * ∂z_b(i, j, k, grid, buoyancy, tracers) / g) + @inbounds input[9, i, j, k] = ∂σᵢ₊₁ = scaling.∂ρ∂z(-ρ₀ * ∂z_b(i, j, k+1, grid, buoyancy, tracers) / g) + + @inbounds input[10, i, j, k] = Jᵇ = scaling.wb(top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers))) + + @inbounds ∂ρ²∂z²[i, j, k] = abs(-ρ₀ * ∂zᶜᶜᶜ(i, j, k, grid, ∂z_b, buoyancy, tracers) / g) +end + +@inline function find_field_max!(i, j, field, grid, h) + kmax = grid.Nz + @inbounds maxf = field[i, j, grid.Nz-1] + + @inbounds for k in grid.Nz-2:-1:2 + kmax = ifelse(field[i, j, k] > maxf, k, kmax) + maxf = ifelse(field[i, j, k] > maxf, field[i, j, k], maxf) + end + + @inbounds h[i, j, 1] = kmax +end + +@kernel function _find_base_boundary_layer!(∂ρ²∂z², grid, h) + i, j = @index(Global, NTuple) + find_field_max!(i, j, ∂ρ²∂z², grid, h) + + @inbounds h[i, j, 1] = ifelse(h[i, j, 1] < 6, ifelse(h[i, j, 1] == 1, grid.Nz+1, 4), h[i, j, 1] - 2) +end + +@kernel function _adjust_nn_fluxes!(diffusivities, grid, closure::NNFluxClosure, tracers, velocities, buoyancy, top_tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + scaling = closure.scaling + convecting = top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers)) > 0 + @inbounds above_base_boundary_layer = k > diffusivities.BBL_index[i, j, 1] & k <= grid.Nz - 1 + + @inbounds diffusivities.wT[i, j, k] = ifelse(convecting & above_base_boundary_layer, scaling.wT.σ * diffusivities.wT[i, j, k], 0) + @inbounds diffusivities.wS[i, j, k] = ifelse(convecting & above_base_boundary_layer, scaling.wS.σ * diffusivities.wS[i, j, k], 0) +end + +# Write here your constructor +# NNFluxClosure() = ... insert NN here ... (make sure it is on GPU if you need it on GPU!) + +const NNC = NNFluxClosure + +##### +##### Abstract Smagorinsky functionality +##### + +# Horizontal fluxes are zero! +@inline viscous_flux_wz( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_ux( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_uy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline diffusive_flux_x(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) +@inline diffusive_flux_y(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) + +# Viscous fluxes are zero (for now) +@inline viscous_flux_uz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) + +# The only function extended by NNFluxClosure +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{1}, c, clock, fields, buoyancy) = @inbounds K.wT[i, j, k] +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{2}, c, clock, fields, buoyancy) = @inbounds K.wS[i, j, k] \ No newline at end of file From 8eaaaaeb94de17ab982f96300ece10c475cb46f7 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 19 Sep 2024 22:21:54 -0400 Subject: [PATCH 051/114] fix bug in NN closure implementation --- NN_closure_global_nof_BBL.jl | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/NN_closure_global_nof_BBL.jl b/NN_closure_global_nof_BBL.jl index 9695cc8dfd..2cbd6ad3a5 100644 --- a/NN_closure_global_nof_BBL.jl +++ b/NN_closure_global_nof_BBL.jl @@ -176,17 +176,23 @@ end i, j = @index(Global, NTuple) find_field_max!(i, j, ∂ρ²∂z², grid, h) - @inbounds h[i, j, 1] = ifelse(h[i, j, 1] < 6, ifelse(h[i, j, 1] == 1, grid.Nz+1, 4), h[i, j, 1] - 2) + @inbounds h[i, j, 1] = ifelse(h[i, j, 1] < 7, ifelse(h[i, j, 1] == 2, grid.Nz+1, 4), h[i, j, 1] - 3) end @kernel function _adjust_nn_fluxes!(diffusivities, grid, closure::NNFluxClosure, tracers, velocities, buoyancy, top_tracer_bcs, clock) i, j, k = @index(Global, NTuple) scaling = closure.scaling + @inbounds BBL_index = diffusivities.BBL_index[i, j, 1] + convecting = top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers)) > 0 - @inbounds above_base_boundary_layer = k > diffusivities.BBL_index[i, j, 1] & k <= grid.Nz - 1 + above_base_boundary_layer = k > BBL_index + below_top = k <= grid.Nz - 1 + above_bottom = k >= 3 + + NN_active = convecting & above_base_boundary_layer & below_top & above_bottom - @inbounds diffusivities.wT[i, j, k] = ifelse(convecting & above_base_boundary_layer, scaling.wT.σ * diffusivities.wT[i, j, k], 0) - @inbounds diffusivities.wS[i, j, k] = ifelse(convecting & above_base_boundary_layer, scaling.wS.σ * diffusivities.wS[i, j, k], 0) + @inbounds diffusivities.wT[i, j, k] = ifelse(NN_active, scaling.wT.σ * diffusivities.wT[i, j, k], 0) + @inbounds diffusivities.wS[i, j, k] = ifelse(NN_active, scaling.wS.σ * diffusivities.wS[i, j, k], 0) end # Write here your constructor From 6535be52f6f318639fa3bf416bbf1c06868f57c2 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Mon, 23 Sep 2024 14:04:54 -0400 Subject: [PATCH 052/114] using Grids.total_size --- xin_kai_vertical_diffusivity_local.jl | 2 +- xin_kai_vertical_diffusivity_local_2step.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/xin_kai_vertical_diffusivity_local.jl b/xin_kai_vertical_diffusivity_local.jl index 6388d67a98..9398dd002e 100644 --- a/xin_kai_vertical_diffusivity_local.jl +++ b/xin_kai_vertical_diffusivity_local.jl @@ -2,7 +2,7 @@ using Oceananigans using Oceananigans.Architectures: architecture using Oceananigans.BuoyancyModels: ∂z_b using Oceananigans.Operators -using Oceananigans.Grids: inactive_node +using Oceananigans.Grids: inactive_node, total_size using Oceananigans.Operators: ℑzᵃᵃᶜ, ℑxyᶠᶠᵃ, ℑxyᶜᶜᵃ using Adapt diff --git a/xin_kai_vertical_diffusivity_local_2step.jl b/xin_kai_vertical_diffusivity_local_2step.jl index e1f995bc00..640f07c6d0 100644 --- a/xin_kai_vertical_diffusivity_local_2step.jl +++ b/xin_kai_vertical_diffusivity_local_2step.jl @@ -2,7 +2,7 @@ using Oceananigans using Oceananigans.Architectures: architecture using Oceananigans.BuoyancyModels: ∂z_b using Oceananigans.Operators -using Oceananigans.Grids: inactive_node +using Oceananigans.Grids: inactive_node, total_size using Oceananigans.Operators: ℑzᵃᵃᶜ, ℑxyᶠᶠᵃ, ℑxyᶜᶜᵃ using Adapt From 812df887ea4ae112b1613269fabd2bf6b4a937c8 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Mon, 23 Sep 2024 14:27:00 -0400 Subject: [PATCH 053/114] add using KernelParameters --- xin_kai_vertical_diffusivity_local.jl | 1 + xin_kai_vertical_diffusivity_local_2step.jl | 1 + 2 files changed, 2 insertions(+) diff --git a/xin_kai_vertical_diffusivity_local.jl b/xin_kai_vertical_diffusivity_local.jl index 9398dd002e..f911b40f2d 100644 --- a/xin_kai_vertical_diffusivity_local.jl +++ b/xin_kai_vertical_diffusivity_local.jl @@ -4,6 +4,7 @@ using Oceananigans.BuoyancyModels: ∂z_b using Oceananigans.Operators using Oceananigans.Grids: inactive_node, total_size using Oceananigans.Operators: ℑzᵃᵃᶜ, ℑxyᶠᶠᵃ, ℑxyᶜᶜᵃ +using Oceananigans.Utils: KernelParameters using Adapt diff --git a/xin_kai_vertical_diffusivity_local_2step.jl b/xin_kai_vertical_diffusivity_local_2step.jl index 640f07c6d0..2c26e09c9a 100644 --- a/xin_kai_vertical_diffusivity_local_2step.jl +++ b/xin_kai_vertical_diffusivity_local_2step.jl @@ -4,6 +4,7 @@ using Oceananigans.BuoyancyModels: ∂z_b using Oceananigans.Operators using Oceananigans.Grids: inactive_node, total_size using Oceananigans.Operators: ℑzᵃᵃᶜ, ℑxyᶠᶠᵃ, ℑxyᶜᶜᵃ +using Oceananigans.Utils: KernelParameters using Adapt From eb50cca98bfaa65206d97186c783d72c93c6d29b Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 24 Sep 2024 12:11:08 -0400 Subject: [PATCH 054/114] update NN model --- NN_closure_global_nof_BBL.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NN_closure_global_nof_BBL.jl b/NN_closure_global_nof_BBL.jl index 2cbd6ad3a5..995e3aa2b9 100644 --- a/NN_closure_global_nof_BBL.jl +++ b/NN_closure_global_nof_BBL.jl @@ -68,7 +68,7 @@ Adapt.adapt_structure(to, nn :: NN) = function NNFluxClosure(arch) dev = ifelse(arch == GPU(), gpu_device(), cpu_device()) - nn_path = "./NDE_FC_Qb_nof_BBL_trainFC24new_scalingtrain54new_2layer_64_relu_2Pr_model_temp.jld2" + nn_path = "./NDE_FC_Qb_nof_BBL_trainFC24new_scalingtrain54new_2layer_64_relu_2Pr_model.jld2" ps, sts, scaling_params, wT_model, wS_model = jldopen(nn_path, "r") do file ps = file["u"] |> dev |> f64 From b2cd2fa6896d19ec0ae741f38dd0554119bc5e1e Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 24 Sep 2024 12:11:33 -0400 Subject: [PATCH 055/114] use BBL integral metric to compute base of boundary layer --- NN_closure_global_nof_BBLintegral.jl | 224 +++++++++++++++++++++++++++ 1 file changed, 224 insertions(+) create mode 100644 NN_closure_global_nof_BBLintegral.jl diff --git a/NN_closure_global_nof_BBLintegral.jl b/NN_closure_global_nof_BBLintegral.jl new file mode 100644 index 0000000000..781b8e8902 --- /dev/null +++ b/NN_closure_global_nof_BBLintegral.jl @@ -0,0 +1,224 @@ +import Oceananigans.TurbulenceClosures: + compute_diffusivities!, + DiffusivityFields, + viscosity, + diffusivity, + diffusive_flux_x, + diffusive_flux_y, + diffusive_flux_z, + viscous_flux_ux, + viscous_flux_uy, + viscous_flux_uz, + viscous_flux_vx, + viscous_flux_vy, + viscous_flux_vz, + viscous_flux_wx, + viscous_flux_wy, + viscous_flux_wz + +using Oceananigans.BuoyancyModels: ∂x_b, ∂y_b, ∂z_b +using Oceananigans.Coriolis +using Oceananigans.Grids: φnode +using Oceananigans.Grids: total_size +using Oceananigans.Utils: KernelParameters +using Oceananigans: architecture, on_architecture +using Lux, LuxCUDA +using JLD2 +using ComponentArrays +using OffsetArrays +using SeawaterPolynomials.TEOS10 + +using KernelAbstractions: @index, @kernel, @private + +import Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure, ExplicitTimeDiscretization + +using Adapt + +include("./feature_scaling.jl") + +@inline hack_sind(φ) = sin(φ * π / 180) + +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::HydrostaticSphericalCoriolis) = 2 * coriolis.rotation_rate * hack_sind(φnode(i, j, k, grid, Center(), Center(), Center())) +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::FPlane) = coriolis.f +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::BetaPlane) = coriolis.f₀ + coriolis.β * ynode(i, j, k, grid, Center(), Center(), Center()) + +struct NN{M, P, S} + model :: M + ps :: P + st :: S +end + +@inline (neural_network::NN)(input) = first(neural_network.model(input, neural_network.ps, neural_network.st)) +@inline tosarray(x::AbstractArray) = SArray{Tuple{size(x)...}}(x) + +struct NNFluxClosure{A <: NN, S} <: AbstractTurbulenceClosure{ExplicitTimeDiscretization, 3} + wT :: A + wS :: A + scaling :: S +end + +Adapt.adapt_structure(to, nn :: NNFluxClosure) = + NNFluxClosure(Adapt.adapt(to, nn.wT), + Adapt.adapt(to, nn.wS), + Adapt.adapt(to, nn.scaling)) + +Adapt.adapt_structure(to, nn :: NN) = + NN(Adapt.adapt(to, nn.model), + Adapt.adapt(to, nn.ps), + Adapt.adapt(to, nn.st)) + +function NNFluxClosure(arch) + dev = ifelse(arch == GPU(), gpu_device(), cpu_device()) + nn_path = "./NDE_FC_Qb_nof_BBLintegral_trainFC24new_scalingtrain54new_2layer_64_relu_2Pr_model.jld2" + + ps, sts, scaling_params, wT_model, wS_model = jldopen(nn_path, "r") do file + ps = file["u"] |> dev |> f64 + sts = file["sts"] |> dev |> f64 + scaling_params = file["scaling"] + wT_model = file["model"].wT + wS_model = file["model"].wS + return ps, sts, scaling_params, wT_model, wS_model + end + + scaling = construct_zeromeanunitvariance_scaling(scaling_params) + + wT_NN = NN(wT_model, ps.wT, sts.wT) + wS_NN = NN(wS_model, ps.wS, sts.wS) + + return NNFluxClosure(wT_NN, wS_NN, scaling) +end + +function DiffusivityFields(grid, tracer_names, bcs, ::NNFluxClosure) + arch = architecture(grid) + wT = ZFaceField(grid) + wS = ZFaceField(grid) + BBL_constraint = CenterField(grid) + BBL_index = Field((Center, Center, Nothing), grid, Int32) + + Nx_in, Ny_in, Nz_in = total_size(wT) + ox_in, oy_in, oz_in = wT.data.offsets + + wrk_in = OffsetArray(zeros(10, Nx_in, Ny_in, Nz_in), 0, ox_in, oy_in, oz_in) + wrk_in = on_architecture(arch, wrk_in) + + return (; wrk_in, wT, wS, BBL_constraint, BBL_index) +end + +function compute_diffusivities!(diffusivities, closure::NNFluxClosure, model; parameters = :xyz) + arch = model.architecture + grid = model.grid + velocities = model.velocities + tracers = model.tracers + buoyancy = model.buoyancy + coriolis = model.coriolis + clock = model.clock + top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) + input = diffusivities.wrk_in + wT = diffusivities.wT + wS = diffusivities.wS + BBL_constraint = diffusivities.BBL_constraint + BBL_index = diffusivities.BBL_index + + Nx_in, Ny_in, Nz_in = total_size(wT) + ox_in, oy_in, oz_in = wT.data.offsets + kp = KernelParameters((Nx_in, Ny_in, Nz_in), (ox_in, oy_in, oz_in)) + kp_2D = KernelParameters((Nx_in, Ny_in), (ox_in, oy_in)) + + launch!(arch, grid, kp, + _populate_input!, input, BBL_constraint, grid, closure, tracers, velocities, buoyancy, coriolis, top_tracer_bcs, clock) + + launch!(arch, grid, kp_2D, _find_base_boundary_layer!, BBL_constraint, grid, BBL_index) + + wT.data.parent .= dropdims(closure.wT(input.parent), dims=1) + wS.data.parent .= dropdims(closure.wS(input.parent), dims=1) + + launch!(arch, grid, kp, _adjust_nn_fluxes!, diffusivities, grid, closure, tracers, velocities, buoyancy, top_tracer_bcs, clock) + return nothing +end + +@kernel function _populate_input!(input, BBL_constraint, grid, closure::NNFluxClosure, tracers, velocities, buoyancy, coriolis, top_tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + + scaling = closure.scaling + + ρ₀ = buoyancy.model.equation_of_state.reference_density + g = buoyancy.model.gravitational_acceleration + eos = TEOS10.TEOS10EquationOfState() + + T, S = tracers.T, tracers.S + + @inbounds input[1, i, j, k] = ∂Tᵢ₋₁ = scaling.∂T∂z(∂zᶜᶜᶠ(i, j, k-1, grid, T)) + @inbounds input[2, i, j, k] = ∂Tᵢ = scaling.∂T∂z(∂zᶜᶜᶠ(i, j, k, grid, T)) + @inbounds input[3, i, j, k] = ∂Tᵢ₊₁ = scaling.∂T∂z(∂zᶜᶜᶠ(i, j, k+1, grid, T)) + + @inbounds input[4, i, j, k] = ∂Sᵢ₋₁ = scaling.∂S∂z(∂zᶜᶜᶠ(i, j, k-1, grid, S)) + @inbounds input[5, i, j, k] = ∂Sᵢ = scaling.∂S∂z(∂zᶜᶜᶠ(i, j, k, grid, S)) + @inbounds input[6, i, j, k] = ∂Sᵢ₊₁ = scaling.∂S∂z(∂zᶜᶜᶠ(i, j, k+1, grid, S)) + + @inbounds input[7, i, j, k] = ∂σᵢ₋₁ = scaling.∂ρ∂z(-ρ₀ * ∂z_b(i, j, k-1, grid, buoyancy, tracers) / g) + @inbounds input[8, i, j, k] = ∂σᵢ = scaling.∂ρ∂z(-ρ₀ * ∂z_b(i, j, k, grid, buoyancy, tracers) / g) + @inbounds input[9, i, j, k] = ∂σᵢ₊₁ = scaling.∂ρ∂z(-ρ₀ * ∂z_b(i, j, k+1, grid, buoyancy, tracers) / g) + + @inbounds input[10, i, j, k] = Jᵇ = scaling.wb(top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers))) + + @inbounds BBL_constraint[i, j, k] = -ρ₀ * ∂z_b(i, j, k+1, grid, buoyancy, tracers) / g - 8 / grid.zᵃᵃᶜ[k] * (TEOS10.ρ(T[i, j, k], S[i, j, k], 0, eos) - TEOS10.ρ(T[i, j, grid.Nz], S[i, j, grid.Nz], 0, eos)) +end + +@inline function find_based_boundary_layer_index!(i, j, field, grid, h) + kmax = 3 + + @inbounds for k in 6:grid.Nz-1 + kmax = ifelse(field[i, j, k] > 0, k-2, kmax) + end + + @inbounds h[i, j, 1] = kmax +end + +@kernel function _find_base_boundary_layer!(BBL_constraint, grid, h) + i, j = @index(Global, NTuple) + find_based_boundary_layer_index!(i, j, BBL_constraint, grid, h) +end + +@kernel function _adjust_nn_fluxes!(diffusivities, grid, closure::NNFluxClosure, tracers, velocities, buoyancy, top_tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + scaling = closure.scaling + @inbounds BBL_index = diffusivities.BBL_index[i, j, 1] + + convecting = top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers)) > 0 + above_base_boundary_layer = k > BBL_index + below_top = k <= grid.Nz - 1 + above_bottom = k >= 3 + + NN_active = convecting & above_base_boundary_layer & below_top & above_bottom + + @inbounds diffusivities.wT[i, j, k] = ifelse(NN_active, scaling.wT.σ * diffusivities.wT[i, j, k], 0) + @inbounds diffusivities.wS[i, j, k] = ifelse(NN_active, scaling.wS.σ * diffusivities.wS[i, j, k], 0) +end + +# Write here your constructor +# NNFluxClosure() = ... insert NN here ... (make sure it is on GPU if you need it on GPU!) + +const NNC = NNFluxClosure + +##### +##### Abstract Smagorinsky functionality +##### + +# Horizontal fluxes are zero! +@inline viscous_flux_wz( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_ux( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_uy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline diffusive_flux_x(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) +@inline diffusive_flux_y(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) + +# Viscous fluxes are zero (for now) +@inline viscous_flux_uz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) + +# The only function extended by NNFluxClosure +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{1}, c, clock, fields, buoyancy) = @inbounds K.wT[i, j, k] +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{2}, c, clock, fields, buoyancy) = @inbounds K.wS[i, j, k] \ No newline at end of file From d3fa8b8e5bcfccf409fd0b257cef32a32540f315 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee <50624521+xkykai@users.noreply.github.com> Date: Fri, 4 Oct 2024 12:39:28 -0400 Subject: [PATCH 056/114] Update xin_kai_vertical_diffusivity.jl Co-authored-by: Gregory L. Wagner --- xin_kai_vertical_diffusivity.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xin_kai_vertical_diffusivity.jl b/xin_kai_vertical_diffusivity.jl index fed133a780..29a482aa75 100644 --- a/xin_kai_vertical_diffusivity.jl +++ b/xin_kai_vertical_diffusivity.jl @@ -248,8 +248,8 @@ end ν_nonlocal = ifelse(entraining, Cᵉⁿ * νᶜⁿ * 0.5 * (tanh((x - Q₀) / δQ) + 1), 0) # Update by averaging in time - @inbounds diffusivities.κᵘ[i, j, k] = ifelse(k <= 1 || k >= grid.Nz+1, 0, ν_local + ν_nonlocal) - @inbounds diffusivities.κᶜ[i, j, k] = ifelse(k <= 1 || k >= grid.Nz+1, 0, (ν_local + ν_nonlocal) / Prₜ) + @inbounds diffusivities.κᵘ[i, j, k] = ifelse((k <= 1) | (k >= grid.Nz+1), 0, ν_local + ν_nonlocal) + @inbounds diffusivities.κᶜ[i, j, k] = ifelse((k <= 1) | (k >= grid.Nz+1), 0, (ν_local + ν_nonlocal) / Prₜ) return nothing end From 2448003c7833544bb3dff069b0f14fe3ced58bdf Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Sat, 5 Oct 2024 21:09:26 -0400 Subject: [PATCH 057/114] remove type piracy TEOS10.s --- 2D_model_LES_sin.jl | 5 +- 3D_model_LES_sin.jl | 7 +- NN_1D_model.jl | 57 +- physicalclosure_doublegyre_model.jl | 53 +- ...calclosure_doublegyre_model_initialized.jl | 592 ++++++++++++++++++ validate_NN_1D_model.jl | 3 +- xk_physicalclosure_doublegyre_model.jl | 31 +- 7 files changed, 709 insertions(+), 39 deletions(-) create mode 100644 physicalclosure_doublegyre_model_initialized.jl diff --git a/2D_model_LES_sin.jl b/2D_model_LES_sin.jl index 3fda63741c..259d370719 100644 --- a/2D_model_LES_sin.jl +++ b/2D_model_LES_sin.jl @@ -9,8 +9,7 @@ using Oceananigans.OutputReaders: FieldTimeSeries using Oceananigans.Grids: xnode, ynode, znode using SeawaterPolynomials using SeawaterPolynomials:TEOS10 -import SeawaterPolynomials.TEOS10: s, ΔS, Sₐᵤ -s(Sᴬ::Number) = Sᴬ + ΔS >= 0 ? √((Sᴬ + ΔS) / Sₐᵤ) : NaN + using Glob # Architecture @@ -40,7 +39,7 @@ const dSdz = 0.0021 const T_surface = 20.0 const S_surface = 36.6 -const max_temperature_flux = 3e-4 +const max_temperature_flux = 2e-4 FILE_DIR = "./LES/NN_2D_channel_sin_cooling_$(max_temperature_flux)_LES" mkpath(FILE_DIR) diff --git a/3D_model_LES_sin.jl b/3D_model_LES_sin.jl index c886878ceb..8e3d0f7766 100644 --- a/3D_model_LES_sin.jl +++ b/3D_model_LES_sin.jl @@ -9,8 +9,7 @@ using Oceananigans.OutputReaders: FieldTimeSeries using Oceananigans.Grids: xnode, ynode, znode using SeawaterPolynomials using SeawaterPolynomials:TEOS10 -import SeawaterPolynomials.TEOS10: s, ΔS, Sₐᵤ -s(Sᴬ::Number) = Sᴬ + ΔS >= 0 ? √((Sᴬ + ΔS) / Sₐᵤ) : NaN + using Glob # Architecture @@ -18,11 +17,11 @@ model_architecture = GPU() # number of grid points Nx = 125 -Ny = 3000 +Ny = 1000 Nz = 250 const Lx = 250meters -const Ly = 6kilometers +const Ly = 2kilometers const Lz = 500meters grid = RectilinearGrid(model_architecture, diff --git a/NN_1D_model.jl b/NN_1D_model.jl index fc8615247b..532cb6cdf4 100644 --- a/NN_1D_model.jl +++ b/NN_1D_model.jl @@ -18,15 +18,13 @@ using Oceananigans.OutputReaders: FieldTimeSeries using Oceananigans.Grids: xnode, ynode, znode using SeawaterPolynomials using SeawaterPolynomials:TEOS10 -import SeawaterPolynomials.TEOS10: s, ΔS, Sₐᵤ -s(Sᴬ::Number) = Sᴬ + ΔS >= 0 ? √((Sᴬ + ΔS) / Sₐᵤ) : NaN + # Architecture model_architecture = CPU() # number of grid points Nz = 64 - const Lz = 512 grid = RectilinearGrid(model_architecture, @@ -46,13 +44,14 @@ const dSdz = 0.0021 const T_surface = 20.0 const S_surface = 36.6 -T_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(1e-4), bottom=GradientBoundaryCondition(dTdz)) +T_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(3e-4)) +u_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(2e-4)) ##### ##### Coriolis ##### -const f₀ = 8e-5 +const f₀ = 1e-4 const β = 1e-11 # coriolis = BetaPlane(f₀=f₀, β = β) coriolis = FPlane(f=f₀) @@ -82,7 +81,8 @@ model = HydrostaticFreeSurfaceModel( closure = (nn_closure, base_closure), # closure = base_closure, tracers = (:T, :S), - boundary_conditions = (; T = T_bcs), + # boundary_conditions = (; T = T_bcs), + boundary_conditions = (; T = T_bcs, u = u_bcs), ) @info "Built $model." @@ -139,10 +139,12 @@ simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterv u, v, w = model.velocities T, S = model.tracers.T, model.tracers.S +ubar = Field(Average(u, dims = (1,2))) +vbar = Field(Average(v, dims = (1,2))) Tbar = Field(Average(T, dims = (1,2))) Sbar = Field(Average(S, dims = (1,2))) -averaged_outputs = (; Tbar, Sbar) +averaged_outputs = (; ubar, vbar, Tbar, Sbar) ##### ##### Build checkpointer and output writer @@ -164,9 +166,11 @@ end # ##### # ##### Visualization # ##### - +#%% using CairoMakie +ubar_data = FieldTimeSeries("./NN_1D_channel_averages.jld2", "ubar") +vbar_data = FieldTimeSeries("./NN_1D_channel_averages.jld2", "vbar") Tbar_data = FieldTimeSeries("./NN_1D_channel_averages.jld2", "Tbar") Sbar_data = FieldTimeSeries("./NN_1D_channel_averages.jld2", "Sbar") @@ -174,27 +178,42 @@ zC = znodes(Tbar_data.grid, Center()) zF = znodes(Tbar_data.grid, Face()) Nt = length(Tbar_data.times) -#%% -fig = Figure(size = (900, 600)) -axT = CairoMakie.Axis(fig[1, 1], xlabel = "T (°C)", ylabel = "z (m)") -axS = CairoMakie.Axis(fig[1, 2], xlabel = "S (g kg⁻¹)", ylabel = "z (m)") -n = Observable(Nt) +fig = Figure(size = (1800, 600)) +axu = CairoMakie.Axis(fig[1, 1], xlabel = "u (m s⁻¹)", ylabel = "z (m)") +axv = CairoMakie.Axis(fig[1, 2], xlabel = "v (m s⁻¹)", ylabel = "z (m)") +axT = CairoMakie.Axis(fig[1, 3], xlabel = "T (°C)", ylabel = "z (m)") +axS = CairoMakie.Axis(fig[1, 4], xlabel = "S (g kg⁻¹)", ylabel = "z (m)") +# slider = Slider(fig[2, :], range=1:Nt) +n = Observable(1) + +ubarₙ = @lift interior(ubar_data[$n], 1, 1, :) +vbarₙ = @lift interior(vbar_data[$n], 1, 1, :) Tbarₙ = @lift interior(Tbar_data[$n], 1, 1, :) Sbarₙ = @lift interior(Sbar_data[$n], 1, 1, :) +ulim = (minimum(ubar_data), maximum(ubar_data)) +vlim = (minimum(vbar_data), maximum(vbar_data)) +Tlim = (minimum(Tbar_data), maximum(Tbar_data)) +Slim = (minimum(Sbar_data), maximum(Sbar_data)) + title_str = @lift "Time: $(round(Tbar_data.times[$n] / 86400, digits=3)) days" +lines!(axu, ubarₙ, zC) +lines!(axv, vbarₙ, zC) lines!(axT, Tbarₙ, zC) lines!(axS, Sbarₙ, zC) +xlims!(axu, ulim) +xlims!(axv, vlim) +xlims!(axT, Tlim) +xlims!(axS, Slim) + Label(fig[0, :], title_str, tellwidth = false) -# CairoMakie.record(fig, "./NN_1D_fields.mp4", 1:Nt, framerate=10) do nn -# n[] = nn -# xlims!(axT, nothing, nothing) -# xlims!(axS, nothing, nothing) -# end +CairoMakie.record(fig, "./NN_1D_fields.mp4", 1:Nt, framerate=60) do nn + n[] = nn +end -display(fig) +# display(fig) #%% \ No newline at end of file diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index 4583664fa8..c037e47cdd 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -25,11 +25,10 @@ using SeawaterPolynomials using SeawaterPolynomials:TEOS10 using ColorSchemes using Glob -import SeawaterPolynomials.TEOS10: s, ΔS, Sₐᵤ -s(Sᴬ::Number) = Sᴬ + ΔS >= 0 ? √((Sᴬ + ΔS) / Sₐᵤ) : NaN + #%% -filename = "doublegyre_CATKEVerticalDiffusivity" +filename = "doublegyre_relaxation_30days_CATKEVerticalDiffusivity" FILE_DIR = "./Output/$(filename)" mkpath(FILE_DIR) @@ -82,7 +81,7 @@ const S_mid = (S_north + S_south) / 2 const τ₀ = 1e-4 const μ_drag = 1/30days -const μ_T = 1/8days +const μ_T = 1/30days ##### ##### Forcing and initial condition @@ -212,7 +211,6 @@ simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterv ##### ##### Diagnostics ##### - u, v, w = model.velocities T, S = model.tracers.T, model.tracers.S U_bt = Field(Integral(u, dims=3)) @@ -263,6 +261,51 @@ simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, indices = (:, 1, :), schedule = TimeInterval(10days)) +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + simulation.output_writers[:xz_south] = JLD2OutputWriter(model, outputs, filename = "$(FILE_DIR)/instantaneous_fields_xz_south", indices = (:, 25, :), diff --git a/physicalclosure_doublegyre_model_initialized.jl b/physicalclosure_doublegyre_model_initialized.jl new file mode 100644 index 0000000000..df33e85250 --- /dev/null +++ b/physicalclosure_doublegyre_model_initialized.jl @@ -0,0 +1,592 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity +using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: CATKEVerticalDiffusivity, CATKEMixingLength, CATKEEquation + +using Oceananigans.BuoyancyModels: ∂z_b +# include("NN_closure_global.jl") +# include("xin_kai_vertical_diffusivity_local.jl") +# include("xin_kai_vertical_diffusivity_2Pr.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes +using Glob + + +#%% +filename = "doublegyre_relaxation_30days_CATKEVerticalDiffusivity_initialized_test" +FILE_DIR = "./Output/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +# vertical_base_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +# convection_closure = XinKaiVerticalDiffusivity() +function CATKE_ocean_closure() + mixing_length = CATKEMixingLength(Cᵇ=0.01) + turbulent_kinetic_energy_equation = CATKEEquation(Cᵂϵ=1.0) + return CATKEVerticalDiffusivity(; mixing_length, turbulent_kinetic_energy_equation) +end +convection_closure = CATKE_ocean_closure() +closure = convection_closure +# closure = vertical_base_closure + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/30days + +##### +##### Forcing and initial condition +##### + +@inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S, :e), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S, :e), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +# noise(z) = rand() * exp(z / 8) + +# T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +# S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) +DATA_DIR = "./Output/doublegyre_30Cwarmflush_relaxation_8days_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" + +u_data = FieldTimeSeries("$(DATA_DIR)/instantaneous_fields.jld2", "u", backend=OnDisk()) +v_data = FieldTimeSeries("$(DATA_DIR)/instantaneous_fields.jld2", "v", backend=OnDisk()) +T_data = FieldTimeSeries("$(DATA_DIR)/instantaneous_fields.jld2", "T", backend=OnDisk()) +S_data = FieldTimeSeries("$(DATA_DIR)/instantaneous_fields.jld2", "S", backend=OnDisk()) + +ntimes = length(u_data.times) + +set!(model, T=T_data[ntimes], S=S_data[ntimes], u=u_data[ntimes], v=v_data[ntimes]) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10950days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)) +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) +compute!(ρ) + +ρᶠ = @at (Center, Center, Face) ρ +∂ρ∂z = ∂z(ρ) +∂²ρ∂z² = ∂z(∂ρ∂z) + +κc = model.diffusivity_fields.κc +wT = κc * ∂z(T) +wS = κc * ∂z(S) + +outputs = (; u, v, w, T, S, ρ, ρᶠ, ∂ρ∂z, ∂²ρ∂z², N², wT, wS) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_south] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_south", + indices = (:, 25, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_north] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_north", + indices = (:, 75, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn + n[] = nn +end + +@info "Done!" +#%% +Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +Nt = length(Ψ_data) +times = Ψ_data.times / 24 / 60^2 / 365 +#%% +timeframe = Nt +Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +clim = maximum(abs, Ψ_frame) + 1e-13 +N_levels = 16 +levels = range(-clim, stop=clim, length=N_levels) +fig = Figure(size=(800, 800)) +ax = Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +tightlimits!(ax) +save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +display(fig) +#%% +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +N²_xz_north_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_north.jld2", "N²") +N²_xz_south_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_south.jld2", "N²") + +xC = N²_xz_north_data.grid.xᶜᵃᵃ[1:N²_xz_north_data.grid.Nx] +zf = N²_xz_north_data.grid.zᵃᵃᶠ[1:N²_xz_north_data.grid.Nz+1] + +yloc_north = N²_xz_north_data.grid.yᵃᶠᵃ[N²_xz_north_data.indices[2][1]] +yloc_south = N²_xz_south_data.grid.yᵃᶠᵃ[N²_xz_south_data.indices[2][1]] + +Nt = length(N²_xz_north_data) +times = N²_xz_north_data.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +N²_lim = (find_min(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) - 1e-13, + find_max(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) + 1e-13) +#%% +fig = Figure(size=(800, 800)) +ax_north = Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_south / 1e3)) km") +ax_south = Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_north / 1e3)) km") + +n = Observable(2) + +N²_north = @lift interior(N²_xz_north_data[$n], :, 1, :) +N²_south = @lift interior(N²_xz_south_data[$n], :, 1, :) + +colorscheme = colorschemes[:jet] + +N²_north_surface = heatmap!(ax_north, xC, zf, N²_north, colormap=colorscheme, colorrange=N²_lim) +N²_south_surface = heatmap!(ax_south, xC, zf, N²_south, colormap=colorscheme, colorrange=N²_lim) + +Colorbar(fig[1:2, 2], N²_north_surface, label="N² (s⁻²)") + +title_str = @lift "Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +@info "Recording buoyancy frequency xz slice" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_buoyancy_frequency_xz_slice.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn + n[] = nn +end + +@info "Done!" +#%% \ No newline at end of file diff --git a/validate_NN_1D_model.jl b/validate_NN_1D_model.jl index f2df7292b8..c7b93d2cd6 100644 --- a/validate_NN_1D_model.jl +++ b/validate_NN_1D_model.jl @@ -18,8 +18,7 @@ using Oceananigans.OutputReaders: FieldTimeSeries using Oceananigans.Grids: xnode, ynode, znode using SeawaterPolynomials using SeawaterPolynomials:TEOS10 -import SeawaterPolynomials.TEOS10: s, ΔS, Sₐᵤ -s(Sᴬ::Number) = Sᴬ + ΔS >= 0 ? √((Sᴬ + ΔS) / Sₐᵤ) : NaN + # Architecture model_architecture = CPU() diff --git a/xk_physicalclosure_doublegyre_model.jl b/xk_physicalclosure_doublegyre_model.jl index 488c109cd0..0001b36b20 100644 --- a/xk_physicalclosure_doublegyre_model.jl +++ b/xk_physicalclosure_doublegyre_model.jl @@ -21,11 +21,11 @@ using SeawaterPolynomials using SeawaterPolynomials:TEOS10 using ColorSchemes using Glob -import SeawaterPolynomials.TEOS10: s, ΔS, Sₐᵤ -s(Sᴬ::Number) = Sᴬ + ΔS >= 0 ? √((Sᴬ + ΔS) / Sₐᵤ) : NaN + #%% -FILE_DIR = "./Output/doublegyre_XinKaiVerticalDiffusivity" +filename = "doublegyre_XinKaiVerticalDiffusivity_streamfunction" +FILE_DIR = "./Output/$(filename)" mkpath(FILE_DIR) # Architecture @@ -175,7 +175,7 @@ update_state!(model) ##### Simulation building ##### Δt₀ = 5minutes -stop_time = 3650days +stop_time = 10950days simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) @@ -214,6 +214,8 @@ T, S = model.tracers.T, model.tracers.S # ν, κ = model.diffusivity_fields[2].κᵘ, model.diffusivity_fields[2].κᶜ # Ri = model.diffusivity_fields[2].Ri # wT, wS = model.diffusivity_fields[2].wT, model.diffusivity_fields[2].wS +U_bt = Field(Integral(u, dims=3)) +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)) # outputs = (; u, v, w, T, S, ν, κ, Ri, wT, wS) # outputs = (; u, v, w, T, S, ν, κ, Ri) @@ -237,8 +239,25 @@ simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, indices = (:, 1, :), schedule = TimeInterval(5days)) +simulation.output_writers[:xz_south] = JLD2OutputWriter(model, outputs, + # filename = "NN_closure_2D_channel_NDE_FC_Qb_18simnew_2layer_128_relu_2Pr", + filename = "$(FILE_DIR)/instantaneous_fields_xz_south", + indices = (:, 25, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_north] = JLD2OutputWriter(model, outputs, + # filename = "NN_closure_2D_channel_NDE_FC_Qb_18simnew_2layer_128_relu_2Pr", + filename = "$(FILE_DIR)/instantaneous_fields_xz_north", + indices = (:, 75, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + # filename = "NN_closure_2D_channel_NDE_FC_Qb_18simnew_2layer_128_relu_2Pr", + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(365days, window=365days)) + simulation.output_writers[:checkpointer] = Checkpointer(model, - schedule = TimeInterval(365days), + schedule = TimeInterval(730days), prefix = "$(FILE_DIR)/checkpointer") @info "Running the simulation..." @@ -419,7 +438,7 @@ zlims!(axS, (-Lz, 0)) zlims!(axu, (-Lz, 0)) zlims!(axv, (-Lz, 0)) -CairoMakie.record(fig, "$(FILE_DIR)/3D_instantaneous_fields.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn @info nn n[] = nn end From cb5f86455fee07f3dd7f0560d77bdebaa4d099a0 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Sat, 5 Oct 2024 21:12:28 -0400 Subject: [PATCH 058/114] NN closure using BBL zone below nonbabkground kappa --- ...zonelast41_doublegyre_model_initialized.jl | 588 ++++++++++++++++++ 1 file changed, 588 insertions(+) create mode 100644 NNclosure_nof_BBLkappazonelast41_doublegyre_model_initialized.jl diff --git a/NNclosure_nof_BBLkappazonelast41_doublegyre_model_initialized.jl b/NNclosure_nof_BBLkappazonelast41_doublegyre_model_initialized.jl new file mode 100644 index 0000000000..b39892e74b --- /dev/null +++ b/NNclosure_nof_BBLkappazonelast41_doublegyre_model_initialized.jl @@ -0,0 +1,588 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("NN_closure_global_nof_BBLkappazonelast41.jl") +include("xin_kai_vertical_diffusivity_local_2step.jl") +# include("xin_kai_vertical_diffusivity_2Pr.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + + +#%% +filename = "doublegyre_relaxation_30days_NN_closure_2D_channel_NDE_FC_Qb_nof_BBLkappazonelast41_trainFC24new_scalingtrain54new_2layer_64_relu_2Pr_initialized_41" +FILE_DIR = "./Output/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +nn_closure = NNFluxClosure(model_architecture) +base_closure = XinKaiLocalVerticalDiffusivity() +vertical_scalar_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +closure = (base_closure, nn_closure, vertical_scalar_closure) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/30days + +##### +##### Forcing and initial condition +##### + +@inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +# noise(z) = rand() * exp(z / 8) + +# T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +# S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +# set!(model, T=T_initial_noisy, S=S_initial_noisy) +# DATA_DIR = "./Output/doublegyre_30Cwarmflush_relaxation_8days_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" +DATA_DIR = "./Output/doublegyre_30Cwarmflush_relaxation_30days_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" + +u_data = FieldTimeSeries("$(DATA_DIR)/instantaneous_fields.jld2", "u", backend=OnDisk()) +v_data = FieldTimeSeries("$(DATA_DIR)/instantaneous_fields.jld2", "v", backend=OnDisk()) +T_data = FieldTimeSeries("$(DATA_DIR)/instantaneous_fields.jld2", "T", backend=OnDisk()) +S_data = FieldTimeSeries("$(DATA_DIR)/instantaneous_fields.jld2", "S", backend=OnDisk()) + +ntimes = length(u_data.times) + +set!(model, T=T_data[ntimes], S=S_data[ntimes], u=u_data[ntimes], v=v_data[ntimes]) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 5110days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); +BBL_index = model.diffusivity_fields[2].BBL_index +wT_NN = model.diffusivity_fields[2].wT +wS_NN = model.diffusivity_fields[2].wS + +κ = model.diffusivity_fields[1].κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +outputs = (; u, v, w, T, S, ρ, N², wT_NN, wS_NN, wT_base, wS_base) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_south] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_south", + indices = (:, 25, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_north] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_north", + indices = (:, 75, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; BBL_index=BBL_index,), + filename = "$(FILE_DIR)/instantaneous_fields_BBL_index", + indices = (:, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 365 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +N²_xz_north_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_north.jld2", "N²") +N²_xz_south_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_south.jld2", "N²") + +xC = N²_xz_north_data.grid.xᶜᵃᵃ[1:N²_xz_north_data.grid.Nx] +zf = N²_xz_north_data.grid.zᵃᵃᶠ[1:N²_xz_north_data.grid.Nz+1] + +yloc_north = N²_xz_north_data.grid.yᵃᶠᵃ[N²_xz_north_data.indices[2][1]] +yloc_south = N²_xz_south_data.grid.yᵃᶠᵃ[N²_xz_south_data.indices[2][1]] + +Nt = length(N²_xz_north_data) +times = N²_xz_north_data.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +N²_lim = (find_min(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) - 1e-13, + find_max(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) + 1e-13) +#%% +fig = Figure(size=(800, 800)) +ax_north = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_south / 1e3)) km") +ax_south = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_north / 1e3)) km") + +n = Observable(2) + +N²_north = @lift interior(N²_xz_north_data[$n], :, 1, :) +N²_south = @lift interior(N²_xz_south_data[$n], :, 1, :) + +colorscheme = colorschemes[:jet] + +N²_north_surface = heatmap!(ax_north, xC, zf, N²_north, colormap=colorscheme, colorrange=N²_lim) +N²_south_surface = heatmap!(ax_south, xC, zf, N²_south, colormap=colorscheme, colorrange=N²_lim) + +Colorbar(fig[1:2, 2], N²_north_surface, label="N² (s⁻²)") + +title_str = @lift "Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +@info "Recording buoyancy frequency xz slice" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_buoyancy_frequency_xz_slice.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end + +@info "Done!" +#%% \ No newline at end of file From f443296d5e48f5e25421e7732e955faf393eef4b Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Sat, 5 Oct 2024 21:59:05 -0400 Subject: [PATCH 059/114] run double gyre with NDE BBLkappazonelast41 --- ...kappazonelast41_doublegyre_model_initialized.jl | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/NNclosure_nof_BBLkappazonelast41_doublegyre_model_initialized.jl b/NNclosure_nof_BBLkappazonelast41_doublegyre_model_initialized.jl index b39892e74b..1ab838295d 100644 --- a/NNclosure_nof_BBLkappazonelast41_doublegyre_model_initialized.jl +++ b/NNclosure_nof_BBLkappazonelast41_doublegyre_model_initialized.jl @@ -215,7 +215,8 @@ u, v, w = model.velocities T, S = model.tracers.T, model.tracers.S U_bt = Field(Integral(u, dims=3)); Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); -BBL_index = model.diffusivity_fields[2].BBL_index +first_index = model.diffusivity_fields[2].first_index +last_index = model.diffusivity_fields[2].last_index wT_NN = model.diffusivity_fields[2].wT wS_NN = model.diffusivity_fields[2].wS @@ -239,6 +240,13 @@ end ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) ρ = Field(ρ_op) +@inline function get_top_buoyancy_flux(i, j, k, grid, buoyancy, T_bc, S_bc, velocities, tracers, clock) + return top_buoyancy_flux(i, j, grid, buoyancy, (; T=T_bc, S=S_bc), clock, merge(velocities, tracers)) +end + +Qb = KernelFunctionOperation{Center, Center, Nothing}(get_top_buoyancy_flux, model.grid, model.buoyancy, T.boundary_conditions.top, S.boundary_conditions.top, model.velocities, model.tracers, model.clock) +Qb = Field(Qb) + outputs = (; u, v, w, T, S, ρ, N², wT_NN, wS_NN, wT_base, wS_base) ##### @@ -314,8 +322,8 @@ simulation.output_writers[:xz_north] = JLD2OutputWriter(model, outputs, indices = (:, 75, :), schedule = TimeInterval(10days)) -simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; BBL_index=BBL_index,), - filename = "$(FILE_DIR)/instantaneous_fields_BBL_index", +simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; first_index, last_index, Qb), + filename = "$(FILE_DIR)/instantaneous_fields_NN_active_diagnostics", indices = (:, :, :), schedule = TimeInterval(10days)) From 2f3eed929b18fa341e1ac2d384620b1d3aebed3e Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Sat, 5 Oct 2024 22:00:04 -0400 Subject: [PATCH 060/114] change initialized state to 8day restoration forcing --- ...ure_nof_BBLkappazonelast41_doublegyre_model_initialized.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/NNclosure_nof_BBLkappazonelast41_doublegyre_model_initialized.jl b/NNclosure_nof_BBLkappazonelast41_doublegyre_model_initialized.jl index 1ab838295d..eacddcbe77 100644 --- a/NNclosure_nof_BBLkappazonelast41_doublegyre_model_initialized.jl +++ b/NNclosure_nof_BBLkappazonelast41_doublegyre_model_initialized.jl @@ -161,8 +161,8 @@ model = HydrostaticFreeSurfaceModel( # S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) # set!(model, T=T_initial_noisy, S=S_initial_noisy) -# DATA_DIR = "./Output/doublegyre_30Cwarmflush_relaxation_8days_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" -DATA_DIR = "./Output/doublegyre_30Cwarmflush_relaxation_30days_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" +DATA_DIR = "./Output/doublegyre_30Cwarmflush_relaxation_8days_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" +# DATA_DIR = "./Output/doublegyre_30Cwarmflush_relaxation_30days_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" u_data = FieldTimeSeries("$(DATA_DIR)/instantaneous_fields.jld2", "u", backend=OnDisk()) v_data = FieldTimeSeries("$(DATA_DIR)/instantaneous_fields.jld2", "v", backend=OnDisk()) From 987f1ed4ee4b51ab5737bf6335246a6d6e5d22b3 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Sat, 5 Oct 2024 22:16:54 -0400 Subject: [PATCH 061/114] run double gyre withj baseclosure initialized --- baseclosure_doublegyre_model_initialized.jl | 575 ++++++++++++++++++++ 1 file changed, 575 insertions(+) create mode 100644 baseclosure_doublegyre_model_initialized.jl diff --git a/baseclosure_doublegyre_model_initialized.jl b/baseclosure_doublegyre_model_initialized.jl new file mode 100644 index 0000000000..a376834270 --- /dev/null +++ b/baseclosure_doublegyre_model_initialized.jl @@ -0,0 +1,575 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("xin_kai_vertical_diffusivity_local_2step.jl") +# include("xin_kai_vertical_diffusivity_2Pr.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + + +#%% +filename = "doublegyre_relaxation_30days_baseclosure_2step_initialized" +FILE_DIR = "./Output/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +base_closure = XinKaiLocalVerticalDiffusivity() +vertical_scalar_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +closure = (base_closure, vertical_scalar_closure) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/30days + +##### +##### Forcing and initial condition +##### + +# @inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) +@inline T_initial(x, y, z) = (T_north + T_south / 2) + 5 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### +# resting initial condition +# noise(z) = rand() * exp(z / 8) + +# T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +# S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) +DATA_DIR = "./Output/doublegyre_30Cwarmflush_relaxation_8days_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" + +u_data = FieldTimeSeries("$(DATA_DIR)/instantaneous_fields.jld2", "u", backend=OnDisk()) +v_data = FieldTimeSeries("$(DATA_DIR)/instantaneous_fields.jld2", "v", backend=OnDisk()) +T_data = FieldTimeSeries("$(DATA_DIR)/instantaneous_fields.jld2", "T", backend=OnDisk()) +S_data = FieldTimeSeries("$(DATA_DIR)/instantaneous_fields.jld2", "S", backend=OnDisk()) + +ntimes = length(u_data.times) + +set!(model, T=T_data[ntimes], S=S_data[ntimes], u=u_data[ntimes], v=v_data[ntimes]) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10950days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); + +κ = model.diffusivity_fields[1].κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +outputs = (; u, v, w, T, S, ρ, N², wT_base, wS_base) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_south] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_south", + indices = (:, 25, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_north] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_north", + indices = (:, 75, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 365 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% +# function find_min(a...) +# return minimum(minimum.([a...])) +# end + +# function find_max(a...) +# return maximum(maximum.([a...])) +# end + +# N²_xz_north_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_north.jld2", "N²") +# N²_xz_south_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_south.jld2", "N²") + +# xC = N²_xz_north_data.grid.xᶜᵃᵃ[1:N²_xz_north_data.grid.Nx] +# zf = N²_xz_north_data.grid.zᵃᵃᶠ[1:N²_xz_north_data.grid.Nz+1] + +# yloc_north = N²_xz_north_data.grid.yᵃᶠᵃ[N²_xz_north_data.indices[2][1]] +# yloc_south = N²_xz_south_data.grid.yᵃᶠᵃ[N²_xz_south_data.indices[2][1]] + +# Nt = length(N²_xz_north_data) +# times = N²_xz_north_data.times / 24 / 60^2 / 365 +# timeframes = 1:Nt + +# N²_lim = (find_min(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) - 1e-13, +# find_max(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) + 1e-13) +# #%% +# fig = Figure(size=(800, 800)) +# ax_north = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_south / 1e3)) km") +# ax_south = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_north / 1e3)) km") + +# n = Observable(2) + +# N²_north = @lift interior(N²_xz_north_data[$n], :, 1, :) +# N²_south = @lift interior(N²_xz_south_data[$n], :, 1, :) + +# colorscheme = colorschemes[:jet] + +# N²_north_surface = heatmap!(ax_north, xC, zf, N²_north, colormap=colorscheme, colorrange=N²_lim) +# N²_south_surface = heatmap!(ax_south, xC, zf, N²_south, colormap=colorscheme, colorrange=N²_lim) + +# Colorbar(fig[1:2, 2], N²_north_surface, label="N² (s⁻²)") + +# title_str = @lift "Time = $(round(times[$n], digits=2)) years" +# Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +# trim!(fig.layout) + +# @info "Recording buoyancy frequency xz slice" +# CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_buoyancy_frequency_xz_slice.mp4", 1:Nt, framerate=3, px_per_unit=2) do nn +# n[] = nn +# end + +# @info "Done!" +# #%% \ No newline at end of file From c4351f79082ad68bb35bbc4951d6b995b5440013 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Sun, 6 Oct 2024 00:20:31 -0400 Subject: [PATCH 062/114] NN closure for augmenting flux in a zone below MLD --- NN_closure_global_nof_BBLkappazonelast41.jl | 245 ++++++++++++++++++++ 1 file changed, 245 insertions(+) create mode 100644 NN_closure_global_nof_BBLkappazonelast41.jl diff --git a/NN_closure_global_nof_BBLkappazonelast41.jl b/NN_closure_global_nof_BBLkappazonelast41.jl new file mode 100644 index 0000000000..9ee4c3c063 --- /dev/null +++ b/NN_closure_global_nof_BBLkappazonelast41.jl @@ -0,0 +1,245 @@ +import Oceananigans.TurbulenceClosures: + compute_diffusivities!, + DiffusivityFields, + viscosity, + diffusivity, + diffusive_flux_x, + diffusive_flux_y, + diffusive_flux_z, + viscous_flux_ux, + viscous_flux_uy, + viscous_flux_uz, + viscous_flux_vx, + viscous_flux_vy, + viscous_flux_vz, + viscous_flux_wx, + viscous_flux_wy, + viscous_flux_wz + +using Oceananigans.BuoyancyModels: ∂x_b, ∂y_b, ∂z_b +using Oceananigans.Coriolis +using Oceananigans.Grids: φnode +using Oceananigans.Grids: total_size +using Oceananigans.Utils: KernelParameters +using Oceananigans: architecture, on_architecture +using Lux, LuxCUDA +using JLD2 +using ComponentArrays +using OffsetArrays +using SeawaterPolynomials.TEOS10 + +using KernelAbstractions: @index, @kernel, @private + +import Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure, ExplicitTimeDiscretization + +using Adapt + +include("./feature_scaling.jl") + +@inline hack_sind(φ) = sin(φ * π / 180) + +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::HydrostaticSphericalCoriolis) = 2 * coriolis.rotation_rate * hack_sind(φnode(i, j, k, grid, Center(), Center(), Center())) +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::FPlane) = coriolis.f +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::BetaPlane) = coriolis.f₀ + coriolis.β * ynode(i, j, k, grid, Center(), Center(), Center()) + +struct NN{M, P, S} + model :: M + ps :: P + st :: S +end + +@inline (neural_network::NN)(input) = first(neural_network.model(input, neural_network.ps, neural_network.st)) +@inline tosarray(x::AbstractArray) = SArray{Tuple{size(x)...}}(x) + +struct NNFluxClosure{A <: NN, S} <: AbstractTurbulenceClosure{ExplicitTimeDiscretization, 3} + wT :: A + wS :: A + scaling :: S +end + +Adapt.adapt_structure(to, nn :: NNFluxClosure) = + NNFluxClosure(Adapt.adapt(to, nn.wT), + Adapt.adapt(to, nn.wS), + Adapt.adapt(to, nn.scaling)) + +Adapt.adapt_structure(to, nn :: NN) = + NN(Adapt.adapt(to, nn.model), + Adapt.adapt(to, nn.ps), + Adapt.adapt(to, nn.st)) + +function NNFluxClosure(arch) + dev = ifelse(arch == GPU(), gpu_device(), cpu_device()) + nn_path = "./NDE_Qb_dt20min_nof_BBLkappazonelast41_wTwS_64simnew_2layer_128_relu_123seed_1.0e-5lr_localbaseclosure_2Pr_6simstableRi_model_temp.jld2" + + ps, sts, scaling_params, wT_model, wS_model = jldopen(nn_path, "r") do file + ps = file["u"] |> dev |> f64 + sts = file["sts"] |> dev |> f64 + scaling_params = file["scaling"] + wT_model = file["model"].wT + wS_model = file["model"].wS + return ps, sts, scaling_params, wT_model, wS_model + end + + scaling = construct_zeromeanunitvariance_scaling(scaling_params) + + wT_NN = NN(wT_model, ps.wT, sts.wT) + wS_NN = NN(wS_model, ps.wS, sts.wS) + + return NNFluxClosure(wT_NN, wS_NN, scaling) +end + +function DiffusivityFields(grid, tracer_names, bcs, ::NNFluxClosure) + arch = architecture(grid) + wT = ZFaceField(grid) + wS = ZFaceField(grid) + first_index = Field((Center, Center, Nothing), grid, Int32) + last_index = Field((Center, Center, Nothing), grid, Int32) + + Nx_in, Ny_in, _ = size(wT) + + wrk_in = zeros(10, Nx_in, Ny_in, 5) + wrk_in = on_architecture(arch, wrk_in) + + wrk_wT = zeros(Nx_in, Ny_in, 5) + wrk_wS = zeros(Nx_in, Ny_in, 5) + wrk_wT = on_architecture(arch, wrk_wT) + wrk_wS = on_architecture(arch, wrk_wS) + + return (; wrk_in, wrk_wT, wrk_wS, wT, wS, first_index, last_index) +end + +function compute_diffusivities!(diffusivities, closure::NNFluxClosure, model; parameters = :xyz) + arch = model.architecture + grid = model.grid + velocities = model.velocities + tracers = model.tracers + buoyancy = model.buoyancy + coriolis = model.coriolis + clock = model.clock + top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) + + κᶜ = model.diffusivity_fields[1].κᶜ + κ₀ = model.closure[1].ν₀ / model.closure[1].Pr_shearₜ + + wrk_in = diffusivities.wrk_in + wrk_wT = diffusivities.wrk_wT + wrk_wS = diffusivities.wrk_wS + wT = diffusivities.wT + wS = diffusivities.wS + + first_index = diffusivities.first_index + last_index = diffusivities.last_index + + Nx_in, Ny_in, Nz_in = total_size(wT) + ox_in, oy_in, oz_in = wT.data.offsets + kp = KernelParameters((Nx_in, Ny_in, Nz_in), (ox_in, oy_in, oz_in)) + kp_2D = KernelParameters((Nx_in, Ny_in), (ox_in, oy_in)) + + kp_wrk = KernelParameters((Nx_in, Ny_in, 5), (0, 0, 0)) + + launch!(arch, grid, kp_2D, _find_NN_active_region!, κᶜ, grid, κ₀, first_index, last_index) + + launch!(arch, grid, kp_wrk, + _populate_input!, wrk_in, first_index, last_index, grid, closure, tracers, velocities, buoyancy, top_tracer_bcs, clock) + + wrk_wT .= dropdims(closure.wT(wrk_in), dims=1) + wrk_wS .= dropdims(closure.wS(wrk_in), dims=1) + + launch!(arch, grid, kp, _fill_adjust_nn_fluxes!, diffusivities, first_index, last_index, wrk_wT, wrk_wS, grid, closure, tracers, velocities, buoyancy, top_tracer_bcs, clock) + return nothing +end + +@kernel function _populate_input!(input, first_index, last_index, grid, closure::NNFluxClosure, tracers, velocities, buoyancy, top_tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + + scaling = closure.scaling + + ρ₀ = buoyancy.model.equation_of_state.reference_density + g = buoyancy.model.gravitational_acceleration + eos = TEOS10.TEOS10EquationOfState() + + @inbounds quiescent = quiescent_condition(first_index[i, j, 1], last_index[i, j, 1]) + @inbounds k_tracer = first_index[i, j, 1] + k - 1 + + T, S = tracers.T, tracers.S + + @inbounds input[1, i, j, k] = ∂Tᵢ₋₁ = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer-1, grid, T))) + @inbounds input[2, i, j, k] = ∂Tᵢ = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer, grid, T))) + @inbounds input[3, i, j, k] = ∂Tᵢ₊₁ = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer+1, grid, T))) + + @inbounds input[4, i, j, k] = ∂Sᵢ₋₁ = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer-1, grid, S))) + @inbounds input[5, i, j, k] = ∂Sᵢ = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer, grid, S))) + @inbounds input[6, i, j, k] = ∂Sᵢ₊₁ = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer+1, grid, S))) + + @inbounds input[7, i, j, k] = ∂σᵢ₋₁ = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k_tracer-1, grid, buoyancy, tracers) / g)) + @inbounds input[8, i, j, k] = ∂σᵢ = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k_tracer, grid, buoyancy, tracers) / g)) + @inbounds input[9, i, j, k] = ∂σᵢ₊₁ = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k_tracer+1, grid, buoyancy, tracers) / g)) + + @inbounds input[10, i, j, k] = Jᵇ = scaling.wb(ifelse(quiescent, 0, top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers)))) + +end + +@kernel function _find_NN_active_region!(κᶜ, grid, κ₀, first_index, last_index) + i, j = @index(Global, NTuple) + top_index = grid.Nz + 1 + + # Find the last index of the background κᶜ + kmax = 1 + @inbounds for k in 2:grid.Nz + kmax = ifelse(κᶜ[i, j, k] ≈ κ₀, k, kmax) + end + + @inbounds last_index[i, j, 1] = ifelse(kmax == top_index, grid.Nz, min(kmax + 1, grid.Nz)) + @inbounds first_index[i, j, 1] = ifelse(kmax == top_index, top_index, max(kmax - 3, 2)) +end + +@inline function quiescent_condition(lo, hi) + return hi - lo != 4 +end + +@kernel function _fill_adjust_nn_fluxes!(diffusivities, first_index, last_index, wrk_wT, wrk_wS, grid, closure::NNFluxClosure, tracers, velocities, buoyancy, top_tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + scaling = closure.scaling + + k_first = first_index[i, j, 1] + k_last = last_index[i, j, 1] + + convecting = top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers)) > 0 + @inbounds quiescent = quiescent_condition(k_first, k_last) + within_zone = (k >= k_first) & (k <= k_last) + + @inbounds k_wrk = clamp(k - k_first + 1, 1, 5) + + NN_active = convecting & !quiescent & within_zone + + @inbounds diffusivities.wT[i, j, k] = ifelse(NN_active, scaling.wT.σ * wrk_wT[i, j, k_wrk], 0) + @inbounds diffusivities.wS[i, j, k] = ifelse(NN_active, scaling.wS.σ * wrk_wS[i, j, k_wrk], 0) +end + +# Write here your constructor +# NNFluxClosure() = ... insert NN here ... (make sure it is on GPU if you need it on GPU!) + +const NNC = NNFluxClosure + +##### +##### Abstract Smagorinsky functionality +##### + +# Horizontal fluxes are zero! +@inline viscous_flux_wz( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_ux( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_uy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline diffusive_flux_x(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) +@inline diffusive_flux_y(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) + +# Viscous fluxes are zero (for now) +@inline viscous_flux_uz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) + +# The only function extended by NNFluxClosure +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{1}, c, clock, fields, buoyancy) = @inbounds K.wT[i, j, k] +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{2}, c, clock, fields, buoyancy) = @inbounds K.wS[i, j, k] \ No newline at end of file From 19f71d3e11322c98772b71ec68c46df4af0f10fc Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Mon, 7 Oct 2024 22:53:43 -0400 Subject: [PATCH 063/114] 8 day relaxation double gyre for baseclosure --- baseclosure_doublegyre_model.jl | 571 ++++++++++++++++++++++++++++++++ 1 file changed, 571 insertions(+) create mode 100644 baseclosure_doublegyre_model.jl diff --git a/baseclosure_doublegyre_model.jl b/baseclosure_doublegyre_model.jl new file mode 100644 index 0000000000..b2e5ab75d5 --- /dev/null +++ b/baseclosure_doublegyre_model.jl @@ -0,0 +1,571 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("xin_kai_vertical_diffusivity_local_2step.jl") +# include("xin_kai_vertical_diffusivity_2Pr.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +# using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + +#%% +filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" +FILE_DIR = "./Output/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +base_closure = XinKaiLocalVerticalDiffusivity() +vertical_scalar_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +closure = (base_closure, vertical_scalar_closure) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/8days + +##### +##### Forcing and initial condition +##### +# @inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) +# @inline T_initial(x, y, z) = (T_north + T_south / 2) + 5 * (1 + z / Lz) +@inline T_initial(x, y, z) = 30 + 20 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10950days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); + +κ = model.diffusivity_fields[1].κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +outputs = (; u, v, w, T, S, ρ, N², wT_base, wS_base) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_south] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_south", + indices = (:, 25, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_north] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_north", + indices = (:, 75, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:complete_fields] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields", + schedule = TimeInterval(1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 365 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% +# function find_min(a...) +# return minimum(minimum.([a...])) +# end + +# function find_max(a...) +# return maximum(maximum.([a...])) +# end + +# N²_xz_north_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_north.jld2", "N²") +# N²_xz_south_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_south.jld2", "N²") + +# xC = N²_xz_north_data.grid.xᶜᵃᵃ[1:N²_xz_north_data.grid.Nx] +# zf = N²_xz_north_data.grid.zᵃᵃᶠ[1:N²_xz_north_data.grid.Nz+1] + +# yloc_north = N²_xz_north_data.grid.yᵃᶠᵃ[N²_xz_north_data.indices[2][1]] +# yloc_south = N²_xz_south_data.grid.yᵃᶠᵃ[N²_xz_south_data.indices[2][1]] + +# Nt = length(N²_xz_north_data) +# times = N²_xz_north_data.times / 24 / 60^2 / 365 +# timeframes = 1:Nt + +# N²_lim = (find_min(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) - 1e-13, +# find_max(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) + 1e-13) +# #%% +# fig = Figure(size=(800, 800)) +# ax_north = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_south / 1e3)) km") +# ax_south = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_north / 1e3)) km") + +# n = Observable(2) + +# N²_north = @lift interior(N²_xz_north_data[$n], :, 1, :) +# N²_south = @lift interior(N²_xz_south_data[$n], :, 1, :) + +# colorscheme = colorschemes[:jet] + +# N²_north_surface = heatmap!(ax_north, xC, zf, N²_north, colormap=colorscheme, colorrange=N²_lim) +# N²_south_surface = heatmap!(ax_south, xC, zf, N²_south, colormap=colorscheme, colorrange=N²_lim) + +# Colorbar(fig[1:2, 2], N²_north_surface, label="N² (s⁻²)") + +# title_str = @lift "Time = $(round(times[$n], digits=2)) years" +# Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +# trim!(fig.layout) + +# @info "Recording buoyancy frequency xz slice" +# CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_buoyancy_frequency_xz_slice.mp4", 1:Nt, framerate=3, px_per_unit=2) do nn +# n[] = nn +# end + +# @info "Done!" +# #%% \ No newline at end of file From 3c4102a319d040c3268f1a5a3e37b1eaa19e049e Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Mon, 7 Oct 2024 22:58:30 -0400 Subject: [PATCH 064/114] uncomment CairoMakie --- baseclosure_doublegyre_model.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/baseclosure_doublegyre_model.jl b/baseclosure_doublegyre_model.jl index b2e5ab75d5..92a1c93832 100644 --- a/baseclosure_doublegyre_model.jl +++ b/baseclosure_doublegyre_model.jl @@ -10,7 +10,7 @@ pushfirst!(LOAD_PATH, @__DIR__) using Printf using Statistics -# using CairoMakie +using CairoMakie using Oceananigans using Oceananigans.Units From b122e7b6787b6d53123eb9dfbe9064afaaa09397 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Mon, 7 Oct 2024 23:04:54 -0400 Subject: [PATCH 065/114] NN closure and CATKE with 8 day restoration and warm flush --- ...nof_BBLkappazonelast41_doublegyre_model.jl | 583 ++++++++++++++++++ physicalclosure_doublegyre_model.jl | 8 +- 2 files changed, 587 insertions(+), 4 deletions(-) create mode 100644 NNclosure_nof_BBLkappazonelast41_doublegyre_model.jl diff --git a/NNclosure_nof_BBLkappazonelast41_doublegyre_model.jl b/NNclosure_nof_BBLkappazonelast41_doublegyre_model.jl new file mode 100644 index 0000000000..701c62fce1 --- /dev/null +++ b/NNclosure_nof_BBLkappazonelast41_doublegyre_model.jl @@ -0,0 +1,583 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("NN_closure_global_nof_BBLkappazonelast41.jl") +include("xin_kai_vertical_diffusivity_local_2step.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + + +#%% +filename = "doublegyre_relaxation_30Cwarmflush_8days_NN_closure_NDE_Qb_dt20min_nof_BBLkappazonelast41_wTwS_64simnew_2layer_128_relu_123seed_1.0e-5lr_localbaseclosure_2Pr_6simstableRi_temp" +FILE_DIR = "./Output/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +nn_closure = NNFluxClosure(model_architecture) +base_closure = XinKaiLocalVerticalDiffusivity() +vertical_scalar_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +closure = (base_closure, nn_closure, vertical_scalar_closure) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/8days + +##### +##### Forcing and initial condition +##### +@inline T_initial(x, y, z) = 30 + 20 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10950days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); +first_index = model.diffusivity_fields[2].first_index +last_index = model.diffusivity_fields[2].last_index +wT_NN = model.diffusivity_fields[2].wT +wS_NN = model.diffusivity_fields[2].wS + +κ = model.diffusivity_fields[1].κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +@inline function get_top_buoyancy_flux(i, j, k, grid, buoyancy, T_bc, S_bc, velocities, tracers, clock) + return top_buoyancy_flux(i, j, grid, buoyancy, (; T=T_bc, S=S_bc), clock, merge(velocities, tracers)) +end + +Qb = KernelFunctionOperation{Center, Center, Nothing}(get_top_buoyancy_flux, model.grid, model.buoyancy, T.boundary_conditions.top, S.boundary_conditions.top, model.velocities, model.tracers, model.clock) +Qb = Field(Qb) + +outputs = (; u, v, w, T, S, ρ, N², wT_NN, wS_NN, wT_base, wS_base) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_south] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_south", + indices = (:, 25, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_north] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_north", + indices = (:, 75, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; first_index, last_index, Qb), + filename = "$(FILE_DIR)/instantaneous_fields_NN_active_diagnostics", + indices = (:, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 365 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +N²_xz_north_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_north.jld2", "N²") +N²_xz_south_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_south.jld2", "N²") + +xC = N²_xz_north_data.grid.xᶜᵃᵃ[1:N²_xz_north_data.grid.Nx] +zf = N²_xz_north_data.grid.zᵃᵃᶠ[1:N²_xz_north_data.grid.Nz+1] + +yloc_north = N²_xz_north_data.grid.yᵃᶠᵃ[N²_xz_north_data.indices[2][1]] +yloc_south = N²_xz_south_data.grid.yᵃᶠᵃ[N²_xz_south_data.indices[2][1]] + +Nt = length(N²_xz_north_data) +times = N²_xz_north_data.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +N²_lim = (find_min(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) - 1e-13, + find_max(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) + 1e-13) +#%% +fig = Figure(size=(800, 800)) +ax_north = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_south / 1e3)) km") +ax_south = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_north / 1e3)) km") + +n = Observable(2) + +N²_north = @lift interior(N²_xz_north_data[$n], :, 1, :) +N²_south = @lift interior(N²_xz_south_data[$n], :, 1, :) + +colorscheme = colorschemes[:jet] + +N²_north_surface = heatmap!(ax_north, xC, zf, N²_north, colormap=colorscheme, colorrange=N²_lim) +N²_south_surface = heatmap!(ax_south, xC, zf, N²_south, colormap=colorscheme, colorrange=N²_lim) + +Colorbar(fig[1:2, 2], N²_north_surface, label="N² (s⁻²)") + +title_str = @lift "Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +@info "Recording buoyancy frequency xz slice" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_buoyancy_frequency_xz_slice.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end + +@info "Done!" +#%% \ No newline at end of file diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index c037e47cdd..44d257a330 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -28,7 +28,7 @@ using Glob #%% -filename = "doublegyre_relaxation_30days_CATKEVerticalDiffusivity" +filename = "doublegyre_relaxation_30Cwarmflush_8days_CATKEVerticalDiffusivity" FILE_DIR = "./Output/$(filename)" mkpath(FILE_DIR) @@ -81,13 +81,13 @@ const S_mid = (S_north + S_south) / 2 const τ₀ = 1e-4 const μ_drag = 1/30days -const μ_T = 1/30days +const μ_T = 1/8days ##### ##### Forcing and initial condition ##### - -@inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) +# @inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) +@inline T_initial(x, y, z) = 30 + 20 * (1 + z / Lz) @inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) From e0bd65a246b1abe2f056afd3fca0763e8ad21fc1 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 8 Oct 2024 01:08:48 -0400 Subject: [PATCH 066/114] fix initial temperature issue --- NNclosure_nof_BBLkappazonelast41_doublegyre_model.jl | 7 ++++--- baseclosure_doublegyre_model.jl | 2 +- physicalclosure_doublegyre_model.jl | 5 +++-- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/NNclosure_nof_BBLkappazonelast41_doublegyre_model.jl b/NNclosure_nof_BBLkappazonelast41_doublegyre_model.jl index 701c62fce1..a8181f78ee 100644 --- a/NNclosure_nof_BBLkappazonelast41_doublegyre_model.jl +++ b/NNclosure_nof_BBLkappazonelast41_doublegyre_model.jl @@ -22,8 +22,9 @@ using ColorSchemes #%% -filename = "doublegyre_relaxation_30Cwarmflush_8days_NN_closure_NDE_Qb_dt20min_nof_BBLkappazonelast41_wTwS_64simnew_2layer_128_relu_123seed_1.0e-5lr_localbaseclosure_2Pr_6simstableRi_temp" -FILE_DIR = "./Output/$(filename)" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_NN_closure_NDE_Qb_dt20min_nof_BBLkappazonelast41_wTwS_64simnew_2layer_128_relu_123seed_1.0e-5lr_localbaseclosure_2Pr_6simstableRi_temp" +# FILE_DIR = "./Output/$(filename)" +FILE_DIR = "~/storage6/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) # Architecture @@ -74,7 +75,7 @@ const μ_T = 1/8days ##### ##### Forcing and initial condition ##### -@inline T_initial(x, y, z) = 30 + 20 * (1 + z / Lz) +@inline T_initial(x, y, z) = 10 + 20 * (1 + z / Lz) @inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) diff --git a/baseclosure_doublegyre_model.jl b/baseclosure_doublegyre_model.jl index 92a1c93832..1e63665e04 100644 --- a/baseclosure_doublegyre_model.jl +++ b/baseclosure_doublegyre_model.jl @@ -74,7 +74,7 @@ const μ_T = 1/8days ##### # @inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) # @inline T_initial(x, y, z) = (T_north + T_south / 2) + 5 * (1 + z / Lz) -@inline T_initial(x, y, z) = 30 + 20 * (1 + z / Lz) +@inline T_initial(x, y, z) = 10 + 20 * (1 + z / Lz) @inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index 44d257a330..144846481a 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -28,8 +28,9 @@ using Glob #%% -filename = "doublegyre_relaxation_30Cwarmflush_8days_CATKEVerticalDiffusivity" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_CATKEVerticalDiffusivity" FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) # Architecture @@ -87,7 +88,7 @@ const μ_T = 1/8days ##### Forcing and initial condition ##### # @inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) -@inline T_initial(x, y, z) = 30 + 20 * (1 + z / Lz) +@inline T_initial(x, y, z) = 10 + 20 * (1 + z / Lz) @inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) From 8dcb6d37ee1a805dc7f9ba0186583309f63e6783 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Wed, 9 Oct 2024 16:21:32 -0400 Subject: [PATCH 067/114] add zonal average calculations to double gyre simulation --- ...nof_BBLkappazonelast41_doublegyre_model.jl | 78 +++--------- baseclosure_doublegyre_model.jl | 21 ++-- physicalclosure_doublegyre_model.jl | 112 +++++------------- 3 files changed, 57 insertions(+), 154 deletions(-) diff --git a/NNclosure_nof_BBLkappazonelast41_doublegyre_model.jl b/NNclosure_nof_BBLkappazonelast41_doublegyre_model.jl index a8181f78ee..7658778581 100644 --- a/NNclosure_nof_BBLkappazonelast41_doublegyre_model.jl +++ b/NNclosure_nof_BBLkappazonelast41_doublegyre_model.jl @@ -22,9 +22,9 @@ using ColorSchemes #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_NN_closure_NDE_Qb_dt20min_nof_BBLkappazonelast41_wTwS_64simnew_2layer_128_relu_123seed_1.0e-5lr_localbaseclosure_2Pr_6simstableRi_temp" -# FILE_DIR = "./Output/$(filename)" -FILE_DIR = "~/storage6/NN_Oceananigans/$(filename)" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_NN_closure_NDE_BBLkappazonelast41_temp" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) # Architecture @@ -235,7 +235,15 @@ end Qb = KernelFunctionOperation{Center, Center, Nothing}(get_top_buoyancy_flux, model.grid, model.buoyancy, T.boundary_conditions.top, S.boundary_conditions.top, model.velocities, model.tracers, model.clock) Qb = Field(Qb) +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + outputs = (; u, v, w, T, S, ρ, N², wT_NN, wS_NN, wT_base, wS_base) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal) ##### ##### Build checkpointer and output writer @@ -300,15 +308,9 @@ simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, indices = (90, :, :), schedule = TimeInterval(10days)) -simulation.output_writers[:xz_south] = JLD2OutputWriter(model, outputs, - filename = "$(FILE_DIR)/instantaneous_fields_xz_south", - indices = (:, 25, :), - schedule = TimeInterval(10days)) - -simulation.output_writers[:xz_north] = JLD2OutputWriter(model, outputs, - filename = "$(FILE_DIR)/instantaneous_fields_xz_north", - indices = (:, 75, :), - schedule = TimeInterval(10days)) +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; first_index, last_index, Qb), filename = "$(FILE_DIR)/instantaneous_fields_NN_active_diagnostics", @@ -529,56 +531,4 @@ end # tightlimits!(ax) # save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) # display(fig) -#%% -function find_min(a...) - return minimum(minimum.([a...])) -end - -function find_max(a...) - return maximum(maximum.([a...])) -end - -N²_xz_north_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_north.jld2", "N²") -N²_xz_south_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_south.jld2", "N²") - -xC = N²_xz_north_data.grid.xᶜᵃᵃ[1:N²_xz_north_data.grid.Nx] -zf = N²_xz_north_data.grid.zᵃᵃᶠ[1:N²_xz_north_data.grid.Nz+1] - -yloc_north = N²_xz_north_data.grid.yᵃᶠᵃ[N²_xz_north_data.indices[2][1]] -yloc_south = N²_xz_south_data.grid.yᵃᶠᵃ[N²_xz_south_data.indices[2][1]] - -Nt = length(N²_xz_north_data) -times = N²_xz_north_data.times / 24 / 60^2 / 365 -timeframes = 1:Nt - -N²_lim = (find_min(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) - 1e-13, - find_max(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) + 1e-13) -#%% -fig = Figure(size=(800, 800)) -ax_north = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_south / 1e3)) km") -ax_south = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_north / 1e3)) km") - -n = Observable(2) - -N²_north = @lift interior(N²_xz_north_data[$n], :, 1, :) -N²_south = @lift interior(N²_xz_south_data[$n], :, 1, :) - -colorscheme = colorschemes[:jet] - -N²_north_surface = heatmap!(ax_north, xC, zf, N²_north, colormap=colorscheme, colorrange=N²_lim) -N²_south_surface = heatmap!(ax_south, xC, zf, N²_south, colormap=colorscheme, colorrange=N²_lim) - -Colorbar(fig[1:2, 2], N²_north_surface, label="N² (s⁻²)") - -title_str = @lift "Time = $(round(times[$n], digits=2)) years" -Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) - -trim!(fig.layout) - -@info "Recording buoyancy frequency xz slice" -CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_buoyancy_frequency_xz_slice.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn - n[] = nn -end - -@info "Done!" #%% \ No newline at end of file diff --git a/baseclosure_doublegyre_model.jl b/baseclosure_doublegyre_model.jl index 1e63665e04..b66cc37c4e 100644 --- a/baseclosure_doublegyre_model.jl +++ b/baseclosure_doublegyre_model.jl @@ -23,6 +23,7 @@ using ColorSchemes #%% filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) # Architecture @@ -223,7 +224,15 @@ end ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) ρ = Field(ρ_op) +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + outputs = (; u, v, w, T, S, ρ, N², wT_base, wS_base) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal) ##### ##### Build checkpointer and output writer @@ -288,15 +297,9 @@ simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, indices = (90, :, :), schedule = TimeInterval(10days)) -simulation.output_writers[:xz_south] = JLD2OutputWriter(model, outputs, - filename = "$(FILE_DIR)/instantaneous_fields_xz_south", - indices = (:, 25, :), - schedule = TimeInterval(10days)) - -simulation.output_writers[:xz_north] = JLD2OutputWriter(model, outputs, - filename = "$(FILE_DIR)/instantaneous_fields_xz_north", - indices = (:, 75, :), - schedule = TimeInterval(10days)) +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), filename = "$(FILE_DIR)/averaged_fields_streamfunction", diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index 144846481a..c2738938a4 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -242,7 +242,15 @@ compute!(ρ) wT = κc * ∂z(T) wS = κc * ∂z(S) +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + outputs = (; u, v, w, T, S, ρ, ρᶠ, ∂ρ∂z, ∂²ρ∂z², N², wT, wS) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal) ##### ##### Build checkpointer and output writer @@ -307,15 +315,9 @@ simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, indices = (90, :, :), schedule = TimeInterval(10days)) -simulation.output_writers[:xz_south] = JLD2OutputWriter(model, outputs, - filename = "$(FILE_DIR)/instantaneous_fields_xz_south", - indices = (:, 25, :), - schedule = TimeInterval(10days)) - -simulation.output_writers[:xz_north] = JLD2OutputWriter(model, outputs, - filename = "$(FILE_DIR)/instantaneous_fields_xz_north", - indices = (:, 75, :), - schedule = TimeInterval(10days)) +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), filename = "$(FILE_DIR)/averaged_fields_streamfunction", @@ -510,76 +512,24 @@ end @info "Done!" #%% -Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") - -xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] -yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] - -Nt = length(Ψ_data) -times = Ψ_data.times / 24 / 60^2 / 365 -#%% -timeframe = Nt -Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 -clim = maximum(abs, Ψ_frame) + 1e-13 -N_levels = 16 -levels = range(-clim, stop=clim, length=N_levels) -fig = Figure(size=(800, 800)) -ax = Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") -cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) -Colorbar(fig[1, 2], cf, label="Ψ (Sv)") -tightlimits!(ax) -save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) -display(fig) -#%% -function find_min(a...) - return minimum(minimum.([a...])) -end - -function find_max(a...) - return maximum(maximum.([a...])) -end - -N²_xz_north_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_north.jld2", "N²") -N²_xz_south_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_south.jld2", "N²") - -xC = N²_xz_north_data.grid.xᶜᵃᵃ[1:N²_xz_north_data.grid.Nx] -zf = N²_xz_north_data.grid.zᵃᵃᶠ[1:N²_xz_north_data.grid.Nz+1] - -yloc_north = N²_xz_north_data.grid.yᵃᶠᵃ[N²_xz_north_data.indices[2][1]] -yloc_south = N²_xz_south_data.grid.yᵃᶠᵃ[N²_xz_south_data.indices[2][1]] - -Nt = length(N²_xz_north_data) -times = N²_xz_north_data.times / 24 / 60^2 / 365 -timeframes = 1:Nt - -N²_lim = (find_min(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) - 1e-13, - find_max(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) + 1e-13) -#%% -fig = Figure(size=(800, 800)) -ax_north = Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_south / 1e3)) km") -ax_south = Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_north / 1e3)) km") - -n = Observable(2) - -N²_north = @lift interior(N²_xz_north_data[$n], :, 1, :) -N²_south = @lift interior(N²_xz_south_data[$n], :, 1, :) - -colorscheme = colorschemes[:jet] - -N²_north_surface = heatmap!(ax_north, xC, zf, N²_north, colormap=colorscheme, colorrange=N²_lim) -N²_south_surface = heatmap!(ax_south, xC, zf, N²_south, colormap=colorscheme, colorrange=N²_lim) - -Colorbar(fig[1:2, 2], N²_north_surface, label="N² (s⁻²)") - -title_str = @lift "Time = $(round(times[$n], digits=2)) days" -Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) - -trim!(fig.layout) - -@info "Recording buoyancy frequency xz slice" -CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_buoyancy_frequency_xz_slice.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn - n[] = nn -end - -@info "Done!" +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 365 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) #%% \ No newline at end of file From 132188a36c7a545bb3a4bed33aad35e061134347 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 10 Oct 2024 01:41:05 -0400 Subject: [PATCH 068/114] run double gyre with seasonal forcing --- ...last41_doublegyre_model_seasonalforcing.jl | 539 ++++++++++++++++ ...losure_doublegyre_model_seasonalforcing.jl | 577 ++++++++++++++++++ ...losure_doublegyre_model_seasonalforcing.jl | 539 ++++++++++++++++ 3 files changed, 1655 insertions(+) create mode 100644 NNclosure_nof_BBLkappazonelast41_doublegyre_model_seasonalforcing.jl create mode 100644 baseclosure_doublegyre_model_seasonalforcing.jl create mode 100644 physicalclosure_doublegyre_model_seasonalforcing.jl diff --git a/NNclosure_nof_BBLkappazonelast41_doublegyre_model_seasonalforcing.jl b/NNclosure_nof_BBLkappazonelast41_doublegyre_model_seasonalforcing.jl new file mode 100644 index 0000000000..c100a5a516 --- /dev/null +++ b/NNclosure_nof_BBLkappazonelast41_doublegyre_model_seasonalforcing.jl @@ -0,0 +1,539 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("NN_closure_global_nof_BBLkappazonelast41.jl") +include("xin_kai_vertical_diffusivity_local_2step.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + + +#%% +filename = "doublegyre_seasonalforcing_30C-20C_relaxation_8days_NN_closure_NDE_BBLkappazonelast41_temp" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +nn_closure = NNFluxClosure(model_architecture) +base_closure = XinKaiLocalVerticalDiffusivity() +vertical_scalar_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +closure = (base_closure, nn_closure, vertical_scalar_closure) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/8days + +const seasonal_period = 360days +const seasonal_forcing_width = Ly / 6 +const seasonal_T_amplitude = 15 + +##### +##### Forcing and initial condition +##### +@inline T_initial(x, y, z) = 20 + 10 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_seasonal(y, t) = seasonal_T_amplitude * exp(-y^2/(2 * seasonal_forcing_width^2)) * sin(2π * t / seasonal_period) +@inline T_ref(y, t) = T_mid - ΔT / Ly * y + T_seasonal(y, t) +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y, t)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10800days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); +first_index = model.diffusivity_fields[2].first_index +last_index = model.diffusivity_fields[2].last_index +wT_NN = model.diffusivity_fields[2].wT +wS_NN = model.diffusivity_fields[2].wS + +κ = model.diffusivity_fields[1].κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +@inline function get_top_buoyancy_flux(i, j, k, grid, buoyancy, T_bc, S_bc, velocities, tracers, clock) + return top_buoyancy_flux(i, j, grid, buoyancy, (; T=T_bc, S=S_bc), clock, merge(velocities, tracers)) +end + +Qb = KernelFunctionOperation{Center, Center, Nothing}(get_top_buoyancy_flux, model.grid, model.buoyancy, T.boundary_conditions.top, S.boundary_conditions.top, model.velocities, model.tracers, model.clock) +Qb = Field(Qb) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + +outputs = (; u, v, w, T, S, ρ, N², wT_NN, wS_NN, wT_base, wS_base) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; first_index, last_index, Qb), + filename = "$(FILE_DIR)/instantaneous_fields_NN_active_diagnostics", + indices = (:, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 360 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% \ No newline at end of file diff --git a/baseclosure_doublegyre_model_seasonalforcing.jl b/baseclosure_doublegyre_model_seasonalforcing.jl new file mode 100644 index 0000000000..9655817341 --- /dev/null +++ b/baseclosure_doublegyre_model_seasonalforcing.jl @@ -0,0 +1,577 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("xin_kai_vertical_diffusivity_local_2step.jl") +# include("xin_kai_vertical_diffusivity_2Pr.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + +#%% +filename = "doublegyre_seasonalforcing_30C-20C_relaxation_8days_baseclosure_test" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +base_closure = XinKaiLocalVerticalDiffusivity() +vertical_scalar_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +closure = (base_closure, vertical_scalar_closure) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/8days + +const seasonal_period = 360days +const seasonal_forcing_width = Ly / 6 +const seasonal_T_amplitude = 15 + +##### +##### Forcing and initial condition +##### +@inline T_initial(x, y, z) = 20 + 10 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_seasonal(y, t) = seasonal_T_amplitude * exp(-y^2/(2 * seasonal_forcing_width^2)) * sin(2π * t / seasonal_period) +@inline T_ref(y, t) = T_mid - ΔT / Ly * y + T_seasonal(y, t) +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y, t)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10800days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); + +κ = model.diffusivity_fields[1].κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + +outputs = (; u, v, w, T, S, ρ, N², wT_base, wS_base) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:complete_fields] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields", + schedule = TimeInterval(1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 360 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% +# function find_min(a...) +# return minimum(minimum.([a...])) +# end + +# function find_max(a...) +# return maximum(maximum.([a...])) +# end + +# N²_xz_north_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_north.jld2", "N²") +# N²_xz_south_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_south.jld2", "N²") + +# xC = N²_xz_north_data.grid.xᶜᵃᵃ[1:N²_xz_north_data.grid.Nx] +# zf = N²_xz_north_data.grid.zᵃᵃᶠ[1:N²_xz_north_data.grid.Nz+1] + +# yloc_north = N²_xz_north_data.grid.yᵃᶠᵃ[N²_xz_north_data.indices[2][1]] +# yloc_south = N²_xz_south_data.grid.yᵃᶠᵃ[N²_xz_south_data.indices[2][1]] + +# Nt = length(N²_xz_north_data) +# times = N²_xz_north_data.times / 24 / 60^2 / 360 +# timeframes = 1:Nt + +# N²_lim = (find_min(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) - 1e-13, +# find_max(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) + 1e-13) +# #%% +# fig = Figure(size=(800, 800)) +# ax_north = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_south / 1e3)) km") +# ax_south = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_north / 1e3)) km") + +# n = Observable(2) + +# N²_north = @lift interior(N²_xz_north_data[$n], :, 1, :) +# N²_south = @lift interior(N²_xz_south_data[$n], :, 1, :) + +# colorscheme = colorschemes[:jet] + +# N²_north_surface = heatmap!(ax_north, xC, zf, N²_north, colormap=colorscheme, colorrange=N²_lim) +# N²_south_surface = heatmap!(ax_south, xC, zf, N²_south, colormap=colorscheme, colorrange=N²_lim) + +# Colorbar(fig[1:2, 2], N²_north_surface, label="N² (s⁻²)") + +# title_str = @lift "Time = $(round(times[$n], digits=2)) years" +# Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +# trim!(fig.layout) + +# @info "Recording buoyancy frequency xz slice" +# CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_buoyancy_frequency_xz_slice.mp4", 1:Nt, framerate=3, px_per_unit=2) do nn +# n[] = nn +# end + +# @info "Done!" +# #%% \ No newline at end of file diff --git a/physicalclosure_doublegyre_model_seasonalforcing.jl b/physicalclosure_doublegyre_model_seasonalforcing.jl new file mode 100644 index 0000000000..b2e63e60d9 --- /dev/null +++ b/physicalclosure_doublegyre_model_seasonalforcing.jl @@ -0,0 +1,539 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity +using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: CATKEVerticalDiffusivity, CATKEMixingLength, CATKEEquation + +using Oceananigans.BuoyancyModels: ∂z_b +# include("NN_closure_global.jl") +# include("xin_kai_vertical_diffusivity_local.jl") +# include("xin_kai_vertical_diffusivity_2Pr.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes +using Glob + + +#%% +filename = "doublegyre_seasonalforcing_30C-20C_relaxation_8days_CATKEVerticalDiffusivity" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +# vertical_base_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +# convection_closure = XinKaiVerticalDiffusivity() +function CATKE_ocean_closure() + mixing_length = CATKEMixingLength(Cᵇ=0.01) + turbulent_kinetic_energy_equation = CATKEEquation(Cᵂϵ=1.0) + return CATKEVerticalDiffusivity(; mixing_length, turbulent_kinetic_energy_equation) +end +convection_closure = CATKE_ocean_closure() +closure = convection_closure +# closure = vertical_base_closure + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/8days + +const seasonal_period = 360days +const seasonal_forcing_width = Ly / 6 +const seasonal_T_amplitude = 15 + +##### +##### Forcing and initial condition +##### +@inline T_initial(x, y, z) = 20 + 10 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_seasonal(y, t) = seasonal_T_amplitude * exp(-y^2/(2 * seasonal_forcing_width^2)) * sin(2π * t / seasonal_period) +@inline T_ref(y, t) = T_mid - ΔT / Ly * y + T_seasonal(y, t) +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y, t)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S, :e), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S, :e), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10800days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)) +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) +compute!(ρ) + +ρᶠ = @at (Center, Center, Face) ρ +∂ρ∂z = ∂z(ρ) +∂²ρ∂z² = ∂z(∂ρ∂z) + +κc = model.diffusivity_fields.κc +wT = κc * ∂z(T) +wS = κc * ∂z(S) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + +outputs = (; u, v, w, T, S, ρ, ρᶠ, ∂ρ∂z, ∂²ρ∂z², N², wT, wS) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 360 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% \ No newline at end of file From 57c3ac55f7b77e4965c3e3ca52acf4db2340b283 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Wed, 16 Oct 2024 16:48:19 -0400 Subject: [PATCH 069/114] increase simulation run time --- baseclosure_doublegyre_model_seasonalforcing.jl | 2 +- physicalclosure_doublegyre_model_seasonalforcing.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/baseclosure_doublegyre_model_seasonalforcing.jl b/baseclosure_doublegyre_model_seasonalforcing.jl index 9655817341..5ffeb56cf6 100644 --- a/baseclosure_doublegyre_model_seasonalforcing.jl +++ b/baseclosure_doublegyre_model_seasonalforcing.jl @@ -169,7 +169,7 @@ update_state!(model) ##### Simulation building ##### Δt₀ = 5minutes -stop_time = 10800days +stop_time = 36000days simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) diff --git a/physicalclosure_doublegyre_model_seasonalforcing.jl b/physicalclosure_doublegyre_model_seasonalforcing.jl index b2e63e60d9..4c692e8cb9 100644 --- a/physicalclosure_doublegyre_model_seasonalforcing.jl +++ b/physicalclosure_doublegyre_model_seasonalforcing.jl @@ -183,7 +183,7 @@ update_state!(model) ##### Simulation building ##### Δt₀ = 5minutes -stop_time = 10800days +stop_time = 36000days simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) From d5a320e1d31915ae38479c80b249b9f5a13162b0 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 17 Oct 2024 15:00:42 -0400 Subject: [PATCH 070/114] wall restoration to maintain strratification --- ...e_model_seasonalforcing_wallrestoration.jl | 546 ++++++++++++++++++ 1 file changed, 546 insertions(+) create mode 100644 physicalclosure_doublegyre_model_seasonalforcing_wallrestoration.jl diff --git a/physicalclosure_doublegyre_model_seasonalforcing_wallrestoration.jl b/physicalclosure_doublegyre_model_seasonalforcing_wallrestoration.jl new file mode 100644 index 0000000000..069d71fc7f --- /dev/null +++ b/physicalclosure_doublegyre_model_seasonalforcing_wallrestoration.jl @@ -0,0 +1,546 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity +using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: CATKEVerticalDiffusivity, CATKEMixingLength, CATKEEquation + +using Oceananigans.BuoyancyModels: ∂z_b +# include("NN_closure_global.jl") +# include("xin_kai_vertical_diffusivity_local.jl") +# include("xin_kai_vertical_diffusivity_2Pr.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes +using Glob + + +#%% +filename = "doublegyre_seasonalforcing_30C-20C_relaxation_8days_CATKEVerticalDiffusivity" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +# vertical_base_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +# convection_closure = XinKaiVerticalDiffusivity() +function CATKE_ocean_closure() + mixing_length = CATKEMixingLength(Cᵇ=0.01) + turbulent_kinetic_energy_equation = CATKEEquation(Cᵂϵ=1.0) + return CATKEVerticalDiffusivity(; mixing_length, turbulent_kinetic_energy_equation) +end +convection_closure = CATKE_ocean_closure() +closure = convection_closure +# closure = vertical_base_closure + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +const Δy = Ly / Ny + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 10 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/8days + +const seasonal_period = 360days +const seasonal_forcing_width = Ly / 6 +const seasonal_T_amplitude = 20 + +##### +##### Forcing and initial condition +##### +@inline T_initial(x, y, z) = 20 + 10 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_seasonal(y, t) = seasonal_T_amplitude * exp(-y^2/(2 * seasonal_forcing_width^2)) * sin(2π * t / seasonal_period) +@inline T_ref(y, t) = T_mid - ΔT / Ly * y + T_seasonal(y, t) +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y, t)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) + +@inline T_north_ref(z) = 10 * (1 + z / Lz) +@inline north_T_flux(x, z, t, T) = μ_T * Δy * (T - T_north_ref(z)) +north_T_flux_bc = FluxBoundaryCondition(north_T_flux; field_dependencies=:T) + +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc, north = north_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S, :e), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S, :e), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10800days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)) +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) +compute!(ρ) + +ρᶠ = @at (Center, Center, Face) ρ +∂ρ∂z = ∂z(ρ) +∂²ρ∂z² = ∂z(∂ρ∂z) + +κc = model.diffusivity_fields.κc +wT = κc * ∂z(T) +wS = κc * ∂z(S) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + +outputs = (; u, v, w, T, S, ρ, ρᶠ, ∂ρ∂z, ∂²ρ∂z², N², wT, wS) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 360 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% \ No newline at end of file From 0cd618afdab4b474a53d6b5f864543223553b826 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 17 Oct 2024 15:04:50 -0400 Subject: [PATCH 071/114] baseclosure doublegyre with wallrestoration --- ...e_model_seasonalforcing_wallrestoration.jl | 584 ++++++++++++++++++ 1 file changed, 584 insertions(+) create mode 100644 baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl diff --git a/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl b/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl new file mode 100644 index 0000000000..097f7a08fa --- /dev/null +++ b/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl @@ -0,0 +1,584 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("xin_kai_vertical_diffusivity_local_2step.jl") +# include("xin_kai_vertical_diffusivity_2Pr.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + +#%% +filename = "doublegyre_seasonalforcing_30C-20C_relaxation_8days_baseclosure_test" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +base_closure = XinKaiLocalVerticalDiffusivity() +vertical_scalar_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +closure = (base_closure, vertical_scalar_closure) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +const Δy = Ly / Ny + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 10 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/8days + +const seasonal_period = 360days +const seasonal_forcing_width = Ly / 6 +const seasonal_T_amplitude = 20 + +##### +##### Forcing and initial condition +##### +@inline T_initial(x, y, z) = 20 + 10 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_seasonal(y, t) = seasonal_T_amplitude * exp(-y^2/(2 * seasonal_forcing_width^2)) * sin(2π * t / seasonal_period) +@inline T_ref(y, t) = T_mid - ΔT / Ly * y + T_seasonal(y, t) +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y, t)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) + +@inline T_north_ref(z) = 10 * (1 + z / Lz) +@inline north_T_flux(x, z, t, T) = μ_T * Δy * (T - T_north_ref(z)) +north_T_flux_bc = FluxBoundaryCondition(north_T_flux; field_dependencies=:T) + +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc, north = north_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10800days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); + +κ = model.diffusivity_fields[1].κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + +outputs = (; u, v, w, T, S, ρ, N², wT_base, wS_base) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:complete_fields] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields", + schedule = TimeInterval(1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 360 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% +# function find_min(a...) +# return minimum(minimum.([a...])) +# end + +# function find_max(a...) +# return maximum(maximum.([a...])) +# end + +# N²_xz_north_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_north.jld2", "N²") +# N²_xz_south_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_south.jld2", "N²") + +# xC = N²_xz_north_data.grid.xᶜᵃᵃ[1:N²_xz_north_data.grid.Nx] +# zf = N²_xz_north_data.grid.zᵃᵃᶠ[1:N²_xz_north_data.grid.Nz+1] + +# yloc_north = N²_xz_north_data.grid.yᵃᶠᵃ[N²_xz_north_data.indices[2][1]] +# yloc_south = N²_xz_south_data.grid.yᵃᶠᵃ[N²_xz_south_data.indices[2][1]] + +# Nt = length(N²_xz_north_data) +# times = N²_xz_north_data.times / 24 / 60^2 / 360 +# timeframes = 1:Nt + +# N²_lim = (find_min(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) - 1e-13, +# find_max(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) + 1e-13) +# #%% +# fig = Figure(size=(800, 800)) +# ax_north = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_south / 1e3)) km") +# ax_south = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_north / 1e3)) km") + +# n = Observable(2) + +# N²_north = @lift interior(N²_xz_north_data[$n], :, 1, :) +# N²_south = @lift interior(N²_xz_south_data[$n], :, 1, :) + +# colorscheme = colorschemes[:jet] + +# N²_north_surface = heatmap!(ax_north, xC, zf, N²_north, colormap=colorscheme, colorrange=N²_lim) +# N²_south_surface = heatmap!(ax_south, xC, zf, N²_south, colormap=colorscheme, colorrange=N²_lim) + +# Colorbar(fig[1:2, 2], N²_north_surface, label="N² (s⁻²)") + +# title_str = @lift "Time = $(round(times[$n], digits=2)) years" +# Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +# trim!(fig.layout) + +# @info "Recording buoyancy frequency xz slice" +# CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_buoyancy_frequency_xz_slice.mp4", 1:Nt, framerate=3, px_per_unit=2) do nn +# n[] = nn +# end + +# @info "Done!" +# #%% \ No newline at end of file From 4abf7e02588854f4a0e9a415abd22788c40eea72 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 17 Oct 2024 15:07:08 -0400 Subject: [PATCH 072/114] change filenames --- baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl | 2 +- ...lclosure_doublegyre_model_seasonalforcing_wallrestoration.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl b/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl index 097f7a08fa..35d38a2b4d 100644 --- a/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl +++ b/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl @@ -21,7 +21,7 @@ using SeawaterPolynomials:TEOS10 using ColorSchemes #%% -filename = "doublegyre_seasonalforcing_30C-20C_relaxation_8days_baseclosure_test" +filename = "doublegyre_seasonalforcing_30C-20C_relaxation_wallrestoration_8days_baseclosure_test" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) diff --git a/physicalclosure_doublegyre_model_seasonalforcing_wallrestoration.jl b/physicalclosure_doublegyre_model_seasonalforcing_wallrestoration.jl index 069d71fc7f..355b8317f7 100644 --- a/physicalclosure_doublegyre_model_seasonalforcing_wallrestoration.jl +++ b/physicalclosure_doublegyre_model_seasonalforcing_wallrestoration.jl @@ -28,7 +28,7 @@ using Glob #%% -filename = "doublegyre_seasonalforcing_30C-20C_relaxation_8days_CATKEVerticalDiffusivity" +filename = "doublegyre_seasonalforcing_30C-20C_relaxation_wallrestoration_8days_CATKEVerticalDiffusivity" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) From 7cf70906d06891710824e8a9b108e14f5b394ac8 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 17 Oct 2024 17:07:21 -0400 Subject: [PATCH 073/114] using wider zone for NN closure --- ...losure_global_Ri_nof_BBLkappazonelast55.jl | 268 ++++++++++++++++++ 1 file changed, 268 insertions(+) create mode 100644 NN_closure_global_Ri_nof_BBLkappazonelast55.jl diff --git a/NN_closure_global_Ri_nof_BBLkappazonelast55.jl b/NN_closure_global_Ri_nof_BBLkappazonelast55.jl new file mode 100644 index 0000000000..9a654f081c --- /dev/null +++ b/NN_closure_global_Ri_nof_BBLkappazonelast55.jl @@ -0,0 +1,268 @@ +import Oceananigans.TurbulenceClosures: + compute_diffusivities!, + DiffusivityFields, + viscosity, + diffusivity, + diffusive_flux_x, + diffusive_flux_y, + diffusive_flux_z, + viscous_flux_ux, + viscous_flux_uy, + viscous_flux_uz, + viscous_flux_vx, + viscous_flux_vy, + viscous_flux_vz, + viscous_flux_wx, + viscous_flux_wy, + viscous_flux_wz + +using Oceananigans.BuoyancyModels: ∂x_b, ∂y_b, ∂z_b +using Oceananigans.Coriolis +using Oceananigans.Grids: φnode +using Oceananigans.Grids: total_size +using Oceananigans.Utils: KernelParameters +using Oceananigans: architecture, on_architecture +using Lux, LuxCUDA +using JLD2 +using ComponentArrays +using OffsetArrays +using SeawaterPolynomials.TEOS10 + +using KernelAbstractions: @index, @kernel, @private + +import Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure, ExplicitTimeDiscretization + +using Adapt + +include("./feature_scaling.jl") + +@inline hack_sind(φ) = sin(φ * π / 180) + +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::HydrostaticSphericalCoriolis) = 2 * coriolis.rotation_rate * hack_sind(φnode(i, j, k, grid, Center(), Center(), Center())) +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::FPlane) = coriolis.f +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::BetaPlane) = coriolis.f₀ + coriolis.β * ynode(i, j, k, grid, Center(), Center(), Center()) + +struct NN{M, P, S} + model :: M + ps :: P + st :: S +end + +@inline (neural_network::NN)(input) = first(neural_network.model(input, neural_network.ps, neural_network.st)) +@inline tosarray(x::AbstractArray) = SArray{Tuple{size(x)...}}(x) + +struct NNFluxClosure{A <: NN, S} <: AbstractTurbulenceClosure{ExplicitTimeDiscretization, 3} + wT :: A + wS :: A + scaling :: S +end + +Adapt.adapt_structure(to, nn :: NNFluxClosure) = + NNFluxClosure(Adapt.adapt(to, nn.wT), + Adapt.adapt(to, nn.wS), + Adapt.adapt(to, nn.scaling)) + +Adapt.adapt_structure(to, nn :: NN) = + NN(Adapt.adapt(to, nn.model), + Adapt.adapt(to, nn.ps), + Adapt.adapt(to, nn.st)) + +function NNFluxClosure(arch) + dev = ifelse(arch == GPU(), gpu_device(), cpu_device()) + nn_path = "./NDE_FC_Qb_Ri_nof_BBLkappazonelast55_trainFC19new_scalingtrain59new_1layer_512_relu_10seed_2Pr_model_temp.jld2" + + ps, sts, scaling_params, wT_model, wS_model = jldopen(nn_path, "r") do file + ps = file["u"] |> dev |> f64 + sts = file["sts"] |> dev |> f64 + scaling_params = file["scaling"] + wT_model = file["model"].wT + wS_model = file["model"].wS + return ps, sts, scaling_params, wT_model, wS_model + end + + scaling = construct_zeromeanunitvariance_scaling(scaling_params) + + wT_NN = NN(wT_model, ps.wT, sts.wT) + wS_NN = NN(wS_model, ps.wS, sts.wS) + + return NNFluxClosure(wT_NN, wS_NN, scaling) +end + +function DiffusivityFields(grid, tracer_names, bcs, ::NNFluxClosure) + arch = architecture(grid) + wT = ZFaceField(grid) + wS = ZFaceField(grid) + first_index = Field((Center, Center, Nothing), grid, Int32) + last_index = Field((Center, Center, Nothing), grid, Int32) + + Nx_in, Ny_in, _ = size(wT) + wrk_in = zeros(13, Nx_in, Ny_in, 10) + wrk_in = on_architecture(arch, wrk_in) + + wrk_wT = zeros(Nx_in, Ny_in, 10) + wrk_wS = zeros(Nx_in, Ny_in, 10) + wrk_wT = on_architecture(arch, wrk_wT) + wrk_wS = on_architecture(arch, wrk_wS) + + return (; wrk_in, wrk_wT, wrk_wS, wT, wS, first_index, last_index) +end + +function compute_diffusivities!(diffusivities, closure::NNFluxClosure, model; parameters = :xyz) + arch = model.architecture + grid = model.grid + velocities = model.velocities + tracers = model.tracers + buoyancy = model.buoyancy + coriolis = model.coriolis + clock = model.clock + top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) + + κᶜ = model.diffusivity_fields[1].κᶜ + κ₀ = model.closure[1].ν₀ / model.closure[1].Pr_shearₜ + Ri = model.diffusivity_fields[1].Ri + + wrk_in = diffusivities.wrk_in + wrk_wT = diffusivities.wrk_wT + wrk_wS = diffusivities.wrk_wS + wT = diffusivities.wT + wS = diffusivities.wS + + first_index = diffusivities.first_index + last_index = diffusivities.last_index + + Nx_in, Ny_in, Nz_in = total_size(wT) + ox_in, oy_in, oz_in = wT.data.offsets + kp = KernelParameters((Nx_in, Ny_in, Nz_in), (ox_in, oy_in, oz_in)) + kp_2D = KernelParameters((Nx_in, Ny_in), (ox_in, oy_in)) + + kp_wrk = KernelParameters((Nx_in, Ny_in, 10), (0, 0, 0)) + + launch!(arch, grid, kp_2D, _find_NN_active_region!, κᶜ, grid, κ₀, first_index, last_index) + + launch!(arch, grid, kp_wrk, + _populate_input!, wrk_in, first_index, last_index, grid, closure, tracers, velocities, buoyancy, top_tracer_bcs, Ri, clock) + + wrk_wT .= dropdims(closure.wT(wrk_in), dims=1) + wrk_wS .= dropdims(closure.wS(wrk_in), dims=1) + + launch!(arch, grid, kp, _fill_adjust_nn_fluxes!, diffusivities, first_index, last_index, wrk_wT, wrk_wS, grid, closure, tracers, velocities, buoyancy, top_tracer_bcs, clock) + return nothing +end + +@kernel function _populate_input!(input, first_index, last_index, grid, closure::NNFluxClosure, tracers, velocities, buoyancy, top_tracer_bcs, Ri, clock) + i, j, k = @index(Global, NTuple) + + scaling = closure.scaling + + ρ₀ = buoyancy.model.equation_of_state.reference_density + g = buoyancy.model.gravitational_acceleration + eos = TEOS10.TEOS10EquationOfState() + + @inbounds k_first = first_index[i, j, 1] + @inbounds k_last = last_index[i, j, 1] + + quiescent = quiescent_condition(k_first, k_last) + + @inbounds k_tracer = clamp_k_interior(k_first + k - 1, grid) + + T, S = tracers.T, tracers.S + + @inbounds input[1, i, j, k] = Riᵢ₋₁ = ifelse(quiescent, 0, atan(Ri[i, j, k_tracer-1])) + @inbounds input[2, i, j, k] = Riᵢ = ifelse(quiescent, 0, atan(Ri[i, j, k_tracer])) + @inbounds input[3, i, j, k] = Riᵢ₊₁ = ifelse(quiescent, 0, atan(Ri[i, j, k_tracer+1])) + + @inbounds input[4, i, j, k] = ∂Tᵢ₋₁ = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer-1, grid, T))) + @inbounds input[5, i, j, k] = ∂Tᵢ = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer, grid, T))) + @inbounds input[6, i, j, k] = ∂Tᵢ₊₁ = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer+1, grid, T))) + + @inbounds input[7, i, j, k] = ∂Sᵢ₋₁ = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer-1, grid, S))) + @inbounds input[8, i, j, k] = ∂Sᵢ = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer, grid, S))) + @inbounds input[9, i, j, k] = ∂Sᵢ₊₁ = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer+1, grid, S))) + + @inbounds input[10, i, j, k] = ∂σᵢ₋₁ = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k_tracer-1, grid, buoyancy, tracers) / g)) + @inbounds input[11, i, j, k] = ∂σᵢ = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k_tracer, grid, buoyancy, tracers) / g)) + @inbounds input[12, i, j, k] = ∂σᵢ₊₁ = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k_tracer+1, grid, buoyancy, tracers) / g)) + + @inbounds input[13, i, j, k] = Jᵇ = scaling.wb(top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers))) + +end + +@kernel function _find_NN_active_region!(κᶜ, grid, κ₀, first_index, last_index) + i, j = @index(Global, NTuple) + top_index = grid.Nz + 1 + grid_point_above_kappa = 5 + grid_point_below_kappa = 5 + + # Find the last index of the background κᶜ + kloc = 1 + @inbounds for k in 2:grid.Nz + kloc = ifelse(κᶜ[i, j, k] ≈ κ₀, k, kloc) + end + + nonbackground_κ_index = kloc + 1 + + @inbounds last_index[i, j, 1] = ifelse(nonbackground_κ_index == top_index, top_index - 1, clamp_k_interior(kloc + grid_point_above_kappa, grid)) + @inbounds first_index[i, j, 1] = ifelse(nonbackground_κ_index == top_index, top_index, clamp_k_interior(kloc - grid_point_below_kappa + 1, grid)) +end + +@inline function quiescent_condition(lo, hi) + return hi < lo +end + +@inline function within_zone_condition(k, lo, hi) + return (k >= lo) & (k <= hi) +end + +@inline function clamp_k_interior(k, grid) + kmax = grid.Nz - 1 + kmin = 3 + + return clamp(k, kmin, kmax) +end + +@kernel function _fill_adjust_nn_fluxes!(diffusivities, first_index, last_index, wrk_wT, wrk_wS, grid, closure::NNFluxClosure, tracers, velocities, buoyancy, top_tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + scaling = closure.scaling + + @inbounds k_first = first_index[i, j, 1] + @inbounds k_last = last_index[i, j, 1] + + convecting = top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers)) > 0 + quiescent = quiescent_condition(k_first, k_last) + within_zone = within_zone_condition(k, k_first, k_last) + + @inbounds k_wrk = clamp(k - k_first + 1, 1, 10) + + NN_active = convecting & !quiescent & within_zone + + @inbounds diffusivities.wT[i, j, k] = ifelse(NN_active, scaling.wT.σ * wrk_wT[i, j, k_wrk], 0) + @inbounds diffusivities.wS[i, j, k] = ifelse(NN_active, scaling.wS.σ * wrk_wS[i, j, k_wrk], 0) +end + +# Write here your constructor +# NNFluxClosure() = ... insert NN here ... (make sure it is on GPU if you need it on GPU!) + +const NNC = NNFluxClosure + +##### +##### Abstract Smagorinsky functionality +##### + +# Horizontal fluxes are zero! +@inline viscous_flux_wz( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_ux( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_uy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline diffusive_flux_x(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) +@inline diffusive_flux_y(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) + +# Viscous fluxes are zero (for now) +@inline viscous_flux_uz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) + +# The only function extended by NNFluxClosure +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{1}, c, clock, fields, buoyancy) = @inbounds K.wT[i, j, k] +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{2}, c, clock, fields, buoyancy) = @inbounds K.wS[i, j, k] \ No newline at end of file From 70dbe2fb2939e05ab1299e64a7678323b94cea20 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 17 Oct 2024 17:12:13 -0400 Subject: [PATCH 074/114] fix variable error --- ...losure_doublegyre_model_seasonalforcing_wallrestoration.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl b/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl index 35d38a2b4d..d665d645d5 100644 --- a/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl +++ b/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl @@ -43,7 +43,7 @@ const Lx = 4000kilometers const Ly = 6000kilometers const Lz = Nz * Δz -const Δy = Ly / Ny +const δy = Ly / Ny grid = RectilinearGrid(model_architecture, Float64, topology = (Bounded, Bounded, Bounded), @@ -107,7 +107,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) @inline T_north_ref(z) = 10 * (1 + z / Lz) -@inline north_T_flux(x, z, t, T) = μ_T * Δy * (T - T_north_ref(z)) +@inline north_T_flux(x, z, t, T) = μ_T * δy * (T - T_north_ref(z)) north_T_flux_bc = FluxBoundaryCondition(north_T_flux; field_dependencies=:T) T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc, north = north_T_flux_bc) From fd57bf884ec9957e519ed966ca7b2c49b9a05dee Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Fri, 18 Oct 2024 00:29:49 -0400 Subject: [PATCH 075/114] update NN configuration --- NN_closure_global_Ri_nof_BBLkappazonelast55.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NN_closure_global_Ri_nof_BBLkappazonelast55.jl b/NN_closure_global_Ri_nof_BBLkappazonelast55.jl index 9a654f081c..2255a8697b 100644 --- a/NN_closure_global_Ri_nof_BBLkappazonelast55.jl +++ b/NN_closure_global_Ri_nof_BBLkappazonelast55.jl @@ -69,7 +69,7 @@ Adapt.adapt_structure(to, nn :: NN) = function NNFluxClosure(arch) dev = ifelse(arch == GPU(), gpu_device(), cpu_device()) - nn_path = "./NDE_FC_Qb_Ri_nof_BBLkappazonelast55_trainFC19new_scalingtrain59new_1layer_512_relu_10seed_2Pr_model_temp.jld2" + nn_path = "./NDE_FC_Qb_Ri_nof_BBLkappazonelast55_trainFC19new_scalingtrain59new_2layer_256_relu_10seed_2Pr_model_temp.jld2" ps, sts, scaling_params, wT_model, wS_model = jldopen(nn_path, "r") do file ps = file["u"] |> dev |> f64 From b582447fc216d16a5f51ac7399ee5204f3e85493 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Fri, 18 Oct 2024 00:35:42 -0400 Subject: [PATCH 076/114] NDE double gyre script for seasonal forcing and wall restoration --- ...last55_doublegyre_model_seasonalforcing.jl | 539 +++++++++++++++++ ...e_model_seasonalforcing_wallrestoration.jl | 549 ++++++++++++++++++ 2 files changed, 1088 insertions(+) create mode 100644 NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model_seasonalforcing.jl create mode 100644 NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model_seasonalforcing_wallrestoration.jl diff --git a/NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model_seasonalforcing.jl b/NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model_seasonalforcing.jl new file mode 100644 index 0000000000..11dd334c24 --- /dev/null +++ b/NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model_seasonalforcing.jl @@ -0,0 +1,539 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("NN_closure_global_Ri_nof_BBLkappazonelast55.jl") +include("xin_kai_vertical_diffusivity_local_2step.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + + +#%% +filename = "doublegyre_seasonalforcing_30C-20C_relaxation_8days_NN_closure_NDE_Ri_BBLkappazonelast55_temp" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +nn_closure = NNFluxClosure(model_architecture) +base_closure = XinKaiLocalVerticalDiffusivity() +vertical_scalar_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +closure = (base_closure, nn_closure, vertical_scalar_closure) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/8days + +const seasonal_period = 360days +const seasonal_forcing_width = Ly / 6 +const seasonal_T_amplitude = 15 + +##### +##### Forcing and initial condition +##### +@inline T_initial(x, y, z) = 20 + 10 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_seasonal(y, t) = seasonal_T_amplitude * exp(-y^2/(2 * seasonal_forcing_width^2)) * sin(2π * t / seasonal_period) +@inline T_ref(y, t) = T_mid - ΔT / Ly * y + T_seasonal(y, t) +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y, t)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 36000days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); +first_index = model.diffusivity_fields[2].first_index +last_index = model.diffusivity_fields[2].last_index +wT_NN = model.diffusivity_fields[2].wT +wS_NN = model.diffusivity_fields[2].wS + +κ = model.diffusivity_fields[1].κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +@inline function get_top_buoyancy_flux(i, j, k, grid, buoyancy, T_bc, S_bc, velocities, tracers, clock) + return top_buoyancy_flux(i, j, grid, buoyancy, (; T=T_bc, S=S_bc), clock, merge(velocities, tracers)) +end + +Qb = KernelFunctionOperation{Center, Center, Nothing}(get_top_buoyancy_flux, model.grid, model.buoyancy, T.boundary_conditions.top, S.boundary_conditions.top, model.velocities, model.tracers, model.clock) +Qb = Field(Qb) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + +outputs = (; u, v, w, T, S, ρ, N², wT_NN, wS_NN, wT_base, wS_base) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; first_index, last_index, Qb), + filename = "$(FILE_DIR)/instantaneous_fields_NN_active_diagnostics", + indices = (:, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 360 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% \ No newline at end of file diff --git a/NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model_seasonalforcing_wallrestoration.jl b/NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model_seasonalforcing_wallrestoration.jl new file mode 100644 index 0000000000..0d0aea8d81 --- /dev/null +++ b/NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model_seasonalforcing_wallrestoration.jl @@ -0,0 +1,549 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("NN_closure_global_Ri_nof_BBLkappazonelast55.jl") +include("xin_kai_vertical_diffusivity_local_2step.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + + +#%% +filename = "doublegyre_seasonalforcing_30C-20C_relaxation_wallrestoration_8days_NN_closure_NDE_Ri_BBLkappazonelast55_temp" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +nn_closure = NNFluxClosure(model_architecture) +base_closure = XinKaiLocalVerticalDiffusivity() +vertical_scalar_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +closure = (base_closure, nn_closure, vertical_scalar_closure) + +# number of grid points +##### +##### Boundary conditions +##### +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +const δy = Ly / Ny + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 10 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/8days + +const seasonal_period = 360days +const seasonal_forcing_width = Ly / 6 +const seasonal_T_amplitude = 20 + +##### +##### Forcing and initial condition +##### +@inline T_initial(x, y, z) = 20 + 10 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_seasonal(y, t) = seasonal_T_amplitude * exp(-y^2/(2 * seasonal_forcing_width^2)) * sin(2π * t / seasonal_period) +@inline T_ref(y, t) = T_mid - ΔT / Ly * y + T_seasonal(y, t) +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y, t)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) + +@inline T_north_ref(z) = 10 * (1 + z / Lz) +@inline north_T_flux(x, z, t, T) = μ_T * δy * (T - T_north_ref(z)) +north_T_flux_bc = FluxBoundaryCondition(north_T_flux; field_dependencies=:T) + +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc, north = north_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10800days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); +first_index = model.diffusivity_fields[2].first_index +last_index = model.diffusivity_fields[2].last_index +wT_NN = model.diffusivity_fields[2].wT +wS_NN = model.diffusivity_fields[2].wS + +κ = model.diffusivity_fields[1].κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +@inline function get_top_buoyancy_flux(i, j, k, grid, buoyancy, T_bc, S_bc, velocities, tracers, clock) + return top_buoyancy_flux(i, j, grid, buoyancy, (; T=T_bc, S=S_bc), clock, merge(velocities, tracers)) +end + +Qb = KernelFunctionOperation{Center, Center, Nothing}(get_top_buoyancy_flux, model.grid, model.buoyancy, T.boundary_conditions.top, S.boundary_conditions.top, model.velocities, model.tracers, model.clock) +Qb = Field(Qb) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + +outputs = (; u, v, w, T, S, ρ, N², wT_NN, wS_NN, wT_base, wS_base) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; first_index, last_index, Qb), + filename = "$(FILE_DIR)/instantaneous_fields_NN_active_diagnostics", + indices = (:, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 360 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% \ No newline at end of file From 6235bc4d3bbea3a49ee97d0c25dd555549c1af04 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Sun, 20 Oct 2024 00:36:43 -0400 Subject: [PATCH 077/114] run NNclosure with Ri nof BBLkappazonelast55 --- ...nof_BBLkappazonelast55_doublegyre_model.jl | 533 ++++++++++++++++++ 1 file changed, 533 insertions(+) create mode 100644 NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model.jl diff --git a/NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model.jl b/NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model.jl new file mode 100644 index 0000000000..f380820afe --- /dev/null +++ b/NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model.jl @@ -0,0 +1,533 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("NN_closure_global_Ri_nof_BBLkappazonelast55.jl") +include("xin_kai_vertical_diffusivity_local_2step.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + +#%% +filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_NN_closure_NDE_Ri_BBLkappazonelast55_temp" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +nn_closure = NNFluxClosure(model_architecture) +base_closure = XinKaiLocalVerticalDiffusivity() +vertical_scalar_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +closure = (base_closure, nn_closure, vertical_scalar_closure) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/8days + +##### +##### Forcing and initial condition +##### +@inline T_initial(x, y, z) = 10 + 20 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10950days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); +first_index = model.diffusivity_fields[2].first_index +last_index = model.diffusivity_fields[2].last_index +wT_NN = model.diffusivity_fields[2].wT +wS_NN = model.diffusivity_fields[2].wS + +κ = model.diffusivity_fields[1].κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +@inline function get_top_buoyancy_flux(i, j, k, grid, buoyancy, T_bc, S_bc, velocities, tracers, clock) + return top_buoyancy_flux(i, j, grid, buoyancy, (; T=T_bc, S=S_bc), clock, merge(velocities, tracers)) +end + +Qb = KernelFunctionOperation{Center, Center, Nothing}(get_top_buoyancy_flux, model.grid, model.buoyancy, T.boundary_conditions.top, S.boundary_conditions.top, model.velocities, model.tracers, model.clock) +Qb = Field(Qb) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + +outputs = (; u, v, w, T, S, ρ, N², wT_NN, wS_NN, wT_base, wS_base) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; first_index, last_index, Qb), + filename = "$(FILE_DIR)/instantaneous_fields_NN_active_diagnostics", + indices = (:, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 365 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% \ No newline at end of file From 3609c212f98b54018ee9535e5c48433fce14d9d6 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Mon, 21 Oct 2024 23:15:32 -0400 Subject: [PATCH 078/114] updated NN model with no Ri --- NN_closure_global_nof_BBLkappazonelast55.jl | 263 ++++++++++++++++++++ 1 file changed, 263 insertions(+) create mode 100644 NN_closure_global_nof_BBLkappazonelast55.jl diff --git a/NN_closure_global_nof_BBLkappazonelast55.jl b/NN_closure_global_nof_BBLkappazonelast55.jl new file mode 100644 index 0000000000..5602b767f9 --- /dev/null +++ b/NN_closure_global_nof_BBLkappazonelast55.jl @@ -0,0 +1,263 @@ +import Oceananigans.TurbulenceClosures: + compute_diffusivities!, + DiffusivityFields, + viscosity, + diffusivity, + diffusive_flux_x, + diffusive_flux_y, + diffusive_flux_z, + viscous_flux_ux, + viscous_flux_uy, + viscous_flux_uz, + viscous_flux_vx, + viscous_flux_vy, + viscous_flux_vz, + viscous_flux_wx, + viscous_flux_wy, + viscous_flux_wz + +using Oceananigans.BuoyancyModels: ∂x_b, ∂y_b, ∂z_b +using Oceananigans.Coriolis +using Oceananigans.Grids: φnode +using Oceananigans.Grids: total_size +using Oceananigans.Utils: KernelParameters +using Oceananigans: architecture, on_architecture +using Lux, LuxCUDA +using JLD2 +using ComponentArrays +using OffsetArrays +using SeawaterPolynomials.TEOS10 + +using KernelAbstractions: @index, @kernel, @private + +import Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure, ExplicitTimeDiscretization + +using Adapt + +include("./feature_scaling.jl") + +@inline hack_sind(φ) = sin(φ * π / 180) + +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::HydrostaticSphericalCoriolis) = 2 * coriolis.rotation_rate * hack_sind(φnode(i, j, k, grid, Center(), Center(), Center())) +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::FPlane) = coriolis.f +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::BetaPlane) = coriolis.f₀ + coriolis.β * ynode(i, j, k, grid, Center(), Center(), Center()) + +struct NN{M, P, S} + model :: M + ps :: P + st :: S +end + +@inline (neural_network::NN)(input) = first(neural_network.model(input, neural_network.ps, neural_network.st)) +@inline tosarray(x::AbstractArray) = SArray{Tuple{size(x)...}}(x) + +struct NNFluxClosure{A <: NN, S} <: AbstractTurbulenceClosure{ExplicitTimeDiscretization, 3} + wT :: A + wS :: A + scaling :: S +end + +Adapt.adapt_structure(to, nn :: NNFluxClosure) = + NNFluxClosure(Adapt.adapt(to, nn.wT), + Adapt.adapt(to, nn.wS), + Adapt.adapt(to, nn.scaling)) + +Adapt.adapt_structure(to, nn :: NN) = + NN(Adapt.adapt(to, nn.model), + Adapt.adapt(to, nn.ps), + Adapt.adapt(to, nn.st)) + +function NNFluxClosure(arch) + dev = ifelse(arch == GPU(), gpu_device(), cpu_device()) + nn_path = "./NDE_FC_Qb_nof_BBLkappazonelast55_trainFC23new_scalingtrain53new_2layer_128_relu_20seed_2Pr_model_temp.jld2" + + ps, sts, scaling_params, wT_model, wS_model = jldopen(nn_path, "r") do file + ps = file["u"] |> dev |> f64 + sts = file["sts"] |> dev |> f64 + scaling_params = file["scaling"] + wT_model = file["model"].wT + wS_model = file["model"].wS + return ps, sts, scaling_params, wT_model, wS_model + end + + scaling = construct_zeromeanunitvariance_scaling(scaling_params) + + wT_NN = NN(wT_model, ps.wT, sts.wT) + wS_NN = NN(wS_model, ps.wS, sts.wS) + + return NNFluxClosure(wT_NN, wS_NN, scaling) +end + +function DiffusivityFields(grid, tracer_names, bcs, ::NNFluxClosure) + arch = architecture(grid) + wT = ZFaceField(grid) + wS = ZFaceField(grid) + first_index = Field((Center, Center, Nothing), grid, Int32) + last_index = Field((Center, Center, Nothing), grid, Int32) + + Nx_in, Ny_in, _ = size(wT) + wrk_in = zeros(10, Nx_in, Ny_in, 10) + wrk_in = on_architecture(arch, wrk_in) + + wrk_wT = zeros(Nx_in, Ny_in, 10) + wrk_wS = zeros(Nx_in, Ny_in, 10) + wrk_wT = on_architecture(arch, wrk_wT) + wrk_wS = on_architecture(arch, wrk_wS) + + return (; wrk_in, wrk_wT, wrk_wS, wT, wS, first_index, last_index) +end + +function compute_diffusivities!(diffusivities, closure::NNFluxClosure, model; parameters = :xyz) + arch = model.architecture + grid = model.grid + velocities = model.velocities + tracers = model.tracers + buoyancy = model.buoyancy + coriolis = model.coriolis + clock = model.clock + top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) + + κᶜ = model.diffusivity_fields[1].κᶜ + κ₀ = model.closure[1].ν₀ / model.closure[1].Pr_shearₜ + Ri = model.diffusivity_fields[1].Ri + + wrk_in = diffusivities.wrk_in + wrk_wT = diffusivities.wrk_wT + wrk_wS = diffusivities.wrk_wS + wT = diffusivities.wT + wS = diffusivities.wS + + first_index = diffusivities.first_index + last_index = diffusivities.last_index + + Nx_in, Ny_in, Nz_in = total_size(wT) + ox_in, oy_in, oz_in = wT.data.offsets + kp = KernelParameters((Nx_in, Ny_in, Nz_in), (ox_in, oy_in, oz_in)) + kp_2D = KernelParameters((Nx_in, Ny_in), (ox_in, oy_in)) + + kp_wrk = KernelParameters((Nx_in, Ny_in, 10), (0, 0, 0)) + + launch!(arch, grid, kp_2D, _find_NN_active_region!, κᶜ, grid, κ₀, first_index, last_index) + + launch!(arch, grid, kp_wrk, + _populate_input!, wrk_in, first_index, last_index, grid, closure, tracers, velocities, buoyancy, top_tracer_bcs, Ri, clock) + + wrk_wT .= dropdims(closure.wT(wrk_in), dims=1) + wrk_wS .= dropdims(closure.wS(wrk_in), dims=1) + + launch!(arch, grid, kp, _fill_adjust_nn_fluxes!, diffusivities, first_index, last_index, wrk_wT, wrk_wS, grid, closure, tracers, velocities, buoyancy, top_tracer_bcs, clock) + return nothing +end + +@kernel function _populate_input!(input, first_index, last_index, grid, closure::NNFluxClosure, tracers, velocities, buoyancy, top_tracer_bcs, Ri, clock) + i, j, k = @index(Global, NTuple) + + scaling = closure.scaling + + ρ₀ = buoyancy.model.equation_of_state.reference_density + g = buoyancy.model.gravitational_acceleration + + @inbounds k_first = first_index[i, j, 1] + @inbounds k_last = last_index[i, j, 1] + + quiescent = quiescent_condition(k_first, k_last) + + @inbounds k_tracer = clamp_k_interior(k_first + k - 1, grid) + + T, S = tracers.T, tracers.S + + @inbounds input[1, i, j, k] = ∂Tᵢ₋₁ = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer-1, grid, T))) + @inbounds input[2, i, j, k] = ∂Tᵢ = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer, grid, T))) + @inbounds input[3, i, j, k] = ∂Tᵢ₊₁ = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer+1, grid, T))) + + @inbounds input[4, i, j, k] = ∂Sᵢ₋₁ = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer-1, grid, S))) + @inbounds input[5, i, j, k] = ∂Sᵢ = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer, grid, S))) + @inbounds input[6, i, j, k] = ∂Sᵢ₊₁ = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k_tracer+1, grid, S))) + + @inbounds input[7, i, j, k] = ∂σᵢ₋₁ = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k_tracer-1, grid, buoyancy, tracers) / g)) + @inbounds input[8, i, j, k] = ∂σᵢ = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k_tracer, grid, buoyancy, tracers) / g)) + @inbounds input[9, i, j, k] = ∂σᵢ₊₁ = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k_tracer+1, grid, buoyancy, tracers) / g)) + + @inbounds input[10, i, j, k] = Jᵇ = scaling.wb(top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers))) + +end + +@kernel function _find_NN_active_region!(κᶜ, grid, κ₀, first_index, last_index) + i, j = @index(Global, NTuple) + top_index = grid.Nz + 1 + grid_point_above_kappa = 5 + grid_point_below_kappa = 5 + + # Find the last index of the background κᶜ + kloc = 1 + @inbounds for k in 2:grid.Nz + kloc = ifelse(κᶜ[i, j, k] ≈ κ₀, k, kloc) + end + + nonbackground_κ_index = kloc + 1 + + @inbounds last_index[i, j, 1] = ifelse(nonbackground_κ_index == top_index, top_index - 1, clamp_k_interior(kloc + grid_point_above_kappa, grid)) + @inbounds first_index[i, j, 1] = ifelse(nonbackground_κ_index == top_index, top_index, clamp_k_interior(kloc - grid_point_below_kappa + 1, grid)) +end + +@inline function quiescent_condition(lo, hi) + return hi < lo +end + +@inline function within_zone_condition(k, lo, hi) + return (k >= lo) & (k <= hi) +end + +@inline function clamp_k_interior(k, grid) + kmax = grid.Nz - 1 + kmin = 3 + + return clamp(k, kmin, kmax) +end + +@kernel function _fill_adjust_nn_fluxes!(diffusivities, first_index, last_index, wrk_wT, wrk_wS, grid, closure::NNFluxClosure, tracers, velocities, buoyancy, top_tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + scaling = closure.scaling + + @inbounds k_first = first_index[i, j, 1] + @inbounds k_last = last_index[i, j, 1] + + convecting = top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers)) > 0 + quiescent = quiescent_condition(k_first, k_last) + within_zone = within_zone_condition(k, k_first, k_last) + + @inbounds k_wrk = clamp(k - k_first + 1, 1, 10) + + NN_active = convecting & !quiescent & within_zone + + @inbounds diffusivities.wT[i, j, k] = ifelse(NN_active, scaling.wT.σ * wrk_wT[i, j, k_wrk], 0) + @inbounds diffusivities.wS[i, j, k] = ifelse(NN_active, scaling.wS.σ * wrk_wS[i, j, k_wrk], 0) +end + +# Write here your constructor +# NNFluxClosure() = ... insert NN here ... (make sure it is on GPU if you need it on GPU!) + +const NNC = NNFluxClosure + +##### +##### Abstract Smagorinsky functionality +##### + +# Horizontal fluxes are zero! +@inline viscous_flux_wz( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_ux( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_uy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline diffusive_flux_x(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) +@inline diffusive_flux_y(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) + +# Viscous fluxes are zero (for now) +@inline viscous_flux_uz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) + +# The only function extended by NNFluxClosure +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{1}, c, clock, fields, buoyancy) = @inbounds K.wT[i, j, k] +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{2}, c, clock, fields, buoyancy) = @inbounds K.wS[i, j, k] \ No newline at end of file From 823e8c7be123aaf93f0677eaafea490c54ad364c Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Mon, 21 Oct 2024 23:17:41 -0400 Subject: [PATCH 079/114] run double gyre with NNclosure with no Ri input kappazonelast55 --- ...nof_BBLkappazonelast55_doublegyre_model.jl | 534 ++++++++++++++++++ 1 file changed, 534 insertions(+) create mode 100644 NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl diff --git a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl new file mode 100644 index 0000000000..0fe56edfc7 --- /dev/null +++ b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl @@ -0,0 +1,534 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("NN_closure_global_nof_BBLkappazonelast55.jl") +include("xin_kai_vertical_diffusivity_local_2step.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + + +#%% +filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_NN_closure_NDE_BBLkappazonelast55_temp" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +nn_closure = NNFluxClosure(model_architecture) +base_closure = XinKaiLocalVerticalDiffusivity() +vertical_scalar_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +closure = (base_closure, nn_closure, vertical_scalar_closure) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/8days + +##### +##### Forcing and initial condition +##### +@inline T_initial(x, y, z) = 10 + 20 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10950days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); +first_index = model.diffusivity_fields[2].first_index +last_index = model.diffusivity_fields[2].last_index +wT_NN = model.diffusivity_fields[2].wT +wS_NN = model.diffusivity_fields[2].wS + +κ = model.diffusivity_fields[1].κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +@inline function get_top_buoyancy_flux(i, j, k, grid, buoyancy, T_bc, S_bc, velocities, tracers, clock) + return top_buoyancy_flux(i, j, grid, buoyancy, (; T=T_bc, S=S_bc), clock, merge(velocities, tracers)) +end + +Qb = KernelFunctionOperation{Center, Center, Nothing}(get_top_buoyancy_flux, model.grid, model.buoyancy, T.boundary_conditions.top, S.boundary_conditions.top, model.velocities, model.tracers, model.clock) +Qb = Field(Qb) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + +outputs = (; u, v, w, T, S, ρ, N², wT_NN, wS_NN, wT_base, wS_base) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; first_index, last_index, Qb), + filename = "$(FILE_DIR)/instantaneous_fields_NN_active_diagnostics", + indices = (:, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 365 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% \ No newline at end of file From 893b7bccebb94706572d465b0859e537fa85ad06 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 24 Oct 2024 00:33:58 -0400 Subject: [PATCH 080/114] add fluxes calculations and zonal average --- ...re_nof_BBLkappazonelast55_doublegyre_model.jl | 16 ++++++++++++++-- baseclosure_doublegyre_model.jl | 6 ++++-- physicalclosure_doublegyre_model.jl | 4 +++- 3 files changed, 21 insertions(+), 5 deletions(-) diff --git a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl index 0fe56edfc7..5bad8f8cda 100644 --- a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl +++ b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl @@ -212,6 +212,9 @@ wS_NN = model.diffusivity_fields[2].wS wT_base = κ * ∂z(T) wS_base = κ * ∂z(S) +wT = wT_NN + wT_base +wS = wS_NN + wS_base + @inline function get_N²(i, j, k, grid, b, C) return ∂z_b(i, j, k, grid, b, C) end @@ -242,8 +245,17 @@ Tbar_zonal = Average(T, dims=1) Sbar_zonal = Average(S, dims=1) ρbar_zonal = Average(ρ, dims=1) -outputs = (; u, v, w, T, S, ρ, N², wT_NN, wS_NN, wT_base, wS_base) -zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal) +wT_NNbar_zonal = Average(wT_NN, dims=1) +wS_NNbar_zonal = Average(wS_NN, dims=1) + +wT_basebar_zonal = Average(wT_base, dims=1) +wS_basebar_zonal = Average(wS_base, dims=1) + +wTbar_zonal = Average(wT, dims=1) +wSbar_zonal = Average(wS, dims=1) + +outputs = (; u, v, w, T, S, ρ, N², wT_NN, wS_NN, wT_base, wS_base, wT, wS) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal, wT_NNbar_zonal, wS_NNbar_zonal, wT_basebar_zonal, wS_basebar_zonal, wTbar_zonal, wSbar_zonal) ##### ##### Build checkpointer and output writer diff --git a/baseclosure_doublegyre_model.jl b/baseclosure_doublegyre_model.jl index b66cc37c4e..d6323be557 100644 --- a/baseclosure_doublegyre_model.jl +++ b/baseclosure_doublegyre_model.jl @@ -207,6 +207,8 @@ U_bt = Field(Integral(u, dims=3)); κ = model.diffusivity_fields[1].κᶜ wT_base = κ * ∂z(T) wS_base = κ * ∂z(S) +wTbar_zonal = Average(wT_base, dims=1) +wSbar_zonal = Average(wS_base, dims=1) @inline function get_N²(i, j, k, grid, b, C) return ∂z_b(i, j, k, grid, b, C) @@ -231,8 +233,8 @@ Tbar_zonal = Average(T, dims=1) Sbar_zonal = Average(S, dims=1) ρbar_zonal = Average(ρ, dims=1) -outputs = (; u, v, w, T, S, ρ, N², wT_base, wS_base) -zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal) +outputs = (; u, v, w, T, S, ρ, N², wT=wT_base, wS=wS_base) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal, wTbar_zonal, wSbar_zonal) ##### ##### Build checkpointer and output writer diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index c2738938a4..7681207e8e 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -248,9 +248,11 @@ wbar_zonal = Average(w, dims=1) Tbar_zonal = Average(T, dims=1) Sbar_zonal = Average(S, dims=1) ρbar_zonal = Average(ρ, dims=1) +wTbar_zonal = Average(wT, dims=1) +wSbar_zonal = Average(wS, dims=1) outputs = (; u, v, w, T, S, ρ, ρᶠ, ∂ρ∂z, ∂²ρ∂z², N², wT, wS) -zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal, wTbar_zonal, wSbar_zonal) ##### ##### Build checkpointer and output writer From 4a5a3e1db05714b9e026e673f50649b1ec17fb2d Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 24 Oct 2024 01:08:11 -0400 Subject: [PATCH 081/114] using centered second order in z instead of WENO --- ...osure_nof_BBLkappazonelast55_doublegyre_model.jl | 12 +++++++----- baseclosure_doublegyre_model.jl | 12 +++++++----- physicalclosure_doublegyre_model.jl | 13 +++++++------ 3 files changed, 21 insertions(+), 16 deletions(-) diff --git a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl index 5bad8f8cda..b86c022fcb 100644 --- a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl +++ b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl @@ -22,7 +22,7 @@ using ColorSchemes #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_NN_closure_NDE_BBLkappazonelast55_temp" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_zC2O_NN_closure_NDE_BBLkappazonelast55_temp" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -35,6 +35,8 @@ base_closure = XinKaiLocalVerticalDiffusivity() vertical_scalar_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) closure = (base_closure, nn_closure, vertical_scalar_closure) +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) + # number of grid points const Nx = 100 const Ny = 100 @@ -126,8 +128,8 @@ coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) model = HydrostaticFreeSurfaceModel( grid = grid, free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), - momentum_advection = WENO(order=5), - tracer_advection = WENO(order=5), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), coriolis = coriolis, closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), @@ -138,8 +140,8 @@ model = HydrostaticFreeSurfaceModel( model = HydrostaticFreeSurfaceModel( grid = grid, free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), - momentum_advection = WENO(order=5), - tracer_advection = WENO(order=5), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), coriolis = coriolis, closure = closure, diff --git a/baseclosure_doublegyre_model.jl b/baseclosure_doublegyre_model.jl index d6323be557..ff51f46d05 100644 --- a/baseclosure_doublegyre_model.jl +++ b/baseclosure_doublegyre_model.jl @@ -21,7 +21,7 @@ using SeawaterPolynomials:TEOS10 using ColorSchemes #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_zC2O_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -33,6 +33,8 @@ base_closure = XinKaiLocalVerticalDiffusivity() vertical_scalar_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) closure = (base_closure, vertical_scalar_closure) +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) + # number of grid points const Nx = 100 const Ny = 100 @@ -126,8 +128,8 @@ coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) model = HydrostaticFreeSurfaceModel( grid = grid, free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), - momentum_advection = WENO(order=5), - tracer_advection = WENO(order=5), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), coriolis = coriolis, closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), @@ -138,8 +140,8 @@ model = HydrostaticFreeSurfaceModel( model = HydrostaticFreeSurfaceModel( grid = grid, free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), - momentum_advection = WENO(order=5), - tracer_advection = WENO(order=5), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), coriolis = coriolis, closure = closure, diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index 7681207e8e..d11a9788ac 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -26,9 +26,8 @@ using SeawaterPolynomials:TEOS10 using ColorSchemes using Glob - #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_CATKEVerticalDiffusivity" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_zC2O_CATKEVerticalDiffusivity" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -47,6 +46,8 @@ convection_closure = CATKE_ocean_closure() closure = convection_closure # closure = vertical_base_closure +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) + # number of grid points const Nx = 100 const Ny = 100 @@ -139,8 +140,8 @@ coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) model = HydrostaticFreeSurfaceModel( grid = grid, free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), - momentum_advection = WENO(order=5), - tracer_advection = WENO(order=5), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), coriolis = coriolis, closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), @@ -151,8 +152,8 @@ model = HydrostaticFreeSurfaceModel( model = HydrostaticFreeSurfaceModel( grid = grid, free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), - momentum_advection = WENO(order=5), - tracer_advection = WENO(order=5), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), coriolis = coriolis, closure = closure, From 0f6bf66486a2b2f9092db22edea8fffa7ff02843 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 24 Oct 2024 01:15:31 -0400 Subject: [PATCH 082/114] remove background vertical scalar diffusivity --- NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl | 3 +-- baseclosure_doublegyre_model.jl | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl index b86c022fcb..2603f8cba3 100644 --- a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl +++ b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl @@ -32,8 +32,7 @@ model_architecture = GPU() nn_closure = NNFluxClosure(model_architecture) base_closure = XinKaiLocalVerticalDiffusivity() -vertical_scalar_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) -closure = (base_closure, nn_closure, vertical_scalar_closure) +closure = (base_closure, nn_closure) advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) diff --git a/baseclosure_doublegyre_model.jl b/baseclosure_doublegyre_model.jl index ff51f46d05..0d16106321 100644 --- a/baseclosure_doublegyre_model.jl +++ b/baseclosure_doublegyre_model.jl @@ -30,8 +30,7 @@ mkpath(FILE_DIR) model_architecture = GPU() base_closure = XinKaiLocalVerticalDiffusivity() -vertical_scalar_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) -closure = (base_closure, vertical_scalar_closure) +closure = base_closure advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) From 838522911b19ab3c265d8afb67ef910053668fd7 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 24 Oct 2024 02:03:09 -0400 Subject: [PATCH 083/114] diffusivity fields indexing for base closure --- baseclosure_doublegyre_model.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/baseclosure_doublegyre_model.jl b/baseclosure_doublegyre_model.jl index 0d16106321..34b763fa0d 100644 --- a/baseclosure_doublegyre_model.jl +++ b/baseclosure_doublegyre_model.jl @@ -205,7 +205,7 @@ T, S = model.tracers.T, model.tracers.S U_bt = Field(Integral(u, dims=3)); Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); -κ = model.diffusivity_fields[1].κᶜ +κ = model.diffusivity_fields.κᶜ wT_base = κ * ∂z(T) wS_base = κ * ∂z(S) wTbar_zonal = Average(wT_base, dims=1) From 45eef0a4bc317c18661873805a0a53c6f7afb85e Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Wed, 30 Oct 2024 16:31:27 -0400 Subject: [PATCH 084/114] change temperature restoration to 30days --- NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl | 4 ++-- baseclosure_doublegyre_model.jl | 4 ++-- physicalclosure_doublegyre_model.jl | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl index 2603f8cba3..739a94c22e 100644 --- a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl +++ b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl @@ -22,7 +22,7 @@ using ColorSchemes #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_zC2O_NN_closure_NDE_BBLkappazonelast55_temp" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zC2O_NN_closure_NDE_BBLkappazonelast55_temp" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -71,7 +71,7 @@ const S_mid = (S_north + S_south) / 2 const τ₀ = 1e-4 const μ_drag = 1/30days -const μ_T = 1/8days +const μ_T = 1/30days ##### ##### Forcing and initial condition diff --git a/baseclosure_doublegyre_model.jl b/baseclosure_doublegyre_model.jl index 34b763fa0d..33204c9529 100644 --- a/baseclosure_doublegyre_model.jl +++ b/baseclosure_doublegyre_model.jl @@ -21,7 +21,7 @@ using SeawaterPolynomials:TEOS10 using ColorSchemes #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_zC2O_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zC2O_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -69,7 +69,7 @@ const S_mid = (S_north + S_south) / 2 const τ₀ = 1e-4 const μ_drag = 1/30days -const μ_T = 1/8days +const μ_T = 1/30days ##### ##### Forcing and initial condition diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index d11a9788ac..d79034e061 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -27,7 +27,7 @@ using ColorSchemes using Glob #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_zC2O_CATKEVerticalDiffusivity" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zC2O_CATKEVerticalDiffusivity" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -83,7 +83,7 @@ const S_mid = (S_north + S_south) / 2 const τ₀ = 1e-4 const μ_drag = 1/30days -const μ_T = 1/8days +const μ_T = 1/30days ##### ##### Forcing and initial condition From e29bfd9bb5e8c39ce784995ccad66597ff1d6232 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Wed, 30 Oct 2024 16:33:29 -0400 Subject: [PATCH 085/114] change vertical advection scheme to WENO5 --- NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl | 6 ++++-- baseclosure_doublegyre_model.jl | 5 +++-- physicalclosure_doublegyre_model.jl | 5 +++-- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl index 739a94c22e..778599c4e7 100644 --- a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl +++ b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl @@ -22,7 +22,8 @@ using ColorSchemes #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zC2O_NN_closure_NDE_BBLkappazonelast55_temp" +# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zC2O_NN_closure_NDE_BBLkappazonelast55_temp" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_NN_closure_NDE_BBLkappazonelast55_temp" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -34,7 +35,8 @@ nn_closure = NNFluxClosure(model_architecture) base_closure = XinKaiLocalVerticalDiffusivity() closure = (base_closure, nn_closure) -advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) +# advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) # number of grid points const Nx = 100 diff --git a/baseclosure_doublegyre_model.jl b/baseclosure_doublegyre_model.jl index 33204c9529..51cff7e579 100644 --- a/baseclosure_doublegyre_model.jl +++ b/baseclosure_doublegyre_model.jl @@ -21,7 +21,8 @@ using SeawaterPolynomials:TEOS10 using ColorSchemes #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zC2O_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" +# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zC2O_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -32,7 +33,7 @@ model_architecture = GPU() base_closure = XinKaiLocalVerticalDiffusivity() closure = base_closure -advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) # number of grid points const Nx = 100 diff --git a/physicalclosure_doublegyre_model.jl b/physicalclosure_doublegyre_model.jl index d79034e061..c11ec722de 100644 --- a/physicalclosure_doublegyre_model.jl +++ b/physicalclosure_doublegyre_model.jl @@ -27,7 +27,8 @@ using ColorSchemes using Glob #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zC2O_CATKEVerticalDiffusivity" +# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zC2O_CATKEVerticalDiffusivity" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_CATKEVerticalDiffusivity" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -46,7 +47,7 @@ convection_closure = CATKE_ocean_closure() closure = convection_closure # closure = vertical_base_closure -advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) # number of grid points const Nx = 100 From 0fba1275bf6534246c03937bdf4b51439d8f7986 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Wed, 30 Oct 2024 16:53:04 -0400 Subject: [PATCH 086/114] update neural network to new weights --- NN_closure_global_nof_BBLkappazonelast55.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NN_closure_global_nof_BBLkappazonelast55.jl b/NN_closure_global_nof_BBLkappazonelast55.jl index 5602b767f9..f7452dae84 100644 --- a/NN_closure_global_nof_BBLkappazonelast55.jl +++ b/NN_closure_global_nof_BBLkappazonelast55.jl @@ -69,7 +69,7 @@ Adapt.adapt_structure(to, nn :: NN) = function NNFluxClosure(arch) dev = ifelse(arch == GPU(), gpu_device(), cpu_device()) - nn_path = "./NDE_FC_Qb_nof_BBLkappazonelast55_trainFC23new_scalingtrain53new_2layer_128_relu_20seed_2Pr_model_temp.jld2" + nn_path = "./NDE_FC_Qb_nof_BBLkappazonelast55_trainFC23new_scalingtrain53new_2layer_128_relu_20seed_2Pr_model.jld2" ps, sts, scaling_params, wT_model, wS_model = jldopen(nn_path, "r") do file ps = file["u"] |> dev |> f64 From 73afa68a3d3d3c44f9f8dc3d6cc391869d1daf5c Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Wed, 30 Oct 2024 16:56:36 -0400 Subject: [PATCH 087/114] fix file name change --- NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl index 778599c4e7..0338a14cdb 100644 --- a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl +++ b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl @@ -22,8 +22,8 @@ using ColorSchemes #%% -# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zC2O_NN_closure_NDE_BBLkappazonelast55_temp" -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_NN_closure_NDE_BBLkappazonelast55_temp" +# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zC2O_NN_closure_NDE_BBLkappazonelast55" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_NN_closure_NDE_BBLkappazonelast55" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) From 1fab19f133ed3953edc1cbec5a0e551cbdaa0a3f Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 31 Oct 2024 01:06:01 -0400 Subject: [PATCH 088/114] run double gyre with centered second order and wall restoration --- ...last55_doublegyre_model_wallrestoration.jl | 554 +++++++++++++++++ ...losure_doublegyre_model_wallrestoration.jl | 586 ++++++++++++++++++ ...losure_doublegyre_model_wallrestoration.jl | 545 ++++++++++++++++ 3 files changed, 1685 insertions(+) create mode 100644 NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl create mode 100644 baseclosure_doublegyre_model_wallrestoration.jl create mode 100644 physicalclosure_doublegyre_model_wallrestoration.jl diff --git a/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl b/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl new file mode 100644 index 0000000000..e6b2db482a --- /dev/null +++ b/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl @@ -0,0 +1,554 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("NN_closure_global_nof_BBLkappazonelast55.jl") +include("xin_kai_vertical_diffusivity_local_2step.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + + +#%% +filename = "doublegyre_30Cwarmflushbottom10_relaxation_wallrestoration_30days_zC2O_NN_closure_NDE_BBLkappazonelast55" +# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_NN_closure_NDE_BBLkappazonelast55" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +nn_closure = NNFluxClosure(model_architecture) +base_closure = XinKaiLocalVerticalDiffusivity() +closure = (base_closure, nn_closure) + +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) +# advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/30days + +##### +##### Forcing and initial condition +##### +@inline T_initial(x, y, z) = 10 + 20 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) + +@inline T_north_ref(z) = ifelse(z >= -500, 0, -5 + 5 * (1 + (z + 500) / (Lz - 500))) +@inline north_T_flux(x, z, t, T) = μ_T * δy * (T - T_north_ref(z)) +north_T_flux_bc = FluxBoundaryCondition(north_T_flux; field_dependencies=:T) + +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc, north = north_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10950days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); +first_index = model.diffusivity_fields[2].first_index +last_index = model.diffusivity_fields[2].last_index +wT_NN = model.diffusivity_fields[2].wT +wS_NN = model.diffusivity_fields[2].wS + +κ = model.diffusivity_fields[1].κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) + +wT = wT_NN + wT_base +wS = wS_NN + wS_base + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +@inline function get_top_buoyancy_flux(i, j, k, grid, buoyancy, T_bc, S_bc, velocities, tracers, clock) + return top_buoyancy_flux(i, j, grid, buoyancy, (; T=T_bc, S=S_bc), clock, merge(velocities, tracers)) +end + +Qb = KernelFunctionOperation{Center, Center, Nothing}(get_top_buoyancy_flux, model.grid, model.buoyancy, T.boundary_conditions.top, S.boundary_conditions.top, model.velocities, model.tracers, model.clock) +Qb = Field(Qb) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + +wT_NNbar_zonal = Average(wT_NN, dims=1) +wS_NNbar_zonal = Average(wS_NN, dims=1) + +wT_basebar_zonal = Average(wT_base, dims=1) +wS_basebar_zonal = Average(wS_base, dims=1) + +wTbar_zonal = Average(wT, dims=1) +wSbar_zonal = Average(wS, dims=1) + +outputs = (; u, v, w, T, S, ρ, N², wT_NN, wS_NN, wT_base, wS_base, wT, wS) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal, wT_NNbar_zonal, wS_NNbar_zonal, wT_basebar_zonal, wS_basebar_zonal, wTbar_zonal, wSbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; first_index, last_index, Qb), + filename = "$(FILE_DIR)/instantaneous_fields_NN_active_diagnostics", + indices = (:, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 365 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% \ No newline at end of file diff --git a/baseclosure_doublegyre_model_wallrestoration.jl b/baseclosure_doublegyre_model_wallrestoration.jl new file mode 100644 index 0000000000..d24375cc01 --- /dev/null +++ b/baseclosure_doublegyre_model_wallrestoration.jl @@ -0,0 +1,586 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("xin_kai_vertical_diffusivity_local_2step.jl") +# include("xin_kai_vertical_diffusivity_2Pr.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + +#%% +filename = "doublegyre_30Cwarmflushbottom10_relaxation_wallrestoration_30days_zC2O_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" +# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +base_closure = XinKaiLocalVerticalDiffusivity() +closure = base_closure + +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) +# advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +const δy = Ly / Ny + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/30days + +##### +##### Forcing and initial condition +##### +# @inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) +# @inline T_initial(x, y, z) = (T_north + T_south / 2) + 5 * (1 + z / Lz) +@inline T_initial(x, y, z) = 10 + 20 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) + +@inline T_north_ref(z) = ifelse(z >= -500, 0, -5 + 5 * (1 + (z + 500) / (Lz - 500))) +@inline north_T_flux(x, z, t, T) = μ_T * δy * (T - T_north_ref(z)) +north_T_flux_bc = FluxBoundaryCondition(north_T_flux; field_dependencies=:T) + +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc, north = north_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10950days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); + +κ = model.diffusivity_fields.κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) +wTbar_zonal = Average(wT_base, dims=1) +wSbar_zonal = Average(wS_base, dims=1) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + +outputs = (; u, v, w, T, S, ρ, N², wT=wT_base, wS=wS_base) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal, wTbar_zonal, wSbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:complete_fields] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields", + schedule = TimeInterval(1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 365 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% +# function find_min(a...) +# return minimum(minimum.([a...])) +# end + +# function find_max(a...) +# return maximum(maximum.([a...])) +# end + +# N²_xz_north_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_north.jld2", "N²") +# N²_xz_south_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_south.jld2", "N²") + +# xC = N²_xz_north_data.grid.xᶜᵃᵃ[1:N²_xz_north_data.grid.Nx] +# zf = N²_xz_north_data.grid.zᵃᵃᶠ[1:N²_xz_north_data.grid.Nz+1] + +# yloc_north = N²_xz_north_data.grid.yᵃᶠᵃ[N²_xz_north_data.indices[2][1]] +# yloc_south = N²_xz_south_data.grid.yᵃᶠᵃ[N²_xz_south_data.indices[2][1]] + +# Nt = length(N²_xz_north_data) +# times = N²_xz_north_data.times / 24 / 60^2 / 365 +# timeframes = 1:Nt + +# N²_lim = (find_min(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) - 1e-13, +# find_max(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) + 1e-13) +# #%% +# fig = Figure(size=(800, 800)) +# ax_north = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_south / 1e3)) km") +# ax_south = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_north / 1e3)) km") + +# n = Observable(2) + +# N²_north = @lift interior(N²_xz_north_data[$n], :, 1, :) +# N²_south = @lift interior(N²_xz_south_data[$n], :, 1, :) + +# colorscheme = colorschemes[:jet] + +# N²_north_surface = heatmap!(ax_north, xC, zf, N²_north, colormap=colorscheme, colorrange=N²_lim) +# N²_south_surface = heatmap!(ax_south, xC, zf, N²_south, colormap=colorscheme, colorrange=N²_lim) + +# Colorbar(fig[1:2, 2], N²_north_surface, label="N² (s⁻²)") + +# title_str = @lift "Time = $(round(times[$n], digits=2)) years" +# Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +# trim!(fig.layout) + +# @info "Recording buoyancy frequency xz slice" +# CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_buoyancy_frequency_xz_slice.mp4", 1:Nt, framerate=3, px_per_unit=2) do nn +# n[] = nn +# end + +# @info "Done!" +# #%% \ No newline at end of file diff --git a/physicalclosure_doublegyre_model_wallrestoration.jl b/physicalclosure_doublegyre_model_wallrestoration.jl new file mode 100644 index 0000000000..041299ba46 --- /dev/null +++ b/physicalclosure_doublegyre_model_wallrestoration.jl @@ -0,0 +1,545 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity +using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: CATKEVerticalDiffusivity, CATKEMixingLength, CATKEEquation + +using Oceananigans.BuoyancyModels: ∂z_b +# include("NN_closure_global.jl") +# include("xin_kai_vertical_diffusivity_local.jl") +# include("xin_kai_vertical_diffusivity_2Pr.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes +using Glob + +#%% +filename = "doublegyre_30Cwarmflushbottom10_relaxation_wallrestoration_30days_zC2O_CATKEVerticalDiffusivity" +# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_CATKEVerticalDiffusivity" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +# vertical_base_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +# convection_closure = XinKaiVerticalDiffusivity() +function CATKE_ocean_closure() + mixing_length = CATKEMixingLength(Cᵇ=0.01) + turbulent_kinetic_energy_equation = CATKEEquation(Cᵂϵ=1.0) + return CATKEVerticalDiffusivity(; mixing_length, turbulent_kinetic_energy_equation) +end +convection_closure = CATKE_ocean_closure() +closure = convection_closure +# closure = vertical_base_closure + +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) +# advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/30days + +##### +##### Forcing and initial condition +##### +# @inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) +@inline T_initial(x, y, z) = 10 + 20 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) + +@inline T_north_ref(z) = ifelse(z >= -500, 0, -5 + 5 * (1 + (z + 500) / (Lz - 500))) +@inline north_T_flux(x, z, t, T) = μ_T * δy * (T - T_north_ref(z)) +north_T_flux_bc = FluxBoundaryCondition(north_T_flux; field_dependencies=:T) + +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc, north = north_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S, :e), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S, :e), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10950days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)) +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) +compute!(ρ) + +ρᶠ = @at (Center, Center, Face) ρ +∂ρ∂z = ∂z(ρ) +∂²ρ∂z² = ∂z(∂ρ∂z) + +κc = model.diffusivity_fields.κc +wT = κc * ∂z(T) +wS = κc * ∂z(S) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) +wTbar_zonal = Average(wT, dims=1) +wSbar_zonal = Average(wS, dims=1) + +outputs = (; u, v, w, T, S, ρ, ρᶠ, ∂ρ∂z, ∂²ρ∂z², N², wT, wS) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal, wTbar_zonal, wSbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 365 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% \ No newline at end of file From 49d901fda640aa040335b3251794de430d479105 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 31 Oct 2024 12:55:16 -0400 Subject: [PATCH 089/114] fix dynamic function invocation that doesn't affect the baseclosure script somehow --- ...e_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl | 2 +- baseclosure_doublegyre_model_wallrestoration.jl | 2 +- physicalclosure_doublegyre_model_wallrestoration.jl | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl b/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl index e6b2db482a..fb81b9fe9b 100644 --- a/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl +++ b/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl @@ -104,7 +104,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), @inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) -@inline T_north_ref(z) = ifelse(z >= -500, 0, -5 + 5 * (1 + (z + 500) / (Lz - 500))) +@inline T_north_ref(z) = min(0, -5 + 5 * (1 + (z + 500) / (Lz - 500))) @inline north_T_flux(x, z, t, T) = μ_T * δy * (T - T_north_ref(z)) north_T_flux_bc = FluxBoundaryCondition(north_T_flux; field_dependencies=:T) diff --git a/baseclosure_doublegyre_model_wallrestoration.jl b/baseclosure_doublegyre_model_wallrestoration.jl index d24375cc01..b675786dec 100644 --- a/baseclosure_doublegyre_model_wallrestoration.jl +++ b/baseclosure_doublegyre_model_wallrestoration.jl @@ -106,7 +106,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), @inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) -@inline T_north_ref(z) = ifelse(z >= -500, 0, -5 + 5 * (1 + (z + 500) / (Lz - 500))) +@inline T_north_ref(z) = min(0, -5 + 5 * (1 + (z + 500) / (Lz - 500))) @inline north_T_flux(x, z, t, T) = μ_T * δy * (T - T_north_ref(z)) north_T_flux_bc = FluxBoundaryCondition(north_T_flux; field_dependencies=:T) diff --git a/physicalclosure_doublegyre_model_wallrestoration.jl b/physicalclosure_doublegyre_model_wallrestoration.jl index 041299ba46..434d248a59 100644 --- a/physicalclosure_doublegyre_model_wallrestoration.jl +++ b/physicalclosure_doublegyre_model_wallrestoration.jl @@ -117,7 +117,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), @inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) -@inline T_north_ref(z) = ifelse(z >= -500, 0, -5 + 5 * (1 + (z + 500) / (Lz - 500))) +@inline T_north_ref(z) = min(0, -5 + 5 * (1 + (z + 500) / (Lz - 500))) @inline north_T_flux(x, z, t, T) = μ_T * δy * (T - T_north_ref(z)) north_T_flux_bc = FluxBoundaryCondition(north_T_flux; field_dependencies=:T) From 534d21d0db4f5a934e5f512e97f7bed34957f73f Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 31 Oct 2024 13:12:14 -0400 Subject: [PATCH 090/114] add y resolution variable --- ...e_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl | 2 ++ physicalclosure_doublegyre_model_wallrestoration.jl | 2 ++ 2 files changed, 4 insertions(+) diff --git a/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl b/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl index fb81b9fe9b..9555398289 100644 --- a/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl +++ b/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl @@ -48,6 +48,8 @@ const Lx = 4000kilometers const Ly = 6000kilometers const Lz = Nz * Δz +const δy = Ly / Ny + grid = RectilinearGrid(model_architecture, Float64, topology = (Bounded, Bounded, Bounded), size = (Nx, Ny, Nz), diff --git a/physicalclosure_doublegyre_model_wallrestoration.jl b/physicalclosure_doublegyre_model_wallrestoration.jl index 434d248a59..e5130a0841 100644 --- a/physicalclosure_doublegyre_model_wallrestoration.jl +++ b/physicalclosure_doublegyre_model_wallrestoration.jl @@ -60,6 +60,8 @@ const Lx = 4000kilometers const Ly = 6000kilometers const Lz = Nz * Δz +const δy = Ly / Ny + grid = RectilinearGrid(model_architecture, Float64, topology = (Bounded, Bounded, Bounded), size = (Nx, Ny, Nz), From 2551c388b2ab4ff39344d435b632649eeda9a324 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Fri, 1 Nov 2024 16:57:09 -0400 Subject: [PATCH 091/114] run NN closure recording xz slices --- ...nof_BBLkappazonelast55_doublegyre_model.jl | 58 +++++++++++++++++-- ...last55_doublegyre_model_wallrestoration.jl | 50 ++++++++++++++++ 2 files changed, 104 insertions(+), 4 deletions(-) diff --git a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl index 0338a14cdb..b0c8c00eab 100644 --- a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl +++ b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl @@ -22,8 +22,8 @@ using ColorSchemes #%% -# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zC2O_NN_closure_NDE_BBLkappazonelast55" -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_NN_closure_NDE_BBLkappazonelast55" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zC2O_NN_closure_NDE_BBLkappazonelast55" +# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_NN_closure_NDE_BBLkappazonelast55" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -35,8 +35,8 @@ nn_closure = NNFluxClosure(model_architecture) base_closure = XinKaiLocalVerticalDiffusivity() closure = (base_closure, nn_closure) -# advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) -advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) +# advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) # number of grid points const Nx = 100 @@ -323,6 +323,56 @@ simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, indices = (90, :, :), schedule = TimeInterval(10days)) +simulation.output_writers[:xz_5] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_5", + indices = (:, 5, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_15] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_15", + indices = (:, 15, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_25] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_25", + indices = (:, 25, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_35] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_35", + indices = (:, 35, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_45] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_45", + indices = (:, 45, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_55] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_55", + indices = (:, 55, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_65] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_65", + indices = (:, 65, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_75] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_75", + indices = (:, 75, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_85] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_85", + indices = (:, 85, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_95] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_95", + indices = (:, 95, :), + schedule = TimeInterval(10days)) + simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, filename = "$(FILE_DIR)/averaged_fields_zonal", schedule = TimeInterval(10days)) diff --git a/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl b/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl index 9555398289..0e09921124 100644 --- a/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl +++ b/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl @@ -330,6 +330,56 @@ simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, indices = (90, :, :), schedule = TimeInterval(10days)) +simulation.output_writers[:xz_5] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_5", + indices = (:, 5, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_15] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_15", + indices = (:, 15, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_25] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_25", + indices = (:, 25, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_35] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_35", + indices = (:, 35, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_45] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_45", + indices = (:, 45, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_55] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_55", + indices = (:, 55, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_65] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_65", + indices = (:, 65, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_75] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_75", + indices = (:, 75, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_85] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_85", + indices = (:, 85, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_95] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_95", + indices = (:, 95, :), + schedule = TimeInterval(10days)) + simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, filename = "$(FILE_DIR)/averaged_fields_zonal", schedule = TimeInterval(10days)) From ff272c4e45137f06bd8bb81b078940c335df9537 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Fri, 1 Nov 2024 17:20:03 -0400 Subject: [PATCH 092/114] recording and plotting xz yz slices and fluxes --- ...nof_BBLkappazonelast55_doublegyre_model.jl | 708 +++++++++++++++++- ...last55_doublegyre_model_wallrestoration.jl | 707 ++++++++++++++++- 2 files changed, 1373 insertions(+), 42 deletions(-) diff --git a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl index b0c8c00eab..89a98b92a1 100644 --- a/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl +++ b/NNclosure_nof_BBLkappazonelast55_doublegyre_model.jl @@ -570,30 +570,696 @@ zlims!(axv, (-Lz, 0)) @info "Recording 3D fields" CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn - @info "Recording frame $nn" n[] = nn end -@info "Done!" #%% -# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") - -# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] -# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] - -# Nt = length(Ψ_data) -# times = Ψ_data.times / 24 / 60^2 / 365 -# #%% -# timeframe = Nt -# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 -# clim = maximum(abs, Ψ_frame) + 1e-13 -# N_levels = 16 -# levels = range(-clim, stop=clim, length=N_levels) -# fig = Figure(size=(800, 800)) -# ax = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") -# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) -# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") -# tightlimits!(ax) -# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +@info "Recording T fields and fluxes in yz" + +fieldname = "T" +fluxname = "wT_NN" +field_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fieldname, backend=OnDisk()) +field_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fieldname, backend=OnDisk()) +field_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fieldname, backend=OnDisk()) +field_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fieldname, backend=OnDisk()) +field_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fieldname, backend=OnDisk()) +field_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fieldname, backend=OnDisk()) +field_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fieldname, backend=OnDisk()) +field_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fieldname, backend=OnDisk()) +field_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fieldname, backend=OnDisk()) +field_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fluxname, backend=OnDisk()) +flux_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fluxname, backend=OnDisk()) +flux_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fluxname, backend=OnDisk()) +flux_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fluxname, backend=OnDisk()) +flux_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fluxname, backend=OnDisk()) +flux_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fluxname, backend=OnDisk()) +flux_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fluxname, backend=OnDisk()) +flux_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fluxname, backend=OnDisk()) +flux_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fluxname, backend=OnDisk()) +flux_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fluxname, backend=OnDisk()) + +xC = field_NN_data_00.grid.xᶜᵃᵃ[1:field_NN_data_00.grid.Nx] +yC = field_NN_data_00.grid.yᵃᶜᵃ[1:field_NN_data_00.grid.Ny] +zC = field_NN_data_00.grid.zᵃᵃᶜ[1:field_NN_data_00.grid.Nz] +zF = field_NN_data_00.grid.zᵃᵃᶠ[1:field_NN_data_00.grid.Nz+1] + +Nt = length(field_NN_data_90) +times = field_NN_data_00.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_00 = CairoMakie.Axis(fig[1, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_00.grid.xᶜᵃᵃ[field_NN_data_00.indices[1][1]] / 1000) km") +axfield_10 = CairoMakie.Axis(fig[1, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_10.grid.xᶜᵃᵃ[field_NN_data_10.indices[1][1]] / 1000) km") +axfield_20 = CairoMakie.Axis(fig[2, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_20.grid.xᶜᵃᵃ[field_NN_data_20.indices[1][1]] / 1000) km") +axfield_30 = CairoMakie.Axis(fig[2, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_30.grid.xᶜᵃᵃ[field_NN_data_30.indices[1][1]] / 1000) km") +axfield_40 = CairoMakie.Axis(fig[3, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_40.grid.xᶜᵃᵃ[field_NN_data_40.indices[1][1]] / 1000) km") +axfield_50 = CairoMakie.Axis(fig[3, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_50.grid.xᶜᵃᵃ[field_NN_data_50.indices[1][1]] / 1000) km") +axfield_60 = CairoMakie.Axis(fig[4, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_60.grid.xᶜᵃᵃ[field_NN_data_60.indices[1][1]] / 1000) km") +axfield_70 = CairoMakie.Axis(fig[4, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_70.grid.xᶜᵃᵃ[field_NN_data_70.indices[1][1]] / 1000) km") +axfield_80 = CairoMakie.Axis(fig[5, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_80.grid.xᶜᵃᵃ[field_NN_data_80.indices[1][1]] / 1000) km") +axfield_90 = CairoMakie.Axis(fig[5, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_90.grid.xᶜᵃᵃ[field_NN_data_90.indices[1][1]] / 1000) km") + +axflux_00 = CairoMakie.Axis(fig[1, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_00.grid.xᶜᵃᵃ[flux_NN_data_00.indices[1][1]] / 1000) km") +axflux_10 = CairoMakie.Axis(fig[1, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_10.grid.xᶜᵃᵃ[flux_NN_data_10.indices[1][1]] / 1000) km") +axflux_20 = CairoMakie.Axis(fig[2, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_20.grid.xᶜᵃᵃ[flux_NN_data_20.indices[1][1]] / 1000) km") +axflux_30 = CairoMakie.Axis(fig[2, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_30.grid.xᶜᵃᵃ[flux_NN_data_30.indices[1][1]] / 1000) km") +axflux_40 = CairoMakie.Axis(fig[3, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_40.grid.xᶜᵃᵃ[flux_NN_data_40.indices[1][1]] / 1000) km") +axflux_50 = CairoMakie.Axis(fig[3, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_50.grid.xᶜᵃᵃ[flux_NN_data_50.indices[1][1]] / 1000) km") +axflux_60 = CairoMakie.Axis(fig[4, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_60.grid.xᶜᵃᵃ[flux_NN_data_60.indices[1][1]] / 1000) km") +axflux_70 = CairoMakie.Axis(fig[4, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_70.grid.xᶜᵃᵃ[flux_NN_data_70.indices[1][1]] / 1000) km") +axflux_80 = CairoMakie.Axis(fig[5, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_80.grid.xᶜᵃᵃ[flux_NN_data_80.indices[1][1]] / 1000) km") +axflux_90 = CairoMakie.Axis(fig[5, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_90.grid.xᶜᵃᵃ[flux_NN_data_90.indices[1][1]] / 1000) km") + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lim = (find_min(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (find_min(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (-maximum(abs, flux_lim), maximum(abs, flux_lim)) + +NN_00ₙ = @lift interior(field_NN_data_00[$n], 1, :, zC_indices) +NN_10ₙ = @lift interior(field_NN_data_10[$n], 1, :, zC_indices) +NN_20ₙ = @lift interior(field_NN_data_20[$n], 1, :, zC_indices) +NN_30ₙ = @lift interior(field_NN_data_30[$n], 1, :, zC_indices) +NN_40ₙ = @lift interior(field_NN_data_40[$n], 1, :, zC_indices) +NN_50ₙ = @lift interior(field_NN_data_50[$n], 1, :, zC_indices) +NN_60ₙ = @lift interior(field_NN_data_60[$n], 1, :, zC_indices) +NN_70ₙ = @lift interior(field_NN_data_70[$n], 1, :, zC_indices) +NN_80ₙ = @lift interior(field_NN_data_80[$n], 1, :, zC_indices) +NN_90ₙ = @lift interior(field_NN_data_90[$n], 1, :, zC_indices) + +flux_00ₙ = @lift interior(flux_NN_data_00[$n], 1, :, zF_indices) +flux_10ₙ = @lift interior(flux_NN_data_10[$n], 1, :, zF_indices) +flux_20ₙ = @lift interior(flux_NN_data_20[$n], 1, :, zF_indices) +flux_30ₙ = @lift interior(flux_NN_data_30[$n], 1, :, zF_indices) +flux_40ₙ = @lift interior(flux_NN_data_40[$n], 1, :, zF_indices) +flux_50ₙ = @lift interior(flux_NN_data_50[$n], 1, :, zF_indices) +flux_60ₙ = @lift interior(flux_NN_data_60[$n], 1, :, zF_indices) +flux_70ₙ = @lift interior(flux_NN_data_70[$n], 1, :, zF_indices) +flux_80ₙ = @lift interior(flux_NN_data_80[$n], 1, :, zF_indices) +flux_90ₙ = @lift interior(flux_NN_data_90[$n], 1, :, zF_indices) + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_00_surface = heatmap!(axfield_00, yC, zC[zC_indices], NN_00ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_10_surface = heatmap!(axfield_10, yC, zC[zC_indices], NN_10ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_20_surface = heatmap!(axfield_20, yC, zC[zC_indices], NN_20ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_30_surface = heatmap!(axfield_30, yC, zC[zC_indices], NN_30ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_40_surface = heatmap!(axfield_40, yC, zC[zC_indices], NN_40ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_50_surface = heatmap!(axfield_50, yC, zC[zC_indices], NN_50ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_60_surface = heatmap!(axfield_60, yC, zC[zC_indices], NN_60ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_70_surface = heatmap!(axfield_70, yC, zC[zC_indices], NN_70ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_80_surface = heatmap!(axfield_80, yC, zC[zC_indices], NN_80ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_90_surface = heatmap!(axfield_90, yC, zC[zC_indices], NN_90ₙ, colormap=colorscheme_field, colorrange=field_lim) + +flux_00_surface = heatmap!(axflux_00, yC, zC[zF_indices], flux_00ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_10_surface = heatmap!(axflux_10, yC, zC[zF_indices], flux_10ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_20_surface = heatmap!(axflux_20, yC, zC[zF_indices], flux_20ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_30_surface = heatmap!(axflux_30, yC, zC[zF_indices], flux_30ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_40_surface = heatmap!(axflux_40, yC, zC[zF_indices], flux_40ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_50_surface = heatmap!(axflux_50, yC, zC[zF_indices], flux_50ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_60_surface = heatmap!(axflux_60, yC, zC[zF_indices], flux_60ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_70_surface = heatmap!(axflux_70, yC, zC[zF_indices], flux_70ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_80_surface = heatmap!(axflux_80, yC, zC[zF_indices], flux_80ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_90_surface = heatmap!(axflux_90, yC, zC[zF_indices], flux_90ₙ, colormap=colorscheme_flux, colorrange=flux_lim) + +Colorbar(fig[1:5, 2], field_00_surface, label="Field") +Colorbar(fig[1:5, 4], flux_00_surface, label="NN Flux") +Colorbar(fig[1:5, 6], field_00_surface, label="Field") +Colorbar(fig[1:5, 8], flux_00_surface, label="NN Flux") + +xlims!(axfield_00, minimum(yC), maximum(yC)) +xlims!(axfield_10, minimum(yC), maximum(yC)) +xlims!(axfield_20, minimum(yC), maximum(yC)) +xlims!(axfield_30, minimum(yC), maximum(yC)) +xlims!(axfield_40, minimum(yC), maximum(yC)) +xlims!(axfield_50, minimum(yC), maximum(yC)) +xlims!(axfield_60, minimum(yC), maximum(yC)) +xlims!(axfield_70, minimum(yC), maximum(yC)) +xlims!(axfield_80, minimum(yC), maximum(yC)) +xlims!(axfield_90, minimum(yC), maximum(yC)) + +ylims!(axfield_00, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_10, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_20, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_30, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_40, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_50, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_60, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_70, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_80, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_90, minimum(zC[zC_indices]), maximum(zC[zC_indices])) + +xlims!(axflux_00, minimum(yC), maximum(yC)) +xlims!(axflux_10, minimum(yC), maximum(yC)) +xlims!(axflux_20, minimum(yC), maximum(yC)) +xlims!(axflux_30, minimum(yC), maximum(yC)) +xlims!(axflux_40, minimum(yC), maximum(yC)) +xlims!(axflux_50, minimum(yC), maximum(yC)) +xlims!(axflux_60, minimum(yC), maximum(yC)) +xlims!(axflux_70, minimum(yC), maximum(yC)) +xlims!(axflux_80, minimum(yC), maximum(yC)) +xlims!(axflux_90, minimum(yC), maximum(yC)) + +ylims!(axflux_00, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_10, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_20, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_30, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_40, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_50, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_60, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_70, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_80, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_90, minimum(zF[zF_indices]), maximum(zF[zF_indices])) + +title_str = @lift "Temperature (°C), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +# save("./Output/compare_3D_instantaneous_fields_slices_NNclosure_fluxes.png", fig) # display(fig) +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_yzslices_fluxes_T.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end +#%% +@info "Recording S fields and fluxes in yz" + +fieldname = "S" +fluxname = "wS_NN" +field_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fieldname, backend=OnDisk()) +field_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fieldname, backend=OnDisk()) +field_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fieldname, backend=OnDisk()) +field_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fieldname, backend=OnDisk()) +field_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fieldname, backend=OnDisk()) +field_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fieldname, backend=OnDisk()) +field_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fieldname, backend=OnDisk()) +field_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fieldname, backend=OnDisk()) +field_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fieldname, backend=OnDisk()) +field_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fluxname, backend=OnDisk()) +flux_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fluxname, backend=OnDisk()) +flux_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fluxname, backend=OnDisk()) +flux_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fluxname, backend=OnDisk()) +flux_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fluxname, backend=OnDisk()) +flux_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fluxname, backend=OnDisk()) +flux_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fluxname, backend=OnDisk()) +flux_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fluxname, backend=OnDisk()) +flux_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fluxname, backend=OnDisk()) +flux_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fluxname, backend=OnDisk()) + +xC = field_NN_data_00.grid.xᶜᵃᵃ[1:field_NN_data_00.grid.Nx] +yC = field_NN_data_00.grid.yᵃᶜᵃ[1:field_NN_data_00.grid.Ny] +zC = field_NN_data_00.grid.zᵃᵃᶜ[1:field_NN_data_00.grid.Nz] +zF = field_NN_data_00.grid.zᵃᵃᶠ[1:field_NN_data_00.grid.Nz+1] + +Nt = length(field_NN_data_90) +times = field_NN_data_00.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_00 = CairoMakie.Axis(fig[1, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_00.grid.xᶜᵃᵃ[field_NN_data_00.indices[1][1]] / 1000) km") +axfield_10 = CairoMakie.Axis(fig[1, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_10.grid.xᶜᵃᵃ[field_NN_data_10.indices[1][1]] / 1000) km") +axfield_20 = CairoMakie.Axis(fig[2, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_20.grid.xᶜᵃᵃ[field_NN_data_20.indices[1][1]] / 1000) km") +axfield_30 = CairoMakie.Axis(fig[2, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_30.grid.xᶜᵃᵃ[field_NN_data_30.indices[1][1]] / 1000) km") +axfield_40 = CairoMakie.Axis(fig[3, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_40.grid.xᶜᵃᵃ[field_NN_data_40.indices[1][1]] / 1000) km") +axfield_50 = CairoMakie.Axis(fig[3, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_50.grid.xᶜᵃᵃ[field_NN_data_50.indices[1][1]] / 1000) km") +axfield_60 = CairoMakie.Axis(fig[4, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_60.grid.xᶜᵃᵃ[field_NN_data_60.indices[1][1]] / 1000) km") +axfield_70 = CairoMakie.Axis(fig[4, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_70.grid.xᶜᵃᵃ[field_NN_data_70.indices[1][1]] / 1000) km") +axfield_80 = CairoMakie.Axis(fig[5, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_80.grid.xᶜᵃᵃ[field_NN_data_80.indices[1][1]] / 1000) km") +axfield_90 = CairoMakie.Axis(fig[5, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_90.grid.xᶜᵃᵃ[field_NN_data_90.indices[1][1]] / 1000) km") + +axflux_00 = CairoMakie.Axis(fig[1, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_00.grid.xᶜᵃᵃ[flux_NN_data_00.indices[1][1]] / 1000) km") +axflux_10 = CairoMakie.Axis(fig[1, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_10.grid.xᶜᵃᵃ[flux_NN_data_10.indices[1][1]] / 1000) km") +axflux_20 = CairoMakie.Axis(fig[2, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_20.grid.xᶜᵃᵃ[flux_NN_data_20.indices[1][1]] / 1000) km") +axflux_30 = CairoMakie.Axis(fig[2, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_30.grid.xᶜᵃᵃ[flux_NN_data_30.indices[1][1]] / 1000) km") +axflux_40 = CairoMakie.Axis(fig[3, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_40.grid.xᶜᵃᵃ[flux_NN_data_40.indices[1][1]] / 1000) km") +axflux_50 = CairoMakie.Axis(fig[3, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_50.grid.xᶜᵃᵃ[flux_NN_data_50.indices[1][1]] / 1000) km") +axflux_60 = CairoMakie.Axis(fig[4, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_60.grid.xᶜᵃᵃ[flux_NN_data_60.indices[1][1]] / 1000) km") +axflux_70 = CairoMakie.Axis(fig[4, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_70.grid.xᶜᵃᵃ[flux_NN_data_70.indices[1][1]] / 1000) km") +axflux_80 = CairoMakie.Axis(fig[5, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_80.grid.xᶜᵃᵃ[flux_NN_data_80.indices[1][1]] / 1000) km") +axflux_90 = CairoMakie.Axis(fig[5, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_90.grid.xᶜᵃᵃ[flux_NN_data_90.indices[1][1]] / 1000) km") + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lim = (find_min(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (find_min(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (-maximum(abs, flux_lim), maximum(abs, flux_lim)) + +NN_00ₙ = @lift interior(field_NN_data_00[$n], 1, :, zC_indices) +NN_10ₙ = @lift interior(field_NN_data_10[$n], 1, :, zC_indices) +NN_20ₙ = @lift interior(field_NN_data_20[$n], 1, :, zC_indices) +NN_30ₙ = @lift interior(field_NN_data_30[$n], 1, :, zC_indices) +NN_40ₙ = @lift interior(field_NN_data_40[$n], 1, :, zC_indices) +NN_50ₙ = @lift interior(field_NN_data_50[$n], 1, :, zC_indices) +NN_60ₙ = @lift interior(field_NN_data_60[$n], 1, :, zC_indices) +NN_70ₙ = @lift interior(field_NN_data_70[$n], 1, :, zC_indices) +NN_80ₙ = @lift interior(field_NN_data_80[$n], 1, :, zC_indices) +NN_90ₙ = @lift interior(field_NN_data_90[$n], 1, :, zC_indices) + +flux_00ₙ = @lift interior(flux_NN_data_00[$n], 1, :, zF_indices) +flux_10ₙ = @lift interior(flux_NN_data_10[$n], 1, :, zF_indices) +flux_20ₙ = @lift interior(flux_NN_data_20[$n], 1, :, zF_indices) +flux_30ₙ = @lift interior(flux_NN_data_30[$n], 1, :, zF_indices) +flux_40ₙ = @lift interior(flux_NN_data_40[$n], 1, :, zF_indices) +flux_50ₙ = @lift interior(flux_NN_data_50[$n], 1, :, zF_indices) +flux_60ₙ = @lift interior(flux_NN_data_60[$n], 1, :, zF_indices) +flux_70ₙ = @lift interior(flux_NN_data_70[$n], 1, :, zF_indices) +flux_80ₙ = @lift interior(flux_NN_data_80[$n], 1, :, zF_indices) +flux_90ₙ = @lift interior(flux_NN_data_90[$n], 1, :, zF_indices) + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_00_surface = heatmap!(axfield_00, yC, zC[zC_indices], NN_00ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_10_surface = heatmap!(axfield_10, yC, zC[zC_indices], NN_10ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_20_surface = heatmap!(axfield_20, yC, zC[zC_indices], NN_20ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_30_surface = heatmap!(axfield_30, yC, zC[zC_indices], NN_30ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_40_surface = heatmap!(axfield_40, yC, zC[zC_indices], NN_40ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_50_surface = heatmap!(axfield_50, yC, zC[zC_indices], NN_50ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_60_surface = heatmap!(axfield_60, yC, zC[zC_indices], NN_60ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_70_surface = heatmap!(axfield_70, yC, zC[zC_indices], NN_70ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_80_surface = heatmap!(axfield_80, yC, zC[zC_indices], NN_80ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_90_surface = heatmap!(axfield_90, yC, zC[zC_indices], NN_90ₙ, colormap=colorscheme_field, colorrange=field_lim) + +flux_00_surface = heatmap!(axflux_00, yC, zC[zF_indices], flux_00ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_10_surface = heatmap!(axflux_10, yC, zC[zF_indices], flux_10ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_20_surface = heatmap!(axflux_20, yC, zC[zF_indices], flux_20ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_30_surface = heatmap!(axflux_30, yC, zC[zF_indices], flux_30ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_40_surface = heatmap!(axflux_40, yC, zC[zF_indices], flux_40ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_50_surface = heatmap!(axflux_50, yC, zC[zF_indices], flux_50ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_60_surface = heatmap!(axflux_60, yC, zC[zF_indices], flux_60ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_70_surface = heatmap!(axflux_70, yC, zC[zF_indices], flux_70ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_80_surface = heatmap!(axflux_80, yC, zC[zF_indices], flux_80ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_90_surface = heatmap!(axflux_90, yC, zC[zF_indices], flux_90ₙ, colormap=colorscheme_flux, colorrange=flux_lim) + +Colorbar(fig[1:5, 2], field_00_surface, label="Field") +Colorbar(fig[1:5, 4], flux_00_surface, label="NN Flux") +Colorbar(fig[1:5, 6], field_00_surface, label="Field") +Colorbar(fig[1:5, 8], flux_00_surface, label="NN Flux") + +xlims!(axfield_00, minimum(yC), maximum(yC)) +xlims!(axfield_10, minimum(yC), maximum(yC)) +xlims!(axfield_20, minimum(yC), maximum(yC)) +xlims!(axfield_30, minimum(yC), maximum(yC)) +xlims!(axfield_40, minimum(yC), maximum(yC)) +xlims!(axfield_50, minimum(yC), maximum(yC)) +xlims!(axfield_60, minimum(yC), maximum(yC)) +xlims!(axfield_70, minimum(yC), maximum(yC)) +xlims!(axfield_80, minimum(yC), maximum(yC)) +xlims!(axfield_90, minimum(yC), maximum(yC)) + +ylims!(axfield_00, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_10, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_20, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_30, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_40, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_50, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_60, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_70, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_80, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_90, minimum(zC[zC_indices]), maximum(zC[zC_indices])) + +xlims!(axflux_00, minimum(yC), maximum(yC)) +xlims!(axflux_10, minimum(yC), maximum(yC)) +xlims!(axflux_20, minimum(yC), maximum(yC)) +xlims!(axflux_30, minimum(yC), maximum(yC)) +xlims!(axflux_40, minimum(yC), maximum(yC)) +xlims!(axflux_50, minimum(yC), maximum(yC)) +xlims!(axflux_60, minimum(yC), maximum(yC)) +xlims!(axflux_70, minimum(yC), maximum(yC)) +xlims!(axflux_80, minimum(yC), maximum(yC)) +xlims!(axflux_90, minimum(yC), maximum(yC)) + +ylims!(axflux_00, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_10, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_20, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_30, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_40, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_50, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_60, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_70, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_80, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_90, minimum(zF[zF_indices]), maximum(zF[zF_indices])) + +title_str = @lift "Salinity (psu), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_yzslices_fluxes_S.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +#%% +@info "Recording T fields and fluxes in xz" + +fieldname = "T" +fluxname = "wT_NN" + +field_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fieldname, backend=OnDisk()) +field_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fieldname, backend=OnDisk()) +field_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fieldname, backend=OnDisk()) +field_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fieldname, backend=OnDisk()) +field_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fieldname, backend=OnDisk()) +field_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fieldname, backend=OnDisk()) +field_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fieldname, backend=OnDisk()) +field_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fieldname, backend=OnDisk()) +field_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fieldname, backend=OnDisk()) +field_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fluxname, backend=OnDisk()) +flux_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fluxname, backend=OnDisk()) +flux_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fluxname, backend=OnDisk()) +flux_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fluxname, backend=OnDisk()) +flux_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fluxname, backend=OnDisk()) +flux_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fluxname, backend=OnDisk()) +flux_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fluxname, backend=OnDisk()) +flux_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fluxname, backend=OnDisk()) +flux_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fluxname, backend=OnDisk()) +flux_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fluxname, backend=OnDisk()) + +field_datas = [field_NN_data_5, field_NN_data_15, field_NN_data_25, field_NN_data_35, field_NN_data_45, field_NN_data_55, field_NN_data_65, field_NN_data_75, field_NN_data_85, field_NN_data_95] +flux_datas = [flux_NN_data_5, flux_NN_data_15, flux_NN_data_25, flux_NN_data_35, flux_NN_data_45, flux_NN_data_55, flux_NN_data_65, flux_NN_data_75, flux_NN_data_85, flux_NN_data_95] + +xC = field_NN_data_5.grid.xᶜᵃᵃ[1:field_NN_data_5.grid.Nx] +yC = field_NN_data_5.grid.yᵃᶜᵃ[1:field_NN_data_5.grid.Ny] +zC = field_NN_data_5.grid.zᵃᵃᶜ[1:field_NN_data_5.grid.Nz] +zF = field_NN_data_5.grid.zᵃᵃᶠ[1:field_NN_data_5.grid.Nz+1] + +Nt = length(field_NN_data_95) +times = field_NN_data_5.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_5 = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_5.grid.yᵃᶜᵃ[field_NN_data_5.indices[2][1]] / 1000) km") +axfield_15 = CairoMakie.Axis(fig[1, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_15.grid.yᵃᶜᵃ[field_NN_data_15.indices[2][1]] / 1000) km") +axfield_25 = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_25.grid.yᵃᶜᵃ[field_NN_data_25.indices[2][1]] / 1000) km") +axfield_35 = CairoMakie.Axis(fig[2, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_35.grid.yᵃᶜᵃ[field_NN_data_35.indices[2][1]] / 1000) km") +axfield_45 = CairoMakie.Axis(fig[3, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_45.grid.yᵃᶜᵃ[field_NN_data_45.indices[2][1]] / 1000) km") +axfield_55 = CairoMakie.Axis(fig[3, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_55.grid.yᵃᶜᵃ[field_NN_data_55.indices[2][1]] / 1000) km") +axfield_65 = CairoMakie.Axis(fig[4, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_65.grid.yᵃᶜᵃ[field_NN_data_65.indices[2][1]] / 1000) km") +axfield_75 = CairoMakie.Axis(fig[4, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_75.grid.yᵃᶜᵃ[field_NN_data_75.indices[2][1]] / 1000) km") +axfield_85 = CairoMakie.Axis(fig[5, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_85.grid.yᵃᶜᵃ[field_NN_data_85.indices[2][1]] / 1000) km") +axfield_95 = CairoMakie.Axis(fig[5, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_95.grid.yᵃᶜᵃ[field_NN_data_95.indices[2][1]] / 1000) km") + +axflux_5 = CairoMakie.Axis(fig[1, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_5.grid.yᵃᶜᵃ[flux_NN_data_5.indices[2][1]] / 1000) km") +axflux_15 = CairoMakie.Axis(fig[1, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_15.grid.yᵃᶜᵃ[flux_NN_data_15.indices[2][1]] / 1000) km") +axflux_25 = CairoMakie.Axis(fig[2, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_25.grid.yᵃᶜᵃ[flux_NN_data_25.indices[2][1]] / 1000) km") +axflux_35 = CairoMakie.Axis(fig[2, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_35.grid.yᵃᶜᵃ[flux_NN_data_35.indices[2][1]] / 1000) km") +axflux_45 = CairoMakie.Axis(fig[3, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_45.grid.yᵃᶜᵃ[flux_NN_data_45.indices[2][1]] / 1000) km") +axflux_55 = CairoMakie.Axis(fig[3, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_55.grid.yᵃᶜᵃ[flux_NN_data_55.indices[2][1]] / 1000) km") +axflux_65 = CairoMakie.Axis(fig[4, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_65.grid.yᵃᶜᵃ[flux_NN_data_65.indices[2][1]] / 1000) km") +axflux_75 = CairoMakie.Axis(fig[4, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_75.grid.yᵃᶜᵃ[flux_NN_data_75.indices[2][1]] / 1000) km") +axflux_85 = CairoMakie.Axis(fig[5, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_85.grid.yᵃᶜᵃ[flux_NN_data_85.indices[2][1]] / 1000) km") +axflux_95 = CairoMakie.Axis(fig[5, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_95.grid.yᵃᶜᵃ[flux_NN_data_95.indices[2][1]] / 1000) km") + +axfields = [axfield_5, axfield_15, axfield_25, axfield_35, axfield_45, axfield_55, axfield_65, axfield_75, axfield_85, axfield_95] +axfluxes = [axflux_5, axflux_15, axflux_25, axflux_35, axflux_45, axflux_55, axflux_65, axflux_75, axflux_85, axflux_95] + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lims = [(find_min(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices))) for field_data in field_datas] + +flux_lims = [(find_min(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices))) for flux_data in flux_datas] + +flux_lims = [(-maximum(abs, flux_lim), maximum(abs, flux_lim)) for flux_lim in flux_lims] + +NNₙs = [@lift interior(field_data[$n], :, 1, zC_indices) for field_data in field_datas] +fluxₙs = [@lift interior(flux_data[$n], :, 1, zF_indices) for flux_data in flux_datas] + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_5_surface = heatmap!(axfield_5, xC, zC[zC_indices], NN_5ₙ, colormap=colorscheme_field, colorrange=field_lims[1]) +field_15_surface = heatmap!(axfield_15, xC, zC[zC_indices], NN_15ₙ, colormap=colorscheme_field, colorrange=field_lims[2]) +field_25_surface = heatmap!(axfield_25, xC, zC[zC_indices], NN_25ₙ, colormap=colorscheme_field, colorrange=field_lims[3]) +field_35_surface = heatmap!(axfield_35, xC, zC[zC_indices], NN_35ₙ, colormap=colorscheme_field, colorrange=field_lims[4]) +field_45_surface = heatmap!(axfield_45, xC, zC[zC_indices], NN_45ₙ, colormap=colorscheme_field, colorrange=field_lims[5]) +field_55_surface = heatmap!(axfield_55, xC, zC[zC_indices], NN_55ₙ, colormap=colorscheme_field, colorrange=field_lims[6]) +field_65_surface = heatmap!(axfield_65, xC, zC[zC_indices], NN_65ₙ, colormap=colorscheme_field, colorrange=field_lims[7]) +field_75_surface = heatmap!(axfield_75, xC, zC[zC_indices], NN_75ₙ, colormap=colorscheme_field, colorrange=field_lims[8]) +field_85_surface = heatmap!(axfield_85, xC, zC[zC_indices], NN_85ₙ, colormap=colorscheme_field, colorrange=field_lims[9]) +field_95_surface = heatmap!(axfield_95, xC, zC[zC_indices], NN_95ₙ, colormap=colorscheme_field, colorrange=field_lims[10]) + +flux_5_surface = heatmap!(axflux_5, xC, zC[zF_indices], flux_5ₙ, colormap=colorscheme_flux, colorrange=flux_lims[1]) +flux_15_surface = heatmap!(axflux_15, xC, zC[zF_indices], flux_15ₙ, colormap=colorscheme_flux, colorrange=flux_lims[2]) +flux_25_surface = heatmap!(axflux_25, xC, zC[zF_indices], flux_25ₙ, colormap=colorscheme_flux, colorrange=flux_lims[3]) +flux_35_surface = heatmap!(axflux_35, xC, zC[zF_indices], flux_35ₙ, colormap=colorscheme_flux, colorrange=flux_lims[4]) +flux_45_surface = heatmap!(axflux_45, xC, zC[zF_indices], flux_45ₙ, colormap=colorscheme_flux, colorrange=flux_lims[5]) +flux_55_surface = heatmap!(axflux_55, xC, zC[zF_indices], flux_55ₙ, colormap=colorscheme_flux, colorrange=flux_lims[6]) +flux_65_surface = heatmap!(axflux_65, xC, zC[zF_indices], flux_65ₙ, colormap=colorscheme_flux, colorrange=flux_lims[7]) +flux_75_surface = heatmap!(axflux_75, xC, zC[zF_indices], flux_75ₙ, colormap=colorscheme_flux, colorrange=flux_lims[8]) +flux_85_surface = heatmap!(axflux_85, xC, zC[zF_indices], flux_85ₙ, colormap=colorscheme_flux, colorrange=flux_lims[9]) +flux_95_surface = heatmap!(axflux_95, xC, zC[zF_indices], flux_95ₙ, colormap=colorscheme_flux, colorrange=flux_lims[10]) + +Colorbar(fig[1, 2], field_5_surface, label="Field") +Colorbar(fig[1, 4], flux_5_surface, label="NN Flux") +Colorbar(fig[1, 6], field_15_surface, label="Field") +Colorbar(fig[1, 8], flux_15_surface, label="NN Flux") +Colorbar(fig[2, 2], field_25_surface, label="Field") +Colorbar(fig[2, 4], flux_25_surface, label="NN Flux") +Colorbar(fig[2, 6], field_35_surface, label="Field") +Colorbar(fig[2, 8], flux_35_surface, label="NN Flux") +Colorbar(fig[3, 2], field_45_surface, label="Field") +Colorbar(fig[3, 4], flux_45_surface, label="NN Flux") +Colorbar(fig[3, 6], field_55_surface, label="Field") +Colorbar(fig[3, 8], flux_55_surface, label="NN Flux") +Colorbar(fig[4, 2], field_65_surface, label="Field") +Colorbar(fig[4, 4], flux_65_surface, label="NN Flux") +Colorbar(fig[4, 6], field_75_surface, label="Field") +Colorbar(fig[4, 8], flux_75_surface, label="NN Flux") +Colorbar(fig[5, 2], field_85_surface, label="Field") +Colorbar(fig[5, 4], flux_85_surface, label="NN Flux") +Colorbar(fig[5, 6], field_95_surface, label="Field") +Colorbar(fig[5, 8], flux_95_surface, label="NN Flux") + +for axfield in axfields + xlims!(axfield, minimum(xC), maximum(xC)) + ylims!(axfield, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +end + +for axflux in axfluxes + xlims!(axflux, minimum(xC), maximum(xC)) + ylims!(axflux, minimum(zC[zF_indices]), maximum(zC[zF_indices])) +end + +title_str = @lift "Temperature (°C), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_xzslices_fluxes_T.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end + +#%% +@info "Recording S fields and fluxes in xz" + +fieldname = "S" +fluxname = "wS_NN" + +field_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fieldname, backend=OnDisk()) +field_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fieldname, backend=OnDisk()) +field_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fieldname, backend=OnDisk()) +field_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fieldname, backend=OnDisk()) +field_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fieldname, backend=OnDisk()) +field_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fieldname, backend=OnDisk()) +field_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fieldname, backend=OnDisk()) +field_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fieldname, backend=OnDisk()) +field_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fieldname, backend=OnDisk()) +field_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fluxname, backend=OnDisk()) +flux_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fluxname, backend=OnDisk()) +flux_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fluxname, backend=OnDisk()) +flux_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fluxname, backend=OnDisk()) +flux_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fluxname, backend=OnDisk()) +flux_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fluxname, backend=OnDisk()) +flux_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fluxname, backend=OnDisk()) +flux_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fluxname, backend=OnDisk()) +flux_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fluxname, backend=OnDisk()) +flux_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fluxname, backend=OnDisk()) + +field_datas = [field_NN_data_5, field_NN_data_15, field_NN_data_25, field_NN_data_35, field_NN_data_45, field_NN_data_55, field_NN_data_65, field_NN_data_75, field_NN_data_85, field_NN_data_95] +flux_datas = [flux_NN_data_5, flux_NN_data_15, flux_NN_data_25, flux_NN_data_35, flux_NN_data_45, flux_NN_data_55, flux_NN_data_65, flux_NN_data_75, flux_NN_data_85, flux_NN_data_95] + +xC = field_NN_data_5.grid.xᶜᵃᵃ[1:field_NN_data_5.grid.Nx] +yC = field_NN_data_5.grid.yᵃᶜᵃ[1:field_NN_data_5.grid.Ny] +zC = field_NN_data_5.grid.zᵃᵃᶜ[1:field_NN_data_5.grid.Nz] +zF = field_NN_data_5.grid.zᵃᵃᶠ[1:field_NN_data_5.grid.Nz+1] + +Nt = length(field_NN_data_95) +times = field_NN_data_5.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_5 = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_5.grid.yᵃᶜᵃ[field_NN_data_5.indices[2][1]] / 1000) km") +axfield_15 = CairoMakie.Axis(fig[1, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_15.grid.yᵃᶜᵃ[field_NN_data_15.indices[2][1]] / 1000) km") +axfield_25 = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_25.grid.yᵃᶜᵃ[field_NN_data_25.indices[2][1]] / 1000) km") +axfield_35 = CairoMakie.Axis(fig[2, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_35.grid.yᵃᶜᵃ[field_NN_data_35.indices[2][1]] / 1000) km") +axfield_45 = CairoMakie.Axis(fig[3, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_45.grid.yᵃᶜᵃ[field_NN_data_45.indices[2][1]] / 1000) km") +axfield_55 = CairoMakie.Axis(fig[3, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_55.grid.yᵃᶜᵃ[field_NN_data_55.indices[2][1]] / 1000) km") +axfield_65 = CairoMakie.Axis(fig[4, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_65.grid.yᵃᶜᵃ[field_NN_data_65.indices[2][1]] / 1000) km") +axfield_75 = CairoMakie.Axis(fig[4, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_75.grid.yᵃᶜᵃ[field_NN_data_75.indices[2][1]] / 1000) km") +axfield_85 = CairoMakie.Axis(fig[5, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_85.grid.yᵃᶜᵃ[field_NN_data_85.indices[2][1]] / 1000) km") +axfield_95 = CairoMakie.Axis(fig[5, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_95.grid.yᵃᶜᵃ[field_NN_data_95.indices[2][1]] / 1000) km") + +axflux_5 = CairoMakie.Axis(fig[1, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_5.grid.yᵃᶜᵃ[flux_NN_data_5.indices[2][1]] / 1000) km") +axflux_15 = CairoMakie.Axis(fig[1, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_15.grid.yᵃᶜᵃ[flux_NN_data_15.indices[2][1]] / 1000) km") +axflux_25 = CairoMakie.Axis(fig[2, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_25.grid.yᵃᶜᵃ[flux_NN_data_25.indices[2][1]] / 1000) km") +axflux_35 = CairoMakie.Axis(fig[2, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_35.grid.yᵃᶜᵃ[flux_NN_data_35.indices[2][1]] / 1000) km") +axflux_45 = CairoMakie.Axis(fig[3, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_45.grid.yᵃᶜᵃ[flux_NN_data_45.indices[2][1]] / 1000) km") +axflux_55 = CairoMakie.Axis(fig[3, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_55.grid.yᵃᶜᵃ[flux_NN_data_55.indices[2][1]] / 1000) km") +axflux_65 = CairoMakie.Axis(fig[4, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_65.grid.yᵃᶜᵃ[flux_NN_data_65.indices[2][1]] / 1000) km") +axflux_75 = CairoMakie.Axis(fig[4, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_75.grid.yᵃᶜᵃ[flux_NN_data_75.indices[2][1]] / 1000) km") +axflux_85 = CairoMakie.Axis(fig[5, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_85.grid.yᵃᶜᵃ[flux_NN_data_85.indices[2][1]] / 1000) km") +axflux_95 = CairoMakie.Axis(fig[5, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_95.grid.yᵃᶜᵃ[flux_NN_data_95.indices[2][1]] / 1000) km") + +axfields = [axfield_5, axfield_15, axfield_25, axfield_35, axfield_45, axfield_55, axfield_65, axfield_75, axfield_85, axfield_95] +axfluxes = [axflux_5, axflux_15, axflux_25, axflux_35, axflux_45, axflux_55, axflux_65, axflux_75, axflux_85, axflux_95] + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lims = [(find_min(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices))) for field_data in field_datas] + +flux_lims = [(find_min(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices))) for flux_data in flux_datas] + +flux_lims = [(-maximum(abs, flux_lim), maximum(abs, flux_lim)) for flux_lim in flux_lims] + +NNₙs = [@lift interior(field_data[$n], :, 1, zC_indices) for field_data in field_datas] +fluxₙs = [@lift interior(flux_data[$n], :, 1, zF_indices) for flux_data in flux_datas] + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_5_surface = heatmap!(axfield_5, xC, zC[zC_indices], NN_5ₙ, colormap=colorscheme_field, colorrange=field_lims[1]) +field_15_surface = heatmap!(axfield_15, xC, zC[zC_indices], NN_15ₙ, colormap=colorscheme_field, colorrange=field_lims[2]) +field_25_surface = heatmap!(axfield_25, xC, zC[zC_indices], NN_25ₙ, colormap=colorscheme_field, colorrange=field_lims[3]) +field_35_surface = heatmap!(axfield_35, xC, zC[zC_indices], NN_35ₙ, colormap=colorscheme_field, colorrange=field_lims[4]) +field_45_surface = heatmap!(axfield_45, xC, zC[zC_indices], NN_45ₙ, colormap=colorscheme_field, colorrange=field_lims[5]) +field_55_surface = heatmap!(axfield_55, xC, zC[zC_indices], NN_55ₙ, colormap=colorscheme_field, colorrange=field_lims[6]) +field_65_surface = heatmap!(axfield_65, xC, zC[zC_indices], NN_65ₙ, colormap=colorscheme_field, colorrange=field_lims[7]) +field_75_surface = heatmap!(axfield_75, xC, zC[zC_indices], NN_75ₙ, colormap=colorscheme_field, colorrange=field_lims[8]) +field_85_surface = heatmap!(axfield_85, xC, zC[zC_indices], NN_85ₙ, colormap=colorscheme_field, colorrange=field_lims[9]) +field_95_surface = heatmap!(axfield_95, xC, zC[zC_indices], NN_95ₙ, colormap=colorscheme_field, colorrange=field_lims[10]) + +flux_5_surface = heatmap!(axflux_5, xC, zC[zF_indices], flux_5ₙ, colormap=colorscheme_flux, colorrange=flux_lims[1]) +flux_15_surface = heatmap!(axflux_15, xC, zC[zF_indices], flux_15ₙ, colormap=colorscheme_flux, colorrange=flux_lims[2]) +flux_25_surface = heatmap!(axflux_25, xC, zC[zF_indices], flux_25ₙ, colormap=colorscheme_flux, colorrange=flux_lims[3]) +flux_35_surface = heatmap!(axflux_35, xC, zC[zF_indices], flux_35ₙ, colormap=colorscheme_flux, colorrange=flux_lims[4]) +flux_45_surface = heatmap!(axflux_45, xC, zC[zF_indices], flux_45ₙ, colormap=colorscheme_flux, colorrange=flux_lims[5]) +flux_55_surface = heatmap!(axflux_55, xC, zC[zF_indices], flux_55ₙ, colormap=colorscheme_flux, colorrange=flux_lims[6]) +flux_65_surface = heatmap!(axflux_65, xC, zC[zF_indices], flux_65ₙ, colormap=colorscheme_flux, colorrange=flux_lims[7]) +flux_75_surface = heatmap!(axflux_75, xC, zC[zF_indices], flux_75ₙ, colormap=colorscheme_flux, colorrange=flux_lims[8]) +flux_85_surface = heatmap!(axflux_85, xC, zC[zF_indices], flux_85ₙ, colormap=colorscheme_flux, colorrange=flux_lims[9]) +flux_95_surface = heatmap!(axflux_95, xC, zC[zF_indices], flux_95ₙ, colormap=colorscheme_flux, colorrange=flux_lims[10]) + +Colorbar(fig[1, 2], field_5_surface, label="Field") +Colorbar(fig[1, 4], flux_5_surface, label="NN Flux") +Colorbar(fig[1, 6], field_15_surface, label="Field") +Colorbar(fig[1, 8], flux_15_surface, label="NN Flux") +Colorbar(fig[2, 2], field_25_surface, label="Field") +Colorbar(fig[2, 4], flux_25_surface, label="NN Flux") +Colorbar(fig[2, 6], field_35_surface, label="Field") +Colorbar(fig[2, 8], flux_35_surface, label="NN Flux") +Colorbar(fig[3, 2], field_45_surface, label="Field") +Colorbar(fig[3, 4], flux_45_surface, label="NN Flux") +Colorbar(fig[3, 6], field_55_surface, label="Field") +Colorbar(fig[3, 8], flux_55_surface, label="NN Flux") +Colorbar(fig[4, 2], field_65_surface, label="Field") +Colorbar(fig[4, 4], flux_65_surface, label="NN Flux") +Colorbar(fig[4, 6], field_75_surface, label="Field") +Colorbar(fig[4, 8], flux_75_surface, label="NN Flux") +Colorbar(fig[5, 2], field_85_surface, label="Field") +Colorbar(fig[5, 4], flux_85_surface, label="NN Flux") +Colorbar(fig[5, 6], field_95_surface, label="Field") +Colorbar(fig[5, 8], flux_95_surface, label="NN Flux") + +for axfield in axfields + xlims!(axfield, minimum(xC), maximum(xC)) + ylims!(axfield, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +end + +for axflux in axfluxes + xlims!(axflux, minimum(xC), maximum(xC)) + ylims!(axflux, minimum(zC[zF_indices]), maximum(zC[zF_indices])) +end + +title_str = @lift "Salinity (psu), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_xzslices_fluxes_S.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end #%% \ No newline at end of file diff --git a/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl b/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl index 0e09921124..7c95a8f5e0 100644 --- a/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl +++ b/NNclosure_nof_BBLkappazonelast55_doublegyre_model_wallrestoration.jl @@ -577,30 +577,695 @@ zlims!(axv, (-Lz, 0)) @info "Recording 3D fields" CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn - @info "Recording frame $nn" n[] = nn end -@info "Done!" #%% -# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") - -# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] -# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] - -# Nt = length(Ψ_data) -# times = Ψ_data.times / 24 / 60^2 / 365 -# #%% -# timeframe = Nt -# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 -# clim = maximum(abs, Ψ_frame) + 1e-13 -# N_levels = 16 -# levels = range(-clim, stop=clim, length=N_levels) -# fig = Figure(size=(800, 800)) -# ax = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") -# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) -# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") -# tightlimits!(ax) -# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +@info "Recording T fields and fluxes in yz" + +fieldname = "T" +fluxname = "wT_NN" +field_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fieldname, backend=OnDisk()) +field_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fieldname, backend=OnDisk()) +field_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fieldname, backend=OnDisk()) +field_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fieldname, backend=OnDisk()) +field_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fieldname, backend=OnDisk()) +field_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fieldname, backend=OnDisk()) +field_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fieldname, backend=OnDisk()) +field_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fieldname, backend=OnDisk()) +field_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fieldname, backend=OnDisk()) +field_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fluxname, backend=OnDisk()) +flux_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fluxname, backend=OnDisk()) +flux_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fluxname, backend=OnDisk()) +flux_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fluxname, backend=OnDisk()) +flux_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fluxname, backend=OnDisk()) +flux_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fluxname, backend=OnDisk()) +flux_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fluxname, backend=OnDisk()) +flux_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fluxname, backend=OnDisk()) +flux_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fluxname, backend=OnDisk()) +flux_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fluxname, backend=OnDisk()) + +xC = field_NN_data_00.grid.xᶜᵃᵃ[1:field_NN_data_00.grid.Nx] +yC = field_NN_data_00.grid.yᵃᶜᵃ[1:field_NN_data_00.grid.Ny] +zC = field_NN_data_00.grid.zᵃᵃᶜ[1:field_NN_data_00.grid.Nz] +zF = field_NN_data_00.grid.zᵃᵃᶠ[1:field_NN_data_00.grid.Nz+1] + +Nt = length(field_NN_data_90) +times = field_NN_data_00.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_00 = CairoMakie.Axis(fig[1, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_00.grid.xᶜᵃᵃ[field_NN_data_00.indices[1][1]] / 1000) km") +axfield_10 = CairoMakie.Axis(fig[1, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_10.grid.xᶜᵃᵃ[field_NN_data_10.indices[1][1]] / 1000) km") +axfield_20 = CairoMakie.Axis(fig[2, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_20.grid.xᶜᵃᵃ[field_NN_data_20.indices[1][1]] / 1000) km") +axfield_30 = CairoMakie.Axis(fig[2, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_30.grid.xᶜᵃᵃ[field_NN_data_30.indices[1][1]] / 1000) km") +axfield_40 = CairoMakie.Axis(fig[3, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_40.grid.xᶜᵃᵃ[field_NN_data_40.indices[1][1]] / 1000) km") +axfield_50 = CairoMakie.Axis(fig[3, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_50.grid.xᶜᵃᵃ[field_NN_data_50.indices[1][1]] / 1000) km") +axfield_60 = CairoMakie.Axis(fig[4, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_60.grid.xᶜᵃᵃ[field_NN_data_60.indices[1][1]] / 1000) km") +axfield_70 = CairoMakie.Axis(fig[4, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_70.grid.xᶜᵃᵃ[field_NN_data_70.indices[1][1]] / 1000) km") +axfield_80 = CairoMakie.Axis(fig[5, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_80.grid.xᶜᵃᵃ[field_NN_data_80.indices[1][1]] / 1000) km") +axfield_90 = CairoMakie.Axis(fig[5, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_90.grid.xᶜᵃᵃ[field_NN_data_90.indices[1][1]] / 1000) km") + +axflux_00 = CairoMakie.Axis(fig[1, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_00.grid.xᶜᵃᵃ[flux_NN_data_00.indices[1][1]] / 1000) km") +axflux_10 = CairoMakie.Axis(fig[1, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_10.grid.xᶜᵃᵃ[flux_NN_data_10.indices[1][1]] / 1000) km") +axflux_20 = CairoMakie.Axis(fig[2, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_20.grid.xᶜᵃᵃ[flux_NN_data_20.indices[1][1]] / 1000) km") +axflux_30 = CairoMakie.Axis(fig[2, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_30.grid.xᶜᵃᵃ[flux_NN_data_30.indices[1][1]] / 1000) km") +axflux_40 = CairoMakie.Axis(fig[3, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_40.grid.xᶜᵃᵃ[flux_NN_data_40.indices[1][1]] / 1000) km") +axflux_50 = CairoMakie.Axis(fig[3, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_50.grid.xᶜᵃᵃ[flux_NN_data_50.indices[1][1]] / 1000) km") +axflux_60 = CairoMakie.Axis(fig[4, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_60.grid.xᶜᵃᵃ[flux_NN_data_60.indices[1][1]] / 1000) km") +axflux_70 = CairoMakie.Axis(fig[4, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_70.grid.xᶜᵃᵃ[flux_NN_data_70.indices[1][1]] / 1000) km") +axflux_80 = CairoMakie.Axis(fig[5, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_80.grid.xᶜᵃᵃ[flux_NN_data_80.indices[1][1]] / 1000) km") +axflux_90 = CairoMakie.Axis(fig[5, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_90.grid.xᶜᵃᵃ[flux_NN_data_90.indices[1][1]] / 1000) km") + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lim = (find_min(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (find_min(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (-maximum(abs, flux_lim), maximum(abs, flux_lim)) + +NN_00ₙ = @lift interior(field_NN_data_00[$n], 1, :, zC_indices) +NN_10ₙ = @lift interior(field_NN_data_10[$n], 1, :, zC_indices) +NN_20ₙ = @lift interior(field_NN_data_20[$n], 1, :, zC_indices) +NN_30ₙ = @lift interior(field_NN_data_30[$n], 1, :, zC_indices) +NN_40ₙ = @lift interior(field_NN_data_40[$n], 1, :, zC_indices) +NN_50ₙ = @lift interior(field_NN_data_50[$n], 1, :, zC_indices) +NN_60ₙ = @lift interior(field_NN_data_60[$n], 1, :, zC_indices) +NN_70ₙ = @lift interior(field_NN_data_70[$n], 1, :, zC_indices) +NN_80ₙ = @lift interior(field_NN_data_80[$n], 1, :, zC_indices) +NN_90ₙ = @lift interior(field_NN_data_90[$n], 1, :, zC_indices) + +flux_00ₙ = @lift interior(flux_NN_data_00[$n], 1, :, zF_indices) +flux_10ₙ = @lift interior(flux_NN_data_10[$n], 1, :, zF_indices) +flux_20ₙ = @lift interior(flux_NN_data_20[$n], 1, :, zF_indices) +flux_30ₙ = @lift interior(flux_NN_data_30[$n], 1, :, zF_indices) +flux_40ₙ = @lift interior(flux_NN_data_40[$n], 1, :, zF_indices) +flux_50ₙ = @lift interior(flux_NN_data_50[$n], 1, :, zF_indices) +flux_60ₙ = @lift interior(flux_NN_data_60[$n], 1, :, zF_indices) +flux_70ₙ = @lift interior(flux_NN_data_70[$n], 1, :, zF_indices) +flux_80ₙ = @lift interior(flux_NN_data_80[$n], 1, :, zF_indices) +flux_90ₙ = @lift interior(flux_NN_data_90[$n], 1, :, zF_indices) + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_00_surface = heatmap!(axfield_00, yC, zC[zC_indices], NN_00ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_10_surface = heatmap!(axfield_10, yC, zC[zC_indices], NN_10ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_20_surface = heatmap!(axfield_20, yC, zC[zC_indices], NN_20ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_30_surface = heatmap!(axfield_30, yC, zC[zC_indices], NN_30ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_40_surface = heatmap!(axfield_40, yC, zC[zC_indices], NN_40ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_50_surface = heatmap!(axfield_50, yC, zC[zC_indices], NN_50ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_60_surface = heatmap!(axfield_60, yC, zC[zC_indices], NN_60ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_70_surface = heatmap!(axfield_70, yC, zC[zC_indices], NN_70ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_80_surface = heatmap!(axfield_80, yC, zC[zC_indices], NN_80ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_90_surface = heatmap!(axfield_90, yC, zC[zC_indices], NN_90ₙ, colormap=colorscheme_field, colorrange=field_lim) + +flux_00_surface = heatmap!(axflux_00, yC, zC[zF_indices], flux_00ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_10_surface = heatmap!(axflux_10, yC, zC[zF_indices], flux_10ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_20_surface = heatmap!(axflux_20, yC, zC[zF_indices], flux_20ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_30_surface = heatmap!(axflux_30, yC, zC[zF_indices], flux_30ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_40_surface = heatmap!(axflux_40, yC, zC[zF_indices], flux_40ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_50_surface = heatmap!(axflux_50, yC, zC[zF_indices], flux_50ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_60_surface = heatmap!(axflux_60, yC, zC[zF_indices], flux_60ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_70_surface = heatmap!(axflux_70, yC, zC[zF_indices], flux_70ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_80_surface = heatmap!(axflux_80, yC, zC[zF_indices], flux_80ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_90_surface = heatmap!(axflux_90, yC, zC[zF_indices], flux_90ₙ, colormap=colorscheme_flux, colorrange=flux_lim) + +Colorbar(fig[1:5, 2], field_00_surface, label="Field") +Colorbar(fig[1:5, 4], flux_00_surface, label="NN Flux") +Colorbar(fig[1:5, 6], field_00_surface, label="Field") +Colorbar(fig[1:5, 8], flux_00_surface, label="NN Flux") + +xlims!(axfield_00, minimum(yC), maximum(yC)) +xlims!(axfield_10, minimum(yC), maximum(yC)) +xlims!(axfield_20, minimum(yC), maximum(yC)) +xlims!(axfield_30, minimum(yC), maximum(yC)) +xlims!(axfield_40, minimum(yC), maximum(yC)) +xlims!(axfield_50, minimum(yC), maximum(yC)) +xlims!(axfield_60, minimum(yC), maximum(yC)) +xlims!(axfield_70, minimum(yC), maximum(yC)) +xlims!(axfield_80, minimum(yC), maximum(yC)) +xlims!(axfield_90, minimum(yC), maximum(yC)) + +ylims!(axfield_00, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_10, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_20, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_30, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_40, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_50, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_60, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_70, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_80, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_90, minimum(zC[zC_indices]), maximum(zC[zC_indices])) + +xlims!(axflux_00, minimum(yC), maximum(yC)) +xlims!(axflux_10, minimum(yC), maximum(yC)) +xlims!(axflux_20, minimum(yC), maximum(yC)) +xlims!(axflux_30, minimum(yC), maximum(yC)) +xlims!(axflux_40, minimum(yC), maximum(yC)) +xlims!(axflux_50, minimum(yC), maximum(yC)) +xlims!(axflux_60, minimum(yC), maximum(yC)) +xlims!(axflux_70, minimum(yC), maximum(yC)) +xlims!(axflux_80, minimum(yC), maximum(yC)) +xlims!(axflux_90, minimum(yC), maximum(yC)) + +ylims!(axflux_00, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_10, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_20, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_30, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_40, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_50, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_60, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_70, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_80, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_90, minimum(zF[zF_indices]), maximum(zF[zF_indices])) + +title_str = @lift "Temperature (°C), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +# save("./Output/compare_3D_instantaneous_fields_slices_NNclosure_fluxes.png", fig) # display(fig) +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_yzslices_fluxes_T.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end +#%% +@info "Recording S fields and fluxes in yz" + +fieldname = "S" +fluxname = "wS_NN" +field_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fieldname, backend=OnDisk()) +field_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fieldname, backend=OnDisk()) +field_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fieldname, backend=OnDisk()) +field_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fieldname, backend=OnDisk()) +field_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fieldname, backend=OnDisk()) +field_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fieldname, backend=OnDisk()) +field_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fieldname, backend=OnDisk()) +field_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fieldname, backend=OnDisk()) +field_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fieldname, backend=OnDisk()) +field_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fluxname, backend=OnDisk()) +flux_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fluxname, backend=OnDisk()) +flux_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fluxname, backend=OnDisk()) +flux_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fluxname, backend=OnDisk()) +flux_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fluxname, backend=OnDisk()) +flux_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fluxname, backend=OnDisk()) +flux_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fluxname, backend=OnDisk()) +flux_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fluxname, backend=OnDisk()) +flux_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fluxname, backend=OnDisk()) +flux_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fluxname, backend=OnDisk()) + +xC = field_NN_data_00.grid.xᶜᵃᵃ[1:field_NN_data_00.grid.Nx] +yC = field_NN_data_00.grid.yᵃᶜᵃ[1:field_NN_data_00.grid.Ny] +zC = field_NN_data_00.grid.zᵃᵃᶜ[1:field_NN_data_00.grid.Nz] +zF = field_NN_data_00.grid.zᵃᵃᶠ[1:field_NN_data_00.grid.Nz+1] + +Nt = length(field_NN_data_90) +times = field_NN_data_00.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_00 = CairoMakie.Axis(fig[1, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_00.grid.xᶜᵃᵃ[field_NN_data_00.indices[1][1]] / 1000) km") +axfield_10 = CairoMakie.Axis(fig[1, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_10.grid.xᶜᵃᵃ[field_NN_data_10.indices[1][1]] / 1000) km") +axfield_20 = CairoMakie.Axis(fig[2, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_20.grid.xᶜᵃᵃ[field_NN_data_20.indices[1][1]] / 1000) km") +axfield_30 = CairoMakie.Axis(fig[2, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_30.grid.xᶜᵃᵃ[field_NN_data_30.indices[1][1]] / 1000) km") +axfield_40 = CairoMakie.Axis(fig[3, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_40.grid.xᶜᵃᵃ[field_NN_data_40.indices[1][1]] / 1000) km") +axfield_50 = CairoMakie.Axis(fig[3, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_50.grid.xᶜᵃᵃ[field_NN_data_50.indices[1][1]] / 1000) km") +axfield_60 = CairoMakie.Axis(fig[4, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_60.grid.xᶜᵃᵃ[field_NN_data_60.indices[1][1]] / 1000) km") +axfield_70 = CairoMakie.Axis(fig[4, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_70.grid.xᶜᵃᵃ[field_NN_data_70.indices[1][1]] / 1000) km") +axfield_80 = CairoMakie.Axis(fig[5, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_80.grid.xᶜᵃᵃ[field_NN_data_80.indices[1][1]] / 1000) km") +axfield_90 = CairoMakie.Axis(fig[5, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_90.grid.xᶜᵃᵃ[field_NN_data_90.indices[1][1]] / 1000) km") + +axflux_00 = CairoMakie.Axis(fig[1, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_00.grid.xᶜᵃᵃ[flux_NN_data_00.indices[1][1]] / 1000) km") +axflux_10 = CairoMakie.Axis(fig[1, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_10.grid.xᶜᵃᵃ[flux_NN_data_10.indices[1][1]] / 1000) km") +axflux_20 = CairoMakie.Axis(fig[2, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_20.grid.xᶜᵃᵃ[flux_NN_data_20.indices[1][1]] / 1000) km") +axflux_30 = CairoMakie.Axis(fig[2, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_30.grid.xᶜᵃᵃ[flux_NN_data_30.indices[1][1]] / 1000) km") +axflux_40 = CairoMakie.Axis(fig[3, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_40.grid.xᶜᵃᵃ[flux_NN_data_40.indices[1][1]] / 1000) km") +axflux_50 = CairoMakie.Axis(fig[3, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_50.grid.xᶜᵃᵃ[flux_NN_data_50.indices[1][1]] / 1000) km") +axflux_60 = CairoMakie.Axis(fig[4, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_60.grid.xᶜᵃᵃ[flux_NN_data_60.indices[1][1]] / 1000) km") +axflux_70 = CairoMakie.Axis(fig[4, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_70.grid.xᶜᵃᵃ[flux_NN_data_70.indices[1][1]] / 1000) km") +axflux_80 = CairoMakie.Axis(fig[5, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_80.grid.xᶜᵃᵃ[flux_NN_data_80.indices[1][1]] / 1000) km") +axflux_90 = CairoMakie.Axis(fig[5, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_90.grid.xᶜᵃᵃ[flux_NN_data_90.indices[1][1]] / 1000) km") + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lim = (find_min(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (find_min(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (-maximum(abs, flux_lim), maximum(abs, flux_lim)) + +NN_00ₙ = @lift interior(field_NN_data_00[$n], 1, :, zC_indices) +NN_10ₙ = @lift interior(field_NN_data_10[$n], 1, :, zC_indices) +NN_20ₙ = @lift interior(field_NN_data_20[$n], 1, :, zC_indices) +NN_30ₙ = @lift interior(field_NN_data_30[$n], 1, :, zC_indices) +NN_40ₙ = @lift interior(field_NN_data_40[$n], 1, :, zC_indices) +NN_50ₙ = @lift interior(field_NN_data_50[$n], 1, :, zC_indices) +NN_60ₙ = @lift interior(field_NN_data_60[$n], 1, :, zC_indices) +NN_70ₙ = @lift interior(field_NN_data_70[$n], 1, :, zC_indices) +NN_80ₙ = @lift interior(field_NN_data_80[$n], 1, :, zC_indices) +NN_90ₙ = @lift interior(field_NN_data_90[$n], 1, :, zC_indices) + +flux_00ₙ = @lift interior(flux_NN_data_00[$n], 1, :, zF_indices) +flux_10ₙ = @lift interior(flux_NN_data_10[$n], 1, :, zF_indices) +flux_20ₙ = @lift interior(flux_NN_data_20[$n], 1, :, zF_indices) +flux_30ₙ = @lift interior(flux_NN_data_30[$n], 1, :, zF_indices) +flux_40ₙ = @lift interior(flux_NN_data_40[$n], 1, :, zF_indices) +flux_50ₙ = @lift interior(flux_NN_data_50[$n], 1, :, zF_indices) +flux_60ₙ = @lift interior(flux_NN_data_60[$n], 1, :, zF_indices) +flux_70ₙ = @lift interior(flux_NN_data_70[$n], 1, :, zF_indices) +flux_80ₙ = @lift interior(flux_NN_data_80[$n], 1, :, zF_indices) +flux_90ₙ = @lift interior(flux_NN_data_90[$n], 1, :, zF_indices) + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_00_surface = heatmap!(axfield_00, yC, zC[zC_indices], NN_00ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_10_surface = heatmap!(axfield_10, yC, zC[zC_indices], NN_10ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_20_surface = heatmap!(axfield_20, yC, zC[zC_indices], NN_20ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_30_surface = heatmap!(axfield_30, yC, zC[zC_indices], NN_30ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_40_surface = heatmap!(axfield_40, yC, zC[zC_indices], NN_40ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_50_surface = heatmap!(axfield_50, yC, zC[zC_indices], NN_50ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_60_surface = heatmap!(axfield_60, yC, zC[zC_indices], NN_60ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_70_surface = heatmap!(axfield_70, yC, zC[zC_indices], NN_70ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_80_surface = heatmap!(axfield_80, yC, zC[zC_indices], NN_80ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_90_surface = heatmap!(axfield_90, yC, zC[zC_indices], NN_90ₙ, colormap=colorscheme_field, colorrange=field_lim) + +flux_00_surface = heatmap!(axflux_00, yC, zC[zF_indices], flux_00ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_10_surface = heatmap!(axflux_10, yC, zC[zF_indices], flux_10ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_20_surface = heatmap!(axflux_20, yC, zC[zF_indices], flux_20ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_30_surface = heatmap!(axflux_30, yC, zC[zF_indices], flux_30ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_40_surface = heatmap!(axflux_40, yC, zC[zF_indices], flux_40ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_50_surface = heatmap!(axflux_50, yC, zC[zF_indices], flux_50ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_60_surface = heatmap!(axflux_60, yC, zC[zF_indices], flux_60ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_70_surface = heatmap!(axflux_70, yC, zC[zF_indices], flux_70ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_80_surface = heatmap!(axflux_80, yC, zC[zF_indices], flux_80ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_90_surface = heatmap!(axflux_90, yC, zC[zF_indices], flux_90ₙ, colormap=colorscheme_flux, colorrange=flux_lim) + +Colorbar(fig[1:5, 2], field_00_surface, label="Field") +Colorbar(fig[1:5, 4], flux_00_surface, label="NN Flux") +Colorbar(fig[1:5, 6], field_00_surface, label="Field") +Colorbar(fig[1:5, 8], flux_00_surface, label="NN Flux") + +xlims!(axfield_00, minimum(yC), maximum(yC)) +xlims!(axfield_10, minimum(yC), maximum(yC)) +xlims!(axfield_20, minimum(yC), maximum(yC)) +xlims!(axfield_30, minimum(yC), maximum(yC)) +xlims!(axfield_40, minimum(yC), maximum(yC)) +xlims!(axfield_50, minimum(yC), maximum(yC)) +xlims!(axfield_60, minimum(yC), maximum(yC)) +xlims!(axfield_70, minimum(yC), maximum(yC)) +xlims!(axfield_80, minimum(yC), maximum(yC)) +xlims!(axfield_90, minimum(yC), maximum(yC)) + +ylims!(axfield_00, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_10, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_20, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_30, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_40, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_50, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_60, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_70, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_80, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_90, minimum(zC[zC_indices]), maximum(zC[zC_indices])) + +xlims!(axflux_00, minimum(yC), maximum(yC)) +xlims!(axflux_10, minimum(yC), maximum(yC)) +xlims!(axflux_20, minimum(yC), maximum(yC)) +xlims!(axflux_30, minimum(yC), maximum(yC)) +xlims!(axflux_40, minimum(yC), maximum(yC)) +xlims!(axflux_50, minimum(yC), maximum(yC)) +xlims!(axflux_60, minimum(yC), maximum(yC)) +xlims!(axflux_70, minimum(yC), maximum(yC)) +xlims!(axflux_80, minimum(yC), maximum(yC)) +xlims!(axflux_90, minimum(yC), maximum(yC)) + +ylims!(axflux_00, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_10, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_20, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_30, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_40, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_50, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_60, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_70, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_80, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_90, minimum(zF[zF_indices]), maximum(zF[zF_indices])) + +title_str = @lift "Salinity (psu), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_yzslices_fluxes_S.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end + +#%% +@info "Recording T fields and fluxes in xz" + +fieldname = "T" +fluxname = "wT_NN" + +field_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fieldname, backend=OnDisk()) +field_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fieldname, backend=OnDisk()) +field_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fieldname, backend=OnDisk()) +field_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fieldname, backend=OnDisk()) +field_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fieldname, backend=OnDisk()) +field_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fieldname, backend=OnDisk()) +field_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fieldname, backend=OnDisk()) +field_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fieldname, backend=OnDisk()) +field_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fieldname, backend=OnDisk()) +field_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fluxname, backend=OnDisk()) +flux_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fluxname, backend=OnDisk()) +flux_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fluxname, backend=OnDisk()) +flux_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fluxname, backend=OnDisk()) +flux_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fluxname, backend=OnDisk()) +flux_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fluxname, backend=OnDisk()) +flux_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fluxname, backend=OnDisk()) +flux_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fluxname, backend=OnDisk()) +flux_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fluxname, backend=OnDisk()) +flux_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fluxname, backend=OnDisk()) + +field_datas = [field_NN_data_5, field_NN_data_15, field_NN_data_25, field_NN_data_35, field_NN_data_45, field_NN_data_55, field_NN_data_65, field_NN_data_75, field_NN_data_85, field_NN_data_95] +flux_datas = [flux_NN_data_5, flux_NN_data_15, flux_NN_data_25, flux_NN_data_35, flux_NN_data_45, flux_NN_data_55, flux_NN_data_65, flux_NN_data_75, flux_NN_data_85, flux_NN_data_95] + +xC = field_NN_data_5.grid.xᶜᵃᵃ[1:field_NN_data_5.grid.Nx] +yC = field_NN_data_5.grid.yᵃᶜᵃ[1:field_NN_data_5.grid.Ny] +zC = field_NN_data_5.grid.zᵃᵃᶜ[1:field_NN_data_5.grid.Nz] +zF = field_NN_data_5.grid.zᵃᵃᶠ[1:field_NN_data_5.grid.Nz+1] + +Nt = length(field_NN_data_95) +times = field_NN_data_5.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_5 = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_5.grid.yᵃᶜᵃ[field_NN_data_5.indices[2][1]] / 1000) km") +axfield_15 = CairoMakie.Axis(fig[1, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_15.grid.yᵃᶜᵃ[field_NN_data_15.indices[2][1]] / 1000) km") +axfield_25 = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_25.grid.yᵃᶜᵃ[field_NN_data_25.indices[2][1]] / 1000) km") +axfield_35 = CairoMakie.Axis(fig[2, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_35.grid.yᵃᶜᵃ[field_NN_data_35.indices[2][1]] / 1000) km") +axfield_45 = CairoMakie.Axis(fig[3, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_45.grid.yᵃᶜᵃ[field_NN_data_45.indices[2][1]] / 1000) km") +axfield_55 = CairoMakie.Axis(fig[3, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_55.grid.yᵃᶜᵃ[field_NN_data_55.indices[2][1]] / 1000) km") +axfield_65 = CairoMakie.Axis(fig[4, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_65.grid.yᵃᶜᵃ[field_NN_data_65.indices[2][1]] / 1000) km") +axfield_75 = CairoMakie.Axis(fig[4, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_75.grid.yᵃᶜᵃ[field_NN_data_75.indices[2][1]] / 1000) km") +axfield_85 = CairoMakie.Axis(fig[5, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_85.grid.yᵃᶜᵃ[field_NN_data_85.indices[2][1]] / 1000) km") +axfield_95 = CairoMakie.Axis(fig[5, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_95.grid.yᵃᶜᵃ[field_NN_data_95.indices[2][1]] / 1000) km") + +axflux_5 = CairoMakie.Axis(fig[1, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_5.grid.yᵃᶜᵃ[flux_NN_data_5.indices[2][1]] / 1000) km") +axflux_15 = CairoMakie.Axis(fig[1, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_15.grid.yᵃᶜᵃ[flux_NN_data_15.indices[2][1]] / 1000) km") +axflux_25 = CairoMakie.Axis(fig[2, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_25.grid.yᵃᶜᵃ[flux_NN_data_25.indices[2][1]] / 1000) km") +axflux_35 = CairoMakie.Axis(fig[2, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_35.grid.yᵃᶜᵃ[flux_NN_data_35.indices[2][1]] / 1000) km") +axflux_45 = CairoMakie.Axis(fig[3, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_45.grid.yᵃᶜᵃ[flux_NN_data_45.indices[2][1]] / 1000) km") +axflux_55 = CairoMakie.Axis(fig[3, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_55.grid.yᵃᶜᵃ[flux_NN_data_55.indices[2][1]] / 1000) km") +axflux_65 = CairoMakie.Axis(fig[4, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_65.grid.yᵃᶜᵃ[flux_NN_data_65.indices[2][1]] / 1000) km") +axflux_75 = CairoMakie.Axis(fig[4, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_75.grid.yᵃᶜᵃ[flux_NN_data_75.indices[2][1]] / 1000) km") +axflux_85 = CairoMakie.Axis(fig[5, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_85.grid.yᵃᶜᵃ[flux_NN_data_85.indices[2][1]] / 1000) km") +axflux_95 = CairoMakie.Axis(fig[5, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_95.grid.yᵃᶜᵃ[flux_NN_data_95.indices[2][1]] / 1000) km") + +axfields = [axfield_5, axfield_15, axfield_25, axfield_35, axfield_45, axfield_55, axfield_65, axfield_75, axfield_85, axfield_95] +axfluxes = [axflux_5, axflux_15, axflux_25, axflux_35, axflux_45, axflux_55, axflux_65, axflux_75, axflux_85, axflux_95] + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lims = [(find_min(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices))) for field_data in field_datas] + +flux_lims = [(find_min(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices))) for flux_data in flux_datas] + +flux_lims = [(-maximum(abs, flux_lim), maximum(abs, flux_lim)) for flux_lim in flux_lims] + +NNₙs = [@lift interior(field_data[$n], :, 1, zC_indices) for field_data in field_datas] +fluxₙs = [@lift interior(flux_data[$n], :, 1, zF_indices) for flux_data in flux_datas] + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_5_surface = heatmap!(axfield_5, xC, zC[zC_indices], NN_5ₙ, colormap=colorscheme_field, colorrange=field_lims[1]) +field_15_surface = heatmap!(axfield_15, xC, zC[zC_indices], NN_15ₙ, colormap=colorscheme_field, colorrange=field_lims[2]) +field_25_surface = heatmap!(axfield_25, xC, zC[zC_indices], NN_25ₙ, colormap=colorscheme_field, colorrange=field_lims[3]) +field_35_surface = heatmap!(axfield_35, xC, zC[zC_indices], NN_35ₙ, colormap=colorscheme_field, colorrange=field_lims[4]) +field_45_surface = heatmap!(axfield_45, xC, zC[zC_indices], NN_45ₙ, colormap=colorscheme_field, colorrange=field_lims[5]) +field_55_surface = heatmap!(axfield_55, xC, zC[zC_indices], NN_55ₙ, colormap=colorscheme_field, colorrange=field_lims[6]) +field_65_surface = heatmap!(axfield_65, xC, zC[zC_indices], NN_65ₙ, colormap=colorscheme_field, colorrange=field_lims[7]) +field_75_surface = heatmap!(axfield_75, xC, zC[zC_indices], NN_75ₙ, colormap=colorscheme_field, colorrange=field_lims[8]) +field_85_surface = heatmap!(axfield_85, xC, zC[zC_indices], NN_85ₙ, colormap=colorscheme_field, colorrange=field_lims[9]) +field_95_surface = heatmap!(axfield_95, xC, zC[zC_indices], NN_95ₙ, colormap=colorscheme_field, colorrange=field_lims[10]) + +flux_5_surface = heatmap!(axflux_5, xC, zC[zF_indices], flux_5ₙ, colormap=colorscheme_flux, colorrange=flux_lims[1]) +flux_15_surface = heatmap!(axflux_15, xC, zC[zF_indices], flux_15ₙ, colormap=colorscheme_flux, colorrange=flux_lims[2]) +flux_25_surface = heatmap!(axflux_25, xC, zC[zF_indices], flux_25ₙ, colormap=colorscheme_flux, colorrange=flux_lims[3]) +flux_35_surface = heatmap!(axflux_35, xC, zC[zF_indices], flux_35ₙ, colormap=colorscheme_flux, colorrange=flux_lims[4]) +flux_45_surface = heatmap!(axflux_45, xC, zC[zF_indices], flux_45ₙ, colormap=colorscheme_flux, colorrange=flux_lims[5]) +flux_55_surface = heatmap!(axflux_55, xC, zC[zF_indices], flux_55ₙ, colormap=colorscheme_flux, colorrange=flux_lims[6]) +flux_65_surface = heatmap!(axflux_65, xC, zC[zF_indices], flux_65ₙ, colormap=colorscheme_flux, colorrange=flux_lims[7]) +flux_75_surface = heatmap!(axflux_75, xC, zC[zF_indices], flux_75ₙ, colormap=colorscheme_flux, colorrange=flux_lims[8]) +flux_85_surface = heatmap!(axflux_85, xC, zC[zF_indices], flux_85ₙ, colormap=colorscheme_flux, colorrange=flux_lims[9]) +flux_95_surface = heatmap!(axflux_95, xC, zC[zF_indices], flux_95ₙ, colormap=colorscheme_flux, colorrange=flux_lims[10]) + +Colorbar(fig[1, 2], field_5_surface, label="Field") +Colorbar(fig[1, 4], flux_5_surface, label="NN Flux") +Colorbar(fig[1, 6], field_15_surface, label="Field") +Colorbar(fig[1, 8], flux_15_surface, label="NN Flux") +Colorbar(fig[2, 2], field_25_surface, label="Field") +Colorbar(fig[2, 4], flux_25_surface, label="NN Flux") +Colorbar(fig[2, 6], field_35_surface, label="Field") +Colorbar(fig[2, 8], flux_35_surface, label="NN Flux") +Colorbar(fig[3, 2], field_45_surface, label="Field") +Colorbar(fig[3, 4], flux_45_surface, label="NN Flux") +Colorbar(fig[3, 6], field_55_surface, label="Field") +Colorbar(fig[3, 8], flux_55_surface, label="NN Flux") +Colorbar(fig[4, 2], field_65_surface, label="Field") +Colorbar(fig[4, 4], flux_65_surface, label="NN Flux") +Colorbar(fig[4, 6], field_75_surface, label="Field") +Colorbar(fig[4, 8], flux_75_surface, label="NN Flux") +Colorbar(fig[5, 2], field_85_surface, label="Field") +Colorbar(fig[5, 4], flux_85_surface, label="NN Flux") +Colorbar(fig[5, 6], field_95_surface, label="Field") +Colorbar(fig[5, 8], flux_95_surface, label="NN Flux") + +for axfield in axfields + xlims!(axfield, minimum(xC), maximum(xC)) + ylims!(axfield, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +end + +for axflux in axfluxes + xlims!(axflux, minimum(xC), maximum(xC)) + ylims!(axflux, minimum(zC[zF_indices]), maximum(zC[zF_indices])) +end + +title_str = @lift "Temperature (°C), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_xzslices_fluxes_T.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end + +#%% +@info "Recording S fields and fluxes in xz" + +fieldname = "S" +fluxname = "wS_NN" + +field_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fieldname, backend=OnDisk()) +field_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fieldname, backend=OnDisk()) +field_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fieldname, backend=OnDisk()) +field_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fieldname, backend=OnDisk()) +field_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fieldname, backend=OnDisk()) +field_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fieldname, backend=OnDisk()) +field_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fieldname, backend=OnDisk()) +field_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fieldname, backend=OnDisk()) +field_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fieldname, backend=OnDisk()) +field_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fluxname, backend=OnDisk()) +flux_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fluxname, backend=OnDisk()) +flux_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fluxname, backend=OnDisk()) +flux_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fluxname, backend=OnDisk()) +flux_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fluxname, backend=OnDisk()) +flux_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fluxname, backend=OnDisk()) +flux_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fluxname, backend=OnDisk()) +flux_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fluxname, backend=OnDisk()) +flux_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fluxname, backend=OnDisk()) +flux_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fluxname, backend=OnDisk()) + +field_datas = [field_NN_data_5, field_NN_data_15, field_NN_data_25, field_NN_data_35, field_NN_data_45, field_NN_data_55, field_NN_data_65, field_NN_data_75, field_NN_data_85, field_NN_data_95] +flux_datas = [flux_NN_data_5, flux_NN_data_15, flux_NN_data_25, flux_NN_data_35, flux_NN_data_45, flux_NN_data_55, flux_NN_data_65, flux_NN_data_75, flux_NN_data_85, flux_NN_data_95] + +xC = field_NN_data_5.grid.xᶜᵃᵃ[1:field_NN_data_5.grid.Nx] +yC = field_NN_data_5.grid.yᵃᶜᵃ[1:field_NN_data_5.grid.Ny] +zC = field_NN_data_5.grid.zᵃᵃᶜ[1:field_NN_data_5.grid.Nz] +zF = field_NN_data_5.grid.zᵃᵃᶠ[1:field_NN_data_5.grid.Nz+1] + +Nt = length(field_NN_data_95) +times = field_NN_data_5.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_5 = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_5.grid.yᵃᶜᵃ[field_NN_data_5.indices[2][1]] / 1000) km") +axfield_15 = CairoMakie.Axis(fig[1, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_15.grid.yᵃᶜᵃ[field_NN_data_15.indices[2][1]] / 1000) km") +axfield_25 = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_25.grid.yᵃᶜᵃ[field_NN_data_25.indices[2][1]] / 1000) km") +axfield_35 = CairoMakie.Axis(fig[2, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_35.grid.yᵃᶜᵃ[field_NN_data_35.indices[2][1]] / 1000) km") +axfield_45 = CairoMakie.Axis(fig[3, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_45.grid.yᵃᶜᵃ[field_NN_data_45.indices[2][1]] / 1000) km") +axfield_55 = CairoMakie.Axis(fig[3, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_55.grid.yᵃᶜᵃ[field_NN_data_55.indices[2][1]] / 1000) km") +axfield_65 = CairoMakie.Axis(fig[4, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_65.grid.yᵃᶜᵃ[field_NN_data_65.indices[2][1]] / 1000) km") +axfield_75 = CairoMakie.Axis(fig[4, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_75.grid.yᵃᶜᵃ[field_NN_data_75.indices[2][1]] / 1000) km") +axfield_85 = CairoMakie.Axis(fig[5, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_85.grid.yᵃᶜᵃ[field_NN_data_85.indices[2][1]] / 1000) km") +axfield_95 = CairoMakie.Axis(fig[5, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_95.grid.yᵃᶜᵃ[field_NN_data_95.indices[2][1]] / 1000) km") + +axflux_5 = CairoMakie.Axis(fig[1, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_5.grid.yᵃᶜᵃ[flux_NN_data_5.indices[2][1]] / 1000) km") +axflux_15 = CairoMakie.Axis(fig[1, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_15.grid.yᵃᶜᵃ[flux_NN_data_15.indices[2][1]] / 1000) km") +axflux_25 = CairoMakie.Axis(fig[2, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_25.grid.yᵃᶜᵃ[flux_NN_data_25.indices[2][1]] / 1000) km") +axflux_35 = CairoMakie.Axis(fig[2, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_35.grid.yᵃᶜᵃ[flux_NN_data_35.indices[2][1]] / 1000) km") +axflux_45 = CairoMakie.Axis(fig[3, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_45.grid.yᵃᶜᵃ[flux_NN_data_45.indices[2][1]] / 1000) km") +axflux_55 = CairoMakie.Axis(fig[3, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_55.grid.yᵃᶜᵃ[flux_NN_data_55.indices[2][1]] / 1000) km") +axflux_65 = CairoMakie.Axis(fig[4, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_65.grid.yᵃᶜᵃ[flux_NN_data_65.indices[2][1]] / 1000) km") +axflux_75 = CairoMakie.Axis(fig[4, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_75.grid.yᵃᶜᵃ[flux_NN_data_75.indices[2][1]] / 1000) km") +axflux_85 = CairoMakie.Axis(fig[5, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_85.grid.yᵃᶜᵃ[flux_NN_data_85.indices[2][1]] / 1000) km") +axflux_95 = CairoMakie.Axis(fig[5, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_95.grid.yᵃᶜᵃ[flux_NN_data_95.indices[2][1]] / 1000) km") + +axfields = [axfield_5, axfield_15, axfield_25, axfield_35, axfield_45, axfield_55, axfield_65, axfield_75, axfield_85, axfield_95] +axfluxes = [axflux_5, axflux_15, axflux_25, axflux_35, axflux_45, axflux_55, axflux_65, axflux_75, axflux_85, axflux_95] + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lims = [(find_min(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices))) for field_data in field_datas] + +flux_lims = [(find_min(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices))) for flux_data in flux_datas] + +flux_lims = [(-maximum(abs, flux_lim), maximum(abs, flux_lim)) for flux_lim in flux_lims] + +NNₙs = [@lift interior(field_data[$n], :, 1, zC_indices) for field_data in field_datas] +fluxₙs = [@lift interior(flux_data[$n], :, 1, zF_indices) for flux_data in flux_datas] + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_5_surface = heatmap!(axfield_5, xC, zC[zC_indices], NN_5ₙ, colormap=colorscheme_field, colorrange=field_lims[1]) +field_15_surface = heatmap!(axfield_15, xC, zC[zC_indices], NN_15ₙ, colormap=colorscheme_field, colorrange=field_lims[2]) +field_25_surface = heatmap!(axfield_25, xC, zC[zC_indices], NN_25ₙ, colormap=colorscheme_field, colorrange=field_lims[3]) +field_35_surface = heatmap!(axfield_35, xC, zC[zC_indices], NN_35ₙ, colormap=colorscheme_field, colorrange=field_lims[4]) +field_45_surface = heatmap!(axfield_45, xC, zC[zC_indices], NN_45ₙ, colormap=colorscheme_field, colorrange=field_lims[5]) +field_55_surface = heatmap!(axfield_55, xC, zC[zC_indices], NN_55ₙ, colormap=colorscheme_field, colorrange=field_lims[6]) +field_65_surface = heatmap!(axfield_65, xC, zC[zC_indices], NN_65ₙ, colormap=colorscheme_field, colorrange=field_lims[7]) +field_75_surface = heatmap!(axfield_75, xC, zC[zC_indices], NN_75ₙ, colormap=colorscheme_field, colorrange=field_lims[8]) +field_85_surface = heatmap!(axfield_85, xC, zC[zC_indices], NN_85ₙ, colormap=colorscheme_field, colorrange=field_lims[9]) +field_95_surface = heatmap!(axfield_95, xC, zC[zC_indices], NN_95ₙ, colormap=colorscheme_field, colorrange=field_lims[10]) + +flux_5_surface = heatmap!(axflux_5, xC, zC[zF_indices], flux_5ₙ, colormap=colorscheme_flux, colorrange=flux_lims[1]) +flux_15_surface = heatmap!(axflux_15, xC, zC[zF_indices], flux_15ₙ, colormap=colorscheme_flux, colorrange=flux_lims[2]) +flux_25_surface = heatmap!(axflux_25, xC, zC[zF_indices], flux_25ₙ, colormap=colorscheme_flux, colorrange=flux_lims[3]) +flux_35_surface = heatmap!(axflux_35, xC, zC[zF_indices], flux_35ₙ, colormap=colorscheme_flux, colorrange=flux_lims[4]) +flux_45_surface = heatmap!(axflux_45, xC, zC[zF_indices], flux_45ₙ, colormap=colorscheme_flux, colorrange=flux_lims[5]) +flux_55_surface = heatmap!(axflux_55, xC, zC[zF_indices], flux_55ₙ, colormap=colorscheme_flux, colorrange=flux_lims[6]) +flux_65_surface = heatmap!(axflux_65, xC, zC[zF_indices], flux_65ₙ, colormap=colorscheme_flux, colorrange=flux_lims[7]) +flux_75_surface = heatmap!(axflux_75, xC, zC[zF_indices], flux_75ₙ, colormap=colorscheme_flux, colorrange=flux_lims[8]) +flux_85_surface = heatmap!(axflux_85, xC, zC[zF_indices], flux_85ₙ, colormap=colorscheme_flux, colorrange=flux_lims[9]) +flux_95_surface = heatmap!(axflux_95, xC, zC[zF_indices], flux_95ₙ, colormap=colorscheme_flux, colorrange=flux_lims[10]) + +Colorbar(fig[1, 2], field_5_surface, label="Field") +Colorbar(fig[1, 4], flux_5_surface, label="NN Flux") +Colorbar(fig[1, 6], field_15_surface, label="Field") +Colorbar(fig[1, 8], flux_15_surface, label="NN Flux") +Colorbar(fig[2, 2], field_25_surface, label="Field") +Colorbar(fig[2, 4], flux_25_surface, label="NN Flux") +Colorbar(fig[2, 6], field_35_surface, label="Field") +Colorbar(fig[2, 8], flux_35_surface, label="NN Flux") +Colorbar(fig[3, 2], field_45_surface, label="Field") +Colorbar(fig[3, 4], flux_45_surface, label="NN Flux") +Colorbar(fig[3, 6], field_55_surface, label="Field") +Colorbar(fig[3, 8], flux_55_surface, label="NN Flux") +Colorbar(fig[4, 2], field_65_surface, label="Field") +Colorbar(fig[4, 4], flux_65_surface, label="NN Flux") +Colorbar(fig[4, 6], field_75_surface, label="Field") +Colorbar(fig[4, 8], flux_75_surface, label="NN Flux") +Colorbar(fig[5, 2], field_85_surface, label="Field") +Colorbar(fig[5, 4], flux_85_surface, label="NN Flux") +Colorbar(fig[5, 6], field_95_surface, label="Field") +Colorbar(fig[5, 8], flux_95_surface, label="NN Flux") + +for axfield in axfields + xlims!(axfield, minimum(xC), maximum(xC)) + ylims!(axfield, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +end + +for axflux in axfluxes + xlims!(axflux, minimum(xC), maximum(xC)) + ylims!(axflux, minimum(zC[zF_indices]), maximum(zC[zF_indices])) +end + +title_str = @lift "Salinity (psu), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_xzslices_fluxes_S.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end #%% \ No newline at end of file From 93bf0bc4941d4c48f54d32256af18935113e90bf Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Fri, 1 Nov 2024 17:35:35 -0400 Subject: [PATCH 093/114] change default z advection scheme to weno 5 --- baseclosure_doublegyre_model_wallrestoration.jl | 8 ++++---- physicalclosure_doublegyre_model_wallrestoration.jl | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/baseclosure_doublegyre_model_wallrestoration.jl b/baseclosure_doublegyre_model_wallrestoration.jl index b675786dec..22e3f67a05 100644 --- a/baseclosure_doublegyre_model_wallrestoration.jl +++ b/baseclosure_doublegyre_model_wallrestoration.jl @@ -21,8 +21,8 @@ using SeawaterPolynomials:TEOS10 using ColorSchemes #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_wallrestoration_30days_zC2O_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" -# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" +# filename = "doublegyre_30Cwarmflushbottom10_relaxation_wallrestoration_30days_zC2O_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_wallrestoration_30days_zWENO5_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -33,8 +33,8 @@ model_architecture = GPU() base_closure = XinKaiLocalVerticalDiffusivity() closure = base_closure -advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) -# advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) +# advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) # number of grid points const Nx = 100 diff --git a/physicalclosure_doublegyre_model_wallrestoration.jl b/physicalclosure_doublegyre_model_wallrestoration.jl index e5130a0841..63921f0a42 100644 --- a/physicalclosure_doublegyre_model_wallrestoration.jl +++ b/physicalclosure_doublegyre_model_wallrestoration.jl @@ -27,8 +27,8 @@ using ColorSchemes using Glob #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_wallrestoration_30days_zC2O_CATKEVerticalDiffusivity" -# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_CATKEVerticalDiffusivity" +# filename = "doublegyre_30Cwarmflushbottom10_relaxation_wallrestoration_30days_zC2O_CATKEVerticalDiffusivity" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_wallrestoration_30days_zWENO5_CATKEVerticalDiffusivity" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -47,8 +47,8 @@ convection_closure = CATKE_ocean_closure() closure = convection_closure # closure = vertical_base_closure -advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) -# advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) +# advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) # number of grid points const Nx = 100 From d36677f4a48257138be5adfef4e1f63a8ecfcd83 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Wed, 6 Nov 2024 17:04:22 -0500 Subject: [PATCH 094/114] run with linear ramp seasonal forcing --- ...losure_doublegyre_model_seasonalforcing.jl | 58 ++++++++++++++++-- ...e_model_seasonalforcing_wallrestoration.jl | 60 +++++++++++++++++-- ...losure_doublegyre_model_seasonalforcing.jl | 58 ++++++++++++++++-- ...e_model_seasonalforcing_wallrestoration.jl | 60 +++++++++++++++++-- 4 files changed, 218 insertions(+), 18 deletions(-) diff --git a/baseclosure_doublegyre_model_seasonalforcing.jl b/baseclosure_doublegyre_model_seasonalforcing.jl index 5ffeb56cf6..f9719ed68b 100644 --- a/baseclosure_doublegyre_model_seasonalforcing.jl +++ b/baseclosure_doublegyre_model_seasonalforcing.jl @@ -21,7 +21,7 @@ using SeawaterPolynomials:TEOS10 using ColorSchemes #%% -filename = "doublegyre_seasonalforcing_30C-20C_relaxation_8days_baseclosure_test" +filename = "doublegyre_linearseasonalforcing_10C_relaxation_30days_baseclosure" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -68,11 +68,11 @@ const S_mid = (S_north + S_south) / 2 const τ₀ = 1e-4 const μ_drag = 1/30days -const μ_T = 1/8days +const μ_T = 1/30days const seasonal_period = 360days const seasonal_forcing_width = Ly / 6 -const seasonal_T_amplitude = 15 +const seasonal_T_amplitude = 10 ##### ##### Forcing and initial condition @@ -99,7 +99,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), east = ValueBoundaryCondition(0), west = ValueBoundaryCondition(0)) -@inline T_seasonal(y, t) = seasonal_T_amplitude * exp(-y^2/(2 * seasonal_forcing_width^2)) * sin(2π * t / seasonal_period) +@inline T_seasonal(y, t) = seasonal_T_amplitude * (y/Ly + 1/2) * sin(2π * t / seasonal_period) @inline T_ref(y, t) = T_mid - ΔT / Ly * y + T_seasonal(y, t) @inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y, t)) surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) @@ -300,6 +300,56 @@ simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, indices = (90, :, :), schedule = TimeInterval(10days)) +simulation.output_writers[:xz_5] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_5", + indices = (:, 5, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_15] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_15", + indices = (:, 15, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_25] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_25", + indices = (:, 25, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_35] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_35", + indices = (:, 35, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_45] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_45", + indices = (:, 45, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_55] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_55", + indices = (:, 55, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_65] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_65", + indices = (:, 65, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_75] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_75", + indices = (:, 75, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_85] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_85", + indices = (:, 85, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_95] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_95", + indices = (:, 95, :), + schedule = TimeInterval(10days)) + simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, filename = "$(FILE_DIR)/averaged_fields_zonal", schedule = TimeInterval(10days)) diff --git a/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl b/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl index d665d645d5..187b88dc22 100644 --- a/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl +++ b/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl @@ -21,7 +21,7 @@ using SeawaterPolynomials:TEOS10 using ColorSchemes #%% -filename = "doublegyre_seasonalforcing_30C-20C_relaxation_wallrestoration_8days_baseclosure_test" +filename = "doublegyre_linearseasonalforcing_10C_relaxation_wallrestoration_30days_baseclosure" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -70,11 +70,11 @@ const S_mid = (S_north + S_south) / 2 const τ₀ = 1e-4 const μ_drag = 1/30days -const μ_T = 1/8days +const μ_T = 1/30days const seasonal_period = 360days const seasonal_forcing_width = Ly / 6 -const seasonal_T_amplitude = 20 +const seasonal_T_amplitude = 10 ##### ##### Forcing and initial condition @@ -101,12 +101,12 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), east = ValueBoundaryCondition(0), west = ValueBoundaryCondition(0)) -@inline T_seasonal(y, t) = seasonal_T_amplitude * exp(-y^2/(2 * seasonal_forcing_width^2)) * sin(2π * t / seasonal_period) +@inline T_north_ref(z) = min(0, -5 + 5 * (1 + (z + 500) / (Lz - 500))) @inline T_ref(y, t) = T_mid - ΔT / Ly * y + T_seasonal(y, t) @inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y, t)) surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) -@inline T_north_ref(z) = 10 * (1 + z / Lz) +@inline T_north_ref(z) = min(0, -5 + 5 * (1 + (z + 500) / (Lz - 500))) @inline north_T_flux(x, z, t, T) = μ_T * δy * (T - T_north_ref(z)) north_T_flux_bc = FluxBoundaryCondition(north_T_flux; field_dependencies=:T) @@ -307,6 +307,56 @@ simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, indices = (90, :, :), schedule = TimeInterval(10days)) +simulation.output_writers[:xz_5] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_5", + indices = (:, 5, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_15] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_15", + indices = (:, 15, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_25] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_25", + indices = (:, 25, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_35] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_35", + indices = (:, 35, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_45] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_45", + indices = (:, 45, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_55] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_55", + indices = (:, 55, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_65] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_65", + indices = (:, 65, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_75] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_75", + indices = (:, 75, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_85] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_85", + indices = (:, 85, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_95] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_95", + indices = (:, 95, :), + schedule = TimeInterval(10days)) + simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, filename = "$(FILE_DIR)/averaged_fields_zonal", schedule = TimeInterval(10days)) diff --git a/physicalclosure_doublegyre_model_seasonalforcing.jl b/physicalclosure_doublegyre_model_seasonalforcing.jl index 4c692e8cb9..7e811d1c3b 100644 --- a/physicalclosure_doublegyre_model_seasonalforcing.jl +++ b/physicalclosure_doublegyre_model_seasonalforcing.jl @@ -28,7 +28,7 @@ using Glob #%% -filename = "doublegyre_seasonalforcing_30C-20C_relaxation_8days_CATKEVerticalDiffusivity" +filename = "doublegyre_linearseasonalforcing_10C_relaxation_30days_CATKEVerticalDiffusivity" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -82,11 +82,11 @@ const S_mid = (S_north + S_south) / 2 const τ₀ = 1e-4 const μ_drag = 1/30days -const μ_T = 1/8days +const μ_T = 1/30days const seasonal_period = 360days const seasonal_forcing_width = Ly / 6 -const seasonal_T_amplitude = 15 +const seasonal_T_amplitude = 10 ##### ##### Forcing and initial condition @@ -113,7 +113,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), east = ValueBoundaryCondition(0), west = ValueBoundaryCondition(0)) -@inline T_seasonal(y, t) = seasonal_T_amplitude * exp(-y^2/(2 * seasonal_forcing_width^2)) * sin(2π * t / seasonal_period) +@inline T_seasonal(y, t) = seasonal_T_amplitude * (y/Ly + 1/2) * sin(2π * t / seasonal_period) @inline T_ref(y, t) = T_mid - ΔT / Ly * y + T_seasonal(y, t) @inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y, t)) surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) @@ -319,6 +319,56 @@ simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, indices = (90, :, :), schedule = TimeInterval(10days)) +simulation.output_writers[:xz_5] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_5", + indices = (:, 5, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_15] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_15", + indices = (:, 15, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_25] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_25", + indices = (:, 25, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_35] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_35", + indices = (:, 35, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_45] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_45", + indices = (:, 45, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_55] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_55", + indices = (:, 55, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_65] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_65", + indices = (:, 65, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_75] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_75", + indices = (:, 75, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_85] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_85", + indices = (:, 85, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_95] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_95", + indices = (:, 95, :), + schedule = TimeInterval(10days)) + simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, filename = "$(FILE_DIR)/averaged_fields_zonal", schedule = TimeInterval(10days)) diff --git a/physicalclosure_doublegyre_model_seasonalforcing_wallrestoration.jl b/physicalclosure_doublegyre_model_seasonalforcing_wallrestoration.jl index 355b8317f7..04dd63ba0d 100644 --- a/physicalclosure_doublegyre_model_seasonalforcing_wallrestoration.jl +++ b/physicalclosure_doublegyre_model_seasonalforcing_wallrestoration.jl @@ -28,7 +28,7 @@ using Glob #%% -filename = "doublegyre_seasonalforcing_30C-20C_relaxation_wallrestoration_8days_CATKEVerticalDiffusivity" +filename = "doublegyre_linearseasonalforcing_10C_relaxation_wallrestoration_30days_CATKEVerticalDiffusivity" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -84,11 +84,11 @@ const S_mid = (S_north + S_south) / 2 const τ₀ = 1e-4 const μ_drag = 1/30days -const μ_T = 1/8days +const μ_T = 1/30days const seasonal_period = 360days const seasonal_forcing_width = Ly / 6 -const seasonal_T_amplitude = 20 +const seasonal_T_amplitude = 10 ##### ##### Forcing and initial condition @@ -115,12 +115,12 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), east = ValueBoundaryCondition(0), west = ValueBoundaryCondition(0)) -@inline T_seasonal(y, t) = seasonal_T_amplitude * exp(-y^2/(2 * seasonal_forcing_width^2)) * sin(2π * t / seasonal_period) +@inline T_seasonal(y, t) = seasonal_T_amplitude * (y/Ly + 1/2) * sin(2π * t / seasonal_period) @inline T_ref(y, t) = T_mid - ΔT / Ly * y + T_seasonal(y, t) @inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y, t)) surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) -@inline T_north_ref(z) = 10 * (1 + z / Lz) +@inline T_north_ref(z) = min(0, -5 + 5 * (1 + (z + 500) / (Lz - 500))) @inline north_T_flux(x, z, t, T) = μ_T * Δy * (T - T_north_ref(z)) north_T_flux_bc = FluxBoundaryCondition(north_T_flux; field_dependencies=:T) @@ -326,6 +326,56 @@ simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, indices = (90, :, :), schedule = TimeInterval(10days)) +simulation.output_writers[:xz_5] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_5", + indices = (:, 5, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_15] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_15", + indices = (:, 15, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_25] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_25", + indices = (:, 25, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_35] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_35", + indices = (:, 35, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_45] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_45", + indices = (:, 45, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_55] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_55", + indices = (:, 55, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_65] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_65", + indices = (:, 65, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_75] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_75", + indices = (:, 75, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_85] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_85", + indices = (:, 85, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_95] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_95", + indices = (:, 95, :), + schedule = TimeInterval(10days)) + simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, filename = "$(FILE_DIR)/averaged_fields_zonal", schedule = TimeInterval(10days)) From ce68cc79b9f0a3e11f68107e0865e6dc49947e6f Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 7 Nov 2024 12:19:26 -0500 Subject: [PATCH 095/114] fix T_seasonal --- ...losure_doublegyre_model_seasonalforcing_wallrestoration.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl b/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl index 187b88dc22..96528f677c 100644 --- a/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl +++ b/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl @@ -100,8 +100,8 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), bottom = v_drag_bc, east = ValueBoundaryCondition(0), west = ValueBoundaryCondition(0)) - -@inline T_north_ref(z) = min(0, -5 + 5 * (1 + (z + 500) / (Lz - 500))) + +@inline T_seasonal(y, t) = seasonal_T_amplitude * (y/Ly + 1/2) * sin(2π * t / seasonal_period) @inline T_ref(y, t) = T_mid - ΔT / Ly * y + T_seasonal(y, t) @inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y, t)) surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) From e1603caebc3063c6bf015307050add68c630abc8 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Fri, 8 Nov 2024 13:40:46 -0500 Subject: [PATCH 096/114] fix temperature restoration --- ...blegyre_model_seasonalforcing_wallrestoration.jl | 13 ++++++------- ...blegyre_model_seasonalforcing_wallrestoration.jl | 11 +++++------ 2 files changed, 11 insertions(+), 13 deletions(-) diff --git a/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl b/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl index 96528f677c..c484faa1af 100644 --- a/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl +++ b/baseclosure_doublegyre_model_seasonalforcing_wallrestoration.jl @@ -58,7 +58,7 @@ grid = RectilinearGrid(model_architecture, Float64, ##### ##### Boundary conditions ##### -const T_north = 10 +const T_north = 0 const T_south = 30 const T_mid = (T_north + T_south) / 2 const ΔT = T_south - T_north @@ -100,7 +100,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), bottom = v_drag_bc, east = ValueBoundaryCondition(0), west = ValueBoundaryCondition(0)) - + @inline T_seasonal(y, t) = seasonal_T_amplitude * (y/Ly + 1/2) * sin(2π * t / seasonal_period) @inline T_ref(y, t) = T_mid - ΔT / Ly * y + T_seasonal(y, t) @inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y, t)) @@ -475,11 +475,10 @@ ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) -colorscheme = colorschemes[:balance] -T_colormap = colorscheme -S_colormap = colorscheme -u_colormap = colorscheme -v_colormap = colorscheme +T_colormap = colorschemes[:viridis] +S_colormap = colorschemes[:viridis] +u_colormap = colorschemes[:balance] +v_colormap = colorschemes[:balance] T_color_range = Tlim S_color_range = Slim diff --git a/physicalclosure_doublegyre_model_seasonalforcing_wallrestoration.jl b/physicalclosure_doublegyre_model_seasonalforcing_wallrestoration.jl index 04dd63ba0d..b0efdf5f82 100644 --- a/physicalclosure_doublegyre_model_seasonalforcing_wallrestoration.jl +++ b/physicalclosure_doublegyre_model_seasonalforcing_wallrestoration.jl @@ -72,7 +72,7 @@ grid = RectilinearGrid(model_architecture, Float64, ##### ##### Boundary conditions ##### -const T_north = 10 +const T_north = 0 const T_south = 30 const T_mid = (T_north + T_south) / 2 const ΔT = T_south - T_north @@ -490,11 +490,10 @@ ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) -colorscheme = colorschemes[:balance] -T_colormap = colorscheme -S_colormap = colorscheme -u_colormap = colorscheme -v_colormap = colorscheme +T_colormap = colorschemes[:viridis] +S_colormap = colorschemes[:viridis] +u_colormap = colorschemes[:balance] +v_colormap = colorschemes[:balance] T_color_range = Tlim S_color_range = Slim From a3ab871ded9e0b5aadef9c4b8900ece7d071ac12 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 21 Nov 2024 15:51:55 -0500 Subject: [PATCH 097/114] implement different NNclosure with Ri zone --- ...ntaneous_fields_slices_BBLkappazonelast.jl | 316 ++++++++++++++++++ validate_NN_nof_BBLRifirstzone510_1D_model.jl | 247 ++++++++++++++ 2 files changed, 563 insertions(+) create mode 100644 compare_3D_instantaneous_fields_slices_BBLkappazonelast.jl create mode 100644 validate_NN_nof_BBLRifirstzone510_1D_model.jl diff --git a/compare_3D_instantaneous_fields_slices_BBLkappazonelast.jl b/compare_3D_instantaneous_fields_slices_BBLkappazonelast.jl new file mode 100644 index 0000000000..0695fc815f --- /dev/null +++ b/compare_3D_instantaneous_fields_slices_BBLkappazonelast.jl @@ -0,0 +1,316 @@ +using GLMakie +using Oceananigans +using ColorSchemes +using SeawaterPolynomials +using SeawaterPolynomials.TEOS10 + +NN_FILE_DIR = "./Output/doublegyre_30Cwarmflushbottom10_relaxation_8days_NN_closure_BBLkappazonelast41_temp" +CATKE_FILE_DIR = "./Output/doublegyre_30Cwarmflushbottom10_relaxation_8days_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" + +fieldname = "S" +ρ_NN_data_00 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz.jld2", fieldname, backend=OnDisk()) +ρ_NN_data_10 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_10.jld2", fieldname, backend=OnDisk()) +ρ_NN_data_20 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_20.jld2", fieldname, backend=OnDisk()) +ρ_NN_data_30 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_30.jld2", fieldname, backend=OnDisk()) +ρ_NN_data_40 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_40.jld2", fieldname, backend=OnDisk()) +ρ_NN_data_50 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_50.jld2", fieldname, backend=OnDisk()) +ρ_NN_data_60 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_60.jld2", fieldname, backend=OnDisk()) +ρ_NN_data_70 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_70.jld2", fieldname, backend=OnDisk()) +ρ_NN_data_80 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_80.jld2", fieldname, backend=OnDisk()) +ρ_NN_data_90 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_90.jld2", fieldname, backend=OnDisk()) + +ρ_CATKE_data_00 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz.jld2", fieldname, backend=OnDisk()) +ρ_CATKE_data_10 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_10.jld2", fieldname, backend=OnDisk()) +ρ_CATKE_data_20 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_20.jld2", fieldname, backend=OnDisk()) +ρ_CATKE_data_30 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_30.jld2", fieldname, backend=OnDisk()) +ρ_CATKE_data_40 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_40.jld2", fieldname, backend=OnDisk()) +ρ_CATKE_data_50 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_50.jld2", fieldname, backend=OnDisk()) +ρ_CATKE_data_60 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_60.jld2", fieldname, backend=OnDisk()) +ρ_CATKE_data_70 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_70.jld2", fieldname, backend=OnDisk()) +ρ_CATKE_data_80 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_80.jld2", fieldname, backend=OnDisk()) +ρ_CATKE_data_90 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_90.jld2", fieldname, backend=OnDisk()) + +# ρ_NN_data_00 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz.jld2", fieldname) +# ρ_NN_data_10 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_10.jld2", fieldname) +# ρ_NN_data_20 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_20.jld2", fieldname) +# ρ_NN_data_30 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_30.jld2", fieldname) +# ρ_NN_data_40 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_40.jld2", fieldname) +# ρ_NN_data_50 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_50.jld2", fieldname) +# ρ_NN_data_60 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_60.jld2", fieldname) +# ρ_NN_data_70 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_70.jld2", fieldname) +# ρ_NN_data_80 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_80.jld2", fieldname) +# ρ_NN_data_90 = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_yz_90.jld2", fieldname) + +# ρ_CATKE_data_00 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz.jld2", fieldname) +# ρ_CATKE_data_10 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_10.jld2", fieldname) +# ρ_CATKE_data_20 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_20.jld2", fieldname) +# ρ_CATKE_data_30 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_30.jld2", fieldname) +# ρ_CATKE_data_40 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_40.jld2", fieldname) +# ρ_CATKE_data_50 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_50.jld2", fieldname) +# ρ_CATKE_data_60 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_60.jld2", fieldname) +# ρ_CATKE_data_70 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_70.jld2", fieldname) +# ρ_CATKE_data_80 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_80.jld2", fieldname) +# ρ_CATKE_data_90 = FieldTimeSeries("$(CATKE_FILE_DIR)/instantaneous_fields_yz_90.jld2", fieldname) + +first_index_data = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_NN_active_diagnostics.jld2", "first_index") +last_index_data = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_NN_active_diagnostics.jld2", "last_index") +Qb_data = FieldTimeSeries("$(NN_FILE_DIR)/instantaneous_fields_NN_active_diagnostics.jld2", "Qb") + +Nx, Ny, Nz = ρ_NN_data_00.grid.Nx, ρ_NN_data_00.grid.Ny, ρ_NN_data_00.grid.Nz + +xC = ρ_NN_data_00.grid.xᶜᵃᵃ[1:ρ_NN_data_00.grid.Nx] +yC = ρ_NN_data_00.grid.yᵃᶜᵃ[1:ρ_NN_data_00.grid.Ny] +zC = ρ_NN_data_00.grid.zᵃᵃᶜ[1:ρ_NN_data_00.grid.Nz] +zF = ρ_NN_data_00.grid.zᵃᵃᶠ[1:ρ_NN_data_00.grid.Nz+1] + +Nt = length(ρ_NN_data_90) +times = ρ_NN_data_00.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(2400, 2400)) + +axNN_00 = GLMakie.Axis(fig[1, 1], xlabel="y (m)", ylabel="z (m)", title="NN, x = $(ρ_NN_data_00.grid.xᶜᵃᵃ[ρ_NN_data_00.indices[1][1]] / 1000) km") +axNN_10 = GLMakie.Axis(fig[1, 3], xlabel="y (m)", ylabel="z (m)", title="NN, x = $(ρ_NN_data_10.grid.xᶜᵃᵃ[ρ_NN_data_10.indices[1][1]] / 1000) km") +axNN_20 = GLMakie.Axis(fig[2, 1], xlabel="y (m)", ylabel="z (m)", title="NN, x = $(ρ_NN_data_20.grid.xᶜᵃᵃ[ρ_NN_data_20.indices[1][1]] / 1000) km") +axNN_30 = GLMakie.Axis(fig[2, 3], xlabel="y (m)", ylabel="z (m)", title="NN, x = $(ρ_NN_data_30.grid.xᶜᵃᵃ[ρ_NN_data_30.indices[1][1]] / 1000) km") +axNN_40 = GLMakie.Axis(fig[3, 1], xlabel="y (m)", ylabel="z (m)", title="NN, x = $(ρ_NN_data_40.grid.xᶜᵃᵃ[ρ_NN_data_40.indices[1][1]] / 1000) km") +axNN_50 = GLMakie.Axis(fig[3, 3], xlabel="y (m)", ylabel="z (m)", title="NN, x = $(ρ_NN_data_50.grid.xᶜᵃᵃ[ρ_NN_data_50.indices[1][1]] / 1000) km") +axNN_60 = GLMakie.Axis(fig[4, 1], xlabel="y (m)", ylabel="z (m)", title="NN, x = $(ρ_NN_data_60.grid.xᶜᵃᵃ[ρ_NN_data_60.indices[1][1]] / 1000) km") +axNN_70 = GLMakie.Axis(fig[4, 3], xlabel="y (m)", ylabel="z (m)", title="NN, x = $(ρ_NN_data_70.grid.xᶜᵃᵃ[ρ_NN_data_70.indices[1][1]] / 1000) km") +axNN_80 = GLMakie.Axis(fig[5, 1], xlabel="y (m)", ylabel="z (m)", title="NN, x = $(ρ_NN_data_80.grid.xᶜᵃᵃ[ρ_NN_data_80.indices[1][1]] / 1000) km") +axNN_90 = GLMakie.Axis(fig[5, 3], xlabel="y (m)", ylabel="z (m)", title="NN, x = $(ρ_NN_data_90.grid.xᶜᵃᵃ[ρ_NN_data_90.indices[1][1]] / 1000) km") + +axCATKE_00 = GLMakie.Axis(fig[1, 2], xlabel="y (m)", ylabel="z (m)", title="Base Closure, x = $(ρ_CATKE_data_00.grid.xᶜᵃᵃ[ρ_CATKE_data_00.indices[1][1]] / 1000) km") +axCATKE_10 = GLMakie.Axis(fig[1, 4], xlabel="y (m)", ylabel="z (m)", title="Base Closure, x = $(ρ_CATKE_data_10.grid.xᶜᵃᵃ[ρ_CATKE_data_10.indices[1][1]] / 1000) km") +axCATKE_20 = GLMakie.Axis(fig[2, 2], xlabel="y (m)", ylabel="z (m)", title="Base Closure, x = $(ρ_CATKE_data_20.grid.xᶜᵃᵃ[ρ_CATKE_data_20.indices[1][1]] / 1000) km") +axCATKE_30 = GLMakie.Axis(fig[2, 4], xlabel="y (m)", ylabel="z (m)", title="Base Closure, x = $(ρ_CATKE_data_30.grid.xᶜᵃᵃ[ρ_CATKE_data_30.indices[1][1]] / 1000) km") +axCATKE_40 = GLMakie.Axis(fig[3, 2], xlabel="y (m)", ylabel="z (m)", title="Base Closure, x = $(ρ_CATKE_data_40.grid.xᶜᵃᵃ[ρ_CATKE_data_40.indices[1][1]] / 1000) km") +axCATKE_50 = GLMakie.Axis(fig[3, 4], xlabel="y (m)", ylabel="z (m)", title="Base Closure, x = $(ρ_CATKE_data_50.grid.xᶜᵃᵃ[ρ_CATKE_data_50.indices[1][1]] / 1000) km") +axCATKE_60 = GLMakie.Axis(fig[4, 2], xlabel="y (m)", ylabel="z (m)", title="Base Closure, x = $(ρ_CATKE_data_60.grid.xᶜᵃᵃ[ρ_CATKE_data_60.indices[1][1]] / 1000) km") +axCATKE_70 = GLMakie.Axis(fig[4, 4], xlabel="y (m)", ylabel="z (m)", title="Base Closure, x = $(ρ_CATKE_data_70.grid.xᶜᵃᵃ[ρ_CATKE_data_70.indices[1][1]] / 1000) km") +axCATKE_80 = GLMakie.Axis(fig[5, 2], xlabel="y (m)", ylabel="z (m)", title="Base Closure, x = $(ρ_CATKE_data_80.grid.xᶜᵃᵃ[ρ_CATKE_data_80.indices[1][1]] / 1000) km") +axCATKE_90 = GLMakie.Axis(fig[5, 4], xlabel="y (m)", ylabel="z (m)", title="Base Closure, x = $(ρ_CATKE_data_90.grid.xᶜᵃᵃ[ρ_CATKE_data_90.indices[1][1]] / 1000) km") + +n = Observable(1096) + +z_indices = 1:200 + +ρlim = (find_min(interior(ρ_NN_data_00[timeframes[1]], :, :, z_indices), interior(ρ_NN_data_00[timeframes[end]], :, :, z_indices), interior(ρ_CATKE_data_00[timeframes[1]], :, :, z_indices), interior(ρ_CATKE_data_00[timeframes[end]], :, :, z_indices)), + find_max(interior(ρ_NN_data_00[timeframes[1]], :, :, z_indices), interior(ρ_NN_data_00[timeframes[end]], :, :, z_indices), interior(ρ_CATKE_data_00[timeframes[1]], :, :, z_indices), interior(ρ_CATKE_data_00[timeframes[end]], :, :, z_indices))) + +NN_00ₙ = @lift interior(ρ_NN_data_00[$n], 1, :, z_indices) +NN_10ₙ = @lift interior(ρ_NN_data_10[$n], 1, :, z_indices) +NN_20ₙ = @lift interior(ρ_NN_data_20[$n], 1, :, z_indices) +NN_30ₙ = @lift interior(ρ_NN_data_30[$n], 1, :, z_indices) +NN_40ₙ = @lift interior(ρ_NN_data_40[$n], 1, :, z_indices) +NN_50ₙ = @lift interior(ρ_NN_data_50[$n], 1, :, z_indices) +NN_60ₙ = @lift interior(ρ_NN_data_60[$n], 1, :, z_indices) +NN_70ₙ = @lift interior(ρ_NN_data_70[$n], 1, :, z_indices) +NN_80ₙ = @lift interior(ρ_NN_data_80[$n], 1, :, z_indices) +NN_90ₙ = @lift interior(ρ_NN_data_90[$n], 1, :, z_indices) + +CATKE_00ₙ = @lift interior(ρ_CATKE_data_00[$n], 1, :, z_indices) +CATKE_10ₙ = @lift interior(ρ_CATKE_data_10[$n], 1, :, z_indices) +CATKE_20ₙ = @lift interior(ρ_CATKE_data_20[$n], 1, :, z_indices) +CATKE_30ₙ = @lift interior(ρ_CATKE_data_30[$n], 1, :, z_indices) +CATKE_40ₙ = @lift interior(ρ_CATKE_data_40[$n], 1, :, z_indices) +CATKE_50ₙ = @lift interior(ρ_CATKE_data_50[$n], 1, :, z_indices) +CATKE_60ₙ = @lift interior(ρ_CATKE_data_60[$n], 1, :, z_indices) +CATKE_70ₙ = @lift interior(ρ_CATKE_data_70[$n], 1, :, z_indices) +CATKE_80ₙ = @lift interior(ρ_CATKE_data_80[$n], 1, :, z_indices) +CATKE_90ₙ = @lift interior(ρ_CATKE_data_90[$n], 1, :, z_indices) + +zs_first_index_00ₙ = @lift zF[Int.(vec(interior(first_index_data[$n], ρ_NN_data_00.indices[1][1], :, :)))] +zs_first_index_10ₙ = @lift zF[Int.(vec(interior(first_index_data[$n], ρ_NN_data_10.indices[1][1], :, :)))] +zs_first_index_20ₙ = @lift zF[Int.(vec(interior(first_index_data[$n], ρ_NN_data_20.indices[1][1], :, :)))] +zs_first_index_30ₙ = @lift zF[Int.(vec(interior(first_index_data[$n], ρ_NN_data_30.indices[1][1], :, :)))] +zs_first_index_40ₙ = @lift zF[Int.(vec(interior(first_index_data[$n], ρ_NN_data_40.indices[1][1], :, :)))] +zs_first_index_50ₙ = @lift zF[Int.(vec(interior(first_index_data[$n], ρ_NN_data_50.indices[1][1], :, :)))] +zs_first_index_60ₙ = @lift zF[Int.(vec(interior(first_index_data[$n], ρ_NN_data_60.indices[1][1], :, :)))] +zs_first_index_70ₙ = @lift zF[Int.(vec(interior(first_index_data[$n], ρ_NN_data_70.indices[1][1], :, :)))] +zs_first_index_80ₙ = @lift zF[Int.(vec(interior(first_index_data[$n], ρ_NN_data_80.indices[1][1], :, :)))] +zs_first_index_90ₙ = @lift zF[Int.(vec(interior(first_index_data[$n], ρ_NN_data_90.indices[1][1], :, :)))] + +zs_last_index_00ₙ = @lift zF[Int.(vec(interior(last_index_data[$n], ρ_NN_data_00.indices[1][1], :, :)))] +zs_last_index_10ₙ = @lift zF[Int.(vec(interior(last_index_data[$n], ρ_NN_data_10.indices[1][1], :, :)))] +zs_last_index_20ₙ = @lift zF[Int.(vec(interior(last_index_data[$n], ρ_NN_data_20.indices[1][1], :, :)))] +zs_last_index_30ₙ = @lift zF[Int.(vec(interior(last_index_data[$n], ρ_NN_data_30.indices[1][1], :, :)))] +zs_last_index_40ₙ = @lift zF[Int.(vec(interior(last_index_data[$n], ρ_NN_data_40.indices[1][1], :, :)))] +zs_last_index_50ₙ = @lift zF[Int.(vec(interior(last_index_data[$n], ρ_NN_data_50.indices[1][1], :, :)))] +zs_last_index_60ₙ = @lift zF[Int.(vec(interior(last_index_data[$n], ρ_NN_data_60.indices[1][1], :, :)))] +zs_last_index_70ₙ = @lift zF[Int.(vec(interior(last_index_data[$n], ρ_NN_data_70.indices[1][1], :, :)))] +zs_last_index_80ₙ = @lift zF[Int.(vec(interior(last_index_data[$n], ρ_NN_data_80.indices[1][1], :, :)))] +zs_last_index_90ₙ = @lift zF[Int.(vec(interior(last_index_data[$n], ρ_NN_data_90.indices[1][1], :, :)))] + +ys_convection_00ₙ = @lift yC[interior(Qb_data[$n], ρ_NN_data_00.indices[1][1], :, 1) .> 0] +ys_convection_10ₙ = @lift yC[interior(Qb_data[$n], ρ_NN_data_10.indices[1][1], :, 1) .> 0] +ys_convection_20ₙ = @lift yC[interior(Qb_data[$n], ρ_NN_data_20.indices[1][1], :, 1) .> 0] +ys_convection_30ₙ = @lift yC[interior(Qb_data[$n], ρ_NN_data_30.indices[1][1], :, 1) .> 0] +ys_convection_40ₙ = @lift yC[interior(Qb_data[$n], ρ_NN_data_40.indices[1][1], :, 1) .> 0] +ys_convection_50ₙ = @lift yC[interior(Qb_data[$n], ρ_NN_data_50.indices[1][1], :, 1) .> 0] +ys_convection_60ₙ = @lift yC[interior(Qb_data[$n], ρ_NN_data_60.indices[1][1], :, 1) .> 0] +ys_convection_70ₙ = @lift yC[interior(Qb_data[$n], ρ_NN_data_70.indices[1][1], :, 1) .> 0] +ys_convection_80ₙ = @lift yC[interior(Qb_data[$n], ρ_NN_data_80.indices[1][1], :, 1) .> 0] +ys_convection_90ₙ = @lift yC[interior(Qb_data[$n], ρ_NN_data_90.indices[1][1], :, 1) .> 0] + +zs_convection_00ₙ = @lift fill(zC[1], length($ys_convection_00ₙ)) +zs_convection_10ₙ = @lift fill(zC[1], length($ys_convection_10ₙ)) +zs_convection_20ₙ = @lift fill(zC[1], length($ys_convection_20ₙ)) +zs_convection_30ₙ = @lift fill(zC[1], length($ys_convection_30ₙ)) +zs_convection_40ₙ = @lift fill(zC[1], length($ys_convection_40ₙ)) +zs_convection_50ₙ = @lift fill(zC[1], length($ys_convection_50ₙ)) +zs_convection_60ₙ = @lift fill(zC[1], length($ys_convection_60ₙ)) +zs_convection_70ₙ = @lift fill(zC[1], length($ys_convection_70ₙ)) +zs_convection_80ₙ = @lift fill(zC[1], length($ys_convection_80ₙ)) +zs_convection_90ₙ = @lift fill(zC[1], length($ys_convection_90ₙ)) + +# colorscheme = Reverse(colorschemes[:jet]) +colorscheme = colorschemes[:jet] + +NN_00_surface = heatmap!(axNN_00, yC, zC[z_indices], NN_00ₙ, colormap=colorscheme, colorrange=ρlim) +NN_10_surface = heatmap!(axNN_10, yC, zC[z_indices], NN_10ₙ, colormap=colorscheme, colorrange=ρlim) +NN_20_surface = heatmap!(axNN_20, yC, zC[z_indices], NN_20ₙ, colormap=colorscheme, colorrange=ρlim) +NN_30_surface = heatmap!(axNN_30, yC, zC[z_indices], NN_30ₙ, colormap=colorscheme, colorrange=ρlim) +NN_40_surface = heatmap!(axNN_40, yC, zC[z_indices], NN_40ₙ, colormap=colorscheme, colorrange=ρlim) +NN_50_surface = heatmap!(axNN_50, yC, zC[z_indices], NN_50ₙ, colormap=colorscheme, colorrange=ρlim) +NN_60_surface = heatmap!(axNN_60, yC, zC[z_indices], NN_60ₙ, colormap=colorscheme, colorrange=ρlim) +NN_70_surface = heatmap!(axNN_70, yC, zC[z_indices], NN_70ₙ, colormap=colorscheme, colorrange=ρlim) +NN_80_surface = heatmap!(axNN_80, yC, zC[z_indices], NN_80ₙ, colormap=colorscheme, colorrange=ρlim) +NN_90_surface = heatmap!(axNN_90, yC, zC[z_indices], NN_90ₙ, colormap=colorscheme, colorrange=ρlim) + +CATKE_00_surface = heatmap!(axCATKE_00, yC, zC[z_indices], CATKE_00ₙ, colormap=colorscheme, colorrange=ρlim) +CATKE_10_surface = heatmap!(axCATKE_10, yC, zC[z_indices], CATKE_10ₙ, colormap=colorscheme, colorrange=ρlim) +CATKE_20_surface = heatmap!(axCATKE_20, yC, zC[z_indices], CATKE_20ₙ, colormap=colorscheme, colorrange=ρlim) +CATKE_30_surface = heatmap!(axCATKE_30, yC, zC[z_indices], CATKE_30ₙ, colormap=colorscheme, colorrange=ρlim) +CATKE_40_surface = heatmap!(axCATKE_40, yC, zC[z_indices], CATKE_40ₙ, colormap=colorscheme, colorrange=ρlim) +CATKE_50_surface = heatmap!(axCATKE_50, yC, zC[z_indices], CATKE_50ₙ, colormap=colorscheme, colorrange=ρlim) +CATKE_60_surface = heatmap!(axCATKE_60, yC, zC[z_indices], CATKE_60ₙ, colormap=colorscheme, colorrange=ρlim) +CATKE_70_surface = heatmap!(axCATKE_70, yC, zC[z_indices], CATKE_70ₙ, colormap=colorscheme, colorrange=ρlim) +CATKE_80_surface = heatmap!(axCATKE_80, yC, zC[z_indices], CATKE_80ₙ, colormap=colorscheme, colorrange=ρlim) +CATKE_90_surface = heatmap!(axCATKE_90, yC, zC[z_indices], CATKE_90ₙ, colormap=colorscheme, colorrange=ρlim) + +# contourlevels = range(ρlim[1], ρlim[2], length=10) + +# NN_00_surface = contourf!(axNN_00, yC, zC, NN_00ₙ, colormap=colorscheme, levels=contourlevels) +# NN_10_surface = contourf!(axNN_10, yC, zC, NN_10ₙ, colormap=colorscheme, levels=contourlevels) +# NN_20_surface = contourf!(axNN_20, yC, zC, NN_20ₙ, colormap=colorscheme, levels=contourlevels) +# NN_30_surface = contourf!(axNN_30, yC, zC, NN_30ₙ, colormap=colorscheme, levels=contourlevels) +# NN_40_surface = contourf!(axNN_40, yC, zC, NN_40ₙ, colormap=colorscheme, levels=contourlevels) +# NN_50_surface = contourf!(axNN_50, yC, zC, NN_50ₙ, colormap=colorscheme, levels=contourlevels) +# NN_60_surface = contourf!(axNN_60, yC, zC, NN_60ₙ, colormap=colorscheme, levels=contourlevels) +# NN_70_surface = contourf!(axNN_70, yC, zC, NN_70ₙ, colormap=colorscheme, levels=contourlevels) +# NN_80_surface = contourf!(axNN_80, yC, zC, NN_80ₙ, colormap=colorscheme, levels=contourlevels) +# NN_90_surface = contourf!(axNN_90, yC, zC, NN_90ₙ, colormap=colorscheme, levels=contourlevels) + +# CATKE_00_surface = contourf!(axCATKE_00, yC, zC, CATKE_00ₙ, colormap=colorscheme, levels=contourlevels) +# CATKE_10_surface = contourf!(axCATKE_10, yC, zC, CATKE_10ₙ, colormap=colorscheme, levels=contourlevels) +# CATKE_20_surface = contourf!(axCATKE_20, yC, zC, CATKE_20ₙ, colormap=colorscheme, levels=contourlevels) +# CATKE_30_surface = contourf!(axCATKE_30, yC, zC, CATKE_30ₙ, colormap=colorscheme, levels=contourlevels) +# CATKE_40_surface = contourf!(axCATKE_40, yC, zC, CATKE_40ₙ, colormap=colorscheme, levels=contourlevels) +# CATKE_50_surface = contourf!(axCATKE_50, yC, zC, CATKE_50ₙ, colormap=colorscheme, levels=contourlevels) +# CATKE_60_surface = contourf!(axCATKE_60, yC, zC, CATKE_60ₙ, colormap=colorscheme, levels=contourlevels) +# CATKE_70_surface = contourf!(axCATKE_70, yC, zC, CATKE_70ₙ, colormap=colorscheme, levels=contourlevels) +# CATKE_80_surface = contourf!(axCATKE_80, yC, zC, CATKE_80ₙ, colormap=colorscheme, levels=contourlevels) +# CATKE_90_surface = contourf!(axCATKE_90, yC, zC, CATKE_90ₙ, colormap=colorscheme, levels=contourlevels) + +Colorbar(fig[1:5, 5], NN_00_surface) + +# lines!(axNN_00, yC, zs_first_index_00ₙ, color=:black) +# lines!(axNN_10, yC, zs_first_index_10ₙ, color=:black) +# lines!(axNN_20, yC, zs_first_index_20ₙ, color=:black) +# lines!(axNN_30, yC, zs_first_index_30ₙ, color=:black) +# lines!(axNN_40, yC, zs_first_index_40ₙ, color=:black) +# lines!(axNN_50, yC, zs_first_index_50ₙ, color=:black) +# lines!(axNN_60, yC, zs_first_index_60ₙ, color=:black) +# lines!(axNN_70, yC, zs_first_index_70ₙ, color=:black) +# lines!(axNN_80, yC, zs_first_index_80ₙ, color=:black) +# lines!(axNN_90, yC, zs_first_index_90ₙ, color=:black) + +# lines!(axNN_00, yC, zs_last_index_00ₙ, color=:black) +# lines!(axNN_10, yC, zs_last_index_10ₙ, color=:black) +# lines!(axNN_20, yC, zs_last_index_20ₙ, color=:black) +# lines!(axNN_30, yC, zs_last_index_30ₙ, color=:black) +# lines!(axNN_40, yC, zs_last_index_40ₙ, color=:black) +# lines!(axNN_50, yC, zs_last_index_50ₙ, color=:black) +# lines!(axNN_60, yC, zs_last_index_60ₙ, color=:black) +# lines!(axNN_70, yC, zs_last_index_70ₙ, color=:black) +# lines!(axNN_80, yC, zs_last_index_80ₙ, color=:black) +# lines!(axNN_90, yC, zs_last_index_90ₙ, color=:black) + +scatter!(axNN_00, ys_convection_00ₙ, zs_convection_00ₙ, color=:red, markersize=10) +scatter!(axNN_10, ys_convection_10ₙ, zs_convection_10ₙ, color=:red, markersize=10) +scatter!(axNN_20, ys_convection_20ₙ, zs_convection_20ₙ, color=:red, markersize=10) +scatter!(axNN_30, ys_convection_30ₙ, zs_convection_30ₙ, color=:red, markersize=10) +scatter!(axNN_40, ys_convection_40ₙ, zs_convection_40ₙ, color=:red, markersize=10) +scatter!(axNN_50, ys_convection_50ₙ, zs_convection_50ₙ, color=:red, markersize=10) +scatter!(axNN_60, ys_convection_60ₙ, zs_convection_60ₙ, color=:red, markersize=10) +scatter!(axNN_70, ys_convection_70ₙ, zs_convection_70ₙ, color=:red, markersize=10) +scatter!(axNN_80, ys_convection_80ₙ, zs_convection_80ₙ, color=:red, markersize=10) +scatter!(axNN_90, ys_convection_90ₙ, zs_convection_90ₙ, color=:red, markersize=10) + +xlims!(axNN_00, minimum(yC), maximum(yC)) +xlims!(axNN_10, minimum(yC), maximum(yC)) +xlims!(axNN_20, minimum(yC), maximum(yC)) +xlims!(axNN_30, minimum(yC), maximum(yC)) +xlims!(axNN_40, minimum(yC), maximum(yC)) +xlims!(axNN_50, minimum(yC), maximum(yC)) +xlims!(axNN_60, minimum(yC), maximum(yC)) +xlims!(axNN_70, minimum(yC), maximum(yC)) +xlims!(axNN_80, minimum(yC), maximum(yC)) +xlims!(axNN_90, minimum(yC), maximum(yC)) + +ylims!(axNN_00, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axNN_10, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axNN_20, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axNN_30, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axNN_40, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axNN_50, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axNN_60, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axNN_70, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axNN_80, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axNN_90, minimum(zC[z_indices]), maximum(zC[z_indices])) + +xlims!(axCATKE_00, minimum(yC), maximum(yC)) +xlims!(axCATKE_10, minimum(yC), maximum(yC)) +xlims!(axCATKE_20, minimum(yC), maximum(yC)) +xlims!(axCATKE_30, minimum(yC), maximum(yC)) +xlims!(axCATKE_40, minimum(yC), maximum(yC)) +xlims!(axCATKE_50, minimum(yC), maximum(yC)) +xlims!(axCATKE_60, minimum(yC), maximum(yC)) +xlims!(axCATKE_70, minimum(yC), maximum(yC)) +xlims!(axCATKE_80, minimum(yC), maximum(yC)) +xlims!(axCATKE_90, minimum(yC), maximum(yC)) + +ylims!(axCATKE_00, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axCATKE_10, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axCATKE_20, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axCATKE_30, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axCATKE_40, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axCATKE_50, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axCATKE_60, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axCATKE_70, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axCATKE_80, minimum(zC[z_indices]), maximum(zC[z_indices])) +ylims!(axCATKE_90, minimum(zC[z_indices]), maximum(zC[z_indices])) + +# title_str = @lift "Temperature (°C), Time = $(round(times[$n], digits=2)) years" +title_str = @lift "Salinity (psu), Time = $(round(times[$n], digits=2)) years" +# title_str = @lift "Potential Density (kg m⁻³), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +display(fig) +GLMakie.record(fig, "./Output/doublegyre_relaxation_8days_30Cwarmflush10bottom_NNclosure_BBLkappazonelast41_baseclosure_S_BBLlines_yzslices.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end +#%% \ No newline at end of file diff --git a/validate_NN_nof_BBLRifirstzone510_1D_model.jl b/validate_NN_nof_BBLRifirstzone510_1D_model.jl new file mode 100644 index 0000000000..57c94e3937 --- /dev/null +++ b/validate_NN_nof_BBLRifirstzone510_1D_model.jl @@ -0,0 +1,247 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("NN_closure_global_nof_BBLRifirstzone510.jl") +include("xin_kai_vertical_diffusivity_local_2step.jl") +include("feature_scaling.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using Oceananigans.TimeSteppers: update_state! +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 + + +# Architecture +model_architecture = CPU() + +file = jldopen("model_inference_run_nof_BBLRifirstzone510.jld2", "r") + +# number of grid points +const Nz = file["Nz"] +const Lz = file["Lz"] + +grid = RectilinearGrid(model_architecture, + topology = (Flat, Flat, Bounded), + size = Nz, + halo = 3, + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const dTdz = file["dTdz"] +const dSdz = file["dSdz"] + +const T_surface = file["T_surface"] +const S_surface = file["S_surface"] + +T_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(file["wT_top"])) +S_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(file["wS_top"])) + +##### +##### Coriolis +##### + +const f₀ = file["f₀"] +coriolis = FPlane(f=f₀) + +##### +##### Forcing and initial condition +##### +T_initial(z) = dTdz * z + T_surface +S_initial(z) = dSdz * z + S_surface + +nn_closure = NNFluxClosure(model_architecture) +base_closure = XinKaiLocalVerticalDiffusivity() + +##### +##### Model building +##### + +@info "Building a model..." + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = ImplicitFreeSurface(), + momentum_advection = WENO(grid = grid), + tracer_advection = WENO(grid = grid), + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = (base_closure, nn_closure), + tracers = (:T, :S), + boundary_conditions = (; T = T_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(z) = T_initial(z) + 1e-6 * noise(z) +S_initial_noisy(z) = S_initial(z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = file["Δt"] +stop_time = file["τ"] + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(20)) + +##### +##### Diagnostics +##### + +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +wT_residual, wS_residual = model.diffusivity_fields[2].wT, model.diffusivity_fields[2].wS +ν, κ = model.diffusivity_fields[1].κᵘ, model.diffusivity_fields[1].κᶜ + +Tbar = Field(Average(T, dims = (1,2))) +Sbar = Field(Average(S, dims = (1,2))) + +averaged_outputs = (; Tbar, Sbar, wT_residual, wS_residual, ν, κ) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:jld2] = JLD2OutputWriter(model, averaged_outputs, + filename = "NN_1D_channel_averages_nof_BBLRifirstzone510", + schedule = TimeInterval(Δt₀), + overwrite_existing = true) + +@info "Running the simulation..." + +try + run!(simulation, pickup = false) +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +# ##### +# ##### Visualization +# ##### +#%% +using GLMakie + +Tbar_data = FieldTimeSeries("./NN_1D_channel_averages_nof_BBLRifirstzone510.jld2", "Tbar") +Sbar_data = FieldTimeSeries("./NN_1D_channel_averages_nof_BBLRifirstzone510.jld2", "Sbar") +wT_residual_data = FieldTimeSeries("./NN_1D_channel_averages_nof_BBLRifirstzone510.jld2", "wT_residual") +wS_residual_data = FieldTimeSeries("./NN_1D_channel_averages_nof_BBLRifirstzone510.jld2", "wS_residual") +ν_data = FieldTimeSeries("./NN_1D_channel_averages_nof_BBLRifirstzone510.jld2", "ν") +κ_data = FieldTimeSeries("./NN_1D_channel_averages_nof_BBLRifirstzone510.jld2", "κ") + +#%% +zC = znodes(Tbar_data.grid, Center()) +zF = znodes(Tbar_data.grid, Face()) + +Nt = length(Tbar_data.times) + +fig = Figure(size = (1500, 1000)) +axT = GLMakie.Axis(fig[1, 1], xlabel = "T (°C)", ylabel = "z (m)") +axS = GLMakie.Axis(fig[2, 1], xlabel = "S (g kg⁻¹)", ylabel = "z (m)") +axwT_residual = GLMakie.Axis(fig[1, 2], xlabel = "wT residual", ylabel = "z (m)") +axwS_residual = GLMakie.Axis(fig[2, 2], xlabel = "wS residual", ylabel = "z (m)") +axν = GLMakie.Axis(fig[1, 3], xlabel = "ν (m² s⁻¹)", ylabel = "z (m)", xscale=log10) +axκ = GLMakie.Axis(fig[2, 3], xlabel = "κ (m² s⁻¹)", ylabel = "z (m)", xscale=log10) + +slider = Slider(fig[3, :], range=2:Nt) +n = slider.value + +Tbarₙ = @lift interior(Tbar_data[$n], 1, 1, :) +Sbarₙ = @lift interior(Sbar_data[$n], 1, 1, :) +wT_residualₙ = @lift interior(wT_residual_data[$n], 1, 1, :) +wS_residualₙ = @lift interior(wS_residual_data[$n], 1, 1, :) +νₙ = @lift interior(ν_data[$n], 1, 1, 2:32) +κₙ = @lift interior(κ_data[$n], 1, 1, 2:32) + +Tbar_truthₙ = @lift file["sol_T"][:, $n] +Sbar_truthₙ = @lift file["sol_S"][:, $n] +wT_residual_truthₙ = @lift file["sol_wT_residual_unscaled"][:, $n] +wS_residual_truthₙ = @lift file["sol_wS_residual_unscaled"][:, $n] +ν_truthₙ = @lift file["sol_ν"][2:32, $n] +κ_truthₙ = @lift file["sol_κ"][2:32, $n] + +title_str = @lift "Time: $(round(Tbar_data.times[$n] / 86400, digits=3)) days" + +wTlim = (minimum(interior(wT_residual_data)), maximum(interior(wT_residual_data))) +wSlim = (minimum(interior(wS_residual_data)), maximum(interior(wS_residual_data))) + +νlim = (1e-6, 10) +κlim = (1e-6, 10) + +lines!(axT, Tbarₙ, zC, label="Oceananigans") +lines!(axS, Sbarₙ, zC, label="Oceananigans") + +lines!(axwT_residual, wT_residualₙ, zF, label="Oceananigans") +lines!(axwS_residual, wS_residualₙ, zF, label="Oceananigans") + +lines!(axν, νₙ, zF[2:32], label="Oceananigans") +lines!(axκ, κₙ, zF[2:32], label="Oceananigans") + +lines!(axT, Tbar_truthₙ, zC, label="Truth") +lines!(axS, Sbar_truthₙ, zC, label="Truth") + +lines!(axwT_residual, wT_residual_truthₙ, zF, label="Truth") +lines!(axwS_residual, wS_residual_truthₙ, zF, label="Truth") + +lines!(axν, ν_truthₙ, zF[2:32], label="Truth") +lines!(axκ, κ_truthₙ, zF[2:32], label="Truth") + +xlims!(axwT_residual, wTlim) +xlims!(axwS_residual, wSlim) +xlims!(axν, νlim) +xlims!(axκ, κlim) + +linkyaxes!(axT, axS, axwT_residual, axwS_residual, axν, axκ) + +axislegend(axT, position = :lb) +Label(fig[0, :], title_str, tellwidth = false) + +# GLMakie.record(fig, "./NN_1D_validation_nof_BBL.mp4", 1:Nt, framerate=60, px_per_unit=4) do nn +# @info nn +# n[] = nn +# end + +display(fig) +#%% +close(file) \ No newline at end of file From 51acb9d125b66bc892d34c3eca5938c2bc2fe627 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 21 Nov 2024 15:57:38 -0500 Subject: [PATCH 098/114] run double gyre on updated closure --- ..._nof_BBLRifirstzone510_doublegyre_model.jl | 1265 +++++++++++++++++ 1 file changed, 1265 insertions(+) create mode 100644 NNclosure_nof_BBLRifirstzone510_doublegyre_model.jl diff --git a/NNclosure_nof_BBLRifirstzone510_doublegyre_model.jl b/NNclosure_nof_BBLRifirstzone510_doublegyre_model.jl new file mode 100644 index 0000000000..64af3bda2f --- /dev/null +++ b/NNclosure_nof_BBLRifirstzone510_doublegyre_model.jl @@ -0,0 +1,1265 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("NN_closure_global_nof_BBLRifirstzone510.jl") +include("xin_kai_vertical_diffusivity_local_2step.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + + +#%% +# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zC2O_NN_closure_NDE_BBLRifirztzone510" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_NN_closure_NDE_BBLRifirztzone510" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +nn_closure = NNFluxClosure(model_architecture) +base_closure = XinKaiLocalVerticalDiffusivity() +closure = (base_closure, nn_closure) + +# advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/30days + +##### +##### Forcing and initial condition +##### +@inline T_initial(x, y, z) = 10 + 20 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 7300days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); +first_index = model.diffusivity_fields[2].first_index +last_index = model.diffusivity_fields[2].last_index +wT_NN = model.diffusivity_fields[2].wT +wS_NN = model.diffusivity_fields[2].wS + +κ = model.diffusivity_fields[1].κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) + +wT = wT_NN + wT_base +wS = wS_NN + wS_base + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +@inline function get_top_buoyancy_flux(i, j, k, grid, buoyancy, T_bc, S_bc, velocities, tracers, clock) + return top_buoyancy_flux(i, j, grid, buoyancy, (; T=T_bc, S=S_bc), clock, merge(velocities, tracers)) +end + +Qb = KernelFunctionOperation{Center, Center, Nothing}(get_top_buoyancy_flux, model.grid, model.buoyancy, T.boundary_conditions.top, S.boundary_conditions.top, model.velocities, model.tracers, model.clock) +Qb = Field(Qb) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + +wT_NNbar_zonal = Average(wT_NN, dims=1) +wS_NNbar_zonal = Average(wS_NN, dims=1) + +wT_basebar_zonal = Average(wT_base, dims=1) +wS_basebar_zonal = Average(wS_base, dims=1) + +wTbar_zonal = Average(wT, dims=1) +wSbar_zonal = Average(wS, dims=1) + +outputs = (; u, v, w, T, S, ρ, N², wT_NN, wS_NN, wT_base, wS_base, wT, wS) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal, wT_NNbar_zonal, wS_NNbar_zonal, wT_basebar_zonal, wS_basebar_zonal, wTbar_zonal, wSbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_5] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_5", + indices = (:, 5, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_15] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_15", + indices = (:, 15, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_25] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_25", + indices = (:, 25, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_35] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_35", + indices = (:, 35, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_45] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_45", + indices = (:, 45, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_55] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_55", + indices = (:, 55, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_65] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_65", + indices = (:, 65, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_75] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_75", + indices = (:, 75, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_85] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_85", + indices = (:, 85, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_95] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_95", + indices = (:, 95, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; first_index, last_index, Qb), + filename = "$(FILE_DIR)/instantaneous_fields_NN_active_diagnostics", + indices = (:, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end + +#%% +@info "Recording T fields and fluxes in yz" + +fieldname = "T" +fluxname = "wT_NN" +field_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fieldname, backend=OnDisk()) +field_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fieldname, backend=OnDisk()) +field_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fieldname, backend=OnDisk()) +field_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fieldname, backend=OnDisk()) +field_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fieldname, backend=OnDisk()) +field_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fieldname, backend=OnDisk()) +field_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fieldname, backend=OnDisk()) +field_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fieldname, backend=OnDisk()) +field_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fieldname, backend=OnDisk()) +field_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fluxname, backend=OnDisk()) +flux_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fluxname, backend=OnDisk()) +flux_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fluxname, backend=OnDisk()) +flux_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fluxname, backend=OnDisk()) +flux_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fluxname, backend=OnDisk()) +flux_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fluxname, backend=OnDisk()) +flux_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fluxname, backend=OnDisk()) +flux_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fluxname, backend=OnDisk()) +flux_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fluxname, backend=OnDisk()) +flux_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fluxname, backend=OnDisk()) + +xC = field_NN_data_00.grid.xᶜᵃᵃ[1:field_NN_data_00.grid.Nx] +yC = field_NN_data_00.grid.yᵃᶜᵃ[1:field_NN_data_00.grid.Ny] +zC = field_NN_data_00.grid.zᵃᵃᶜ[1:field_NN_data_00.grid.Nz] +zF = field_NN_data_00.grid.zᵃᵃᶠ[1:field_NN_data_00.grid.Nz+1] + +Nt = length(field_NN_data_90) +times = field_NN_data_00.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_00 = CairoMakie.Axis(fig[1, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_00.grid.xᶜᵃᵃ[field_NN_data_00.indices[1][1]] / 1000) km") +axfield_10 = CairoMakie.Axis(fig[1, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_10.grid.xᶜᵃᵃ[field_NN_data_10.indices[1][1]] / 1000) km") +axfield_20 = CairoMakie.Axis(fig[2, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_20.grid.xᶜᵃᵃ[field_NN_data_20.indices[1][1]] / 1000) km") +axfield_30 = CairoMakie.Axis(fig[2, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_30.grid.xᶜᵃᵃ[field_NN_data_30.indices[1][1]] / 1000) km") +axfield_40 = CairoMakie.Axis(fig[3, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_40.grid.xᶜᵃᵃ[field_NN_data_40.indices[1][1]] / 1000) km") +axfield_50 = CairoMakie.Axis(fig[3, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_50.grid.xᶜᵃᵃ[field_NN_data_50.indices[1][1]] / 1000) km") +axfield_60 = CairoMakie.Axis(fig[4, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_60.grid.xᶜᵃᵃ[field_NN_data_60.indices[1][1]] / 1000) km") +axfield_70 = CairoMakie.Axis(fig[4, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_70.grid.xᶜᵃᵃ[field_NN_data_70.indices[1][1]] / 1000) km") +axfield_80 = CairoMakie.Axis(fig[5, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_80.grid.xᶜᵃᵃ[field_NN_data_80.indices[1][1]] / 1000) km") +axfield_90 = CairoMakie.Axis(fig[5, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_90.grid.xᶜᵃᵃ[field_NN_data_90.indices[1][1]] / 1000) km") + +axflux_00 = CairoMakie.Axis(fig[1, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_00.grid.xᶜᵃᵃ[flux_NN_data_00.indices[1][1]] / 1000) km") +axflux_10 = CairoMakie.Axis(fig[1, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_10.grid.xᶜᵃᵃ[flux_NN_data_10.indices[1][1]] / 1000) km") +axflux_20 = CairoMakie.Axis(fig[2, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_20.grid.xᶜᵃᵃ[flux_NN_data_20.indices[1][1]] / 1000) km") +axflux_30 = CairoMakie.Axis(fig[2, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_30.grid.xᶜᵃᵃ[flux_NN_data_30.indices[1][1]] / 1000) km") +axflux_40 = CairoMakie.Axis(fig[3, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_40.grid.xᶜᵃᵃ[flux_NN_data_40.indices[1][1]] / 1000) km") +axflux_50 = CairoMakie.Axis(fig[3, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_50.grid.xᶜᵃᵃ[flux_NN_data_50.indices[1][1]] / 1000) km") +axflux_60 = CairoMakie.Axis(fig[4, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_60.grid.xᶜᵃᵃ[flux_NN_data_60.indices[1][1]] / 1000) km") +axflux_70 = CairoMakie.Axis(fig[4, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_70.grid.xᶜᵃᵃ[flux_NN_data_70.indices[1][1]] / 1000) km") +axflux_80 = CairoMakie.Axis(fig[5, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_80.grid.xᶜᵃᵃ[flux_NN_data_80.indices[1][1]] / 1000) km") +axflux_90 = CairoMakie.Axis(fig[5, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_90.grid.xᶜᵃᵃ[flux_NN_data_90.indices[1][1]] / 1000) km") + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lim = (find_min(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (find_min(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (-maximum(abs, flux_lim), maximum(abs, flux_lim)) + +NN_00ₙ = @lift interior(field_NN_data_00[$n], 1, :, zC_indices) +NN_10ₙ = @lift interior(field_NN_data_10[$n], 1, :, zC_indices) +NN_20ₙ = @lift interior(field_NN_data_20[$n], 1, :, zC_indices) +NN_30ₙ = @lift interior(field_NN_data_30[$n], 1, :, zC_indices) +NN_40ₙ = @lift interior(field_NN_data_40[$n], 1, :, zC_indices) +NN_50ₙ = @lift interior(field_NN_data_50[$n], 1, :, zC_indices) +NN_60ₙ = @lift interior(field_NN_data_60[$n], 1, :, zC_indices) +NN_70ₙ = @lift interior(field_NN_data_70[$n], 1, :, zC_indices) +NN_80ₙ = @lift interior(field_NN_data_80[$n], 1, :, zC_indices) +NN_90ₙ = @lift interior(field_NN_data_90[$n], 1, :, zC_indices) + +flux_00ₙ = @lift interior(flux_NN_data_00[$n], 1, :, zF_indices) +flux_10ₙ = @lift interior(flux_NN_data_10[$n], 1, :, zF_indices) +flux_20ₙ = @lift interior(flux_NN_data_20[$n], 1, :, zF_indices) +flux_30ₙ = @lift interior(flux_NN_data_30[$n], 1, :, zF_indices) +flux_40ₙ = @lift interior(flux_NN_data_40[$n], 1, :, zF_indices) +flux_50ₙ = @lift interior(flux_NN_data_50[$n], 1, :, zF_indices) +flux_60ₙ = @lift interior(flux_NN_data_60[$n], 1, :, zF_indices) +flux_70ₙ = @lift interior(flux_NN_data_70[$n], 1, :, zF_indices) +flux_80ₙ = @lift interior(flux_NN_data_80[$n], 1, :, zF_indices) +flux_90ₙ = @lift interior(flux_NN_data_90[$n], 1, :, zF_indices) + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_00_surface = heatmap!(axfield_00, yC, zC[zC_indices], NN_00ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_10_surface = heatmap!(axfield_10, yC, zC[zC_indices], NN_10ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_20_surface = heatmap!(axfield_20, yC, zC[zC_indices], NN_20ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_30_surface = heatmap!(axfield_30, yC, zC[zC_indices], NN_30ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_40_surface = heatmap!(axfield_40, yC, zC[zC_indices], NN_40ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_50_surface = heatmap!(axfield_50, yC, zC[zC_indices], NN_50ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_60_surface = heatmap!(axfield_60, yC, zC[zC_indices], NN_60ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_70_surface = heatmap!(axfield_70, yC, zC[zC_indices], NN_70ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_80_surface = heatmap!(axfield_80, yC, zC[zC_indices], NN_80ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_90_surface = heatmap!(axfield_90, yC, zC[zC_indices], NN_90ₙ, colormap=colorscheme_field, colorrange=field_lim) + +flux_00_surface = heatmap!(axflux_00, yC, zC[zF_indices], flux_00ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_10_surface = heatmap!(axflux_10, yC, zC[zF_indices], flux_10ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_20_surface = heatmap!(axflux_20, yC, zC[zF_indices], flux_20ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_30_surface = heatmap!(axflux_30, yC, zC[zF_indices], flux_30ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_40_surface = heatmap!(axflux_40, yC, zC[zF_indices], flux_40ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_50_surface = heatmap!(axflux_50, yC, zC[zF_indices], flux_50ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_60_surface = heatmap!(axflux_60, yC, zC[zF_indices], flux_60ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_70_surface = heatmap!(axflux_70, yC, zC[zF_indices], flux_70ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_80_surface = heatmap!(axflux_80, yC, zC[zF_indices], flux_80ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_90_surface = heatmap!(axflux_90, yC, zC[zF_indices], flux_90ₙ, colormap=colorscheme_flux, colorrange=flux_lim) + +Colorbar(fig[1:5, 2], field_00_surface, label="Field") +Colorbar(fig[1:5, 4], flux_00_surface, label="NN Flux") +Colorbar(fig[1:5, 6], field_00_surface, label="Field") +Colorbar(fig[1:5, 8], flux_00_surface, label="NN Flux") + +xlims!(axfield_00, minimum(yC), maximum(yC)) +xlims!(axfield_10, minimum(yC), maximum(yC)) +xlims!(axfield_20, minimum(yC), maximum(yC)) +xlims!(axfield_30, minimum(yC), maximum(yC)) +xlims!(axfield_40, minimum(yC), maximum(yC)) +xlims!(axfield_50, minimum(yC), maximum(yC)) +xlims!(axfield_60, minimum(yC), maximum(yC)) +xlims!(axfield_70, minimum(yC), maximum(yC)) +xlims!(axfield_80, minimum(yC), maximum(yC)) +xlims!(axfield_90, minimum(yC), maximum(yC)) + +ylims!(axfield_00, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_10, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_20, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_30, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_40, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_50, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_60, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_70, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_80, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_90, minimum(zC[zC_indices]), maximum(zC[zC_indices])) + +xlims!(axflux_00, minimum(yC), maximum(yC)) +xlims!(axflux_10, minimum(yC), maximum(yC)) +xlims!(axflux_20, minimum(yC), maximum(yC)) +xlims!(axflux_30, minimum(yC), maximum(yC)) +xlims!(axflux_40, minimum(yC), maximum(yC)) +xlims!(axflux_50, minimum(yC), maximum(yC)) +xlims!(axflux_60, minimum(yC), maximum(yC)) +xlims!(axflux_70, minimum(yC), maximum(yC)) +xlims!(axflux_80, minimum(yC), maximum(yC)) +xlims!(axflux_90, minimum(yC), maximum(yC)) + +ylims!(axflux_00, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_10, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_20, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_30, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_40, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_50, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_60, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_70, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_80, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_90, minimum(zF[zF_indices]), maximum(zF[zF_indices])) + +title_str = @lift "Temperature (°C), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +# save("./Output/compare_3D_instantaneous_fields_slices_NNclosure_fluxes.png", fig) +# display(fig) +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_yzslices_fluxes_T.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end +#%% +@info "Recording S fields and fluxes in yz" + +fieldname = "S" +fluxname = "wS_NN" +field_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fieldname, backend=OnDisk()) +field_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fieldname, backend=OnDisk()) +field_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fieldname, backend=OnDisk()) +field_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fieldname, backend=OnDisk()) +field_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fieldname, backend=OnDisk()) +field_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fieldname, backend=OnDisk()) +field_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fieldname, backend=OnDisk()) +field_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fieldname, backend=OnDisk()) +field_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fieldname, backend=OnDisk()) +field_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fluxname, backend=OnDisk()) +flux_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fluxname, backend=OnDisk()) +flux_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fluxname, backend=OnDisk()) +flux_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fluxname, backend=OnDisk()) +flux_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fluxname, backend=OnDisk()) +flux_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fluxname, backend=OnDisk()) +flux_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fluxname, backend=OnDisk()) +flux_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fluxname, backend=OnDisk()) +flux_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fluxname, backend=OnDisk()) +flux_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fluxname, backend=OnDisk()) + +xC = field_NN_data_00.grid.xᶜᵃᵃ[1:field_NN_data_00.grid.Nx] +yC = field_NN_data_00.grid.yᵃᶜᵃ[1:field_NN_data_00.grid.Ny] +zC = field_NN_data_00.grid.zᵃᵃᶜ[1:field_NN_data_00.grid.Nz] +zF = field_NN_data_00.grid.zᵃᵃᶠ[1:field_NN_data_00.grid.Nz+1] + +Nt = length(field_NN_data_90) +times = field_NN_data_00.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_00 = CairoMakie.Axis(fig[1, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_00.grid.xᶜᵃᵃ[field_NN_data_00.indices[1][1]] / 1000) km") +axfield_10 = CairoMakie.Axis(fig[1, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_10.grid.xᶜᵃᵃ[field_NN_data_10.indices[1][1]] / 1000) km") +axfield_20 = CairoMakie.Axis(fig[2, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_20.grid.xᶜᵃᵃ[field_NN_data_20.indices[1][1]] / 1000) km") +axfield_30 = CairoMakie.Axis(fig[2, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_30.grid.xᶜᵃᵃ[field_NN_data_30.indices[1][1]] / 1000) km") +axfield_40 = CairoMakie.Axis(fig[3, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_40.grid.xᶜᵃᵃ[field_NN_data_40.indices[1][1]] / 1000) km") +axfield_50 = CairoMakie.Axis(fig[3, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_50.grid.xᶜᵃᵃ[field_NN_data_50.indices[1][1]] / 1000) km") +axfield_60 = CairoMakie.Axis(fig[4, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_60.grid.xᶜᵃᵃ[field_NN_data_60.indices[1][1]] / 1000) km") +axfield_70 = CairoMakie.Axis(fig[4, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_70.grid.xᶜᵃᵃ[field_NN_data_70.indices[1][1]] / 1000) km") +axfield_80 = CairoMakie.Axis(fig[5, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_80.grid.xᶜᵃᵃ[field_NN_data_80.indices[1][1]] / 1000) km") +axfield_90 = CairoMakie.Axis(fig[5, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_90.grid.xᶜᵃᵃ[field_NN_data_90.indices[1][1]] / 1000) km") + +axflux_00 = CairoMakie.Axis(fig[1, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_00.grid.xᶜᵃᵃ[flux_NN_data_00.indices[1][1]] / 1000) km") +axflux_10 = CairoMakie.Axis(fig[1, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_10.grid.xᶜᵃᵃ[flux_NN_data_10.indices[1][1]] / 1000) km") +axflux_20 = CairoMakie.Axis(fig[2, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_20.grid.xᶜᵃᵃ[flux_NN_data_20.indices[1][1]] / 1000) km") +axflux_30 = CairoMakie.Axis(fig[2, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_30.grid.xᶜᵃᵃ[flux_NN_data_30.indices[1][1]] / 1000) km") +axflux_40 = CairoMakie.Axis(fig[3, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_40.grid.xᶜᵃᵃ[flux_NN_data_40.indices[1][1]] / 1000) km") +axflux_50 = CairoMakie.Axis(fig[3, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_50.grid.xᶜᵃᵃ[flux_NN_data_50.indices[1][1]] / 1000) km") +axflux_60 = CairoMakie.Axis(fig[4, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_60.grid.xᶜᵃᵃ[flux_NN_data_60.indices[1][1]] / 1000) km") +axflux_70 = CairoMakie.Axis(fig[4, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_70.grid.xᶜᵃᵃ[flux_NN_data_70.indices[1][1]] / 1000) km") +axflux_80 = CairoMakie.Axis(fig[5, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_80.grid.xᶜᵃᵃ[flux_NN_data_80.indices[1][1]] / 1000) km") +axflux_90 = CairoMakie.Axis(fig[5, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_90.grid.xᶜᵃᵃ[flux_NN_data_90.indices[1][1]] / 1000) km") + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lim = (find_min(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (find_min(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (-maximum(abs, flux_lim), maximum(abs, flux_lim)) + +NN_00ₙ = @lift interior(field_NN_data_00[$n], 1, :, zC_indices) +NN_10ₙ = @lift interior(field_NN_data_10[$n], 1, :, zC_indices) +NN_20ₙ = @lift interior(field_NN_data_20[$n], 1, :, zC_indices) +NN_30ₙ = @lift interior(field_NN_data_30[$n], 1, :, zC_indices) +NN_40ₙ = @lift interior(field_NN_data_40[$n], 1, :, zC_indices) +NN_50ₙ = @lift interior(field_NN_data_50[$n], 1, :, zC_indices) +NN_60ₙ = @lift interior(field_NN_data_60[$n], 1, :, zC_indices) +NN_70ₙ = @lift interior(field_NN_data_70[$n], 1, :, zC_indices) +NN_80ₙ = @lift interior(field_NN_data_80[$n], 1, :, zC_indices) +NN_90ₙ = @lift interior(field_NN_data_90[$n], 1, :, zC_indices) + +flux_00ₙ = @lift interior(flux_NN_data_00[$n], 1, :, zF_indices) +flux_10ₙ = @lift interior(flux_NN_data_10[$n], 1, :, zF_indices) +flux_20ₙ = @lift interior(flux_NN_data_20[$n], 1, :, zF_indices) +flux_30ₙ = @lift interior(flux_NN_data_30[$n], 1, :, zF_indices) +flux_40ₙ = @lift interior(flux_NN_data_40[$n], 1, :, zF_indices) +flux_50ₙ = @lift interior(flux_NN_data_50[$n], 1, :, zF_indices) +flux_60ₙ = @lift interior(flux_NN_data_60[$n], 1, :, zF_indices) +flux_70ₙ = @lift interior(flux_NN_data_70[$n], 1, :, zF_indices) +flux_80ₙ = @lift interior(flux_NN_data_80[$n], 1, :, zF_indices) +flux_90ₙ = @lift interior(flux_NN_data_90[$n], 1, :, zF_indices) + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_00_surface = heatmap!(axfield_00, yC, zC[zC_indices], NN_00ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_10_surface = heatmap!(axfield_10, yC, zC[zC_indices], NN_10ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_20_surface = heatmap!(axfield_20, yC, zC[zC_indices], NN_20ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_30_surface = heatmap!(axfield_30, yC, zC[zC_indices], NN_30ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_40_surface = heatmap!(axfield_40, yC, zC[zC_indices], NN_40ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_50_surface = heatmap!(axfield_50, yC, zC[zC_indices], NN_50ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_60_surface = heatmap!(axfield_60, yC, zC[zC_indices], NN_60ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_70_surface = heatmap!(axfield_70, yC, zC[zC_indices], NN_70ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_80_surface = heatmap!(axfield_80, yC, zC[zC_indices], NN_80ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_90_surface = heatmap!(axfield_90, yC, zC[zC_indices], NN_90ₙ, colormap=colorscheme_field, colorrange=field_lim) + +flux_00_surface = heatmap!(axflux_00, yC, zC[zF_indices], flux_00ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_10_surface = heatmap!(axflux_10, yC, zC[zF_indices], flux_10ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_20_surface = heatmap!(axflux_20, yC, zC[zF_indices], flux_20ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_30_surface = heatmap!(axflux_30, yC, zC[zF_indices], flux_30ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_40_surface = heatmap!(axflux_40, yC, zC[zF_indices], flux_40ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_50_surface = heatmap!(axflux_50, yC, zC[zF_indices], flux_50ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_60_surface = heatmap!(axflux_60, yC, zC[zF_indices], flux_60ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_70_surface = heatmap!(axflux_70, yC, zC[zF_indices], flux_70ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_80_surface = heatmap!(axflux_80, yC, zC[zF_indices], flux_80ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_90_surface = heatmap!(axflux_90, yC, zC[zF_indices], flux_90ₙ, colormap=colorscheme_flux, colorrange=flux_lim) + +Colorbar(fig[1:5, 2], field_00_surface, label="Field") +Colorbar(fig[1:5, 4], flux_00_surface, label="NN Flux") +Colorbar(fig[1:5, 6], field_00_surface, label="Field") +Colorbar(fig[1:5, 8], flux_00_surface, label="NN Flux") + +xlims!(axfield_00, minimum(yC), maximum(yC)) +xlims!(axfield_10, minimum(yC), maximum(yC)) +xlims!(axfield_20, minimum(yC), maximum(yC)) +xlims!(axfield_30, minimum(yC), maximum(yC)) +xlims!(axfield_40, minimum(yC), maximum(yC)) +xlims!(axfield_50, minimum(yC), maximum(yC)) +xlims!(axfield_60, minimum(yC), maximum(yC)) +xlims!(axfield_70, minimum(yC), maximum(yC)) +xlims!(axfield_80, minimum(yC), maximum(yC)) +xlims!(axfield_90, minimum(yC), maximum(yC)) + +ylims!(axfield_00, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_10, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_20, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_30, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_40, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_50, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_60, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_70, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_80, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_90, minimum(zC[zC_indices]), maximum(zC[zC_indices])) + +xlims!(axflux_00, minimum(yC), maximum(yC)) +xlims!(axflux_10, minimum(yC), maximum(yC)) +xlims!(axflux_20, minimum(yC), maximum(yC)) +xlims!(axflux_30, minimum(yC), maximum(yC)) +xlims!(axflux_40, minimum(yC), maximum(yC)) +xlims!(axflux_50, minimum(yC), maximum(yC)) +xlims!(axflux_60, minimum(yC), maximum(yC)) +xlims!(axflux_70, minimum(yC), maximum(yC)) +xlims!(axflux_80, minimum(yC), maximum(yC)) +xlims!(axflux_90, minimum(yC), maximum(yC)) + +ylims!(axflux_00, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_10, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_20, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_30, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_40, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_50, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_60, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_70, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_80, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_90, minimum(zF[zF_indices]), maximum(zF[zF_indices])) + +title_str = @lift "Salinity (psu), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_yzslices_fluxes_S.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +#%% +@info "Recording T fields and fluxes in xz" + +fieldname = "T" +fluxname = "wT_NN" + +field_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fieldname, backend=OnDisk()) +field_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fieldname, backend=OnDisk()) +field_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fieldname, backend=OnDisk()) +field_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fieldname, backend=OnDisk()) +field_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fieldname, backend=OnDisk()) +field_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fieldname, backend=OnDisk()) +field_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fieldname, backend=OnDisk()) +field_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fieldname, backend=OnDisk()) +field_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fieldname, backend=OnDisk()) +field_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fluxname, backend=OnDisk()) +flux_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fluxname, backend=OnDisk()) +flux_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fluxname, backend=OnDisk()) +flux_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fluxname, backend=OnDisk()) +flux_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fluxname, backend=OnDisk()) +flux_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fluxname, backend=OnDisk()) +flux_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fluxname, backend=OnDisk()) +flux_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fluxname, backend=OnDisk()) +flux_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fluxname, backend=OnDisk()) +flux_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fluxname, backend=OnDisk()) + +field_datas = [field_NN_data_5, field_NN_data_15, field_NN_data_25, field_NN_data_35, field_NN_data_45, field_NN_data_55, field_NN_data_65, field_NN_data_75, field_NN_data_85, field_NN_data_95] +flux_datas = [flux_NN_data_5, flux_NN_data_15, flux_NN_data_25, flux_NN_data_35, flux_NN_data_45, flux_NN_data_55, flux_NN_data_65, flux_NN_data_75, flux_NN_data_85, flux_NN_data_95] + +xC = field_NN_data_5.grid.xᶜᵃᵃ[1:field_NN_data_5.grid.Nx] +yC = field_NN_data_5.grid.yᵃᶜᵃ[1:field_NN_data_5.grid.Ny] +zC = field_NN_data_5.grid.zᵃᵃᶜ[1:field_NN_data_5.grid.Nz] +zF = field_NN_data_5.grid.zᵃᵃᶠ[1:field_NN_data_5.grid.Nz+1] + +Nt = length(field_NN_data_95) +times = field_NN_data_5.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_5 = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_5.grid.yᵃᶜᵃ[field_NN_data_5.indices[2][1]] / 1000) km") +axfield_15 = CairoMakie.Axis(fig[1, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_15.grid.yᵃᶜᵃ[field_NN_data_15.indices[2][1]] / 1000) km") +axfield_25 = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_25.grid.yᵃᶜᵃ[field_NN_data_25.indices[2][1]] / 1000) km") +axfield_35 = CairoMakie.Axis(fig[2, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_35.grid.yᵃᶜᵃ[field_NN_data_35.indices[2][1]] / 1000) km") +axfield_45 = CairoMakie.Axis(fig[3, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_45.grid.yᵃᶜᵃ[field_NN_data_45.indices[2][1]] / 1000) km") +axfield_55 = CairoMakie.Axis(fig[3, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_55.grid.yᵃᶜᵃ[field_NN_data_55.indices[2][1]] / 1000) km") +axfield_65 = CairoMakie.Axis(fig[4, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_65.grid.yᵃᶜᵃ[field_NN_data_65.indices[2][1]] / 1000) km") +axfield_75 = CairoMakie.Axis(fig[4, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_75.grid.yᵃᶜᵃ[field_NN_data_75.indices[2][1]] / 1000) km") +axfield_85 = CairoMakie.Axis(fig[5, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_85.grid.yᵃᶜᵃ[field_NN_data_85.indices[2][1]] / 1000) km") +axfield_95 = CairoMakie.Axis(fig[5, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_95.grid.yᵃᶜᵃ[field_NN_data_95.indices[2][1]] / 1000) km") + +axflux_5 = CairoMakie.Axis(fig[1, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_5.grid.yᵃᶜᵃ[flux_NN_data_5.indices[2][1]] / 1000) km") +axflux_15 = CairoMakie.Axis(fig[1, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_15.grid.yᵃᶜᵃ[flux_NN_data_15.indices[2][1]] / 1000) km") +axflux_25 = CairoMakie.Axis(fig[2, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_25.grid.yᵃᶜᵃ[flux_NN_data_25.indices[2][1]] / 1000) km") +axflux_35 = CairoMakie.Axis(fig[2, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_35.grid.yᵃᶜᵃ[flux_NN_data_35.indices[2][1]] / 1000) km") +axflux_45 = CairoMakie.Axis(fig[3, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_45.grid.yᵃᶜᵃ[flux_NN_data_45.indices[2][1]] / 1000) km") +axflux_55 = CairoMakie.Axis(fig[3, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_55.grid.yᵃᶜᵃ[flux_NN_data_55.indices[2][1]] / 1000) km") +axflux_65 = CairoMakie.Axis(fig[4, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_65.grid.yᵃᶜᵃ[flux_NN_data_65.indices[2][1]] / 1000) km") +axflux_75 = CairoMakie.Axis(fig[4, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_75.grid.yᵃᶜᵃ[flux_NN_data_75.indices[2][1]] / 1000) km") +axflux_85 = CairoMakie.Axis(fig[5, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_85.grid.yᵃᶜᵃ[flux_NN_data_85.indices[2][1]] / 1000) km") +axflux_95 = CairoMakie.Axis(fig[5, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_95.grid.yᵃᶜᵃ[flux_NN_data_95.indices[2][1]] / 1000) km") + +axfields = [axfield_5, axfield_15, axfield_25, axfield_35, axfield_45, axfield_55, axfield_65, axfield_75, axfield_85, axfield_95] +axfluxes = [axflux_5, axflux_15, axflux_25, axflux_35, axflux_45, axflux_55, axflux_65, axflux_75, axflux_85, axflux_95] + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lims = [(find_min(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices))) for field_data in field_datas] + +flux_lims = [(find_min(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices))) for flux_data in flux_datas] + +flux_lims = [(-maximum(abs, flux_lim), maximum(abs, flux_lim)) for flux_lim in flux_lims] + +NNₙs = [@lift interior(field_data[$n], :, 1, zC_indices) for field_data in field_datas] +fluxₙs = [@lift interior(flux_data[$n], :, 1, zF_indices) for flux_data in flux_datas] + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_5_surface = heatmap!(axfield_5, xC, zC[zC_indices], NN_5ₙ, colormap=colorscheme_field, colorrange=field_lims[1]) +field_15_surface = heatmap!(axfield_15, xC, zC[zC_indices], NN_15ₙ, colormap=colorscheme_field, colorrange=field_lims[2]) +field_25_surface = heatmap!(axfield_25, xC, zC[zC_indices], NN_25ₙ, colormap=colorscheme_field, colorrange=field_lims[3]) +field_35_surface = heatmap!(axfield_35, xC, zC[zC_indices], NN_35ₙ, colormap=colorscheme_field, colorrange=field_lims[4]) +field_45_surface = heatmap!(axfield_45, xC, zC[zC_indices], NN_45ₙ, colormap=colorscheme_field, colorrange=field_lims[5]) +field_55_surface = heatmap!(axfield_55, xC, zC[zC_indices], NN_55ₙ, colormap=colorscheme_field, colorrange=field_lims[6]) +field_65_surface = heatmap!(axfield_65, xC, zC[zC_indices], NN_65ₙ, colormap=colorscheme_field, colorrange=field_lims[7]) +field_75_surface = heatmap!(axfield_75, xC, zC[zC_indices], NN_75ₙ, colormap=colorscheme_field, colorrange=field_lims[8]) +field_85_surface = heatmap!(axfield_85, xC, zC[zC_indices], NN_85ₙ, colormap=colorscheme_field, colorrange=field_lims[9]) +field_95_surface = heatmap!(axfield_95, xC, zC[zC_indices], NN_95ₙ, colormap=colorscheme_field, colorrange=field_lims[10]) + +flux_5_surface = heatmap!(axflux_5, xC, zC[zF_indices], flux_5ₙ, colormap=colorscheme_flux, colorrange=flux_lims[1]) +flux_15_surface = heatmap!(axflux_15, xC, zC[zF_indices], flux_15ₙ, colormap=colorscheme_flux, colorrange=flux_lims[2]) +flux_25_surface = heatmap!(axflux_25, xC, zC[zF_indices], flux_25ₙ, colormap=colorscheme_flux, colorrange=flux_lims[3]) +flux_35_surface = heatmap!(axflux_35, xC, zC[zF_indices], flux_35ₙ, colormap=colorscheme_flux, colorrange=flux_lims[4]) +flux_45_surface = heatmap!(axflux_45, xC, zC[zF_indices], flux_45ₙ, colormap=colorscheme_flux, colorrange=flux_lims[5]) +flux_55_surface = heatmap!(axflux_55, xC, zC[zF_indices], flux_55ₙ, colormap=colorscheme_flux, colorrange=flux_lims[6]) +flux_65_surface = heatmap!(axflux_65, xC, zC[zF_indices], flux_65ₙ, colormap=colorscheme_flux, colorrange=flux_lims[7]) +flux_75_surface = heatmap!(axflux_75, xC, zC[zF_indices], flux_75ₙ, colormap=colorscheme_flux, colorrange=flux_lims[8]) +flux_85_surface = heatmap!(axflux_85, xC, zC[zF_indices], flux_85ₙ, colormap=colorscheme_flux, colorrange=flux_lims[9]) +flux_95_surface = heatmap!(axflux_95, xC, zC[zF_indices], flux_95ₙ, colormap=colorscheme_flux, colorrange=flux_lims[10]) + +Colorbar(fig[1, 2], field_5_surface, label="Field") +Colorbar(fig[1, 4], flux_5_surface, label="NN Flux") +Colorbar(fig[1, 6], field_15_surface, label="Field") +Colorbar(fig[1, 8], flux_15_surface, label="NN Flux") +Colorbar(fig[2, 2], field_25_surface, label="Field") +Colorbar(fig[2, 4], flux_25_surface, label="NN Flux") +Colorbar(fig[2, 6], field_35_surface, label="Field") +Colorbar(fig[2, 8], flux_35_surface, label="NN Flux") +Colorbar(fig[3, 2], field_45_surface, label="Field") +Colorbar(fig[3, 4], flux_45_surface, label="NN Flux") +Colorbar(fig[3, 6], field_55_surface, label="Field") +Colorbar(fig[3, 8], flux_55_surface, label="NN Flux") +Colorbar(fig[4, 2], field_65_surface, label="Field") +Colorbar(fig[4, 4], flux_65_surface, label="NN Flux") +Colorbar(fig[4, 6], field_75_surface, label="Field") +Colorbar(fig[4, 8], flux_75_surface, label="NN Flux") +Colorbar(fig[5, 2], field_85_surface, label="Field") +Colorbar(fig[5, 4], flux_85_surface, label="NN Flux") +Colorbar(fig[5, 6], field_95_surface, label="Field") +Colorbar(fig[5, 8], flux_95_surface, label="NN Flux") + +for axfield in axfields + xlims!(axfield, minimum(xC), maximum(xC)) + ylims!(axfield, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +end + +for axflux in axfluxes + xlims!(axflux, minimum(xC), maximum(xC)) + ylims!(axflux, minimum(zC[zF_indices]), maximum(zC[zF_indices])) +end + +title_str = @lift "Temperature (°C), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_xzslices_fluxes_T.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end + +#%% +@info "Recording S fields and fluxes in xz" + +fieldname = "S" +fluxname = "wS_NN" + +field_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fieldname, backend=OnDisk()) +field_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fieldname, backend=OnDisk()) +field_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fieldname, backend=OnDisk()) +field_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fieldname, backend=OnDisk()) +field_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fieldname, backend=OnDisk()) +field_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fieldname, backend=OnDisk()) +field_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fieldname, backend=OnDisk()) +field_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fieldname, backend=OnDisk()) +field_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fieldname, backend=OnDisk()) +field_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fluxname, backend=OnDisk()) +flux_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fluxname, backend=OnDisk()) +flux_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fluxname, backend=OnDisk()) +flux_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fluxname, backend=OnDisk()) +flux_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fluxname, backend=OnDisk()) +flux_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fluxname, backend=OnDisk()) +flux_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fluxname, backend=OnDisk()) +flux_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fluxname, backend=OnDisk()) +flux_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fluxname, backend=OnDisk()) +flux_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fluxname, backend=OnDisk()) + +field_datas = [field_NN_data_5, field_NN_data_15, field_NN_data_25, field_NN_data_35, field_NN_data_45, field_NN_data_55, field_NN_data_65, field_NN_data_75, field_NN_data_85, field_NN_data_95] +flux_datas = [flux_NN_data_5, flux_NN_data_15, flux_NN_data_25, flux_NN_data_35, flux_NN_data_45, flux_NN_data_55, flux_NN_data_65, flux_NN_data_75, flux_NN_data_85, flux_NN_data_95] + +xC = field_NN_data_5.grid.xᶜᵃᵃ[1:field_NN_data_5.grid.Nx] +yC = field_NN_data_5.grid.yᵃᶜᵃ[1:field_NN_data_5.grid.Ny] +zC = field_NN_data_5.grid.zᵃᵃᶜ[1:field_NN_data_5.grid.Nz] +zF = field_NN_data_5.grid.zᵃᵃᶠ[1:field_NN_data_5.grid.Nz+1] + +Nt = length(field_NN_data_95) +times = field_NN_data_5.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_5 = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_5.grid.yᵃᶜᵃ[field_NN_data_5.indices[2][1]] / 1000) km") +axfield_15 = CairoMakie.Axis(fig[1, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_15.grid.yᵃᶜᵃ[field_NN_data_15.indices[2][1]] / 1000) km") +axfield_25 = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_25.grid.yᵃᶜᵃ[field_NN_data_25.indices[2][1]] / 1000) km") +axfield_35 = CairoMakie.Axis(fig[2, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_35.grid.yᵃᶜᵃ[field_NN_data_35.indices[2][1]] / 1000) km") +axfield_45 = CairoMakie.Axis(fig[3, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_45.grid.yᵃᶜᵃ[field_NN_data_45.indices[2][1]] / 1000) km") +axfield_55 = CairoMakie.Axis(fig[3, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_55.grid.yᵃᶜᵃ[field_NN_data_55.indices[2][1]] / 1000) km") +axfield_65 = CairoMakie.Axis(fig[4, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_65.grid.yᵃᶜᵃ[field_NN_data_65.indices[2][1]] / 1000) km") +axfield_75 = CairoMakie.Axis(fig[4, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_75.grid.yᵃᶜᵃ[field_NN_data_75.indices[2][1]] / 1000) km") +axfield_85 = CairoMakie.Axis(fig[5, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_85.grid.yᵃᶜᵃ[field_NN_data_85.indices[2][1]] / 1000) km") +axfield_95 = CairoMakie.Axis(fig[5, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_95.grid.yᵃᶜᵃ[field_NN_data_95.indices[2][1]] / 1000) km") + +axflux_5 = CairoMakie.Axis(fig[1, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_5.grid.yᵃᶜᵃ[flux_NN_data_5.indices[2][1]] / 1000) km") +axflux_15 = CairoMakie.Axis(fig[1, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_15.grid.yᵃᶜᵃ[flux_NN_data_15.indices[2][1]] / 1000) km") +axflux_25 = CairoMakie.Axis(fig[2, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_25.grid.yᵃᶜᵃ[flux_NN_data_25.indices[2][1]] / 1000) km") +axflux_35 = CairoMakie.Axis(fig[2, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_35.grid.yᵃᶜᵃ[flux_NN_data_35.indices[2][1]] / 1000) km") +axflux_45 = CairoMakie.Axis(fig[3, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_45.grid.yᵃᶜᵃ[flux_NN_data_45.indices[2][1]] / 1000) km") +axflux_55 = CairoMakie.Axis(fig[3, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_55.grid.yᵃᶜᵃ[flux_NN_data_55.indices[2][1]] / 1000) km") +axflux_65 = CairoMakie.Axis(fig[4, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_65.grid.yᵃᶜᵃ[flux_NN_data_65.indices[2][1]] / 1000) km") +axflux_75 = CairoMakie.Axis(fig[4, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_75.grid.yᵃᶜᵃ[flux_NN_data_75.indices[2][1]] / 1000) km") +axflux_85 = CairoMakie.Axis(fig[5, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_85.grid.yᵃᶜᵃ[flux_NN_data_85.indices[2][1]] / 1000) km") +axflux_95 = CairoMakie.Axis(fig[5, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_95.grid.yᵃᶜᵃ[flux_NN_data_95.indices[2][1]] / 1000) km") + +axfields = [axfield_5, axfield_15, axfield_25, axfield_35, axfield_45, axfield_55, axfield_65, axfield_75, axfield_85, axfield_95] +axfluxes = [axflux_5, axflux_15, axflux_25, axflux_35, axflux_45, axflux_55, axflux_65, axflux_75, axflux_85, axflux_95] + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lims = [(find_min(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices))) for field_data in field_datas] + +flux_lims = [(find_min(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices))) for flux_data in flux_datas] + +flux_lims = [(-maximum(abs, flux_lim), maximum(abs, flux_lim)) for flux_lim in flux_lims] + +NNₙs = [@lift interior(field_data[$n], :, 1, zC_indices) for field_data in field_datas] +fluxₙs = [@lift interior(flux_data[$n], :, 1, zF_indices) for flux_data in flux_datas] + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_5_surface = heatmap!(axfield_5, xC, zC[zC_indices], NN_5ₙ, colormap=colorscheme_field, colorrange=field_lims[1]) +field_15_surface = heatmap!(axfield_15, xC, zC[zC_indices], NN_15ₙ, colormap=colorscheme_field, colorrange=field_lims[2]) +field_25_surface = heatmap!(axfield_25, xC, zC[zC_indices], NN_25ₙ, colormap=colorscheme_field, colorrange=field_lims[3]) +field_35_surface = heatmap!(axfield_35, xC, zC[zC_indices], NN_35ₙ, colormap=colorscheme_field, colorrange=field_lims[4]) +field_45_surface = heatmap!(axfield_45, xC, zC[zC_indices], NN_45ₙ, colormap=colorscheme_field, colorrange=field_lims[5]) +field_55_surface = heatmap!(axfield_55, xC, zC[zC_indices], NN_55ₙ, colormap=colorscheme_field, colorrange=field_lims[6]) +field_65_surface = heatmap!(axfield_65, xC, zC[zC_indices], NN_65ₙ, colormap=colorscheme_field, colorrange=field_lims[7]) +field_75_surface = heatmap!(axfield_75, xC, zC[zC_indices], NN_75ₙ, colormap=colorscheme_field, colorrange=field_lims[8]) +field_85_surface = heatmap!(axfield_85, xC, zC[zC_indices], NN_85ₙ, colormap=colorscheme_field, colorrange=field_lims[9]) +field_95_surface = heatmap!(axfield_95, xC, zC[zC_indices], NN_95ₙ, colormap=colorscheme_field, colorrange=field_lims[10]) + +flux_5_surface = heatmap!(axflux_5, xC, zC[zF_indices], flux_5ₙ, colormap=colorscheme_flux, colorrange=flux_lims[1]) +flux_15_surface = heatmap!(axflux_15, xC, zC[zF_indices], flux_15ₙ, colormap=colorscheme_flux, colorrange=flux_lims[2]) +flux_25_surface = heatmap!(axflux_25, xC, zC[zF_indices], flux_25ₙ, colormap=colorscheme_flux, colorrange=flux_lims[3]) +flux_35_surface = heatmap!(axflux_35, xC, zC[zF_indices], flux_35ₙ, colormap=colorscheme_flux, colorrange=flux_lims[4]) +flux_45_surface = heatmap!(axflux_45, xC, zC[zF_indices], flux_45ₙ, colormap=colorscheme_flux, colorrange=flux_lims[5]) +flux_55_surface = heatmap!(axflux_55, xC, zC[zF_indices], flux_55ₙ, colormap=colorscheme_flux, colorrange=flux_lims[6]) +flux_65_surface = heatmap!(axflux_65, xC, zC[zF_indices], flux_65ₙ, colormap=colorscheme_flux, colorrange=flux_lims[7]) +flux_75_surface = heatmap!(axflux_75, xC, zC[zF_indices], flux_75ₙ, colormap=colorscheme_flux, colorrange=flux_lims[8]) +flux_85_surface = heatmap!(axflux_85, xC, zC[zF_indices], flux_85ₙ, colormap=colorscheme_flux, colorrange=flux_lims[9]) +flux_95_surface = heatmap!(axflux_95, xC, zC[zF_indices], flux_95ₙ, colormap=colorscheme_flux, colorrange=flux_lims[10]) + +Colorbar(fig[1, 2], field_5_surface, label="Field") +Colorbar(fig[1, 4], flux_5_surface, label="NN Flux") +Colorbar(fig[1, 6], field_15_surface, label="Field") +Colorbar(fig[1, 8], flux_15_surface, label="NN Flux") +Colorbar(fig[2, 2], field_25_surface, label="Field") +Colorbar(fig[2, 4], flux_25_surface, label="NN Flux") +Colorbar(fig[2, 6], field_35_surface, label="Field") +Colorbar(fig[2, 8], flux_35_surface, label="NN Flux") +Colorbar(fig[3, 2], field_45_surface, label="Field") +Colorbar(fig[3, 4], flux_45_surface, label="NN Flux") +Colorbar(fig[3, 6], field_55_surface, label="Field") +Colorbar(fig[3, 8], flux_55_surface, label="NN Flux") +Colorbar(fig[4, 2], field_65_surface, label="Field") +Colorbar(fig[4, 4], flux_65_surface, label="NN Flux") +Colorbar(fig[4, 6], field_75_surface, label="Field") +Colorbar(fig[4, 8], flux_75_surface, label="NN Flux") +Colorbar(fig[5, 2], field_85_surface, label="Field") +Colorbar(fig[5, 4], flux_85_surface, label="NN Flux") +Colorbar(fig[5, 6], field_95_surface, label="Field") +Colorbar(fig[5, 8], flux_95_surface, label="NN Flux") + +for axfield in axfields + xlims!(axfield, minimum(xC), maximum(xC)) + ylims!(axfield, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +end + +for axflux in axfluxes + xlims!(axflux, minimum(xC), maximum(xC)) + ylims!(axflux, minimum(zC[zF_indices]), maximum(zC[zF_indices])) +end + +title_str = @lift "Salinity (psu), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_xzslices_fluxes_S.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end +#%% \ No newline at end of file From 33f2ea823a4abbb1584ad0aa779dcad8d36be27e Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 21 Nov 2024 21:41:25 -0500 Subject: [PATCH 099/114] new NN closure witth Rifirstzone --- NN_closure_global_nof_BBLRifirstzone510.jl | 270 +++++++++++++++++++++ 1 file changed, 270 insertions(+) create mode 100644 NN_closure_global_nof_BBLRifirstzone510.jl diff --git a/NN_closure_global_nof_BBLRifirstzone510.jl b/NN_closure_global_nof_BBLRifirstzone510.jl new file mode 100644 index 0000000000..265a4b700e --- /dev/null +++ b/NN_closure_global_nof_BBLRifirstzone510.jl @@ -0,0 +1,270 @@ +import Oceananigans.TurbulenceClosures: + compute_diffusivities!, + DiffusivityFields, + viscosity, + diffusivity, + diffusive_flux_x, + diffusive_flux_y, + diffusive_flux_z, + viscous_flux_ux, + viscous_flux_uy, + viscous_flux_uz, + viscous_flux_vx, + viscous_flux_vy, + viscous_flux_vz, + viscous_flux_wx, + viscous_flux_wy, + viscous_flux_wz + +using Oceananigans.BuoyancyModels: ∂x_b, ∂y_b, ∂z_b +using Oceananigans.Coriolis +using Oceananigans.Grids: φnode +using Oceananigans.Grids: total_size +using Oceananigans.Utils: KernelParameters +using Oceananigans: architecture, on_architecture +using Lux, LuxCUDA +using JLD2 +using ComponentArrays +using OffsetArrays +using SeawaterPolynomials.TEOS10 + +using KernelAbstractions: @index, @kernel, @private + +import Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure, ExplicitTimeDiscretization + +using Adapt + +include("./feature_scaling.jl") + +@inline hack_sind(φ) = sin(φ * π / 180) + +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::HydrostaticSphericalCoriolis) = 2 * coriolis.rotation_rate * hack_sind(φnode(i, j, k, grid, Center(), Center(), Center())) +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::FPlane) = coriolis.f +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::BetaPlane) = coriolis.f₀ + coriolis.β * ynode(i, j, k, grid, Center(), Center(), Center()) + +struct NN{M, P, S} + model :: M + ps :: P + st :: S +end + +@inline (neural_network::NN)(input) = first(neural_network.model(input, neural_network.ps, neural_network.st)) +@inline tosarray(x::AbstractArray) = SArray{Tuple{size(x)...}}(x) + +struct NNFluxClosure{A <: NN, S} <: AbstractTurbulenceClosure{ExplicitTimeDiscretization, 3} + wT :: A + wS :: A + scaling :: S +end + +Adapt.adapt_structure(to, nn :: NNFluxClosure) = + NNFluxClosure(Adapt.adapt(to, nn.wT), + Adapt.adapt(to, nn.wS), + Adapt.adapt(to, nn.scaling)) + +Adapt.adapt_structure(to, nn :: NN) = + NN(Adapt.adapt(to, nn.model), + Adapt.adapt(to, nn.ps), + Adapt.adapt(to, nn.st)) + +function NNFluxClosure(arch) + dev = ifelse(arch == GPU(), gpu_device(), cpu_device()) + nn_path = "./NDE3_FC_Qb_nof_BBLRifirst510_trainFC26new_model.jld2" + + ps, sts, scaling_params, wT_model, wS_model = jldopen(nn_path, "r") do file + ps = file["u_train"] |> dev |> f64 + sts = file["sts"] |> dev |> f64 + scaling_params = file["scaling"] + wT_model = file["model"].wT + wS_model = file["model"].wS + return ps, sts, scaling_params, wT_model, wS_model + end + + scaling = construct_zeromeanunitvariance_scaling(scaling_params) + + wT_NN = NN(wT_model, ps.wT, sts.wT) + wS_NN = NN(wS_model, ps.wS, sts.wS) + + return NNFluxClosure(wT_NN, wS_NN, scaling) +end + +function DiffusivityFields(grid, tracer_names, bcs, ::NNFluxClosure) + arch = architecture(grid) + wT = ZFaceField(grid) + wS = ZFaceField(grid) + first_index = Field((Center, Center, Nothing), grid, Int32) + last_index = Field((Center, Center, Nothing), grid, Int32) + + Nx_in, Ny_in, _ = size(wT) + wrk_in = zeros(10, Nx_in, Ny_in, 15) + wrk_in = on_architecture(arch, wrk_in) + + wrk_wT = zeros(Nx_in, Ny_in, 15) + wrk_wS = zeros(Nx_in, Ny_in, 15) + wrk_wT = on_architecture(arch, wrk_wT) + wrk_wS = on_architecture(arch, wrk_wS) + + return (; wrk_in, wrk_wT, wrk_wS, wT, wS, first_index, last_index) +end + +function compute_diffusivities!(diffusivities, closure::NNFluxClosure, model; parameters = :xyz) + arch = model.architecture + grid = model.grid + velocities = model.velocities + tracers = model.tracers + buoyancy = model.buoyancy + coriolis = model.coriolis + clock = model.clock + top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) + + κᶜ = model.diffusivity_fields[1].κᶜ + Riᶜ = model.closure[1].Riᶜ + Ri = model.diffusivity_fields[1].Ri + + wrk_in = diffusivities.wrk_in + wrk_wT = diffusivities.wrk_wT + wrk_wS = diffusivities.wrk_wS + wT = diffusivities.wT + wS = diffusivities.wS + + first_index = diffusivities.first_index + last_index = diffusivities.last_index + + Nx_in, Ny_in, Nz_in = total_size(wT) + ox_in, oy_in, oz_in = wT.data.offsets + kp = KernelParameters((Nx_in, Ny_in, Nz_in), (ox_in, oy_in, oz_in)) + kp_2D = KernelParameters((Nx_in, Ny_in), (ox_in, oy_in)) + + kp_wrk = KernelParameters((Nx_in, Ny_in, 10), (0, 0, 0)) + + launch!(arch, grid, kp_2D, _find_NN_active_region!, Ri, grid, Riᶜ, first_index, last_index) + + launch!(arch, grid, kp_wrk, + _populate_input!, wrk_in, first_index, last_index, grid, closure, tracers, velocities, buoyancy, top_tracer_bcs, Ri, clock) + + wrk_wT .= dropdims(closure.wT(wrk_in), dims=1) + wrk_wS .= dropdims(closure.wS(wrk_in), dims=1) + + launch!(arch, grid, kp, _fill_adjust_nn_fluxes!, diffusivities, first_index, last_index, wrk_wT, wrk_wS, grid, closure, tracers, velocities, buoyancy, top_tracer_bcs, clock) + return nothing +end + +@kernel function _populate_input!(input, first_index, last_index, grid, closure::NNFluxClosure, tracers, velocities, buoyancy, top_tracer_bcs, Ri, clock) + i, j, k = @index(Global, NTuple) + + scaling = closure.scaling + + ρ₀ = buoyancy.model.equation_of_state.reference_density + g = buoyancy.model.gravitational_acceleration + + @inbounds k_first = first_index[i, j, 1] + @inbounds k_last = last_index[i, j, 1] + + quiescent = quiescent_condition(k_first, k_last) + + @inbounds k_tracer = clamp_k_interior(k_first + k - 1, grid) + + @inbounds k₋ = clamp_k_interior(k_tracer - 1, grid) + @inbounds k₀ = clamp_k_interior(k_tracer, grid) + @inbounds k₊ = clamp_k_interior(k_tracer + 1, grid) + + T, S = tracers.T, tracers.S + + @inbounds input[1, i, j, k] = ∂Tᵢ₋₁ = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k₋, grid, T))) + @inbounds input[2, i, j, k] = ∂Tᵢ = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k₀, grid, T))) + @inbounds input[3, i, j, k] = ∂Tᵢ₊₁ = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k₊, grid, T))) + + @inbounds input[4, i, j, k] = ∂Sᵢ₋₁ = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k₋, grid, S))) + @inbounds input[5, i, j, k] = ∂Sᵢ = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k₀, grid, S))) + @inbounds input[6, i, j, k] = ∂Sᵢ₊₁ = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k₊, grid, S))) + + @inbounds input[7, i, j, k] = ∂σᵢ₋₁ = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k₋, grid, buoyancy, tracers) / g)) + @inbounds input[8, i, j, k] = ∂σᵢ = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k₀, grid, buoyancy, tracers) / g)) + @inbounds input[9, i, j, k] = ∂σᵢ₊₁ = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k₊, grid, buoyancy, tracers) / g)) + + @inbounds input[10, i, j, k] = Jᵇ = scaling.wb(top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers))) + +end + +@kernel function _find_NN_active_region!(Ri, grid, Riᶜ, first_index, last_index) + i, j = @index(Global, NTuple) + top_index = grid.Nz + 1 + grid_point_above_kappa = 10 + grid_point_below_kappa = 5 + + background_κ_index = findfirst(Ri[2:end] .< Riᶜ) + + # Find the first index of the background κᶜ + kloc = grid.Nz+1 + @inbounds for k in grid.Nz:-1:2 + kloc = ifelse(Ri[i, j, k] < Riᶜ, k, kloc) + end + + background_κ_index = kloc - 1 + nonbackground_κ_index = background_κ_index + 1 + + @inbounds last_index[i, j, 1] = ifelse(nonbackground_κ_index == top_index, top_index - 1, clamp_k_interior(background_κ_index + grid_point_above_kappa, grid)) + @inbounds first_index[i, j, 1] = ifelse(nonbackground_κ_index == top_index, top_index, clamp_k_interior(background_κ_index - grid_point_below_kappa + 1, grid)) +end + +@inline function quiescent_condition(lo, hi) + return hi < lo +end + +@inline function within_zone_condition(k, lo, hi) + return (k >= lo) & (k <= hi) +end + +@inline function clamp_k_interior(k, grid) + kmax = grid.Nz + kmin = 2 + + return clamp(k, kmin, kmax) +end + +@kernel function _fill_adjust_nn_fluxes!(diffusivities, first_index, last_index, wrk_wT, wrk_wS, grid, closure::NNFluxClosure, tracers, velocities, buoyancy, top_tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + scaling = closure.scaling + + @inbounds k_first = first_index[i, j, 1] + @inbounds k_last = last_index[i, j, 1] + + convecting = top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers)) > 0 + quiescent = quiescent_condition(k_first, k_last) + within_zone = within_zone_condition(k, k_first, k_last) + + @inbounds k_wrk = clamp(k - k_first + 1, 1, 10) + + NN_active = convecting & !quiescent & within_zone + + @inbounds diffusivities.wT[i, j, k] = ifelse(NN_active, scaling.wT.σ * wrk_wT[i, j, k_wrk], 0) + @inbounds diffusivities.wS[i, j, k] = ifelse(NN_active, scaling.wS.σ * wrk_wS[i, j, k_wrk], 0) +end + +# Write here your constructor +# NNFluxClosure() = ... insert NN here ... (make sure it is on GPU if you need it on GPU!) + +const NNC = NNFluxClosure + +##### +##### Abstract Smagorinsky functionality +##### + +# Horizontal fluxes are zero! +@inline viscous_flux_wz( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_ux( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_uy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline diffusive_flux_x(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) +@inline diffusive_flux_y(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) + +# Viscous fluxes are zero (for now) +@inline viscous_flux_uz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) + +# The only function extended by NNFluxClosure +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{1}, c, clock, fields, buoyancy) = @inbounds K.wT[i, j, k] +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{2}, c, clock, fields, buoyancy) = @inbounds K.wS[i, j, k] \ No newline at end of file From 36cecb656c292c1860cace8c91aa99af20b882b7 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Thu, 21 Nov 2024 22:16:08 -0500 Subject: [PATCH 100/114] fix NN closure --- NN_closure_global_nof_BBLRifirstzone510.jl | 2 -- 1 file changed, 2 deletions(-) diff --git a/NN_closure_global_nof_BBLRifirstzone510.jl b/NN_closure_global_nof_BBLRifirstzone510.jl index 265a4b700e..5251c95873 100644 --- a/NN_closure_global_nof_BBLRifirstzone510.jl +++ b/NN_closure_global_nof_BBLRifirstzone510.jl @@ -192,8 +192,6 @@ end grid_point_above_kappa = 10 grid_point_below_kappa = 5 - background_κ_index = findfirst(Ri[2:end] .< Riᶜ) - # Find the first index of the background κᶜ kloc = grid.Nz+1 @inbounds for k in grid.Nz:-1:2 From 946d7101a4ae02ba3c8d118a062f2f9898d6f157 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Fri, 22 Nov 2024 14:45:17 -0500 Subject: [PATCH 101/114] fix NN_closure --- NN_closure_global_nof_BBLRifirstzone510.jl | 42 ++++++++++++------- validate_NN_nof_BBLRifirstzone510_1D_model.jl | 2 + 2 files changed, 29 insertions(+), 15 deletions(-) diff --git a/NN_closure_global_nof_BBLRifirstzone510.jl b/NN_closure_global_nof_BBLRifirstzone510.jl index 5251c95873..ec31e26dec 100644 --- a/NN_closure_global_nof_BBLRifirstzone510.jl +++ b/NN_closure_global_nof_BBLRifirstzone510.jl @@ -51,16 +51,20 @@ end @inline (neural_network::NN)(input) = first(neural_network.model(input, neural_network.ps, neural_network.st)) @inline tosarray(x::AbstractArray) = SArray{Tuple{size(x)...}}(x) -struct NNFluxClosure{A <: NN, S} <: AbstractTurbulenceClosure{ExplicitTimeDiscretization, 3} - wT :: A - wS :: A - scaling :: S +struct NNFluxClosure{A <: NN, S, G} <: AbstractTurbulenceClosure{ExplicitTimeDiscretization, 3} + wT :: A + wS :: A + scaling :: S + grid_point_above :: G + grid_point_below :: G end Adapt.adapt_structure(to, nn :: NNFluxClosure) = NNFluxClosure(Adapt.adapt(to, nn.wT), Adapt.adapt(to, nn.wS), - Adapt.adapt(to, nn.scaling)) + Adapt.adapt(to, nn.scaling), + Adapt.adapt(to, nn.grid_point_above), + Adapt.adapt(to, nn.grid_point_below)) Adapt.adapt_structure(to, nn :: NN) = NN(Adapt.adapt(to, nn.model), @@ -85,18 +89,24 @@ function NNFluxClosure(arch) wT_NN = NN(wT_model, ps.wT, sts.wT) wS_NN = NN(wS_model, ps.wS, sts.wS) - return NNFluxClosure(wT_NN, wS_NN, scaling) + grid_point_above = 10 + grid_point_below = 5 + + return NNFluxClosure(wT_NN, wS_NN, scaling, grid_point_above, grid_point_below) end -function DiffusivityFields(grid, tracer_names, bcs, ::NNFluxClosure) +function DiffusivityFields(grid, tracer_names, bcs, closure::NNFluxClosure) arch = architecture(grid) wT = ZFaceField(grid) wS = ZFaceField(grid) first_index = Field((Center, Center, Nothing), grid, Int32) last_index = Field((Center, Center, Nothing), grid, Int32) + N_input = closure.wT.model.layers.layer_1.in_dims + N_levels = closure.grid_point_above + closure.grid_point_below + Nx_in, Ny_in, _ = size(wT) - wrk_in = zeros(10, Nx_in, Ny_in, 15) + wrk_in = zeros(N_input, Nx_in, Ny_in, N_levels) wrk_in = on_architecture(arch, wrk_in) wrk_wT = zeros(Nx_in, Ny_in, 15) @@ -117,7 +127,6 @@ function compute_diffusivities!(diffusivities, closure::NNFluxClosure, model; pa clock = model.clock top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) - κᶜ = model.diffusivity_fields[1].κᶜ Riᶜ = model.closure[1].Riᶜ Ri = model.diffusivity_fields[1].Ri @@ -135,9 +144,11 @@ function compute_diffusivities!(diffusivities, closure::NNFluxClosure, model; pa kp = KernelParameters((Nx_in, Ny_in, Nz_in), (ox_in, oy_in, oz_in)) kp_2D = KernelParameters((Nx_in, Ny_in), (ox_in, oy_in)) - kp_wrk = KernelParameters((Nx_in, Ny_in, 10), (0, 0, 0)) + N_levels = closure.grid_point_above + closure.grid_point_below + + kp_wrk = KernelParameters((Nx_in, Ny_in, N_levels), (0, 0, 0)) - launch!(arch, grid, kp_2D, _find_NN_active_region!, Ri, grid, Riᶜ, first_index, last_index) + launch!(arch, grid, kp_2D, _find_NN_active_region!, Ri, grid, Riᶜ, first_index, last_index, closure) launch!(arch, grid, kp_wrk, _populate_input!, wrk_in, first_index, last_index, grid, closure, tracers, velocities, buoyancy, top_tracer_bcs, Ri, clock) @@ -186,11 +197,11 @@ end end -@kernel function _find_NN_active_region!(Ri, grid, Riᶜ, first_index, last_index) +@kernel function _find_NN_active_region!(Ri, grid, Riᶜ, first_index, last_index, closure::NNFluxClosure) i, j = @index(Global, NTuple) top_index = grid.Nz + 1 - grid_point_above_kappa = 10 - grid_point_below_kappa = 5 + grid_point_above_kappa = closure.grid_point_above + grid_point_below_kappa = closure.grid_point_below # Find the first index of the background κᶜ kloc = grid.Nz+1 @@ -231,7 +242,8 @@ end quiescent = quiescent_condition(k_first, k_last) within_zone = within_zone_condition(k, k_first, k_last) - @inbounds k_wrk = clamp(k - k_first + 1, 1, 10) + N_levels = closure.grid_point_above + closure.grid_point_below + @inbounds k_wrk = clamp(k - k_first + 1, 1, N_levels) NN_active = convecting & !quiescent & within_zone diff --git a/validate_NN_nof_BBLRifirstzone510_1D_model.jl b/validate_NN_nof_BBLRifirstzone510_1D_model.jl index 57c94e3937..54f4fdd9df 100644 --- a/validate_NN_nof_BBLRifirstzone510_1D_model.jl +++ b/validate_NN_nof_BBLRifirstzone510_1D_model.jl @@ -103,6 +103,7 @@ update_state!(model) ##### Δt₀ = file["Δt"] stop_time = file["τ"] +# stop_time = 100minutes simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) @@ -125,6 +126,7 @@ function print_progress(sim) end simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(20)) +# simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(1)) ##### ##### Diagnostics From 0c55c448cd112b3f0a74e7a0dbc02bbd26827f74 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 26 Nov 2024 11:44:53 -0500 Subject: [PATCH 102/114] updated base closure with new calibration --- ...ai_vertical_diffusivity_local_2step_new.jl | 242 ++++++++++++++++++ 1 file changed, 242 insertions(+) create mode 100644 xin_kai_vertical_diffusivity_local_2step_new.jl diff --git a/xin_kai_vertical_diffusivity_local_2step_new.jl b/xin_kai_vertical_diffusivity_local_2step_new.jl new file mode 100644 index 0000000000..9dda064e79 --- /dev/null +++ b/xin_kai_vertical_diffusivity_local_2step_new.jl @@ -0,0 +1,242 @@ +using Oceananigans +using Oceananigans.Architectures: architecture +using Oceananigans.BuoyancyModels: ∂z_b +using Oceananigans.Operators +using Oceananigans.Grids: inactive_node, total_size +using Oceananigans.Operators: ℑzᵃᵃᶜ, ℑxyᶠᶠᵃ, ℑxyᶜᶜᵃ +using Oceananigans.Utils: KernelParameters + +using Adapt + +using KernelAbstractions: @index, @kernel +using KernelAbstractions.Extras.LoopInfo: @unroll + +using Oceananigans.TurbulenceClosures: + tapering_factorᶠᶜᶜ, + tapering_factorᶜᶠᶜ, + tapering_factorᶜᶜᶠ, + tapering_factor, + SmallSlopeIsopycnalTensor, + AbstractScalarDiffusivity, + ExplicitTimeDiscretization, + FluxTapering, + isopycnal_rotation_tensor_xz_ccf, + isopycnal_rotation_tensor_yz_ccf, + isopycnal_rotation_tensor_zz_ccf + +import Oceananigans.TurbulenceClosures: + compute_diffusivities!, + DiffusivityFields, + viscosity, + diffusivity, + getclosure, + top_buoyancy_flux, + diffusive_flux_x, + diffusive_flux_y, + diffusive_flux_z, + viscous_flux_ux, + viscous_flux_vx, + viscous_flux_uy, + viscous_flux_vy + +using Oceananigans.Utils: launch! +using Oceananigans.Coriolis: fᶠᶠᵃ +using Oceananigans.Operators +using Oceananigans.BuoyancyModels: ∂x_b, ∂y_b, ∂z_b + +using Oceananigans.TurbulenceClosures +using Oceananigans.TurbulenceClosures: HorizontalFormulation, VerticalFormulation, AbstractScalarDiffusivity +using Oceananigans.TurbulenceClosures: AbstractScalarBiharmonicDiffusivity +using Oceananigans.Operators +using Oceananigans.Operators: Δxᶜᶜᶜ, Δyᶜᶜᶜ, ℑxyᶜᶜᵃ, ζ₃ᶠᶠᶜ, div_xyᶜᶜᶜ +using Oceananigans.Operators: Δx, Δy +using Oceananigans.Operators: ℑxyz + +using Oceananigans.Operators: ℑxyzᶜᶜᶠ, ℑyzᵃᶜᶠ, ℑxzᶜᵃᶠ, Δxᶜᶜᶜ, Δyᶜᶜᶜ + +using Oceananigans.BoundaryConditions + +struct XinKaiLocalVerticalDiffusivity{TD, FT} <: AbstractScalarDiffusivity{TD, VerticalFormulation, 2} + ν₀ :: FT + νˢʰ :: FT + νᶜⁿ :: FT + Pr_convₜ :: FT + Pr_shearₜ :: FT + Riᶜ :: FT + δRi :: FT +end + +function XinKaiLocalVerticalDiffusivity{TD}(ν₀ :: FT, + νˢʰ :: FT, + νᶜⁿ :: FT, + Pr_convₜ :: FT, + Pr_shearₜ :: FT, + Riᶜ :: FT, + δRi :: FT) where {TD, FT} + + return XinKaiLocalVerticalDiffusivity{TD, FT}(ν₀, νˢʰ, νᶜⁿ, Pr_convₜ, Pr_shearₜ, Riᶜ, δRi) +end + +function XinKaiLocalVerticalDiffusivity(time_discretization = VerticallyImplicitTimeDiscretization(), + FT = Float64; + ν₀ = 1e-5, + νˢʰ = 0.0615914063656973, + νᶜⁿ = 1.0514706176740092, + Pr_convₜ = 0.2684497234339729, + Pr_shearₜ = 1.0842017486284887, + Riᶜ = 0.4366901962987793, + δRi = 0.001484740489701266) + + TD = typeof(time_discretization) + + return XinKaiLocalVerticalDiffusivity{TD}(convert(FT, ν₀), + convert(FT, νˢʰ), + convert(FT, νᶜⁿ), + convert(FT, Pr_convₜ), + convert(FT, Pr_shearₜ), + convert(FT, Riᶜ), + convert(FT, δRi)) +end + +XinKaiLocalVerticalDiffusivity(FT::DataType; kw...) = + XinKaiLocalVerticalDiffusivity(VerticallyImplicitTimeDiscretization(), FT; kw...) + +Adapt.adapt_structure(to, clo::XinKaiLocalVerticalDiffusivity{TD, FT}) where {TD, FT} = + XinKaiLocalVerticalDiffusivity{TD, FT}(clo.ν₀, clo.νˢʰ, clo.νᶜⁿ, clo.Pr_convₜ, clo.Pr_shearₜ, clo.Riᶜ, clo.δRi) + +##### +##### Diffusivity field utilities +##### + +const RBVD = XinKaiLocalVerticalDiffusivity +const RBVDArray = AbstractArray{<:RBVD} +const FlavorOfXKVD = Union{RBVD, RBVDArray} +const c = Center() +const f = Face() + +@inline viscosity_location(::FlavorOfXKVD) = (c, c, f) +@inline diffusivity_location(::FlavorOfXKVD) = (c, c, f) + +@inline viscosity(::FlavorOfXKVD, diffusivities) = diffusivities.κᵘ +@inline diffusivity(::FlavorOfXKVD, diffusivities, id) = diffusivities.κᶜ + +with_tracers(tracers, closure::FlavorOfXKVD) = closure + +# Note: computing diffusivities at cell centers for now. +function DiffusivityFields(grid, tracer_names, bcs, closure::FlavorOfXKVD) + κᶜ = Field((Center, Center, Face), grid) + κᵘ = Field((Center, Center, Face), grid) + Ri = Field((Center, Center, Face), grid) + return (; κᶜ, κᵘ, Ri) +end + +function compute_diffusivities!(diffusivities, closure::FlavorOfXKVD, model; parameters = :xyz) + arch = model.architecture + grid = model.grid + clock = model.clock + tracers = model.tracers + buoyancy = model.buoyancy + velocities = model.velocities + top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) + + Nx_in, Ny_in, Nz_in = total_size(diffusivities.κᶜ) + ox_in, oy_in, oz_in = diffusivities.κᶜ.data.offsets + + kp = KernelParameters((Nx_in, Ny_in, Nz_in), (ox_in, oy_in, oz_in)) + + launch!(arch, grid, kp, + compute_ri_number!, + diffusivities, + grid, + closure, + velocities, + tracers, + buoyancy, + top_tracer_bcs, + clock) + + # Use `only_local_halos` to ensure that no communication occurs during + # this call to fill_halo_regions! + fill_halo_regions!(diffusivities.Ri; only_local_halos=true) + + launch!(arch, grid, kp, + compute_xinkai_diffusivities!, + diffusivities, + grid, + closure, + velocities, + tracers, + buoyancy, + top_tracer_bcs, + clock) + + return nothing +end + +@inline ϕ²(i, j, k, grid, ϕ, args...) = ϕ(i, j, k, grid, args...)^2 + +@inline function shear_squaredᶜᶜᶠ(i, j, k, grid, velocities) + ∂z_u² = ℑxᶜᵃᵃ(i, j, k, grid, ϕ², ∂zᶠᶜᶠ, velocities.u) + ∂z_v² = ℑyᵃᶜᵃ(i, j, k, grid, ϕ², ∂zᶜᶠᶠ, velocities.v) + return ∂z_u² + ∂z_v² +end + +@inline function Riᶜᶜᶠ(i, j, k, grid, velocities, buoyancy, tracers) + S² = shear_squaredᶜᶜᶠ(i, j, k, grid, velocities) + N² = ∂z_b(i, j, k, grid, buoyancy, tracers) + Ri = N² / S² + + # Clip N² and avoid NaN + return ifelse(N² == 0, zero(grid), Ri) +end + +const c = Center() +const f = Face() + +@kernel function compute_ri_number!(diffusivities, grid, closure::FlavorOfXKVD, + velocities, tracers, buoyancy, tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + @inbounds diffusivities.Ri[i, j, k] = Riᶜᶜᶠ(i, j, k, grid, velocities, buoyancy, tracers) +end + +@kernel function compute_xinkai_diffusivities!(diffusivities, grid, closure::FlavorOfXKVD, + velocities, tracers, buoyancy, tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + _compute_xinkai_diffusivities!(i, j, k, diffusivities, grid, closure, + velocities, tracers, buoyancy, tracer_bcs, clock) +end + +@inline function _compute_xinkai_diffusivities!(i, j, k, diffusivities, grid, closure, + velocities, tracers, buoyancy, tracer_bcs, clock) + + # Ensure this works with "ensembles" of closures, in addition to ordinary single closures + closure_ij = getclosure(i, j, closure) + + ν₀ = closure_ij.ν₀ + νˢʰ = closure_ij.νˢʰ + νᶜⁿ = closure_ij.νᶜⁿ + Pr_convₜ = closure_ij.Pr_convₜ + Pr_shearₜ = closure_ij.Pr_shearₜ + Riᶜ = closure_ij.Riᶜ + δRi = closure_ij.δRi + + κ₀ = ν₀ / Pr_shearₜ + κˢʰ = νˢʰ / Pr_shearₜ + κᶜⁿ = νᶜⁿ / Pr_convₜ + + # (Potentially) apply a horizontal filter to the Richardson number + Ri = ℑxyᶜᶜᵃ(i, j, k, grid, ℑxyᶠᶠᵃ, diffusivities.Ri) + + # Conditions + convecting = Ri < 0 # applies regardless of Qᵇ + + # Convective adjustment diffusivity + ν_local = ifelse(convecting, (νˢʰ - νᶜⁿ) * tanh(Ri / δRi) + νˢʰ, clamp((ν₀ - νˢʰ) * Ri / Riᶜ + νˢʰ, ν₀, νˢʰ)) + κ_local = ifelse(convecting, (κˢʰ - κᶜⁿ) * tanh(Ri / δRi) + κˢʰ, clamp((κ₀ - κˢʰ) * Ri / Riᶜ + κˢʰ, κ₀, κˢʰ)) + + # Update by averaging in time + @inbounds diffusivities.κᵘ[i, j, k] = ifelse(k <= 1 || k >= grid.Nz+1, 0, ν_local) + @inbounds diffusivities.κᶜ[i, j, k] = ifelse(k <= 1 || k >= grid.Nz+1, 0, κ_local) + + return nothing +end From dada2c5282749217083036378c886642d8235722 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 26 Nov 2024 11:46:56 -0500 Subject: [PATCH 103/114] run base closure and CATKE with mode waters --- baseclosure_doublegyre_model_modewater.jl | 589 ++++++++++++++++++ physicalclosure_doublegyre_model_modewater.jl | 551 ++++++++++++++++ 2 files changed, 1140 insertions(+) create mode 100644 baseclosure_doublegyre_model_modewater.jl create mode 100644 physicalclosure_doublegyre_model_modewater.jl diff --git a/baseclosure_doublegyre_model_modewater.jl b/baseclosure_doublegyre_model_modewater.jl new file mode 100644 index 0000000000..d1a18db90a --- /dev/null +++ b/baseclosure_doublegyre_model_modewater.jl @@ -0,0 +1,589 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("xin_kai_vertical_diffusivity_local_2step_new.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + +#%% +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zC2O_newbaseclosure_trainFC24new_scalingtrain54new_2Pr_2step" +# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +base_closure = XinKaiLocalVerticalDiffusivity() +closure = base_closure + +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) +# advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/30days + +const X₀ = -Lx/2 + 800kilometers +const Y₀ = -Ly/2 + 1500kilometers +const R₀ = 700kilometers +const Qᵀ_mode = 4.5e-4 +const σ_mode = 20kilometers + +##### +##### Forcing and initial condition +##### +# @inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) +# @inline T_initial(x, y, z) = (T_north + T_south / 2) + 5 * (1 + z / Lz) +@inline T_initial(x, y, z) = 10 + 20 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y + +@inline Qᵀ_winter(t) = max(0, -Qᵀ_mode * sin(2π * t / 360days)) +@inline Qᵀ_subpolar(x, y, t) = ifelse((x - X₀)^2 + (y - Y₀)^2 <= R₀^2, Qᵀ_winter(t), + exp(-(sqrt((x - X₀)^2 + (y - Y₀)^2) - R₀)^2 / (2 * σ_mode^2)) * Qᵀ_winter(t)) + +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) + Qᵀ_subpolar(x, y, t) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10800days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); + +κ = model.diffusivity_fields.κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) +wTbar_zonal = Average(wT_base, dims=1) +wSbar_zonal = Average(wS_base, dims=1) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + +outputs = (; u, v, w, T, S, ρ, N², wT=wT_base, wS=wS_base) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal, wTbar_zonal, wSbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:complete_fields] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields", + schedule = TimeInterval(1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 365 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% +# function find_min(a...) +# return minimum(minimum.([a...])) +# end + +# function find_max(a...) +# return maximum(maximum.([a...])) +# end + +# N²_xz_north_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_north.jld2", "N²") +# N²_xz_south_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_south.jld2", "N²") + +# xC = N²_xz_north_data.grid.xᶜᵃᵃ[1:N²_xz_north_data.grid.Nx] +# zf = N²_xz_north_data.grid.zᵃᵃᶠ[1:N²_xz_north_data.grid.Nz+1] + +# yloc_north = N²_xz_north_data.grid.yᵃᶠᵃ[N²_xz_north_data.indices[2][1]] +# yloc_south = N²_xz_south_data.grid.yᵃᶠᵃ[N²_xz_south_data.indices[2][1]] + +# Nt = length(N²_xz_north_data) +# times = N²_xz_north_data.times / 24 / 60^2 / 365 +# timeframes = 1:Nt + +# N²_lim = (find_min(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) - 1e-13, +# find_max(interior(N²_xz_north_data, :, 1, :, timeframes), interior(N²_xz_south_data, :, 1, :, timeframes)) + 1e-13) +# #%% +# fig = Figure(size=(800, 800)) +# ax_north = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_south / 1e3)) km") +# ax_south = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="Buoyancy Frequency N² at y = $(round(yloc_north / 1e3)) km") + +# n = Observable(2) + +# N²_north = @lift interior(N²_xz_north_data[$n], :, 1, :) +# N²_south = @lift interior(N²_xz_south_data[$n], :, 1, :) + +# colorscheme = colorschemes[:jet] + +# N²_north_surface = heatmap!(ax_north, xC, zf, N²_north, colormap=colorscheme, colorrange=N²_lim) +# N²_south_surface = heatmap!(ax_south, xC, zf, N²_south, colormap=colorscheme, colorrange=N²_lim) + +# Colorbar(fig[1:2, 2], N²_north_surface, label="N² (s⁻²)") + +# title_str = @lift "Time = $(round(times[$n], digits=2)) years" +# Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +# trim!(fig.layout) + +# @info "Recording buoyancy frequency xz slice" +# CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_buoyancy_frequency_xz_slice.mp4", 1:Nt, framerate=3, px_per_unit=2) do nn +# n[] = nn +# end + +# @info "Done!" +# #%% \ No newline at end of file diff --git a/physicalclosure_doublegyre_model_modewater.jl b/physicalclosure_doublegyre_model_modewater.jl new file mode 100644 index 0000000000..571a1d21e0 --- /dev/null +++ b/physicalclosure_doublegyre_model_modewater.jl @@ -0,0 +1,551 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity +using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: CATKEVerticalDiffusivity, CATKEMixingLength, CATKEEquation + +using Oceananigans.BuoyancyModels: ∂z_b +# include("NN_closure_global.jl") +# include("xin_kai_vertical_diffusivity_local.jl") +# include("xin_kai_vertical_diffusivity_2Pr.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes +using Glob + +#%% +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zC2O_CATKEVerticalDiffusivity" +# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_CATKEVerticalDiffusivity" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +# vertical_base_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) +# convection_closure = XinKaiVerticalDiffusivity() +function CATKE_ocean_closure() + mixing_length = CATKEMixingLength(Cᵇ=0.01) + turbulent_kinetic_energy_equation = CATKEEquation(Cᵂϵ=1.0) + return CATKEVerticalDiffusivity(; mixing_length, turbulent_kinetic_energy_equation) +end +convection_closure = CATKE_ocean_closure() +closure = convection_closure +# closure = vertical_base_closure + +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) +# advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/30days + +const X₀ = -Lx/2 + 800kilometers +const Y₀ = -Ly/2 + 1500kilometers +const R₀ = 700kilometers +const Qᵀ_mode = 4.5e-4 +const σ_mode = 20kilometers + +##### +##### Forcing and initial condition +##### +# @inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) +@inline T_initial(x, y, z) = 10 + 20 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y + +@inline Qᵀ_winter(t) = max(0, -Qᵀ_mode * sin(2π * t / 360days)) +@inline Qᵀ_subpolar(x, y, t) = ifelse((x - X₀)^2 + (y - Y₀)^2 <= R₀^2, Qᵀ_winter(t), + exp(-(sqrt((x - X₀)^2 + (y - Y₀)^2) - R₀)^2 / (2 * σ_mode^2)) * Qᵀ_winter(t)) + +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) + Qᵀ_subpolar(x, y, t) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S, :e), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S, :e), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10950days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)) +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)) + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) +compute!(ρ) + +ρᶠ = @at (Center, Center, Face) ρ +∂ρ∂z = ∂z(ρ) +∂²ρ∂z² = ∂z(∂ρ∂z) + +κc = model.diffusivity_fields.κc +wT = κc * ∂z(T) +wS = κc * ∂z(S) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) +wTbar_zonal = Average(wT, dims=1) +wSbar_zonal = Average(wS, dims=1) + +outputs = (; u, v, w, T, S, ρ, ρᶠ, ∂ρ∂z, ∂²ρ∂z², N², wT, wS) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal, wTbar_zonal, wSbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=20, px_per_unit=2) do nn + n[] = nn +end + +@info "Done!" +#%% +# Ψ_data = FieldTimeSeries("$(FILE_DIR)/averaged_fields_streamfunction.jld2", "Ψ") + +# xF = Ψ_data.grid.xᶠᵃᵃ[1:Ψ_data.grid.Nx+1] +# yC = Ψ_data.grid.yᵃᶜᵃ[1:Ψ_data.grid.Ny] + +# Nt = length(Ψ_data) +# times = Ψ_data.times / 24 / 60^2 / 365 +# #%% +# timeframe = Nt +# Ψ_frame = interior(Ψ_data[timeframe], :, :, 1) ./ 1e6 +# clim = maximum(abs, Ψ_frame) + 1e-13 +# N_levels = 16 +# levels = range(-clim, stop=clim, length=N_levels) +# fig = Figure(size=(800, 800)) +# ax = Axis(fig[1, 1], xlabel="x (m)", ylabel="y (m)", title="CATKE Vertical Diffusivity, Yearly-Averaged Barotropic streamfunction Ψ, Year $(times[timeframe])") +# cf = contourf!(ax, xF, yC, Ψ_frame, levels=levels, colormap=Reverse(:RdBu_11)) +# Colorbar(fig[1, 2], cf, label="Ψ (Sv)") +# tightlimits!(ax) +# save("$(FILE_DIR)/barotropic_streamfunction_$(timeframe).png", fig, px_per_unit=4) +# display(fig) +#%% \ No newline at end of file From 1d0ef1f76cd356ab2ae8e2d9d93137ed3bf51baa Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 26 Nov 2024 12:59:49 -0500 Subject: [PATCH 104/114] fix advection scheme and file name --- baseclosure_doublegyre_model_modewater.jl | 6 ++---- physicalclosure_doublegyre_model_modewater.jl | 9 ++------- 2 files changed, 4 insertions(+), 11 deletions(-) diff --git a/baseclosure_doublegyre_model_modewater.jl b/baseclosure_doublegyre_model_modewater.jl index d1a18db90a..11328e9ede 100644 --- a/baseclosure_doublegyre_model_modewater.jl +++ b/baseclosure_doublegyre_model_modewater.jl @@ -20,8 +20,7 @@ using SeawaterPolynomials:TEOS10 using ColorSchemes #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zC2O_newbaseclosure_trainFC24new_scalingtrain54new_2Pr_2step" -# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_baseclosure_trainFC24new_scalingtrain54new_2Pr_2step" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_newbaseclosure" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -32,8 +31,7 @@ model_architecture = GPU() base_closure = XinKaiLocalVerticalDiffusivity() closure = base_closure -advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) -# advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) # number of grid points const Nx = 100 diff --git a/physicalclosure_doublegyre_model_modewater.jl b/physicalclosure_doublegyre_model_modewater.jl index 571a1d21e0..ed5a62c0d9 100644 --- a/physicalclosure_doublegyre_model_modewater.jl +++ b/physicalclosure_doublegyre_model_modewater.jl @@ -27,8 +27,7 @@ using ColorSchemes using Glob #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zC2O_CATKEVerticalDiffusivity" -# filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_CATKEVerticalDiffusivity" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_CATKEVerticalDiffusivity" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -36,8 +35,6 @@ mkpath(FILE_DIR) # Architecture model_architecture = GPU() -# vertical_base_closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5) -# convection_closure = XinKaiVerticalDiffusivity() function CATKE_ocean_closure() mixing_length = CATKEMixingLength(Cᵇ=0.01) turbulent_kinetic_energy_equation = CATKEEquation(Cᵂϵ=1.0) @@ -45,10 +42,8 @@ function CATKE_ocean_closure() end convection_closure = CATKE_ocean_closure() closure = convection_closure -# closure = vertical_base_closure -advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), CenteredSecondOrder()) -# advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) # number of grid points const Nx = 100 From 4898be4c8b0854a9a5285994b8c6afdf1e827b41 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 26 Nov 2024 13:02:02 -0500 Subject: [PATCH 105/114] new NN closure including Ri --- NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model.jl b/NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model.jl index f380820afe..6e6e425e94 100644 --- a/NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model.jl +++ b/NNclosure_Ri_nof_BBLkappazonelast55_doublegyre_model.jl @@ -21,7 +21,7 @@ using SeawaterPolynomials:TEOS10 using ColorSchemes #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_8days_NN_closure_NDE_Ri_BBLkappazonelast55_temp" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_NN_closure_NDE_Ri_BBLkappazonelast55_temp" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" mkpath(FILE_DIR) @@ -69,7 +69,7 @@ const S_mid = (S_north + S_south) / 2 const τ₀ = 1e-4 const μ_drag = 1/30days -const μ_T = 1/8days +const μ_T = 1/30days ##### ##### Forcing and initial condition From 3e84aaff49a9a8f04caaba3f674fe1609958144e Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 26 Nov 2024 13:03:35 -0500 Subject: [PATCH 106/114] actual new closure including Ri --- NN_closure_global_Ri_nof_BBLRifirstzone510.jl | 294 ++++++++++++++++++ 1 file changed, 294 insertions(+) create mode 100644 NN_closure_global_Ri_nof_BBLRifirstzone510.jl diff --git a/NN_closure_global_Ri_nof_BBLRifirstzone510.jl b/NN_closure_global_Ri_nof_BBLRifirstzone510.jl new file mode 100644 index 0000000000..bde30dfe58 --- /dev/null +++ b/NN_closure_global_Ri_nof_BBLRifirstzone510.jl @@ -0,0 +1,294 @@ +import Oceananigans.TurbulenceClosures: + compute_diffusivities!, + DiffusivityFields, + viscosity, + diffusivity, + diffusive_flux_x, + diffusive_flux_y, + diffusive_flux_z, + viscous_flux_ux, + viscous_flux_uy, + viscous_flux_uz, + viscous_flux_vx, + viscous_flux_vy, + viscous_flux_vz, + viscous_flux_wx, + viscous_flux_wy, + viscous_flux_wz + +using Oceananigans.BuoyancyModels: ∂x_b, ∂y_b, ∂z_b +using Oceananigans.Coriolis +using Oceananigans.Grids: φnode +using Oceananigans.Grids: total_size +using Oceananigans.Utils: KernelParameters +using Oceananigans: architecture, on_architecture +using Lux, LuxCUDA +using JLD2 +using ComponentArrays +using OffsetArrays +using SeawaterPolynomials.TEOS10 + +using KernelAbstractions: @index, @kernel, @private + +import Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure, ExplicitTimeDiscretization + +using Adapt + +include("./feature_scaling.jl") + +@inline hack_sind(φ) = sin(φ * π / 180) + +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::HydrostaticSphericalCoriolis) = 2 * coriolis.rotation_rate * hack_sind(φnode(i, j, k, grid, Center(), Center(), Center())) +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::FPlane) = coriolis.f +@inline fᶜᶜᵃ(i, j, k, grid, coriolis::BetaPlane) = coriolis.f₀ + coriolis.β * ynode(i, j, k, grid, Center(), Center(), Center()) + +struct NN{M, P, S} + model :: M + ps :: P + st :: S +end + +@inline (neural_network::NN)(input) = first(neural_network.model(input, neural_network.ps, neural_network.st)) +@inline tosarray(x::AbstractArray) = SArray{Tuple{size(x)...}}(x) + +struct NNFluxClosure{A <: NN, S, G} <: AbstractTurbulenceClosure{ExplicitTimeDiscretization, 3} + wT :: A + wS :: A + scaling :: S + grid_point_above :: G + grid_point_below :: G +end + +Adapt.adapt_structure(to, nn :: NNFluxClosure) = + NNFluxClosure(Adapt.adapt(to, nn.wT), + Adapt.adapt(to, nn.wS), + Adapt.adapt(to, nn.scaling), + Adapt.adapt(to, nn.grid_point_above), + Adapt.adapt(to, nn.grid_point_below)) + +Adapt.adapt_structure(to, nn :: NN) = + NN(Adapt.adapt(to, nn.model), + Adapt.adapt(to, nn.ps), + Adapt.adapt(to, nn.st)) + +function NNFluxClosure(arch) + dev = ifelse(arch == GPU(), gpu_device(), cpu_device()) + nn_path = "./NDE5_FC_Qb_Ri_nof_BBLRifirst510_train62newnohighrotation_scalingtrain62newnohighrotation_validate30new_3layer_128_relu_30seed_2Pr_ls5_model_temp.jld2" + + ps, sts, scaling_params, wT_model, wS_model = jldopen(nn_path, "r") do file + ps = file["u_validation"] |> dev |> f64 + sts = file["sts"] |> dev |> f64 + scaling_params = file["scaling"] + wT_model = file["model"].wT + wS_model = file["model"].wS + return ps, sts, scaling_params, wT_model, wS_model + end + + scaling = construct_zeromeanunitvariance_scaling(scaling_params) + + wT_NN = NN(wT_model, ps.wT, sts.wT) + wS_NN = NN(wS_model, ps.wS, sts.wS) + + grid_point_above = 10 + grid_point_below = 5 + + return NNFluxClosure(wT_NN, wS_NN, scaling, grid_point_above, grid_point_below) +end + +function DiffusivityFields(grid, tracer_names, bcs, closure::NNFluxClosure) + arch = architecture(grid) + wT = ZFaceField(grid) + wS = ZFaceField(grid) + first_index = Field((Center, Center, Nothing), grid, Int32) + last_index = Field((Center, Center, Nothing), grid, Int32) + + N_input = closure.wT.model.layers.layer_1.in_dims + N_levels = closure.grid_point_above + closure.grid_point_below + + Nx_in, Ny_in, _ = size(wT) + wrk_in = zeros(N_input, Nx_in, Ny_in, N_levels) + wrk_in = on_architecture(arch, wrk_in) + + wrk_wT = zeros(Nx_in, Ny_in, 15) + wrk_wS = zeros(Nx_in, Ny_in, 15) + wrk_wT = on_architecture(arch, wrk_wT) + wrk_wS = on_architecture(arch, wrk_wS) + + return (; wrk_in, wrk_wT, wrk_wS, wT, wS, first_index, last_index) +end + +function compute_diffusivities!(diffusivities, closure::NNFluxClosure, model; parameters = :xyz) + arch = model.architecture + grid = model.grid + velocities = model.velocities + tracers = model.tracers + buoyancy = model.buoyancy + coriolis = model.coriolis + clock = model.clock + top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers)) + + Riᶜ = model.closure[1].Riᶜ + Ri = model.diffusivity_fields[1].Ri + + wrk_in = diffusivities.wrk_in + wrk_wT = diffusivities.wrk_wT + wrk_wS = diffusivities.wrk_wS + wT = diffusivities.wT + wS = diffusivities.wS + + first_index = diffusivities.first_index + last_index = diffusivities.last_index + + Nx_in, Ny_in, Nz_in = total_size(wT) + ox_in, oy_in, oz_in = wT.data.offsets + kp = KernelParameters((Nx_in, Ny_in, Nz_in), (ox_in, oy_in, oz_in)) + kp_2D = KernelParameters((Nx_in, Ny_in), (ox_in, oy_in)) + + N_levels = closure.grid_point_above + closure.grid_point_below + + kp_wrk = KernelParameters((Nx_in, Ny_in, N_levels), (0, 0, 0)) + + launch!(arch, grid, kp_2D, _find_NN_active_region!, Ri, grid, Riᶜ, first_index, last_index, closure) + + launch!(arch, grid, kp_wrk, + _populate_input!, wrk_in, first_index, last_index, grid, closure, tracers, velocities, buoyancy, top_tracer_bcs, Ri, clock) + + wrk_wT .= dropdims(closure.wT(wrk_in), dims=1) + wrk_wS .= dropdims(closure.wS(wrk_in), dims=1) + + launch!(arch, grid, kp, _fill_adjust_nn_fluxes!, diffusivities, first_index, last_index, wrk_wT, wrk_wS, grid, closure, tracers, velocities, buoyancy, top_tracer_bcs, clock) + return nothing +end + +@kernel function _populate_input!(input, first_index, last_index, grid, closure::NNFluxClosure, tracers, velocities, buoyancy, top_tracer_bcs, Ri, clock) + i, j, k = @index(Global, NTuple) + + scaling = closure.scaling + + ρ₀ = buoyancy.model.equation_of_state.reference_density + g = buoyancy.model.gravitational_acceleration + + @inbounds k_first = first_index[i, j, 1] + @inbounds k_last = last_index[i, j, 1] + + quiescent = quiescent_condition(k_first, k_last) + + @inbounds k_tracer = clamp_k_interior(k_first + k - 1, grid) + + @inbounds k₋₂ = clamp_k_interior(k_tracer - 2, grid) + @inbounds k₋₁ = clamp_k_interior(k_tracer - 1, grid) + @inbounds k₀ = clamp_k_interior(k_tracer, grid) + @inbounds k₊₁ = clamp_k_interior(k_tracer + 1, grid) + @inbounds k₊₂ = clamp_k_interior(k_tracer + 2, grid) + + T, S = tracers.T, tracers.S + + @inbounds input[1, i, j, k] = ifelse(quiescent, 0, atan(Ri[i, j, k₋₂])) + @inbounds input[2, i, j, k] = ifelse(quiescent, 0, atan(Ri[i, j, k₋₁])) + @inbounds input[3, i, j, k] = ifelse(quiescent, 0, atan(Ri[i, j, k₀])) + @inbounds input[4, i, j, k] = ifelse(quiescent, 0, atan(Ri[i, j, k₊₁])) + @inbounds input[5, i, j, k] = ifelse(quiescent, 0, atan(Ri[i, j, k₊₂])) + + @inbounds input[6, i, j, k] = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k₋₂, grid, T))) + @inbounds input[7, i, j, k] = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k₋₁, grid, T))) + @inbounds input[8, i, j, k] = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k₀, grid, T))) + @inbounds input[9, i, j, k] = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k₊₁, grid, T))) + @inbounds input[10, i, j, k] = scaling.∂T∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k₊₂, grid, T))) + + @inbounds input[11, i, j, k] = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k₋₂, grid, S))) + @inbounds input[12, i, j, k] = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k₋₁, grid, S))) + @inbounds input[13, i, j, k] = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k₀, grid, S))) + @inbounds input[14, i, j, k] = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k₊₁, grid, S))) + @inbounds input[15, i, j, k] = scaling.∂S∂z(ifelse(quiescent, 0, ∂zᶜᶜᶠ(i, j, k₊₂, grid, S))) + + @inbounds input[16, i, j, k] = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k₋₂, grid, buoyancy, tracers) / g)) + @inbounds input[17, i, j, k] = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k₋₁, grid, buoyancy, tracers) / g)) + @inbounds input[18, i, j, k] = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k₀, grid, buoyancy, tracers) / g)) + @inbounds input[19, i, j, k] = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k₊₁, grid, buoyancy, tracers) / g)) + @inbounds input[20, i, j, k] = scaling.∂ρ∂z(ifelse(quiescent, 0, -ρ₀ * ∂z_b(i, j, k₊₂, grid, buoyancy, tracers) / g)) + + @inbounds input[21, i, j, k] = scaling.wb(top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers))) + +end + +@kernel function _find_NN_active_region!(Ri, grid, Riᶜ, first_index, last_index, closure::NNFluxClosure) + i, j = @index(Global, NTuple) + top_index = grid.Nz + 1 + grid_point_above_kappa = closure.grid_point_above + grid_point_below_kappa = closure.grid_point_below + + # Find the first index of the background κᶜ + kloc = grid.Nz+1 + @inbounds for k in grid.Nz:-1:2 + kloc = ifelse(Ri[i, j, k] < Riᶜ, k, kloc) + end + + background_κ_index = kloc - 1 + nonbackground_κ_index = background_κ_index + 1 + + @inbounds last_index[i, j, 1] = ifelse(nonbackground_κ_index == top_index, top_index - 1, clamp_k_interior(background_κ_index + grid_point_above_kappa, grid)) + @inbounds first_index[i, j, 1] = ifelse(nonbackground_κ_index == top_index, top_index, clamp_k_interior(background_κ_index - grid_point_below_kappa + 1, grid)) +end + +@inline function quiescent_condition(lo, hi) + return hi < lo +end + +@inline function within_zone_condition(k, lo, hi) + return (k >= lo) & (k <= hi) +end + +@inline function clamp_k_interior(k, grid) + kmax = grid.Nz + kmin = 2 + + return clamp(k, kmin, kmax) +end + +@kernel function _fill_adjust_nn_fluxes!(diffusivities, first_index, last_index, wrk_wT, wrk_wS, grid, closure::NNFluxClosure, tracers, velocities, buoyancy, top_tracer_bcs, clock) + i, j, k = @index(Global, NTuple) + scaling = closure.scaling + + @inbounds k_first = first_index[i, j, 1] + @inbounds k_last = last_index[i, j, 1] + + convecting = top_buoyancy_flux(i, j, grid, buoyancy, top_tracer_bcs, clock, merge(velocities, tracers)) > 0 + quiescent = quiescent_condition(k_first, k_last) + within_zone = within_zone_condition(k, k_first, k_last) + + N_levels = closure.grid_point_above + closure.grid_point_below + @inbounds k_wrk = clamp(k - k_first + 1, 1, N_levels) + + NN_active = convecting & !quiescent & within_zone + + @inbounds diffusivities.wT[i, j, k] = ifelse(NN_active, scaling.wT.σ * wrk_wT[i, j, k_wrk], 0) + @inbounds diffusivities.wS[i, j, k] = ifelse(NN_active, scaling.wS.σ * wrk_wS[i, j, k_wrk], 0) +end + +# Write here your constructor +# NNFluxClosure() = ... insert NN here ... (make sure it is on GPU if you need it on GPU!) + +const NNC = NNFluxClosure + +##### +##### Abstract Smagorinsky functionality +##### + +# Horizontal fluxes are zero! +@inline viscous_flux_wz( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_wy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_ux( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vx( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_uy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vy( i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline diffusive_flux_x(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) +@inline diffusive_flux_y(i, j, k, grid, clo::NNC, K, ::Val{tracer_index}, c, clock, fields, buoyancy) where tracer_index = zero(grid) + +# Viscous fluxes are zero (for now) +@inline viscous_flux_uz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) +@inline viscous_flux_vz(i, j, k, grid, clo::NNC, K, clk, fields, b) = zero(grid) + +# The only function extended by NNFluxClosure +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{1}, c, clock, fields, buoyancy) = @inbounds K.wT[i, j, k] +@inline diffusive_flux_z(i, j, k, grid, clo::NNC, K, ::Val{2}, c, clock, fields, buoyancy) = @inbounds K.wS[i, j, k] \ No newline at end of file From 88f7fa486a3357df9f76216a3892b49fd295948e Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 26 Nov 2024 13:06:31 -0500 Subject: [PATCH 107/114] run double gyre with new NN closure --- ..._nof_BBLRifirstzone510_doublegyre_model.jl | 1263 ++++++++++++++++ ...firstzone510_doublegyre_model_modewater.jl | 1275 +++++++++++++++++ 2 files changed, 2538 insertions(+) create mode 100644 NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model.jl create mode 100644 NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl diff --git a/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model.jl b/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model.jl new file mode 100644 index 0000000000..7a5f9d2a2e --- /dev/null +++ b/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model.jl @@ -0,0 +1,1263 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("NN_closure_global_Ri_nof_BBLRifirstzone510.jl") +include("xin_kai_vertical_diffusivity_local_2step_new.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + + +#%% +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_zWENO5_NN_closure_NDE5_Ri_BBLRifirztzone510_temp" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +nn_closure = NNFluxClosure(model_architecture) +base_closure = XinKaiLocalVerticalDiffusivity() +closure = (base_closure, nn_closure) + +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/30days + +##### +##### Forcing and initial condition +##### +@inline T_initial(x, y, z) = 10 + 20 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10800days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); +first_index = model.diffusivity_fields[2].first_index +last_index = model.diffusivity_fields[2].last_index +wT_NN = model.diffusivity_fields[2].wT +wS_NN = model.diffusivity_fields[2].wS + +κ = model.diffusivity_fields[1].κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) + +wT = wT_NN + wT_base +wS = wS_NN + wS_base + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +@inline function get_top_buoyancy_flux(i, j, k, grid, buoyancy, T_bc, S_bc, velocities, tracers, clock) + return top_buoyancy_flux(i, j, grid, buoyancy, (; T=T_bc, S=S_bc), clock, merge(velocities, tracers)) +end + +Qb = KernelFunctionOperation{Center, Center, Nothing}(get_top_buoyancy_flux, model.grid, model.buoyancy, T.boundary_conditions.top, S.boundary_conditions.top, model.velocities, model.tracers, model.clock) +Qb = Field(Qb) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + +wT_NNbar_zonal = Average(wT_NN, dims=1) +wS_NNbar_zonal = Average(wS_NN, dims=1) + +wT_basebar_zonal = Average(wT_base, dims=1) +wS_basebar_zonal = Average(wS_base, dims=1) + +wTbar_zonal = Average(wT, dims=1) +wSbar_zonal = Average(wS, dims=1) + +outputs = (; u, v, w, T, S, ρ, N², wT_NN, wS_NN, wT_base, wS_base, wT, wS) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal, wT_NNbar_zonal, wS_NNbar_zonal, wT_basebar_zonal, wS_basebar_zonal, wTbar_zonal, wSbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_5] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_5", + indices = (:, 5, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_15] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_15", + indices = (:, 15, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_25] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_25", + indices = (:, 25, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_35] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_35", + indices = (:, 35, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_45] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_45", + indices = (:, 45, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_55] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_55", + indices = (:, 55, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_65] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_65", + indices = (:, 65, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_75] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_75", + indices = (:, 75, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_85] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_85", + indices = (:, 85, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_95] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_95", + indices = (:, 95, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; first_index, last_index, Qb), + filename = "$(FILE_DIR)/instantaneous_fields_NN_active_diagnostics", + indices = (:, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end + +#%% +@info "Recording T fields and fluxes in yz" + +fieldname = "T" +fluxname = "wT_NN" +field_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fieldname, backend=OnDisk()) +field_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fieldname, backend=OnDisk()) +field_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fieldname, backend=OnDisk()) +field_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fieldname, backend=OnDisk()) +field_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fieldname, backend=OnDisk()) +field_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fieldname, backend=OnDisk()) +field_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fieldname, backend=OnDisk()) +field_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fieldname, backend=OnDisk()) +field_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fieldname, backend=OnDisk()) +field_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fluxname, backend=OnDisk()) +flux_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fluxname, backend=OnDisk()) +flux_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fluxname, backend=OnDisk()) +flux_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fluxname, backend=OnDisk()) +flux_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fluxname, backend=OnDisk()) +flux_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fluxname, backend=OnDisk()) +flux_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fluxname, backend=OnDisk()) +flux_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fluxname, backend=OnDisk()) +flux_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fluxname, backend=OnDisk()) +flux_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fluxname, backend=OnDisk()) + +xC = field_NN_data_00.grid.xᶜᵃᵃ[1:field_NN_data_00.grid.Nx] +yC = field_NN_data_00.grid.yᵃᶜᵃ[1:field_NN_data_00.grid.Ny] +zC = field_NN_data_00.grid.zᵃᵃᶜ[1:field_NN_data_00.grid.Nz] +zF = field_NN_data_00.grid.zᵃᵃᶠ[1:field_NN_data_00.grid.Nz+1] + +Nt = length(field_NN_data_90) +times = field_NN_data_00.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_00 = CairoMakie.Axis(fig[1, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_00.grid.xᶜᵃᵃ[field_NN_data_00.indices[1][1]] / 1000) km") +axfield_10 = CairoMakie.Axis(fig[1, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_10.grid.xᶜᵃᵃ[field_NN_data_10.indices[1][1]] / 1000) km") +axfield_20 = CairoMakie.Axis(fig[2, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_20.grid.xᶜᵃᵃ[field_NN_data_20.indices[1][1]] / 1000) km") +axfield_30 = CairoMakie.Axis(fig[2, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_30.grid.xᶜᵃᵃ[field_NN_data_30.indices[1][1]] / 1000) km") +axfield_40 = CairoMakie.Axis(fig[3, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_40.grid.xᶜᵃᵃ[field_NN_data_40.indices[1][1]] / 1000) km") +axfield_50 = CairoMakie.Axis(fig[3, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_50.grid.xᶜᵃᵃ[field_NN_data_50.indices[1][1]] / 1000) km") +axfield_60 = CairoMakie.Axis(fig[4, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_60.grid.xᶜᵃᵃ[field_NN_data_60.indices[1][1]] / 1000) km") +axfield_70 = CairoMakie.Axis(fig[4, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_70.grid.xᶜᵃᵃ[field_NN_data_70.indices[1][1]] / 1000) km") +axfield_80 = CairoMakie.Axis(fig[5, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_80.grid.xᶜᵃᵃ[field_NN_data_80.indices[1][1]] / 1000) km") +axfield_90 = CairoMakie.Axis(fig[5, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_90.grid.xᶜᵃᵃ[field_NN_data_90.indices[1][1]] / 1000) km") + +axflux_00 = CairoMakie.Axis(fig[1, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_00.grid.xᶜᵃᵃ[flux_NN_data_00.indices[1][1]] / 1000) km") +axflux_10 = CairoMakie.Axis(fig[1, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_10.grid.xᶜᵃᵃ[flux_NN_data_10.indices[1][1]] / 1000) km") +axflux_20 = CairoMakie.Axis(fig[2, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_20.grid.xᶜᵃᵃ[flux_NN_data_20.indices[1][1]] / 1000) km") +axflux_30 = CairoMakie.Axis(fig[2, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_30.grid.xᶜᵃᵃ[flux_NN_data_30.indices[1][1]] / 1000) km") +axflux_40 = CairoMakie.Axis(fig[3, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_40.grid.xᶜᵃᵃ[flux_NN_data_40.indices[1][1]] / 1000) km") +axflux_50 = CairoMakie.Axis(fig[3, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_50.grid.xᶜᵃᵃ[flux_NN_data_50.indices[1][1]] / 1000) km") +axflux_60 = CairoMakie.Axis(fig[4, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_60.grid.xᶜᵃᵃ[flux_NN_data_60.indices[1][1]] / 1000) km") +axflux_70 = CairoMakie.Axis(fig[4, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_70.grid.xᶜᵃᵃ[flux_NN_data_70.indices[1][1]] / 1000) km") +axflux_80 = CairoMakie.Axis(fig[5, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_80.grid.xᶜᵃᵃ[flux_NN_data_80.indices[1][1]] / 1000) km") +axflux_90 = CairoMakie.Axis(fig[5, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_90.grid.xᶜᵃᵃ[flux_NN_data_90.indices[1][1]] / 1000) km") + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lim = (find_min(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (find_min(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (-maximum(abs, flux_lim), maximum(abs, flux_lim)) + +NN_00ₙ = @lift interior(field_NN_data_00[$n], 1, :, zC_indices) +NN_10ₙ = @lift interior(field_NN_data_10[$n], 1, :, zC_indices) +NN_20ₙ = @lift interior(field_NN_data_20[$n], 1, :, zC_indices) +NN_30ₙ = @lift interior(field_NN_data_30[$n], 1, :, zC_indices) +NN_40ₙ = @lift interior(field_NN_data_40[$n], 1, :, zC_indices) +NN_50ₙ = @lift interior(field_NN_data_50[$n], 1, :, zC_indices) +NN_60ₙ = @lift interior(field_NN_data_60[$n], 1, :, zC_indices) +NN_70ₙ = @lift interior(field_NN_data_70[$n], 1, :, zC_indices) +NN_80ₙ = @lift interior(field_NN_data_80[$n], 1, :, zC_indices) +NN_90ₙ = @lift interior(field_NN_data_90[$n], 1, :, zC_indices) + +flux_00ₙ = @lift interior(flux_NN_data_00[$n], 1, :, zF_indices) +flux_10ₙ = @lift interior(flux_NN_data_10[$n], 1, :, zF_indices) +flux_20ₙ = @lift interior(flux_NN_data_20[$n], 1, :, zF_indices) +flux_30ₙ = @lift interior(flux_NN_data_30[$n], 1, :, zF_indices) +flux_40ₙ = @lift interior(flux_NN_data_40[$n], 1, :, zF_indices) +flux_50ₙ = @lift interior(flux_NN_data_50[$n], 1, :, zF_indices) +flux_60ₙ = @lift interior(flux_NN_data_60[$n], 1, :, zF_indices) +flux_70ₙ = @lift interior(flux_NN_data_70[$n], 1, :, zF_indices) +flux_80ₙ = @lift interior(flux_NN_data_80[$n], 1, :, zF_indices) +flux_90ₙ = @lift interior(flux_NN_data_90[$n], 1, :, zF_indices) + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_00_surface = heatmap!(axfield_00, yC, zC[zC_indices], NN_00ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_10_surface = heatmap!(axfield_10, yC, zC[zC_indices], NN_10ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_20_surface = heatmap!(axfield_20, yC, zC[zC_indices], NN_20ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_30_surface = heatmap!(axfield_30, yC, zC[zC_indices], NN_30ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_40_surface = heatmap!(axfield_40, yC, zC[zC_indices], NN_40ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_50_surface = heatmap!(axfield_50, yC, zC[zC_indices], NN_50ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_60_surface = heatmap!(axfield_60, yC, zC[zC_indices], NN_60ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_70_surface = heatmap!(axfield_70, yC, zC[zC_indices], NN_70ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_80_surface = heatmap!(axfield_80, yC, zC[zC_indices], NN_80ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_90_surface = heatmap!(axfield_90, yC, zC[zC_indices], NN_90ₙ, colormap=colorscheme_field, colorrange=field_lim) + +flux_00_surface = heatmap!(axflux_00, yC, zC[zF_indices], flux_00ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_10_surface = heatmap!(axflux_10, yC, zC[zF_indices], flux_10ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_20_surface = heatmap!(axflux_20, yC, zC[zF_indices], flux_20ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_30_surface = heatmap!(axflux_30, yC, zC[zF_indices], flux_30ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_40_surface = heatmap!(axflux_40, yC, zC[zF_indices], flux_40ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_50_surface = heatmap!(axflux_50, yC, zC[zF_indices], flux_50ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_60_surface = heatmap!(axflux_60, yC, zC[zF_indices], flux_60ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_70_surface = heatmap!(axflux_70, yC, zC[zF_indices], flux_70ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_80_surface = heatmap!(axflux_80, yC, zC[zF_indices], flux_80ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_90_surface = heatmap!(axflux_90, yC, zC[zF_indices], flux_90ₙ, colormap=colorscheme_flux, colorrange=flux_lim) + +Colorbar(fig[1:5, 2], field_00_surface, label="Field") +Colorbar(fig[1:5, 4], flux_00_surface, label="NN Flux") +Colorbar(fig[1:5, 6], field_00_surface, label="Field") +Colorbar(fig[1:5, 8], flux_00_surface, label="NN Flux") + +xlims!(axfield_00, minimum(yC), maximum(yC)) +xlims!(axfield_10, minimum(yC), maximum(yC)) +xlims!(axfield_20, minimum(yC), maximum(yC)) +xlims!(axfield_30, minimum(yC), maximum(yC)) +xlims!(axfield_40, minimum(yC), maximum(yC)) +xlims!(axfield_50, minimum(yC), maximum(yC)) +xlims!(axfield_60, minimum(yC), maximum(yC)) +xlims!(axfield_70, minimum(yC), maximum(yC)) +xlims!(axfield_80, minimum(yC), maximum(yC)) +xlims!(axfield_90, minimum(yC), maximum(yC)) + +ylims!(axfield_00, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_10, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_20, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_30, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_40, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_50, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_60, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_70, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_80, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_90, minimum(zC[zC_indices]), maximum(zC[zC_indices])) + +xlims!(axflux_00, minimum(yC), maximum(yC)) +xlims!(axflux_10, minimum(yC), maximum(yC)) +xlims!(axflux_20, minimum(yC), maximum(yC)) +xlims!(axflux_30, minimum(yC), maximum(yC)) +xlims!(axflux_40, minimum(yC), maximum(yC)) +xlims!(axflux_50, minimum(yC), maximum(yC)) +xlims!(axflux_60, minimum(yC), maximum(yC)) +xlims!(axflux_70, minimum(yC), maximum(yC)) +xlims!(axflux_80, minimum(yC), maximum(yC)) +xlims!(axflux_90, minimum(yC), maximum(yC)) + +ylims!(axflux_00, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_10, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_20, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_30, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_40, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_50, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_60, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_70, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_80, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_90, minimum(zF[zF_indices]), maximum(zF[zF_indices])) + +title_str = @lift "Temperature (°C), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +# save("./Output/compare_3D_instantaneous_fields_slices_NNclosure_fluxes.png", fig) +# display(fig) +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_yzslices_fluxes_T.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end +#%% +@info "Recording S fields and fluxes in yz" + +fieldname = "S" +fluxname = "wS_NN" +field_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fieldname, backend=OnDisk()) +field_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fieldname, backend=OnDisk()) +field_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fieldname, backend=OnDisk()) +field_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fieldname, backend=OnDisk()) +field_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fieldname, backend=OnDisk()) +field_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fieldname, backend=OnDisk()) +field_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fieldname, backend=OnDisk()) +field_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fieldname, backend=OnDisk()) +field_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fieldname, backend=OnDisk()) +field_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fluxname, backend=OnDisk()) +flux_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fluxname, backend=OnDisk()) +flux_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fluxname, backend=OnDisk()) +flux_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fluxname, backend=OnDisk()) +flux_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fluxname, backend=OnDisk()) +flux_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fluxname, backend=OnDisk()) +flux_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fluxname, backend=OnDisk()) +flux_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fluxname, backend=OnDisk()) +flux_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fluxname, backend=OnDisk()) +flux_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fluxname, backend=OnDisk()) + +xC = field_NN_data_00.grid.xᶜᵃᵃ[1:field_NN_data_00.grid.Nx] +yC = field_NN_data_00.grid.yᵃᶜᵃ[1:field_NN_data_00.grid.Ny] +zC = field_NN_data_00.grid.zᵃᵃᶜ[1:field_NN_data_00.grid.Nz] +zF = field_NN_data_00.grid.zᵃᵃᶠ[1:field_NN_data_00.grid.Nz+1] + +Nt = length(field_NN_data_90) +times = field_NN_data_00.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_00 = CairoMakie.Axis(fig[1, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_00.grid.xᶜᵃᵃ[field_NN_data_00.indices[1][1]] / 1000) km") +axfield_10 = CairoMakie.Axis(fig[1, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_10.grid.xᶜᵃᵃ[field_NN_data_10.indices[1][1]] / 1000) km") +axfield_20 = CairoMakie.Axis(fig[2, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_20.grid.xᶜᵃᵃ[field_NN_data_20.indices[1][1]] / 1000) km") +axfield_30 = CairoMakie.Axis(fig[2, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_30.grid.xᶜᵃᵃ[field_NN_data_30.indices[1][1]] / 1000) km") +axfield_40 = CairoMakie.Axis(fig[3, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_40.grid.xᶜᵃᵃ[field_NN_data_40.indices[1][1]] / 1000) km") +axfield_50 = CairoMakie.Axis(fig[3, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_50.grid.xᶜᵃᵃ[field_NN_data_50.indices[1][1]] / 1000) km") +axfield_60 = CairoMakie.Axis(fig[4, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_60.grid.xᶜᵃᵃ[field_NN_data_60.indices[1][1]] / 1000) km") +axfield_70 = CairoMakie.Axis(fig[4, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_70.grid.xᶜᵃᵃ[field_NN_data_70.indices[1][1]] / 1000) km") +axfield_80 = CairoMakie.Axis(fig[5, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_80.grid.xᶜᵃᵃ[field_NN_data_80.indices[1][1]] / 1000) km") +axfield_90 = CairoMakie.Axis(fig[5, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_90.grid.xᶜᵃᵃ[field_NN_data_90.indices[1][1]] / 1000) km") + +axflux_00 = CairoMakie.Axis(fig[1, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_00.grid.xᶜᵃᵃ[flux_NN_data_00.indices[1][1]] / 1000) km") +axflux_10 = CairoMakie.Axis(fig[1, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_10.grid.xᶜᵃᵃ[flux_NN_data_10.indices[1][1]] / 1000) km") +axflux_20 = CairoMakie.Axis(fig[2, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_20.grid.xᶜᵃᵃ[flux_NN_data_20.indices[1][1]] / 1000) km") +axflux_30 = CairoMakie.Axis(fig[2, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_30.grid.xᶜᵃᵃ[flux_NN_data_30.indices[1][1]] / 1000) km") +axflux_40 = CairoMakie.Axis(fig[3, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_40.grid.xᶜᵃᵃ[flux_NN_data_40.indices[1][1]] / 1000) km") +axflux_50 = CairoMakie.Axis(fig[3, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_50.grid.xᶜᵃᵃ[flux_NN_data_50.indices[1][1]] / 1000) km") +axflux_60 = CairoMakie.Axis(fig[4, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_60.grid.xᶜᵃᵃ[flux_NN_data_60.indices[1][1]] / 1000) km") +axflux_70 = CairoMakie.Axis(fig[4, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_70.grid.xᶜᵃᵃ[flux_NN_data_70.indices[1][1]] / 1000) km") +axflux_80 = CairoMakie.Axis(fig[5, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_80.grid.xᶜᵃᵃ[flux_NN_data_80.indices[1][1]] / 1000) km") +axflux_90 = CairoMakie.Axis(fig[5, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_90.grid.xᶜᵃᵃ[flux_NN_data_90.indices[1][1]] / 1000) km") + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lim = (find_min(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (find_min(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (-maximum(abs, flux_lim), maximum(abs, flux_lim)) + +NN_00ₙ = @lift interior(field_NN_data_00[$n], 1, :, zC_indices) +NN_10ₙ = @lift interior(field_NN_data_10[$n], 1, :, zC_indices) +NN_20ₙ = @lift interior(field_NN_data_20[$n], 1, :, zC_indices) +NN_30ₙ = @lift interior(field_NN_data_30[$n], 1, :, zC_indices) +NN_40ₙ = @lift interior(field_NN_data_40[$n], 1, :, zC_indices) +NN_50ₙ = @lift interior(field_NN_data_50[$n], 1, :, zC_indices) +NN_60ₙ = @lift interior(field_NN_data_60[$n], 1, :, zC_indices) +NN_70ₙ = @lift interior(field_NN_data_70[$n], 1, :, zC_indices) +NN_80ₙ = @lift interior(field_NN_data_80[$n], 1, :, zC_indices) +NN_90ₙ = @lift interior(field_NN_data_90[$n], 1, :, zC_indices) + +flux_00ₙ = @lift interior(flux_NN_data_00[$n], 1, :, zF_indices) +flux_10ₙ = @lift interior(flux_NN_data_10[$n], 1, :, zF_indices) +flux_20ₙ = @lift interior(flux_NN_data_20[$n], 1, :, zF_indices) +flux_30ₙ = @lift interior(flux_NN_data_30[$n], 1, :, zF_indices) +flux_40ₙ = @lift interior(flux_NN_data_40[$n], 1, :, zF_indices) +flux_50ₙ = @lift interior(flux_NN_data_50[$n], 1, :, zF_indices) +flux_60ₙ = @lift interior(flux_NN_data_60[$n], 1, :, zF_indices) +flux_70ₙ = @lift interior(flux_NN_data_70[$n], 1, :, zF_indices) +flux_80ₙ = @lift interior(flux_NN_data_80[$n], 1, :, zF_indices) +flux_90ₙ = @lift interior(flux_NN_data_90[$n], 1, :, zF_indices) + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_00_surface = heatmap!(axfield_00, yC, zC[zC_indices], NN_00ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_10_surface = heatmap!(axfield_10, yC, zC[zC_indices], NN_10ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_20_surface = heatmap!(axfield_20, yC, zC[zC_indices], NN_20ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_30_surface = heatmap!(axfield_30, yC, zC[zC_indices], NN_30ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_40_surface = heatmap!(axfield_40, yC, zC[zC_indices], NN_40ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_50_surface = heatmap!(axfield_50, yC, zC[zC_indices], NN_50ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_60_surface = heatmap!(axfield_60, yC, zC[zC_indices], NN_60ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_70_surface = heatmap!(axfield_70, yC, zC[zC_indices], NN_70ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_80_surface = heatmap!(axfield_80, yC, zC[zC_indices], NN_80ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_90_surface = heatmap!(axfield_90, yC, zC[zC_indices], NN_90ₙ, colormap=colorscheme_field, colorrange=field_lim) + +flux_00_surface = heatmap!(axflux_00, yC, zC[zF_indices], flux_00ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_10_surface = heatmap!(axflux_10, yC, zC[zF_indices], flux_10ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_20_surface = heatmap!(axflux_20, yC, zC[zF_indices], flux_20ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_30_surface = heatmap!(axflux_30, yC, zC[zF_indices], flux_30ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_40_surface = heatmap!(axflux_40, yC, zC[zF_indices], flux_40ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_50_surface = heatmap!(axflux_50, yC, zC[zF_indices], flux_50ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_60_surface = heatmap!(axflux_60, yC, zC[zF_indices], flux_60ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_70_surface = heatmap!(axflux_70, yC, zC[zF_indices], flux_70ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_80_surface = heatmap!(axflux_80, yC, zC[zF_indices], flux_80ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_90_surface = heatmap!(axflux_90, yC, zC[zF_indices], flux_90ₙ, colormap=colorscheme_flux, colorrange=flux_lim) + +Colorbar(fig[1:5, 2], field_00_surface, label="Field") +Colorbar(fig[1:5, 4], flux_00_surface, label="NN Flux") +Colorbar(fig[1:5, 6], field_00_surface, label="Field") +Colorbar(fig[1:5, 8], flux_00_surface, label="NN Flux") + +xlims!(axfield_00, minimum(yC), maximum(yC)) +xlims!(axfield_10, minimum(yC), maximum(yC)) +xlims!(axfield_20, minimum(yC), maximum(yC)) +xlims!(axfield_30, minimum(yC), maximum(yC)) +xlims!(axfield_40, minimum(yC), maximum(yC)) +xlims!(axfield_50, minimum(yC), maximum(yC)) +xlims!(axfield_60, minimum(yC), maximum(yC)) +xlims!(axfield_70, minimum(yC), maximum(yC)) +xlims!(axfield_80, minimum(yC), maximum(yC)) +xlims!(axfield_90, minimum(yC), maximum(yC)) + +ylims!(axfield_00, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_10, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_20, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_30, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_40, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_50, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_60, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_70, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_80, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_90, minimum(zC[zC_indices]), maximum(zC[zC_indices])) + +xlims!(axflux_00, minimum(yC), maximum(yC)) +xlims!(axflux_10, minimum(yC), maximum(yC)) +xlims!(axflux_20, minimum(yC), maximum(yC)) +xlims!(axflux_30, minimum(yC), maximum(yC)) +xlims!(axflux_40, minimum(yC), maximum(yC)) +xlims!(axflux_50, minimum(yC), maximum(yC)) +xlims!(axflux_60, minimum(yC), maximum(yC)) +xlims!(axflux_70, minimum(yC), maximum(yC)) +xlims!(axflux_80, minimum(yC), maximum(yC)) +xlims!(axflux_90, minimum(yC), maximum(yC)) + +ylims!(axflux_00, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_10, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_20, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_30, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_40, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_50, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_60, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_70, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_80, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_90, minimum(zF[zF_indices]), maximum(zF[zF_indices])) + +title_str = @lift "Salinity (psu), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_yzslices_fluxes_S.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +#%% +@info "Recording T fields and fluxes in xz" + +fieldname = "T" +fluxname = "wT_NN" + +field_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fieldname, backend=OnDisk()) +field_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fieldname, backend=OnDisk()) +field_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fieldname, backend=OnDisk()) +field_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fieldname, backend=OnDisk()) +field_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fieldname, backend=OnDisk()) +field_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fieldname, backend=OnDisk()) +field_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fieldname, backend=OnDisk()) +field_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fieldname, backend=OnDisk()) +field_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fieldname, backend=OnDisk()) +field_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fluxname, backend=OnDisk()) +flux_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fluxname, backend=OnDisk()) +flux_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fluxname, backend=OnDisk()) +flux_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fluxname, backend=OnDisk()) +flux_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fluxname, backend=OnDisk()) +flux_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fluxname, backend=OnDisk()) +flux_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fluxname, backend=OnDisk()) +flux_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fluxname, backend=OnDisk()) +flux_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fluxname, backend=OnDisk()) +flux_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fluxname, backend=OnDisk()) + +field_datas = [field_NN_data_5, field_NN_data_15, field_NN_data_25, field_NN_data_35, field_NN_data_45, field_NN_data_55, field_NN_data_65, field_NN_data_75, field_NN_data_85, field_NN_data_95] +flux_datas = [flux_NN_data_5, flux_NN_data_15, flux_NN_data_25, flux_NN_data_35, flux_NN_data_45, flux_NN_data_55, flux_NN_data_65, flux_NN_data_75, flux_NN_data_85, flux_NN_data_95] + +xC = field_NN_data_5.grid.xᶜᵃᵃ[1:field_NN_data_5.grid.Nx] +yC = field_NN_data_5.grid.yᵃᶜᵃ[1:field_NN_data_5.grid.Ny] +zC = field_NN_data_5.grid.zᵃᵃᶜ[1:field_NN_data_5.grid.Nz] +zF = field_NN_data_5.grid.zᵃᵃᶠ[1:field_NN_data_5.grid.Nz+1] + +Nt = length(field_NN_data_95) +times = field_NN_data_5.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_5 = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_5.grid.yᵃᶜᵃ[field_NN_data_5.indices[2][1]] / 1000) km") +axfield_15 = CairoMakie.Axis(fig[1, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_15.grid.yᵃᶜᵃ[field_NN_data_15.indices[2][1]] / 1000) km") +axfield_25 = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_25.grid.yᵃᶜᵃ[field_NN_data_25.indices[2][1]] / 1000) km") +axfield_35 = CairoMakie.Axis(fig[2, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_35.grid.yᵃᶜᵃ[field_NN_data_35.indices[2][1]] / 1000) km") +axfield_45 = CairoMakie.Axis(fig[3, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_45.grid.yᵃᶜᵃ[field_NN_data_45.indices[2][1]] / 1000) km") +axfield_55 = CairoMakie.Axis(fig[3, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_55.grid.yᵃᶜᵃ[field_NN_data_55.indices[2][1]] / 1000) km") +axfield_65 = CairoMakie.Axis(fig[4, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_65.grid.yᵃᶜᵃ[field_NN_data_65.indices[2][1]] / 1000) km") +axfield_75 = CairoMakie.Axis(fig[4, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_75.grid.yᵃᶜᵃ[field_NN_data_75.indices[2][1]] / 1000) km") +axfield_85 = CairoMakie.Axis(fig[5, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_85.grid.yᵃᶜᵃ[field_NN_data_85.indices[2][1]] / 1000) km") +axfield_95 = CairoMakie.Axis(fig[5, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_95.grid.yᵃᶜᵃ[field_NN_data_95.indices[2][1]] / 1000) km") + +axflux_5 = CairoMakie.Axis(fig[1, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_5.grid.yᵃᶜᵃ[flux_NN_data_5.indices[2][1]] / 1000) km") +axflux_15 = CairoMakie.Axis(fig[1, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_15.grid.yᵃᶜᵃ[flux_NN_data_15.indices[2][1]] / 1000) km") +axflux_25 = CairoMakie.Axis(fig[2, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_25.grid.yᵃᶜᵃ[flux_NN_data_25.indices[2][1]] / 1000) km") +axflux_35 = CairoMakie.Axis(fig[2, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_35.grid.yᵃᶜᵃ[flux_NN_data_35.indices[2][1]] / 1000) km") +axflux_45 = CairoMakie.Axis(fig[3, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_45.grid.yᵃᶜᵃ[flux_NN_data_45.indices[2][1]] / 1000) km") +axflux_55 = CairoMakie.Axis(fig[3, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_55.grid.yᵃᶜᵃ[flux_NN_data_55.indices[2][1]] / 1000) km") +axflux_65 = CairoMakie.Axis(fig[4, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_65.grid.yᵃᶜᵃ[flux_NN_data_65.indices[2][1]] / 1000) km") +axflux_75 = CairoMakie.Axis(fig[4, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_75.grid.yᵃᶜᵃ[flux_NN_data_75.indices[2][1]] / 1000) km") +axflux_85 = CairoMakie.Axis(fig[5, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_85.grid.yᵃᶜᵃ[flux_NN_data_85.indices[2][1]] / 1000) km") +axflux_95 = CairoMakie.Axis(fig[5, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_95.grid.yᵃᶜᵃ[flux_NN_data_95.indices[2][1]] / 1000) km") + +axfields = [axfield_5, axfield_15, axfield_25, axfield_35, axfield_45, axfield_55, axfield_65, axfield_75, axfield_85, axfield_95] +axfluxes = [axflux_5, axflux_15, axflux_25, axflux_35, axflux_45, axflux_55, axflux_65, axflux_75, axflux_85, axflux_95] + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lims = [(find_min(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices))) for field_data in field_datas] + +flux_lims = [(find_min(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices))) for flux_data in flux_datas] + +flux_lims = [(-maximum(abs, flux_lim), maximum(abs, flux_lim)) for flux_lim in flux_lims] + +NNₙs = [@lift interior(field_data[$n], :, 1, zC_indices) for field_data in field_datas] +fluxₙs = [@lift interior(flux_data[$n], :, 1, zF_indices) for flux_data in flux_datas] + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_5_surface = heatmap!(axfield_5, xC, zC[zC_indices], NN_5ₙ, colormap=colorscheme_field, colorrange=field_lims[1]) +field_15_surface = heatmap!(axfield_15, xC, zC[zC_indices], NN_15ₙ, colormap=colorscheme_field, colorrange=field_lims[2]) +field_25_surface = heatmap!(axfield_25, xC, zC[zC_indices], NN_25ₙ, colormap=colorscheme_field, colorrange=field_lims[3]) +field_35_surface = heatmap!(axfield_35, xC, zC[zC_indices], NN_35ₙ, colormap=colorscheme_field, colorrange=field_lims[4]) +field_45_surface = heatmap!(axfield_45, xC, zC[zC_indices], NN_45ₙ, colormap=colorscheme_field, colorrange=field_lims[5]) +field_55_surface = heatmap!(axfield_55, xC, zC[zC_indices], NN_55ₙ, colormap=colorscheme_field, colorrange=field_lims[6]) +field_65_surface = heatmap!(axfield_65, xC, zC[zC_indices], NN_65ₙ, colormap=colorscheme_field, colorrange=field_lims[7]) +field_75_surface = heatmap!(axfield_75, xC, zC[zC_indices], NN_75ₙ, colormap=colorscheme_field, colorrange=field_lims[8]) +field_85_surface = heatmap!(axfield_85, xC, zC[zC_indices], NN_85ₙ, colormap=colorscheme_field, colorrange=field_lims[9]) +field_95_surface = heatmap!(axfield_95, xC, zC[zC_indices], NN_95ₙ, colormap=colorscheme_field, colorrange=field_lims[10]) + +flux_5_surface = heatmap!(axflux_5, xC, zC[zF_indices], flux_5ₙ, colormap=colorscheme_flux, colorrange=flux_lims[1]) +flux_15_surface = heatmap!(axflux_15, xC, zC[zF_indices], flux_15ₙ, colormap=colorscheme_flux, colorrange=flux_lims[2]) +flux_25_surface = heatmap!(axflux_25, xC, zC[zF_indices], flux_25ₙ, colormap=colorscheme_flux, colorrange=flux_lims[3]) +flux_35_surface = heatmap!(axflux_35, xC, zC[zF_indices], flux_35ₙ, colormap=colorscheme_flux, colorrange=flux_lims[4]) +flux_45_surface = heatmap!(axflux_45, xC, zC[zF_indices], flux_45ₙ, colormap=colorscheme_flux, colorrange=flux_lims[5]) +flux_55_surface = heatmap!(axflux_55, xC, zC[zF_indices], flux_55ₙ, colormap=colorscheme_flux, colorrange=flux_lims[6]) +flux_65_surface = heatmap!(axflux_65, xC, zC[zF_indices], flux_65ₙ, colormap=colorscheme_flux, colorrange=flux_lims[7]) +flux_75_surface = heatmap!(axflux_75, xC, zC[zF_indices], flux_75ₙ, colormap=colorscheme_flux, colorrange=flux_lims[8]) +flux_85_surface = heatmap!(axflux_85, xC, zC[zF_indices], flux_85ₙ, colormap=colorscheme_flux, colorrange=flux_lims[9]) +flux_95_surface = heatmap!(axflux_95, xC, zC[zF_indices], flux_95ₙ, colormap=colorscheme_flux, colorrange=flux_lims[10]) + +Colorbar(fig[1, 2], field_5_surface, label="Field") +Colorbar(fig[1, 4], flux_5_surface, label="NN Flux") +Colorbar(fig[1, 6], field_15_surface, label="Field") +Colorbar(fig[1, 8], flux_15_surface, label="NN Flux") +Colorbar(fig[2, 2], field_25_surface, label="Field") +Colorbar(fig[2, 4], flux_25_surface, label="NN Flux") +Colorbar(fig[2, 6], field_35_surface, label="Field") +Colorbar(fig[2, 8], flux_35_surface, label="NN Flux") +Colorbar(fig[3, 2], field_45_surface, label="Field") +Colorbar(fig[3, 4], flux_45_surface, label="NN Flux") +Colorbar(fig[3, 6], field_55_surface, label="Field") +Colorbar(fig[3, 8], flux_55_surface, label="NN Flux") +Colorbar(fig[4, 2], field_65_surface, label="Field") +Colorbar(fig[4, 4], flux_65_surface, label="NN Flux") +Colorbar(fig[4, 6], field_75_surface, label="Field") +Colorbar(fig[4, 8], flux_75_surface, label="NN Flux") +Colorbar(fig[5, 2], field_85_surface, label="Field") +Colorbar(fig[5, 4], flux_85_surface, label="NN Flux") +Colorbar(fig[5, 6], field_95_surface, label="Field") +Colorbar(fig[5, 8], flux_95_surface, label="NN Flux") + +for axfield in axfields + xlims!(axfield, minimum(xC), maximum(xC)) + ylims!(axfield, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +end + +for axflux in axfluxes + xlims!(axflux, minimum(xC), maximum(xC)) + ylims!(axflux, minimum(zC[zF_indices]), maximum(zC[zF_indices])) +end + +title_str = @lift "Temperature (°C), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_xzslices_fluxes_T.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end + +#%% +@info "Recording S fields and fluxes in xz" + +fieldname = "S" +fluxname = "wS_NN" + +field_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fieldname, backend=OnDisk()) +field_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fieldname, backend=OnDisk()) +field_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fieldname, backend=OnDisk()) +field_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fieldname, backend=OnDisk()) +field_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fieldname, backend=OnDisk()) +field_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fieldname, backend=OnDisk()) +field_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fieldname, backend=OnDisk()) +field_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fieldname, backend=OnDisk()) +field_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fieldname, backend=OnDisk()) +field_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fluxname, backend=OnDisk()) +flux_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fluxname, backend=OnDisk()) +flux_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fluxname, backend=OnDisk()) +flux_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fluxname, backend=OnDisk()) +flux_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fluxname, backend=OnDisk()) +flux_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fluxname, backend=OnDisk()) +flux_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fluxname, backend=OnDisk()) +flux_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fluxname, backend=OnDisk()) +flux_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fluxname, backend=OnDisk()) +flux_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fluxname, backend=OnDisk()) + +field_datas = [field_NN_data_5, field_NN_data_15, field_NN_data_25, field_NN_data_35, field_NN_data_45, field_NN_data_55, field_NN_data_65, field_NN_data_75, field_NN_data_85, field_NN_data_95] +flux_datas = [flux_NN_data_5, flux_NN_data_15, flux_NN_data_25, flux_NN_data_35, flux_NN_data_45, flux_NN_data_55, flux_NN_data_65, flux_NN_data_75, flux_NN_data_85, flux_NN_data_95] + +xC = field_NN_data_5.grid.xᶜᵃᵃ[1:field_NN_data_5.grid.Nx] +yC = field_NN_data_5.grid.yᵃᶜᵃ[1:field_NN_data_5.grid.Ny] +zC = field_NN_data_5.grid.zᵃᵃᶜ[1:field_NN_data_5.grid.Nz] +zF = field_NN_data_5.grid.zᵃᵃᶠ[1:field_NN_data_5.grid.Nz+1] + +Nt = length(field_NN_data_95) +times = field_NN_data_5.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_5 = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_5.grid.yᵃᶜᵃ[field_NN_data_5.indices[2][1]] / 1000) km") +axfield_15 = CairoMakie.Axis(fig[1, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_15.grid.yᵃᶜᵃ[field_NN_data_15.indices[2][1]] / 1000) km") +axfield_25 = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_25.grid.yᵃᶜᵃ[field_NN_data_25.indices[2][1]] / 1000) km") +axfield_35 = CairoMakie.Axis(fig[2, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_35.grid.yᵃᶜᵃ[field_NN_data_35.indices[2][1]] / 1000) km") +axfield_45 = CairoMakie.Axis(fig[3, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_45.grid.yᵃᶜᵃ[field_NN_data_45.indices[2][1]] / 1000) km") +axfield_55 = CairoMakie.Axis(fig[3, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_55.grid.yᵃᶜᵃ[field_NN_data_55.indices[2][1]] / 1000) km") +axfield_65 = CairoMakie.Axis(fig[4, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_65.grid.yᵃᶜᵃ[field_NN_data_65.indices[2][1]] / 1000) km") +axfield_75 = CairoMakie.Axis(fig[4, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_75.grid.yᵃᶜᵃ[field_NN_data_75.indices[2][1]] / 1000) km") +axfield_85 = CairoMakie.Axis(fig[5, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_85.grid.yᵃᶜᵃ[field_NN_data_85.indices[2][1]] / 1000) km") +axfield_95 = CairoMakie.Axis(fig[5, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_95.grid.yᵃᶜᵃ[field_NN_data_95.indices[2][1]] / 1000) km") + +axflux_5 = CairoMakie.Axis(fig[1, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_5.grid.yᵃᶜᵃ[flux_NN_data_5.indices[2][1]] / 1000) km") +axflux_15 = CairoMakie.Axis(fig[1, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_15.grid.yᵃᶜᵃ[flux_NN_data_15.indices[2][1]] / 1000) km") +axflux_25 = CairoMakie.Axis(fig[2, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_25.grid.yᵃᶜᵃ[flux_NN_data_25.indices[2][1]] / 1000) km") +axflux_35 = CairoMakie.Axis(fig[2, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_35.grid.yᵃᶜᵃ[flux_NN_data_35.indices[2][1]] / 1000) km") +axflux_45 = CairoMakie.Axis(fig[3, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_45.grid.yᵃᶜᵃ[flux_NN_data_45.indices[2][1]] / 1000) km") +axflux_55 = CairoMakie.Axis(fig[3, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_55.grid.yᵃᶜᵃ[flux_NN_data_55.indices[2][1]] / 1000) km") +axflux_65 = CairoMakie.Axis(fig[4, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_65.grid.yᵃᶜᵃ[flux_NN_data_65.indices[2][1]] / 1000) km") +axflux_75 = CairoMakie.Axis(fig[4, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_75.grid.yᵃᶜᵃ[flux_NN_data_75.indices[2][1]] / 1000) km") +axflux_85 = CairoMakie.Axis(fig[5, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_85.grid.yᵃᶜᵃ[flux_NN_data_85.indices[2][1]] / 1000) km") +axflux_95 = CairoMakie.Axis(fig[5, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_95.grid.yᵃᶜᵃ[flux_NN_data_95.indices[2][1]] / 1000) km") + +axfields = [axfield_5, axfield_15, axfield_25, axfield_35, axfield_45, axfield_55, axfield_65, axfield_75, axfield_85, axfield_95] +axfluxes = [axflux_5, axflux_15, axflux_25, axflux_35, axflux_45, axflux_55, axflux_65, axflux_75, axflux_85, axflux_95] + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lims = [(find_min(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices))) for field_data in field_datas] + +flux_lims = [(find_min(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices))) for flux_data in flux_datas] + +flux_lims = [(-maximum(abs, flux_lim), maximum(abs, flux_lim)) for flux_lim in flux_lims] + +NNₙs = [@lift interior(field_data[$n], :, 1, zC_indices) for field_data in field_datas] +fluxₙs = [@lift interior(flux_data[$n], :, 1, zF_indices) for flux_data in flux_datas] + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_5_surface = heatmap!(axfield_5, xC, zC[zC_indices], NN_5ₙ, colormap=colorscheme_field, colorrange=field_lims[1]) +field_15_surface = heatmap!(axfield_15, xC, zC[zC_indices], NN_15ₙ, colormap=colorscheme_field, colorrange=field_lims[2]) +field_25_surface = heatmap!(axfield_25, xC, zC[zC_indices], NN_25ₙ, colormap=colorscheme_field, colorrange=field_lims[3]) +field_35_surface = heatmap!(axfield_35, xC, zC[zC_indices], NN_35ₙ, colormap=colorscheme_field, colorrange=field_lims[4]) +field_45_surface = heatmap!(axfield_45, xC, zC[zC_indices], NN_45ₙ, colormap=colorscheme_field, colorrange=field_lims[5]) +field_55_surface = heatmap!(axfield_55, xC, zC[zC_indices], NN_55ₙ, colormap=colorscheme_field, colorrange=field_lims[6]) +field_65_surface = heatmap!(axfield_65, xC, zC[zC_indices], NN_65ₙ, colormap=colorscheme_field, colorrange=field_lims[7]) +field_75_surface = heatmap!(axfield_75, xC, zC[zC_indices], NN_75ₙ, colormap=colorscheme_field, colorrange=field_lims[8]) +field_85_surface = heatmap!(axfield_85, xC, zC[zC_indices], NN_85ₙ, colormap=colorscheme_field, colorrange=field_lims[9]) +field_95_surface = heatmap!(axfield_95, xC, zC[zC_indices], NN_95ₙ, colormap=colorscheme_field, colorrange=field_lims[10]) + +flux_5_surface = heatmap!(axflux_5, xC, zC[zF_indices], flux_5ₙ, colormap=colorscheme_flux, colorrange=flux_lims[1]) +flux_15_surface = heatmap!(axflux_15, xC, zC[zF_indices], flux_15ₙ, colormap=colorscheme_flux, colorrange=flux_lims[2]) +flux_25_surface = heatmap!(axflux_25, xC, zC[zF_indices], flux_25ₙ, colormap=colorscheme_flux, colorrange=flux_lims[3]) +flux_35_surface = heatmap!(axflux_35, xC, zC[zF_indices], flux_35ₙ, colormap=colorscheme_flux, colorrange=flux_lims[4]) +flux_45_surface = heatmap!(axflux_45, xC, zC[zF_indices], flux_45ₙ, colormap=colorscheme_flux, colorrange=flux_lims[5]) +flux_55_surface = heatmap!(axflux_55, xC, zC[zF_indices], flux_55ₙ, colormap=colorscheme_flux, colorrange=flux_lims[6]) +flux_65_surface = heatmap!(axflux_65, xC, zC[zF_indices], flux_65ₙ, colormap=colorscheme_flux, colorrange=flux_lims[7]) +flux_75_surface = heatmap!(axflux_75, xC, zC[zF_indices], flux_75ₙ, colormap=colorscheme_flux, colorrange=flux_lims[8]) +flux_85_surface = heatmap!(axflux_85, xC, zC[zF_indices], flux_85ₙ, colormap=colorscheme_flux, colorrange=flux_lims[9]) +flux_95_surface = heatmap!(axflux_95, xC, zC[zF_indices], flux_95ₙ, colormap=colorscheme_flux, colorrange=flux_lims[10]) + +Colorbar(fig[1, 2], field_5_surface, label="Field") +Colorbar(fig[1, 4], flux_5_surface, label="NN Flux") +Colorbar(fig[1, 6], field_15_surface, label="Field") +Colorbar(fig[1, 8], flux_15_surface, label="NN Flux") +Colorbar(fig[2, 2], field_25_surface, label="Field") +Colorbar(fig[2, 4], flux_25_surface, label="NN Flux") +Colorbar(fig[2, 6], field_35_surface, label="Field") +Colorbar(fig[2, 8], flux_35_surface, label="NN Flux") +Colorbar(fig[3, 2], field_45_surface, label="Field") +Colorbar(fig[3, 4], flux_45_surface, label="NN Flux") +Colorbar(fig[3, 6], field_55_surface, label="Field") +Colorbar(fig[3, 8], flux_55_surface, label="NN Flux") +Colorbar(fig[4, 2], field_65_surface, label="Field") +Colorbar(fig[4, 4], flux_65_surface, label="NN Flux") +Colorbar(fig[4, 6], field_75_surface, label="Field") +Colorbar(fig[4, 8], flux_75_surface, label="NN Flux") +Colorbar(fig[5, 2], field_85_surface, label="Field") +Colorbar(fig[5, 4], flux_85_surface, label="NN Flux") +Colorbar(fig[5, 6], field_95_surface, label="Field") +Colorbar(fig[5, 8], flux_95_surface, label="NN Flux") + +for axfield in axfields + xlims!(axfield, minimum(xC), maximum(xC)) + ylims!(axfield, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +end + +for axflux in axfluxes + xlims!(axflux, minimum(xC), maximum(xC)) + ylims!(axflux, minimum(zC[zF_indices]), maximum(zC[zF_indices])) +end + +title_str = @lift "Salinity (psu), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_xzslices_fluxes_S.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end +#%% \ No newline at end of file diff --git a/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl b/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl new file mode 100644 index 0000000000..1f107c5cb9 --- /dev/null +++ b/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl @@ -0,0 +1,1275 @@ +#using Pkg +# pkg"add Oceananigans CairoMakie" +using Oceananigans +include("NN_closure_global_Ri_nof_BBLRifirstzone510.jl") +include("xin_kai_vertical_diffusivity_local_2step_new.jl") + +ENV["GKSwstype"] = "100" + +pushfirst!(LOAD_PATH, @__DIR__) + +using Printf +using Statistics +using CairoMakie + +using Oceananigans +using Oceananigans.Units +using Oceananigans.OutputReaders: FieldTimeSeries +using Oceananigans.Grids: xnode, ynode, znode +using SeawaterPolynomials +using SeawaterPolynomials:TEOS10 +using ColorSchemes + +#%% +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_NN_closure_NDE5_Ri_BBLRifirztzone510_temp" +FILE_DIR = "./Output/$(filename)" +# FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +mkpath(FILE_DIR) + +# Architecture +model_architecture = GPU() + +nn_closure = NNFluxClosure(model_architecture) +base_closure = XinKaiLocalVerticalDiffusivity() +closure = (base_closure, nn_closure) + +advection_scheme = FluxFormAdvection(WENO(order=5), WENO(order=5), WENO(order=5)) + +# number of grid points +const Nx = 100 +const Ny = 100 +const Nz = 200 + +const Δz = 8meters +const Lx = 4000kilometers +const Ly = 6000kilometers +const Lz = Nz * Δz + +grid = RectilinearGrid(model_architecture, Float64, + topology = (Bounded, Bounded, Bounded), + size = (Nx, Ny, Nz), + halo = (4, 4, 4), + x = (-Lx/2, Lx/2), + y = (-Ly/2, Ly/2), + z = (-Lz, 0)) + +@info "Built a grid: $grid." + +##### +##### Boundary conditions +##### +const T_north = 0 +const T_south = 30 +const T_mid = (T_north + T_south) / 2 +const ΔT = T_south - T_north + +const S_north = 34 +const S_south = 37 +const S_mid = (S_north + S_south) / 2 + +const τ₀ = 1e-4 + +const μ_drag = 1/30days +const μ_T = 1/30days + +const X₀ = -Lx/2 + 800kilometers +const Y₀ = -Ly/2 + 1500kilometers +const R₀ = 700kilometers +const Qᵀ_mode = 4.5e-4 +const σ_mode = 20kilometers + +##### +##### Forcing and initial condition +##### +# @inline T_initial(x, y, z) = T_north + ΔT / 2 * (1 + z / Lz) +# @inline T_initial(x, y, z) = (T_north + T_south / 2) + 5 * (1 + z / Lz) +@inline T_initial(x, y, z) = 10 + 20 * (1 + z / Lz) + +@inline surface_u_flux(x, y, t) = -τ₀ * cos(2π * y / Ly) + +surface_u_flux_bc = FluxBoundaryCondition(surface_u_flux) + +@inline u_drag(x, y, t, u) = @inbounds -μ_drag * Lz * u +@inline v_drag(x, y, t, v) = @inbounds -μ_drag * Lz * v + +u_drag_bc = FluxBoundaryCondition(u_drag; field_dependencies=:u) +v_drag_bc = FluxBoundaryCondition(v_drag; field_dependencies=:v) + +u_bcs = FieldBoundaryConditions( top = surface_u_flux_bc, + bottom = u_drag_bc, + north = ValueBoundaryCondition(0), + south = ValueBoundaryCondition(0)) + +v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), + bottom = v_drag_bc, + east = ValueBoundaryCondition(0), + west = ValueBoundaryCondition(0)) + +@inline T_ref(y) = T_mid - ΔT / Ly * y + +@inline Qᵀ_winter(t) = max(0, -Qᵀ_mode * sin(2π * t / 360days)) +@inline Qᵀ_subpolar(x, y, t) = ifelse((x - X₀)^2 + (y - Y₀)^2 <= R₀^2, Qᵀ_winter(t), + exp(-(sqrt((x - X₀)^2 + (y - Y₀)^2) - R₀)^2 / (2 * σ_mode^2)) * Qᵀ_winter(t)) + +@inline surface_T_flux(x, y, t, T) = μ_T * Δz * (T - T_ref(y)) + Qᵀ_subpolar(x, y, t) +surface_T_flux_bc = FluxBoundaryCondition(surface_T_flux; field_dependencies=:T) +T_bcs = FieldBoundaryConditions(top = surface_T_flux_bc) + +@inline S_ref(y) = (S_north - S_south) / Ly * y + S_mid +@inline S_initial(x, y, z) = S_ref(y) +@inline surface_S_flux(x, y, t, S) = μ_T * Δz * (S - S_ref(y)) +surface_S_flux_bc = FluxBoundaryCondition(surface_S_flux; field_dependencies=:S) +S_bcs = FieldBoundaryConditions(top = surface_S_flux_bc) + +##### +##### Coriolis +##### +coriolis = BetaPlane(rotation_rate=7.292115e-5, latitude=45, radius=6371e3) + +##### +##### Model building +##### + +@info "Building a model..." + +# This is a weird bug. If a model is not initialized with a closure other than XinKaiVerticalDiffusivity, +# the code will throw a CUDA: illegal memory access error for models larger than a small size. +# This is a workaround to initialize the model with a closure other than XinKaiVerticalDiffusivity first, +# then the code will run without any issues. +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = VerticalScalarDiffusivity(ν=1e-5, κ=1e-5), + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +model = HydrostaticFreeSurfaceModel( + grid = grid, + free_surface = SplitExplicitFreeSurface(grid, cfl=0.75), + momentum_advection = advection_scheme, + tracer_advection = advection_scheme, + buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()), + coriolis = coriolis, + closure = closure, + tracers = (:T, :S), + boundary_conditions = (; u = u_bcs, v = v_bcs, T = T_bcs, S = S_bcs), +) + +@info "Built $model." + +##### +##### Initial conditions +##### + +# resting initial condition +noise(z) = rand() * exp(z / 8) + +T_initial_noisy(x, y, z) = T_initial(x, y, z) + 1e-6 * noise(z) +S_initial_noisy(x, y, z) = S_initial(x, y, z) + 1e-6 * noise(z) + +set!(model, T=T_initial_noisy, S=S_initial_noisy) +using Oceananigans.TimeSteppers: update_state! +update_state!(model) +##### +##### Simulation building +##### +Δt₀ = 5minutes +stop_time = 10800days + +simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) + +# add timestep wizard callback +# wizard = TimeStepWizard(cfl=0.25, max_change=1.05, max_Δt=12minutes) +# simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + +# add progress callback +wall_clock = [time_ns()] + +function print_progress(sim) + @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): %6.3e, max(v): %6.3e, max(T): %6.3e, max(S): %6.3e, next Δt: %s\n", + 100 * (sim.model.clock.time / sim.stop_time), + sim.model.clock.iteration, + prettytime(sim.model.clock.time), + prettytime(1e-9 * (time_ns() - wall_clock[1])), + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + maximum(abs, sim.model.tracers.T), + maximum(abs, sim.model.tracers.S), + prettytime(sim.Δt)) + + wall_clock[1] = time_ns() + + return nothing +end + +simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(100)) + +##### +##### Diagnostics +##### +u, v, w = model.velocities +T, S = model.tracers.T, model.tracers.S +U_bt = Field(Integral(u, dims=3)); +Ψ = Field(CumulativeIntegral(-U_bt, dims=2)); +first_index = model.diffusivity_fields[2].first_index +last_index = model.diffusivity_fields[2].last_index +wT_NN = model.diffusivity_fields[2].wT +wS_NN = model.diffusivity_fields[2].wS + +κ = model.diffusivity_fields[1].κᶜ +wT_base = κ * ∂z(T) +wS_base = κ * ∂z(S) + +wT = wT_NN + wT_base +wS = wS_NN + wS_base + +@inline function get_N²(i, j, k, grid, b, C) + return ∂z_b(i, j, k, grid, b, C) +end + +N²_op = KernelFunctionOperation{Center, Center, Face}(get_N², model.grid, model.buoyancy.model, model.tracers) +N² = Field(N²_op) + +@inline function get_density(i, j, k, grid, b, C) + T, S = Oceananigans.BuoyancyModels.get_temperature_and_salinity(b, C) + @inbounds ρ = TEOS10.ρ(T[i, j, k], S[i, j, k], 0, b.model.equation_of_state) + return ρ +end + +ρ_op = KernelFunctionOperation{Center, Center, Center}(get_density, model.grid, model.buoyancy, model.tracers) +ρ = Field(ρ_op) + +@inline function get_top_buoyancy_flux(i, j, k, grid, buoyancy, T_bc, S_bc, velocities, tracers, clock) + return top_buoyancy_flux(i, j, grid, buoyancy, (; T=T_bc, S=S_bc), clock, merge(velocities, tracers)) +end + +Qb = KernelFunctionOperation{Center, Center, Nothing}(get_top_buoyancy_flux, model.grid, model.buoyancy, T.boundary_conditions.top, S.boundary_conditions.top, model.velocities, model.tracers, model.clock) +Qb = Field(Qb) + +ubar_zonal = Average(u, dims=1) +vbar_zonal = Average(v, dims=1) +wbar_zonal = Average(w, dims=1) +Tbar_zonal = Average(T, dims=1) +Sbar_zonal = Average(S, dims=1) +ρbar_zonal = Average(ρ, dims=1) + +wT_NNbar_zonal = Average(wT_NN, dims=1) +wS_NNbar_zonal = Average(wS_NN, dims=1) + +wT_basebar_zonal = Average(wT_base, dims=1) +wS_basebar_zonal = Average(wS_base, dims=1) + +wTbar_zonal = Average(wT, dims=1) +wSbar_zonal = Average(wS, dims=1) + +outputs = (; u, v, w, T, S, ρ, N², wT_NN, wS_NN, wT_base, wS_base, wT, wS) +zonal_outputs = (; ubar_zonal, vbar_zonal, wbar_zonal, Tbar_zonal, Sbar_zonal, ρbar_zonal, wT_NNbar_zonal, wS_NNbar_zonal, wT_basebar_zonal, wS_basebar_zonal, wTbar_zonal, wSbar_zonal) + +##### +##### Build checkpointer and output writer +##### +simulation.output_writers[:xy] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xy", + indices = (:, :, Nz), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz", + indices = (1, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz", + indices = (:, 1, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_10] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_10", + indices = (10, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_20] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_20", + indices = (20, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_30] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_30", + indices = (30, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_40] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_40", + indices = (40, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_50] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_50", + indices = (50, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_60] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_60", + indices = (60, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_70] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_70", + indices = (70, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_80] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_80", + indices = (80, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:yz_90] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_yz_90", + indices = (90, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_5] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_5", + indices = (:, 5, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_15] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_15", + indices = (:, 15, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_25] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_25", + indices = (:, 25, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_35] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_35", + indices = (:, 35, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_45] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_45", + indices = (:, 45, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_55] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_55", + indices = (:, 55, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_65] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_65", + indices = (:, 65, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_75] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_75", + indices = (:, 75, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_85] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_85", + indices = (:, 85, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:xz_95] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields_xz_95", + indices = (:, 95, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_outputs, + filename = "$(FILE_DIR)/averaged_fields_zonal", + schedule = TimeInterval(10days)) + +simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; first_index, last_index, Qb), + filename = "$(FILE_DIR)/instantaneous_fields_NN_active_diagnostics", + indices = (:, :, :), + schedule = TimeInterval(10days)) + +simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), + filename = "$(FILE_DIR)/averaged_fields_streamfunction", + schedule = AveragedTimeInterval(1825days, window=1825days)) + +simulation.output_writers[:checkpointer] = Checkpointer(model, + schedule = TimeInterval(730days), + prefix = "$(FILE_DIR)/checkpointer") + +@info "Running the simulation..." + +try + files = readdir(FILE_DIR) + checkpoint_files = files[occursin.("checkpointer_iteration", files)] + if !isempty(checkpoint_files) + checkpoint_iters = parse.(Int, [filename[findfirst("iteration", filename)[end]+1:findfirst(".jld2", filename)[1]-1] for filename in checkpoint_files]) + pickup_iter = maximum(checkpoint_iters) + run!(simulation, pickup="$(FILE_DIR)/checkpointer_iteration$(pickup_iter).jld2") + else + run!(simulation) + end +catch err + @info "run! threw an error! The error message is" + showerror(stdout, err) +end + +#%% +T_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "T") +T_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "T") +T_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "T") + +S_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "S") +S_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "S") +S_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "S") + +u_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "u") +u_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "u") +u_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "u") + +v_xy_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xy.jld2", "v") +v_xz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz.jld2", "v") +v_yz_data = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", "v") + +times = T_xy_data.times ./ 24 ./ 60^2 +Nt = length(times) +timeframes = 1:Nt + +# Nx, Ny, Nz = T_xy_data.grid.Nx, T_xy_data.grid.Ny, T_xy_data.grid.Nz +xC, yC, zC = T_xy_data.grid.xᶜᵃᵃ[1:Nx], T_xy_data.grid.yᵃᶜᵃ[1:Ny], T_xy_data.grid.zᵃᵃᶜ[1:Nz] +zF = T_xy_data.grid.zᵃᵃᶠ[1:Nz+1] + +# Lx, Ly, Lz = T_xy_data.grid.Lx, T_xy_data.grid.Ly, T_xy_data.grid.Lz + +xCs_xy = xC +yCs_xy = yC +zCs_xy = [zC[Nz] for x in xCs_xy, y in yCs_xy] + +yCs_yz = yC +xCs_yz = range(xC[1], stop=xC[1], length=length(zC)) +zCs_yz = zeros(length(xCs_yz), length(yCs_yz)) +for j in axes(zCs_yz, 2) + zCs_yz[:, j] .= zC +end + +xCs_xz = xC +yCs_xz = range(yC[1], stop=yC[1], length=length(zC)) +zCs_xz = zeros(length(xCs_xz), length(yCs_xz)) +for i in axes(zCs_xz, 1) + zCs_xz[i, :] .= zC +end + +xFs_xy = xC +yFs_xy = yC +zFs_xy = [zF[Nz+1] for x in xFs_xy, y in yFs_xy] + +yFs_yz = yC +xFs_yz = range(xC[1], stop=xC[1], length=length(zF)) +zFs_yz = zeros(length(xFs_yz), length(yFs_yz)) +for j in axes(zFs_yz, 2) + zFs_yz[:, j] .= zF +end + +xFs_xz = xC +yFs_xz = range(yC[1], stop=yC[1], length=length(zF)) +zFs_xz = zeros(length(xFs_xz), length(yFs_xz)) +for i in axes(zFs_xz, 1) + zFs_xz[i, :] .= zF +end + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +# for freeconvection +# startheight = 64 + +# for wind mixing +startheight = 1 +Tlim = (find_min(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(T_xy_data, :, :, 1, timeframes), interior(T_yz_data, 1, :, startheight:Nz, timeframes), interior(T_xz_data, :, 1, startheight:Nz, timeframes))) +Slim = (find_min(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(S_xy_data, :, :, 1, timeframes), interior(S_yz_data, 1, :, startheight:Nz, timeframes), interior(S_xz_data, :, 1, startheight:Nz, timeframes))) +ulim = (-find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(u_xy_data, :, :, 1, timeframes), interior(u_yz_data, 1, :, startheight:Nz, timeframes), interior(u_xz_data, :, 1, startheight:Nz, timeframes))) +vlim = (-find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes)), + find_max(interior(v_xy_data, :, :, 1, timeframes), interior(v_yz_data, 1, :, startheight:Nz, timeframes), interior(v_xz_data, :, 1, startheight:Nz, timeframes))) + +colorscheme = colorschemes[:balance] +T_colormap = colorscheme +S_colormap = colorscheme +u_colormap = colorscheme +v_colormap = colorscheme + +T_color_range = Tlim +S_color_range = Slim +u_color_range = ulim +v_color_range = vlim +#%% +plot_aspect = (2, 3, 0.5) +fig = Figure(size=(1500, 700)) +axT = CairoMakie.Axis3(fig[1, 1], title="Temperature (°C)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axS = CairoMakie.Axis3(fig[1, 3], title="Salinity (g kg⁻¹)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axu = CairoMakie.Axis3(fig[2, 1], title="u (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) +axv = CairoMakie.Axis3(fig[2, 3], title="v (m/s)", xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", viewmode=:fitzoom, aspect=plot_aspect) + +n = Observable(1) + +T_xy = @lift interior(T_xy_data[$n], :, :, 1) +T_yz = @lift transpose(interior(T_yz_data[$n], 1, :, :)) +T_xz = @lift interior(T_xz_data[$n], :, 1, :) + +S_xy = @lift interior(S_xy_data[$n], :, :, 1) +S_yz = @lift transpose(interior(S_yz_data[$n], 1, :, :)) +S_xz = @lift interior(S_xz_data[$n], :, 1, :) + +u_xy = @lift interior(u_xy_data[$n], :, :, 1) +u_yz = @lift transpose(interior(u_yz_data[$n], 1, :, :)) +u_xz = @lift interior(u_xz_data[$n], :, 1, :) + +v_xy = @lift interior(v_xy_data[$n], :, :, 1) +v_yz = @lift transpose(interior(v_yz_data[$n], 1, :, :)) +v_xz = @lift interior(v_xz_data[$n], :, 1, :) + +# time_str = @lift "Surface Cooling, Time = $(round(times[$n], digits=2)) hours" +time_str = @lift "Surface Wind Stress, Time = $(round(times[$n], digits=2)) days" +Label(fig[0, :], text=time_str, tellwidth=false, font=:bold) + +T_xy_surface = surface!(axT, xCs_xy, yCs_xy, zCs_xy, color=T_xy, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_yz_surface = surface!(axT, xCs_yz, yCs_yz, zCs_yz, color=T_yz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) +T_xz_surface = surface!(axT, xCs_xz, yCs_xz, zCs_xz, color=T_xz, colormap=T_colormap, colorrange = T_color_range, lowclip=T_colormap[1]) + +S_xy_surface = surface!(axS, xCs_xy, yCs_xy, zCs_xy, color=S_xy, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_yz_surface = surface!(axS, xCs_yz, yCs_yz, zCs_yz, color=S_yz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) +S_xz_surface = surface!(axS, xCs_xz, yCs_xz, zCs_xz, color=S_xz, colormap=S_colormap, colorrange = S_color_range, lowclip=S_colormap[1]) + +u_xy_surface = surface!(axu, xCs_xy, yCs_xy, zCs_xy, color=u_xy, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_yz_surface = surface!(axu, xCs_yz, yCs_yz, zCs_yz, color=u_yz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) +u_xz_surface = surface!(axu, xCs_xz, yCs_xz, zCs_xz, color=u_xz, colormap=u_colormap, colorrange = u_color_range, lowclip=u_colormap[1], highclip=u_colormap[end]) + +v_xy_surface = surface!(axv, xCs_xy, yCs_xy, zCs_xy, color=v_xy, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_yz_surface = surface!(axv, xCs_yz, yCs_yz, zCs_yz, color=v_yz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) +v_xz_surface = surface!(axv, xCs_xz, yCs_xz, zCs_xz, color=v_xz, colormap=v_colormap, colorrange = v_color_range, lowclip=v_colormap[1], highclip=v_colormap[end]) + +Colorbar(fig[1,2], T_xy_surface) +Colorbar(fig[1,4], S_xy_surface) +Colorbar(fig[2,2], u_xy_surface) +Colorbar(fig[2,4], v_xy_surface) + +xlims!(axT, (-Lx/2, Lx/2)) +xlims!(axS, (-Lx/2, Lx/2)) +xlims!(axu, (-Lx/2, Lx/2)) +xlims!(axv, (-Lx/2, Lx/2)) + +ylims!(axT, (-Ly/2, Ly/2)) +ylims!(axS, (-Ly/2, Ly/2)) +ylims!(axu, (-Ly/2, Ly/2)) +ylims!(axv, (-Ly/2, Ly/2)) + +zlims!(axT, (-Lz, 0)) +zlims!(axS, (-Lz, 0)) +zlims!(axu, (-Lz, 0)) +zlims!(axv, (-Lz, 0)) + +@info "Recording 3D fields" +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_3D_instantaneous_fields.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end + +#%% +@info "Recording T fields and fluxes in yz" + +fieldname = "T" +fluxname = "wT_NN" +field_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fieldname, backend=OnDisk()) +field_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fieldname, backend=OnDisk()) +field_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fieldname, backend=OnDisk()) +field_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fieldname, backend=OnDisk()) +field_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fieldname, backend=OnDisk()) +field_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fieldname, backend=OnDisk()) +field_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fieldname, backend=OnDisk()) +field_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fieldname, backend=OnDisk()) +field_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fieldname, backend=OnDisk()) +field_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fluxname, backend=OnDisk()) +flux_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fluxname, backend=OnDisk()) +flux_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fluxname, backend=OnDisk()) +flux_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fluxname, backend=OnDisk()) +flux_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fluxname, backend=OnDisk()) +flux_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fluxname, backend=OnDisk()) +flux_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fluxname, backend=OnDisk()) +flux_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fluxname, backend=OnDisk()) +flux_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fluxname, backend=OnDisk()) +flux_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fluxname, backend=OnDisk()) + +xC = field_NN_data_00.grid.xᶜᵃᵃ[1:field_NN_data_00.grid.Nx] +yC = field_NN_data_00.grid.yᵃᶜᵃ[1:field_NN_data_00.grid.Ny] +zC = field_NN_data_00.grid.zᵃᵃᶜ[1:field_NN_data_00.grid.Nz] +zF = field_NN_data_00.grid.zᵃᵃᶠ[1:field_NN_data_00.grid.Nz+1] + +Nt = length(field_NN_data_90) +times = field_NN_data_00.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_00 = CairoMakie.Axis(fig[1, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_00.grid.xᶜᵃᵃ[field_NN_data_00.indices[1][1]] / 1000) km") +axfield_10 = CairoMakie.Axis(fig[1, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_10.grid.xᶜᵃᵃ[field_NN_data_10.indices[1][1]] / 1000) km") +axfield_20 = CairoMakie.Axis(fig[2, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_20.grid.xᶜᵃᵃ[field_NN_data_20.indices[1][1]] / 1000) km") +axfield_30 = CairoMakie.Axis(fig[2, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_30.grid.xᶜᵃᵃ[field_NN_data_30.indices[1][1]] / 1000) km") +axfield_40 = CairoMakie.Axis(fig[3, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_40.grid.xᶜᵃᵃ[field_NN_data_40.indices[1][1]] / 1000) km") +axfield_50 = CairoMakie.Axis(fig[3, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_50.grid.xᶜᵃᵃ[field_NN_data_50.indices[1][1]] / 1000) km") +axfield_60 = CairoMakie.Axis(fig[4, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_60.grid.xᶜᵃᵃ[field_NN_data_60.indices[1][1]] / 1000) km") +axfield_70 = CairoMakie.Axis(fig[4, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_70.grid.xᶜᵃᵃ[field_NN_data_70.indices[1][1]] / 1000) km") +axfield_80 = CairoMakie.Axis(fig[5, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_80.grid.xᶜᵃᵃ[field_NN_data_80.indices[1][1]] / 1000) km") +axfield_90 = CairoMakie.Axis(fig[5, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_90.grid.xᶜᵃᵃ[field_NN_data_90.indices[1][1]] / 1000) km") + +axflux_00 = CairoMakie.Axis(fig[1, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_00.grid.xᶜᵃᵃ[flux_NN_data_00.indices[1][1]] / 1000) km") +axflux_10 = CairoMakie.Axis(fig[1, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_10.grid.xᶜᵃᵃ[flux_NN_data_10.indices[1][1]] / 1000) km") +axflux_20 = CairoMakie.Axis(fig[2, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_20.grid.xᶜᵃᵃ[flux_NN_data_20.indices[1][1]] / 1000) km") +axflux_30 = CairoMakie.Axis(fig[2, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_30.grid.xᶜᵃᵃ[flux_NN_data_30.indices[1][1]] / 1000) km") +axflux_40 = CairoMakie.Axis(fig[3, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_40.grid.xᶜᵃᵃ[flux_NN_data_40.indices[1][1]] / 1000) km") +axflux_50 = CairoMakie.Axis(fig[3, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_50.grid.xᶜᵃᵃ[flux_NN_data_50.indices[1][1]] / 1000) km") +axflux_60 = CairoMakie.Axis(fig[4, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_60.grid.xᶜᵃᵃ[flux_NN_data_60.indices[1][1]] / 1000) km") +axflux_70 = CairoMakie.Axis(fig[4, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_70.grid.xᶜᵃᵃ[flux_NN_data_70.indices[1][1]] / 1000) km") +axflux_80 = CairoMakie.Axis(fig[5, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_80.grid.xᶜᵃᵃ[flux_NN_data_80.indices[1][1]] / 1000) km") +axflux_90 = CairoMakie.Axis(fig[5, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_90.grid.xᶜᵃᵃ[flux_NN_data_90.indices[1][1]] / 1000) km") + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lim = (find_min(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (find_min(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (-maximum(abs, flux_lim), maximum(abs, flux_lim)) + +NN_00ₙ = @lift interior(field_NN_data_00[$n], 1, :, zC_indices) +NN_10ₙ = @lift interior(field_NN_data_10[$n], 1, :, zC_indices) +NN_20ₙ = @lift interior(field_NN_data_20[$n], 1, :, zC_indices) +NN_30ₙ = @lift interior(field_NN_data_30[$n], 1, :, zC_indices) +NN_40ₙ = @lift interior(field_NN_data_40[$n], 1, :, zC_indices) +NN_50ₙ = @lift interior(field_NN_data_50[$n], 1, :, zC_indices) +NN_60ₙ = @lift interior(field_NN_data_60[$n], 1, :, zC_indices) +NN_70ₙ = @lift interior(field_NN_data_70[$n], 1, :, zC_indices) +NN_80ₙ = @lift interior(field_NN_data_80[$n], 1, :, zC_indices) +NN_90ₙ = @lift interior(field_NN_data_90[$n], 1, :, zC_indices) + +flux_00ₙ = @lift interior(flux_NN_data_00[$n], 1, :, zF_indices) +flux_10ₙ = @lift interior(flux_NN_data_10[$n], 1, :, zF_indices) +flux_20ₙ = @lift interior(flux_NN_data_20[$n], 1, :, zF_indices) +flux_30ₙ = @lift interior(flux_NN_data_30[$n], 1, :, zF_indices) +flux_40ₙ = @lift interior(flux_NN_data_40[$n], 1, :, zF_indices) +flux_50ₙ = @lift interior(flux_NN_data_50[$n], 1, :, zF_indices) +flux_60ₙ = @lift interior(flux_NN_data_60[$n], 1, :, zF_indices) +flux_70ₙ = @lift interior(flux_NN_data_70[$n], 1, :, zF_indices) +flux_80ₙ = @lift interior(flux_NN_data_80[$n], 1, :, zF_indices) +flux_90ₙ = @lift interior(flux_NN_data_90[$n], 1, :, zF_indices) + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_00_surface = heatmap!(axfield_00, yC, zC[zC_indices], NN_00ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_10_surface = heatmap!(axfield_10, yC, zC[zC_indices], NN_10ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_20_surface = heatmap!(axfield_20, yC, zC[zC_indices], NN_20ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_30_surface = heatmap!(axfield_30, yC, zC[zC_indices], NN_30ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_40_surface = heatmap!(axfield_40, yC, zC[zC_indices], NN_40ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_50_surface = heatmap!(axfield_50, yC, zC[zC_indices], NN_50ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_60_surface = heatmap!(axfield_60, yC, zC[zC_indices], NN_60ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_70_surface = heatmap!(axfield_70, yC, zC[zC_indices], NN_70ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_80_surface = heatmap!(axfield_80, yC, zC[zC_indices], NN_80ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_90_surface = heatmap!(axfield_90, yC, zC[zC_indices], NN_90ₙ, colormap=colorscheme_field, colorrange=field_lim) + +flux_00_surface = heatmap!(axflux_00, yC, zC[zF_indices], flux_00ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_10_surface = heatmap!(axflux_10, yC, zC[zF_indices], flux_10ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_20_surface = heatmap!(axflux_20, yC, zC[zF_indices], flux_20ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_30_surface = heatmap!(axflux_30, yC, zC[zF_indices], flux_30ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_40_surface = heatmap!(axflux_40, yC, zC[zF_indices], flux_40ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_50_surface = heatmap!(axflux_50, yC, zC[zF_indices], flux_50ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_60_surface = heatmap!(axflux_60, yC, zC[zF_indices], flux_60ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_70_surface = heatmap!(axflux_70, yC, zC[zF_indices], flux_70ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_80_surface = heatmap!(axflux_80, yC, zC[zF_indices], flux_80ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_90_surface = heatmap!(axflux_90, yC, zC[zF_indices], flux_90ₙ, colormap=colorscheme_flux, colorrange=flux_lim) + +Colorbar(fig[1:5, 2], field_00_surface, label="Field") +Colorbar(fig[1:5, 4], flux_00_surface, label="NN Flux") +Colorbar(fig[1:5, 6], field_00_surface, label="Field") +Colorbar(fig[1:5, 8], flux_00_surface, label="NN Flux") + +xlims!(axfield_00, minimum(yC), maximum(yC)) +xlims!(axfield_10, minimum(yC), maximum(yC)) +xlims!(axfield_20, minimum(yC), maximum(yC)) +xlims!(axfield_30, minimum(yC), maximum(yC)) +xlims!(axfield_40, minimum(yC), maximum(yC)) +xlims!(axfield_50, minimum(yC), maximum(yC)) +xlims!(axfield_60, minimum(yC), maximum(yC)) +xlims!(axfield_70, minimum(yC), maximum(yC)) +xlims!(axfield_80, minimum(yC), maximum(yC)) +xlims!(axfield_90, minimum(yC), maximum(yC)) + +ylims!(axfield_00, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_10, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_20, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_30, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_40, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_50, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_60, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_70, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_80, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_90, minimum(zC[zC_indices]), maximum(zC[zC_indices])) + +xlims!(axflux_00, minimum(yC), maximum(yC)) +xlims!(axflux_10, minimum(yC), maximum(yC)) +xlims!(axflux_20, minimum(yC), maximum(yC)) +xlims!(axflux_30, minimum(yC), maximum(yC)) +xlims!(axflux_40, minimum(yC), maximum(yC)) +xlims!(axflux_50, minimum(yC), maximum(yC)) +xlims!(axflux_60, minimum(yC), maximum(yC)) +xlims!(axflux_70, minimum(yC), maximum(yC)) +xlims!(axflux_80, minimum(yC), maximum(yC)) +xlims!(axflux_90, minimum(yC), maximum(yC)) + +ylims!(axflux_00, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_10, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_20, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_30, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_40, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_50, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_60, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_70, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_80, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_90, minimum(zF[zF_indices]), maximum(zF[zF_indices])) + +title_str = @lift "Temperature (°C), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +# save("./Output/compare_3D_instantaneous_fields_slices_NNclosure_fluxes.png", fig) +# display(fig) +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_yzslices_fluxes_T.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end +#%% +@info "Recording S fields and fluxes in yz" + +fieldname = "S" +fluxname = "wS_NN" +field_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fieldname, backend=OnDisk()) +field_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fieldname, backend=OnDisk()) +field_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fieldname, backend=OnDisk()) +field_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fieldname, backend=OnDisk()) +field_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fieldname, backend=OnDisk()) +field_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fieldname, backend=OnDisk()) +field_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fieldname, backend=OnDisk()) +field_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fieldname, backend=OnDisk()) +field_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fieldname, backend=OnDisk()) +field_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_00 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz.jld2", fluxname, backend=OnDisk()) +flux_NN_data_10 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_10.jld2", fluxname, backend=OnDisk()) +flux_NN_data_20 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_20.jld2", fluxname, backend=OnDisk()) +flux_NN_data_30 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_30.jld2", fluxname, backend=OnDisk()) +flux_NN_data_40 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_40.jld2", fluxname, backend=OnDisk()) +flux_NN_data_50 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_50.jld2", fluxname, backend=OnDisk()) +flux_NN_data_60 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_60.jld2", fluxname, backend=OnDisk()) +flux_NN_data_70 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_70.jld2", fluxname, backend=OnDisk()) +flux_NN_data_80 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_80.jld2", fluxname, backend=OnDisk()) +flux_NN_data_90 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_yz_90.jld2", fluxname, backend=OnDisk()) + +xC = field_NN_data_00.grid.xᶜᵃᵃ[1:field_NN_data_00.grid.Nx] +yC = field_NN_data_00.grid.yᵃᶜᵃ[1:field_NN_data_00.grid.Ny] +zC = field_NN_data_00.grid.zᵃᵃᶜ[1:field_NN_data_00.grid.Nz] +zF = field_NN_data_00.grid.zᵃᵃᶠ[1:field_NN_data_00.grid.Nz+1] + +Nt = length(field_NN_data_90) +times = field_NN_data_00.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_00 = CairoMakie.Axis(fig[1, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_00.grid.xᶜᵃᵃ[field_NN_data_00.indices[1][1]] / 1000) km") +axfield_10 = CairoMakie.Axis(fig[1, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_10.grid.xᶜᵃᵃ[field_NN_data_10.indices[1][1]] / 1000) km") +axfield_20 = CairoMakie.Axis(fig[2, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_20.grid.xᶜᵃᵃ[field_NN_data_20.indices[1][1]] / 1000) km") +axfield_30 = CairoMakie.Axis(fig[2, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_30.grid.xᶜᵃᵃ[field_NN_data_30.indices[1][1]] / 1000) km") +axfield_40 = CairoMakie.Axis(fig[3, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_40.grid.xᶜᵃᵃ[field_NN_data_40.indices[1][1]] / 1000) km") +axfield_50 = CairoMakie.Axis(fig[3, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_50.grid.xᶜᵃᵃ[field_NN_data_50.indices[1][1]] / 1000) km") +axfield_60 = CairoMakie.Axis(fig[4, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_60.grid.xᶜᵃᵃ[field_NN_data_60.indices[1][1]] / 1000) km") +axfield_70 = CairoMakie.Axis(fig[4, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_70.grid.xᶜᵃᵃ[field_NN_data_70.indices[1][1]] / 1000) km") +axfield_80 = CairoMakie.Axis(fig[5, 1], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_80.grid.xᶜᵃᵃ[field_NN_data_80.indices[1][1]] / 1000) km") +axfield_90 = CairoMakie.Axis(fig[5, 5], xlabel="y (m)", ylabel="z (m)", title="x = $(field_NN_data_90.grid.xᶜᵃᵃ[field_NN_data_90.indices[1][1]] / 1000) km") + +axflux_00 = CairoMakie.Axis(fig[1, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_00.grid.xᶜᵃᵃ[flux_NN_data_00.indices[1][1]] / 1000) km") +axflux_10 = CairoMakie.Axis(fig[1, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_10.grid.xᶜᵃᵃ[flux_NN_data_10.indices[1][1]] / 1000) km") +axflux_20 = CairoMakie.Axis(fig[2, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_20.grid.xᶜᵃᵃ[flux_NN_data_20.indices[1][1]] / 1000) km") +axflux_30 = CairoMakie.Axis(fig[2, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_30.grid.xᶜᵃᵃ[flux_NN_data_30.indices[1][1]] / 1000) km") +axflux_40 = CairoMakie.Axis(fig[3, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_40.grid.xᶜᵃᵃ[flux_NN_data_40.indices[1][1]] / 1000) km") +axflux_50 = CairoMakie.Axis(fig[3, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_50.grid.xᶜᵃᵃ[flux_NN_data_50.indices[1][1]] / 1000) km") +axflux_60 = CairoMakie.Axis(fig[4, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_60.grid.xᶜᵃᵃ[flux_NN_data_60.indices[1][1]] / 1000) km") +axflux_70 = CairoMakie.Axis(fig[4, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_70.grid.xᶜᵃᵃ[flux_NN_data_70.indices[1][1]] / 1000) km") +axflux_80 = CairoMakie.Axis(fig[5, 3], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_80.grid.xᶜᵃᵃ[flux_NN_data_80.indices[1][1]] / 1000) km") +axflux_90 = CairoMakie.Axis(fig[5, 7], xlabel="y (m)", ylabel="z (m)", title="x = $(flux_NN_data_90.grid.xᶜᵃᵃ[flux_NN_data_90.indices[1][1]] / 1000) km") + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lim = (find_min(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_NN_data_00[timeframes[1]], :, :, zC_indices), interior(field_NN_data_00[timeframes[end]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (find_min(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_NN_data_00[timeframes[1]], :, :, zC_indices), interior(flux_NN_data_00[timeframes[end]], :, :, zC_indices))) + +flux_lim = (-maximum(abs, flux_lim), maximum(abs, flux_lim)) + +NN_00ₙ = @lift interior(field_NN_data_00[$n], 1, :, zC_indices) +NN_10ₙ = @lift interior(field_NN_data_10[$n], 1, :, zC_indices) +NN_20ₙ = @lift interior(field_NN_data_20[$n], 1, :, zC_indices) +NN_30ₙ = @lift interior(field_NN_data_30[$n], 1, :, zC_indices) +NN_40ₙ = @lift interior(field_NN_data_40[$n], 1, :, zC_indices) +NN_50ₙ = @lift interior(field_NN_data_50[$n], 1, :, zC_indices) +NN_60ₙ = @lift interior(field_NN_data_60[$n], 1, :, zC_indices) +NN_70ₙ = @lift interior(field_NN_data_70[$n], 1, :, zC_indices) +NN_80ₙ = @lift interior(field_NN_data_80[$n], 1, :, zC_indices) +NN_90ₙ = @lift interior(field_NN_data_90[$n], 1, :, zC_indices) + +flux_00ₙ = @lift interior(flux_NN_data_00[$n], 1, :, zF_indices) +flux_10ₙ = @lift interior(flux_NN_data_10[$n], 1, :, zF_indices) +flux_20ₙ = @lift interior(flux_NN_data_20[$n], 1, :, zF_indices) +flux_30ₙ = @lift interior(flux_NN_data_30[$n], 1, :, zF_indices) +flux_40ₙ = @lift interior(flux_NN_data_40[$n], 1, :, zF_indices) +flux_50ₙ = @lift interior(flux_NN_data_50[$n], 1, :, zF_indices) +flux_60ₙ = @lift interior(flux_NN_data_60[$n], 1, :, zF_indices) +flux_70ₙ = @lift interior(flux_NN_data_70[$n], 1, :, zF_indices) +flux_80ₙ = @lift interior(flux_NN_data_80[$n], 1, :, zF_indices) +flux_90ₙ = @lift interior(flux_NN_data_90[$n], 1, :, zF_indices) + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_00_surface = heatmap!(axfield_00, yC, zC[zC_indices], NN_00ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_10_surface = heatmap!(axfield_10, yC, zC[zC_indices], NN_10ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_20_surface = heatmap!(axfield_20, yC, zC[zC_indices], NN_20ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_30_surface = heatmap!(axfield_30, yC, zC[zC_indices], NN_30ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_40_surface = heatmap!(axfield_40, yC, zC[zC_indices], NN_40ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_50_surface = heatmap!(axfield_50, yC, zC[zC_indices], NN_50ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_60_surface = heatmap!(axfield_60, yC, zC[zC_indices], NN_60ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_70_surface = heatmap!(axfield_70, yC, zC[zC_indices], NN_70ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_80_surface = heatmap!(axfield_80, yC, zC[zC_indices], NN_80ₙ, colormap=colorscheme_field, colorrange=field_lim) +field_90_surface = heatmap!(axfield_90, yC, zC[zC_indices], NN_90ₙ, colormap=colorscheme_field, colorrange=field_lim) + +flux_00_surface = heatmap!(axflux_00, yC, zC[zF_indices], flux_00ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_10_surface = heatmap!(axflux_10, yC, zC[zF_indices], flux_10ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_20_surface = heatmap!(axflux_20, yC, zC[zF_indices], flux_20ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_30_surface = heatmap!(axflux_30, yC, zC[zF_indices], flux_30ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_40_surface = heatmap!(axflux_40, yC, zC[zF_indices], flux_40ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_50_surface = heatmap!(axflux_50, yC, zC[zF_indices], flux_50ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_60_surface = heatmap!(axflux_60, yC, zC[zF_indices], flux_60ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_70_surface = heatmap!(axflux_70, yC, zC[zF_indices], flux_70ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_80_surface = heatmap!(axflux_80, yC, zC[zF_indices], flux_80ₙ, colormap=colorscheme_flux, colorrange=flux_lim) +flux_90_surface = heatmap!(axflux_90, yC, zC[zF_indices], flux_90ₙ, colormap=colorscheme_flux, colorrange=flux_lim) + +Colorbar(fig[1:5, 2], field_00_surface, label="Field") +Colorbar(fig[1:5, 4], flux_00_surface, label="NN Flux") +Colorbar(fig[1:5, 6], field_00_surface, label="Field") +Colorbar(fig[1:5, 8], flux_00_surface, label="NN Flux") + +xlims!(axfield_00, minimum(yC), maximum(yC)) +xlims!(axfield_10, minimum(yC), maximum(yC)) +xlims!(axfield_20, minimum(yC), maximum(yC)) +xlims!(axfield_30, minimum(yC), maximum(yC)) +xlims!(axfield_40, minimum(yC), maximum(yC)) +xlims!(axfield_50, minimum(yC), maximum(yC)) +xlims!(axfield_60, minimum(yC), maximum(yC)) +xlims!(axfield_70, minimum(yC), maximum(yC)) +xlims!(axfield_80, minimum(yC), maximum(yC)) +xlims!(axfield_90, minimum(yC), maximum(yC)) + +ylims!(axfield_00, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_10, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_20, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_30, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_40, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_50, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_60, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_70, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_80, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +ylims!(axfield_90, minimum(zC[zC_indices]), maximum(zC[zC_indices])) + +xlims!(axflux_00, minimum(yC), maximum(yC)) +xlims!(axflux_10, minimum(yC), maximum(yC)) +xlims!(axflux_20, minimum(yC), maximum(yC)) +xlims!(axflux_30, minimum(yC), maximum(yC)) +xlims!(axflux_40, minimum(yC), maximum(yC)) +xlims!(axflux_50, minimum(yC), maximum(yC)) +xlims!(axflux_60, minimum(yC), maximum(yC)) +xlims!(axflux_70, minimum(yC), maximum(yC)) +xlims!(axflux_80, minimum(yC), maximum(yC)) +xlims!(axflux_90, minimum(yC), maximum(yC)) + +ylims!(axflux_00, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_10, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_20, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_30, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_40, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_50, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_60, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_70, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_80, minimum(zF[zF_indices]), maximum(zF[zF_indices])) +ylims!(axflux_90, minimum(zF[zF_indices]), maximum(zF[zF_indices])) + +title_str = @lift "Salinity (psu), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_yzslices_fluxes_S.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + @info "Recording frame $nn" + n[] = nn +end + +#%% +@info "Recording T fields and fluxes in xz" + +fieldname = "T" +fluxname = "wT_NN" + +field_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fieldname, backend=OnDisk()) +field_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fieldname, backend=OnDisk()) +field_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fieldname, backend=OnDisk()) +field_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fieldname, backend=OnDisk()) +field_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fieldname, backend=OnDisk()) +field_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fieldname, backend=OnDisk()) +field_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fieldname, backend=OnDisk()) +field_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fieldname, backend=OnDisk()) +field_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fieldname, backend=OnDisk()) +field_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fluxname, backend=OnDisk()) +flux_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fluxname, backend=OnDisk()) +flux_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fluxname, backend=OnDisk()) +flux_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fluxname, backend=OnDisk()) +flux_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fluxname, backend=OnDisk()) +flux_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fluxname, backend=OnDisk()) +flux_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fluxname, backend=OnDisk()) +flux_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fluxname, backend=OnDisk()) +flux_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fluxname, backend=OnDisk()) +flux_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fluxname, backend=OnDisk()) + +field_datas = [field_NN_data_5, field_NN_data_15, field_NN_data_25, field_NN_data_35, field_NN_data_45, field_NN_data_55, field_NN_data_65, field_NN_data_75, field_NN_data_85, field_NN_data_95] +flux_datas = [flux_NN_data_5, flux_NN_data_15, flux_NN_data_25, flux_NN_data_35, flux_NN_data_45, flux_NN_data_55, flux_NN_data_65, flux_NN_data_75, flux_NN_data_85, flux_NN_data_95] + +xC = field_NN_data_5.grid.xᶜᵃᵃ[1:field_NN_data_5.grid.Nx] +yC = field_NN_data_5.grid.yᵃᶜᵃ[1:field_NN_data_5.grid.Ny] +zC = field_NN_data_5.grid.zᵃᵃᶜ[1:field_NN_data_5.grid.Nz] +zF = field_NN_data_5.grid.zᵃᵃᶠ[1:field_NN_data_5.grid.Nz+1] + +Nt = length(field_NN_data_95) +times = field_NN_data_5.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_5 = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_5.grid.yᵃᶜᵃ[field_NN_data_5.indices[2][1]] / 1000) km") +axfield_15 = CairoMakie.Axis(fig[1, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_15.grid.yᵃᶜᵃ[field_NN_data_15.indices[2][1]] / 1000) km") +axfield_25 = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_25.grid.yᵃᶜᵃ[field_NN_data_25.indices[2][1]] / 1000) km") +axfield_35 = CairoMakie.Axis(fig[2, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_35.grid.yᵃᶜᵃ[field_NN_data_35.indices[2][1]] / 1000) km") +axfield_45 = CairoMakie.Axis(fig[3, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_45.grid.yᵃᶜᵃ[field_NN_data_45.indices[2][1]] / 1000) km") +axfield_55 = CairoMakie.Axis(fig[3, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_55.grid.yᵃᶜᵃ[field_NN_data_55.indices[2][1]] / 1000) km") +axfield_65 = CairoMakie.Axis(fig[4, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_65.grid.yᵃᶜᵃ[field_NN_data_65.indices[2][1]] / 1000) km") +axfield_75 = CairoMakie.Axis(fig[4, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_75.grid.yᵃᶜᵃ[field_NN_data_75.indices[2][1]] / 1000) km") +axfield_85 = CairoMakie.Axis(fig[5, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_85.grid.yᵃᶜᵃ[field_NN_data_85.indices[2][1]] / 1000) km") +axfield_95 = CairoMakie.Axis(fig[5, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_95.grid.yᵃᶜᵃ[field_NN_data_95.indices[2][1]] / 1000) km") + +axflux_5 = CairoMakie.Axis(fig[1, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_5.grid.yᵃᶜᵃ[flux_NN_data_5.indices[2][1]] / 1000) km") +axflux_15 = CairoMakie.Axis(fig[1, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_15.grid.yᵃᶜᵃ[flux_NN_data_15.indices[2][1]] / 1000) km") +axflux_25 = CairoMakie.Axis(fig[2, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_25.grid.yᵃᶜᵃ[flux_NN_data_25.indices[2][1]] / 1000) km") +axflux_35 = CairoMakie.Axis(fig[2, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_35.grid.yᵃᶜᵃ[flux_NN_data_35.indices[2][1]] / 1000) km") +axflux_45 = CairoMakie.Axis(fig[3, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_45.grid.yᵃᶜᵃ[flux_NN_data_45.indices[2][1]] / 1000) km") +axflux_55 = CairoMakie.Axis(fig[3, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_55.grid.yᵃᶜᵃ[flux_NN_data_55.indices[2][1]] / 1000) km") +axflux_65 = CairoMakie.Axis(fig[4, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_65.grid.yᵃᶜᵃ[flux_NN_data_65.indices[2][1]] / 1000) km") +axflux_75 = CairoMakie.Axis(fig[4, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_75.grid.yᵃᶜᵃ[flux_NN_data_75.indices[2][1]] / 1000) km") +axflux_85 = CairoMakie.Axis(fig[5, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_85.grid.yᵃᶜᵃ[flux_NN_data_85.indices[2][1]] / 1000) km") +axflux_95 = CairoMakie.Axis(fig[5, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_95.grid.yᵃᶜᵃ[flux_NN_data_95.indices[2][1]] / 1000) km") + +axfields = [axfield_5, axfield_15, axfield_25, axfield_35, axfield_45, axfield_55, axfield_65, axfield_75, axfield_85, axfield_95] +axfluxes = [axflux_5, axflux_15, axflux_25, axflux_35, axflux_45, axflux_55, axflux_65, axflux_75, axflux_85, axflux_95] + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lims = [(find_min(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices))) for field_data in field_datas] + +flux_lims = [(find_min(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices))) for flux_data in flux_datas] + +flux_lims = [(-maximum(abs, flux_lim), maximum(abs, flux_lim)) for flux_lim in flux_lims] + +NNₙs = [@lift interior(field_data[$n], :, 1, zC_indices) for field_data in field_datas] +fluxₙs = [@lift interior(flux_data[$n], :, 1, zF_indices) for flux_data in flux_datas] + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_5_surface = heatmap!(axfield_5, xC, zC[zC_indices], NN_5ₙ, colormap=colorscheme_field, colorrange=field_lims[1]) +field_15_surface = heatmap!(axfield_15, xC, zC[zC_indices], NN_15ₙ, colormap=colorscheme_field, colorrange=field_lims[2]) +field_25_surface = heatmap!(axfield_25, xC, zC[zC_indices], NN_25ₙ, colormap=colorscheme_field, colorrange=field_lims[3]) +field_35_surface = heatmap!(axfield_35, xC, zC[zC_indices], NN_35ₙ, colormap=colorscheme_field, colorrange=field_lims[4]) +field_45_surface = heatmap!(axfield_45, xC, zC[zC_indices], NN_45ₙ, colormap=colorscheme_field, colorrange=field_lims[5]) +field_55_surface = heatmap!(axfield_55, xC, zC[zC_indices], NN_55ₙ, colormap=colorscheme_field, colorrange=field_lims[6]) +field_65_surface = heatmap!(axfield_65, xC, zC[zC_indices], NN_65ₙ, colormap=colorscheme_field, colorrange=field_lims[7]) +field_75_surface = heatmap!(axfield_75, xC, zC[zC_indices], NN_75ₙ, colormap=colorscheme_field, colorrange=field_lims[8]) +field_85_surface = heatmap!(axfield_85, xC, zC[zC_indices], NN_85ₙ, colormap=colorscheme_field, colorrange=field_lims[9]) +field_95_surface = heatmap!(axfield_95, xC, zC[zC_indices], NN_95ₙ, colormap=colorscheme_field, colorrange=field_lims[10]) + +flux_5_surface = heatmap!(axflux_5, xC, zC[zF_indices], flux_5ₙ, colormap=colorscheme_flux, colorrange=flux_lims[1]) +flux_15_surface = heatmap!(axflux_15, xC, zC[zF_indices], flux_15ₙ, colormap=colorscheme_flux, colorrange=flux_lims[2]) +flux_25_surface = heatmap!(axflux_25, xC, zC[zF_indices], flux_25ₙ, colormap=colorscheme_flux, colorrange=flux_lims[3]) +flux_35_surface = heatmap!(axflux_35, xC, zC[zF_indices], flux_35ₙ, colormap=colorscheme_flux, colorrange=flux_lims[4]) +flux_45_surface = heatmap!(axflux_45, xC, zC[zF_indices], flux_45ₙ, colormap=colorscheme_flux, colorrange=flux_lims[5]) +flux_55_surface = heatmap!(axflux_55, xC, zC[zF_indices], flux_55ₙ, colormap=colorscheme_flux, colorrange=flux_lims[6]) +flux_65_surface = heatmap!(axflux_65, xC, zC[zF_indices], flux_65ₙ, colormap=colorscheme_flux, colorrange=flux_lims[7]) +flux_75_surface = heatmap!(axflux_75, xC, zC[zF_indices], flux_75ₙ, colormap=colorscheme_flux, colorrange=flux_lims[8]) +flux_85_surface = heatmap!(axflux_85, xC, zC[zF_indices], flux_85ₙ, colormap=colorscheme_flux, colorrange=flux_lims[9]) +flux_95_surface = heatmap!(axflux_95, xC, zC[zF_indices], flux_95ₙ, colormap=colorscheme_flux, colorrange=flux_lims[10]) + +Colorbar(fig[1, 2], field_5_surface, label="Field") +Colorbar(fig[1, 4], flux_5_surface, label="NN Flux") +Colorbar(fig[1, 6], field_15_surface, label="Field") +Colorbar(fig[1, 8], flux_15_surface, label="NN Flux") +Colorbar(fig[2, 2], field_25_surface, label="Field") +Colorbar(fig[2, 4], flux_25_surface, label="NN Flux") +Colorbar(fig[2, 6], field_35_surface, label="Field") +Colorbar(fig[2, 8], flux_35_surface, label="NN Flux") +Colorbar(fig[3, 2], field_45_surface, label="Field") +Colorbar(fig[3, 4], flux_45_surface, label="NN Flux") +Colorbar(fig[3, 6], field_55_surface, label="Field") +Colorbar(fig[3, 8], flux_55_surface, label="NN Flux") +Colorbar(fig[4, 2], field_65_surface, label="Field") +Colorbar(fig[4, 4], flux_65_surface, label="NN Flux") +Colorbar(fig[4, 6], field_75_surface, label="Field") +Colorbar(fig[4, 8], flux_75_surface, label="NN Flux") +Colorbar(fig[5, 2], field_85_surface, label="Field") +Colorbar(fig[5, 4], flux_85_surface, label="NN Flux") +Colorbar(fig[5, 6], field_95_surface, label="Field") +Colorbar(fig[5, 8], flux_95_surface, label="NN Flux") + +for axfield in axfields + xlims!(axfield, minimum(xC), maximum(xC)) + ylims!(axfield, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +end + +for axflux in axfluxes + xlims!(axflux, minimum(xC), maximum(xC)) + ylims!(axflux, minimum(zC[zF_indices]), maximum(zC[zF_indices])) +end + +title_str = @lift "Temperature (°C), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_xzslices_fluxes_T.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end + +#%% +@info "Recording S fields and fluxes in xz" + +fieldname = "S" +fluxname = "wS_NN" + +field_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fieldname, backend=OnDisk()) +field_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fieldname, backend=OnDisk()) +field_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fieldname, backend=OnDisk()) +field_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fieldname, backend=OnDisk()) +field_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fieldname, backend=OnDisk()) +field_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fieldname, backend=OnDisk()) +field_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fieldname, backend=OnDisk()) +field_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fieldname, backend=OnDisk()) +field_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fieldname, backend=OnDisk()) +field_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fieldname, backend=OnDisk()) + +flux_NN_data_5 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_5.jld2", fluxname, backend=OnDisk()) +flux_NN_data_15 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_15.jld2", fluxname, backend=OnDisk()) +flux_NN_data_25 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_25.jld2", fluxname, backend=OnDisk()) +flux_NN_data_35 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_35.jld2", fluxname, backend=OnDisk()) +flux_NN_data_45 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_45.jld2", fluxname, backend=OnDisk()) +flux_NN_data_55 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_55.jld2", fluxname, backend=OnDisk()) +flux_NN_data_65 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_65.jld2", fluxname, backend=OnDisk()) +flux_NN_data_75 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_75.jld2", fluxname, backend=OnDisk()) +flux_NN_data_85 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_85.jld2", fluxname, backend=OnDisk()) +flux_NN_data_95 = FieldTimeSeries("$(FILE_DIR)/instantaneous_fields_xz_95.jld2", fluxname, backend=OnDisk()) + +field_datas = [field_NN_data_5, field_NN_data_15, field_NN_data_25, field_NN_data_35, field_NN_data_45, field_NN_data_55, field_NN_data_65, field_NN_data_75, field_NN_data_85, field_NN_data_95] +flux_datas = [flux_NN_data_5, flux_NN_data_15, flux_NN_data_25, flux_NN_data_35, flux_NN_data_45, flux_NN_data_55, flux_NN_data_65, flux_NN_data_75, flux_NN_data_85, flux_NN_data_95] + +xC = field_NN_data_5.grid.xᶜᵃᵃ[1:field_NN_data_5.grid.Nx] +yC = field_NN_data_5.grid.yᵃᶜᵃ[1:field_NN_data_5.grid.Ny] +zC = field_NN_data_5.grid.zᵃᵃᶜ[1:field_NN_data_5.grid.Nz] +zF = field_NN_data_5.grid.zᵃᵃᶠ[1:field_NN_data_5.grid.Nz+1] + +Nt = length(field_NN_data_95) +times = field_NN_data_5.times / 24 / 60^2 / 365 +timeframes = 1:Nt + +function find_min(a...) + return minimum(minimum.([a...])) +end + +function find_max(a...) + return maximum(maximum.([a...])) +end + +#%% +fig = Figure(size=(3000, 1200)) + +axfield_5 = CairoMakie.Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_5.grid.yᵃᶜᵃ[field_NN_data_5.indices[2][1]] / 1000) km") +axfield_15 = CairoMakie.Axis(fig[1, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_15.grid.yᵃᶜᵃ[field_NN_data_15.indices[2][1]] / 1000) km") +axfield_25 = CairoMakie.Axis(fig[2, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_25.grid.yᵃᶜᵃ[field_NN_data_25.indices[2][1]] / 1000) km") +axfield_35 = CairoMakie.Axis(fig[2, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_35.grid.yᵃᶜᵃ[field_NN_data_35.indices[2][1]] / 1000) km") +axfield_45 = CairoMakie.Axis(fig[3, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_45.grid.yᵃᶜᵃ[field_NN_data_45.indices[2][1]] / 1000) km") +axfield_55 = CairoMakie.Axis(fig[3, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_55.grid.yᵃᶜᵃ[field_NN_data_55.indices[2][1]] / 1000) km") +axfield_65 = CairoMakie.Axis(fig[4, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_65.grid.yᵃᶜᵃ[field_NN_data_65.indices[2][1]] / 1000) km") +axfield_75 = CairoMakie.Axis(fig[4, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_75.grid.yᵃᶜᵃ[field_NN_data_75.indices[2][1]] / 1000) km") +axfield_85 = CairoMakie.Axis(fig[5, 1], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_85.grid.yᵃᶜᵃ[field_NN_data_85.indices[2][1]] / 1000) km") +axfield_95 = CairoMakie.Axis(fig[5, 5], xlabel="x (m)", ylabel="z (m)", title="y = $(field_NN_data_95.grid.yᵃᶜᵃ[field_NN_data_95.indices[2][1]] / 1000) km") + +axflux_5 = CairoMakie.Axis(fig[1, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_5.grid.yᵃᶜᵃ[flux_NN_data_5.indices[2][1]] / 1000) km") +axflux_15 = CairoMakie.Axis(fig[1, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_15.grid.yᵃᶜᵃ[flux_NN_data_15.indices[2][1]] / 1000) km") +axflux_25 = CairoMakie.Axis(fig[2, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_25.grid.yᵃᶜᵃ[flux_NN_data_25.indices[2][1]] / 1000) km") +axflux_35 = CairoMakie.Axis(fig[2, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_35.grid.yᵃᶜᵃ[flux_NN_data_35.indices[2][1]] / 1000) km") +axflux_45 = CairoMakie.Axis(fig[3, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_45.grid.yᵃᶜᵃ[flux_NN_data_45.indices[2][1]] / 1000) km") +axflux_55 = CairoMakie.Axis(fig[3, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_55.grid.yᵃᶜᵃ[flux_NN_data_55.indices[2][1]] / 1000) km") +axflux_65 = CairoMakie.Axis(fig[4, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_65.grid.yᵃᶜᵃ[flux_NN_data_65.indices[2][1]] / 1000) km") +axflux_75 = CairoMakie.Axis(fig[4, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_75.grid.yᵃᶜᵃ[flux_NN_data_75.indices[2][1]] / 1000) km") +axflux_85 = CairoMakie.Axis(fig[5, 3], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_85.grid.yᵃᶜᵃ[flux_NN_data_85.indices[2][1]] / 1000) km") +axflux_95 = CairoMakie.Axis(fig[5, 7], xlabel="x (m)", ylabel="z (m)", title="y = $(flux_NN_data_95.grid.yᵃᶜᵃ[flux_NN_data_95.indices[2][1]] / 1000) km") + +axfields = [axfield_5, axfield_15, axfield_25, axfield_35, axfield_45, axfield_55, axfield_65, axfield_75, axfield_85, axfield_95] +axfluxes = [axflux_5, axflux_15, axflux_25, axflux_35, axflux_45, axflux_55, axflux_65, axflux_75, axflux_85, axflux_95] + +n = Observable(1096) + +zC_indices = 1:200 +zF_indices = 2:200 + +field_lims = [(find_min(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(field_data[timeframes[1]], :, :, zC_indices), interior(field_data[timeframes[end]], :, :, zC_indices))) for field_data in field_datas] + +flux_lims = [(find_min(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices)), + find_max(interior(flux_data[timeframes[1]], :, :, zC_indices), interior(flux_data[timeframes[end]], :, :, zC_indices))) for flux_data in flux_datas] + +flux_lims = [(-maximum(abs, flux_lim), maximum(abs, flux_lim)) for flux_lim in flux_lims] + +NNₙs = [@lift interior(field_data[$n], :, 1, zC_indices) for field_data in field_datas] +fluxₙs = [@lift interior(flux_data[$n], :, 1, zF_indices) for flux_data in flux_datas] + +colorscheme_field = colorschemes[:viridis] +colorscheme_flux = colorschemes[:balance] + +field_5_surface = heatmap!(axfield_5, xC, zC[zC_indices], NN_5ₙ, colormap=colorscheme_field, colorrange=field_lims[1]) +field_15_surface = heatmap!(axfield_15, xC, zC[zC_indices], NN_15ₙ, colormap=colorscheme_field, colorrange=field_lims[2]) +field_25_surface = heatmap!(axfield_25, xC, zC[zC_indices], NN_25ₙ, colormap=colorscheme_field, colorrange=field_lims[3]) +field_35_surface = heatmap!(axfield_35, xC, zC[zC_indices], NN_35ₙ, colormap=colorscheme_field, colorrange=field_lims[4]) +field_45_surface = heatmap!(axfield_45, xC, zC[zC_indices], NN_45ₙ, colormap=colorscheme_field, colorrange=field_lims[5]) +field_55_surface = heatmap!(axfield_55, xC, zC[zC_indices], NN_55ₙ, colormap=colorscheme_field, colorrange=field_lims[6]) +field_65_surface = heatmap!(axfield_65, xC, zC[zC_indices], NN_65ₙ, colormap=colorscheme_field, colorrange=field_lims[7]) +field_75_surface = heatmap!(axfield_75, xC, zC[zC_indices], NN_75ₙ, colormap=colorscheme_field, colorrange=field_lims[8]) +field_85_surface = heatmap!(axfield_85, xC, zC[zC_indices], NN_85ₙ, colormap=colorscheme_field, colorrange=field_lims[9]) +field_95_surface = heatmap!(axfield_95, xC, zC[zC_indices], NN_95ₙ, colormap=colorscheme_field, colorrange=field_lims[10]) + +flux_5_surface = heatmap!(axflux_5, xC, zC[zF_indices], flux_5ₙ, colormap=colorscheme_flux, colorrange=flux_lims[1]) +flux_15_surface = heatmap!(axflux_15, xC, zC[zF_indices], flux_15ₙ, colormap=colorscheme_flux, colorrange=flux_lims[2]) +flux_25_surface = heatmap!(axflux_25, xC, zC[zF_indices], flux_25ₙ, colormap=colorscheme_flux, colorrange=flux_lims[3]) +flux_35_surface = heatmap!(axflux_35, xC, zC[zF_indices], flux_35ₙ, colormap=colorscheme_flux, colorrange=flux_lims[4]) +flux_45_surface = heatmap!(axflux_45, xC, zC[zF_indices], flux_45ₙ, colormap=colorscheme_flux, colorrange=flux_lims[5]) +flux_55_surface = heatmap!(axflux_55, xC, zC[zF_indices], flux_55ₙ, colormap=colorscheme_flux, colorrange=flux_lims[6]) +flux_65_surface = heatmap!(axflux_65, xC, zC[zF_indices], flux_65ₙ, colormap=colorscheme_flux, colorrange=flux_lims[7]) +flux_75_surface = heatmap!(axflux_75, xC, zC[zF_indices], flux_75ₙ, colormap=colorscheme_flux, colorrange=flux_lims[8]) +flux_85_surface = heatmap!(axflux_85, xC, zC[zF_indices], flux_85ₙ, colormap=colorscheme_flux, colorrange=flux_lims[9]) +flux_95_surface = heatmap!(axflux_95, xC, zC[zF_indices], flux_95ₙ, colormap=colorscheme_flux, colorrange=flux_lims[10]) + +Colorbar(fig[1, 2], field_5_surface, label="Field") +Colorbar(fig[1, 4], flux_5_surface, label="NN Flux") +Colorbar(fig[1, 6], field_15_surface, label="Field") +Colorbar(fig[1, 8], flux_15_surface, label="NN Flux") +Colorbar(fig[2, 2], field_25_surface, label="Field") +Colorbar(fig[2, 4], flux_25_surface, label="NN Flux") +Colorbar(fig[2, 6], field_35_surface, label="Field") +Colorbar(fig[2, 8], flux_35_surface, label="NN Flux") +Colorbar(fig[3, 2], field_45_surface, label="Field") +Colorbar(fig[3, 4], flux_45_surface, label="NN Flux") +Colorbar(fig[3, 6], field_55_surface, label="Field") +Colorbar(fig[3, 8], flux_55_surface, label="NN Flux") +Colorbar(fig[4, 2], field_65_surface, label="Field") +Colorbar(fig[4, 4], flux_65_surface, label="NN Flux") +Colorbar(fig[4, 6], field_75_surface, label="Field") +Colorbar(fig[4, 8], flux_75_surface, label="NN Flux") +Colorbar(fig[5, 2], field_85_surface, label="Field") +Colorbar(fig[5, 4], flux_85_surface, label="NN Flux") +Colorbar(fig[5, 6], field_95_surface, label="Field") +Colorbar(fig[5, 8], flux_95_surface, label="NN Flux") + +for axfield in axfields + xlims!(axfield, minimum(xC), maximum(xC)) + ylims!(axfield, minimum(zC[zC_indices]), maximum(zC[zC_indices])) +end + +for axflux in axfluxes + xlims!(axflux, minimum(xC), maximum(xC)) + ylims!(axflux, minimum(zC[zF_indices]), maximum(zC[zF_indices])) +end + +title_str = @lift "Salinity (psu), Time = $(round(times[$n], digits=2)) years" +Label(fig[0, :], text=title_str, tellwidth=false, font=:bold) + +trim!(fig.layout) + +CairoMakie.record(fig, "$(FILE_DIR)/$(filename)_xzslices_fluxes_S.mp4", 1:Nt, framerate=15, px_per_unit=2) do nn + n[] = nn +end +#%% \ No newline at end of file From 030c890e24e9182cd42f8e183ae8d4cd09283d1e Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 26 Nov 2024 13:16:22 -0500 Subject: [PATCH 108/114] run CATKE for 10800days --- physicalclosure_doublegyre_model_modewater.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/physicalclosure_doublegyre_model_modewater.jl b/physicalclosure_doublegyre_model_modewater.jl index ed5a62c0d9..6b80b049fb 100644 --- a/physicalclosure_doublegyre_model_modewater.jl +++ b/physicalclosure_doublegyre_model_modewater.jl @@ -188,7 +188,7 @@ update_state!(model) ##### Simulation building ##### Δt₀ = 5minutes -stop_time = 10950days +stop_time = 10800days simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) From 35dbad313d63222e7da73b688a3a08c1eeb74f84 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Mon, 2 Dec 2024 18:32:12 -0500 Subject: [PATCH 109/114] flush mode water after equilibration --- ...BLRifirstzone510_doublegyre_model_modewater.jl | 15 ++++++++++----- baseclosure_doublegyre_model_modewater.jl | 13 +++++++------ physicalclosure_doublegyre_model_modewater.jl | 15 ++++++++++----- 3 files changed, 27 insertions(+), 16 deletions(-) diff --git a/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl b/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl index 1f107c5cb9..f5a56d045d 100644 --- a/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl +++ b/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl @@ -21,9 +21,11 @@ using SeawaterPolynomials:TEOS10 using ColorSchemes #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_NN_closure_NDE5_Ri_BBLRifirztzone510_temp" +const Qᵀ_mode = 3.5e-4 +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_NN_closure_NDE5_Ri_BBLRifirztzone510_temp_QT$()" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +@info "$(FILE_DIR)" mkpath(FILE_DIR) # Architecture @@ -75,7 +77,6 @@ const μ_T = 1/30days const X₀ = -Lx/2 + 800kilometers const Y₀ = -Ly/2 + 1500kilometers const R₀ = 700kilometers -const Qᵀ_mode = 4.5e-4 const σ_mode = 20kilometers ##### @@ -107,7 +108,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), @inline T_ref(y) = T_mid - ΔT / Ly * y -@inline Qᵀ_winter(t) = max(0, -Qᵀ_mode * sin(2π * t / 360days)) +@inline Qᵀ_winter(t) = ifelse(t < 10800days, 0, max(0, -Qᵀ_mode * sin(2π * t / 360days))) @inline Qᵀ_subpolar(x, y, t) = ifelse((x - X₀)^2 + (y - Y₀)^2 <= R₀^2, Qᵀ_winter(t), exp(-(sqrt((x - X₀)^2 + (y - Y₀)^2) - R₀)^2 / (2 * σ_mode^2)) * Qᵀ_winter(t)) @@ -179,7 +180,7 @@ update_state!(model) ##### Simulation building ##### Δt₀ = 5minutes -stop_time = 10800days +stop_time = 12600days simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) @@ -394,7 +395,11 @@ simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; first_index, last_i simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), filename = "$(FILE_DIR)/averaged_fields_streamfunction", - schedule = AveragedTimeInterval(1825days, window=1825days)) + schedule = AveragedTimeInterval(1800days, window=1800days)) + +simulation.output_writers[:complete_fields] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields", + schedule = TimeInterval(1800days)) simulation.output_writers[:checkpointer] = Checkpointer(model, schedule = TimeInterval(730days), diff --git a/baseclosure_doublegyre_model_modewater.jl b/baseclosure_doublegyre_model_modewater.jl index 11328e9ede..40da3d8fd1 100644 --- a/baseclosure_doublegyre_model_modewater.jl +++ b/baseclosure_doublegyre_model_modewater.jl @@ -20,9 +20,11 @@ using SeawaterPolynomials:TEOS10 using ColorSchemes #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_newbaseclosure" +const Qᵀ_mode = 3.5e-4 +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_newbaseclosure_QT$(Qᵀ_mode)" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +@info "$(FILE_DIR)" mkpath(FILE_DIR) # Architecture @@ -73,7 +75,6 @@ const μ_T = 1/30days const X₀ = -Lx/2 + 800kilometers const Y₀ = -Ly/2 + 1500kilometers const R₀ = 700kilometers -const Qᵀ_mode = 4.5e-4 const σ_mode = 20kilometers ##### @@ -105,7 +106,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), @inline T_ref(y) = T_mid - ΔT / Ly * y -@inline Qᵀ_winter(t) = max(0, -Qᵀ_mode * sin(2π * t / 360days)) +@inline Qᵀ_winter(t) = ifelse(t < 10800days, 0, max(0, -Qᵀ_mode * sin(2π * t / 360days))) @inline Qᵀ_subpolar(x, y, t) = ifelse((x - X₀)^2 + (y - Y₀)^2 <= R₀^2, Qᵀ_winter(t), exp(-(sqrt((x - X₀)^2 + (y - Y₀)^2) - R₀)^2 / (2 * σ_mode^2)) * Qᵀ_winter(t)) @@ -177,7 +178,7 @@ update_state!(model) ##### Simulation building ##### Δt₀ = 5minutes -stop_time = 10800days +stop_time = 12600days simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) @@ -316,11 +317,11 @@ simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_output simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), filename = "$(FILE_DIR)/averaged_fields_streamfunction", - schedule = AveragedTimeInterval(1825days, window=1825days)) + schedule = AveragedTimeInterval(1800days, window=1800days)) simulation.output_writers[:complete_fields] = JLD2OutputWriter(model, outputs, filename = "$(FILE_DIR)/instantaneous_fields", - schedule = TimeInterval(1825days)) + schedule = TimeInterval(1800days)) simulation.output_writers[:checkpointer] = Checkpointer(model, schedule = TimeInterval(730days), diff --git a/physicalclosure_doublegyre_model_modewater.jl b/physicalclosure_doublegyre_model_modewater.jl index 6b80b049fb..9a5777b97a 100644 --- a/physicalclosure_doublegyre_model_modewater.jl +++ b/physicalclosure_doublegyre_model_modewater.jl @@ -27,9 +27,11 @@ using ColorSchemes using Glob #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_CATKEVerticalDiffusivity" +const Qᵀ_mode = 3.5e-4 +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_CATKEVerticalDiffusivity_QT$(Qᵀ_mode)" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +@info "$(FILE_DIR)" mkpath(FILE_DIR) # Architecture @@ -85,7 +87,6 @@ const μ_T = 1/30days const X₀ = -Lx/2 + 800kilometers const Y₀ = -Ly/2 + 1500kilometers const R₀ = 700kilometers -const Qᵀ_mode = 4.5e-4 const σ_mode = 20kilometers ##### @@ -116,7 +117,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), @inline T_ref(y) = T_mid - ΔT / Ly * y -@inline Qᵀ_winter(t) = max(0, -Qᵀ_mode * sin(2π * t / 360days)) +@inline Qᵀ_winter(t) = ifelse(t < 10800days, 0, max(0, -Qᵀ_mode * sin(2π * t / 360days))) @inline Qᵀ_subpolar(x, y, t) = ifelse((x - X₀)^2 + (y - Y₀)^2 <= R₀^2, Qᵀ_winter(t), exp(-(sqrt((x - X₀)^2 + (y - Y₀)^2) - R₀)^2 / (2 * σ_mode^2)) * Qᵀ_winter(t)) @@ -188,7 +189,7 @@ update_state!(model) ##### Simulation building ##### Δt₀ = 5minutes -stop_time = 10800days +stop_time = 12600days simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) @@ -332,7 +333,11 @@ simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_output simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), filename = "$(FILE_DIR)/averaged_fields_streamfunction", - schedule = AveragedTimeInterval(1825days, window=1825days)) + schedule = AveragedTimeInterval(1800days, window=1800days)) + +simulation.output_writers[:complete_fields] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields", + schedule = TimeInterval(1800days)) simulation.output_writers[:checkpointer] = Checkpointer(model, schedule = TimeInterval(730days), From 88e3cf18e4c028d585fee9dd2c20abd7a99b5982 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Mon, 2 Dec 2024 18:33:06 -0500 Subject: [PATCH 110/114] fix filename --- ...osure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl b/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl index f5a56d045d..9fbf9cf1fd 100644 --- a/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl +++ b/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl @@ -22,7 +22,7 @@ using ColorSchemes #%% const Qᵀ_mode = 3.5e-4 -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_NN_closure_NDE5_Ri_BBLRifirztzone510_temp_QT$()" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_NN_closure_NDE5_Ri_BBLRifirztzone510_temp_QT$(Qᵀ_mode)" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" @info "$(FILE_DIR)" From 5b340276f26c246a77da6dcb293c780100916b8e Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 3 Dec 2024 06:47:08 -0500 Subject: [PATCH 111/114] Revert "fix filename" This reverts commit 88e3cf18e4c028d585fee9dd2c20abd7a99b5982. --- ...osure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl b/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl index 9fbf9cf1fd..f5a56d045d 100644 --- a/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl +++ b/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl @@ -22,7 +22,7 @@ using ColorSchemes #%% const Qᵀ_mode = 3.5e-4 -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_NN_closure_NDE5_Ri_BBLRifirztzone510_temp_QT$(Qᵀ_mode)" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_NN_closure_NDE5_Ri_BBLRifirztzone510_temp_QT$()" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" @info "$(FILE_DIR)" From 8e458e0616d916704f1ef9b2d0805ced024c4a8d Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 3 Dec 2024 06:47:41 -0500 Subject: [PATCH 112/114] Revert "flush mode water after equilibration" This reverts commit 35dbad313d63222e7da73b688a3a08c1eeb74f84. --- ...BLRifirstzone510_doublegyre_model_modewater.jl | 15 +++++---------- baseclosure_doublegyre_model_modewater.jl | 13 ++++++------- physicalclosure_doublegyre_model_modewater.jl | 15 +++++---------- 3 files changed, 16 insertions(+), 27 deletions(-) diff --git a/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl b/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl index f5a56d045d..1f107c5cb9 100644 --- a/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl +++ b/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl @@ -21,11 +21,9 @@ using SeawaterPolynomials:TEOS10 using ColorSchemes #%% -const Qᵀ_mode = 3.5e-4 -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_NN_closure_NDE5_Ri_BBLRifirztzone510_temp_QT$()" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_NN_closure_NDE5_Ri_BBLRifirztzone510_temp" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" -@info "$(FILE_DIR)" mkpath(FILE_DIR) # Architecture @@ -77,6 +75,7 @@ const μ_T = 1/30days const X₀ = -Lx/2 + 800kilometers const Y₀ = -Ly/2 + 1500kilometers const R₀ = 700kilometers +const Qᵀ_mode = 4.5e-4 const σ_mode = 20kilometers ##### @@ -108,7 +107,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), @inline T_ref(y) = T_mid - ΔT / Ly * y -@inline Qᵀ_winter(t) = ifelse(t < 10800days, 0, max(0, -Qᵀ_mode * sin(2π * t / 360days))) +@inline Qᵀ_winter(t) = max(0, -Qᵀ_mode * sin(2π * t / 360days)) @inline Qᵀ_subpolar(x, y, t) = ifelse((x - X₀)^2 + (y - Y₀)^2 <= R₀^2, Qᵀ_winter(t), exp(-(sqrt((x - X₀)^2 + (y - Y₀)^2) - R₀)^2 / (2 * σ_mode^2)) * Qᵀ_winter(t)) @@ -180,7 +179,7 @@ update_state!(model) ##### Simulation building ##### Δt₀ = 5minutes -stop_time = 12600days +stop_time = 10800days simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) @@ -395,11 +394,7 @@ simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; first_index, last_i simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), filename = "$(FILE_DIR)/averaged_fields_streamfunction", - schedule = AveragedTimeInterval(1800days, window=1800days)) - -simulation.output_writers[:complete_fields] = JLD2OutputWriter(model, outputs, - filename = "$(FILE_DIR)/instantaneous_fields", - schedule = TimeInterval(1800days)) + schedule = AveragedTimeInterval(1825days, window=1825days)) simulation.output_writers[:checkpointer] = Checkpointer(model, schedule = TimeInterval(730days), diff --git a/baseclosure_doublegyre_model_modewater.jl b/baseclosure_doublegyre_model_modewater.jl index 40da3d8fd1..11328e9ede 100644 --- a/baseclosure_doublegyre_model_modewater.jl +++ b/baseclosure_doublegyre_model_modewater.jl @@ -20,11 +20,9 @@ using SeawaterPolynomials:TEOS10 using ColorSchemes #%% -const Qᵀ_mode = 3.5e-4 -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_newbaseclosure_QT$(Qᵀ_mode)" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_newbaseclosure" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" -@info "$(FILE_DIR)" mkpath(FILE_DIR) # Architecture @@ -75,6 +73,7 @@ const μ_T = 1/30days const X₀ = -Lx/2 + 800kilometers const Y₀ = -Ly/2 + 1500kilometers const R₀ = 700kilometers +const Qᵀ_mode = 4.5e-4 const σ_mode = 20kilometers ##### @@ -106,7 +105,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), @inline T_ref(y) = T_mid - ΔT / Ly * y -@inline Qᵀ_winter(t) = ifelse(t < 10800days, 0, max(0, -Qᵀ_mode * sin(2π * t / 360days))) +@inline Qᵀ_winter(t) = max(0, -Qᵀ_mode * sin(2π * t / 360days)) @inline Qᵀ_subpolar(x, y, t) = ifelse((x - X₀)^2 + (y - Y₀)^2 <= R₀^2, Qᵀ_winter(t), exp(-(sqrt((x - X₀)^2 + (y - Y₀)^2) - R₀)^2 / (2 * σ_mode^2)) * Qᵀ_winter(t)) @@ -178,7 +177,7 @@ update_state!(model) ##### Simulation building ##### Δt₀ = 5minutes -stop_time = 12600days +stop_time = 10800days simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) @@ -317,11 +316,11 @@ simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_output simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), filename = "$(FILE_DIR)/averaged_fields_streamfunction", - schedule = AveragedTimeInterval(1800days, window=1800days)) + schedule = AveragedTimeInterval(1825days, window=1825days)) simulation.output_writers[:complete_fields] = JLD2OutputWriter(model, outputs, filename = "$(FILE_DIR)/instantaneous_fields", - schedule = TimeInterval(1800days)) + schedule = TimeInterval(1825days)) simulation.output_writers[:checkpointer] = Checkpointer(model, schedule = TimeInterval(730days), diff --git a/physicalclosure_doublegyre_model_modewater.jl b/physicalclosure_doublegyre_model_modewater.jl index 9a5777b97a..6b80b049fb 100644 --- a/physicalclosure_doublegyre_model_modewater.jl +++ b/physicalclosure_doublegyre_model_modewater.jl @@ -27,11 +27,9 @@ using ColorSchemes using Glob #%% -const Qᵀ_mode = 3.5e-4 -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_CATKEVerticalDiffusivity_QT$(Qᵀ_mode)" +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_CATKEVerticalDiffusivity" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" -@info "$(FILE_DIR)" mkpath(FILE_DIR) # Architecture @@ -87,6 +85,7 @@ const μ_T = 1/30days const X₀ = -Lx/2 + 800kilometers const Y₀ = -Ly/2 + 1500kilometers const R₀ = 700kilometers +const Qᵀ_mode = 4.5e-4 const σ_mode = 20kilometers ##### @@ -117,7 +116,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), @inline T_ref(y) = T_mid - ΔT / Ly * y -@inline Qᵀ_winter(t) = ifelse(t < 10800days, 0, max(0, -Qᵀ_mode * sin(2π * t / 360days))) +@inline Qᵀ_winter(t) = max(0, -Qᵀ_mode * sin(2π * t / 360days)) @inline Qᵀ_subpolar(x, y, t) = ifelse((x - X₀)^2 + (y - Y₀)^2 <= R₀^2, Qᵀ_winter(t), exp(-(sqrt((x - X₀)^2 + (y - Y₀)^2) - R₀)^2 / (2 * σ_mode^2)) * Qᵀ_winter(t)) @@ -189,7 +188,7 @@ update_state!(model) ##### Simulation building ##### Δt₀ = 5minutes -stop_time = 12600days +stop_time = 10800days simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) @@ -333,11 +332,7 @@ simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_output simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), filename = "$(FILE_DIR)/averaged_fields_streamfunction", - schedule = AveragedTimeInterval(1800days, window=1800days)) - -simulation.output_writers[:complete_fields] = JLD2OutputWriter(model, outputs, - filename = "$(FILE_DIR)/instantaneous_fields", - schedule = TimeInterval(1800days)) + schedule = AveragedTimeInterval(1825days, window=1825days)) simulation.output_writers[:checkpointer] = Checkpointer(model, schedule = TimeInterval(730days), From 34f74371dfa1484bded3497b7bde09fec4149887 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 3 Dec 2024 06:49:42 -0500 Subject: [PATCH 113/114] Revert "Merge branch 'main' into xk/embed-nn" This reverts commit f4f7383138f62100ee6bd28037f193fbf6545e81, reversing changes made to 030c890e24e9182cd42f8e183ae8d4cd09283d1e. --- .buildkite/distributed/pipeline.yml | 60 +-- .buildkite/pipeline.yml | 120 ++--- Dockerfile | 2 +- Project.toml | 13 +- benchmark/benchmark_advection_schemes.jl | 2 +- docs/Project.toml | 2 +- docs/oceananigans.bib | 21 - docs/src/appendix/library.md | 7 - docs/src/fields.md | 14 +- docs/src/grids.md | 25 +- docs/src/index.md | 4 - docs/src/model_setup/boundary_conditions.md | 4 +- .../buoyancy_and_equation_of_state.md | 18 +- docs/src/model_setup/callbacks.md | 1 - docs/src/model_setup/lagrangian_particles.md | 59 ++- docs/src/model_setup/output_writers.md | 89 +++- docs/src/model_setup/tracers.md | 6 +- ...ent_diffusivity_closures_and_les_models.md | 4 +- .../boundary_conditions.md | 4 +- docs/src/operations.md | 2 - examples/convecting_plankton.jl | 3 +- examples/horizontal_convection.jl | 5 + examples/internal_tide.jl | 33 +- examples/internal_wave.jl | 3 +- examples/kelvin_helmholtz_instability.jl | 14 +- examples/langmuir_turbulence.jl | 44 +- examples/ocean_wind_mixing_and_convection.jl | 31 +- examples/tilted_bottom_boundary_layer.jl | 26 +- examples/two_dimensional_turbulence.jl | 8 +- ext/OceananigansMakieExt.jl | 6 +- src/AbstractOperations/AbstractOperations.jl | 8 +- src/AbstractOperations/at.jl | 5 +- src/AbstractOperations/grid_metrics.jl | 176 +------ src/Advection/Advection.jl | 4 +- src/Advection/adapt_advection_order.jl | 34 +- src/Advection/centered_reconstruction.jl | 5 +- src/Advection/flux_form_advection.jl | 13 +- ...y_preserving_tracer_advection_operators.jl | 4 +- src/Advection/upwind_biased_reconstruction.jl | 10 +- src/Advection/vector_invariant_advection.jl | 104 ++-- src/Advection/vector_invariant_upwinding.jl | 12 +- src/Advection/weno_reconstruction.jl | 47 +- src/Biogeochemistry.jl | 2 +- .../field_boundary_conditions.jl | 25 +- src/BoundaryConditions/fill_halo_regions.jl | 22 +- .../fill_halo_regions_open.jl | 135 +++-- .../fill_halo_regions_periodic.jl | 6 +- ...apolation_open_boundary_matching_scheme.jl | 4 +- src/BuoyancyModels/buoyancy.jl | 2 +- .../halo_communication.jl | 29 +- src/Fields/Fields.jl | 10 +- src/Fields/abstract_field.jl | 4 + src/Fields/field.jl | 41 +- src/Fields/interpolate.jl | 64 ++- src/Fields/set!.jl | 25 +- src/Forcings/advective_forcing.jl | 2 +- src/Grids/Grids.jl | 4 +- src/Grids/automatic_halo_sizing.jl | 12 +- src/Grids/grid_generation.jl | 1 + src/Grids/grid_utils.jl | 9 - src/Grids/latitude_longitude_grid.jl | 86 +++- src/Grids/nodes_and_spacings.jl | 107 +++- src/Grids/orthogonal_spherical_shell_grid.jl | 44 +- src/Grids/rectilinear_grid.jl | 19 +- src/ImmersedBoundaries/ImmersedBoundaries.jl | 222 +++++++- .../abstract_grid_fitted_boundary.jl | 18 + .../conditional_differences.jl | 25 +- src/ImmersedBoundaries/grid_fitted_bottom.jl | 82 +-- .../immersed_boundary_grid.jl | 90 ---- .../immersed_boundary_interface.jl | 56 -- .../immersed_boundary_nodes.jl | 69 --- .../immersed_grid_metrics.jl | 23 +- src/ImmersedBoundaries/mask_immersed_field.jl | 51 +- src/ImmersedBoundaries/partial_cell_bottom.jl | 68 +-- .../HydrostaticFreeSurfaceModels.jl | 31 +- .../SplitExplicitFreeSurfaces.jl | 43 -- .../barotropic_split_explicit_corrector.jl | 60 --- .../compute_slow_tendencies.jl | 154 ------ ...distributed_split_explicit_free_surface.jl | 25 - .../initialize_split_explicit_substepping.jl | 60 --- .../split_explicit_free_surface.jl | 327 ------------ .../split_explicit_timesteppers.jl | 160 ------ .../step_split_explicit_free_surface.jl | 168 ------ .../barotropic_pressure_correction.jl | 23 +- ...ute_hydrostatic_free_surface_tendencies.jl | 65 ++- ...distributed_split_explicit_free_surface.jl | 112 ++++ .../explicit_free_surface.jl | 103 +--- .../hydrostatic_free_surface_ab2_step.jl | 17 +- .../hydrostatic_free_surface_field_tuples.jl | 26 - .../hydrostatic_free_surface_model.jl | 120 +++-- .../hydrostatic_free_surface_rk3_step.jl | 122 ----- ..._free_surface_tendency_kernel_functions.jl | 43 +- .../implicit_free_surface.jl | 5 +- .../prescribed_hydrostatic_velocity_fields.jl | 5 +- .../{nothing_free_surface.jl => rigid_lid.jl} | 0 .../show_hydrostatic_free_surface_model.jl | 2 +- .../single_column_model_mode.jl | 8 +- .../split_explicit_free_surface.jl | 424 +++++++++++++++ .../split_explicit_free_surface_kernels.jl | 482 ++++++++++++++++++ ...te_hydrostatic_free_surface_model_state.jl | 1 - .../LagrangianParticleTracking.jl | 2 +- .../lagrangian_particle_advection.jl | 52 +- src/Models/Models.jl | 4 +- .../compute_nonhydrostatic_tendencies.jl | 8 +- .../nonhydrostatic_model.jl | 6 +- ...onhydrostatic_tendency_kernel_functions.jl | 28 +- .../update_nonhydrostatic_model_state.jl | 2 +- .../compute_shallow_water_tendencies.jl | 12 +- .../ShallowWaterModels/shallow_water_model.jl | 8 +- .../solution_and_tracer_tendencies.jl | 22 +- src/Models/seawater_density.jl | 2 +- src/MultiRegion/cubed_sphere_grid.jl | 10 +- .../multi_region_boundary_conditions.jl | 16 +- src/MultiRegion/multi_region_field.jl | 2 - src/MultiRegion/multi_region_grid.jl | 62 ++- src/MultiRegion/multi_region_models.jl | 2 + ...ulti_region_split_explicit_free_surface.jl | 91 ++-- src/Oceananigans.jl | 125 +++-- src/Operators/Operators.jl | 31 +- src/Operators/derivative_operators.jl | 12 +- src/Operators/interpolation_utils.jl | 12 +- .../spacings_and_areas_and_volumes.jl | 167 +++--- src/Operators/topology_aware_operators.jl | 67 --- src/Operators/vector_rotation_operators.jl | 62 ++- src/Operators/vorticity_operators.jl | 35 +- src/OutputReaders/field_dataset.jl | 24 +- src/OutputReaders/field_time_series.jl | 113 ++-- .../field_time_series_indexing.jl | 39 +- src/OutputReaders/set_field_time_series.jl | 9 +- src/OutputWriters/jld2_output_writer.jl | 37 +- src/OutputWriters/netcdf_output_writer.jl | 60 ++- src/OutputWriters/output_writer_utils.jl | 8 +- src/OutputWriters/windowed_time_average.jl | 11 +- src/Simulations/Simulations.jl | 1 + src/Simulations/simulation.jl | 2 +- src/Solvers/Solvers.jl | 2 +- .../conjugate_gradient_poisson_solver.jl | 6 +- src/TimeSteppers/TimeSteppers.jl | 1 - src/TimeSteppers/quasi_adams_bashforth_2.jl | 3 +- .../split_hydrostatic_runge_kutta_3.jl | 175 ------- src/TurbulenceClosures/TurbulenceClosures.jl | 25 +- .../turbulence_closure_diagnostics.jl | 2 +- .../Smagorinskys/Smagorinskys.jl | 15 - .../Smagorinskys/dynamic_coefficient.jl | 266 ---------- .../Smagorinskys/lilly_coefficient.jl | 146 ------ .../Smagorinskys/scale_invariant_operators.jl | 181 ------- .../Smagorinskys/smagorinsky.jl | 147 ------ .../catke_equation.jl | 2 +- .../catke_mixing_length.jl | 13 +- .../catke_vertical_diffusivity.jl | 22 +- .../time_step_catke_equation.jl | 5 +- .../smagorinsky_lilly.jl | 213 ++++++++ src/Utils/kernel_launching.jl | 29 +- test/dependencies_for_runtests.jl | 5 +- ...n_large_eddy_simulation_regression_test.jl | 11 +- test/runtests.jl | 11 - test/test_checkpointer.jl | 6 +- .../test_conjugate_gradient_poisson_solver.jl | 28 - test/test_cubed_spheres.jl | 232 +++++++++ test/test_distributed_poisson_solvers.jl | 4 +- test/test_distributed_transpose.jl | 2 + test/test_enzyme.jl | 84 ++- test/test_field.jl | 45 -- test/test_field_scans.jl | 23 +- test/test_forcings.jl | 12 +- test/test_grids.jl | 188 +++---- ...face_immersed_boundaries_implicit_solve.jl | 4 +- test/test_hydrostatic_free_surface_models.jl | 111 ++-- test/test_hydrostatic_regression.jl | 8 +- test/test_immersed_advection.jl | 12 +- test/test_immersed_boundary_grid.jl | 14 +- test/test_implicit_free_surface_solver.jl | 10 +- test/test_jld2_output_writer.jl | 4 +- test/test_lagrangian_particle_tracking.jl | 18 - test/test_nonhydrostatic_models.jl | 4 +- test/test_output_readers.jl | 74 +-- test/test_poisson_solvers.jl | 14 - test/test_shallow_water_models.jl | 2 +- ...test_split_explicit_free_surface_solver.jl | 111 ++-- .../test_split_explicit_vertical_integrals.jl | 67 +-- test/test_time_stepping.jl | 26 +- test/test_turbulence_closures.jl | 73 +-- test/test_vector_rotation_operators.jl | 178 ------- test/utils_for_runtests.jl | 84 +-- .../plot_rates_convergence_advection.jl | 14 +- ...lankton_continuous_form_biogeochemistry.jl | 4 +- .../OneDimensionalCosineAdvectionDiffusion.jl | 2 +- ...neDimensionalGaussianAdvectionDiffusion.jl | 2 +- ...woDimensionalGaussianAdvectionDiffusion.jl | 2 +- .../distributed_shallow_water_turbulence.jl | 2 +- .../2D_rough_rayleighbenard.jl | 4 +- .../cylinder_flow_with_tracer.jl | 6 +- .../immersed_boundaries/internal_tide.jl | 2 +- .../nonlinear_topography.jl | 4 +- .../resting_stratified_bumpy_ocean.jl | 4 +- .../tracer_advection_over_bump.jl | 2 +- .../3d_turbulence_smagorinsky.jl | 102 ---- .../compare_closures.jl | 143 ------ .../visualize_closure_comparison.jl | 102 ---- .../isotropic_mixing_closures/wall_flow.jl | 136 ----- .../particles_in_convection.jl | 2 +- .../zonally_averaged_abernathey_channel.jl | 2 +- .../multi_region_internal_tide.jl | 4 +- .../open_boundaries/oscillating_flow.jl | 221 ++++---- .../periodic_advection/periodic_advection.jl | 2 +- .../plot_stommel_gyre_advection.py | 2 +- .../stommel_gyre/stommel_gyre_advection.jl | 2 +- .../plot_thermal_bubble_advection.py | 2 +- validation/thermal_bubble/thermal_bubble.jl | 2 +- .../gpu_tkevd_ensemble.jl | 12 +- .../heterogeneous_windy_convection.jl | 12 +- 211 files changed, 3893 insertions(+), 5758 deletions(-) delete mode 100644 src/ImmersedBoundaries/immersed_boundary_grid.jl delete mode 100644 src/ImmersedBoundaries/immersed_boundary_interface.jl delete mode 100644 src/ImmersedBoundaries/immersed_boundary_nodes.jl delete mode 100644 src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl delete mode 100644 src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl delete mode 100644 src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl delete mode 100644 src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/distributed_split_explicit_free_surface.jl delete mode 100644 src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl delete mode 100644 src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl delete mode 100644 src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl delete mode 100644 src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl create mode 100644 src/Models/HydrostaticFreeSurfaceModels/distributed_split_explicit_free_surface.jl delete mode 100644 src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk3_step.jl rename src/Models/HydrostaticFreeSurfaceModels/{nothing_free_surface.jl => rigid_lid.jl} (100%) create mode 100644 src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface.jl create mode 100644 src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl delete mode 100644 src/Operators/topology_aware_operators.jl delete mode 100644 src/TimeSteppers/split_hydrostatic_runge_kutta_3.jl delete mode 100644 src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/Smagorinskys.jl delete mode 100644 src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/dynamic_coefficient.jl delete mode 100644 src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/lilly_coefficient.jl delete mode 100644 src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/scale_invariant_operators.jl delete mode 100644 src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/smagorinsky.jl create mode 100644 src/TurbulenceClosures/turbulence_closure_implementations/smagorinsky_lilly.jl delete mode 100644 test/test_conjugate_gradient_poisson_solver.jl create mode 100644 test/test_cubed_spheres.jl delete mode 100644 test/test_vector_rotation_operators.jl delete mode 100644 validation/isotropic_mixing_closures/3d_turbulence_smagorinsky.jl delete mode 100644 validation/isotropic_mixing_closures/compare_closures.jl delete mode 100644 validation/isotropic_mixing_closures/visualize_closure_comparison.jl delete mode 100644 validation/isotropic_mixing_closures/wall_flow.jl diff --git a/.buildkite/distributed/pipeline.yml b/.buildkite/distributed/pipeline.yml index 9f7869de63..0c06cbb02f 100644 --- a/.buildkite/distributed/pipeline.yml +++ b/.buildkite/distributed/pipeline.yml @@ -1,7 +1,7 @@ agents: queue: new-central - slurm_mem: 8G # Note that the tests run on shared nodes, so limiting the memory usage might help in avoiding long queues - modules: climacommon/2024_10_08 + slurm_mem: 8G + modules: climacommon/2024_05_27 env: JULIA_LOAD_PATH: "${JULIA_LOAD_PATH}:${BUILDKITE_BUILD_CHECKOUT_PATH}/.buildkite/distributed" @@ -16,14 +16,13 @@ steps: key: "init_central" env: TEST_GROUP: "init" - GPU_TEST: "true" command: - - echo "--- Initialize tests" - - "julia -O0 --project -e 'using Pkg; Pkg.test()'" + - echo "--- Instantiate project" + - "julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: - slurm_mem: 8G - slurm_ntasks: 1 - slurm_gpus_per_task: 1 + slurm_mem: 120G + slurm_gpus: 1 + slurm_cpus_per_task: 8 - wait @@ -31,27 +30,20 @@ steps: key: "distributed_cpu" env: TEST_GROUP: "distributed" - MPI_TEST: "true" commands: - "srun julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: - slurm_mem: 8G + slurm_mem: 120G slurm_ntasks: 4 - retry: - automatic: - - exit_status: 1 - limit: 1 - label: "🐲 gpu distributed unit tests" key: "distributed_gpu" env: TEST_GROUP: "distributed" - GPU_TEST: "true" - MPI_TEST: "true" commands: - "srun julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: - slurm_mem: 8G + slurm_mem: 120G slurm_ntasks: 4 slurm_gpus_per_task: 1 retry: @@ -59,31 +51,25 @@ steps: - exit_status: 1 limit: 1 + - label: "🦾 cpu distributed solvers tests" key: "distributed_solvers_cpu" env: TEST_GROUP: "distributed_solvers" - MPI_TEST: "true" commands: - "srun julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: - slurm_mem: 50G + slurm_mem: 120G slurm_ntasks: 4 - retry: - automatic: - - exit_status: 1 - limit: 1 - label: "🛸 gpu distributed solvers tests" key: "distributed_solvers_gpu" env: TEST_GROUP: "distributed_solvers" - GPU_TEST: "true" - MPI_TEST: "true" commands: - "srun julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: - slurm_mem: 50G + slurm_mem: 120G slurm_ntasks: 4 slurm_gpus_per_task: 1 retry: @@ -95,27 +81,20 @@ steps: key: "distributed_hydrostatic_model_cpu" env: TEST_GROUP: "distributed_hydrostatic_model" - MPI_TEST: "true" commands: - "srun julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: - slurm_mem: 50G + slurm_mem: 120G slurm_ntasks: 4 - retry: - automatic: - - exit_status: 1 - limit: 1 - label: "🦏 gpu distributed hydrostatic model tests" key: "distributed_hydrostatic_model_gpu" env: TEST_GROUP: "distributed_hydrostatic_model" - GPU_TEST: "true" - MPI_TEST: "true" commands: - "srun julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: - slurm_mem: 80G # Apparently the GPU tests require more memory + slurm_mem: 120G slurm_ntasks: 4 slurm_gpus_per_task: 1 retry: @@ -127,27 +106,20 @@ steps: key: "distributed_nonhydrostatic_regression_cpu" env: TEST_GROUP: "distributed_nonhydrostatic_regression" - MPI_TEST: "true" commands: - "srun julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: - slurm_mem: 50G + slurm_mem: 120G slurm_ntasks: 4 - retry: - automatic: - - exit_status: 1 - limit: 1 - label: "🕺 gpu distributed nonhydrostatic regression" key: "distributed_nonhydrostatic_regression_gpu" env: TEST_GROUP: "distributed_nonhydrostatic_regression" - GPU_TEST: "true" - MPI_TEST: "true" commands: - "srun julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: - slurm_mem: 50G + slurm_mem: 120G slurm_ntasks: 4 slurm_gpus_per_task: 1 retry: diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index d8dc9987da..32ae563fae 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -1,5 +1,5 @@ env: - JULIA_VERSION: "1.10.6" + JULIA_VERSION: "1.10.5" JULIA_MINOR_VERSION: "1.10" SVERDRUP_HOME: "/data5/glwagner" TARTARUS_HOME: "/storage5/buildkite-agent" @@ -23,7 +23,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 - label: "🏕️ initialize cpu environment" @@ -44,7 +44,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 ##### @@ -55,7 +55,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "unit" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -63,7 +62,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -71,6 +70,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "unit" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -78,7 +78,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -90,7 +90,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "poisson_solvers_1" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -98,7 +97,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -106,6 +105,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "poisson_solvers_1" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -113,7 +113,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -121,7 +121,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "poisson_solvers_2" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -129,7 +128,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -137,6 +136,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "poisson_solvers_2" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -144,7 +144,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -152,7 +152,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "matrix_poisson_solvers" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -160,7 +159,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -168,6 +167,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "matrix_poisson_solvers" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -175,7 +175,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -183,7 +183,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "general_solvers" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -191,7 +190,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -199,6 +198,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "general_solvers" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -206,7 +206,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -218,7 +218,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "time_stepping_1" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -226,7 +225,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -234,6 +233,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "time_stepping_1" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -241,7 +241,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -253,7 +253,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "time_stepping_2" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -261,7 +260,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -269,6 +268,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "time_stepping_2" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -276,7 +276,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -288,7 +288,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "time_stepping_3" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -296,7 +295,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -304,6 +303,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "time_stepping_3" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -311,7 +311,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -323,7 +323,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "turbulence_closures" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -331,7 +330,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -339,6 +338,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "turbulence_closures" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -346,7 +346,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -358,7 +358,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "hydrostatic_free_surface" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -366,7 +365,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -374,6 +373,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "hydrostatic_free_surface" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -381,7 +381,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -393,7 +393,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "shallow_water" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -401,7 +400,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -409,6 +408,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "shallow_water" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -416,7 +416,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -428,7 +428,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "simulation" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -436,7 +435,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -444,6 +443,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "simulation" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -451,7 +451,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -463,7 +463,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "lagrangian" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -471,7 +470,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -479,6 +478,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "lagrangian" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -486,7 +486,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -498,7 +498,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "abstract_operations" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -506,7 +505,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -514,6 +513,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "abstract_operations" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -521,7 +521,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -533,7 +533,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "multi_region" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -541,7 +540,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -549,6 +548,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "multi_region" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -556,7 +556,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -568,7 +568,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "nonhydrostatic_regression" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -576,7 +575,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -584,6 +583,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "nonhydrostatic_regression" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -591,7 +591,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -599,7 +599,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "hydrostatic_regression" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -607,7 +606,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -615,6 +614,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "hydrostatic_regression" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -622,7 +622,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -634,7 +634,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "scripts" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -642,7 +641,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -650,6 +649,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "scripts" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -657,7 +657,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" @@ -669,7 +669,6 @@ steps: env: JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "enzyme" - GPU_TEST: "true" commands: - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -677,7 +676,7 @@ steps: architecture: GPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_gpu" @@ -685,6 +684,7 @@ steps: env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" TEST_GROUP: "enzyme" + CUDA_VISIBLE_DEVICES: "-1" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" agents: @@ -692,7 +692,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init_cpu" diff --git a/Dockerfile b/Dockerfile index 2422f99aec..3c088bc329 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,4 +1,4 @@ -FROM julia:1.10.6 +FROM julia:1.10.5 LABEL maintainer="Ali Ramadhan " RUN apt-get update && apt-get install -y hdf5-tools diff --git a/Project.toml b/Project.toml index d430772f25..ea9cb84d72 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Oceananigans" uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09" authors = ["Climate Modeling Alliance and contributors"] -version = "0.94.3" +version = "0.92.1" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" @@ -46,14 +46,14 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" OceananigansEnzymeExt = "Enzyme" [compat] -Adapt = "4.1.1" +Adapt = "3, 4" CUDA = "4.1.1, 5" Crayons = "4" CubedSphere = "0.1, 0.2" Dates = "1.9" Distances = "0.10" DocStringExtensions = "0.8, 0.9" -Enzyme = "0.13.14" +Enzyme = "0.13.3" FFTW = "1" Glob = "1.3" IncompleteLU = "0.2" @@ -76,17 +76,16 @@ Rotations = "1.0" SeawaterPolynomials = "0.3.4" SparseArrays = "1.9" Statistics = "1.9" -StructArrays = "0.4, 0.5, 0.6, 0.7" +StructArrays = "0.4, 0.5, 0.6" julia = "1.9" [extras] -CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2" DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" -MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267" +OrthogonalSphericalShellGrids = "c2be9673-fb75-4747-82dc-aa2bb9f4aed0" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" TimesDates = "bdfc003b-8df8-5c39-adcd-3a9087f5df4a" [targets] -test = ["DataDeps", "SafeTestsets", "Test", "Enzyme", "CUDA_Runtime_jll", "MPIPreferences", "TimesDates"] +test = ["DataDeps", "Enzyme", "OrthogonalSphericalShellGrids", "SafeTestsets", "Test", "TimesDates"] diff --git a/benchmark/benchmark_advection_schemes.jl b/benchmark/benchmark_advection_schemes.jl index 4d902f480b..26232c17a7 100644 --- a/benchmark/benchmark_advection_schemes.jl +++ b/benchmark/benchmark_advection_schemes.jl @@ -44,7 +44,7 @@ if GPU in Architectures end for Arch in Architectures - suite_arch = speedups_suite(suite[@tagged Arch], base_case=(Arch, Centered)) + suite_arch = speedups_suite(suite[@tagged Arch], base_case=(Arch, CenteredSecondOrder)) df_arch = speedups_dataframe(suite_arch, slowdown=true) sort!(df_arch, :Schemes, by=string) benchmarks_pretty_table(df_arch, title="Advection schemes relative performance ($Arch)") diff --git a/docs/Project.toml b/docs/Project.toml index 9ef6261635..a7e0daed23 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -18,7 +18,7 @@ TimesDates = "bdfc003b-8df8-5c39-adcd-3a9087f5df4a" CairoMakie = "0.11, 0.12" Documenter = "1" DocumenterCitations = "1" -JLD2 = "0.4, 0,5" +JLD2 = "0.4" Literate = "2.9" NCDatasets = "0.12.10, 0.13.1, 0.14" Polynomials = "4" diff --git a/docs/oceananigans.bib b/docs/oceananigans.bib index 1cdc5da818..714804fa04 100644 --- a/docs/oceananigans.bib +++ b/docs/oceananigans.bib @@ -967,25 +967,4 @@ @article{Shchepetkin2005 pages={347--404}, year={2005}, doi={10.1016/j.ocemod.2004.08.002} -} - -@article{BouZeid05, - author = {Bou-Zeid, Elie and Meneveau, Charles and Parlange, Marc}, - title = "{A scale-dependent Lagrangian dynamic model for large eddy simulation of complex turbulent flows}", - journal = {Physics of Fluids}, - volume = {17}, - number = {2}, - pages = {025105}, - year = {2005}, - doi = {10.1063/1.1839152}, -} - -@article{Lan2022, - author = {Lan, Rihui and Ju, Lili and Wanh, Zhu and Gunzburger, Max and Jones, Philip}, - title = {High-order multirate explicit time-stepping schemes for the baroclinic-barotropic split dynamics in primitive equations}, - journal = {Journal of Computational Physics}, - volume = {457}, - pages = {111050}, - year = {2022}, - doi = {https://doi.org/10.1016/j.jcp.2022.111050}, } \ No newline at end of file diff --git a/docs/src/appendix/library.md b/docs/src/appendix/library.md index ce961093ed..8a2119bd03 100644 --- a/docs/src/appendix/library.md +++ b/docs/src/appendix/library.md @@ -122,13 +122,6 @@ Modules = [Oceananigans.Models.HydrostaticFreeSurfaceModels] Private = false ``` -### Split-explicit free-surface - -```@autodocs -Modules = [Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces] -Private = false -``` - ### Shallow-water models ```@autodocs diff --git a/docs/src/fields.md b/docs/src/fields.md index e37215ee48..6d39014d78 100644 --- a/docs/src/fields.md +++ b/docs/src/fields.md @@ -3,7 +3,7 @@ `Field`s and its relatives are core Oceananigans data structures. `Field`s are more or less arrays of `data` located on a `grid`, whose entries correspond to the average value of some quantity over some finite-sized volume. -`Field`s also may contain `boundary_conditions`, may be computed from an `operand` +`Field`s also may contain `boundary_conditions`, may be computed from an `operand` or expression involving other fields, and may cover only a portion of the total `indices` spanned by the grid. @@ -43,10 +43,10 @@ primary mesh. For example, the primary or `Center` cell spacings in ``z`` are ```jldoctest fields -zspacings(grid, Center())[:, :, 1:4] +zspacings(grid, Center()) # output -4-element Vector{Float64}: +4-element view(OffsetArray(::Vector{Float64}, 0:5), 1:4) with eltype Float64: 0.1 0.19999999999999998 0.3 @@ -57,10 +57,10 @@ corresponding to cell interfaces located at `z = [0, 0.1, 0.3, 0.6, 1]`. But then for the grid which is staggered in `z` relative to the primary mesh, ```jldoctest fields -zspacings(grid, Face())[:, :, 1:5] +zspacings(grid, Face()) # output -5-element Vector{Float64}: +5-element view(OffsetArray(::Vector{Float64}, -1:5), 1:5) with eltype Float64: 0.1 0.15000000000000002 0.24999999999999994 @@ -277,7 +277,7 @@ Note that indexing into `c` is the same as indexing into `c.data`. ```jldoctest fields c[:, :, :] == c.data - + # output true ``` @@ -333,7 +333,7 @@ non-`Flat` directions must be included. For example, to `set!` on a one-dimensio one_d_grid = RectilinearGrid(size=7, x=(0, 7), topology=(Periodic, Flat, Flat)) one_d_c = CenterField(one_d_grid) -# The one-dimensional grid varies only in `x` +# The one-dimensional grid varies only in `x` still_pretty_fun(x) = 3x set!(one_d_c, still_pretty_fun) diff --git a/docs/src/grids.md b/docs/src/grids.md index 29c111ba63..71ec02adc0 100644 --- a/docs/src/grids.md +++ b/docs/src/grids.md @@ -174,7 +174,7 @@ mountain_grid = ImmersedBoundaryGrid(grid, GridFittedBottom(mountain)) # output 20×20×20 ImmersedBoundaryGrid{Float64, Bounded, Bounded, Bounded} on CPU with 3×3×3 halo: -├── immersed_boundary: GridFittedBottom(mean(z)=4.5, min(z)=0.0, max(z)=100.0) +├── immersed_boundary: GridFittedBottom(mean(z)=6.28318, min(z)=1.58939e-8, max(z)=93.9413) ├── underlying_grid: 20×20×20 RectilinearGrid{Float64, Bounded, Bounded, Bounded} on CPU with 3×3×3 halo ├── Bounded x ∈ [-5000.0, 5000.0] regularly spaced with Δx=500.0 ├── Bounded y ∈ [-5000.0, 5000.0] regularly spaced with Δy=500.0 @@ -219,7 +219,7 @@ current_figure() ## Once more with feeling -In summary, making a grid requires +In summary, making a grid requires * The machine architecture, or whether data is stored on the CPU, GPU, or distributed across multiple devices or nodes. * Information about the domain geometry. Domains can take a variety of shapes, including @@ -413,7 +413,7 @@ hidespines!(axy) axΔy = Axis(fig[2, 1]; xlabel = "y (m)", ylabel = "y-spacing (m)") scatter!(axΔy, yc, Δy) -hidespines!(axΔy, :t, :r) +hidespines!(axΔy, :t, :r) axz = Axis(fig[3, 1], title="z-grid") lines!(axz, [-Lz, 0], [0, 0], color=:gray) @@ -440,17 +440,18 @@ using Oceananigans ```@example latlon_nodes grid = LatitudeLongitudeGrid(size = (1, 44), - longitude = (0, 1), + longitude = (0, 1), latitude = (0, 88), topology = (Bounded, Bounded, Flat)) +φ = φnodes(grid, Center()) Δx = xspacings(grid, Center(), Center()) using CairoMakie fig = Figure(size=(600, 400)) ax = Axis(fig[1, 1], xlabel="Zonal spacing on 2 degree grid (km)", ylabel="Latitude (degrees)") -scatter!(ax, Δx / 1e3) +scatter!(ax, Δx ./ 1e3, φ) current_figure() ``` @@ -472,7 +473,7 @@ m = 2 # spacing at the equator in degrees function latitude_faces(j) if j == 1 # equator return 0 - else # crudely estimate the location of the jth face + else # crudely estimate the location of the jth face φ₋ = latitude_faces(j-1) φ′ = φ₋ + m * scale_factor(φ₋) / 2 return φ₋ + m * scale_factor(φ′) @@ -492,7 +493,7 @@ grid = LatitudeLongitudeGrid(size = (Nx, Ny), 180×28×1 LatitudeLongitudeGrid{Float64, Bounded, Bounded, Flat} on CPU with 3×3×0 halo and with precomputed metrics ├── longitude: Bounded λ ∈ [0.0, 360.0] regularly spaced with Δλ=2.0 ├── latitude: Bounded φ ∈ [0.0, 77.2679] variably spaced with min(Δφ)=2.0003, max(Δφ)=6.95319 -└── z: Flat z +└── z: Flat z ``` We've also illustrated the construction of a grid that is `Flat` in the vertical direction. @@ -507,7 +508,7 @@ m = 2 # spacing at the equator in degrees function latitude_faces(j) if j == 1 # equator return 0 - else # crudely estimate the location of the jth face + else # crudely estimate the location of the jth face φ₋ = latitude_faces(j-1) φ′ = φ₋ + m * scale_factor(φ₋) / 2 return φ₋ + m * scale_factor(φ′) @@ -529,17 +530,17 @@ grid = LatitudeLongitudeGrid(size = (Nx, Ny), ```@example plot φ = φnodes(grid, Center()) -Δx = xspacings(grid, Center(), Center())[1, 1:Ny] -Δy = yspacings(grid, Center(), Center())[1, 1:Ny] +Δx = xspacings(grid, Center(), Center(), with_halos=true)[1:Ny] +Δy = yspacings(grid, Center())[1:Ny] using CairoMakie fig = Figure(size=(800, 400), title="Spacings on a Mercator grid") axx = Axis(fig[1, 1], xlabel="Zonal spacing (km)", ylabel="Latitude (degrees)") -scatter!(axx, Δx / 1e3, φ) +scatter!(axx, Δx ./ 1e3, φ) axy = Axis(fig[1, 2], xlabel="Meridional spacing (km)") -scatter!(axy, Δy / 1e3, φ) +scatter!(axy, Δy ./ 1e3, φ) hidespines!(axx, :t, :r) hidespines!(axy, :t, :l, :r) diff --git a/docs/src/index.md b/docs/src/index.md index 628b4eebf1..2bbe707aa4 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -79,10 +79,6 @@ cite our work and mention Oceananigans by name. If you have work using Oceananigans that you would like to have listed here, please open a pull request to add it or let us know! -1. Bisits, J. I., Zika, J. D., and Evans, D. G. (2024) [Does cabbeling shape the thermohaline structure of high-latitude oceans?](https://doi.org/10.1175/JPO-D-24-0061.1) _Journal of Physical Oceanography_, in press. DOI: [10.1175/JPO-D-24-0061.1](https://doi.org/10.1175/JPO-D-24-0061.1) - -1. Strong-Wright J. and Taylor, J. R. (2024) [A model of tidal flow and tracer release in a giant kelp forest](https://doi.org/10.1017/flo.2024.13), _Flow_, **4**, E21. DOI: [10.1017/flo.2024.13](https://doi.org/10.1017/flo.2024.13) - 1. Chor, T. and Wenegrat, J. (2024). [The turbulent dynamics of anticyclonic submesoscale headland wakes](https://doi.org/10.31223/X5570C), _Earth arXiv_, DOI: [10.31223/X5570C](https://doi.org/10.31223/X5570C) 1. Wagner, G. L., Hillier, A., Constantinou, N. C., Silvestri, S., Souza, A., Burns, K., Hill, C., Campin, J.-M., Marshall, J., and Ferrari, R. (2024). [Formulation and calibration of CATKE, a one-equation parameterization for microscale ocean mixing](https://doi.org/10.48550/arXiv.2306.13204), _arXiv_, arXiv:2306.13204. DOI: [10.48550/arXiv.2306.13204](https://doi.org/10.48550/arXiv.2306.13204) diff --git a/docs/src/model_setup/boundary_conditions.md b/docs/src/model_setup/boundary_conditions.md index 2a3ebc60d5..a8c0656012 100644 --- a/docs/src/model_setup/boundary_conditions.md +++ b/docs/src/model_setup/boundary_conditions.md @@ -43,7 +43,7 @@ julia> model = NonhydrostaticModel(; grid, boundary_conditions=(u=no_slip_field_ NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0) ├── grid: 16×16×16 RectilinearGrid{Float64, Periodic, Bounded, Bounded} on CPU with 3×3×3 halo ├── timestepper: RungeKutta3TimeStepper -├── advection scheme: Centered(order=2) +├── advection scheme: Centered reconstruction order 2 ├── tracers: () ├── closure: Nothing ├── buoyancy: Nothing @@ -405,7 +405,7 @@ julia> model = NonhydrostaticModel(grid=grid, boundary_conditions=(u=u_bcs, c=c_ NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0) ├── grid: 16×16×16 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo ├── timestepper: RungeKutta3TimeStepper -├── advection scheme: Centered(order=2) +├── advection scheme: Centered reconstruction order 2 ├── tracers: c ├── closure: Nothing ├── buoyancy: Nothing diff --git a/docs/src/model_setup/buoyancy_and_equation_of_state.md b/docs/src/model_setup/buoyancy_and_equation_of_state.md index d21df10c0f..42b670078e 100644 --- a/docs/src/model_setup/buoyancy_and_equation_of_state.md +++ b/docs/src/model_setup/buoyancy_and_equation_of_state.md @@ -29,7 +29,7 @@ julia> model = NonhydrostaticModel(; grid, buoyancy=nothing) NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0) ├── grid: 8×8×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo ├── timestepper: RungeKutta3TimeStepper -├── advection scheme: Centered(order=2) +├── advection scheme: Centered reconstruction order 2 ├── tracers: () ├── closure: Nothing ├── buoyancy: Nothing @@ -44,7 +44,7 @@ julia> model = NonhydrostaticModel(; grid) NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0) ├── grid: 8×8×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo ├── timestepper: RungeKutta3TimeStepper -├── advection scheme: Centered(order=2) +├── advection scheme: Centered reconstruction order 2 ├── tracers: () ├── closure: Nothing ├── buoyancy: Nothing @@ -78,7 +78,7 @@ julia> model = NonhydrostaticModel(; grid, buoyancy=BuoyancyTracer(), tracers=:b NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0) ├── grid: 8×8×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo ├── timestepper: RungeKutta3TimeStepper -├── advection scheme: Centered(order=2) +├── advection scheme: Centered reconstruction order 2 ├── tracers: b ├── closure: Nothing ├── buoyancy: BuoyancyTracer with ĝ = NegativeZDirection() @@ -99,7 +99,7 @@ HydrostaticFreeSurfaceModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = │ └── solver: FFTImplicitFreeSurfaceSolver ├── advection scheme: │ ├── momentum: Vector Invariant, Dimension-by-dimension reconstruction -│ └── b: Centered(order=2) +│ └── b: Centered reconstruction order 2 └── coriolis: Nothing ``` @@ -119,7 +119,7 @@ julia> model = NonhydrostaticModel(; grid, buoyancy=SeawaterBuoyancy(), tracers= NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0) ├── grid: 8×8×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo ├── timestepper: RungeKutta3TimeStepper -├── advection scheme: Centered(order=2) +├── advection scheme: Centered reconstruction order 2 ├── tracers: (T, S) ├── closure: Nothing ├── buoyancy: SeawaterBuoyancy with g=9.80665 and LinearEquationOfState(thermal_expansion=0.000167, haline_contraction=0.00078) with ĝ = NegativeZDirection() @@ -140,8 +140,8 @@ HydrostaticFreeSurfaceModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = │ └── solver: FFTImplicitFreeSurfaceSolver ├── advection scheme: │ ├── momentum: Vector Invariant, Dimension-by-dimension reconstruction -│ ├── T: Centered(order=2) -│ └── S: Centered(order=2) +│ ├── T: Centered reconstruction order 2 +│ └── S: Centered reconstruction order 2 └── coriolis: Nothing ``` @@ -158,7 +158,7 @@ julia> model = NonhydrostaticModel(; grid, buoyancy, tracers=(:T, :S)) NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0) ├── grid: 8×8×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo ├── timestepper: RungeKutta3TimeStepper -├── advection scheme: Centered(order=2) +├── advection scheme: Centered reconstruction order 2 ├── tracers: (T, S) ├── closure: Nothing ├── buoyancy: SeawaterBuoyancy with g=1.3 and LinearEquationOfState(thermal_expansion=0.000167, haline_contraction=0.00078) with ĝ = NegativeZDirection() @@ -239,7 +239,7 @@ julia> model = NonhydrostaticModel(; grid, buoyancy, tracers=:b) NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0) ├── grid: 8×8×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo ├── timestepper: RungeKutta3TimeStepper -├── advection scheme: Centered(order=2) +├── advection scheme: Centered reconstruction order 2 ├── tracers: b ├── closure: Nothing ├── buoyancy: BuoyancyTracer with ĝ = (0.0, 0.707107, 0.707107) diff --git a/docs/src/model_setup/callbacks.md b/docs/src/model_setup/callbacks.md index 23e21fc278..b60d3cb703 100644 --- a/docs/src/model_setup/callbacks.md +++ b/docs/src/model_setup/callbacks.md @@ -51,7 +51,6 @@ the `:u` field using parameters: ```@example checkpointing using Oceananigans -using Oceananigans: TendencyCallsite model = NonhydrostaticModel(grid=RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))) diff --git a/docs/src/model_setup/lagrangian_particles.md b/docs/src/model_setup/lagrangian_particles.md index be10c534b5..f7deb6b515 100644 --- a/docs/src/model_setup/lagrangian_particles.md +++ b/docs/src/model_setup/lagrangian_particles.md @@ -15,15 +15,15 @@ end ``` ```jldoctest particles -grid = RectilinearGrid(size=(10, 10, 10), extent=(1, 1, 1)) +grid = RectilinearGrid(size=(10, 10, 10), extent=(1, 1, 1)); -Nparticles = 10 +n_particles = 10; -x₀ = zeros(Nparticles) +x₀ = zeros(n_particles); -y₀ = rand(Nparticles) +y₀ = rand(n_particles); -z₀ = -0.5 * ones(Nparticles) +z₀ = -0.5 * ones(n_particles); lagrangian_particles = LagrangianParticles(x=x₀, y=y₀, z=z₀) @@ -44,7 +44,7 @@ model = NonhydrostaticModel(grid=grid, particles=lagrangian_particles) NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0) ├── grid: 10×10×10 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo ├── timestepper: RungeKutta3TimeStepper -├── advection scheme: Centered(order=2) +├── advection scheme: Centered reconstruction order 2 ├── tracers: () ├── closure: Nothing ├── buoyancy: Nothing @@ -73,13 +73,13 @@ struct LagrangianMicrobe{T, S, D} dna :: D end -Nparticles = 3 +n_particles = 3; -x₀ = zeros(Nparticles) +x₀ = zeros(n_particles); -y₀ = rand(Nparticles) +y₀ = rand(n_particles); -z₀ = -0.5 * ones(Nparticles) +z₀ = -0.5 * ones(n_particles); species = [:rock, :paper, :scissors] @@ -108,28 +108,43 @@ Particle properties can be written to disk using JLD2 or NetCDF. When writing to JLD2 you can pass `model.particles` as part of the named tuple of outputs. -```@setup particles -using Oceananigans -grid = RectilinearGrid(size=(10, 10, 10), extent=(1, 1, 1)) -Nparticles = 3 -x₀ = zeros(Nparticles) -y₀ = rand(Nparticles) -z₀ = -0.5 * ones(Nparticles) -lagrangian_particles = LagrangianParticles(x=x₀, y=y₀, z=z₀) -model = NonhydrostaticModel(; grid, particles=lagrangian_particles) +```@meta +DocTestFilters = r"└── file size: [0-9]*.[0-9]* KiB" ``` -```@example particles -JLD2OutputWriter(model, (; particles=model.particles), filename="particles", schedule=TimeInterval(15)) +```jldoctest particles +JLD2OutputWriter(model, (particles=model.particles,), filename="particles", schedule=TimeInterval(15)) + +# output +JLD2OutputWriter scheduled on TimeInterval(15 seconds): +├── filepath: ./particles.jld2 +├── 1 outputs: particles +├── array type: Array{Float64} +├── including: [:grid, :coriolis, :buoyancy, :closure] +├── file_splitting: NoFileSplitting +└── file size: 17.6 KiB ``` When writing to NetCDF you should write particles to a separate file as the NetCDF dimensions differ for particle trajectories. You can just pass `model.particles` straight to `NetCDFOutputWriter`: -```@example particles +```jldoctest particles NetCDFOutputWriter(model, model.particles, filename="particles.nc", schedule=TimeInterval(15)) + +# output +NetCDFOutputWriter scheduled on TimeInterval(15 seconds): +├── filepath: ./particles.nc +├── dimensions: particle_id(10), time(0) +├── 1 outputs: particles +└── array type: Array{Float64} +├── file_splitting: NoFileSplitting +└── file size: 9.9 KiB ``` !!! warn "Outputting custom particle properties to NetCDF" NetCDF does not support arbitrary data types. If you need to write custom particle properties to disk that are not supported by NetCDF then you should use JLD2 (which should support almost any Julia data type). + +```@meta +DocTestFilters = nothing +``` diff --git a/docs/src/model_setup/output_writers.md b/docs/src/model_setup/output_writers.md index 2d7698c5a7..999eaaaef1 100644 --- a/docs/src/model_setup/output_writers.md +++ b/docs/src/model_setup/output_writers.md @@ -50,7 +50,7 @@ passing it a dictionary of (label, field) pairs and any indices for slicing if y Saving the u velocity field and temperature fields as full 3D fields, surface 2D slices, and 1D columns to separate NetCDF files: -```@example netcdf1 +```jldoctest netcdf1 using Oceananigans grid = RectilinearGrid(size=(16, 16, 16), extent=(1, 1, 1)) @@ -63,26 +63,53 @@ fields = Dict("u" => model.velocities.u, "c" => model.tracers.c) simulation.output_writers[:field_writer] = NetCDFOutputWriter(model, fields, filename="more_fields.nc", schedule=TimeInterval(60)) + +# output +NetCDFOutputWriter scheduled on TimeInterval(1 minute): +├── filepath: ./more_fields.nc +├── dimensions: zC(16), zF(17), xC(16), yF(16), xF(16), yC(16), time(0) +├── 2 outputs: (c, u) +└── array type: Array{Float64} +├── file_splitting: NoFileSplitting +└── file size: 14.9 KiB ``` -```@example netcdf1 +```jldoctest netcdf1 simulation.output_writers[:surface_slice_writer] = NetCDFOutputWriter(model, fields, filename="another_surface_xy_slice.nc", schedule=TimeInterval(60), indices=(:, :, grid.Nz)) + +# output +NetCDFOutputWriter scheduled on TimeInterval(1 minute): +├── filepath: ./another_surface_xy_slice.nc +├── dimensions: zC(1), zF(1), xC(16), yF(16), xF(16), yC(16), time(0) +├── 2 outputs: (c, u) +└── array type: Array{Float64} +├── file_splitting: NoFileSplitting +└── file size: 14.9 KiB ``` -```@example netcdf1 +```jldoctest netcdf1 simulation.output_writers[:averaged_profile_writer] = NetCDFOutputWriter(model, fields, filename = "another_averaged_z_profile.nc", schedule = AveragedTimeInterval(60, window=20), indices = (1, 1, :)) + +# output +NetCDFOutputWriter scheduled on TimeInterval(1 minute): +├── filepath: ./another_averaged_z_profile.nc +├── dimensions: zC(16), zF(17), xC(1), yF(1), xF(1), yC(1), time(0) +├── 2 outputs: (c, u) averaged on AveragedTimeInterval(window=20 seconds, stride=1, interval=1 minute) +└── array type: Array{Float64} +├── file_splitting: NoFileSplitting +└── file size: 17.6 KiB ``` `NetCDFOutputWriter` also accepts output functions that write scalars and arrays to disk, provided that their `dimensions` are provided: -```@example +```jldoctest using Oceananigans Nx, Ny, Nz = 16, 16, 16 @@ -120,13 +147,22 @@ simulation.output_writers[:things] = NetCDFOutputWriter(model, outputs, schedule=IterationInterval(1), filename="things.nc", dimensions=dims, verbose=true, global_attributes=global_attributes, output_attributes=output_attributes) + +# output +NetCDFOutputWriter scheduled on IterationInterval(1): +├── filepath: ./things.nc +├── dimensions: zC(16), zF(17), xC(16), yF(16), xF(16), yC(16), time(0) +├── 3 outputs: (profile, slice, scalar) +└── array type: Array{Float64} +├── file_splitting: NoFileSplitting +└── file size: 17.8 KiB ``` `NetCDFOutputWriter` can also be configured for `outputs` that are interpolated or regridded to a different grid than `model.grid`. To use this functionality, include the keyword argument `grid = output_grid`. -```@example +```jldoctest using Oceananigans using Oceananigans.Fields: interpolate! @@ -143,6 +179,15 @@ output_writer = NetCDFOutputWriter(model, outputs; grid = coarse_grid, filename = "coarse_u.nc", schedule = IterationInterval(1)) + +# output +NetCDFOutputWriter scheduled on IterationInterval(1): +├── filepath: ./coarse_u.nc +├── dimensions: zC(4), zF(5), xC(1), yF(1), xF(1), yC(1), time(0) +├── 1 outputs: u +└── array type: Array{Float64} +├── file_splitting: NoFileSplitting +└── file size: 14.6 KiB ``` See [`NetCDFOutputWriter`](@ref) for more information. @@ -163,7 +208,7 @@ of the function will be saved to the JLD2 file. Write out 3D fields for u, v, w, and a tracer c, along with a horizontal average: -```@example jld2_output_writer +```jldoctest jld2_output_writer using Oceananigans using Oceananigans.Utils: hour, minute @@ -184,17 +229,36 @@ simulation.output_writers[:velocities] = JLD2OutputWriter(model, model.velocitie filename = "some_more_data.jld2", schedule = TimeInterval(20minute), init = init_save_some_metadata!) + +# output +JLD2OutputWriter scheduled on TimeInterval(20 minutes): +├── filepath: ./some_more_data.jld2 +├── 3 outputs: (u, v, w) +├── array type: Array{Float64} +├── including: [:grid, :coriolis, :buoyancy, :closure] +├── file_splitting: NoFileSplitting +└── file size: 28.5 KiB ``` and a time- and horizontal-average of tracer `c` every 20 minutes of simulation time to a file called `some_more_averaged_data.jld2` -```@example jld2_output_writer +```jldoctest jld2_output_writer simulation.output_writers[:avg_c] = JLD2OutputWriter(model, (; c=c_avg), filename = "some_more_averaged_data.jld2", schedule = AveragedTimeInterval(20minute, window=5minute)) + +# output +JLD2OutputWriter scheduled on TimeInterval(20 minutes): +├── filepath: ./some_more_averaged_data.jld2 +├── 1 outputs: c averaged on AveragedTimeInterval(window=5 minutes, stride=1, interval=20 minutes) +├── array type: Array{Float64} +├── including: [:grid, :coriolis, :buoyancy, :closure] +├── file_splitting: NoFileSplitting +└── file size: 18.3 KiB ``` + See [`JLD2OutputWriter`](@ref) for more information. ## Time-averaged output @@ -229,7 +293,7 @@ AveragedTimeInterval(window=1 day, stride=1, interval=4 days) An `AveragedTimeInterval` schedule directs an output writer to time-average its outputs before writing them to disk: -```@example averaged_time_interval +```jldoctest averaged_time_interval using Oceananigans using Oceananigans.Units @@ -240,4 +304,13 @@ simulation = Simulation(model, Δt=10minutes, stop_time=30days) simulation.output_writers[:velocities] = JLD2OutputWriter(model, model.velocities, filename = "even_more_averaged_velocity_data.jld2", schedule = AveragedTimeInterval(4days, window=1day, stride=2)) + +# output +JLD2OutputWriter scheduled on TimeInterval(4 days): +├── filepath: ./even_more_averaged_velocity_data.jld2 +├── 3 outputs: (u, v, w) averaged on AveragedTimeInterval(window=1 day, stride=2, interval=4 days) +├── array type: Array{Float64} +├── including: [:grid, :coriolis, :buoyancy, :closure] +├── file_splitting: NoFileSplitting +└── file size: 27.6 KiB ``` diff --git a/docs/src/model_setup/tracers.md b/docs/src/model_setup/tracers.md index 2aa1668352..947cb6e9ff 100644 --- a/docs/src/model_setup/tracers.md +++ b/docs/src/model_setup/tracers.md @@ -16,7 +16,7 @@ julia> model = NonhydrostaticModel(; grid) NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0) ├── grid: 16×16×16 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo ├── timestepper: RungeKutta3TimeStepper -├── advection scheme: Centered(order=2) +├── advection scheme: Centered reconstruction order 2 ├── tracers: () ├── closure: Nothing ├── buoyancy: Nothing @@ -31,7 +31,7 @@ julia> model = NonhydrostaticModel(; grid, tracers=(:T, :S)) NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0) ├── grid: 16×16×16 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo ├── timestepper: RungeKutta3TimeStepper -├── advection scheme: Centered(order=2) +├── advection scheme: Centered reconstruction order 2 ├── tracers: (T, S) ├── closure: Nothing ├── buoyancy: Nothing @@ -66,7 +66,7 @@ julia> model = NonhydrostaticModel(; grid, tracers=(:T, :S, :C₁, :CO₂, :nitr NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0) ├── grid: 16×16×16 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo ├── timestepper: RungeKutta3TimeStepper -├── advection scheme: Centered(order=2) +├── advection scheme: Centered reconstruction order 2 ├── tracers: (T, S, C₁, CO₂, nitrogen) ├── closure: Nothing ├── buoyancy: Nothing diff --git a/docs/src/model_setup/turbulent_diffusivity_closures_and_les_models.md b/docs/src/model_setup/turbulent_diffusivity_closures_and_les_models.md index fd8e7b8530..175a3856c2 100644 --- a/docs/src/model_setup/turbulent_diffusivity_closures_and_les_models.md +++ b/docs/src/model_setup/turbulent_diffusivity_closures_and_les_models.md @@ -65,9 +65,7 @@ To use the default Smagorinsky-Lilly LES closure, we write julia> using Oceananigans.TurbulenceClosures julia> closure = SmagorinskyLilly() -Smagorinsky closure with -├── coefficient = LillyCoefficient(smagorinsky = 0.16, reduction_factor = 1.0) -└── Pr = 1.0 +SmagorinskyLilly: C=0.16, Cb=1.0, Pr=1.0 ``` The parameters `C`, `Cb`, and `Pr` may alternatively be specified explicitly. diff --git a/docs/src/numerical_implementation/boundary_conditions.md b/docs/src/numerical_implementation/boundary_conditions.md index d582ee1206..ef5b6d08f4 100644 --- a/docs/src/numerical_implementation/boundary_conditions.md +++ b/docs/src/numerical_implementation/boundary_conditions.md @@ -211,14 +211,16 @@ point to ``` but we then pressure correct the interior so a new ``\mathcal{O}(\Delta t)`` error is introduced as ```math + \begin{equation} \begin{align} u^{n+1}_{1jk} &\approx u^{n+1}_{3jk} + (u^{n+1}_{2jk} - u^{n+1}_{jk4}) / 2 + \mathcal{O}(\Delta x^2),\\ &= u^\star_{1jk} - \Delta t \left(\boldsymbol{\nabla}p^{n+1}_{3jk} + (\boldsymbol{\nabla}p^{n+1}_{2jk} - \boldsymbol{\nabla}p^{n+1}_{4jk}) / 2\right) + \mathcal{O}(\Delta x^2),\\ &\approx u^\star_{1jk} + \mathcal{O}(\Delta x^2) + \mathcal{O}(\Delta t). \end{align} + \end{equation} ``` This is prefered to a divergent interior solution as open boundary conditions (except no penetration) are typlically already unphysical and only used in an attempt to allow information to enter or exit the domain. -Open boundary conditions are represented by the [`Open`](@ref) type. +Open boundary conditions are represented by the [`Open`](@ref) type. \ No newline at end of file diff --git a/docs/src/operations.md b/docs/src/operations.md index 04cbeb35e6..e50e34f1e9 100644 --- a/docs/src/operations.md +++ b/docs/src/operations.md @@ -132,8 +132,6 @@ quadratic[1, 1, 1] = 1.0431614901668005 Also like `Field`s, `AbstractOperation`s have a _location_. For example, ```jldoctest operations -using Oceananigans: location - @show location(c) @show location(dx_c) nothing diff --git a/examples/convecting_plankton.jl b/examples/convecting_plankton.jl index 87335b55d8..d7aa86fb8c 100644 --- a/examples/convecting_plankton.jl +++ b/examples/convecting_plankton.jl @@ -128,7 +128,8 @@ plankton_dynamics = Forcing(growing_and_grazing, field_dependencies = :P, # and Coriolis forces appropriate for planktonic convection at mid-latitudes on Earth. model = NonhydrostaticModel(; grid, - advection = UpwindBiased(order=5), + advection = UpwindBiasedFifthOrder(), + timestepper = :RungeKutta3, closure = ScalarDiffusivity(ν=1e-4, κ=1e-4), coriolis = FPlane(f=1e-4), tracers = (:b, :P), # P for Plankton diff --git a/examples/horizontal_convection.jl b/examples/horizontal_convection.jl index 6f48795792..44fa24b6a3 100644 --- a/examples/horizontal_convection.jl +++ b/examples/horizontal_convection.jl @@ -179,6 +179,10 @@ b_timeseries = FieldTimeSeries(saved_output_filename, "b") ζ_timeseries = FieldTimeSeries(saved_output_filename, "ζ") times = b_timeseries.times + +## Coordinate arrays +xc, yc, zc = nodes(b_timeseries[1]) +xζ, yζ, zζ = nodes(ζ_timeseries[1]) nothing #hide χ_timeseries = deepcopy(b_timeseries) @@ -325,3 +329,4 @@ lines!(ax_Nu, t, Nu; linewidth = 3) current_figure() #hide fig + diff --git a/examples/internal_tide.jl b/examples/internal_tide.jl index 08d2ef4db3..5b31f18418 100644 --- a/examples/internal_tide.jl +++ b/examples/internal_tide.jl @@ -14,6 +14,7 @@ using Oceananigans using Oceananigans.Units +using Oceananigans.ImmersedBoundaries: PartialCellBottom # ## Grid @@ -205,6 +206,20 @@ wmax = maximum(abs, w_t[end]) times = u′_t.times nothing #hide +# We retrieve each field's coordinates and convert from meters to kilometers. + +xu, _, zu = nodes(u′_t[1]) +xw, _, zw = nodes(w_t[1]) +xN², _, zN² = nodes(N²_t[1]) + +xu = xu ./ 1e3 +xw = xw ./ 1e3 +xN² = xN² ./ 1e3 +zu = zu ./ 1e3 +zw = zw ./ 1e3 +zN² = zN² ./ 1e3 +nothing #hide + # ## Visualize # Now we can visualize our resutls! We use `CairoMakie` here. On a system with OpenGL @@ -220,13 +235,13 @@ n = Observable(1) title = @lift @sprintf("t = %1.2f days = %1.2f T₂", round(times[$n] / day, digits=2) , round(times[$n] / T₂, digits=2)) -u′ₙ = @lift u′_t[$n] - wₙ = @lift w_t[$n] -N²ₙ = @lift N²_t[$n] +u′n = @lift u′_t[$n] + wn = @lift w_t[$n] +N²n = @lift N²_t[$n] -axis_kwargs = (xlabel = "x [m]", - ylabel = "z [m]", - limits = ((-grid.Lx/2, grid.Lx/2), (-grid.Lz, 0)), +axis_kwargs = (xlabel = "x [km]", + ylabel = "z [km]", + limits = ((-grid.Lx/2e3, grid.Lx/2e3), (-grid.Lz/1e3, 0)), # note conversion to kilometers titlesize = 20) fig = Figure(size = (700, 900)) @@ -234,15 +249,15 @@ fig = Figure(size = (700, 900)) fig[1, :] = Label(fig, title, fontsize=24, tellwidth=false) ax_u = Axis(fig[2, 1]; title = "u'-velocity", axis_kwargs...) -hm_u = heatmap!(ax_u, u′ₙ; nan_color=:gray, colorrange=(-umax, umax), colormap=:balance) +hm_u = heatmap!(ax_u, xu, zu, u′n; nan_color=:gray, colorrange=(-umax, umax), colormap=:balance) Colorbar(fig[2, 2], hm_u, label = "m s⁻¹") ax_w = Axis(fig[3, 1]; title = "w-velocity", axis_kwargs...) -hm_w = heatmap!(ax_w, wₙ; nan_color=:gray, colorrange=(-wmax, wmax), colormap=:balance) +hm_w = heatmap!(ax_w, xw, zw, wn; nan_color=:gray, colorrange=(-wmax, wmax), colormap=:balance) Colorbar(fig[3, 2], hm_w, label = "m s⁻¹") ax_N² = Axis(fig[4, 1]; title = "stratification N²", axis_kwargs...) -hm_N² = heatmap!(ax_N², N²ₙ; nan_color=:gray, colorrange=(0.9Nᵢ², 1.1Nᵢ²), colormap=:magma) +hm_N² = heatmap!(ax_N², xN², zN², N²n; nan_color=:gray, colorrange=(0.9Nᵢ², 1.1Nᵢ²), colormap=:magma) Colorbar(fig[4, 2], hm_N², label = "s⁻²") fig diff --git a/examples/internal_wave.jl b/examples/internal_wave.jl index 92e5cc2380..9fb2ed12ae 100644 --- a/examples/internal_wave.jl +++ b/examples/internal_wave.jl @@ -53,7 +53,8 @@ B = BackgroundField(B_func, parameters=N) # `b` that we identify as buoyancy by setting `buoyancy=BuoyancyTracer()`. model = NonhydrostaticModel(; grid, coriolis, - advection = Centered(order=4), + advection = CenteredFourthOrder(), + timestepper = :RungeKutta3, closure = ScalarDiffusivity(ν=1e-6, κ=1e-6), tracers = :b, buoyancy = BuoyancyTracer(), diff --git a/examples/kelvin_helmholtz_instability.jl b/examples/kelvin_helmholtz_instability.jl index ac897b4431..843c23d37f 100644 --- a/examples/kelvin_helmholtz_instability.jl +++ b/examples/kelvin_helmholtz_instability.jl @@ -138,12 +138,14 @@ fig # # The model -model = NonhydrostaticModel(; grid, - advection = UpwindBiased(order=5), - background_fields = (u=U, b=B), - closure = ScalarDiffusivity(ν=2e-4, κ=2e-4), - buoyancy = BuoyancyTracer(), - tracers = :b) +model = NonhydrostaticModel(timestepper = :RungeKutta3, + advection = UpwindBiasedFifthOrder(), + grid = grid, + coriolis = nothing, + background_fields = (u=U, b=B), + closure = ScalarDiffusivity(ν=2e-4, κ=2e-4), + buoyancy = BuoyancyTracer(), + tracers = :b) # We have included a "pinch" of viscosity and diffusivity in anticipation of what will follow furtherdown: # viscosity and diffusivity will ensure numerical stability when we evolve the unstable mode to the point diff --git a/examples/langmuir_turbulence.jl b/examples/langmuir_turbulence.jl index 9886c2fca9..9111d80cd9 100644 --- a/examples/langmuir_turbulence.jl +++ b/examples/langmuir_turbulence.jl @@ -90,18 +90,18 @@ uˢ(z) = Uˢ * exp(z / vertical_scale) # # At the surface ``z = 0``, Wagner et al. (2021) impose -τx = -3.72e-5 # m² s⁻², surface kinematic momentum flux +Qᵘ = -3.72e-5 # m² s⁻², surface kinematic momentum flux -u_boundary_conditions = FieldBoundaryConditions(top = FluxBoundaryCondition(τx)) +u_boundary_conditions = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵘ)) # Wagner et al. (2021) impose a linear buoyancy gradient `N²` at the bottom # along with a weak, destabilizing flux of buoyancy at the surface to faciliate # spin-up from rest. -Jᵇ = 2.307e-8 # m² s⁻³, surface buoyancy flux +Qᵇ = 2.307e-8 # m² s⁻³, surface buoyancy flux N² = 1.936e-5 # s⁻², initial and bottom buoyancy gradient -b_boundary_conditions = FieldBoundaryConditions(top = FluxBoundaryCondition(Jᵇ), +b_boundary_conditions = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵇ), bottom = GradientBoundaryCondition(N²)) # !!! info "The flux convention in Oceananigans" @@ -153,7 +153,7 @@ bᵢ(x, y, z) = stratification(z) + 1e-1 * Ξ(z) * N² * model.grid.Lz # This initial condition is consistent with a wavy, quiescent ocean suddenly impacted # by winds. To this quiescent state we add noise scaled by the friction velocity to ``u`` and ``w``. -u★ = sqrt(abs(τx)) +u★ = sqrt(abs(Qᵘ)) uᵢ(x, y, z) = u★ * 1e-1 * Ξ(z) wᵢ(x, y, z) = u★ * 1e-1 * Ξ(z) @@ -254,6 +254,8 @@ time_series = (; wv = FieldTimeSeries("langmuir_turbulence_averages.jld2", "wv")) times = time_series.w.times +xw, yw, zw = nodes(time_series.w) +xu, yu, zu = nodes(time_series.u) nothing #hide # We are now ready to animate using Makie. We use Makie's `Observable` to animate @@ -307,43 +309,43 @@ nothing #hide wₙ = @lift time_series.w[$n] uₙ = @lift time_series.u[$n] -Bₙ = @lift view(time_series.B[$n], 1, 1, :) -Uₙ = @lift view(time_series.U[$n], 1, 1, :) -Vₙ = @lift view(time_series.V[$n], 1, 1, :) -wuₙ = @lift view(time_series.wu[$n], 1, 1, :) -wvₙ = @lift view(time_series.wv[$n], 1, 1, :) +Bₙ = @lift time_series.B[$n][1, 1, :] +Uₙ = @lift time_series.U[$n][1, 1, :] +Vₙ = @lift time_series.V[$n][1, 1, :] +wuₙ = @lift time_series.wu[$n][1, 1, :] +wvₙ = @lift time_series.wv[$n][1, 1, :] k = searchsortedfirst(grid.zᵃᵃᶠ[:], -8) -wxyₙ = @lift view(time_series.w[$n], :, :, k) -wxzₙ = @lift view(time_series.w[$n], :, 1, :) -uxzₙ = @lift view(time_series.u[$n], :, 1, :) +wxyₙ = @lift interior(time_series.w[$n], :, :, k) +wxzₙ = @lift interior(time_series.w[$n], :, 1, :) +uxzₙ = @lift interior(time_series.u[$n], :, 1, :) wlims = (-0.03, 0.03) ulims = (-0.05, 0.05) -lines!(ax_B, Bₙ) +lines!(ax_B, Bₙ, zu) -lines!(ax_U, Uₙ; label = L"\bar{u}") -lines!(ax_U, Vₙ; label = L"\bar{v}") +lines!(ax_U, Uₙ, zu; label = L"\bar{u}") +lines!(ax_U, Vₙ, zu; label = L"\bar{v}") axislegend(ax_U; position = :rb) -lines!(ax_fluxes, wuₙ; label = L"mean $wu$") -lines!(ax_fluxes, wvₙ; label = L"mean $wv$") +lines!(ax_fluxes, wuₙ, zw; label = L"mean $wu$") +lines!(ax_fluxes, wvₙ, zw; label = L"mean $wv$") axislegend(ax_fluxes; position = :rb) -hm_wxy = heatmap!(ax_wxy, wxyₙ; +hm_wxy = heatmap!(ax_wxy, xw, yw, wxyₙ; colorrange = wlims, colormap = :balance) Colorbar(fig[1, 3], hm_wxy; label = "m s⁻¹") -hm_wxz = heatmap!(ax_wxz, wxzₙ; +hm_wxz = heatmap!(ax_wxz, xw, zw, wxzₙ; colorrange = wlims, colormap = :balance) Colorbar(fig[2, 3], hm_wxz; label = "m s⁻¹") -ax_uxz = heatmap!(ax_uxz, uxzₙ; +ax_uxz = heatmap!(ax_uxz, xu, zu, uxzₙ; colorrange = ulims, colormap = :balance) diff --git a/examples/ocean_wind_mixing_and_convection.jl b/examples/ocean_wind_mixing_and_convection.jl index 6daa4c0dc0..795b2eff6e 100644 --- a/examples/ocean_wind_mixing_and_convection.jl +++ b/examples/ocean_wind_mixing_and_convection.jl @@ -50,13 +50,13 @@ h(k) = (k - 1) / Nz ## Linear near-surface generator ζ₀(k) = 1 + (h(k) - 1) / refinement -## Bottom-intensified stretching function +## Bottom-intensified stretching function Σ(k) = (1 - exp(-stretching * h(k))) / (1 - exp(-stretching)) ## Generating function z_faces(k) = Lz * (ζ₀(k) * Σ(k) - 1) -grid = RectilinearGrid(size = (Nx, Nx, Nz), +grid = RectilinearGrid(size = (Nx, Nx, Nz), x = (0, Lx), y = (0, Ly), z = z_faces) @@ -66,8 +66,8 @@ grid = RectilinearGrid(size = (Nx, Nx, Nz), fig = Figure(size=(1200, 800)) ax = Axis(fig[1, 1], ylabel = "Depth (m)", xlabel = "Vertical spacing (m)") -lines!(ax, zspacings(grid, Center())) -scatter!(ax, zspacings(grid, Center())) +lines!(ax, zspacings(grid, Center()), znodes(grid, Center())) +scatter!(ax, zspacings(grid, Center()), znodes(grid, Center())) current_figure() #hide fig @@ -84,18 +84,18 @@ buoyancy = SeawaterBuoyancy(equation_of_state=LinearEquationOfState(thermal_expa # We calculate the surface temperature flux associated with surface cooling of # 200 W m⁻², reference density `ρₒ`, and heat capacity `cᴾ`, -Q = 200.0 # W m⁻², surface _heat_ flux +Qʰ = 200.0 # W m⁻², surface _heat_ flux ρₒ = 1026.0 # kg m⁻³, average density at the surface of the world ocean cᴾ = 3991.0 # J K⁻¹ kg⁻¹, typical heat capacity for seawater -Jᵀ = Q / (ρₒ * cᴾ) # K m s⁻¹, surface _temperature_ flux +Qᵀ = Qʰ / (ρₒ * cᴾ) # K m s⁻¹, surface _temperature_ flux # Finally, we impose a temperature gradient `dTdz` both initially and at the # bottom of the domain, culminating in the boundary conditions on temperature, dTdz = 0.01 # K m⁻¹ -T_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Jᵀ), +T_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵀ), bottom = GradientBoundaryCondition(dTdz)) # Note that a positive temperature flux at the surface of the ocean @@ -111,26 +111,26 @@ u₁₀ = 10 # m s⁻¹, average wind velocity 10 meters above the ocean cᴰ = 2.5e-3 # dimensionless drag coefficient ρₐ = 1.225 # kg m⁻³, average density of air at sea-level -τx = - ρₐ / ρₒ * cᴰ * u₁₀ * abs(u₁₀) # m² s⁻² +Qᵘ = - ρₐ / ρₒ * cᴰ * u₁₀ * abs(u₁₀) # m² s⁻² # The boundary conditions on `u` are thus -u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(τx)) +u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵘ)) # For salinity, `S`, we impose an evaporative flux of the form -@inline Jˢ(x, y, t, S, evaporation_rate) = - evaporation_rate * S # [salinity unit] m s⁻¹ +@inline Qˢ(x, y, t, S, evaporation_rate) = - evaporation_rate * S # [salinity unit] m s⁻¹ nothing #hide # where `S` is salinity. We use an evporation rate of 1 millimeter per hour, evaporation_rate = 1e-3 / hour # m s⁻¹ -# We build the `Flux` evaporation `BoundaryCondition` with the function `Jˢ`, -# indicating that `Jˢ` depends on salinity `S` and passing +# We build the `Flux` evaporation `BoundaryCondition` with the function `Qˢ`, +# indicating that `Qˢ` depends on salinity `S` and passing # the parameter `evaporation_rate`, -evaporation_bc = FluxBoundaryCondition(Jˢ, field_dependencies=:S, parameters=evaporation_rate) +evaporation_bc = FluxBoundaryCondition(Qˢ, field_dependencies=:S, parameters=evaporation_rate) # The full salinity boundary conditions are @@ -145,7 +145,8 @@ S_bcs = FieldBoundaryConditions(top=evaporation_bc) # scales smaller than the grid scale that we cannot explicitly resolve. model = NonhydrostaticModel(; grid, buoyancy, - advection = UpwindBiased(order=5), + advection = UpwindBiasedFifthOrder(), + timestepper = :RungeKutta3, tracers = (:T, :S), coriolis = FPlane(f=1e-4), closure = AnisotropicMinimumDissipation(), @@ -172,7 +173,7 @@ model = NonhydrostaticModel(; grid, buoyancy, Tᵢ(x, y, z) = 20 + dTdz * z + dTdz * model.grid.Lz * 1e-6 * Ξ(z) ## Velocity initial condition: random noise scaled by the friction velocity. -uᵢ(x, y, z) = sqrt(abs(τx)) * 1e-3 * Ξ(z) +uᵢ(x, y, z) = sqrt(abs(Qᵘ)) * 1e-3 * Ξ(z) ## `set!` the `model` fields using functions or constants: set!(model, u=uᵢ, w=uᵢ, T=Tᵢ, S=35) diff --git a/examples/tilted_bottom_boundary_layer.jl b/examples/tilted_bottom_boundary_layer.jl index c6b0b15eb3..0e1efdb500 100644 --- a/examples/tilted_bottom_boundary_layer.jl +++ b/examples/tilted_bottom_boundary_layer.jl @@ -1,9 +1,11 @@ +# # Tilted bottom boundary layer example +# # This example simulates a two-dimensional oceanic bottom boundary layer # in a domain that's tilted with respect to gravity. We simulate the perturbation # away from a constant along-slope (y-direction) velocity constant density stratification. # This perturbation develops into a turbulent bottom boundary layer due to momentum # loss at the bottom boundary modeled with a quadratic drag law. -# +# # This example illustrates # # * changing the direction of gravitational acceleration in the buoyancy model; @@ -33,7 +35,7 @@ Nz = 64 ## Creates a grid with near-constant spacing `refinement * Lz / Nz` ## near the bottom: refinement = 1.8 # controls spacing near surface (higher means finer spaced) -stretching = 10 # controls rate of stretching at bottom +stretching = 10 # controls rate of stretching at bottom ## "Warped" height coordinate h(k) = (Nz + 1 - k) / Nz @@ -41,7 +43,7 @@ h(k) = (Nz + 1 - k) / Nz ## Linear near-surface generator ζ(k) = 1 + (h(k) - 1) / refinement -## Bottom-intensified stretching function +## Bottom-intensified stretching function Σ(k) = (1 - exp(-stretching * h(k))) / (1 - exp(-stretching)) ## Generating function @@ -56,9 +58,11 @@ grid = RectilinearGrid(topology = (Periodic, Flat, Bounded), using CairoMakie -scatterlines(zspacings(grid, Center()), - axis = (ylabel = "Depth (m)", - xlabel = "Vertical spacing (m)")) +lines(zspacings(grid, Center()), znodes(grid, Center()), + axis = (ylabel = "Depth (m)", + xlabel = "Vertical spacing (m)")) + +scatter!(zspacings(grid, Center()), znodes(grid, Center())) current_figure() #hide @@ -137,16 +141,18 @@ V∞_field = BackgroundField(V∞) # ## Create the `NonhydrostaticModel` # -# We are now ready to create the model. We create a `NonhydrostaticModel` with a -# fifth-order `UpwindBiased` advection scheme and a constant viscosity and diffusivity. -# Here we use a smallish value of ``10^{-4} \, \rm{m}^2\, \rm{s}^{-1}``. +# We are now ready to create the model. We create a `NonhydrostaticModel` with an +# `UpwindBiasedFifthOrder` advection scheme, a `RungeKutta3` timestepper, +# and a constant viscosity and diffusivity. Here we use a smallish value +# of ``10^{-4} \, \rm{m}^2\, \rm{s}^{-1}``. ν = 1e-4 κ = 1e-4 closure = ScalarDiffusivity(ν=ν, κ=κ) model = NonhydrostaticModel(; grid, buoyancy, coriolis, closure, - advection = UpwindBiased(order=5), + timestepper = :RungeKutta3, + advection = UpwindBiasedFifthOrder(), tracers = :b, boundary_conditions = (u=u_bcs, v=v_bcs, b=b_bcs), background_fields = (; b=B∞_field, v=V∞_field)) diff --git a/examples/two_dimensional_turbulence.jl b/examples/two_dimensional_turbulence.jl index 6212bc8b15..e6807b963d 100644 --- a/examples/two_dimensional_turbulence.jl +++ b/examples/two_dimensional_turbulence.jl @@ -26,7 +26,8 @@ using Oceananigans grid = RectilinearGrid(size=(128, 128), extent=(2π, 2π), topology=(Periodic, Periodic, Flat)) model = NonhydrostaticModel(; grid, - advection = UpwindBiased(order=5), + timestepper = :RungeKutta3, + advection = UpwindBiasedFifthOrder(), closure = ScalarDiffusivity(ν=1e-5)) # ## Random initial conditions @@ -125,6 +126,11 @@ run!(simulation) s_timeseries = FieldTimeSeries(filename * ".jld2", "s") times = ω_timeseries.times + +# Construct the ``x, y, z`` grid for plotting purposes, + +xω, yω, zω = nodes(ω_timeseries) +xs, ys, zs = nodes(s_timeseries) nothing #hide # and animate the vorticity and fluid speed. diff --git a/ext/OceananigansMakieExt.jl b/ext/OceananigansMakieExt.jl index 6627ce8491..d214b8281c 100644 --- a/ext/OceananigansMakieExt.jl +++ b/ext/OceananigansMakieExt.jl @@ -2,7 +2,6 @@ module OceananigansMakieExt using Oceananigans using Oceananigans.Grids: OrthogonalSphericalShellGrid -using Oceananigans.Fields: AbstractField using Oceananigans.AbstractOperations: AbstractOperation using Oceananigans.Architectures: on_architecture using Oceananigans.ImmersedBoundaries: mask_immersed_field! @@ -15,7 +14,7 @@ import Makie: args_preferred_axis # do not overstate a preference for being plotted in a 3D LScene. # Because often we are trying to plot 1D and 2D Field, even though # (perhaps incorrectly) all Field are AbstractArray{3}. -args_preferred_axis(::AbstractField) = nothing +args_preferred_axis(::Field) = nothing function drop_singleton_indices(N) if N == 1 @@ -39,7 +38,7 @@ end convert_arguments(pl::Type{<:AbstractPlot}, f::Field) = convert_arguments(pl, convert_field_argument(f)...) -function convert_arguments(pl::Type{<:AbstractPlot}, op::AbstractOperation) +function convert_arguments(pl::Type{<:AbstractPlot}, fop::AbstractOperation) f = Field(op) compute!(f) return convert_arguments(pl, f) @@ -147,4 +146,3 @@ function convert_arguments(pl::Type{<:AbstractPlot}, ξ1::AbstractArray, ξ2::Ab end end # module - diff --git a/src/AbstractOperations/AbstractOperations.jl b/src/AbstractOperations/AbstractOperations.jl index f63fed8ae5..576f633179 100644 --- a/src/AbstractOperations/AbstractOperations.jl +++ b/src/AbstractOperations/AbstractOperations.jl @@ -5,10 +5,12 @@ export Δx, Δy, Δz, Ax, Ay, Az, volume export Average, Integral, CumulativeIntegral, KernelFunctionOperation export UnaryOperation, Derivative, BinaryOperation, MultiaryOperation, ConditionalOperation +using Base: @propagate_inbounds +import Adapt using CUDA -using Base: @propagate_inbounds +using Oceananigans using Oceananigans.Architectures using Oceananigans.Grids using Oceananigans.Operators @@ -16,11 +18,9 @@ using Oceananigans.BoundaryConditions using Oceananigans.Fields using Oceananigans.Utils -using Oceananigans: location, AbstractModel using Oceananigans.Operators: interpolation_operator using Oceananigans.Architectures: device - -import Adapt +using Oceananigans: AbstractModel import Oceananigans.Architectures: architecture, on_architecture import Oceananigans.BoundaryConditions: fill_halo_regions! diff --git a/src/AbstractOperations/at.jl b/src/AbstractOperations/at.jl index 18157fcab4..84c85c96ba 100644 --- a/src/AbstractOperations/at.jl +++ b/src/AbstractOperations/at.jl @@ -1,6 +1,3 @@ -using Oceananigans -using Oceananigans.Fields: default_indices - """ insert_location(ex::Expr, location) @@ -54,6 +51,8 @@ macro at(location, abstract_operation) return wrapped_operation end +using Oceananigans.Fields: default_indices + # Numbers and functions do not have index restrictions indices(f::Function) = default_indices(3) indices(f::Number) = default_indices(3) diff --git a/src/AbstractOperations/grid_metrics.jl b/src/AbstractOperations/grid_metrics.jl index 4205135ef0..f873242aa1 100644 --- a/src/AbstractOperations/grid_metrics.jl +++ b/src/AbstractOperations/grid_metrics.jl @@ -1,29 +1,26 @@ using Adapt using Oceananigans.Operators -using Oceananigans.Grids: AbstractGrid -using Oceananigans.Fields: AbstractField, default_indices, location - -import Oceananigans.Grids: xspacings, yspacings, zspacings, λspacings, φspacings +using Oceananigans.Fields: default_indices abstract type AbstractGridMetric end -struct XSpacingMetric <: AbstractGridMetric end -struct YSpacingMetric <: AbstractGridMetric end -struct ZSpacingMetric <: AbstractGridMetric end +struct XSpacingMetric <: AbstractGridMetric end +struct YSpacingMetric <: AbstractGridMetric end +struct ZSpacingMetric <: AbstractGridMetric end metric_function_prefix(::XSpacingMetric) = :Δx metric_function_prefix(::YSpacingMetric) = :Δy metric_function_prefix(::ZSpacingMetric) = :Δz -struct XAreaMetric <: AbstractGridMetric end -struct YAreaMetric <: AbstractGridMetric end -struct ZAreaMetric <: AbstractGridMetric end +struct XAreaMetric <: AbstractGridMetric end +struct YAreaMetric <: AbstractGridMetric end +struct ZAreaMetric <: AbstractGridMetric end metric_function_prefix(::XAreaMetric) = :Ax metric_function_prefix(::YAreaMetric) = :Ay metric_function_prefix(::ZAreaMetric) = :Az -struct VolumeMetric <: AbstractGridMetric end +struct VolumeMetric <: AbstractGridMetric end metric_function_prefix(::VolumeMetric) = :V @@ -36,7 +33,7 @@ const Δy = YSpacingMetric() Instance of `ZSpacingMetric` that generates `BinaryOperation`s between `AbstractField`s and the vertical grid spacing evaluated -at the same location as the `AbstractField`. +at the same location as the `AbstractField`. `Δx` and `Δy` play a similar role for horizontal grid spacings. @@ -136,7 +133,7 @@ Adapt.adapt_structure(to, gm::GridMetricOperation{LX, LY, LZ}) where {LX, LY, LZ on_architecture(to, gm::GridMetricOperation{LX, LY, LZ}) where {LX, LY, LZ} = GridMetricOperation{LX, LY, LZ}(on_architecture(to, gm.metric), on_architecture(to, gm.grid)) - + @inline Base.getindex(gm::GridMetricOperation, i, j, k) = gm.metric(i, j, k, gm.grid) @@ -144,156 +141,3 @@ indices(::GridMetricOperation) = default_indices(3) # Special constructor for BinaryOperation GridMetricOperation(L, metric, grid) = GridMetricOperation{L[1], L[2], L[3]}(metric_function(L, metric), grid) - -##### -##### Spacings -##### - -""" - xspacings(grid, ℓx, ℓy, ℓz) - -Return a `KernelFunctionOperation` that computes the grid spacings for `grid` -in the ``x`` direction at location `ℓx, ℓy, ℓz`. - -Examples -======== -```jldoctest -julia> using Oceananigans - -julia> grid = RectilinearGrid(size=(2, 4, 8), extent=(1, 1, 1)); - -julia> xspacings(grid, Center(), Center(), Center()) -KernelFunctionOperation at (Center, Center, Center) -├── grid: 2×4×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 2×3×3 halo -├── kernel_function: xspacing (generic function with 27 methods) -└── arguments: ("Center", "Center", "Center") -``` -""" -function xspacings(grid, ℓx, ℓy, ℓz) - LX, LY, LZ = map(typeof, (ℓx, ℓy, ℓz)) - Δx_op = KernelFunctionOperation{LX, LY, LZ}(xspacing, grid, ℓx, ℓy, ℓz) - return Δx_op -end - -""" - yspacings(grid, ℓx, ℓy, ℓz) - -Return a `KernelFunctionOperation` that computes the grid spacings for `grid` -in the ``y`` direction at location `ℓx, ℓy, ℓz`. - -Examples -======== -```jldoctest -julia> using Oceananigans - -julia> grid = RectilinearGrid(size=(2, 4, 8), extent=(1, 1, 1)); - -julia> yspacings(grid, Center(), Face(), Center()) -KernelFunctionOperation at (Center, Face, Center) -├── grid: 2×4×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 2×3×3 halo -├── kernel_function: yspacing (generic function with 27 methods) -└── arguments: ("Center", "Face", "Center") -``` -""" -function yspacings(grid, ℓx, ℓy, ℓz) - LX, LY, LZ = map(typeof, (ℓx, ℓy, ℓz)) - Δy_op = KernelFunctionOperation{LX, LY, LZ}(yspacing, grid, ℓx, ℓy, ℓz) - return Δy_op -end - -""" - zspacings(grid, ℓx, ℓy, ℓz) - -Return a `KernelFunctionOperation` that computes the grid spacings for `grid` -in the ``z`` direction at location `ℓx, ℓy, ℓz`. - -Examples -======== -```jldoctest -julia> using Oceananigans - -julia> grid = RectilinearGrid(size=(2, 4, 8), extent=(1, 1, 1)); - -julia> zspacings(grid, Center(), Center(), Face()) -KernelFunctionOperation at (Center, Center, Face) -├── grid: 2×4×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 2×3×3 halo -├── kernel_function: zspacing (generic function with 27 methods) -└── arguments: ("Center", "Center", "Face") -``` -""" -function zspacings(grid, ℓx, ℓy, ℓz) - LX, LY, LZ = map(typeof, (ℓx, ℓy, ℓz)) - Δz_op = KernelFunctionOperation{LX, LY, LZ}(zspacing, grid, ℓx, ℓy, ℓz) - return Δz_op -end - -""" - λspacings(grid, ℓx, ℓy, ℓz) - -Return a `KernelFunctionOperation` that computes the grid spacings for `grid` -in the ``z`` direction at location `ℓx, ℓy, ℓz`. - -Examples -======== -```jldoctest -julia> using Oceananigans - -julia> grid = LatitudeLongitudeGrid(size=(36, 34, 25), - longitude = (-180, 180), - latitude = (-85, 85), - z = (-1000, 0)); - -julia> λspacings(grid, Center(), Face(), Center()) -KernelFunctionOperation at (Center, Face, Center) -├── grid: 36×34×25 LatitudeLongitudeGrid{Float64, Periodic, Bounded, Bounded} on CPU with 3×3×3 halo and with precomputed metrics -├── kernel_function: λspacing (generic function with 5 methods) -└── arguments: ("Center", "Face", "Center") -``` -""" -function λspacings(grid, ℓx, ℓy, ℓz) - LX, LY, LZ = map(typeof, (ℓx, ℓy, ℓz)) - Δλ_op = KernelFunctionOperation{LX, LY, LZ}(λspacing, grid, ℓx, ℓy, ℓz) - return Δλ_op -end - -""" - φspacings(grid, ℓx, ℓy, ℓz) - -Return a `KernelFunctionOperation` that computes the grid spacings for `grid` -in the ``z`` direction at location `ℓx, ℓy, ℓz`. - -Examples -======== -```jldoctest -julia> using Oceananigans - -julia> grid = LatitudeLongitudeGrid(size=(36, 34, 25), - longitude = (-180, 180), - latitude = (-85, 85), - z = (-1000, 0)); - -julia> φspacings(grid, Center(), Face(), Center()) -KernelFunctionOperation at (Center, Face, Center) -├── grid: 36×34×25 LatitudeLongitudeGrid{Float64, Periodic, Bounded, Bounded} on CPU with 3×3×3 halo and with precomputed metrics -├── kernel_function: φspacing (generic function with 5 methods) -└── arguments: ("Center", "Face", "Center") -``` -""" -function φspacings(grid, ℓx, ℓy, ℓz) - LX, LY, LZ = map(typeof, (ℓx, ℓy, ℓz)) - Δφ_op = KernelFunctionOperation{LX, LY, LZ}(φspacing, grid, ℓx, ℓy, ℓz) - return Δφ_op -end - -@inline xspacings(field::AbstractField) = xspacings(field.grid, location(field)...) -@inline yspacings(field::AbstractField) = yspacings(field.grid, location(field)...) -@inline zspacings(field::AbstractField) = zspacings(field.grid, location(field)...) -@inline λspacings(field::AbstractField) = λspacings(field.grid, location(field)...) -@inline φspacings(field::AbstractField) = φspacings(field.grid, location(field)...) - -# Some defaults for e.g. easy CFL computations. -@inline xspacings(grid::AbstractGrid) = xspacings(grid, Center(), Center(), Center()) -@inline yspacings(grid::AbstractGrid) = yspacings(grid, Center(), Center(), Center()) -@inline zspacings(grid::AbstractGrid) = zspacings(grid, Center(), Center(), Center()) -@inline λspacings(grid::AbstractGrid) = λspacings(grid, Center(), Center(), Center()) -@inline φspacings(grid::AbstractGrid) = φspacings(grid, Center(), Center(), Center()) diff --git a/src/Advection/Advection.jl b/src/Advection/Advection.jl index 6473915aed..c0c1d13e11 100644 --- a/src/Advection/Advection.jl +++ b/src/Advection/Advection.jl @@ -17,7 +17,9 @@ export advective_tracer_flux_z, AdvectionScheme, - Centered, UpwindBiased, WENO, + Centered, CenteredSecondOrder, CenteredFourthOrder, + UpwindBiased, UpwindBiasedFirstOrder, UpwindBiasedThirdOrder, UpwindBiasedFifthOrder, + WENO, WENOThirdOrder, WENOFifthOrder, VectorInvariant, WENOVectorInvariant, FluxFormAdvection, EnergyConserving, diff --git a/src/Advection/adapt_advection_order.jl b/src/Advection/adapt_advection_order.jl index 588dc497f2..41f3f60c9a 100644 --- a/src/Advection/adapt_advection_order.jl +++ b/src/Advection/adapt_advection_order.jl @@ -20,14 +20,9 @@ function adapt_advection_order(advection, grid::AbstractGrid) advection_y = y_advection(advection) advection_z = z_advection(advection) - tx = topology(grid, 1)() - ty = topology(grid, 2)() - tz = topology(grid, 3)() - - # Note: no adaptation in Flat directions. - new_advection_x = adapt_advection_order(tx, advection_x, size(grid, 1), grid) - new_advection_y = adapt_advection_order(ty, advection_y, size(grid, 2), grid) - new_advection_z = adapt_advection_order(tz, advection_z, size(grid, 3), grid) + new_advection_x = adapt_advection_order(advection_x, size(grid, 1), grid) + new_advection_y = adapt_advection_order(advection_y, size(grid, 2), grid) + new_advection_z = adapt_advection_order(advection_z, size(grid, 3), grid) # Check that we indeed changed the advection operator changed_x = new_advection_x != advection_x @@ -44,12 +39,13 @@ function adapt_advection_order(advection, grid::AbstractGrid) @info "Using the advection scheme $(summary(new_advection.y)) in the y-direction because size(grid, 2) = $(size(grid, 2))" end if changed_z - @info "Using the advection scheme $(summary(new_advection.z)) in the z-direction because size(grid, 3) = $(size(grid, 3))" + @info "Using the advection scheme $(summary(new_advection.z)) in the x-direction because size(grid, 3) = $(size(grid, 3))" end return ifelse(changed_advection, new_advection, advection) end + x_advection(flux_form::FluxFormAdvection) = flux_form.x y_advection(flux_form::FluxFormAdvection) = flux_form.y z_advection(flux_form::FluxFormAdvection) = flux_form.z @@ -58,15 +54,10 @@ x_advection(advection) = advection y_advection(advection) = advection z_advection(advection) = advection -# Don't adapt advection in Flat directions -adapt_advection_order(topo, advection, N, grid) = adapt_advection_order(advection, N, grid) -adapt_advection_order(::Flat, advection, N, grid) = advection - # For the moment, we do not adapt the advection order for the VectorInvariant advection scheme adapt_advection_order(advection::VectorInvariant, grid::AbstractGrid) = advection adapt_advection_order(advection::Nothing, grid::AbstractGrid) = nothing adapt_advection_order(advection::Nothing, N::Int, grid::AbstractGrid) = nothing -adapt_advection_order(advection, N, grid) = advection ##### ##### Directional adapt advection order @@ -76,7 +67,7 @@ function adapt_advection_order(advection::Centered{B}, N::Int, grid::AbstractGri if N >= B return advection else - return Centered(; order=2N) + return Centered(; order = 2N) end end @@ -84,18 +75,15 @@ function adapt_advection_order(advection::UpwindBiased{B}, N::Int, grid::Abstrac if N >= B return advection else - return UpwindBiased(; order=2N-1) + return UpwindBiased(; order = 2N - 1) end end """ new_weno_scheme(grid, order, bounds, XT, YT, ZT) -Constructs a new WENO scheme based on the given parameters. -`XT`, `YT`, and `ZT` is the type of the precomputed weno coefficients in the -x-direction, y-direction and z-direction. -A _non-stretched_ WENO scheme has `T` equal to `Nothing` everywhere. -In case of a non-stretched WENO scheme, +Constructs a new WENO scheme based on the given parameters. `XT`, `YT`, and `ZT` is the type of the precomputed weno coefficients in the +x-direction, y-direction and z-direction. A _non-stretched_ WENO scheme has `T` equal to `Nothing` everywhere. In case of a non-stretched WENO scheme, we rebuild the advection without passing the grid information, otherwise we use the grid to account for stretched directions. """ new_weno_scheme(::WENO, grid, order, bounds, ::Type{Nothing}, ::Type{Nothing}, ::Type{Nothing},) = WENO(; order, bounds) @@ -105,6 +93,6 @@ function adapt_advection_order(advection::WENO{B, FT, XT, YT, ZT}, N::Int, grid: if N >= B return advection else - return new_weno_scheme(advection, grid, 2N-1, advection.bounds, XT, YT, ZT) + return new_weno_scheme(advection, grid, 2N - 1, advection.bounds, XT, YT, ZT) end -end +end \ No newline at end of file diff --git a/src/Advection/centered_reconstruction.jl b/src/Advection/centered_reconstruction.jl index c0dbd4a22c..5884d37b5a 100644 --- a/src/Advection/centered_reconstruction.jl +++ b/src/Advection/centered_reconstruction.jl @@ -58,7 +58,7 @@ function Centered(FT::DataType = Float64; grid = nothing, order = 2) return Centered{N, FT}(coefficients..., buffer_scheme) end -Base.summary(a::Centered{N}) where N = string("Centered(order=", 2N, ")") +Base.summary(a::Centered{N}) where N = string("Centered reconstruction order ", N*2) Base.show(io::IO, a::Centered{N, FT, XT, YT, ZT}) where {N, FT, XT, YT, ZT} = print(io, summary(a), " \n", @@ -85,6 +85,9 @@ on_architecture(to, scheme::Centered{N, FT}) where {N, FT} = # Useful aliases Centered(grid, FT::DataType=Float64; kwargs...) = Centered(FT; grid, kwargs...) +CenteredSecondOrder(grid=nothing, FT::DataType=Float64) = Centered(grid, FT; order=2) +CenteredFourthOrder(grid=nothing, FT::DataType=Float64) = Centered(grid, FT; order=4) + const ACAS = AbstractCenteredAdvectionScheme # left and right biased for Centered reconstruction are just symmetric! diff --git a/src/Advection/flux_form_advection.jl b/src/Advection/flux_form_advection.jl index 6e4de07af8..faa35fb668 100644 --- a/src/Advection/flux_form_advection.jl +++ b/src/Advection/flux_form_advection.jl @@ -10,10 +10,10 @@ struct FluxFormAdvection{N, FT, A, B, C} <: AbstractAdvectionScheme{N, FT} end """ - FluxFormAdvection(x_advection, y_advection, z_advection) + function FluxFormAdvection(x, y, z) -Return a `FluxFormAdvection` type with reconstructions schemes `x_advection`, `y_advection`, -and `z_advection` to be applied in the ``x``, ``y``, and ``z`` directions, respectively. +Builds a `FluxFormAdvection` type with reconstructions schemes `x`, `y`, and `z` to be applied in +the x, y, and z direction, respectively. """ function FluxFormAdvection(x_advection, y_advection, z_advection) Hx = required_halo_size_x(x_advection) @@ -26,13 +26,8 @@ function FluxFormAdvection(x_advection, y_advection, z_advection) return FluxFormAdvection{H, FT}(x_advection, y_advection, z_advection) end -Base.summary(scheme::FluxFormAdvection) = string("FluxFormAdvection(x=", - summary(scheme.x), ", y=", - summary(scheme.y), ", z=", - summary(scheme.z), ")") - Base.show(io::IO, scheme::FluxFormAdvection) = - print(io, "FluxFormAdvection with direction-based reconstructions:", " \n", + print(io, "FluxFormAdvection with reconstructions: ", " \n", " ├── x: ", summary(scheme.x), "\n", " ├── y: ", summary(scheme.y), "\n", " └── z: ", summary(scheme.z)) diff --git a/src/Advection/positivity_preserving_tracer_advection_operators.jl b/src/Advection/positivity_preserving_tracer_advection_operators.jl index c6a6a0a608..ff652ebabb 100644 --- a/src/Advection/positivity_preserving_tracer_advection_operators.jl +++ b/src/Advection/positivity_preserving_tracer_advection_operators.jl @@ -4,11 +4,9 @@ const ω̂₁ = 5/18 const ω̂ₙ = 5/18 const ε₂ = 1e-20 -# Here in the future we can easily add UpwindBiased +# Here in the future we can easily add UpwindBiasedFifthOrder const BoundPreservingScheme = PositiveWENO -@inline div_Uc(i, j, k, grid, advection::BoundPreservingScheme, U, ::ZeroField) = zero(grid) - # Is this immersed-boundary safe without having to extend it in ImmersedBoundaries.jl? I think so... (velocity on immmersed boundaries is masked to 0) @inline function div_Uc(i, j, k, grid, advection::BoundPreservingScheme, U, c) diff --git a/src/Advection/upwind_biased_reconstruction.jl b/src/Advection/upwind_biased_reconstruction.jl index 303d32aab7..eb3c6805ec 100644 --- a/src/Advection/upwind_biased_reconstruction.jl +++ b/src/Advection/upwind_biased_reconstruction.jl @@ -3,9 +3,9 @@ ##### """ - struct UpwindBiased <: AbstractUpwindBiasedAdvectionScheme{3} + struct UpwindBiasedFifthOrder <: AbstractUpwindBiasedAdvectionScheme{3} -Upwind-biased reconstruction scheme. +Upwind-biased fifth-order advection scheme. """ struct UpwindBiased{N, FT, XT, YT, ZT, CA, SI} <: AbstractUpwindBiasedAdvectionScheme{N, FT} "Coefficient for Upwind reconstruction on stretched ``x``-faces" @@ -65,7 +65,7 @@ function UpwindBiased(FT::DataType = Float64; grid = nothing, order = 3) return UpwindBiased{N, FT}(coefficients..., buffer_scheme, advecting_velocity_scheme) end -Base.summary(a::UpwindBiased{N}) where N = string("UpwindBiased(order=", 2N-1, ")") +Base.summary(a::UpwindBiased{N}) where N = string("Upwind Biased reconstruction order ", N*2-1) Base.show(io::IO, a::UpwindBiased{N, FT, XT, YT, ZT}) where {N, FT, XT, YT, ZT} = print(io, summary(a), " \n", @@ -95,6 +95,10 @@ on_architecture(to, scheme::UpwindBiased{N, FT}) where {N, FT} = # Useful aliases UpwindBiased(grid, FT::DataType=Float64; kwargs...) = UpwindBiased(FT; grid, kwargs...) +UpwindBiasedFirstOrder(grid=nothing, FT::DataType=Float64) = UpwindBiased(grid, FT; order = 1) +UpwindBiasedThirdOrder(grid=nothing, FT::DataType=Float64) = UpwindBiased(grid, FT; order = 3) +UpwindBiasedFifthOrder(grid=nothing, FT::DataType=Float64) = UpwindBiased(grid, FT; order = 5) + const AUAS = AbstractUpwindBiasedAdvectionScheme # symmetric interpolation for UpwindBiased and WENO diff --git a/src/Advection/vector_invariant_advection.jl b/src/Advection/vector_invariant_advection.jl index 5d5becb2f5..e16cd7438a 100644 --- a/src/Advection/vector_invariant_advection.jl +++ b/src/Advection/vector_invariant_advection.jl @@ -36,12 +36,12 @@ end VectorInvariant(; vorticity_scheme = EnstrophyConserving(), vorticity_stencil = VelocityStencil(), vertical_scheme = EnergyConserving(), + kinetic_energy_gradient_scheme = vertical_scheme, divergence_scheme = vertical_scheme, - kinetic_energy_gradient_scheme = divergence_scheme, - upwinding = OnlySelfUpwinding(; cross_scheme = divergence_scheme), + upwinding = OnlySelfUpwinding(; cross_scheme = vertical_scheme), multi_dimensional_stencil = false) -Return a vector-invariant momentum advection scheme. +Return a vector invariant momentum advection scheme. Keyword arguments ================= @@ -65,10 +65,13 @@ Keyword arguments - `upwinding`: Treatment of upwinded reconstruction of divergence and kinetic energy gradient. Default: `OnlySelfUpwinding()`. Options: * `CrossAndSelfUpwinding()` * `OnlySelfUpwinding()` + * `VelocityUpwinding()` -- `multi_dimensional_stencil`: whether or not to use a horizontal two-dimensional stencil for the reconstruction - of vorticity, divergence and kinetic energy gradient. Currently the "tangential" - direction uses 5th-order centered WENO reconstruction. Default: false +- `upwinding` + +- `multi_dimensional_stencil` : whether or not to use a horizontal two-dimensional stencil for the reconstruction + of vorticity, divergence and kinetic energy gradient. Currently the "tangential" + direction uses 5th-order centered WENO reconstruction. Examples ======== @@ -79,25 +82,26 @@ julia> using Oceananigans julia> VectorInvariant() Vector Invariant, Dimension-by-dimension reconstruction Vorticity flux scheme: - └── Oceananigans.Advection.EnstrophyConserving{Float64} + └── EnstrophyConserving{Float64} Vertical advection / Divergence flux scheme: - └── Oceananigans.Advection.EnergyConserving{Float64} + └── EnergyConserving{Float64} + ``` ```jldoctest julia> using Oceananigans julia> VectorInvariant(vorticity_scheme = WENO(), vertical_scheme = WENO(order = 3)) -Vector Invariant, Dimension-by-dimension reconstruction - Vorticity flux scheme: - ├── WENO(order=5) +Vector Invariant, Dimension-by-dimension reconstruction + Vorticity flux scheme: + ├── WENO reconstruction order 5 └── smoothness ζ: Oceananigans.Advection.VelocityStencil() - Vertical advection / Divergence flux scheme: - ├── WENO(order=3) - └── upwinding treatment: OnlySelfUpwinding - KE gradient and Divergence flux cross terms reconstruction: - └── Centered(order=2) - Smoothness measures: + Vertical advection / Divergence flux scheme: + ├── WENO reconstruction order 3 + └── upwinding treatment: OnlySelfUpwinding + KE gradient and Divergence flux cross terms reconstruction: + └── Centered reconstruction order 2 + Smoothness measures: ├── smoothness δU: FunctionStencil f = divergence_smoothness ├── smoothness δV: FunctionStencil f = divergence_smoothness ├── smoothness δu²: FunctionStencil f = u_smoothness @@ -123,9 +127,9 @@ function VectorInvariant(; vorticity_scheme = EnstrophyConserving(), FT = eltype(vorticity_scheme) return VectorInvariant{N, FT, multi_dimensional_stencil}(vorticity_scheme, - vorticity_stencil, + vorticity_stencil, vertical_scheme, - kinetic_energy_gradient_scheme, + kinetic_energy_gradient_scheme, divergence_scheme, upwinding) end @@ -171,26 +175,12 @@ Base.show(io::IO, a::VectorInvariant{N, FT}) where {N, FT} = nothing_to_default(user_value; default) = isnothing(user_value) ? default : user_value """ - WENOVectorInvariant(FT = Float64; - upwinding = nothing, - vorticity_stencil = VelocityStencil(), - order = nothing, - vorticity_order = nothing, - vertical_order = nothing, - divergence_order = nothing, - kinetic_energy_gradient_order = nothing, - multi_dimensional_stencil = false, - weno_kw...) - -Return a vector-invariant weighted essentially non-oscillatory (WENO) scheme. -See [`VectorInvariant`](@ref) and [`WENO`](@ref) for kwargs definitions. - -If `multi_dimensional_stencil = true` is selected, then a 2D horizontal stencil -is implemented for the WENO scheme (instead of a 1D stencil). This 2D horizontal -stencil performs a centered 5th-order WENO reconstruction of vorticity, -divergence and kinetic energy in the horizontal direction tangential to the upwind direction. + WENOVectorInvariant(; upwinding = nothing, + multi_dimensional_stencil = false, + weno_kw...) + """ -function WENOVectorInvariant(FT::DataType = Float64; +function WENOVectorInvariant(FT::DataType = Float64; upwinding = nothing, vorticity_stencil = VelocityStencil(), order = nothing, @@ -221,21 +211,15 @@ function WENOVectorInvariant(FT::DataType = Float64; default_upwinding = OnlySelfUpwinding(cross_scheme = divergence_scheme) upwinding = nothing_to_default(upwinding; default = default_upwinding) - N = max(required_halo_size_x(vorticity_scheme), - required_halo_size_y(vorticity_scheme), - required_halo_size_x(divergence_scheme), - required_halo_size_y(divergence_scheme), - required_halo_size_x(kinetic_energy_gradient_scheme), - required_halo_size_y(kinetic_energy_gradient_scheme), - required_halo_size_z(vertical_scheme)) - + schemes = (vorticity_scheme, vertical_scheme, kinetic_energy_gradient_scheme, divergence_scheme) + N = maximum(required_halo_size(s) for s in schemes) FT = eltype(vorticity_scheme) # assumption return VectorInvariant{N, FT, multi_dimensional_stencil}(vorticity_scheme, - vorticity_stencil, + vorticity_stencil, vertical_scheme, - kinetic_energy_gradient_scheme, - divergence_scheme, + kinetic_energy_gradient_scheme, + divergence_scheme, upwinding) end @@ -250,8 +234,8 @@ end return Hx == 1 ? Hx : Hx + 1 end -@inline required_halo_size_y(scheme::VectorInvariant) = required_halo_size_x(scheme) -@inline required_halo_size_z(scheme::VectorInvariant) = required_halo_size_z(scheme.vertical_scheme) +@inline required_halo_size_y(scheme::VectorInvariant) = required_halo_size_x(scheme) +@inline required_halo_size_z(scheme::VectorInvariant) = required_halo_size_z(scheme.vertical_scheme) Adapt.adapt_structure(to, scheme::VectorInvariant{N, FT, M}) where {N, FT, M} = VectorInvariant{N, FT, M}(Adapt.adapt(to, scheme.vorticity_scheme), @@ -284,10 +268,10 @@ for bias in (:_biased, :_symmetric) multidim_interp = Symbol(:_multi_dimensional_reconstruction_, dir2) @eval begin - @inline $interp_func(i, j, k, grid, ::VectorInvariant, interp_scheme, args...) = + @inline $interp_func(i, j, k, grid, ::VectorInvariant, interp_scheme, args...) = $interp_func(i, j, k, grid, interp_scheme, args...) - @inline $interp_func(i, j, k, grid, ::MultiDimensionalVectorInvariant, interp_scheme, args...) = + @inline $interp_func(i, j, k, grid, ::MultiDimensionalVectorInvariant, interp_scheme, args...) = $multidim_interp(i, j, k, grid, interp_scheme, $interp_func, args...) end end @@ -315,8 +299,8 @@ end ##### Follows https://mitgcm.readthedocs.io/en/latest/algorithm/algorithm.html#vector-invariant-momentum-equations ##### -@inbounds ζ₂wᶠᶜᶠ(i, j, k, grid, u, w) = ℑxᶠᵃᵃ(i, j, k, grid, Az_qᶜᶜᶠ, w) * ∂zᶠᶜᶠ(i, j, k, grid, u) -@inbounds ζ₁wᶜᶠᶠ(i, j, k, grid, v, w) = ℑyᵃᶠᵃ(i, j, k, grid, Az_qᶜᶜᶠ, w) * ∂zᶜᶠᶠ(i, j, k, grid, v) +@inbounds ζ₂wᶠᶜᶠ(i, j, k, grid, u, w) = ℑxᶠᵃᵃ(i, j, k, grid, Az_qᶜᶜᶠ, w) * ∂zᶠᶜᶠ(i, j, k, grid, u) +@inbounds ζ₁wᶜᶠᶠ(i, j, k, grid, v, w) = ℑyᵃᶠᵃ(i, j, k, grid, Az_qᶜᶜᶠ, w) * ∂zᶜᶠᶠ(i, j, k, grid, v) @inline vertical_advection_U(i, j, k, grid, ::VectorInvariantVerticalEnergyConserving, U) = ℑzᵃᵃᶜ(i, j, k, grid, ζ₂wᶠᶜᶠ, U.u, U.w) / Azᶠᶜᶜ(i, j, k, grid) @inline vertical_advection_V(i, j, k, grid, ::VectorInvariantVerticalEnergyConserving, U) = ℑzᵃᵃᶜ(i, j, k, grid, ζ₁wᶜᶠᶠ, U.v, U.w) / Azᶜᶠᶜ(i, j, k, grid) @@ -345,8 +329,8 @@ end ##### Horizontal advection 4 formulations: ##### 1. Energy conservative ##### 2. Enstrophy conservative -##### 3. Dimension-By-Dimension Vorticity upwinding -##### 4. Two-Dimensional (x and y) Vorticity upwinding +##### 3. Dimension-By-Dimension Vorticity upwinding +##### 4. Two-Dimensional (x and y) Vorticity upwinding ##### ##### @@ -377,7 +361,7 @@ end return - v̂ * ζᴿ end -@inline function horizontal_advection_V(i, j, k, grid, scheme::VectorInvariantUpwindVorticity, u, v) +@inline function horizontal_advection_V(i, j, k, grid, scheme::VectorInvariantUpwindVorticity, u, v) Sζ = scheme.vorticity_stencil @@ -391,7 +375,7 @@ end ##### Fallback to flux form advection (LatitudeLongitudeGrid) ##### -@inline function U_dot_∇u(i, j, k, grid, advection::AbstractAdvectionScheme, U) +@inline function U_dot_∇u(i, j, k, grid, advection::AbstractAdvectionScheme, U) v̂ = ℑxᶠᵃᵃ(i, j, k, grid, ℑyᵃᶜᵃ, Δx_qᶜᶠᶜ, U.v) / Δxᶠᶜᶜ(i, j, k, grid) û = @inbounds U.u[i, j, k] @@ -441,7 +425,7 @@ const CZ{N} = Centered{N, <:Any, <:Any, <:Any, <:Nothing} const AS = AbstractSmoothnessStencil -# To adapt passing smoothness stencils to upwind biased schemes and centered schemes (not WENO) +# To adapt passing smoothness stencils to upwind biased schemes and centered schemes (not WENO) for b in 1:6 @eval begin @inline inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, s::C{$b}, f::Function, idx, loc, ::AS, args...) = inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, s, f, idx, loc, args...) diff --git a/src/Advection/vector_invariant_upwinding.jl b/src/Advection/vector_invariant_upwinding.jl index cf184582f2..4949a453e2 100644 --- a/src/Advection/vector_invariant_upwinding.jl +++ b/src/Advection/vector_invariant_upwinding.jl @@ -27,7 +27,7 @@ end @inline extract_centered_scheme(scheme::AUAS) = scheme.advecting_velocity_scheme """ - OnlySelfUpwinding(; cross_scheme = Centered(), + OnlySelfUpwinding(; cross_scheme = CenteredSecondOrder(), δU_stencil = FunctionStencil(divergence_smoothness), δV_stencil = FunctionStencil(divergence_smoothness), δu²_stencil = FunctionStencil(u_smoothness), @@ -43,7 +43,7 @@ Keyword arguments ================= - `cross_scheme`: Advection scheme used for cross-reconstructed terms (tangential velocities) - in the kinetic energy gradient and the divergence flux. Defaults to `Centered()`. + in the kinetic energy gradient and the divergence flux. Defaults to `CenteredSecondOrder()`. - `δU_stencil`: Stencil used for smoothness indicators of `δx_U` in case of a `WENO` upwind reconstruction. Defaults to `FunctionStencil(divergence_smoothness)` - `δV_stencil`: Same as `δU_stencil` but for the smoothness of `δy_V` @@ -52,7 +52,7 @@ Keyword arguments - `δv²_stencil`: Same as `δu²_stencil` but for the smoothness of `δy_v²` Defaults to `FunctionStencil(v_smoothness)` """ -OnlySelfUpwinding(; cross_scheme = Centered(), +OnlySelfUpwinding(; cross_scheme = CenteredSecondOrder(), δU_stencil = FunctionStencil(divergence_smoothness), δV_stencil = FunctionStencil(divergence_smoothness), δu²_stencil = FunctionStencil(u_smoothness), @@ -60,7 +60,7 @@ OnlySelfUpwinding(; cross_scheme = Centered(), ) = OnlySelfUpwinding(extract_centered_scheme(cross_scheme), δU_stencil, δV_stencil, δu²_stencil, δv²_stencil) """ - CrossAndSelfUpwinding(; cross_scheme = Centered(), + CrossAndSelfUpwinding(; cross_scheme = CenteredSecondOrder(), divergence_stencil = DefaultStencil(), δu²_stencil = FunctionStencil(u_smoothness), δv²_stencil = FunctionStencil(v_smoothness)) @@ -74,7 +74,7 @@ Keyword arguments ================= - `cross_scheme`: Advection scheme used for cross-reconstructed terms (tangential velocities) - in the kinetic energy gradient. Defaults to `Centered()`. + in the kinetic energy gradient. Defaults to `CenteredSecondOrder()`. - `divergence_stencil`: Stencil used for smoothness indicators of `δx_U + δy_V` in case of a `WENO` upwind reconstruction. Defaults to `DefaultStencil()`. - `δu²_stencil`: Stencil used for smoothness indicators of `δx_u²` in case of a `WENO` upwind reconstruction. @@ -82,7 +82,7 @@ Keyword arguments - `δv²_stencil`: Same as `δu²_stencil` but for the smoothness of `δy_v²` Defaults to `FunctionStencil(v_smoothness)` """ -CrossAndSelfUpwinding(; cross_scheme = Centered(), +CrossAndSelfUpwinding(; cross_scheme = CenteredSecondOrder(), divergence_stencil = DefaultStencil(), δu²_stencil = FunctionStencil(u_smoothness), δv²_stencil = FunctionStencil(v_smoothness), diff --git a/src/Advection/weno_reconstruction.jl b/src/Advection/weno_reconstruction.jl index 23d2208077..f91c633557 100644 --- a/src/Advection/weno_reconstruction.jl +++ b/src/Advection/weno_reconstruction.jl @@ -58,14 +58,14 @@ Examples julia> using Oceananigans julia> WENO() -WENO(order=5) - Boundary scheme: - └── WENO(order=3) - Symmetric scheme: - └── Centered(order=4) +WENO reconstruction order 5 + Boundary scheme: + └── WENO reconstruction order 3 + Symmetric scheme: + └── Centered reconstruction order 4 Directions: - ├── X regular - ├── Y regular + ├── X regular + ├── Y regular └── Z regular ``` @@ -82,14 +82,14 @@ julia> grid = RectilinearGrid(size = (Nx, Nz), halo = (4, 4), topology=(Periodic x = (0, Lx), z = chebychev_spaced_z_faces); julia> WENO(grid; order=7) -WENO(order=7) - Boundary scheme: - └── WENO(order=5) - Symmetric scheme: - └── Centered(order=6) +WENO reconstruction order 7 + Boundary scheme: + └── WENO reconstruction order 5 + Symmetric scheme: + └── Centered reconstruction order 6 Directions: - ├── X regular - ├── Y regular + ├── X regular + ├── Y regular └── Z stretched ``` """ @@ -104,18 +104,15 @@ function WENO(FT::DataType=Float64; mod(order, 2) == 0 && throw(ArgumentError("WENO reconstruction scheme is defined only for odd orders")) - if !isnothing(bounds) - @warn "Bounds preserving WENO is experimental." - end - if order < 3 - # WENO(order=1) is equivalent to UpwindBiased(order=1) - return UpwindBiased(FT; order=1) + # WENO(order = 1) is equivalent to UpwindBiased(order = 1) + return UpwindBiased(FT; order = 1) else - N = Int((order + 1) ÷ 2) + N = Int((order + 1) ÷ 2) + weno_coefficients = compute_reconstruction_coefficients(grid, FT, :WENO; order = N) advecting_velocity_scheme = Centered(FT; grid, order = order - 1) - buffer_scheme = WENO(FT; grid, order=order-2, bounds) + buffer_scheme = WENO(FT; grid, order = order - 2, bounds) end return WENO{N, FT}(weno_coefficients..., bounds, buffer_scheme, advecting_velocity_scheme) @@ -123,10 +120,14 @@ end WENO(grid, FT::DataType=Float64; kwargs...) = WENO(FT; grid, kwargs...) +# Some usefull aliases +WENOThirdOrder(grid=nothing, FT::DataType=Float64; kwargs...) = WENO(grid, FT; order=3, kwargs...) +WENOFifthOrder(grid=nothing, FT::DataType=Float64; kwargs...) = WENO(grid, FT; order=5, kwargs...) + # Flavours of WENO const PositiveWENO = WENO{<:Any, <:Any, <:Any, <:Any, <:Any, <:Tuple} -Base.summary(a::WENO{N}) where N = string("WENO(order=", N*2-1, ")") +Base.summary(a::WENO{N}) where N = string("WENO reconstruction order ", N*2-1) Base.show(io::IO, a::WENO{N, FT, RX, RY, RZ, PP}) where {N, FT, RX, RY, RZ, PP} = print(io, summary(a), " \n", diff --git a/src/Biogeochemistry.jl b/src/Biogeochemistry.jl index c63997e4d9..21f7e27471 100644 --- a/src/Biogeochemistry.jl +++ b/src/Biogeochemistry.jl @@ -1,7 +1,7 @@ module Biogeochemistry using Oceananigans.Grids: Center, xnode, ynode, znode -using Oceananigans.Advection: div_Uc, Centered +using Oceananigans.Advection: div_Uc, CenteredSecondOrder using Oceananigans.Architectures: device, architecture using Oceananigans.Fields: ZeroField diff --git a/src/BoundaryConditions/field_boundary_conditions.jl b/src/BoundaryConditions/field_boundary_conditions.jl index 92c749af6c..a42beeee2c 100644 --- a/src/BoundaryConditions/field_boundary_conditions.jl +++ b/src/BoundaryConditions/field_boundary_conditions.jl @@ -39,12 +39,12 @@ default_auxiliary_bc(::LeftConnected, ::Face) = nothing ##### mutable struct FieldBoundaryConditions{W, E, S, N, B, T, I} - west :: W - east :: E - south :: S - north :: N - bottom :: B - top :: T + west :: W + east :: E + south :: S + north :: N + bottom :: B + top :: T immersed :: I end @@ -65,15 +65,6 @@ FieldBoundaryConditions(indices::Tuple, ::Nothing) = nothing window_boundary_conditions(::Colon, left, right) = left, right window_boundary_conditions(::UnitRange, left, right) = nothing, nothing -on_architecture(arch, fbcs::FieldBoundaryConditions) = - FieldBoundaryConditions(on_architecture(arch, fbcs.west), - on_architecture(arch, fbcs.east), - on_architecture(arch, fbcs.south), - on_architecture(arch, fbcs.north), - on_architecture(arch, fbcs.bottom), - on_architecture(arch, fbcs.top), - on_architecture(arch, fbcs.immersed)) - """ FieldBoundaryConditions(; kwargs...) @@ -166,9 +157,9 @@ end function regularize_immersed_boundary_condition(ibc, grid, loc, field_name, args...) if !(ibc isa DefaultBoundaryCondition) msg = """$field_name was assigned an immersed boundary condition - $ibc, + $ibc , but this is not supported on - $(summary(grid)). + $(summary(grid)) . The immersed boundary condition on $field_name will have no effect. """ diff --git a/src/BoundaryConditions/fill_halo_regions.jl b/src/BoundaryConditions/fill_halo_regions.jl index 7ac3e79ea5..fa380e7029 100644 --- a/src/BoundaryConditions/fill_halo_regions.jl +++ b/src/BoundaryConditions/fill_halo_regions.jl @@ -64,9 +64,7 @@ function fill_halo_regions!(c::MaybeTupledData, boundary_conditions, indices, lo return nothing end -function fill_halo_event!(c, fill_halos!, bcs, indices, loc, arch, grid, args...; - async = false, # This kwargs is specific to DistributedGrids, here is does nothing - kwargs...) +function fill_halo_event!(c, fill_halos!, bcs, indices, loc, arch, grid, args...; kwargs...) # Calculate size and offset of the fill_halo kernel # We assume that the kernel size is the same for west and east boundaries, @@ -297,27 +295,27 @@ end ##### Kernel launchers for single-sided fill_halos ##### -fill_west_halo!(c, bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) = +fill_west_halo!(c, bc, size, offset, loc, arch, grid, args...; kwargs...) = launch!(arch, grid, KernelParameters(size, offset), _fill_only_west_halo!, c, bc, loc, grid, Tuple(args); kwargs...) -fill_east_halo!(c, bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) = +fill_east_halo!(c, bc, size, offset, loc, arch, grid, args...; kwargs...) = launch!(arch, grid, KernelParameters(size, offset), _fill_only_east_halo!, c, bc, loc, grid, Tuple(args); kwargs...) -fill_south_halo!(c, bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) = +fill_south_halo!(c, bc, size, offset, loc, arch, grid, args...; kwargs...) = launch!(arch, grid, KernelParameters(size, offset), _fill_only_south_halo!, c, bc, loc, grid, Tuple(args); kwargs...) -fill_north_halo!(c, bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) = +fill_north_halo!(c, bc, size, offset, loc, arch, grid, args...; kwargs...) = launch!(arch, grid, KernelParameters(size, offset), _fill_only_north_halo!, c, bc, loc, grid, Tuple(args); kwargs...) -fill_bottom_halo!(c, bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) = +fill_bottom_halo!(c, bc, size, offset, loc, arch, grid, args...; kwargs...) = launch!(arch, grid, KernelParameters(size, offset), _fill_only_bottom_halo!, c, bc, loc, grid, Tuple(args); kwargs...) -fill_top_halo!(c, bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) = +fill_top_halo!(c, bc, size, offset, loc, arch, grid, args...; kwargs...) = launch!(arch, grid, KernelParameters(size, offset), _fill_only_top_halo!, c, bc, loc, grid, Tuple(args); kwargs...) @@ -325,17 +323,17 @@ fill_top_halo!(c, bc, size, offset, loc, arch, grid, args...; only_local_halos = ##### Kernel launchers for double-sided fill_halos ##### -function fill_west_and_east_halo!(c, west_bc, east_bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) +function fill_west_and_east_halo!(c, west_bc, east_bc, size, offset, loc, arch, grid, args...; kwargs...) return launch!(arch, grid, KernelParameters(size, offset), _fill_west_and_east_halo!, c, west_bc, east_bc, loc, grid, Tuple(args); kwargs...) end -function fill_south_and_north_halo!(c, south_bc, north_bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) +function fill_south_and_north_halo!(c, south_bc, north_bc, size, offset, loc, arch, grid, args...; kwargs...) return launch!(arch, grid, KernelParameters(size, offset), _fill_south_and_north_halo!, c, south_bc, north_bc, loc, grid, Tuple(args); kwargs...) end -function fill_bottom_and_top_halo!(c, bottom_bc, top_bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) +function fill_bottom_and_top_halo!(c, bottom_bc, top_bc, size, offset, loc, arch, grid, args...; kwargs...) return launch!(arch, grid, KernelParameters(size, offset), _fill_bottom_and_top_halo!, c, bottom_bc, top_bc, loc, grid, Tuple(args); kwargs...) end diff --git a/src/BoundaryConditions/fill_halo_regions_open.jl b/src/BoundaryConditions/fill_halo_regions_open.jl index 039de899dd..152ebd29b7 100644 --- a/src/BoundaryConditions/fill_halo_regions_open.jl +++ b/src/BoundaryConditions/fill_halo_regions_open.jl @@ -6,63 +6,104 @@ Fill open boundary halo regions by filling boundary conditions on field faces with `open_fill`. """ -function fill_open_boundary_regions!(field, boundary_conditions, indices, loc, grid, args...; only_local_halos = false, kwargs...) +function fill_open_boundary_regions!(field, boundary_conditions, indices, loc, grid, args...; kwargs...) arch = architecture(grid) - # gets `fill_halo!`, the function which fills open boundaries at `loc` - # The underlying assumption is that open boundaries are uniquely defined by the location `loc`: - # (Face, Center, Center) -> fill west and east - # (Center, Face, Center) -> fill south and north - # (Center, Center, Face) -> fill bottom and top - fill_halo! = get_open_halo_filling_functions(loc) + left_bc = left_velocity_open_boundary_condition(boundary_conditions, loc) + right_bc = right_velocity_open_boundary_condition(boundary_conditions, loc) - left_bc = left_open_boundary_condition(boundary_conditions, loc) - right_bc = right_open_boundary_condition(boundary_conditions, loc) + # gets `open_fill`, the function which fills open boundaries at `loc`, as well as `regular_fill` + # which is the function which fills non-open boundaries at `loc` which informs `fill_halo_size` + open_fill, regular_fill = get_open_halo_filling_functions(loc) + fill_size = fill_halo_size(field, regular_fill, indices, boundary_conditions, loc, grid) - bcs_tuple = (left_bc, right_bc) + launch!(arch, grid, fill_size, open_fill, field, left_bc, right_bc, loc, grid, args) - if !isnothing(fill_halo!) && any(!isnothing, bcs_tuple) + return nothing +end - # Overwrite the `only_local_halos` keyword argument, because open boundaries - # are always local boundaries that do not require communication - only_local_halos = true +fill_open_boundary_regions!(fields::NTuple, boundary_conditions, indices, loc, grid, args...; kwargs...) = + [fill_open_boundary_regions!(field, boundary_conditions[n], indices, loc[n], grid, args...; kwargs...) for (n, field) in enumerate(fields)] - fill_halo_event!(field, fill_halo!, bcs_tuple, indices, loc, arch, grid, args...; only_local_halos, kwargs...) - end +# for regular halo fills +@inline left_velocity_open_boundary_condition(boundary_condition, loc) = nothing +@inline left_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Face, Center, Center}) = boundary_conditions.west +@inline left_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Center, Face, Center}) = boundary_conditions.south +@inline left_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Center, Center, Face}) = boundary_conditions.bottom - return nothing +@inline right_velocity_open_boundary_condition(boundary_conditions, loc) = nothing +@inline right_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Face, Center, Center}) = boundary_conditions.east +@inline right_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Center, Face, Center}) = boundary_conditions.north +@inline right_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Center, Center, Face}) = boundary_conditions.top + +# for multi region halo fills +@inline left_velocity_open_boundary_condition(boundary_conditions::Tuple, ::Tuple{Face, Center, Center}) = @inbounds boundary_conditions[1] +@inline left_velocity_open_boundary_condition(boundary_conditions::Tuple, ::Tuple{Center, Face, Center}) = @inbounds boundary_conditions[1] +@inline left_velocity_open_boundary_condition(boundary_conditions::Tuple, ::Tuple{Center, Center, Face}) = @inbounds boundary_conditions[1] + +@inline right_velocity_open_boundary_condition(boundary_conditions::Tuple, ::Tuple{Face, Center, Center}) = @inbounds boundary_conditions[2] +@inline right_velocity_open_boundary_condition(boundary_conditions::Tuple, ::Tuple{Center, Face, Center}) = @inbounds boundary_conditions[2] +@inline right_velocity_open_boundary_condition(boundary_conditions::Tuple, ::Tuple{Center, Center, Face}) = @inbounds boundary_conditions[2] + +@inline get_open_halo_filling_functions(loc) = _no_fill!, _no_fill! +@inline get_open_halo_filling_functions(::Tuple{Face, Center, Center}) = _fill_west_and_east_open_halo!, fill_west_and_east_halo! +@inline get_open_halo_filling_functions(::Tuple{Center, Face, Center}) = _fill_south_and_north_open_halo!, fill_south_and_north_halo! +@inline get_open_halo_filling_functions(::Tuple{Center, Center, Face}) = _fill_bottom_and_top_open_halo!, fill_bottom_and_top_halo! + +@kernel _no_fill!(args...) = nothing + +@inline fill_halo_size(field, ::typeof(_no_fill!), args...) = (0, 0) + +@kernel function _fill_west_and_east_open_halo!(c, west_bc, east_bc, loc, grid, args) + j, k = @index(Global, NTuple) + _fill_west_open_halo!(j, k, grid, c, west_bc, loc, args...) + _fill_east_open_halo!(j, k, grid, c, east_bc, loc, args...) end -@inline get_open_halo_filling_functions(loc) = nothing -@inline get_open_halo_filling_functions(::Tuple{Face, Center, Center}) = fill_west_and_east_halo! -@inline get_open_halo_filling_functions(::Tuple{Center, Face, Center}) = fill_south_and_north_halo! -@inline get_open_halo_filling_functions(::Tuple{Center, Center, Face}) = fill_bottom_and_top_halo! +@kernel function _fill_south_and_north_open_halo!(c, south_bc, north_bc, loc, grid, args) + i, k = @index(Global, NTuple) + _fill_south_open_halo!(i, k, grid, c, south_bc, loc, args...) + _fill_north_open_halo!(i, k, grid, c, north_bc, loc, args...) +end -function fill_open_boundary_regions!(fields::Tuple, boundary_conditions, indices, loc, grid, args...; kwargs...) - for n in eachindex(fields) - fill_open_boundary_regions!(fields[n], boundary_conditions[n], indices, loc[n], grid, args...; kwargs...) - end - return nothing +@kernel function _fill_bottom_and_top_open_halo!(c, bottom_bc, top_bc, loc, grid, args) + i, j = @index(Global, NTuple) + _fill_bottom_open_halo!(i, j, grid, c, bottom_bc, loc, args...) + _fill_top_open_halo!(i, j, grid, c, top_bc, loc, args...) end -@inline retrieve_open_bc(bc::OBC) = bc -@inline retrieve_open_bc(bc) = nothing - -# for regular halo fills, return nothing if the BC is not an OBC -@inline left_open_boundary_condition(boundary_condition, loc) = nothing -@inline left_open_boundary_condition(boundary_conditions, ::Tuple{Face, Center, Center}) = retrieve_open_bc(boundary_conditions.west) -@inline left_open_boundary_condition(boundary_conditions, ::Tuple{Center, Face, Center}) = retrieve_open_bc(boundary_conditions.south) -@inline left_open_boundary_condition(boundary_conditions, ::Tuple{Center, Center, Face}) = retrieve_open_bc(boundary_conditions.bottom) - -@inline right_open_boundary_condition(boundary_conditions, loc) = nothing -@inline right_open_boundary_condition(boundary_conditions, ::Tuple{Face, Center, Center}) = retrieve_open_bc(boundary_conditions.east) -@inline right_open_boundary_condition(boundary_conditions, ::Tuple{Center, Face, Center}) = retrieve_open_bc(boundary_conditions.north) -@inline right_open_boundary_condition(boundary_conditions, ::Tuple{Center, Center, Face}) = retrieve_open_bc(boundary_conditions.top) - -# Opern boundary fill -@inline _fill_west_halo!(j, k, grid, c, bc::OBC, loc, args...) = @inbounds c[1, j, k] = getbc(bc, j, k, grid, args...) -@inline _fill_east_halo!(j, k, grid, c, bc::OBC, loc, args...) = @inbounds c[grid.Nx + 1, j, k] = getbc(bc, j, k, grid, args...) -@inline _fill_south_halo!(i, k, grid, c, bc::OBC, loc, args...) = @inbounds c[i, 1, k] = getbc(bc, i, k, grid, args...) -@inline _fill_north_halo!(i, k, grid, c, bc::OBC, loc, args...) = @inbounds c[i, grid.Ny + 1, k] = getbc(bc, i, k, grid, args...) -@inline _fill_bottom_halo!(i, j, grid, c, bc::OBC, loc, args...) = @inbounds c[i, j, 1] = getbc(bc, i, j, grid, args...) -@inline _fill_top_halo!(i, j, grid, c, bc::OBC, loc, args...) = @inbounds c[i, j, grid.Nz + 1] = getbc(bc, i, j, grid, args...) +# Generic fallback + +@inline _fill_west_open_halo!(j, k, grid, c, bc, loc, args...) = nothing +@inline _fill_east_open_halo!(j, k, grid, c, bc, loc, args...) = nothing +@inline _fill_south_open_halo!(i, k, grid, c, bc, loc, args...) = nothing +@inline _fill_north_open_halo!(i, k, grid, c, bc, loc, args...) = nothing +@inline _fill_bottom_open_halo!(i, j, grid, c, bc, loc, args...) = nothing +@inline _fill_top_open_halo!(i, j, grid, c, bc, loc, args...) = nothing + +# Open boundary condition fallback + +@inline _fill_west_open_halo!(j, k, grid, c, bc::OBC, loc, args...) = @inbounds c[1, j, k] = getbc(bc, j, k, grid, args...) +@inline _fill_east_open_halo!(j, k, grid, c, bc::OBC, loc, args...) = @inbounds c[grid.Nx + 1, j, k] = getbc(bc, j, k, grid, args...) +@inline _fill_south_open_halo!(i, k, grid, c, bc::OBC, loc, args...) = @inbounds c[i, 1, k] = getbc(bc, i, k, grid, args...) +@inline _fill_north_open_halo!(i, k, grid, c, bc::OBC, loc, args...) = @inbounds c[i, grid.Ny + 1, k] = getbc(bc, i, k, grid, args...) +@inline _fill_bottom_open_halo!(i, j, grid, c, bc::OBC, loc, args...) = @inbounds c[i, j, 1] = getbc(bc, i, j, grid, args...) +@inline _fill_top_open_halo!(i, j, grid, c, bc::OBC, loc, args...) = @inbounds c[i, j, grid.Nz + 1] = getbc(bc, i, j, grid, args...) + +# Regular boundary fill defaults + +@inline _fill_west_halo!(j, k, grid, c, bc::OBC, loc, args...) = _fill_west_open_halo!(j, k, grid, c, bc, loc, args...) +@inline _fill_east_halo!(j, k, grid, c, bc::OBC, loc, args...) = _fill_east_open_halo!(j, k, grid, c, bc, loc, args...) +@inline _fill_south_halo!(i, k, grid, c, bc::OBC, loc, args...) = _fill_south_open_halo!(i, k, grid, c, bc, loc, args...) +@inline _fill_north_halo!(i, k, grid, c, bc::OBC, loc, args...) = _fill_north_open_halo!(i, k, grid, c, bc, loc, args...) +@inline _fill_bottom_halo!(i, j, grid, c, bc::OBC, loc, args...) = _fill_bottom_open_halo!(i, j, grid, c, bc, loc, args...) +@inline _fill_top_halo!(i, j, grid, c, bc::OBC, loc, args...) = _fill_top_open_halo!(i, j, grid, c, bc, loc, args...) + +# Regular boundary fill for wall normal velocities + +@inline _fill_west_halo!(j, k, grid, c, bc::OBC, ::Tuple{Face, <:Any, <:Any}, args...) = nothing +@inline _fill_east_halo!(j, k, grid, c, bc::OBC, ::Tuple{Face, <:Any, <:Any}, args...) = nothing +@inline _fill_south_halo!(i, k, grid, c, bc::OBC, ::Tuple{<:Any, Face, <:Any}, args...) = nothing +@inline _fill_north_halo!(i, k, grid, c, bc::OBC, ::Tuple{<:Any, Face, <:Any}, args...) = nothing +@inline _fill_bottom_halo!(i, j, grid, c, bc::OBC, ::Tuple{<:Any, <:Any, Face}, args...) = nothing +@inline _fill_top_halo!(i, j, grid, c, bc::OBC, ::Tuple{<:Any, <:Any, Face}, args...) = nothing diff --git a/src/BoundaryConditions/fill_halo_regions_periodic.jl b/src/BoundaryConditions/fill_halo_regions_periodic.jl index edd97528fa..890b9c2fba 100644 --- a/src/BoundaryConditions/fill_halo_regions_periodic.jl +++ b/src/BoundaryConditions/fill_halo_regions_periodic.jl @@ -15,19 +15,19 @@ end @inline fix_halo_offsets(o, co) = co > 0 ? o - co : o # Windowed fields have only positive offsets to correct -function fill_west_and_east_halo!(c, ::PBCT, ::PBCT, size, offset, loc, arch, grid, args...; only_local_halos = false, kw...) +function fill_west_and_east_halo!(c, ::PBCT, ::PBCT, size, offset, loc, arch, grid, args...; kw...) c_parent, yz_size, offset = parent_size_and_offset(c, 2, 3, size, offset) launch!(arch, grid, KernelParameters(yz_size, offset), fill_periodic_west_and_east_halo!, c_parent, Val(grid.Hx), grid.Nx; kw...) return nothing end -function fill_south_and_north_halo!(c, ::PBCT, ::PBCT, size, offset, loc, arch, grid, args...; only_local_halos = false, kw...) +function fill_south_and_north_halo!(c, ::PBCT, ::PBCT, size, offset, loc, arch, grid, args...; kw...) c_parent, xz_size, offset = parent_size_and_offset(c, 1, 3, size, offset) launch!(arch, grid, KernelParameters(xz_size, offset), fill_periodic_south_and_north_halo!, c_parent, Val(grid.Hy), grid.Ny; kw...) return nothing end -function fill_bottom_and_top_halo!(c, ::PBCT, ::PBCT, size, offset, loc, arch, grid, args...; only_local_halos = false, kw...) +function fill_bottom_and_top_halo!(c, ::PBCT, ::PBCT, size, offset, loc, arch, grid, args...; kw...) c_parent, xy_size, offset = parent_size_and_offset(c, 1, 2, size, offset) launch!(arch, grid, KernelParameters(xy_size, offset), fill_periodic_bottom_and_top_halo!, c_parent, Val(grid.Hz), grid.Nz; kw...) return nothing diff --git a/src/BoundaryConditions/flat_extrapolation_open_boundary_matching_scheme.jl b/src/BoundaryConditions/flat_extrapolation_open_boundary_matching_scheme.jl index b91ac36249..a5bc6435c7 100644 --- a/src/BoundaryConditions/flat_extrapolation_open_boundary_matching_scheme.jl +++ b/src/BoundaryConditions/flat_extrapolation_open_boundary_matching_scheme.jl @@ -128,7 +128,7 @@ end spacing_factor = Δz₁ / (Δz₂ + Δz₃) - gradient_free_ϕ = @inbounds ϕ[i, j, 3] - (ϕ[i, j, 2] - ϕ[i, j, 4]) * spacing_factor + gradient_free_ϕ = @inbounds ϕ[i, j, 3] - (ϕ[i, k, 2] - ϕ[i, j, 4]) * spacing_factor @inbounds ϕ[i, j, 1] = relax(i, j, grid, gradient_free_ϕ, bc, clock, model_fields) @@ -149,4 +149,4 @@ end @inbounds ϕ[i, j, k] = relax(i, j, grid, gradient_free_ϕ, bc, clock, model_fields) return nothing -end +end \ No newline at end of file diff --git a/src/BuoyancyModels/buoyancy.jl b/src/BuoyancyModels/buoyancy.jl index 188a7876a4..3754cb4a4d 100644 --- a/src/BuoyancyModels/buoyancy.jl +++ b/src/BuoyancyModels/buoyancy.jl @@ -32,7 +32,7 @@ model = NonhydrostaticModel(; grid, buoyancy, tracers=:b) NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0) ├── grid: 1×8×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×3×3 halo ├── timestepper: RungeKutta3TimeStepper -├── advection scheme: Centered(order=2) +├── advection scheme: Centered reconstruction order 2 ├── tracers: b ├── closure: Nothing ├── buoyancy: BuoyancyTracer with ĝ = (0.0, -0.707107, -0.707107) diff --git a/src/DistributedComputations/halo_communication.jl b/src/DistributedComputations/halo_communication.jl index 011bb12dbe..872ffa6331 100644 --- a/src/DistributedComputations/halo_communication.jl +++ b/src/DistributedComputations/halo_communication.jl @@ -78,8 +78,8 @@ end ##### Filling halos for halo communication boundary conditions ##### -function tupled_fill_halo_regions!(fields, grid::DistributedGrid, args...; kwargs...) - for field in fields +function tupled_fill_halo_regions!(full_fields, grid::DistributedGrid, args...; kwargs...) + for field in full_fields fill_halo_regions!(field, args...; kwargs...) end end @@ -98,14 +98,12 @@ function fill_halo_regions!(field::DistributedField, args...; kwargs...) kwargs...) end -function fill_halo_regions!(c::OffsetArray, bcs, indices, loc, grid::DistributedGrid, buffers, args...; - fill_boundary_normal_velocities=true, kwargs...) - +function fill_halo_regions!(c::OffsetArray, bcs, indices, loc, grid::DistributedGrid, buffers, args...; fill_boundary_normal_velocities = true, kwargs...) if fill_boundary_normal_velocities - fill_open_boundary_regions!(c, bcs, indices, loc, grid, buffers, args...; kwargs...) + fill_open_boundary_regions!(c, bcs, indices, loc, grid, args...; kwargs...) end - arch = architecture(grid) + arch = architecture(grid) fill_halos!, bcs = permute_boundary_conditions(bcs) number_of_tasks = length(fill_halos!) @@ -138,17 +136,16 @@ end # Syncronous MPI fill_halo_event! cooperative_waitall!(requests) - # Reset MPI tag arch.mpi_tag[] -= arch.mpi_tag[] + recv_from_buffers!(c, buffers, grid, Val(side)) return nothing end # corner passing routine -function fill_corners!(c, connectivity, indices, loc, arch, grid, buffers, args...; - async=false, only_local_halos=false, kw...) +function fill_corners!(c, connectivity, indices, loc, arch, grid, buffers, args...; async = false, only_local_halos = false, kwargs...) # No corner filling needed! only_local_halos && return nothing @@ -159,10 +156,10 @@ function fill_corners!(c, connectivity, indices, loc, arch, grid, buffers, args. requests = MPI.Request[] - reqsw = fill_southwest_halo!(c, connectivity.southwest, indices, loc, arch, grid, buffers, buffers.southwest, args...; kw...) - reqse = fill_southeast_halo!(c, connectivity.southeast, indices, loc, arch, grid, buffers, buffers.southeast, args...; kw...) - reqnw = fill_northwest_halo!(c, connectivity.northwest, indices, loc, arch, grid, buffers, buffers.northwest, args...; kw...) - reqne = fill_northeast_halo!(c, connectivity.northeast, indices, loc, arch, grid, buffers, buffers.northeast, args...; kw...) + reqsw = fill_southwest_halo!(c, connectivity.southwest, indices, loc, arch, grid, buffers, buffers.southwest, args...; kwargs...) + reqse = fill_southeast_halo!(c, connectivity.southeast, indices, loc, arch, grid, buffers, buffers.southeast, args...; kwargs...) + reqnw = fill_northwest_halo!(c, connectivity.northwest, indices, loc, arch, grid, buffers, buffers.northwest, args...; kwargs...) + reqne = fill_northeast_halo!(c, connectivity.northeast, indices, loc, arch, grid, buffers, buffers.northeast, args...; kwargs...) !isnothing(reqsw) && push!(requests, reqsw...) !isnothing(reqse) && push!(requests, reqse...) @@ -190,8 +187,7 @@ cooperative_waitall!(req::Array{MPI.Request}) = MPI.Waitall(req) # There are two additional keyword arguments (with respect to serial `fill_halo_event!`s) that take an effect on `DistributedGrids`: # - only_local_halos: if true, only the local halos are filled, i.e. corresponding to non-communicating boundary conditions # - async: if true, ansynchronous MPI communication is enabled -function fill_halo_event!(c, fill_halos!, bcs, indices, loc, arch, grid::DistributedGrid, buffers, args...; - async = false, only_local_halos = false, kwargs...) +function fill_halo_event!(c, fill_halos!, bcs, indices, loc, arch, grid::DistributedGrid, buffers, args...; async = false, only_local_halos = false, kwargs...) buffer_side = communication_side(Val(fill_halos!)) @@ -206,6 +202,7 @@ function fill_halo_event!(c, fill_halos!, bcs, indices, loc, arch, grid::Distrib offset = fill_halo_offset(size, fill_halos!, indices) requests = fill_halos!(c, bcs..., size, offset, loc, arch, grid, buffers, args...; only_local_halos, kwargs...) + pool_requests_or_complete_comm!(c, arch, grid, buffers, requests, async, buffer_side) return nothing diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl index 8cec48a86c..caa949d0c4 100644 --- a/src/Fields/Fields.jl +++ b/src/Fields/Fields.jl @@ -1,10 +1,10 @@ module Fields -export Face, Center, location +export Face, Center export AbstractField, Field, Average, Integral, Reduction, Accumulation, field export CenterField, XFaceField, YFaceField, ZFaceField export BackgroundField -export interior, data, xnode, ynode, znode +export interior, data, xnode, ynode, znode, location export set!, compute!, @compute, regrid! export VelocityFields, TracerFields, TendencyFields, tracernames export interpolate @@ -15,12 +15,6 @@ using Oceananigans.BoundaryConditions using Oceananigans.Utils import Oceananigans.Architectures: on_architecture -import Oceananigans: location, instantiated_location - -"Return the location `(LX, LY, LZ)` of an `AbstractField{LX, LY, LZ}`." -@inline location(a) = (Nothing, Nothing, Nothing) # used in AbstractOperations for location inference -@inline location(a, i) = location(a)[i] -@inline instantiated_location(a) = (nothing, nothing, nothing) include("abstract_field.jl") include("constant_field.jl") diff --git a/src/Fields/abstract_field.jl b/src/Fields/abstract_field.jl index 8f470663fb..c8c9b7b5f1 100644 --- a/src/Fields/abstract_field.jl +++ b/src/Fields/abstract_field.jl @@ -9,6 +9,7 @@ using Oceananigans.Utils using Oceananigans.Grids: interior_indices, interior_parent_indices import Base: minimum, maximum, extrema +import Oceananigans: location, instantiated_location import Oceananigans.Architectures: architecture, child_architecture import Oceananigans.Grids: interior_x_indices, interior_y_indices, interior_z_indices import Oceananigans.Grids: total_size, topology, nodes, xnodes, ynodes, znodes, node, xnode, ynode, znode @@ -33,6 +34,9 @@ Base.IndexStyle(::AbstractField) = IndexCartesian() ##### AbstractField functionality ##### +"Returns the location `(LX, LY, LZ)` of an `AbstractField{LX, LY, LZ}`." +@inline location(a) = (Nothing, Nothing, Nothing) # used in AbstractOperations for location inference +@inline location(a, i) = location(a)[i] @inline location(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = (LX, LY, LZ) # note no instantiation @inline instantiated_location(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = (LX(), LY(), LZ()) Base.eltype(::AbstractField{<:Any, <:Any, <:Any, <:Any, T}) where T = T diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 974a5fdc12..ee806db849 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -40,7 +40,7 @@ function validate_field_data(loc, data, grid, indices) Fx, Fy, Fz = total_size(grid, loc, indices) if size(data) != (Fx, Fy, Fz) - LX, LY, LZ = loc + LX, LY, LZ = loc e = "Cannot construct field at ($LX, $LY, $LZ) with size(data)=$(size(data)). " * "`data` must have size ($Fx, $Fy, $Fz)." throw(ArgumentError(e)) @@ -179,7 +179,7 @@ function Field(loc::Tuple, return Field(loc, grid, data, boundary_conditions, indices, operand, status) end - + Field(z::ZeroField; kw...) = z Field(f::Field; indices=f.indices) = view(f, indices...) # hmm... @@ -394,7 +394,7 @@ Return a view of `f` that excludes halo points. interior(f::Field) = interior(f.data, location(f), f.grid, f.indices) interior(a::OffsetArray, loc, grid, indices) = interior(a, loc, topology(grid), size(grid), halo_size(grid), indices) interior(f::Field, I...) = view(interior(f), I...) - + # Don't use axes(f) to checkbounds; use axes(f.data) Base.checkbounds(f::Field, I...) = Base.checkbounds(f.data, I...) @@ -419,8 +419,8 @@ total_size(f::Field) = total_size(f.grid, location(f), f.indices) ##### Move Fields between architectures ##### -on_architecture(arch, field::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = - Field{LX, LY, LZ}(on_architecture(arch, field.grid), +on_architecture(arch, field::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = + Field{LX, LY, LZ}(on_architecture(arch, field.grid), on_architecture(arch, field.data), on_architecture(arch, field.boundary_conditions), on_architecture(arch, field.indices), @@ -513,14 +513,14 @@ const ReducedField = Union{XReducedField, XYReducedField, XYZReducedField} -reduced_dimensions(::Field) = () -reduced_dimensions(::XReducedField) = tuple(1) -reduced_dimensions(::YReducedField) = tuple(2) -reduced_dimensions(::ZReducedField) = tuple(3) -reduced_dimensions(::YZReducedField) = (2, 3) -reduced_dimensions(::XZReducedField) = (1, 3) -reduced_dimensions(::XYReducedField) = (1, 2) -reduced_dimensions(::XYZReducedField) = (1, 2, 3) +reduced_dimensions(field::Field) = () +reduced_dimensions(field::XReducedField) = tuple(1) +reduced_dimensions(field::YReducedField) = tuple(2) +reduced_dimensions(field::ZReducedField) = tuple(3) +reduced_dimensions(field::YZReducedField) = (2, 3) +reduced_dimensions(field::XZReducedField) = (1, 3) +reduced_dimensions(field::XYReducedField) = (1, 2) +reduced_dimensions(field::XYZReducedField) = (1, 2, 3) @propagate_inbounds Base.getindex(r::XReducedField, i, j, k) = getindex(r.data, 1, j, k) @propagate_inbounds Base.getindex(r::YReducedField, i, j, k) = getindex(r.data, i, 1, k) @@ -590,15 +590,6 @@ const ReducedAbstractField = Union{XReducedAbstractField, XYReducedAbstractField, XYZReducedAbstractField} -reduced_dimensions(::AbstractField) = () -reduced_dimensions(::XReducedAbstractField) = tuple(1) -reduced_dimensions(::YReducedAbstractField) = tuple(2) -reduced_dimensions(::ZReducedAbstractField) = tuple(3) -reduced_dimensions(::YZReducedAbstractField) = (2, 3) -reduced_dimensions(::XZReducedAbstractField) = (1, 3) -reduced_dimensions(::XYReducedAbstractField) = (1, 2) -reduced_dimensions(::XYZReducedAbstractField) = (1, 2, 3) - # TODO: needs test Statistics.dot(a::Field, b::Field) = mapreduce((x, y) -> x * y, +, interior(a), interior(b)) @@ -614,7 +605,7 @@ const AnyReduction = typeof(Base.any!) initialize_reduced_field!(::SumReduction, f, r::ReducedAbstractField, c) = Base.initarray!(interior(r), f, Base.add_sum, true, interior(c)) initialize_reduced_field!(::ProdReduction, f, r::ReducedAbstractField, c) = Base.initarray!(interior(r), f, Base.mul_prod, true, interior(c)) initialize_reduced_field!(::AllReduction, f, r::ReducedAbstractField, c) = Base.initarray!(interior(r), f, &, true, interior(c)) -initialize_reduced_field!(::AnyReduction, f, r::ReducedAbstractField, c) = Base.initarray!(interior(r), f, |, true, interior(c)) +initialize_reduced_field!(::AnyReduction, f, r::ReducedAbstractField, c) = Base.initarray!(interior(r), f, |, true, interior(c)) initialize_reduced_field!(::MaximumReduction, f, r::ReducedAbstractField, c) = Base.mapfirst!(f, interior(r), interior(c)) initialize_reduced_field!(::MinimumReduction, f, r::ReducedAbstractField, c) = Base.mapfirst!(f, interior(r), interior(c)) @@ -667,7 +658,7 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod) reduction! = Symbol(reduction, '!') @eval begin - + # In-place function Base.$(reduction!)(f::Function, r::ReducedAbstractField, @@ -719,7 +710,7 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod) end end -function Statistics._mean(f, c::AbstractField, ::Colon; condition = nothing, mask = 0) +function Statistics._mean(f, c::AbstractField, ::Colon; condition = nothing, mask = 0) operator = condition_operand(f, c, condition, mask) return sum(operator) / conditional_length(operator) end diff --git a/src/Fields/interpolate.jl b/src/Fields/interpolate.jl index 95de911e8d..6acd5b7ad9 100644 --- a/src/Fields/interpolate.jl +++ b/src/Fields/interpolate.jl @@ -1,4 +1,4 @@ -using Oceananigans.Grids: topology, node, _node, φnode, λnode, +using Oceananigans.Grids: topology, node, _node, xspacings, yspacings, zspacings, λspacings, φspacings, XFlatGrid, YFlatGrid, ZFlatGrid, XYFlatGrid, YZFlatGrid, XZFlatGrid, @@ -7,8 +7,6 @@ using Oceananigans.Grids: topology, node, _node, φnode, λnode, ZRegOrthogonalSphericalShellGrid, RectilinearGrid, LatitudeLongitudeGrid -using Oceananigans.Operators: Δx, Δy, Δz - using Oceananigans.Architectures: child_architecture # GPU-compatile middle point calculation @@ -68,16 +66,16 @@ end @inline function fractional_x_index(x, locs, grid::XRegularRG) x₀ = xnode(1, 1, 1, grid, locs...) - dx = Δx(1, 1, 1, grid, locs...) + Δx = xspacings(grid, locs...) FT = eltype(grid) - return convert(FT, (x - x₀) / dx) + 1 # 1 - based indexing + return convert(FT, (x - x₀) / Δx) end @inline function fractional_x_index(λ, locs, grid::XRegularLLG) λ₀ = λnode(1, 1, 1, grid, locs...) - λ₁ = λnode(2, 1, 1, grid, locs...) + Δλ = λspacings(grid, locs...) FT = eltype(grid) - return convert(FT, (λ - λ₀) / (λ₁ - λ₀)) + 1 # 1 - based indexing + return convert(FT, (λ - λ₀) / Δλ) end @inline function fractional_x_index(x, locs, grid::RectilinearGrid) @@ -85,7 +83,7 @@ end Tx = topology(grid, 1)() Nx = length(loc, Tx, grid.Nx) xn = xnodes(grid, locs...) - return fractional_index(x, xn, Nx) + return fractional_index(x, xn, Nx) - 1 end @inline function fractional_x_index(x, locs, grid::LatitudeLongitudeGrid) @@ -93,23 +91,23 @@ end Tx = topology(grid, 1)() Nx = length(loc, Tx, grid.Nx) xn = λnodes(grid, locs...) - return fractional_index(x, xn, Nx) + return fractional_index(x, xn, Nx) - 1 end @inline fractional_y_index(y, locs, grid::YFlatGrid) = zero(grid) @inline function fractional_y_index(y, locs, grid::YRegularRG) y₀ = ynode(1, 1, 1, grid, locs...) - dy = Δy(1, 1, 1, grid, locs...) + Δy = yspacings(grid, locs...) FT = eltype(grid) - return convert(FT, (y - y₀) / dy) + 1 # 1 - based indexing + return convert(FT, (y - y₀) / Δy) end @inline function fractional_y_index(φ, locs, grid::YRegularLLG) - φ₀ = φnode(1, 1, 1, grid, locs...) - φ₁ = φnode(1, 2, 1, grid, locs...) + φ₀ = φnode(1, 1, 1, grid, locs...) + Δφ = φspacings(grid, locs...) FT = eltype(grid) - return convert(FT, (φ - φ₀) / (φ₁ - φ₀)) + 1 # 1 - based indexing + return convert(FT, (φ - φ₀) / Δφ) end @inline function fractional_y_index(y, locs, grid::RectilinearGrid) @@ -117,7 +115,7 @@ end Ty = topology(grid, 2)() Ny = length(loc, Ty, grid.Ny) yn = ynodes(grid, locs...) - return fractional_index(y, yn, Ny) + return fractional_index(y, yn, Ny) - 1 end @inline function fractional_y_index(y, locs, grid::LatitudeLongitudeGrid) @@ -125,7 +123,7 @@ end Ty = topology(grid, 2)() Ny = length(loc, Ty, grid.Ny) yn = φnodes(grid, locs...) - return fractional_index(y, yn, Ny) + return fractional_index(y, yn, Ny) - 1 end @inline fractional_z_index(z, locs, grid::ZFlatGrid) = zero(grid) @@ -134,8 +132,8 @@ ZRegGrid = Union{ZRegularRG, ZRegularLLG, ZRegOrthogonalSphericalShellGrid} @inline function fractional_z_index(z::FT, locs, grid::ZRegGrid) where FT z₀ = znode(1, 1, 1, grid, locs...) - dz = Δz(1, 1, 1, grid, locs...) - return convert(FT, (z - z₀) / dz) + 1 # 1 - based indexing + Δz = zspacings(grid, locs...) + return convert(FT, (z - z₀) / Δz) end @inline function fractional_z_index(z, locs, grid) @@ -143,7 +141,7 @@ end Tz = topology(grid, 3)() Nz = length(loc, Tz, grid.Nz) zn = znodes(grid, loc) - return fractional_index(z, zn, Nz) + return fractional_index(z, zn, Nz) - 1 end """ @@ -212,7 +210,22 @@ end return (ii, jj, kk) end -@inline _fractional_indices(at_node, grid, ::Nothing, ::Nothing, ::Nothing) = (nothing, nothing, nothing) +""" + truncate_fractional_indices(fi, fj, fk) + +Truncate _fractional_ indices output from fractional indices `fi, fj, fk` to integer indices, dealing +with `nothing` indices for `Flat` domains. +""" +@inline function truncate_fractional_indices(fi, fj, fk) + i = truncate_fractional_index(fi) + j = truncate_fractional_index(fj) + k = truncate_fractional_index(fk) + return (i, j, k) +end + +@inline truncate_fractional_index(::Nothing) = 1 +@inline truncate_fractional_index(fi) = Base.unsafe_trunc(Int, fi) + """ interpolate(at_node, from_field, from_loc, from_grid) @@ -247,12 +260,15 @@ right of `i`, and `ξ` is the fractional distance between `i` and the left bound `i⁻`, such that `ξ ∈ [0, 1)`. """ @inline function interpolator(fractional_idx) + # We use mod and trunc as CUDA.modf is not defined. # For why we use Base.unsafe_trunc instead of trunc see: # https://github.com/CliMA/Oceananigans.jl/issues/828 # https://github.com/CliMA/Oceananigans.jl/pull/997 i⁻ = Base.unsafe_trunc(Int, fractional_idx) - i⁺ = i⁻ + 1 + i⁻ = Int(i⁻ + 1) # convert to "proper" integer? + shift = Int(sign(fractional_idx)) + i⁺ = i⁻ + shift ξ = mod(fractional_idx, 1) return (i⁻, i⁺, ξ) @@ -302,12 +318,6 @@ end @inline flatten_node(x, ::Nothing, z) = flatten_node(x, z) @inline flatten_node(x, y, ::Nothing) = flatten_node(x, y) -@inline flatten_node(x, ::Nothing, ::Nothing) = tuple(x) -@inline flatten_node(::Nothing, y, ::Nothing) = tuple(y) -@inline flatten_node(::Nothing, ::Nothing, z) = tuple(z) - -@inline flatten_node(::Nothing, ::Nothing, ::Nothing) = tuple() - @inline flatten_node(x, y) = (x, y) @inline flatten_node(::Nothing, y) = flatten_node(y) @inline flatten_node(x, ::Nothing) = flatten_node(x) diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl index e311f659f1..e4b8a29cb6 100644 --- a/src/Fields/set!.jl +++ b/src/Fields/set!.jl @@ -3,7 +3,7 @@ using KernelAbstractions: @kernel, @index using Adapt: adapt_structure using Oceananigans.Grids: on_architecture, node_names -using Oceananigans.Architectures: child_architecture, cpu_architecture, device, GPU, CPU +using Oceananigans.Architectures: child_architecture, device, GPU, CPU using Oceananigans.Utils: work_layout ##### @@ -45,16 +45,13 @@ end function set_to_function!(u, f) # Supports serial and distributed - arch = architecture(u) - child_arch = child_architecture(u) + arch = child_architecture(u) # Determine cpu_grid and cpu_u - if child_arch isa GPU - cpu_arch = cpu_architecture(arch) - cpu_grid = on_architecture(cpu_arch, u.grid) - cpu_u = Field(location(u), cpu_grid; indices = indices(u)) - - elseif child_arch isa CPU + if arch isa GPU + cpu_grid = on_architecture(CPU(), u.grid) + cpu_u = Field(location(u), cpu_grid; indices = indices(u)) + elseif arch isa CPU cpu_grid = u.grid cpu_u = u end @@ -68,8 +65,8 @@ function set_to_function!(u, f) catch err u_loc = Tuple(L() for L in location(u)) - arg_str = tuple_string(node_names(u.grid, u_loc...)) - loc_str = tuple_string(location(u)) + arg_str = tuple_string(node_names(u.grid, u_loc...)) + loc_str = tuple_string(location(u)) topo_str = tuple_string(topology(u.grid)) msg = string("An error was encountered within set! while setting the field", '\n', '\n', @@ -84,8 +81,10 @@ function set_to_function!(u, f) end # Transfer data to GPU if u is on the GPU - child_arch isa GPU && set!(u, cpu_u) - + if child_architecture(u) isa GPU + set!(u, cpu_u) + end + return u end diff --git a/src/Forcings/advective_forcing.jl b/src/Forcings/advective_forcing.jl index 03714d4e08..1bf2c2083d 100644 --- a/src/Forcings/advective_forcing.jl +++ b/src/Forcings/advective_forcing.jl @@ -1,4 +1,4 @@ -using Oceananigans.Advection: div_Uc, div_𝐯u, div_𝐯v, div_𝐯w +using Oceananigans.Advection: UpwindBiasedFifthOrder, div_Uc, div_𝐯u, div_𝐯v, div_𝐯w using Oceananigans.Fields: ZeroField, ConstantField using Oceananigans.Utils: SumOfArrays using Adapt diff --git a/src/Grids/Grids.jl b/src/Grids/Grids.jl index 519233561b..fd8aebbf21 100644 --- a/src/Grids/Grids.jl +++ b/src/Grids/Grids.jl @@ -16,10 +16,8 @@ export ξnode, ηnode, rnode export xnode, ynode, znode, λnode, φnode export xnodes, ynodes, znodes, λnodes, φnodes export spacings -export xspacing, yspacing, zspacing, λspacing, φspacing -export xspacings, yspacings, zspacings, λspacings, φspacings +export xspacings, yspacings, zspacings, xspacing, yspacing, zspacing export minimum_xspacing, minimum_yspacing, minimum_zspacing -export static_column_depthᶜᶜᵃ, static_column_depthᶠᶜᵃ, static_column_depthᶜᶠᵃ, static_column_depthᶠᶠᵃ export offset_data, new_data export on_architecture diff --git a/src/Grids/automatic_halo_sizing.jl b/src/Grids/automatic_halo_sizing.jl index 9eb1874654..262728efa1 100644 --- a/src/Grids/automatic_halo_sizing.jl +++ b/src/Grids/automatic_halo_sizing.jl @@ -8,10 +8,10 @@ Example ======= ```jldoctest -using Oceananigans.Advection: Centered(order=4) +using Oceananigans.Advection: CenteredFourthOrder using Oceananigans.Grids: required_halo_size_x -required_halo_size_x(Centered(order=4)) +required_halo_size_x(CenteredFourthOrder()) # output 2 @@ -28,10 +28,10 @@ Example ======= ```jldoctest -using Oceananigans.Advection: Centered(order=4) +using Oceananigans.Advection: CenteredFourthOrder using Oceananigans.Grids: required_halo_size_y -required_halo_size_y(Centered(order=4)) +required_halo_size_y(CenteredFourthOrder()) # output 2 @@ -48,10 +48,10 @@ Example ======= ```jldoctest -using Oceananigans.Advection: Centered(order=4) +using Oceananigans.Advection: CenteredFourthOrder using Oceananigans.Grids: required_halo_size_z -required_halo_size_z(Centered(order=4)) +required_halo_size_z(CenteredFourthOrder()) # output 2 diff --git a/src/Grids/grid_generation.jl b/src/Grids/grid_generation.jl index 84212fb580..b0ba1ad7dc 100644 --- a/src/Grids/grid_generation.jl +++ b/src/Grids/grid_generation.jl @@ -129,3 +129,4 @@ generate_coordinate(FT, ::Flat, N, H, c::Number, coordinate_name, arch) = generate_coordinate(FT, ::Flat, N, H, ::Nothing, coordinate_name, arch) = FT(1), nothing, nothing, FT(1), FT(1) + diff --git a/src/Grids/grid_utils.jl b/src/Grids/grid_utils.jl index 178d7ab87c..c881ebeb86 100644 --- a/src/Grids/grid_utils.jl +++ b/src/Grids/grid_utils.jl @@ -314,15 +314,6 @@ coordinate_summary(topo, Δ::Union{AbstractVector, AbstractMatrix}, name) = name, prettysummary(minimum(parent(Δ))), name, prettysummary(maximum(parent(Δ)))) -##### -##### Static column depth -##### - -@inline static_column_depthᶜᶜᵃ(i, j, grid) = grid.Lz -@inline static_column_depthᶜᶠᵃ(i, j, grid) = grid.Lz -@inline static_column_depthᶠᶜᵃ(i, j, grid) = grid.Lz -@inline static_column_depthᶠᶠᵃ(i, j, grid) = grid.Lz - ##### ##### Spherical geometry ##### diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl index bd411ea5d1..98dc1bc366 100644 --- a/src/Grids/latitude_longitude_grid.jl +++ b/src/Grids/latitude_longitude_grid.jl @@ -278,7 +278,7 @@ function validate_lat_lon_grid_args(topology, size, halo, FT, latitude, longitud φ₁, φ₂ = get_domain_extent(latitude, Nφ) -90 <= φ₁ || throw(ArgumentError("The southernmost latitude cannot be less than -90 degrees.")) - φ₂ <= 90 || throw(ArgumentError("The northern latitude cannot be greater than 90 degrees.")) + φ₂ <= 90 || throw(ArgumentError("The northern latitude cannot be less than -90 degrees.")) φ₁ <= φ₂ || throw(ArgumentError("Latitudes must increase south to north.")) if TX == Flat || TY == Flat @@ -332,9 +332,9 @@ function Base.show(io::IO, grid::LatitudeLongitudeGrid, withsummary=true) "└── ", z_summary) end -@inline x_domain(grid::LLG) = domain(topology(grid, 1)(), grid.Nx, grid.λᶠᵃᵃ) -@inline y_domain(grid::LLG) = domain(topology(grid, 2)(), grid.Ny, grid.φᵃᶠᵃ) -@inline z_domain(grid::LLG) = domain(topology(grid, 3)(), grid.Nz, grid.zᵃᵃᶠ) +@inline x_domain(grid::LLG{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = domain(TX, grid.Nx, grid.λᶠᵃᵃ) +@inline y_domain(grid::LLG{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = domain(TY, grid.Ny, grid.φᵃᶠᵃ) +@inline z_domain(grid::LLG{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = domain(TZ, grid.Nz, grid.zᵃᵃᶠ) @inline cpu_face_constructor_x(grid::XRegularLLG) = x_domain(grid) @inline cpu_face_constructor_y(grid::YRegularLLG) = y_domain(grid) @@ -655,11 +655,6 @@ end @inline xnodes(grid::LLG, ℓx, ℓy, ℓz; with_halos=false) = xnodes(grid, ℓx, ℓy; with_halos) @inline ynodes(grid::LLG, ℓx, ℓy, ℓz; with_halos=false) = ynodes(grid, ℓy; with_halos) -@inline λnodes(grid::LLG, ℓx::F; with_halos=false) = _property(grid.λᶠᵃᵃ, ℓx, topology(grid, 1), size(grid, 1), with_halos) -@inline λnodes(grid::LLG, ℓx::C; with_halos=false) = _property(grid.λᶜᵃᵃ, ℓx, topology(grid, 1), size(grid, 1), with_halos) -@inline φnodes(grid::LLG, ℓy::F; with_halos=false) = _property(grid.φᵃᶠᵃ, ℓy, topology(grid, 2), size(grid, 2), with_halos) -@inline φnodes(grid::LLG, ℓy::C; with_halos=false) = _property(grid.φᵃᶜᵃ, ℓy, topology(grid, 2), size(grid, 2), with_halos) - # Generalized coordinates @inline ξnodes(grid::LLG, ℓx; kwargs...) = λnodes(grid, ℓx; kwargs...) @inline ηnodes(grid::LLG, ℓy; kwargs...) = φnodes(grid, ℓy; kwargs...) @@ -669,10 +664,76 @@ end @inline ηnodes(grid::LLG, ℓx, ℓy, ℓz; kwargs...) = φnodes(grid, ℓy; kwargs...) @inline rnodes(grid::LLG, ℓx, ℓy, ℓz; kwargs...) = znodes(grid, ℓz; kwargs...) +##### +##### Grid spacings in x, y, z (in meters) +##### + +@inline xspacings(grid::LLG, ℓx::C, ℓy::C; with_halos=false) = _property(grid.Δxᶜᶜᵃ, ℓx, ℓy, + topology(grid, 1), topology(grid, 2), + size(grid, 1), size(grid, 2), with_halos) + +@inline xspacings(grid::LLG, ℓx::C, ℓy::F; with_halos=false) = _property(grid.Δxᶜᶠᵃ, ℓx, ℓy, + topology(grid, 1), topology(grid, 2), + size(grid, 1), size(grid, 2), with_halos) + +@inline xspacings(grid::LLG, ℓx::F, ℓy::C; with_halos=false) = _property(grid.Δxᶠᶜᵃ, ℓx, ℓy, + topology(grid, 1), topology(grid, 2), + size(grid, 1), size(grid, 2), with_halos) + +@inline xspacings(grid::LLG, ℓx::F, ℓy::F; with_halos=false) = _property(grid.Δxᶠᶠᵃ, ℓx, ℓy, + topology(grid, 1), topology(grid, 2), + size(grid, 1), size(grid, 2), with_halos) + +@inline xspacings(grid::HRegularLLG, ℓx::C, ℓy::C; with_halos=false) = _property(grid.Δxᶜᶜᵃ, ℓy, topology(grid, 2), + size(grid, 2), with_halos) + +@inline xspacings(grid::HRegularLLG, ℓx::C, ℓy::F; with_halos=false) = _property(grid.Δxᶜᶠᵃ, ℓy, topology(grid, 2), + size(grid, 2), with_halos) + +@inline xspacings(grid::HRegularLLG, ℓx::F, ℓy::C; with_halos=false) = _property(grid.Δxᶠᶜᵃ, ℓy, topology(grid, 2), + size(grid, 2), with_halos) + +@inline xspacings(grid::HRegularLLG, ℓx::F, ℓy::F; with_halos=false) = _property(grid.Δxᶠᶠᵃ, ℓy, topology(grid, 2), + size(grid, 2), with_halos) + + +@inline yspacings(grid::YNonRegularLLG, ℓx::C, ℓy::F; with_halos=false) = _property(grid.Δyᶜᶠᵃ, ℓy, topoloy(grid, 2), + size(grid, 2), with_halos) + +@inline yspacings(grid::YNonRegularLLG, ℓx::F, ℓy::C; with_halos=false) = _property(grid.Δyᶠᶜᵃ, ℓy, topoloy(grid, 2), + size(grid, 2), with_halos) + +@inline yspacings(grid::YRegularLLG, ℓx, ℓy; with_halos=false) = yspacings(grid, ℓy; with_halos) +@inline yspacings(grid, ℓy::C; kwargs...) = grid.Δyᶠᶜᵃ +@inline yspacings(grid, ℓy::F; kwargs...) = grid.Δyᶜᶠᵃ + +@inline xspacings(grid::LLG, ℓx, ℓy, ℓz; kwargs...) = xspacings(grid, ℓx, ℓy; kwargs...) +@inline yspacings(grid::LLG, ℓx, ℓy, ℓz; kwargs...) = yspacings(grid, ℓx, ℓy; kwargs...) +@inline zspacings(grid::LLG, ℓx, ℓy, ℓz; kwargs...) = zspacings(grid, ℓz; kwargs...) + +##### +##### Grid spacings in λ, φ (in degrees) +##### + +@inline λnodes(grid::LLG, ℓx::F; with_halos=false) = _property(grid.λᶠᵃᵃ, ℓx, topology(grid, 1), size(grid, 1), with_halos) +@inline λnodes(grid::LLG, ℓx::C; with_halos=false) = _property(grid.λᶜᵃᵃ, ℓx, topology(grid, 1), size(grid, 1), with_halos) +@inline φnodes(grid::LLG, ℓy::F; with_halos=false) = _property(grid.φᵃᶠᵃ, ℓy, topology(grid, 2), size(grid, 2), with_halos) +@inline φnodes(grid::LLG, ℓy::C; with_halos=false) = _property(grid.φᵃᶜᵃ, ℓy, topology(grid, 2), size(grid, 2), with_halos) + ##### ##### Grid spacings ##### +@inline λspacings(grid::LLG, ℓx::C; with_halos=false) = _property(grid.Δλᶜᵃᵃ, ℓx, topology(grid, 1), size(grid, 1), with_halos) +@inline λspacings(grid::LLG, ℓx::F; with_halos=false) = _property(grid.Δλᶠᵃᵃ, ℓx, topology(grid, 1), size(grid, 1), with_halos) +@inline φspacings(grid::LLG, ℓy::C; with_halos=false) = _property(grid.Δφᵃᶜᵃ, ℓy, topology(grid, 2), size(grid, 2), with_halos) +@inline φspacings(grid::LLG, ℓy::F; with_halos=false) = _property(grid.Δφᵃᶠᵃ, ℓy, topology(grid, 2), size(grid, 2), with_halos) +@inline zspacings(grid::LLG, ℓz::C; with_halos=false) = _property(grid.Δzᵃᵃᶜ, ℓz, topology(grid, 3), size(grid, 3), with_halos) +@inline zspacings(grid::LLG, ℓz::F; with_halos=false) = _property(grid.Δzᵃᵃᶠ, ℓz, topology(grid, 3), size(grid, 3), with_halos) + +@inline λspacings(grid::LLG, ℓx, ℓy, ℓz; kwargs...) = λspacings(grid, ℓx; kwargs...) +@inline φspacings(grid::LLG, ℓx, ℓy, ℓz; kwargs...) = φspacings(grid, ℓy; kwargs...) + @inline λspacing(i, grid::LLG, ::C) = @inbounds grid.Δλᶜᵃᵃ[i] @inline λspacing(i, grid::LLG, ::F) = @inbounds grid.Δλᶠᵃᵃ[i] @inline λspacing(i, grid::XRegularLLG, ::C) = grid.Δλᶜᵃᵃ @@ -685,10 +746,3 @@ end @inline λspacing(i, j, k, grid::LLG, ℓx, ℓy, ℓz) = λspacing(i, grid, ℓx) @inline φspacing(i, j, k, grid::LLG, ℓx, ℓy, ℓz) = φspacing(j, grid, ℓy) - -@inline xspacings(grid::LLG, ℓx, ℓy) = xspacings(grid, ℓx, ℓy, nothing) -@inline yspacings(grid::LLG, ℓx, ℓy) = yspacings(grid, ℓx, ℓy, nothing) -@inline zspacings(grid::LLG, ℓz) = zspacings(grid, nothing, nothing, ℓz) - -@inline λspacings(grid::LLG, ℓx) = λspacings(grid, ℓx, nothing, nothing) -@inline φspacings(grid::LLG, ℓy) = φspacings(grid, nothing, ℓy, nothing) diff --git a/src/Grids/nodes_and_spacings.jl b/src/Grids/nodes_and_spacings.jl index fd814d94e5..e348c0cb47 100644 --- a/src/Grids/nodes_and_spacings.jl +++ b/src/Grids/nodes_and_spacings.jl @@ -161,36 +161,96 @@ nodes(grid::AbstractGrid, (ℓx, ℓy, ℓz); reshape=false, with_halos=false) = ##### << Spacings >> ##### -# placeholders -# see Oceananigans/AbstractOperations/grid_metrics.jl for definitions +# placeholders; see Oceananigans.Operators for x/y/zspacing definitions function xspacing end function yspacing end function zspacing end -function λspacing end -function φspacing end -function xspacings end -function yspacings end -function zspacings end -function λspacings end -function φspacings end +""" + xspacings(grid, ℓx, ℓy, ℓz; with_halos=true) + +Return the spacings over the interior nodes on `grid` in the ``x``-direction for the location `ℓx`, +`ℓy`, `ℓz`. For `Bounded` directions, `Face` nodes include the boundary points. + +```jldoctest xspacings +julia> using Oceananigans + +julia> grid = LatitudeLongitudeGrid(size=(8, 15, 10), longitude=(-20, 60), latitude=(-10, 50), z=(-100, 0)); + +julia> xspacings(grid, Center(), Face(), Center()) +16-element view(OffsetArray(::Vector{Float64}, -2:18), 1:16) with eltype Float64: + 1.0950562585518518e6 + 1.1058578920188267e6 + 1.1112718969963323e6 + 1.1112718969963323e6 + 1.1058578920188267e6 + 1.0950562585518518e6 + 1.0789196210678827e6 + 1.0575265956426917e6 + 1.0309814069457315e6 + 999413.38046802 + 962976.3124613502 + 921847.720658409 + 876227.979424229 + 826339.3435524226 + 772424.8654621692 + 714747.2110712599 +``` +""" +@inline xspacings(grid, ℓx, ℓy, ℓz; with_halos=true) = xspacings(grid, ℓx; with_halos) + +""" + yspacings(grid, ℓx, ℓy, ℓz; with_halos=true) + +Return the spacings over the interior nodes on `grid` in the ``y``-direction for the location `ℓx`, +`ℓy`, `ℓz`. For `Bounded` directions, `Face` nodes include the boundary points. + +```jldoctest yspacings +julia> using Oceananigans + +julia> grid = LatitudeLongitudeGrid(size=(20, 15, 10), longitude=(0, 20), latitude=(-15, 15), z=(-100, 0)); + +julia> yspacings(grid, Center(), Center(), Center()) +222389.85328911748 +``` +""" +@inline yspacings(grid, ℓx, ℓy, ℓz; with_halos=true) = yspacings(grid, ℓy; with_halos) + +""" + zspacings(grid, ℓx, ℓy, ℓz; with_halos=true) + +Return the spacings over the interior nodes on `grid` in the ``z``-direction for the location `ℓx`, +`ℓy`, `ℓz`. For `Bounded` directions, `Face` nodes include the boundary points. + +```jldoctest zspacings +julia> using Oceananigans + +julia> grid = LatitudeLongitudeGrid(size=(20, 15, 10), longitude=(0, 20), latitude=(-15, 15), z=(-100, 0)); + +julia> zspacings(grid, Center(), Center(), Center()) +10.0 +``` +""" +@inline zspacings(grid, ℓx, ℓy, ℓz; with_halos=true) = zspacings(grid, ℓz; with_halos) destantiate(::Face) = Face destantiate(::Center) = Center -spacings_function(::Val{:x}) = xspacings -spacings_function(::Val{:y}) = yspacings -spacings_function(::Val{:z}) = zspacings -spacings_function(::Val{:λ}) = λspacings -spacings_function(::Val{:φ}) = φspacings +spacing_function(::Val{:x}) = xspacing +spacing_function(::Val{:y}) = yspacing +spacing_function(::Val{:z}) = zspacing function minimum_spacing(s, grid, ℓx, ℓy, ℓz) - spacings = spacings_function(s) - return minimum(spacings(grid, ℓx, ℓy, ℓz)) + spacing = spacing_function(s) + LX, LY, LZ = map(destantiate, (ℓx, ℓy, ℓz)) + Δ = KernelFunctionOperation{LX, LY, LZ}(spacing, grid, ℓx, ℓy, ℓz) + + return minimum(Δ) end """ minimum_xspacing(grid, ℓx, ℓy, ℓz) + minimum_xspacing(grid) = minimum_xspacing(grid, Center(), Center(), Center()) Return the minimum spacing for `grid` in ``x`` direction at location `ℓx, ℓy, ℓz`. @@ -205,11 +265,11 @@ julia> minimum_xspacing(grid, Center(), Center(), Center()) 0.5 ``` """ -minimum_xspacing(grid, loc...) = minimum(xspacings(grid, loc...)) -minimum_xspacing(grid) = minimum(xspacings(grid)) - +minimum_xspacing(grid, ℓx, ℓy, ℓz) = minimum_spacing(Val(:x), grid, ℓx, ℓy, ℓz) +minimum_xspacing(grid) = minimum_spacing(Val(:x), grid, Center(), Center(), Center()) """ minimum_yspacing(grid, ℓx, ℓy, ℓz) + minimum_yspacing(grid) = minimum_yspacing(grid, Center(), Center(), Center()) Return the minimum spacing for `grid` in ``y`` direction at location `ℓx, ℓy, ℓz`. @@ -224,8 +284,8 @@ julia> minimum_yspacing(grid, Center(), Center(), Center()) 0.25 ``` """ -minimum_yspacing(grid, loc...) = minimum(yspacings(grid, loc...)) -minimum_yspacing(grid) = minimum(yspacings(grid)) +minimum_yspacing(grid, ℓx, ℓy, ℓz) = minimum_spacing(Val(:y), grid, ℓx, ℓy, ℓz) +minimum_yspacing(grid) = minimum_spacing(Val(:y), grid, Center(), Center(), Center()) """ minimum_zspacing(grid, ℓx, ℓy, ℓz) @@ -244,5 +304,6 @@ julia> minimum_zspacing(grid, Center(), Center(), Center()) 0.125 ``` """ -minimum_zspacing(grid, loc...) = minimum(zspacings(grid, loc...)) -minimum_zspacing(grid) = minimum(zspacings(grid)) +minimum_zspacing(grid, ℓx, ℓy, ℓz) = minimum_spacing(Val(:z), grid, ℓx, ℓy, ℓz) +minimum_zspacing(grid) = minimum_spacing(Val(:z), grid, Center(), Center(), Center()) + diff --git a/src/Grids/orthogonal_spherical_shell_grid.jl b/src/Grids/orthogonal_spherical_shell_grid.jl index 7a30c2b7a4..1567b2a5fb 100644 --- a/src/Grids/orthogonal_spherical_shell_grid.jl +++ b/src/Grids/orthogonal_spherical_shell_grid.jl @@ -76,7 +76,7 @@ end const OSSG = OrthogonalSphericalShellGrid const ZRegOSSG = OrthogonalSphericalShellGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} const ZRegOrthogonalSphericalShellGrid = ZRegOSSG -const ConformalCubedSpherePanel = OrthogonalSphericalShellGrid{<:Any, FullyConnected, FullyConnected, <:Any, <:Any, <:Any, <:Any, <:CubedSphereConformalMapping} +const ConformalCubedSpherePanel = OrthogonalSphericalShellGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:CubedSphereConformalMapping} # convenience constructor for OSSG without any conformal_mapping properties OrthogonalSphericalShellGrid(architecture, Nx, Ny, Nz, Hx, Hy, Hz, Lz, @@ -100,7 +100,7 @@ OrthogonalSphericalShellGrid(architecture, Nx, Ny, Nz, Hx, Hy, Hz, Lz, halo = (1, 1, 1), rotation = nothing) -Create a `OrthogonalSphericalShellGrid` that represents a section of a sphere after it has been +Create a `OrthogonalSphericalShellGrid` that represents a section of a sphere after it has been conformally mapped from the face of a cube. The cube's coordinates are `ξ` and `η` (which, by default, both take values in the range ``[-1, 1]``. @@ -176,7 +176,7 @@ function conformal_cubed_sphere_panel(architecture::AbstractArchitecture = CPU() halo = (1, 1, 1), rotation = nothing) - if architecture == GPU() && !has_cuda() + if architecture == GPU() && !has_cuda() throw(ArgumentError("Cannot create a GPU grid. No CUDA-enabled GPU was detected!")) end @@ -364,9 +364,9 @@ function conformal_cubed_sphere_panel(architecture::AbstractArchitecture = CPU() j = 1 Δyᶠᶠᵃ[i, j] = 2haversine((λᶠᶜᵃ[i, j], φᶠᶜᵃ[i, j]), (λᶠᶠᵃ[i, j ], φᶠᶠᵃ[i, j ]), radius) end - + for i in 1:Nξ+1 - j = Nη+1 + j = Nη+1 Δyᶠᶠᵃ[i, j] = 2haversine((λᶠᶠᵃ[i, j], φᶠᶠᵃ[i, j]), (λᶠᶜᵃ[i, j-1], φᶠᶜᵃ[i, j-1]), radius) end end @@ -1201,9 +1201,31 @@ rname(::OSSG) = :z ##### Grid spacings in x, y, z (in meters) ##### -@inline xspacings(grid::OSSG, ℓx, ℓy) = xspacings(grid, ℓx, ℓy, nothing) -@inline yspacings(grid::OSSG, ℓx, ℓy) = yspacings(grid, ℓx, ℓy, nothing) -@inline zspacings(grid::OSSG, ℓz) = zspacings(grid, nothing, nothing, ℓz) - -@inline λspacings(grid::OSSG, ℓx, ℓy) = λspacings(grid, ℓx, ℓy, nothing) -@inline φspacings(grid::OSSG, ℓx, ℓy) = φspacings(grid, ℓx, ℓy, nothing) +@inline xspacings(grid::OSSG, ℓx::Center, ℓy::Center; with_halos=false) = + with_halos ? grid.Δxᶜᶜᵃ : view(grid.Δxᶜᶜᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline xspacings(grid::OSSG, ℓx::Face , ℓy::Center; with_halos=false) = + with_halos ? grid.Δxᶠᶜᵃ : view(grid.Δxᶠᶜᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline xspacings(grid::OSSG, ℓx::Center, ℓy::Face ; with_halos=false) = + with_halos ? grid.Δxᶜᶠᵃ : view(grid.Δxᶜᶠᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline xspacings(grid::OSSG, ℓx::Face , ℓy::Face ; with_halos=false) = + with_halos ? grid.Δxᶠᶠᵃ : view(grid.Δxᶠᶠᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) + +@inline yspacings(grid::OSSG, ℓx::Center, ℓy::Center; with_halos=false) = + with_halos ? grid.Δyᶜᶜᵃ : view(grid.Δyᶜᶜᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline yspacings(grid::OSSG, ℓx::Face , ℓy::Center; with_halos=false) = + with_halos ? grid.Δyᶠᶜᵃ : view(grid.Δyᶠᶜᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline yspacings(grid::OSSG, ℓx::Center, ℓy::Face ; with_halos=false) = + with_halos ? grid.Δyᶜᶠᵃ : view(grid.Δyᶜᶠᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline yspacings(grid::OSSG, ℓx::Face , ℓy::Face ; with_halos=false) = + with_halos ? grid.Δyᶠᶠᵃ : view(grid.Δyᶠᶠᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) + +@inline zspacings(grid::OSSG, ℓz::Center; with_halos=false) = with_halos ? grid.Δzᵃᵃᶜ : + view(grid.Δzᵃᵃᶜ, interior_indices(ℓz, topology(grid, 3)(), grid.Nz)) +@inline zspacings(grid::ZRegOSSG, ℓz::Center; with_halos=false) = grid.Δzᵃᵃᶜ +@inline zspacings(grid::OSSG, ℓz::Face; with_halos=false) = with_halos ? grid.Δzᵃᵃᶠ : + view(grid.Δzᵃᵃᶠ, interior_indices(ℓz, topology(grid, 3)(), grid.Nz)) +@inline zspacings(grid::ZRegOSSG, ℓz::Face; with_halos=false) = grid.Δzᵃᵃᶠ + +@inline xspacings(grid::OSSG, ℓx, ℓy, ℓz; with_halos=false) = xspacings(grid, ℓx, ℓy; with_halos) +@inline yspacings(grid::OSSG, ℓx, ℓy, ℓz; with_halos=false) = yspacings(grid, ℓx, ℓy; with_halos) +@inline zspacings(grid::OSSG, ℓx, ℓy, ℓz; with_halos=false) = zspacings(grid, ℓz; with_halos) diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index f148625f15..25a30e7cbc 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -518,12 +518,19 @@ const C = Center @inline ηnodes(grid::RG, ℓx, ℓy, ℓz; kwargs...) = ynodes(grid, ℓy; kwargs...) @inline rnodes(grid::RG, ℓx, ℓy, ℓz; kwargs...) = znodes(grid, ℓz; kwargs...) -@inline isrectilinear(::RG) = true - ##### -##### Grid-specific grid spacings +##### Grid spacings ##### -@inline xspacings(grid::RG, ℓx) = xspacings(grid, ℓx, nothing, nothing) -@inline yspacings(grid::RG, ℓy) = yspacings(grid, nothing, ℓy, nothing) -@inline zspacings(grid::RG, ℓz) = zspacings(grid, nothing, nothing, ℓz) +@inline xspacings(grid::RG, ℓx::C; with_halos=false) = _property(grid.Δxᶜᵃᵃ, ℓx, topology(grid, 1), size(grid, 1), with_halos) +@inline xspacings(grid::RG, ℓx::F; with_halos=false) = _property(grid.Δxᶠᵃᵃ, ℓx, topology(grid, 1), size(grid, 1), with_halos) +@inline yspacings(grid::RG, ℓy::C; with_halos=false) = _property(grid.Δyᵃᶜᵃ, ℓy, topology(grid, 2), size(grid, 2), with_halos) +@inline yspacings(grid::RG, ℓy::F; with_halos=false) = _property(grid.Δyᵃᶠᵃ, ℓy, topology(grid, 2), size(grid, 2), with_halos) +@inline zspacings(grid::RG, ℓz::C; with_halos=false) = _property(grid.Δzᵃᵃᶜ, ℓz, topology(grid, 3), size(grid, 3), with_halos) +@inline zspacings(grid::RG, ℓz::F; with_halos=false) = _property(grid.Δzᵃᵃᶠ, ℓz, topology(grid, 3), size(grid, 3), with_halos) + +@inline xspacings(grid::RG, ℓx, ℓy, ℓz; kwargs...) = xspacings(grid, ℓx; kwargs...) +@inline yspacings(grid::RG, ℓx, ℓy, ℓz; kwargs...) = yspacings(grid, ℓy; kwargs...) +@inline zspacings(grid::RG, ℓx, ℓy, ℓz; kwargs...) = zspacings(grid, ℓz; kwargs...) + +@inline isrectilinear(::RG) = true diff --git a/src/ImmersedBoundaries/ImmersedBoundaries.jl b/src/ImmersedBoundaries/ImmersedBoundaries.jl index 2b45f4af9a..b15f6ab591 100644 --- a/src/ImmersedBoundaries/ImmersedBoundaries.jl +++ b/src/ImmersedBoundaries/ImmersedBoundaries.jl @@ -17,23 +17,231 @@ import Base: show, summary import Oceananigans.Grids: cpu_face_constructor_x, cpu_face_constructor_y, cpu_face_constructor_z, x_domain, y_domain, z_domain -import Oceananigans.Grids: architecture, with_halo, inflate_halo_size_one_dimension, +import Oceananigans.Grids: architecture, on_architecture, with_halo, inflate_halo_size_one_dimension, xnode, ynode, znode, λnode, φnode, node, ξnode, ηnode, rnode, ξname, ηname, rname, node_names, xnodes, ynodes, znodes, λnodes, φnodes, nodes, ξnodes, ηnodes, rnodes, - static_column_depthᶜᶜᵃ, static_column_depthᶠᶜᵃ, static_column_depthᶜᶠᵃ, static_column_depthᶠᶠᵃ, inactive_cell -import Oceananigans.Architectures: on_architecture - import Oceananigans.Fields: fractional_x_index, fractional_y_index, fractional_z_index -include("immersed_boundary_grid.jl") -include("immersed_boundary_interface.jl") -include("immersed_boundary_nodes.jl") +""" + abstract type AbstractImmersedBoundary + +Abstract supertype for immersed boundary grids. +""" +abstract type AbstractImmersedBoundary end + +##### +##### ImmersedBoundaryGrid +##### + +struct ImmersedBoundaryGrid{FT, TX, TY, TZ, G, I, M, S, Arch} <: AbstractGrid{FT, TX, TY, TZ, Arch} + architecture :: Arch + underlying_grid :: G + immersed_boundary :: I + interior_active_cells :: M + active_z_columns :: S + + # Internal interface + function ImmersedBoundaryGrid{TX, TY, TZ}(grid::G, ib::I, mi::M, ms::S) where {TX, TY, TZ, G <: AbstractUnderlyingGrid, I, M, S} + FT = eltype(grid) + arch = architecture(grid) + Arch = typeof(arch) + return new{FT, TX, TY, TZ, G, I, M, S, Arch}(arch, grid, ib, mi, ms) + end + + # Constructor with no active map + function ImmersedBoundaryGrid{TX, TY, TZ}(grid::G, ib::I) where {TX, TY, TZ, G <: AbstractUnderlyingGrid, I} + FT = eltype(grid) + arch = architecture(grid) + Arch = typeof(arch) + return new{FT, TX, TY, TZ, G, I, Nothing, Nothing, Arch}(arch, grid, ib, nothing, nothing) + end +end + +const IBG = ImmersedBoundaryGrid + +@inline Base.getproperty(ibg::IBG, property::Symbol) = get_ibg_property(ibg, Val(property)) +@inline get_ibg_property(ibg::IBG, ::Val{property}) where property = getfield(getfield(ibg, :underlying_grid), property) +@inline get_ibg_property(ibg::IBG, ::Val{:immersed_boundary}) = getfield(ibg, :immersed_boundary) +@inline get_ibg_property(ibg::IBG, ::Val{:underlying_grid}) = getfield(ibg, :underlying_grid) +@inline get_ibg_property(ibg::IBG, ::Val{:interior_active_cells}) = getfield(ibg, :interior_active_cells) +@inline get_ibg_property(ibg::IBG, ::Val{:active_z_columns}) = getfield(ibg, :active_z_columns) + +@inline architecture(ibg::IBG) = architecture(ibg.underlying_grid) + +@inline x_domain(ibg::IBG) = x_domain(ibg.underlying_grid) +@inline y_domain(ibg::IBG) = y_domain(ibg.underlying_grid) +@inline z_domain(ibg::IBG) = z_domain(ibg.underlying_grid) + +Adapt.adapt_structure(to, ibg::IBG{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = + ImmersedBoundaryGrid{TX, TY, TZ}(adapt(to, ibg.underlying_grid), + adapt(to, ibg.immersed_boundary), + nothing, + nothing) + +with_halo(halo, ibg::ImmersedBoundaryGrid) = + ImmersedBoundaryGrid(with_halo(halo, ibg.underlying_grid), ibg.immersed_boundary) + +# ImmersedBoundaryGrids require an extra halo point to check the "inactivity" of a `Face` node at N + H +# (which requires checking `Center` nodes at N + H and N + H + 1) +inflate_halo_size_one_dimension(req_H, old_H, _, ::IBG) = max(req_H + 1, old_H) +inflate_halo_size_one_dimension(req_H, old_H, ::Type{Flat}, ::IBG) = 0 + +# Defining the bottom +@inline z_bottom(i, j, grid) = znode(i, j, 1, grid, c, c, f) +@inline z_bottom(i, j, ibg::IBG) = error("The function `bottom` has not been defined for $(summary(ibg))!") + +function Base.summary(grid::ImmersedBoundaryGrid) + FT = eltype(grid) + TX, TY, TZ = topology(grid) + + return string(size_summary(size(grid)), + " ImmersedBoundaryGrid{$FT, $TX, $TY, $TZ} on ", summary(architecture(grid)), + " with ", size_summary(halo_size(grid)), " halo") +end + +function show(io::IO, ibg::ImmersedBoundaryGrid) + print(io, summary(ibg), ":", "\n", + "├── immersed_boundary: ", summary(ibg.immersed_boundary), "\n", + "├── underlying_grid: ", summary(ibg.underlying_grid), "\n") + + return show(io, ibg.underlying_grid, false) +end + +##### +##### Interface for immersed_boundary +##### + +""" + immersed_cell(i, j, k, grid) + +Return true if a `cell` is "completely" immersed, and thus +is not part of the prognostic state. +""" +@inline immersed_cell(i, j, k, grid) = false + +# Unpack to make defining new immersed boundaries more convenient +@inline immersed_cell(i, j, k, grid::ImmersedBoundaryGrid) = + immersed_cell(i, j, k, grid.underlying_grid, grid.immersed_boundary) + +""" + inactive_cell(i, j, k, grid::ImmersedBoundaryGrid) + +Return `true` if the tracer cell at `i, j, k` either (i) lies outside the `Bounded` domain +or (ii) lies within the immersed region of `ImmersedBoundaryGrid`. + +Example +======= + +Consider the configuration + +``` + Immersed Fluid + =========== ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ + + c c + i-1 i + + | ========= | | + × === ∘ === × ∘ × + | ========= | | + +i-1 i + f f f +``` + +We then have + +* `inactive_node(i, 1, 1, grid, f, c, c) = false` + +As well as + +* `inactive_node(i, 1, 1, grid, c, c, c) = false` +* `inactive_node(i-1, 1, 1, grid, c, c, c) = true` +* `inactive_node(i-1, 1, 1, grid, f, c, c) = true` +""" +@inline inactive_cell(i, j, k, ibg::IBG) = immersed_cell(i, j, k, ibg) | inactive_cell(i, j, k, ibg.underlying_grid) + +# Isolate periphery of the immersed boundary +@inline immersed_peripheral_node(i, j, k, ibg::IBG, LX, LY, LZ) = peripheral_node(i, j, k, ibg, LX, LY, LZ) & + !peripheral_node(i, j, k, ibg.underlying_grid, LX, LY, LZ) + +@inline immersed_inactive_node(i, j, k, ibg::IBG, LX, LY, LZ) = inactive_node(i, j, k, ibg, LX, LY, LZ) & + !inactive_node(i, j, k, ibg.underlying_grid, LX, LY, LZ) + +##### +##### Utilities +##### + +const c = Center() +const f = Face() + +@inline Base.zero(ibg::IBG) = zero(ibg.underlying_grid) + +@inline xnode(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = xnode(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) +@inline ynode(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = ynode(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) +@inline znode(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = znode(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) + +@inline λnode(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = λnode(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) +@inline φnode(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = φnode(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) + +@inline ξnode(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = ξnode(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) +@inline ηnode(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = ηnode(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) +@inline rnode(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = rnode(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) + +@inline node(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = node(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) + +nodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = nodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) +nodes(ibg::IBG, (ℓx, ℓy, ℓz); kwargs...) = nodes(ibg, ℓx, ℓy, ℓz; kwargs...) + +xnodes(ibg::IBG, loc; kwargs...) = xnodes(ibg.underlying_grid, loc; kwargs...) +ynodes(ibg::IBG, loc; kwargs...) = ynodes(ibg.underlying_grid, loc; kwargs...) +znodes(ibg::IBG, loc; kwargs...) = znodes(ibg.underlying_grid, loc; kwargs...) + +λnodes(ibg::IBG, loc; kwargs...) = λnodes(ibg.underlying_grid, loc; kwargs...) +φnodes(ibg::IBG, loc; kwargs...) = φnodes(ibg.underlying_grid, loc; kwargs...) + +ξnodes(ibg::IBG, loc; kwargs...) = ξnodes(ibg.underlying_grid, loc; kwargs...) +ηnodes(ibg::IBG, loc; kwargs...) = ηnodes(ibg.underlying_grid, loc; kwargs...) +rnodes(ibg::IBG, loc; kwargs...) = rnodes(ibg.underlying_grid, loc; kwargs...) + +xnodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = xnodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) +ynodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = ynodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) +znodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = znodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) + +λnodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = λnodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) +φnodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = φnodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) + +ξnodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = ξnodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) +ηnodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = ηnodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) +rnodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = rnodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) + +@inline cpu_face_constructor_x(ibg::IBG) = cpu_face_constructor_x(ibg.underlying_grid) +@inline cpu_face_constructor_y(ibg::IBG) = cpu_face_constructor_y(ibg.underlying_grid) +@inline cpu_face_constructor_z(ibg::IBG) = cpu_face_constructor_z(ibg.underlying_grid) + +node_names(ibg::IBG, ℓx, ℓy, ℓz) = node_names(ibg.underlying_grid, ℓx, ℓy, ℓz) +ξname(ibg::IBG) = ξname(ibg.underlying_grid) +ηname(ibg::IBG) = ηname(ibg.underlying_grid) +rname(ibg::IBG) = rname(ibg.underlying_grid) + +function on_architecture(arch, ibg::IBG) + underlying_grid = on_architecture(arch, ibg.underlying_grid) + immersed_boundary = on_architecture(arch, ibg.immersed_boundary) + return ImmersedBoundaryGrid(underlying_grid, immersed_boundary) +end + +isrectilinear(ibg::IBG) = isrectilinear(ibg.underlying_grid) + +@inline fractional_x_index(x, locs, grid::ImmersedBoundaryGrid) = fractional_x_index(x, locs, grid.underlying_grid) +@inline fractional_y_index(x, locs, grid::ImmersedBoundaryGrid) = fractional_y_index(x, locs, grid.underlying_grid) +@inline fractional_z_index(x, locs, grid::ImmersedBoundaryGrid) = fractional_z_index(x, locs, grid.underlying_grid) + include("active_cells_map.jl") include("immersed_grid_metrics.jl") include("abstract_grid_fitted_boundary.jl") diff --git a/src/ImmersedBoundaries/abstract_grid_fitted_boundary.jl b/src/ImmersedBoundaries/abstract_grid_fitted_boundary.jl index 2516b89c3c..64dad56766 100644 --- a/src/ImmersedBoundaries/abstract_grid_fitted_boundary.jl +++ b/src/ImmersedBoundaries/abstract_grid_fitted_boundary.jl @@ -18,3 +18,21 @@ const AGFB = AbstractGridFittedBoundary @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, <:Any, Flat, Flat}, ib::AGFB) = _immersed_cell(i, 1, 1, grid, ib) @inline immersed_cell(i, j, k, grid::AbstractGrid{<:Any, Flat, Flat, Flat}, ib::AGFB) = _immersed_cell(1, 1, 1, grid, ib) end + +function clamp_bottom_height!(bottom_field, grid) + launch!(architecture(grid), grid, :xy, _clamp_bottom_height!, bottom_field, grid) + return nothing +end + +const c = Center() +const f = Face() + +@kernel function _clamp_bottom_height!(z, grid) + i, j = @index(Global, NTuple) + Nz = size(grid, 3) + zmin = znode(i, j, 1, grid, c, c, f) + zmax = znode(i, j, Nz+1, grid, c, c, f) + @inbounds z[i, j, 1] = clamp(z[i, j, 1], zmin, zmax) +end + + diff --git a/src/ImmersedBoundaries/conditional_differences.jl b/src/ImmersedBoundaries/conditional_differences.jl index 43706507d3..1e9ae156bf 100644 --- a/src/ImmersedBoundaries/conditional_differences.jl +++ b/src/ImmersedBoundaries/conditional_differences.jl @@ -1,4 +1,4 @@ -import Oceananigans.Operators: +import Oceananigans.Operators: δxᶠᶜᶜ, δxᶠᶜᶠ, δxᶠᶠᶜ, δxᶠᶠᶠ, δxᶜᶜᶜ, δxᶜᶜᶠ, δxᶜᶠᶜ, δxᶜᶠᶠ, δyᶜᶠᶜ, δyᶜᶠᶠ, δyᶠᶠᶜ, δyᶠᶠᶠ, @@ -6,10 +6,6 @@ import Oceananigans.Operators: δzᶜᶜᶠ, δzᶜᶠᶠ, δzᶠᶜᶠ, δzᶠᶠᶠ, δzᶜᶜᶜ, δzᶜᶠᶜ, δzᶠᶜᶜ, δzᶠᶠᶜ -import Oceananigans.Operators: - δxTᶜᵃᵃ, δyTᵃᶜᵃ, - ∂xTᶠᶜᶠ, ∂yTᶜᶠᶠ - # Conditional differences that are "immersed boundary aware". # Here we return `zero(ibg)` rather than `δx` (for example) when _one_ of the # nodes involved in the difference is `immersed_inactive_node`. @@ -63,30 +59,15 @@ for (d, ξ) in enumerate((:x, :y, :z)) # `other_locs` contains locations in the two "other" directions not being differenced other_locs = [] - for l in 1:3 + for l in 1:3 if l != d push!(other_locs, loc[l]) end end - + @eval begin @inline $δξ(i, j, k, ibg::IBG, args...) = $conditional_δξ($(other_locs[1]), $(other_locs[2]), i, j, k, ibg, $δξ, args...) @inline $δξ(i, j, k, ibg::IBG, f::Function, args...) = $conditional_δξ($(other_locs[1]), $(other_locs[2]), i, j, k, ibg, $δξ, f::Function, args...) end end end - -# Topology-aware immersed boundary operators -# * Velocities are `0` on `peripheral_node`s and -# * Tracers should ensure no-flux on `inactive_node`s - -@inline conditional_U_fcc(i, j, k, grid, ibg::IBG, f::Function, args...) = ifelse(peripheral_node(i, j, k, ibg, f, c, c), zero(ibg), f(i, j, k, grid, args...)) -@inline conditional_V_cfc(i, j, k, grid, ibg::IBG, f::Function, args...) = ifelse(peripheral_node(i, j, k, ibg, c, f, c), zero(ibg), f(i, j, k, grid, args...)) - -@inline conditional_∂xTᶠᶜᶠ(i, j, k, ibg::IBG, args...) = ifelse(inactive_node(i, j, k, ibg, c, c, f) | inactive_node(i-1, j, k, ibg, c, c, f), zero(ibg), ∂xTᶠᶜᶠ(i, j, k, ibg.underlying_grid, args...)) -@inline conditional_∂yTᶜᶠᶠ(i, j, k, ibg::IBG, args...) = ifelse(inactive_node(i, j, k, ibg, c, c, f) | inactive_node(i, j-1, k, ibg, c, c, f), zero(ibg), ∂yTᶜᶠᶠ(i, j, k, ibg.underlying_grid, args...)) - -@inline δxTᶜᵃᵃ(i, j, k, ibg::IBG, f::Function, args...) = δxTᶜᵃᵃ(i, j, k, ibg.underlying_grid, conditional_U_fcc, ibg, f, args...) -@inline δyTᵃᶜᵃ(i, j, k, ibg::IBG, f::Function, args...) = δyTᵃᶜᵃ(i, j, k, ibg.underlying_grid, conditional_V_cfc, ibg, f, args...) -@inline ∂xTᶠᶜᶠ(i, j, k, ibg::IBG, f::Function, args...) = conditional_∂xTᶠᶜᶠ(i, j, k, ibg, f, args...) -@inline ∂yTᶜᶠᶠ(i, j, k, ibg::IBG, f::Function, args...) = conditional_∂yTᶜᶠᶠ(i, j, k, ibg, f, args...) diff --git a/src/ImmersedBoundaries/grid_fitted_bottom.jl b/src/ImmersedBoundaries/grid_fitted_bottom.jl index 9364bba3e3..53e2aa2805 100644 --- a/src/ImmersedBoundaries/grid_fitted_bottom.jl +++ b/src/ImmersedBoundaries/grid_fitted_bottom.jl @@ -18,14 +18,14 @@ abstract type AbstractGridFittedBottom{H} <: AbstractGridFittedBoundary end struct CenterImmersedCondition end struct InterfaceImmersedCondition end +Base.summary(::CenterImmersedCondition) = "CenterImmersedCondition" +Base.summary(::InterfaceImmersedCondition) = "InterfaceImmersedCondition" + struct GridFittedBottom{H, I} <: AbstractGridFittedBottom{H} bottom_height :: H immersed_condition :: I end -Base.summary(::CenterImmersedCondition) = "CenterImmersedCondition" -Base.summary(::InterfaceImmersedCondition) = "InterfaceImmersedCondition" - const GFBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:GridFittedBottom} """ @@ -36,7 +36,6 @@ Return a bottom immersed boundary. Keyword Arguments ================= - * `bottom_height`: an array or function that gives the height of the bottom in absolute ``z`` coordinates. @@ -73,77 +72,50 @@ Base.summary(ib::GridFittedBottom{<:Function}) = @sprintf("GridFittedBottom(%s)" function Base.show(io::IO, ib::GridFittedBottom) print(io, summary(ib), '\n') print(io, "├── bottom_height: ", prettysummary(ib.bottom_height), '\n') + print(io, "└── immersed_condition: ", summary(ib.immersed_condition)) end -on_architecture(arch, ib::GridFittedBottom) = GridFittedBottom(on_architecture(arch, ib.bottom_height), ib.immersed_condition) - -function on_architecture(arch, ib::GridFittedBottom{<:Field}) - architecture(ib.bottom_height) == arch && return ib - arch_grid = on_architecture(arch, ib.bottom_height.grid) - new_bottom_height = Field{Center, Center, Nothing}(arch_grid) - set!(new_bottom_height, ib.bottom_height) - fill_halo_regions!(new_bottom_height) - return GridFittedBottom(new_bottom_height, ib.immersed_condition) -end - -Adapt.adapt_structure(to, ib::GridFittedBottom) = GridFittedBottom(adapt(to, ib.bottom_height), adapt(to, ib.immersed_condition)) +@inline z_bottom(i, j, ibg::GFBIBG) = @inbounds ibg.immersed_boundary.bottom_height[i, j, 1] """ ImmersedBoundaryGrid(grid, ib::GridFittedBottom) Return a grid with `GridFittedBottom` immersed boundary (`ib`). -Computes `ib.bottom_height` and wraps it in a Field. `ib.bottom_height` is the z-coordinate of top-most interface -of the last ``immersed`` cell in the column. +Computes `ib.bottom_height` and wraps it in a Field. """ function ImmersedBoundaryGrid(grid, ib::GridFittedBottom) bottom_field = Field{Center, Center, Nothing}(grid) set!(bottom_field, ib.bottom_height) - @apply_regionally compute_numerical_bottom_height!(bottom_field, grid, ib) + @apply_regionally clamp_bottom_height!(bottom_field, grid) fill_halo_regions!(bottom_field) - new_ib = GridFittedBottom(bottom_field) + new_ib = GridFittedBottom(bottom_field, ib.immersed_condition) TX, TY, TZ = topology(grid) return ImmersedBoundaryGrid{TX, TY, TZ}(grid, new_ib) end -compute_numerical_bottom_height!(bottom_field, grid, ib) = - launch!(architecture(grid), grid, :xy, _compute_numerical_bottom_height!, bottom_field, grid, ib) - -@kernel function _compute_numerical_bottom_height!(bottom_field, grid, ib::GridFittedBottom) - i, j = @index(Global, NTuple) - zb = @inbounds bottom_field[i, j, 1] - @inbounds bottom_field[i, j, 1] = znode(i, j, 1, grid, c, c, f) - condition = ib.immersed_condition - for k in 1:grid.Nz - z⁺ = znode(i, j, k+1, grid, c, c, f) - z = znode(i, j, k, grid, c, c, c) - bottom_cell = ifelse(condition isa CenterImmersedCondition, z ≤ zb, z⁺ ≤ zb) - @inbounds bottom_field[i, j, 1] = ifelse(bottom_cell, z⁺, bottom_field[i, j, 1]) - end +@inline function _immersed_cell(i, j, k, underlying_grid, ib::GridFittedBottom{<:Any, <:InterfaceImmersedCondition}) + z = znode(i, j, k+1, underlying_grid, c, c, f) + h = @inbounds ib.bottom_height[i, j, 1] + return z ≤ h end -@inline function _immersed_cell(i, j, k, underlying_grid, ib::GridFittedBottom) - z = znode(i, j, k, underlying_grid, c, c, c) - zb = @inbounds ib.bottom_height[i, j, 1] - return z ≤ zb +@inline function _immersed_cell(i, j, k, underlying_grid, ib::GridFittedBottom{<:Any, <:CenterImmersedCondition}) + z = znode(i, j, k, underlying_grid, c, c, c) + h = @inbounds ib.bottom_height[i, j, 1] + return z ≤ h end -##### -##### Static column depth -##### +on_architecture(arch, ib::GridFittedBottom) = GridFittedBottom(ib.bottom_height, ib.immersed_condition) -const AGFBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractGridFittedBottom} - -@inline static_column_depthᶜᶜᵃ(i, j, ibg::AGFBIBG) = @inbounds znode(i, j, ibg.Nz+1, ibg, c, c, f) - ibg.immersed_boundary.bottom_height[i, j, 1] -@inline static_column_depthᶜᶠᵃ(i, j, ibg::AGFBIBG) = min(static_column_depthᶜᶜᵃ(i, j-1, ibg), static_column_depthᶜᶜᵃ(i, j, ibg)) -@inline static_column_depthᶠᶜᵃ(i, j, ibg::AGFBIBG) = min(static_column_depthᶜᶜᵃ(i-1, j, ibg), static_column_depthᶜᶜᵃ(i, j, ibg)) -@inline static_column_depthᶠᶠᵃ(i, j, ibg::AGFBIBG) = min(static_column_depthᶠᶜᵃ(i, j-1, ibg), static_column_depthᶠᶜᵃ(i, j, ibg)) - -# Make sure column_height works for horizontally-Flat topologies. -XFlatAGFIBG = ImmersedBoundaryGrid{<:Any, <:Flat, <:Any, <:Any, <:Any, <:AbstractGridFittedBottom} -YFlatAGFIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Flat, <:Any, <:Any, <:AbstractGridFittedBottom} +function on_architecture(arch, ib::GridFittedBottom{<:Field}) + architecture(ib.bottom_height) == arch && return ib + arch_grid = on_architecture(arch, ib.bottom_height.grid) + new_bottom_height = Field{Center, Center, Nothing}(arch_grid) + set!(new_bottom_height, ib.bottom_height) + fill_halo_regions!(new_bottom_height) + return GridFittedBottom(new_bottom_height, ib.immersed_condition) +end -@inline static_column_depthᶠᶜᵃ(i, j, ibg::XFlatAGFIBG) = static_column_depthᶜᶜᵃ(i, j, ibg) -@inline static_column_depthᶜᶠᵃ(i, j, ibg::YFlatAGFIBG) = static_column_depthᶜᶜᵃ(i, j, ibg) -@inline static_column_depthᶠᶠᵃ(i, j, ibg::XFlatAGFIBG) = static_column_depthᶜᶠᵃ(i, j, ibg) -@inline static_column_depthᶠᶠᵃ(i, j, ibg::YFlatAGFIBG) = static_column_depthᶠᶜᵃ(i, j, ibg) +Adapt.adapt_structure(to, ib::GridFittedBottom) = GridFittedBottom(adapt(to, ib.bottom_height), + ib.immersed_condition) diff --git a/src/ImmersedBoundaries/immersed_boundary_grid.jl b/src/ImmersedBoundaries/immersed_boundary_grid.jl deleted file mode 100644 index 664dfb8ecf..0000000000 --- a/src/ImmersedBoundaries/immersed_boundary_grid.jl +++ /dev/null @@ -1,90 +0,0 @@ -""" - abstract type AbstractImmersedBoundary - -Abstract supertype for immersed boundary grids. -""" -abstract type AbstractImmersedBoundary end - -struct ImmersedBoundaryGrid{FT, TX, TY, TZ, G, I, M, S, Arch} <: AbstractGrid{FT, TX, TY, TZ, Arch} - architecture :: Arch - underlying_grid :: G - immersed_boundary :: I - interior_active_cells :: M - active_z_columns :: S - - # Internal interface - function ImmersedBoundaryGrid{TX, TY, TZ}(grid::G, ib::I, mi::M, ms::S) where {TX, TY, TZ, G <: AbstractUnderlyingGrid, I, M, S} - FT = eltype(grid) - arch = architecture(grid) - Arch = typeof(arch) - return new{FT, TX, TY, TZ, G, I, M, S, Arch}(arch, grid, ib, mi, ms) - end - - # Constructor with no active map - function ImmersedBoundaryGrid{TX, TY, TZ}(grid::G, ib::I) where {TX, TY, TZ, G <: AbstractUnderlyingGrid, I} - FT = eltype(grid) - arch = architecture(grid) - Arch = typeof(arch) - return new{FT, TX, TY, TZ, G, I, Nothing, Nothing, Arch}(arch, grid, ib, nothing, nothing) - end -end - -const IBG = ImmersedBoundaryGrid - -@inline Base.getproperty(ibg::IBG, property::Symbol) = get_ibg_property(ibg, Val(property)) -@inline get_ibg_property(ibg::IBG, ::Val{property}) where property = getfield(getfield(ibg, :underlying_grid), property) -@inline get_ibg_property(ibg::IBG, ::Val{:immersed_boundary}) = getfield(ibg, :immersed_boundary) -@inline get_ibg_property(ibg::IBG, ::Val{:underlying_grid}) = getfield(ibg, :underlying_grid) -@inline get_ibg_property(ibg::IBG, ::Val{:interior_active_cells}) = getfield(ibg, :interior_active_cells) -@inline get_ibg_property(ibg::IBG, ::Val{:active_z_columns}) = getfield(ibg, :active_z_columns) - -@inline architecture(ibg::IBG) = architecture(ibg.underlying_grid) - -@inline x_domain(ibg::IBG) = x_domain(ibg.underlying_grid) -@inline y_domain(ibg::IBG) = y_domain(ibg.underlying_grid) -@inline z_domain(ibg::IBG) = z_domain(ibg.underlying_grid) - -Adapt.adapt_structure(to, ibg::IBG{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = - ImmersedBoundaryGrid{TX, TY, TZ}(adapt(to, ibg.underlying_grid), - adapt(to, ibg.immersed_boundary), - nothing, - nothing) - -with_halo(halo, ibg::ImmersedBoundaryGrid) = - ImmersedBoundaryGrid(with_halo(halo, ibg.underlying_grid), ibg.immersed_boundary) - -# ImmersedBoundaryGrids require an extra halo point to check the "inactivity" of a `Face` node at N + H -# (which requires checking `Center` nodes at N + H and N + H + 1) -inflate_halo_size_one_dimension(req_H, old_H, _, ::IBG) = max(req_H + 1, old_H) -inflate_halo_size_one_dimension(req_H, old_H, ::Type{Flat}, ::IBG) = 0 - -# Defining the bottom -@inline z_bottom(i, j, grid) = znode(i, j, 1, grid, c, c, f) -@inline z_bottom(i, j, ibg::IBG) = error("The function `bottom` has not been defined for $(summary(ibg))!") - -function Base.summary(grid::ImmersedBoundaryGrid) - FT = eltype(grid) - TX, TY, TZ = topology(grid) - - return string(size_summary(size(grid)), - " ImmersedBoundaryGrid{$FT, $TX, $TY, $TZ} on ", summary(architecture(grid)), - " with ", size_summary(halo_size(grid)), " halo") -end - -function show(io::IO, ibg::ImmersedBoundaryGrid) - print(io, summary(ibg), ":", "\n", - "├── immersed_boundary: ", summary(ibg.immersed_boundary), "\n", - "├── underlying_grid: ", summary(ibg.underlying_grid), "\n") - - return show(io, ibg.underlying_grid, false) -end - -@inline Base.zero(ibg::IBG) = zero(ibg.underlying_grid) - -function on_architecture(arch, ibg::IBG) - underlying_grid = on_architecture(arch, ibg.underlying_grid) - immersed_boundary = on_architecture(arch, ibg.immersed_boundary) - return ImmersedBoundaryGrid(underlying_grid, immersed_boundary) -end - -isrectilinear(ibg::IBG) = isrectilinear(ibg.underlying_grid) diff --git a/src/ImmersedBoundaries/immersed_boundary_interface.jl b/src/ImmersedBoundaries/immersed_boundary_interface.jl deleted file mode 100644 index 8f6e62153b..0000000000 --- a/src/ImmersedBoundaries/immersed_boundary_interface.jl +++ /dev/null @@ -1,56 +0,0 @@ -""" - immersed_cell(i, j, k, grid) - -Return true if a `cell` is "completely" immersed, and thus -is not part of the prognostic state. -""" -@inline immersed_cell(i, j, k, grid) = false - -# Unpack to make defining new immersed boundaries more convenient -@inline immersed_cell(i, j, k, grid::ImmersedBoundaryGrid) = - immersed_cell(i, j, k, grid.underlying_grid, grid.immersed_boundary) - -""" - inactive_cell(i, j, k, grid::ImmersedBoundaryGrid) - -Return `true` if the tracer cell at `i, j, k` either (i) lies outside the `Bounded` domain -or (ii) lies within the immersed region of `ImmersedBoundaryGrid`. - -Example -======= - -Consider the configuration - -``` - Immersed Fluid - =========== ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ - - c c - i-1 i - - | ========= | | - × === ∘ === × ∘ × - | ========= | | - -i-1 i - f f f -``` - -We then have - -* `inactive_node(i, 1, 1, grid, f, c, c) = false` - -As well as - -* `inactive_node(i, 1, 1, grid, c, c, c) = false` -* `inactive_node(i-1, 1, 1, grid, c, c, c) = true` -* `inactive_node(i-1, 1, 1, grid, f, c, c) = true` -""" -@inline inactive_cell(i, j, k, ibg::IBG) = immersed_cell(i, j, k, ibg) | inactive_cell(i, j, k, ibg.underlying_grid) - -# Isolate periphery of the immersed boundary -@inline immersed_peripheral_node(i, j, k, ibg::IBG, LX, LY, LZ) = peripheral_node(i, j, k, ibg, LX, LY, LZ) & - !peripheral_node(i, j, k, ibg.underlying_grid, LX, LY, LZ) - -@inline immersed_inactive_node(i, j, k, ibg::IBG, LX, LY, LZ) = inactive_node(i, j, k, ibg, LX, LY, LZ) & - !inactive_node(i, j, k, ibg.underlying_grid, LX, LY, LZ) diff --git a/src/ImmersedBoundaries/immersed_boundary_nodes.jl b/src/ImmersedBoundaries/immersed_boundary_nodes.jl deleted file mode 100644 index 62f39894f8..0000000000 --- a/src/ImmersedBoundaries/immersed_boundary_nodes.jl +++ /dev/null @@ -1,69 +0,0 @@ -import Oceananigans.Grids: xspacings, yspacings, zspacings - -const c = Center() -const f = Face() - -@inline xnode(i, ibg::IBG, ℓx) = xnode(i, ibg.underlying_grid, ℓx) -@inline ynode(j, ibg::IBG, ℓy) = ynode(j, ibg.underlying_grid, ℓy) -@inline znode(k, ibg::IBG, ℓz) = znode(k, ibg.underlying_grid, ℓz) - -@inline λnode(i, ibg::IBG, ℓx) = λnode(i, ibg.underlying_grid, ℓx) -@inline φnode(j, ibg::IBG, ℓy) = φnode(i, ibg.underlying_grid, ℓy) - -@inline xnode(i, j, ibg::IBG, ℓx, ℓy) = xnode(i, j, ibg.underlying_grid, ℓx, ℓy) - -@inline xnode(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = xnode(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) -@inline ynode(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = ynode(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) -@inline znode(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = znode(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) - -@inline λnode(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = λnode(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) -@inline φnode(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = φnode(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) - -@inline ξnode(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = ξnode(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) -@inline ηnode(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = ηnode(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) -@inline rnode(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = rnode(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) - -@inline node(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = node(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) - -nodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = nodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) -nodes(ibg::IBG, (ℓx, ℓy, ℓz); kwargs...) = nodes(ibg, ℓx, ℓy, ℓz; kwargs...) - -xnodes(ibg::IBG, loc; kwargs...) = xnodes(ibg.underlying_grid, loc; kwargs...) -ynodes(ibg::IBG, loc; kwargs...) = ynodes(ibg.underlying_grid, loc; kwargs...) -znodes(ibg::IBG, loc; kwargs...) = znodes(ibg.underlying_grid, loc; kwargs...) - -λnodes(ibg::IBG, loc; kwargs...) = λnodes(ibg.underlying_grid, loc; kwargs...) -φnodes(ibg::IBG, loc; kwargs...) = φnodes(ibg.underlying_grid, loc; kwargs...) - -ξnodes(ibg::IBG, loc; kwargs...) = ξnodes(ibg.underlying_grid, loc; kwargs...) -ηnodes(ibg::IBG, loc; kwargs...) = ηnodes(ibg.underlying_grid, loc; kwargs...) -rnodes(ibg::IBG, loc; kwargs...) = rnodes(ibg.underlying_grid, loc; kwargs...) - -xnodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = xnodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) -ynodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = ynodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) -znodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = znodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) - -λnodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = λnodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) -φnodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = φnodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) - -ξnodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = ξnodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) -ηnodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = ηnodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) -rnodes(ibg::IBG, ℓx, ℓy, ℓz; kwargs...) = rnodes(ibg.underlying_grid, ℓx, ℓy, ℓz; kwargs...) - -@inline cpu_face_constructor_x(ibg::IBG) = cpu_face_constructor_x(ibg.underlying_grid) -@inline cpu_face_constructor_y(ibg::IBG) = cpu_face_constructor_y(ibg.underlying_grid) -@inline cpu_face_constructor_z(ibg::IBG) = cpu_face_constructor_z(ibg.underlying_grid) - -node_names(ibg::IBG, ℓx, ℓy, ℓz) = node_names(ibg.underlying_grid, ℓx, ℓy, ℓz) - -ξname(ibg::IBG) = ξname(ibg.underlying_grid) -ηname(ibg::IBG) = ηname(ibg.underlying_grid) -rname(ibg::IBG) = rname(ibg.underlying_grid) - -@inline fractional_x_index(x, locs, grid::ImmersedBoundaryGrid) = fractional_x_index(x, locs, grid.underlying_grid) -@inline fractional_y_index(x, locs, grid::ImmersedBoundaryGrid) = fractional_y_index(x, locs, grid.underlying_grid) -@inline fractional_z_index(x, locs, grid::ImmersedBoundaryGrid) = fractional_z_index(x, locs, grid.underlying_grid) - -xspacings(ibg::ImmersedBoundaryGrid, ℓ...) = xspacings(ibg.underlying_grid, ℓ...) -yspacings(ibg::ImmersedBoundaryGrid, ℓ...) = yspacings(ibg.underlying_grid, ℓ...) -zspacings(ibg::ImmersedBoundaryGrid, ℓ...) = zspacings(ibg.underlying_grid, ℓ...) diff --git a/src/ImmersedBoundaries/immersed_grid_metrics.jl b/src/ImmersedBoundaries/immersed_grid_metrics.jl index 48cb239768..a7dfbcbe10 100644 --- a/src/ImmersedBoundaries/immersed_grid_metrics.jl +++ b/src/ImmersedBoundaries/immersed_grid_metrics.jl @@ -1,7 +1,10 @@ using Oceananigans.AbstractOperations: GridMetricOperation import Oceananigans.Grids: coordinates -import Oceananigans.Operators: Δzᵃᵃᶠ, Δzᵃᵃᶜ, intrinsic_vector, extrinsic_vector + +const c = Center() +const f = Face() +const IBG = ImmersedBoundaryGrid # Grid metrics for ImmersedBoundaryGrid # @@ -9,10 +12,10 @@ import Oceananigans.Operators: Δzᵃᵃᶠ, Δzᵃᵃᶜ, intrinsic_vector, ext # # For non "full-cell" immersed boundaries, grid metric functions # must be extended for the specific immersed boundary grid in question. - +# for LX in (:ᶜ, :ᶠ), LY in (:ᶜ, :ᶠ), LZ in (:ᶜ, :ᶠ) for dir in (:x, :y, :z), operator in (:Δ, :A) - + metric = Symbol(operator, dir, LX, LY, LZ) @eval begin import Oceananigans.Operators: $metric @@ -27,14 +30,10 @@ for LX in (:ᶜ, :ᶠ), LY in (:ᶜ, :ᶠ), LZ in (:ᶜ, :ᶠ) end end -coordinates(grid::IBG) = coordinates(grid.underlying_grid) - -@inline Δzᵃᵃᶠ(i, j, k, ibg::IBG) = Δzᵃᵃᶠ(i, j, k, ibg.underlying_grid) @inline Δzᵃᵃᶜ(i, j, k, ibg::IBG) = Δzᵃᵃᶜ(i, j, k, ibg.underlying_grid) +@inline Δzᵃᵃᶠ(i, j, k, ibg::IBG) = Δzᵃᵃᶠ(i, j, k, ibg.underlying_grid) -# Extend both 2D and 3D methods -@inline intrinsic_vector(i, j, k, ibg::IBG, u, v) = intrinsic_vector(i, j, k, ibg.underlying_grid, u, v) -@inline extrinsic_vector(i, j, k, ibg::IBG, u, v) = extrinsic_vector(i, j, k, ibg.underlying_grid, u, v) - -@inline intrinsic_vector(i, j, k, ibg::IBG, u, v, w) = intrinsic_vector(i, j, k, ibg.underlying_grid, u, v, w) -@inline extrinsic_vector(i, j, k, ibg::IBG, u, v, w) = extrinsic_vector(i, j, k, ibg.underlying_grid, u, v, w) \ No newline at end of file +coordinates(grid::IBG) = coordinates(grid.underlying_grid) +xspacings(X, grid::IBG) = xspacings(X, grid.underlying_grid) +yspacings(Y, grid::IBG) = yspacings(Y, grid.underlying_grid) +zspacings(Z, grid::IBG) = zspacings(Z, grid.underlying_grid) diff --git a/src/ImmersedBoundaries/mask_immersed_field.jl b/src/ImmersedBoundaries/mask_immersed_field.jl index bcfd01645c..9ae23d88c3 100644 --- a/src/ImmersedBoundaries/mask_immersed_field.jl +++ b/src/ImmersedBoundaries/mask_immersed_field.jl @@ -1,7 +1,7 @@ using KernelAbstractions: @kernel, @index using Statistics using Oceananigans.AbstractOperations: BinaryOperation -using Oceananigans.Fields: location, XReducedField, YReducedField, ZReducedField, Field, ReducedField +using Oceananigans.Fields: location, ZReducedField, Field instantiate(T::Type) = T() instantiate(t) = t @@ -89,55 +89,6 @@ end @inbounds field[i, j, k] = scalar_mask(i, j, k, grid, grid.immersed_boundary, loc..., value, field, mask) end -##### -##### Masking a `ReducedField` -##### - -# We mask a `ReducedField` if the entire reduced direction is immersed. -# This requires a sweep over the reduced direction - -function mask_immersed_field!(field::ReducedField, grid::ImmersedBoundaryGrid, loc, value) - loc = instantiate.(loc) - dims = reduced_dimensions(field) - launch!(architecture(field), grid, size(field), _mask_immersed_reduced_field!, field, dims, loc, grid, value) - return nothing -end - -@kernel function _mask_immersed_reduced_field!(field, dims, loc, grid, value) - i, j, k = @index(Global, NTuple) - mask = inactive_dimensions(i, j, k, grid, dims, loc) - @inbounds field[i, j, k] = ifelse(mask, value, field[i, j, k]) -end - -@inline inactive_search_range(i, grid, dim, dims) = ifelse(dim ∈ dims, 1:size(grid, dim), i:i) - -@inline function inactive_dimensions(i₀, j₀, k₀, grid, dims, loc) - mask = true - irange = inactive_search_range(i₀, grid, 1, dims) - jrange = inactive_search_range(j₀, grid, 2, dims) - krange = inactive_search_range(k₀, grid, 3, dims) - - # The loop activates over the whole direction only if reduced directions - for i in irange, j in jrange, k in krange - mask = mask & peripheral_node(i, j, k, grid, loc...) - end - - return mask -end - -### -### Efficient masking for `OnlyZReducedField` and an `AbstractGridFittedBoundary` -### - -const AGFBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractGridFittedBottom} - -const CenterOrFace = Union{Center, Face} -const OnlyZReducedField = Field{<:CenterOrFace, <:CenterOrFace, Nothing} - -# Does not require a sweep -mask_immersed_field!(field::OnlyZReducedField, grid::AGFBIBG, loc, value) = - mask_immersed_field_xy!(field, grid, loc, value; k=size(grid, 3), mask=peripheral_node) - ##### ##### Masking for GridFittedBoundary ##### diff --git a/src/ImmersedBoundaries/partial_cell_bottom.jl b/src/ImmersedBoundaries/partial_cell_bottom.jl index 36cc32487b..d08f6db815 100644 --- a/src/ImmersedBoundaries/partial_cell_bottom.jl +++ b/src/ImmersedBoundaries/partial_cell_bottom.jl @@ -63,42 +63,13 @@ end function ImmersedBoundaryGrid(grid, ib::PartialCellBottom) bottom_field = Field{Center, Center, Nothing}(grid) set!(bottom_field, ib.bottom_height) - @apply_regionally compute_numerical_bottom_height!(bottom_field, grid, ib) + @apply_regionally clamp_bottom_height!(bottom_field, grid) fill_halo_regions!(bottom_field) new_ib = PartialCellBottom(bottom_field, ib.minimum_fractional_cell_height) TX, TY, TZ = topology(grid) return ImmersedBoundaryGrid{TX, TY, TZ}(grid, new_ib) end -@kernel function _compute_numerical_bottom_height!(bottom_field, grid, ib::PartialCellBottom) - i, j = @index(Global, NTuple) - - # Save analytical bottom height - zb = @inbounds bottom_field[i, j, 1] - - # Cap bottom height at Lz and at rnode(i, j, grid.Nz+1, grid, c, c, f) - domain_bottom = znode(i, j, 1, grid, c, c, f) - domain_top = znode(i, j, grid.Nz+1, grid, c, c, f) - @inbounds bottom_field[i, j, 1] = clamp(zb, domain_bottom, domain_top) - adjusted_zb = bottom_field[i, j, 1] - - ϵ = ib.minimum_fractional_cell_height - - for k in 1:grid.Nz - z⁻ = znode(i, j, k, grid, c, c, f) - z⁺ = znode(i, j, k+1, grid, c, c, f) - Δz = Δzᶜᶜᶜ(i, j, k, grid) - bottom_cell = (z⁻ ≤ zb) & (z⁺ ≥ zb) - capped_zb = min(z⁺ - ϵ * Δz, zb) - - # If the size of the bottom cell is less than ϵ Δz, - # we enforce a minimum size of ϵ Δz. - adjusted_zb = ifelse(bottom_cell, capped_zb, zb) - end - - @inbounds bottom_field[i, j, 1] = adjusted_zb -end - function on_architecture(arch, ib::PartialCellBottom{<:Field}) architecture(ib.bottom_height) == arch && return ib arch_grid = on_architecture(arch, ib.bottom_height.grid) @@ -108,20 +79,20 @@ function on_architecture(arch, ib::PartialCellBottom{<:Field}) end Adapt.adapt_structure(to, ib::PartialCellBottom) = PartialCellBottom(adapt(to, ib.bottom_height), - ib.minimum_fractional_cell_height) + ib.minimum_fractional_cell_height) on_architecture(to, ib::PartialCellBottom) = PartialCellBottom(on_architecture(to, ib.bottom_height), - on_architecture(to, ib.minimum_fractional_cell_height)) + on_architecture(to, ib.minimum_fractional_cell_height)) """ immersed underlying --x-- --x-- - - + + ∘ ↑ ∘ k+1 | - | + | k+1 --x-- | k+1 --x-- ↑ <- node z ∘ ↓ | zb ⋅⋅x⋅⋅ | @@ -130,17 +101,18 @@ on_architecture(to, ib::PartialCellBottom) = PartialCellBottom(on_architecture(t | | k --x-- ↓ - + Criterion is zb ≥ z - ϵ Δz """ @inline function _immersed_cell(i, j, k, underlying_grid, ib::PartialCellBottom) - z⁺ = znode(i, j, k+1, underlying_grid, c, c, f) + # Face node below current cell + z = znode(i, j, k, underlying_grid, c, c, f) + zb = @inbounds ib.bottom_height[i, j, 1] ϵ = ib.minimum_fractional_cell_height + # z + Δz is equal to the face above the current cell Δz = Δzᶜᶜᶜ(i, j, k, underlying_grid) - z★ = z⁺ - Δz * ϵ - zb = @inbounds ib.bottom_height[i, j, 1] - return z★ < zb + return (z + Δz * (1 - ϵ)) ≤ zb end @inline function bottom_cell(i, j, k, ibg::PCBIBG) @@ -157,14 +129,15 @@ end # Get node at face above and defining nodes on c,c,f z = znode(i, j, k+1, underlying_grid, c, c, f) - # Get bottom z-coordinate and fractional Δz parameter - zb = @inbounds ib.bottom_height[i, j, 1] + # Get bottom height and fractional Δz parameter + h = @inbounds ib.bottom_height[i, j, 1] + ϵ = ibg.immersed_boundary.minimum_fractional_cell_height # Are we in a bottom cell? at_the_bottom = bottom_cell(i, j, k, ibg) - full_Δz = Δzᶜᶜᶜ(i, j, k, ibg.underlying_grid) - partial_Δz = z - zb + full_Δz = Δzᶜᶜᶜ(i, j, k, ibg.underlying_grid) + partial_Δz = max(ϵ * full_Δz, z - h) return ifelse(at_the_bottom, partial_Δz, full_Δz) end @@ -185,9 +158,9 @@ end @inline Δzᶠᶜᶜ(i, j, k, ibg::PCBIBG) = min(Δzᶜᶜᶜ(i-1, j, k, ibg), Δzᶜᶜᶜ(i, j, k, ibg)) @inline Δzᶜᶠᶜ(i, j, k, ibg::PCBIBG) = min(Δzᶜᶜᶜ(i, j-1, k, ibg), Δzᶜᶜᶜ(i, j, k, ibg)) @inline Δzᶠᶠᶜ(i, j, k, ibg::PCBIBG) = min(Δzᶠᶜᶜ(i, j-1, k, ibg), Δzᶠᶜᶜ(i, j, k, ibg)) - + @inline Δzᶠᶜᶠ(i, j, k, ibg::PCBIBG) = min(Δzᶜᶜᶠ(i-1, j, k, ibg), Δzᶜᶜᶠ(i, j, k, ibg)) -@inline Δzᶜᶠᶠ(i, j, k, ibg::PCBIBG) = min(Δzᶜᶜᶠ(i, j-1, k, ibg), Δzᶜᶜᶠ(i, j, k, ibg)) +@inline Δzᶜᶠᶠ(i, j, k, ibg::PCBIBG) = min(Δzᶜᶜᶠ(i, j-1, k, ibg), Δzᶜᶜᶠ(i, j, k, ibg)) @inline Δzᶠᶠᶠ(i, j, k, ibg::PCBIBG) = min(Δzᶠᶜᶠ(i, j-1, k, ibg), Δzᶠᶜᶠ(i, j, k, ibg)) # Make sure Δz works for horizontally-Flat topologies. @@ -202,3 +175,6 @@ YFlatPCBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Flat, <:Any, <:Any, <:Partial @inline Δzᶜᶠᶠ(i, j, k, ibg::YFlatPCBIBG) = Δzᶜᶜᶠ(i, j, k, ibg) @inline Δzᶠᶠᶜ(i, j, k, ibg::XFlatPCBIBG) = Δzᶜᶠᶜ(i, j, k, ibg) @inline Δzᶠᶠᶜ(i, j, k, ibg::YFlatPCBIBG) = Δzᶠᶜᶜ(i, j, k, ibg) + +@inline z_bottom(i, j, ibg::PCBIBG) = @inbounds ibg.immersed_boundary.bottom_height[i, j, 1] + diff --git a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl index 033e304998..a602a289f7 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl @@ -19,8 +19,6 @@ import Oceananigans.Advection: cell_advection_timescale import Oceananigans.TimeSteppers: step_lagrangian_particles! import Oceananigans.Architectures: on_architecture -using Oceananigans.TimeSteppers: SplitRungeKutta3TimeStepper, QuasiAdamsBashforth2TimeStepper - abstract type AbstractFreeSurface{E, G} end # This is only used by the cubed sphere for now. @@ -33,13 +31,8 @@ fill_horizontal_velocity_halos!(args...) = nothing free_surface_displacement_field(velocities, free_surface, grid) = ZFaceField(grid, indices = (:, :, size(grid, 3)+1)) free_surface_displacement_field(velocities, ::Nothing, grid) = nothing -# free surface initialization functions -initialize_free_surface!(free_surface, grid, velocities) = nothing - include("compute_w_from_continuity.jl") - -# No free surface -include("nothing_free_surface.jl") +include("rigid_lid.jl") # Explicit free-surface solver functionality include("explicit_free_surface.jl") @@ -53,9 +46,9 @@ include("matrix_implicit_free_surface_solver.jl") include("implicit_free_surface.jl") # Split-Explicit free-surface solver functionality -include("SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl") - -using .SplitExplicitFreeSurfaces +include("split_explicit_free_surface.jl") +include("distributed_split_explicit_free_surface.jl") +include("split_explicit_free_surface_kernels.jl") include("hydrostatic_free_surface_field_tuples.jl") include("hydrostatic_free_surface_model.jl") @@ -92,13 +85,6 @@ Return a flattened `NamedTuple` of the prognostic fields associated with `Hydros η = free_surface.η), tracers) -@inline hydrostatic_prognostic_fields(velocities, free_surface::SplitExplicitFreeSurface, tracers) = merge((u = velocities.u, - v = velocities.v, - η = free_surface.η, - U = free_surface.barotropic_velocities.U, - V = free_surface.barotropic_velocities.V), - tracers) - @inline hydrostatic_prognostic_fields(velocities, ::Nothing, tracers) = merge((u = velocities.u, v = velocities.v), tracers) @@ -109,14 +95,6 @@ Return a flattened `NamedTuple` of the prognostic fields associated with `Hydros tracers, (; η = free_surface.η)) -@inline hydrostatic_fields(velocities, free_surface::SplitExplicitFreeSurface, tracers) = merge((u = velocities.u, - v = velocities.v, - w = velocities.w, - η = free_surface.η, - U = free_surface.barotropic_velocities.U, - V = free_surface.barotropic_velocities.V), - tracers) - @inline hydrostatic_fields(velocities, ::Nothing, tracers) = merge((u = velocities.u, v = velocities.v, w = velocities.w), @@ -134,7 +112,6 @@ include("compute_hydrostatic_free_surface_tendencies.jl") include("compute_hydrostatic_free_surface_buffers.jl") include("update_hydrostatic_free_surface_model_state.jl") include("hydrostatic_free_surface_ab2_step.jl") -include("hydrostatic_free_surface_rk3_step.jl") include("store_hydrostatic_free_surface_tendencies.jl") include("prescribed_hydrostatic_velocity_fields.jl") include("single_column_model_mode.jl") diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl deleted file mode 100644 index 6ce20fd0b7..0000000000 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl +++ /dev/null @@ -1,43 +0,0 @@ -module SplitExplicitFreeSurfaces - -export SplitExplicitFreeSurface, ForwardBackwardScheme, AdamsBashforth3Scheme -export FixedSubstepNumber, FixedTimeStepSize - -using Oceananigans -using Oceananigans.Architectures -using Oceananigans.Architectures: convert_args -using Oceananigans.Fields -using Oceananigans.Utils -using Oceananigans.Grids -using Oceananigans.Operators -using Oceananigans.BoundaryConditions -using Oceananigans.Grids: AbstractGrid, topology -using Oceananigans.ImmersedBoundaries: active_linear_index_to_tuple, mask_immersed_field! -using Oceananigans.Models.HydrostaticFreeSurfaceModels: AbstractFreeSurface, - free_surface_displacement_field - -using Adapt -using Base -using KernelAbstractions: @index, @kernel -using KernelAbstractions.Extras.LoopInfo: @unroll - -import Oceananigans.Models.HydrostaticFreeSurfaceModels: initialize_free_surface!, - materialize_free_surface, - step_free_surface!, - compute_free_surface_tendency!, - explicit_barotropic_pressure_x_gradient, - explicit_barotropic_pressure_y_gradient - -include("split_explicit_timesteppers.jl") -include("split_explicit_free_surface.jl") -include("distributed_split_explicit_free_surface.jl") -include("initialize_split_explicit_substepping.jl") -include("compute_slow_tendencies.jl") -include("step_split_explicit_free_surface.jl") -include("barotropic_split_explicit_corrector.jl") - -# extend -@inline explicit_barotropic_pressure_x_gradient(i, j, k, grid, ::SplitExplicitFreeSurface) = zero(grid) -@inline explicit_barotropic_pressure_y_gradient(i, j, k, grid, ::SplitExplicitFreeSurface) = zero(grid) - -end diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl deleted file mode 100644 index 2b7f7845f6..0000000000 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl +++ /dev/null @@ -1,60 +0,0 @@ -# Kernels to compute the vertical integral of the velocities -@kernel function _barotropic_mode_kernel!(U, V, grid, ::Nothing, u, v) - i, j = @index(Global, NTuple) - barotropic_mode_kernel!(U, V, i, j, grid, u, v) -end - -@kernel function _barotropic_mode_kernel!(U, V, grid, active_cells_map, u, v) - idx = @index(Global, Linear) - i, j = active_linear_index_to_tuple(idx, active_cells_map) - barotropic_mode_kernel!(U, V, i, j, grid, u, v) -end - -@inline function barotropic_mode_kernel!(U, V, i, j, grid, u, v) - @inbounds U[i, j, 1] = Δzᶠᶜᶜ(i, j, 1, grid) * u[i, j, 1] - @inbounds V[i, j, 1] = Δzᶜᶠᶜ(i, j, 1, grid) * v[i, j, 1] - - for k in 2:grid.Nz - @inbounds U[i, j, 1] += Δzᶠᶜᶜ(i, j, k, grid) * u[i, j, k] - @inbounds V[i, j, 1] += Δzᶜᶠᶜ(i, j, k, grid) * v[i, j, k] - end - - return nothing -end - -@inline function compute_barotropic_mode!(U, V, grid, u, v) - active_cells_map = retrieve_surface_active_cells_map(grid) - launch!(architecture(grid), grid, :xy, _barotropic_mode_kernel!, U, V, grid, active_cells_map, u, v; active_cells_map) - return nothing -end - -@kernel function _barotropic_split_explicit_corrector!(u, v, U, V, U̅, V̅, grid) - i, j, k = @index(Global, NTuple) - - @inbounds begin - Hᶠᶜ = static_column_depthᶠᶜᵃ(i, j, grid) - Hᶜᶠ = static_column_depthᶜᶠᵃ(i, j, grid) - - u[i, j, k] = u[i, j, k] + (U[i, j, 1] - U̅[i, j, 1]) / Hᶠᶜ - v[i, j, k] = v[i, j, k] + (V[i, j, 1] - V̅[i, j, 1]) / Hᶜᶠ - end -end - -# Correcting `u` and `v` with the barotropic mode computed in `free_surface` -function barotropic_split_explicit_corrector!(u, v, free_surface, grid) - state = free_surface.filtered_state - U, V = free_surface.barotropic_velocities - U̅, V̅ = state.U, state.V - arch = architecture(grid) - - # NOTE: the filtered `U̅` and `V̅` have been copied in the instantaneous `U` and `V`, - # so we use the filtered velocities as "work arrays" to store the vertical integrals - # of the instantaneous velocities `u` and `v`. - compute_barotropic_mode!(U̅, V̅, grid, u, v) - - # add in "good" barotropic mode - launch!(arch, grid, :xyz, _barotropic_split_explicit_corrector!, - u, v, U, V, U̅, V̅, grid) - - return nothing -end diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl deleted file mode 100644 index 30f327ed20..0000000000 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl +++ /dev/null @@ -1,154 +0,0 @@ -##### -##### Compute slow tendencies with an AB2 timestepper -##### - -# Calculate RHS for the barotropic time step. -@kernel function _compute_integrated_ab2_tendencies!(Gᵁ, Gⱽ, grid, ::Nothing, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ) - i, j = @index(Global, NTuple) - ab2_integrate_tendencies!(Gᵁ, Gⱽ, i, j, grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ) -end - -@kernel function _compute_integrated_ab2_tendencies!(Gᵁ, Gⱽ, grid, active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ) - idx = @index(Global, Linear) - i, j = active_linear_index_to_tuple(idx, active_cells_map) - ab2_integrate_tendencies!(Gᵁ, Gⱽ, i, j, grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ) -end - -@inline function ab2_integrate_tendencies!(Gᵁ, Gⱽ, i, j, grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ) - locU = (Face(), Center(), Center()) - locV = (Center(), Face(), Center()) - - @inbounds Gᵁ[i, j, 1] = Δzᶠᶜᶜ(i, j, 1, grid) * ab2_step_G(i, j, 1, grid, locU..., Gu⁻, Guⁿ, χ) - @inbounds Gⱽ[i, j, 1] = Δzᶜᶠᶜ(i, j, 1, grid) * ab2_step_G(i, j, 1, grid, locV..., Gv⁻, Gvⁿ, χ) - - for k in 2:grid.Nz - @inbounds Gᵁ[i, j, 1] += Δzᶠᶜᶜ(i, j, k, grid) * ab2_step_G(i, j, k, grid, locU..., Gu⁻, Guⁿ, χ) - @inbounds Gⱽ[i, j, 1] += Δzᶜᶠᶜ(i, j, k, grid) * ab2_step_G(i, j, k, grid, locV..., Gv⁻, Gvⁿ, χ) - end -end - -@inline function ab2_step_G(i, j, k, grid, ℓx, ℓy, ℓz, G⁻, Gⁿ, χ::FT) where FT - C₁ = convert(FT, 3/2) + χ - C₂ = convert(FT, 1/2) + χ - - # multiply G⁻ by false if C₂ is zero to - # prevent propagationg possible NaNs - not_euler = C₂ != 0 - - Gⁿ⁺¹ = @inbounds C₁ * Gⁿ[i, j, k] - C₂ * G⁻[i, j, k] * not_euler - immersed = peripheral_node(i, j, k, grid, ℓx, ℓy, ℓz) - - return ifelse(immersed, zero(grid), Gⁿ⁺¹) -end - -@inline function compute_split_explicit_forcing!(GUⁿ, GVⁿ, grid, Guⁿ, Gvⁿ, - timestepper::QuasiAdamsBashforth2TimeStepper, stage) - active_cells_map = retrieve_surface_active_cells_map(grid) - - Gu⁻ = timestepper.G⁻.u - Gv⁻ = timestepper.G⁻.v - - launch!(architecture(grid), grid, :xy, _compute_integrated_ab2_tendencies!, GUⁿ, GVⁿ, grid, - active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, timestepper.χ; active_cells_map) - - return nothing -end - -##### -##### Compute slow tendencies with a RK3 timestepper -##### - -@inline function G_vertical_integral(i, j, grid, Gⁿ, ℓx, ℓy, ℓz) - immersed = peripheral_node(i, j, 1, grid, ℓx, ℓy, ℓz) - - Gⁿ⁺¹ = Δz(i, j, 1, grid, ℓx, ℓy, ℓz) * ifelse(immersed, zero(grid), Gⁿ[i, j, 1]) - - @inbounds for k in 2:grid.Nz - immersed = peripheral_node(i, j, k, grid, ℓx, ℓy, ℓz) - Gⁿ⁺¹ += Δz(i, j, k, grid, ℓx, ℓy, ℓz) * ifelse(immersed, zero(grid), Gⁿ[i, j, k]) - end - - return Gⁿ⁺¹ -end - -@kernel function _compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, grid, active_cells_map, Guⁿ, Gvⁿ, stage) - idx = @index(Global, Linear) - i, j = active_linear_index_to_tuple(idx, active_cells_map) - compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, i, j, grid, Guⁿ, Gvⁿ, stage) -end - -@kernel function _compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, grid, ::Nothing, Guⁿ, Gvⁿ, stage) - i, j = @index(Global, NTuple) - compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, i, j, grid, Guⁿ, Gvⁿ, stage) -end - -@inline function compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, i, j, grid, Guⁿ, Gvⁿ, ::Val{1}) - @inbounds GUⁿ[i, j, 1] = G_vertical_integral(i, j, grid, Guⁿ, Face(), Center(), Center()) - @inbounds GVⁿ[i, j, 1] = G_vertical_integral(i, j, grid, Gvⁿ, Center(), Face(), Center()) - - @inbounds GU⁻[i, j, 1] = GUⁿ[i, j, 1] - @inbounds GV⁻[i, j, 1] = GVⁿ[i, j, 1] -end - -@inline function compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, i, j, grid, Guⁿ, Gvⁿ, ::Val{2}) - FT = eltype(GUⁿ) - - @inbounds GUⁿ[i, j, 1] = G_vertical_integral(i, j, grid, Guⁿ, Face(), Center(), Center()) - @inbounds GVⁿ[i, j, 1] = G_vertical_integral(i, j, grid, Gvⁿ, Center(), Face(), Center()) - - @inbounds GU⁻[i, j, 1] = (GUⁿ[i, j, 1] + GU⁻[i, j, 1]) / 6 - @inbounds GV⁻[i, j, 1] = (GVⁿ[i, j, 1] + GV⁻[i, j, 1]) / 6 -end - -@inline function compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, i, j, grid, Guⁿ, Gvⁿ, ::Val{3}) - FT = eltype(GUⁿ) - - GUi = G_vertical_integral(i, j, grid, Guⁿ, Face(), Center(), Center()) - GVi = G_vertical_integral(i, j, grid, Gvⁿ, Center(), Face(), Center()) - - @inbounds GUⁿ[i, j, 1] = convert(FT, 2/3) * GUi + GU⁻[i, j, 1] - @inbounds GVⁿ[i, j, 1] = convert(FT, 2/3) * GVi + GV⁻[i, j, 1] -end - -@inline function compute_split_explicit_forcing!(GUⁿ, GVⁿ, grid, Guⁿ, Gvⁿ, - timestepper::SplitRungeKutta3TimeStepper, stage) - - GU⁻ = timestepper.G⁻.U - GV⁻ = timestepper.G⁻.V - - active_cells_map = retrieve_surface_active_cells_map(grid) - launch!(architecture(grid), grid, :xy, _compute_integrated_rk3_tendencies!, - GUⁿ, GVⁿ, GU⁻, GV⁻, grid, active_cells_map, Guⁿ, Gvⁿ, stage; active_cells_map) - - return nothing -end - -##### -##### Free surface setup -##### - -# Setting up the RHS for the barotropic step (tendencies of the barotropic velocity components) -# This function is called after `calculate_tendency` and before `ab2_step_velocities!` -function compute_free_surface_tendency!(grid, model, free_surface::SplitExplicitFreeSurface) - - Guⁿ = model.timestepper.Gⁿ.u - Gvⁿ = model.timestepper.Gⁿ.v - - GUⁿ = model.timestepper.Gⁿ.U - GVⁿ = model.timestepper.Gⁿ.V - - barotropic_timestepper = free_surface.timestepper - baroclinic_timestepper = model.timestepper - - stage = model.clock.stage - - @apply_regionally begin - compute_split_explicit_forcing!(GUⁿ, GVⁿ, grid, Guⁿ, Gvⁿ, baroclinic_timestepper, Val(stage)) - initialize_free_surface_state!(free_surface, baroclinic_timestepper, barotropic_timestepper, Val(stage)) - end - - fields_to_fill = (GUⁿ, GVⁿ) - fill_halo_regions!(fields_to_fill; async = true) - - return nothing -end diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/distributed_split_explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/distributed_split_explicit_free_surface.jl deleted file mode 100644 index 4560467668..0000000000 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/distributed_split_explicit_free_surface.jl +++ /dev/null @@ -1,25 +0,0 @@ -using Oceananigans.DistributedComputations: DistributedField -using Oceananigans.DistributedComputations: SynchronizedDistributed, synchronize_communication! - -const DistributedSplitExplicit = SplitExplicitFreeSurface{<:DistributedField} - -wait_free_surface_communication!(free_surface, model, arch) = nothing -wait_free_surface_communication!(::DistributedSplitExplicit, model, ::SynchronizedDistributed) = nothing - -function wait_free_surface_communication!(free_surface::DistributedSplitExplicit, model, arch) - - barotropic_velocities = free_surface.barotropic_velocities - - for field in (barotropic_velocities.U, barotropic_velocities.V) - synchronize_communication!(field) - end - - Gᵁ = model.timestepper.Gⁿ.U - Gⱽ = model.timestepper.Gⁿ.V - - for field in (Gᵁ, Gⱽ) - synchronize_communication!(field) - end - - return nothing -end diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl deleted file mode 100644 index 4c2976591f..0000000000 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl +++ /dev/null @@ -1,60 +0,0 @@ -using Oceananigans.ImmersedBoundaries: retrieve_surface_active_cells_map, peripheral_node -using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper, SplitRungeKutta3TimeStepper -using Oceananigans.Operators: Δz - -# This file contains two different initializations methods performed at different stages of the simulation. -# -# - `initialize_free_surface!`: the first initialization, performed only once at the beginning of the simulation, -# calculates the barotropic velocities from the velocity initial conditions. -# -# - `initialize_free_surface_state!`: is performed at the beginning of the substepping procedure, resets the filtered state to zero -# and reinitializes the timestepper auxiliaries from the previous filtered state. - -# `initialize_free_surface!` is called at the beginning of the simulation to initialize the free surface state -# from the initial velocity conditions. -function initialize_free_surface!(sefs::SplitExplicitFreeSurface, grid, velocities) - barotropic_velocities = sefs.barotropic_velocities - @apply_regionally compute_barotropic_mode!(barotropic_velocities.U, barotropic_velocities.V, grid, velocities.u, velocities.v) - fill_halo_regions!((barotropic_velocities.U, barotropic_velocities.V)) - return nothing -end - -# `initialize_free_surface_state!` is called at the beginning of the substepping to -# reset the filtered state to zero and reinitialize the state from the filtered state. -function initialize_free_surface_state!(free_surface, baroclinic_timestepper, timestepper, stage) - - η = free_surface.η - U, V = free_surface.barotropic_velocities - - initialize_free_surface_timestepper!(timestepper, η, U, V) - - fill!(free_surface.filtered_state.η, 0) - fill!(free_surface.filtered_state.U, 0) - fill!(free_surface.filtered_state.V, 0) - - return nothing -end - -# At the last stage we reset the velocities and perform the complete substepping from n to n+1 -function initialize_free_surface_state!(free_surface, baroclinic_ts::SplitRungeKutta3TimeStepper, barotropic_ts, ::Val{3}) - - η = free_surface.η - U, V = free_surface.barotropic_velocities - - Uⁿ⁻¹ = baroclinic_ts.Ψ⁻.U - Vⁿ⁻¹ = baroclinic_ts.Ψ⁻.V - ηⁿ⁻¹ = baroclinic_ts.Ψ⁻.η - - # Restart from the state at baroclinic step n - parent(U) .= parent(Uⁿ⁻¹) - parent(V) .= parent(Vⁿ⁻¹) - parent(η) .= parent(ηⁿ⁻¹) - - initialize_free_surface_timestepper!(barotropic_ts, η, U, V) - - fill!(free_surface.filtered_state.η, 0) - fill!(free_surface.filtered_state.U, 0) - fill!(free_surface.filtered_state.V, 0) - - return nothing -end diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl deleted file mode 100644 index 7c1e920d06..0000000000 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl +++ /dev/null @@ -1,327 +0,0 @@ -using Oceananigans.BuoyancyModels: g_Earth -using Oceananigans.Grids: with_halo -import Oceananigans.Grids: on_architecture - -struct SplitExplicitFreeSurface{H, U, M, FT, K , S, T} <: AbstractFreeSurface{H, FT} - η :: H - barotropic_velocities :: U # A namedtuple with U, V - filtered_state :: M # A namedtuple with η, U, V averaged throughout the substepping - gravitational_acceleration :: FT - kernel_parameters :: K - substepping :: S # Either `FixedSubstepNumber` or `FixedTimeStepSize` - timestepper :: T # Contains all auxiliary field and settings necessary to the particular timestepping -end - -""" - SplitExplicitFreeSurface(grid = nothing; - gravitational_acceleration = g_Earth, - substeps = nothing, - cfl = nothing, - fixed_Δt = nothing, - averaging_kernel = averaging_shape_function, - timestepper = ForwardBackwardScheme()) - -Return a `SplitExplicitFreeSurface` representing an explicit time discretization of -a free surface dynamics with `gravitational_acceleration`. The free surface dynamics are solved by discretizing: -```math -\begin{gather} -∂_t η = - \nabla ⋅ U, \\ -∂_t U = - g H \nabla η + G^U, -\end{gather} -``` -where ``η`` is the free surface displacement, ``U`` is the barotropic velocity vector, calculated as the vertical integral -of the velocity field ``u`` and ``v``, ``H`` is the column depth, ``G^U`` is the slow forcing calculated as the integral of the -tendency of ``u`` and ``v``, and ``g`` is the gravitational acceleration. - -The discretized equations are solved within a baroclinic timestep (``Δt``) by substepping with a ``Δτ < Δt``. -The barotropic velocities are filtered throughout the substepping and, finally, -the barotropic mode of the velocities at the new time step is corrected with the filtered velocities. - -Keyword Arguments -================= - -- `gravitational_acceleration`: the gravitational acceleration (default: `g_Earth`) - -- `substeps`: The number of substeps that divide the range `(t, t + 2Δt)`, where `Δt` is the baroclinic - timestep. Note that some averaging functions do not require substepping until `2Δt`. - The number of substeps is reduced automatically to the last index of `averaging_kernel` - for which `averaging_kernel > 0`. - -- `cfl`: If set then the number of `substeps` are computed based on the advective timescale imposed from - the barotropic gravity-wave speed that corresponds to depth `grid.Lz`. If `fixed_Δt` is provided, - then the number of `substeps` adapts to maintain an exact `cfl`. If not, the effective cfl will - always be lower than the specified `cfl` provided that the baroclinic time step satisfies - `Δt_baroclinic < fixed_Δt`. - -!!! info "Needed keyword arguments" - Either `substeps` _or_ `cfl` need to be prescribed. - - When `cfl` is prescribed then `grid` is also required as a positional argument. - -- `fixed_Δt`: The maximum baroclinic timestep allowed. If `fixed_Δt` is a `nothing` and a cfl is provided, - then the number of substeps will be computed on the fly from the baroclinic time step to - maintain a constant cfl. - -- `averaging_kernel`: A function of `τ` used to average the barotropic transport `U` and the free surface - `η` within the barotropic advancement. `τ` is the fractional substep going from 0 to 2 - with the baroclinic time step `t + Δt` located at `τ = 1`. The `averaging_kernel` - function should be centered at `τ = 1`, that is, ``∑ (aₘ m / M) = 1``, where the - the summation occurs for ``m = 1, ..., M_*``. Here, ``m = 0`` and ``m = M`` correspond - to the two consecutive baroclinic timesteps between which the barotropic timestepping - occurs and ``M_*`` corresponds to the last barotropic time step for which the - `averaging_kernel > 0`. By default, the averaging kernel described by [Shchepetkin2005](@citet) - is used. - -- `timestepper`: Time stepping scheme used for the barotropic advancement. Choose one of: - * `ForwardBackwardScheme()` (default): `η = f(U)` then `U = f(η)`, - * `AdamsBashforth3Scheme()`: `η = f(U, Uᵐ⁻¹, Uᵐ⁻²)` then `U = f(η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻²)`. - -References -========== - -Shchepetkin, A. F., & McWilliams, J. C. (2005). The regional oceanic modeling system (ROMS): a split-explicit, free-surface, topography-following-coordinate oceanic model. Ocean Modelling, 9(4), 347-404. -""" -function SplitExplicitFreeSurface(grid = nothing; - gravitational_acceleration = g_Earth, - substeps = nothing, - cfl = nothing, - fixed_Δt = nothing, - averaging_kernel = averaging_shape_function, - timestepper = ForwardBackwardScheme()) - - if !isnothing(grid) - FT = eltype(grid) - else - # this is a fallback and only used via the outer constructor, - # in case no grid is provided; when afterwards the free surfade - # is materialized via materialize_free_surface - # FT becomes eltype(grid) - FT = Float64 - end - - if isnothing(grid) && !isnothing(cfl) && !isnothing(fixed_Δt) - throw(ArgumentError(string("The grid is a required positional argument to SplitExplicitFreeSurface when cfl is specified and `fixed_Δt != nothing`.", - "For example, SplitExplicitFreeSurface(grid; fixed_Δt=$(fixed_Δt), cfl=$(cfl), ...)"))) - end - - gravitational_acceleration = convert(FT, gravitational_acceleration) - substepping = split_explicit_substepping(cfl, substeps, fixed_Δt, grid, averaging_kernel, gravitational_acceleration) - - return SplitExplicitFreeSurface(nothing, - nothing, - nothing, - gravitational_acceleration, - nothing, - substepping, - timestepper) -end - -# Simplest case: we have the substeps and the averaging kernel -function split_explicit_substepping(::Nothing, substeps, fixed_Δt, grid, averaging_kernel, gravitational_acceleration) - FT = eltype(gravitational_acceleration) - fractional_step_size, averaging_weights = weights_from_substeps(FT, substeps, averaging_kernel) - return FixedSubstepNumber(fractional_step_size, averaging_weights) -end - -# The substeps are calculated dynamically when a cfl without a fixed_Δt is provided -function split_explicit_substepping(cfl, ::Nothing, ::Nothing, grid, averaging_kernel, gravitational_acceleration) - cfl = convert(eltype(grid), cfl) - return FixedTimeStepSize(grid; cfl, averaging_kernel) -end - -# The number of substeps are calculated based on the cfl and the fixed_Δt -function split_explicit_substepping(cfl, ::Nothing, fixed_Δt, grid, averaging_kernel, gravitational_acceleration) - substepping = split_explicit_substepping(cfl, nothing, nothing, grid, averaging_kernel, gravitational_acceleration) - substeps = ceil(Int, 2 * fixed_Δt / substepping.Δt_barotropic) - substepping = split_explicit_substepping(nothing, substeps, nothing, grid, averaging_kernel, gravitational_acceleration) - return substepping -end - -# TODO: When open boundary conditions are online -# We need to calculate the barotropic boundary conditions -# from the baroclinic boundary conditions by integrating the BC upwards -@inline west_barotropic_bc(baroclinic_velocity) = baroclinic_velocity.boundary_conditions.west -@inline east_barotropic_bc(baroclinic_velocity) = baroclinic_velocity.boundary_conditions.east -@inline south_barotropic_bc(baroclinic_velocity) = baroclinic_velocity.boundary_conditions.south -@inline north_barotropic_bc(baroclinic_velocity) = baroclinic_velocity.boundary_conditions.north - -@inline barotropic_bc(baroclinic_velocity) = FieldBoundaryConditions( - west = west_barotropic_bc(baroclinic_velocity), - east = east_barotropic_bc(baroclinic_velocity), - south = south_barotropic_bc(baroclinic_velocity), - north = north_barotropic_bc(baroclinic_velocity), - top = nothing, - bottom = nothing -) - -const ConnectedTopologies = Union{LeftConnected, RightConnected, FullyConnected} - -# Internal function for HydrostaticFreeSurfaceModel -function materialize_free_surface(free_surface::SplitExplicitFreeSurface, velocities, grid) - - TX, TY, _ = topology(grid) - substepping = free_surface.substepping - - if (TX() isa ConnectedTopologies || TY() isa ConnectedTopologies) && substepping isa FixedTimeStepSize - throw(ArgumentError("A variable substepping through a CFL condition is not supported for the `SplitExplicitFreeSurface` on $(summary(grid)). \n - Provide a fixed number of substeps through the `substeps` keyword argument as: \n - `free_surface = SplitExplicitFreeSurface(grid; substeps = N)` where `N::Int`")) - end - - maybe_extended_grid = maybe_extend_halos(TX, TY, grid, substepping) - - η = free_surface_displacement_field(velocities, free_surface, maybe_extended_grid) - η̅ = free_surface_displacement_field(velocities, free_surface, maybe_extended_grid) - - u_baroclinic = velocities.u - v_baroclinic = velocities.v - - u_bc = barotropic_bc(u_baroclinic) - v_bc = barotropic_bc(v_baroclinic) - - U = Field{Face, Center, Nothing}(maybe_extended_grid, boundary_conditions = u_bc) - V = Field{Center, Face, Nothing}(maybe_extended_grid, boundary_conditions = v_bc) - - U̅ = Field{Face, Center, Nothing}(maybe_extended_grid, boundary_conditions = u_bc) - V̅ = Field{Center, Face, Nothing}(maybe_extended_grid, boundary_conditions = v_bc) - - filtered_state = (η = η̅, U = U̅, V = V̅) - barotropic_velocities = (U = U, V = V) - - kernel_parameters = maybe_augmented_kernel_parameters(TX, TY, substepping, maybe_extended_grid) - - gravitational_acceleration = convert(eltype(grid), free_surface.gravitational_acceleration) - timestepper = materialize_timestepper(free_surface.timestepper, maybe_extended_grid, free_surface, velocities, u_bc, v_bc) - - return SplitExplicitFreeSurface(η, - barotropic_velocities, - filtered_state, - gravitational_acceleration, - kernel_parameters, - substepping, - timestepper) -end - -# (p = 2, q = 4, r = 0.18927) minimize dispersion error from Shchepetkin and McWilliams (2005): https://doi.org/10.1016/j.ocemod.2004.08.002 -@inline function averaging_shape_function(τ::FT; p = 2, q = 4, r = FT(0.18927)) where FT - τ₀ = (p + 2) * (p + q + 2) / (p + 1) / (p + q + 1) - - return (τ / τ₀)^p * (1 - (τ / τ₀)^q) - r * (τ / τ₀) -end - -@inline cosine_averaging_kernel(τ::FT) where FT = τ ≥ 0.5 && τ ≤ 1.5 ? convert(FT, 1 + cos(2π * (τ - 1))) : zero(FT) -@inline constant_averaging_kernel(τ::FT) where FT = convert(FT, 1) - -""" An internal type for the `SplitExplicitFreeSurface` that allows substepping with -a fixed `Δt_barotropic` based on a CFL condition """ -struct FixedTimeStepSize{B, F} - Δt_barotropic :: B - averaging_kernel :: F -end - -""" An internal type for the `SplitExplicitFreeSurface` that allows substepping with -a fixed number of substeps with time step size of `fractional_step_size * Δt_baroclinic` """ -struct FixedSubstepNumber{B, F} - fractional_step_size :: B - averaging_weights :: F -end - -function FixedTimeStepSize(grid; - cfl = 0.7, - averaging_kernel = averaging_shape_function, - gravitational_acceleration = g_Earth) - - FT = eltype(grid) - - Δx⁻² = topology(grid)[1] == Flat ? 0 : 1 / minimum_xspacing(grid)^2 - Δy⁻² = topology(grid)[2] == Flat ? 0 : 1 / minimum_yspacing(grid)^2 - Δs = sqrt(1 / (Δx⁻² + Δy⁻²)) - - wave_speed = sqrt(gravitational_acceleration * grid.Lz) - - Δt_barotropic = convert(FT, cfl * Δs / wave_speed) - - return FixedTimeStepSize(Δt_barotropic, averaging_kernel) -end - -@inline function weights_from_substeps(FT, substeps, averaging_kernel) - - τᶠ = range(FT(0), FT(2), length = substeps+1) - Δτ = τᶠ[2] - τᶠ[1] - - averaging_weights = map(averaging_kernel, τᶠ[2:end]) - idx = searchsortedlast(averaging_weights, 0, rev=true) - substeps = idx - - averaging_weights = averaging_weights[1:idx] - averaging_weights ./= sum(averaging_weights) - - return Δτ, tuple(averaging_weights...) -end - -Base.summary(s::FixedTimeStepSize) = string("FixedTimeStepSize($(prettytime(s.Δt_barotropic)))") -Base.summary(s::FixedSubstepNumber) = string("FixedSubstepNumber($(length(s.averaging_weights)))") - -Base.summary(sefs::SplitExplicitFreeSurface) = string("SplitExplicitFreeSurface substepping with $(summary(sefs.substepping))") - -Base.show(io::IO, sefs::SplitExplicitFreeSurface) = print(io, "$(summary(sefs))\n") - -##### -##### Maybe extend halos in Connected topologies -##### - -# Extending halos is not allowed with variable time-stepping -maybe_extend_halos(TX, TY, grid, ::FixedTimeStepSize) = grid - -function maybe_extend_halos(TX, TY, grid, substepping::FixedSubstepNumber) - - old_halos = halo_size(grid) - Nsubsteps = length(substepping.averaging_weights) - step_halo = Nsubsteps + 1 - - Hx = TX() isa ConnectedTopologies ? max(step_halo, old_halos[1]) : old_halos[1] - Hy = TY() isa ConnectedTopologies ? max(step_halo, old_halos[2]) : old_halos[2] - - new_halos = (Hx, Hy, old_halos[3]) - - return with_halo(new_halos, grid) -end - -maybe_augmented_kernel_parameters(TX, TY, ::FixedTimeStepSize, grid) = :xy - -function maybe_augmented_kernel_parameters(TX, TY, ::FixedSubstepNumber, grid) - Nx, Ny, _ = size(grid) - Hx, Hy, _ = halo_size(grid) - - kernel_sizes = map(split_explicit_kernel_size, (TX, TY), (Nx, Ny), (Hx, Hy)) - - return KernelParameters(kernel_sizes...) -end - -split_explicit_kernel_size(topo, N, H) = 1:N -split_explicit_kernel_size(::Type{FullyConnected}, N, H) = -H+2:N+H-1 -split_explicit_kernel_size(::Type{RightConnected}, N, H) = 1:N+H-1 -split_explicit_kernel_size(::Type{LeftConnected}, N, H) = -H+2:N - -# Adapt -Adapt.adapt_structure(to, free_surface::SplitExplicitFreeSurface) = - SplitExplicitFreeSurface(Adapt.adapt(to, free_surface.η), - Adapt.adapt(to, free_surface.barotropic_velocities), - Adapt.adapt(to, free_surface.filtered_state), - free_surface.gravitational_acceleration, - nothing, - Adapt.adapt(to, free_surface.substepping), - Adapt.adapt(to, free_surface.timestepper)) - -for Type in (:SplitExplicitFreeSurface, - :AdamsBashforth3Scheme, - :FixedTimeStepSize, - :FixedSubstepNumber) - - @eval begin - function on_architecture(to, fs::$Type) - args = Tuple(on_architecture(to, prop) for prop in propertynames(fs)) - return $Type(args...) - end - end -end diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl deleted file mode 100644 index e624054d5e..0000000000 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl +++ /dev/null @@ -1,160 +0,0 @@ -""" - struct ForwardBackwardScheme - -A timestepping scheme used for substepping in the split-explicit free surface solver. - -The equations are evolved as follows: -```math -\begin{gather} -η^{m+1} = η^m - Δτ (∂_x U^m + ∂_y V^m), \\ -U^{m+1} = U^m - Δτ (∂_x η^{m+1} - G^U), \\ -V^{m+1} = V^m - Δτ (∂_y η^{m+1} - G^V). -\end{gather} -``` -""" -struct ForwardBackwardScheme end - -materialize_timestepper(::ForwardBackwardScheme, grid, args...) = ForwardBackwardScheme() - -struct AdamsBashforth3Scheme{CC, FC, CF, FT} - ηᵐ :: CC - ηᵐ⁻¹ :: CC - ηᵐ⁻² :: CC - Uᵐ⁻¹ :: FC - Uᵐ⁻² :: FC - Vᵐ⁻¹ :: CF - Vᵐ⁻² :: CF - β :: FT - α :: FT - θ :: FT - γ :: FT - δ :: FT - ϵ :: FT - μ :: FT -end - -""" - AdamsBashforth3Scheme(; β = 0.281105, - α = 1.5 + β, - θ = -0.5 - 2β, - γ = 0.088, - δ = 0.614, - ϵ = 0.013, - μ = 1 - δ - γ - ϵ) - -Create an instance of `AdamsBashforth3Scheme` with the specified parameters. -This scheme is used for substepping in the split-explicit free surface solver, -where an AB3 extrapolation is used to evaluate barotropic velocities and -free surface at time-step `m + 1/2`: - -The equations are evolved as follows: - -```math -\begin{gather} -η^{m+1} = η^m - Δτ g H (∂_x Ũ + ∂y Ṽ), \\ -U^{m+1} = U^m - Δτ (∂_x η̃ - G^U), \\ -V^{m+1} = V^m - Δτ (∂_y η̃ - G^V), -\end{gather} -``` - -where `η̃`, `Ũ` and `Ṽ` are the AB3 time-extrapolated values of free surface, -barotropic zonal and meridional velocities, respectively: - -```math -\begin{gather} -Ũ = α U^m + θ U^{m-1} + β U^{m-2}, \\ -Ṽ = α V^m + θ V^{m-1} + β V^{m-2}, \\ -η̃ = δ η^{m+1} + μ η^m + γ η^{m-1} + ϵ η^{m-2}. -\end{gather} -``` - -The default values for the time-extrapolation coefficients, described by [Shchepetkin2005](@citet), -correspond to the best stability range for the AB3 algorithm. -""" -AdamsBashforth3Scheme(; β = 0.281105, α = 1.5 + β, θ = - 0.5 - 2β, γ = 0.088, δ = 0.614, ϵ = 0.013, μ = 1 - δ - γ - ϵ) = - AdamsBashforth3Scheme(nothing, nothing, nothing, nothing, nothing, nothing, nothing, β, α, θ, γ, δ, ϵ, μ) - -Adapt.adapt_structure(to, t::AdamsBashforth3Scheme) = - AdamsBashforth3Scheme( - Adapt.adapt(to, t.ηᵐ ), - Adapt.adapt(to, t.ηᵐ⁻¹), - Adapt.adapt(to, t.ηᵐ⁻²), - Adapt.adapt(to, t.Uᵐ⁻¹), - Adapt.adapt(to, t.Uᵐ⁻²), - Adapt.adapt(to, t.Vᵐ⁻¹), - Adapt.adapt(to, t.Vᵐ⁻²), - t.β, t.α, t.θ, t.γ, t.δ, t.ϵ, t.μ) - -function materialize_timestepper(t::AdamsBashforth3Scheme, grid, free_surface, velocities, u_bc, v_bc) - ηᵐ = free_surface_displacement_field(velocities, free_surface, grid) - ηᵐ⁻¹ = free_surface_displacement_field(velocities, free_surface, grid) - ηᵐ⁻² = free_surface_displacement_field(velocities, free_surface, grid) - - Uᵐ⁻¹ = Field{Face, Center, Nothing}(grid; boundary_conditions = u_bc) - Uᵐ⁻² = Field{Face, Center, Nothing}(grid; boundary_conditions = u_bc) - Vᵐ⁻¹ = Field{Center, Face, Nothing}(grid; boundary_conditions = v_bc) - Vᵐ⁻² = Field{Center, Face, Nothing}(grid; boundary_conditions = v_bc) - - FT = eltype(grid) - - β = convert(FT, t.β) - α = convert(FT, t.α) - θ = convert(FT, t.θ) - γ = convert(FT, t.γ) - δ = convert(FT, t.δ) - ϵ = convert(FT, t.ϵ) - μ = convert(FT, t.μ) - - return AdamsBashforth3Scheme(ηᵐ, ηᵐ⁻¹, ηᵐ⁻², Uᵐ⁻¹, Uᵐ⁻², Vᵐ⁻¹, Vᵐ⁻², β, α, θ, γ, δ, ϵ, μ) -end - -##### -##### Timestepper extrapolations and utils -##### - -function materialize_timestepper(name::Symbol, args...) - fullname = Symbol(name, :Scheme) - TS = getglobal(@__MODULE__, fullname) - return materialize_timestepper(TS, args...) -end - -initialize_free_surface_timestepper!(::ForwardBackwardScheme, args...) = nothing - -function initialize_free_surface_timestepper!(timestepper::AdamsBashforth3Scheme, η, U, V) - parent(timestepper.Uᵐ⁻¹) .= parent(U) - parent(timestepper.Vᵐ⁻¹) .= parent(V) - - parent(timestepper.Uᵐ⁻²) .= parent(U) - parent(timestepper.Vᵐ⁻²) .= parent(V) - - parent(timestepper.ηᵐ) .= parent(η) - parent(timestepper.ηᵐ⁻¹) .= parent(η) - parent(timestepper.ηᵐ⁻²) .= parent(η) - - return nothing -end - -# The functions `η★` `U★` and `V★` represent the value of free surface, barotropic zonal and meridional velocity at time step m+1/2 -@inline U★(i, j, k, grid, t::ForwardBackwardScheme, Uᵐ) = @inbounds Uᵐ[i, j, k] -@inline U★(i, j, k, grid, t::AdamsBashforth3Scheme, Uᵐ) = @inbounds t.α * Uᵐ[i, j, k] + t.θ * t.Uᵐ⁻¹[i, j, k] + t.β * t.Uᵐ⁻²[i, j, k] - -@inline η★(i, j, k, grid, t::ForwardBackwardScheme, ηᵐ⁺¹) = @inbounds ηᵐ⁺¹[i, j, k] -@inline η★(i, j, k, grid, t::AdamsBashforth3Scheme, ηᵐ⁺¹) = @inbounds t.δ * ηᵐ⁺¹[i, j, k] + t.μ * t.ηᵐ[i, j, k] + t.γ * t.ηᵐ⁻¹[i, j, k] + t.ϵ * t.ηᵐ⁻²[i, j, k] - -@inline store_previous_velocities!(::ForwardBackwardScheme, i, j, k, U) = nothing -@inline store_previous_free_surface!(::ForwardBackwardScheme, i, j, k, η) = nothing - -@inline function store_previous_velocities!(t::AdamsBashforth3Scheme, i, j, k, U) - @inbounds t.Uᵐ⁻²[i, j, k] = t.Uᵐ⁻¹[i, j, k] - @inbounds t.Uᵐ⁻¹[i, j, k] = U[i, j, k] - - return nothing -end - -@inline function store_previous_free_surface!(t::AdamsBashforth3Scheme, i, j, k, η) - @inbounds t.ηᵐ⁻²[i, j, k] = t.ηᵐ⁻¹[i, j, k] - @inbounds t.ηᵐ⁻¹[i, j, k] = t.ηᵐ[i, j, k] - @inbounds t.ηᵐ[i, j, k] = η[i, j, k] - - return nothing -end diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl deleted file mode 100644 index 655e2e9068..0000000000 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl +++ /dev/null @@ -1,168 +0,0 @@ -# Evolution Kernels -# -# ∂t(η) = -∇⋅U -# ∂t(U) = - gH∇η + f -# -# the free surface field η and its average η̄ are located on `Face`s at the surface (grid.Nz +1). All other intermediate variables -# (U, V, Ū, V̄) are barotropic fields (`ReducedField`) for which a k index is not defined - -@kernel function _split_explicit_free_surface!(grid, Δτ, η, U, V, timestepper) - i, j = @index(Global, NTuple) - k_top = grid.Nz+1 - - store_previous_free_surface!(timestepper, i, j, k_top, η) - @inbounds η[i, j, k_top] -= Δτ * (δxTᶜᵃᵃ(i, j, grid.Nz, grid, Δy_qᶠᶜᶠ, U★, timestepper, U) + - δyTᵃᶜᵃ(i, j, grid.Nz, grid, Δx_qᶜᶠᶠ, U★, timestepper, V)) / Azᶜᶜᶠ(i, j, k_top, grid) -end - -@kernel function _split_explicit_barotropic_velocity!(averaging_weight, grid, Δτ, - η, U, V, - η̅, U̅, V̅, - Gᵁ, Gⱽ, g, - timestepper) - i, j = @index(Global, NTuple) - k_top = grid.Nz+1 - - store_previous_velocities!(timestepper, i, j, 1, U) - store_previous_velocities!(timestepper, i, j, 1, V) - - Hᶠᶜ = static_column_depthᶠᶜᵃ(i, j, grid) - Hᶜᶠ = static_column_depthᶜᶠᵃ(i, j, grid) - - @inbounds begin - # ∂τ(U) = - ∇η + G - U[i, j, 1] += Δτ * (- g * Hᶠᶜ * ∂xTᶠᶜᶠ(i, j, k_top, grid, η★, timestepper, η) + Gᵁ[i, j, 1]) - V[i, j, 1] += Δτ * (- g * Hᶜᶠ * ∂yTᶜᶠᶠ(i, j, k_top, grid, η★, timestepper, η) + Gⱽ[i, j, 1]) - - # time-averaging - η̅[i, j, k_top] += averaging_weight * η[i, j, k_top] - U̅[i, j, 1] += averaging_weight * U[i, j, 1] - V̅[i, j, 1] += averaging_weight * V[i, j, 1] - end -end - -# Change name -const FNS = FixedSubstepNumber -const FTS = FixedTimeStepSize - -# since weights can be negative in the first few substeps (as in the default averaging kernel), -# we set a minimum number of substeps to execute to avoid numerical issues -const MINIMUM_SUBSTEPS = 5 - -@inline calculate_substeps(substepping::FNS, Δt=nothing) = length(substepping.averaging_weights) -@inline calculate_substeps(substepping::FTS, Δt) = max(MINIMUM_SUBSTEPS, ceil(Int, 2 * Δt / substepping.Δt_barotropic)) - -@inline calculate_adaptive_settings(substepping::FNS, substeps) = substepping.fractional_step_size, substepping.averaging_weights -@inline calculate_adaptive_settings(substepping::FTS, substeps) = weights_from_substeps(eltype(substepping.Δt_barotropic), - substeps, substepping.averaging_kernel) - -function iterate_split_explicit!(free_surface, grid, GUⁿ, GVⁿ, Δτᴮ, weights, ::Val{Nsubsteps}) where Nsubsteps - arch = architecture(grid) - - η = free_surface.η - grid = free_surface.η.grid - state = free_surface.filtered_state - timestepper = free_surface.timestepper - g = free_surface.gravitational_acceleration - parameters = free_surface.kernel_parameters - - # unpack state quantities, parameters and forcing terms - U, V = free_surface.barotropic_velocities - η̅, U̅, V̅ = state.η, state.U, state.V - - free_surface_kernel!, _ = configure_kernel(arch, grid, parameters, _split_explicit_free_surface!) - barotropic_velocity_kernel!, _ = configure_kernel(arch, grid, parameters, _split_explicit_barotropic_velocity!) - - η_args = (grid, Δτᴮ, η, U, V, - timestepper) - - U_args = (grid, Δτᴮ, η, U, V, - η̅, U̅, V̅, GUⁿ, GVⁿ, g, - timestepper) - - GC.@preserve η_args U_args begin - - # We need to perform ~50 time-steps which means - # launching ~100 very small kernels: we are limited by - # latency of argument conversion to GPU-compatible values. - # To alleviate this penalty we convert first and then we substep! - converted_η_args = convert_args(arch, η_args) - converted_U_args = convert_args(arch, U_args) - - @unroll for substep in 1:Nsubsteps - Base.@_inline_meta - averaging_weight = weights[substep] - free_surface_kernel!(converted_η_args...) - barotropic_velocity_kernel!(averaging_weight, converted_U_args...) - end - end - - return nothing -end - -@kernel function _update_split_explicit_state!(η, U, V, grid, η̅, U̅, V̅) - i, j = @index(Global, NTuple) - k_top = grid.Nz+1 - - @inbounds begin - η[i, j, k_top] = η̅[i, j, k_top] - U[i, j, 1] = U̅[i, j, 1] - V[i, j, 1] = V̅[i, j, 1] - end -end - -##### -##### SplitExplicitFreeSurface barotropic subcylicing -##### - -function step_free_surface!(free_surface::SplitExplicitFreeSurface, model, baroclinic_timestepper, Δt) - - # Note: free_surface.η.grid != model.grid for DistributedSplitExplicitFreeSurface - # since halo_size(free_surface.η.grid) != halo_size(model.grid) - free_surface_grid = free_surface.η.grid - filtered_state = free_surface.filtered_state - substepping = free_surface.substepping - - barotropic_velocities = free_surface.barotropic_velocities - - # Wait for setup step to finish - wait_free_surface_communication!(free_surface, model, architecture(free_surface_grid)) - - # Calculate the substepping parameterers - # barotropic time step as fraction of baroclinic step and averaging weights - Nsubsteps = calculate_substeps(substepping, Δt) - fractional_Δt, weights = calculate_adaptive_settings(substepping, Nsubsteps) - Nsubsteps = length(weights) - - # barotropic time step in seconds - Δτᴮ = fractional_Δt * Δt - - # Slow forcing terms - GUⁿ = model.timestepper.Gⁿ.U - GVⁿ = model.timestepper.Gⁿ.V - - #free surface state - η = free_surface.η - U = barotropic_velocities.U - V = barotropic_velocities.V - η̅ = filtered_state.η - U̅ = filtered_state.U - V̅ = filtered_state.V - - # reset free surface averages - @apply_regionally begin - # Solve for the free surface at tⁿ⁺¹ - iterate_split_explicit!(free_surface, free_surface_grid, GUⁿ, GVⁿ, Δτᴮ, weights, Val(Nsubsteps)) - - # Update eta and velocities for the next timestep - # The halos are updated in the `update_state!` function - launch!(architecture(free_surface_grid), free_surface_grid, :xy, - _update_split_explicit_state!, η, U, V, free_surface_grid, η̅, U̅, V̅) - - # Preparing velocities for the barotropic correction - mask_immersed_field!(model.velocities.u) - mask_immersed_field!(model.velocities.v) - end - - return nothing -end diff --git a/src/Models/HydrostaticFreeSurfaceModels/barotropic_pressure_correction.jl b/src/Models/HydrostaticFreeSurfaceModels/barotropic_pressure_correction.jl index 004f1099d1..382236feec 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/barotropic_pressure_correction.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/barotropic_pressure_correction.jl @@ -1,4 +1,3 @@ -using .SplitExplicitFreeSurfaces: barotropic_split_explicit_corrector! import Oceananigans.TimeSteppers: calculate_pressure_correction!, pressure_correct_velocities! calculate_pressure_correction!(::HydrostaticFreeSurfaceModel, Δt) = nothing @@ -7,20 +6,21 @@ calculate_pressure_correction!(::HydrostaticFreeSurfaceModel, Δt) = nothing ##### Barotropic pressure correction for models with a free surface ##### -pressure_correct_velocities!(model::HydrostaticFreeSurfaceModel, Δt; kwargs...) = - pressure_correct_velocities!(model, model.free_surface, Δt; kwargs...) +const HFSM = HydrostaticFreeSurfaceModel +const ExplicitFreeSurfaceHFSM = HFSM{<:Any, <:Any, <:Any, <:ExplicitFreeSurface} +const ImplicitFreeSurfaceHFSM = HFSM{<:Any, <:Any, <:Any, <:ImplicitFreeSurface} +const SplitExplicitFreeSurfaceHFSM = HFSM{<:Any, <:Any, <:Any, <:SplitExplicitFreeSurface} -# Fallback -pressure_correct_velocities!(model, free_surface, Δt; kwargs...) = nothing +pressure_correct_velocities!(model::ExplicitFreeSurfaceHFSM, Δt; kwargs...) = nothing ##### -##### Barotropic pressure correction for models with an Implicit free surface +##### Barotropic pressure correction for models with a free surface ##### -function pressure_correct_velocities!(model, ::ImplicitFreeSurface, Δt) +function pressure_correct_velocities!(model::ImplicitFreeSurfaceHFSM, Δt) launch!(model.architecture, model.grid, :xyz, - _barotropic_pressure_correction!, + _barotropic_pressure_correction, model.velocities, model.grid, Δt, @@ -30,7 +30,10 @@ function pressure_correct_velocities!(model, ::ImplicitFreeSurface, Δt) return nothing end -function pressure_correct_velocities!(model, ::SplitExplicitFreeSurface, Δt) +calculate_free_surface_tendency!(grid, ::ImplicitFreeSurfaceHFSM , args...) = nothing +calculate_free_surface_tendency!(grid, ::SplitExplicitFreeSurfaceHFSM, args...) = nothing + +function pressure_correct_velocities!(model::SplitExplicitFreeSurfaceHFSM, Δt) u, v, _ = model.velocities grid = model.grid barotropic_split_explicit_corrector!(u, v, model.free_surface, grid) @@ -38,7 +41,7 @@ function pressure_correct_velocities!(model, ::SplitExplicitFreeSurface, Δt) return nothing end -@kernel function _barotropic_pressure_correction!(U, grid, Δt, g, η) +@kernel function _barotropic_pressure_correction(U, grid, Δt, g, η) i, j, k = @index(Global, NTuple) @inbounds begin diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl index 9e21089913..82c71b05dd 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl @@ -38,6 +38,7 @@ function compute_tendencies!(model::HydrostaticFreeSurfaceModel, callbacks) compute_hydrostatic_boundary_tendency_contributions!(model.timestepper.Gⁿ, model.architecture, model.velocities, + model.free_surface, model.tracers, model.clock, fields(model), @@ -53,9 +54,6 @@ function compute_tendencies!(model::HydrostaticFreeSurfaceModel, callbacks) return nothing end -# Fallback -compute_free_surface_tendency!(grid, model, free_surface) = nothing - @inline function top_tracer_boundary_conditions(grid, tracers) names = propertynames(tracers) values = Tuple(tracers[c].boundary_conditions.top for c in names) @@ -91,8 +89,8 @@ function compute_hydrostatic_free_surface_tendency_contributions!(model, kernel_ model.tracers, model.diffusivity_fields, model.auxiliary_fields, - model.clock, - c_forcing) + c_forcing, + model.clock) launch!(arch, grid, kernel_parameters, compute_hydrostatic_free_surface_Gc!, @@ -117,6 +115,24 @@ function apply_flux_bcs!(Gcⁿ, c, arch, args) return nothing end +function compute_free_surface_tendency!(grid, model, kernel_parameters) + + arch = architecture(grid) + + args = tuple(model.velocities, + model.free_surface, + model.tracers, + model.auxiliary_fields, + model.forcing, + model.clock) + + launch!(arch, grid, kernel_parameters, + compute_hydrostatic_free_surface_Gη!, model.timestepper.Gⁿ.η, + grid, args) + + return nothing +end + """ Calculate momentum tendencies if momentum is not prescribed.""" function compute_hydrostatic_momentum_tendencies!(model, velocities, kernel_parameters; active_cells_map = nothing) @@ -137,6 +153,7 @@ function compute_hydrostatic_momentum_tendencies!(model, velocities, kernel_para model.diffusivity_fields, model.pressure.pHY′, model.auxiliary_fields, + model.forcing, model.clock) u_kernel_args = tuple(start_momentum_kernel_args..., u_immersed_bc, end_momentum_kernel_args...) @@ -144,20 +161,21 @@ function compute_hydrostatic_momentum_tendencies!(model, velocities, kernel_para launch!(arch, grid, kernel_parameters, compute_hydrostatic_free_surface_Gu!, model.timestepper.Gⁿ.u, grid, - active_cells_map, u_kernel_args, model.forcing.u; + active_cells_map, u_kernel_args; active_cells_map) launch!(arch, grid, kernel_parameters, compute_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, grid, - active_cells_map, v_kernel_args, model.forcing.v; + active_cells_map, v_kernel_args; active_cells_map) + compute_free_surface_tendency!(grid, model, :xy) + return nothing end - """ Apply boundary conditions by adding flux divergences to the right-hand-side. """ -function compute_hydrostatic_boundary_tendency_contributions!(Gⁿ, arch, velocities, tracers, args...) +function compute_hydrostatic_boundary_tendency_contributions!(Gⁿ, arch, velocities, free_surface, tracers, args...) args = Tuple(args) @@ -166,6 +184,9 @@ function compute_hydrostatic_boundary_tendency_contributions!(Gⁿ, arch, veloci apply_flux_bcs!(Gⁿ[i], velocities[i], arch, args) end + # Free surface + apply_flux_bcs!(Gⁿ.η, displacement(free_surface), arch, args) + # Tracer fields for i in propertynames(tracers) apply_flux_bcs!(Gⁿ[i], tracers[i], arch, args) @@ -179,27 +200,27 @@ end ##### """ Calculate the right-hand-side of the u-velocity equation. """ -@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, ::Nothing, args, forcing) +@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, ::Nothing, args) i, j, k = @index(Global, NTuple) - @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args..., forcing) + @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args...) end -@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, active_cells_map, args, forcing) +@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, active_cells_map, args) idx = @index(Global, Linear) i, j, k = active_linear_index_to_tuple(idx, active_cells_map) - @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args..., forcing) + @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args...) end """ Calculate the right-hand-side of the v-velocity equation. """ -@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, ::Nothing, args, forcing) +@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, ::Nothing, args) i, j, k = @index(Global, NTuple) - @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args..., forcing) + @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args...) end -@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, active_cells_map, args, forcing) +@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, active_cells_map, args) idx = @index(Global, Linear) i, j, k = active_linear_index_to_tuple(idx, active_cells_map) - @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args..., forcing) + @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args...) end ##### @@ -217,3 +238,13 @@ end i, j, k = active_linear_index_to_tuple(idx, active_cells_map) @inbounds Gc[i, j, k] = hydrostatic_free_surface_tracer_tendency(i, j, k, grid, args...) end + +##### +##### Tendency calculators for an explicit free surface +##### + +""" Calculate the right-hand-side of the free surface displacement (``η``) equation. """ +@kernel function compute_hydrostatic_free_surface_Gη!(Gη, grid, args) + i, j = @index(Global, NTuple) + @inbounds Gη[i, j, grid.Nz+1] = free_surface_tendency(i, j, grid, args...) +end diff --git a/src/Models/HydrostaticFreeSurfaceModels/distributed_split_explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/distributed_split_explicit_free_surface.jl new file mode 100644 index 0000000000..6f0afa5a26 --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/distributed_split_explicit_free_surface.jl @@ -0,0 +1,112 @@ +using Oceananigans.AbstractOperations: GridMetricOperation, Δz +using Oceananigans.DistributedComputations: DistributedGrid, DistributedField +using Oceananigans.DistributedComputations: SynchronizedDistributed, synchronize_communication! +using Oceananigans.Models.HydrostaticFreeSurfaceModels: SplitExplicitState, SplitExplicitFreeSurface + +import Oceananigans.Models.HydrostaticFreeSurfaceModels: materialize_free_surface, SplitExplicitAuxiliaryFields + +function SplitExplicitAuxiliaryFields(grid::DistributedGrid) + + Gᵁ = Field((Face, Center, Nothing), grid) + Gⱽ = Field((Center, Face, Nothing), grid) + + Hᶠᶜ = Field((Face, Center, Nothing), grid) + Hᶜᶠ = Field((Center, Face, Nothing), grid) + + calculate_column_height!(Hᶠᶜ, (Face, Center, Center)) + calculate_column_height!(Hᶜᶠ, (Center, Face, Center)) + + fill_halo_regions!((Hᶠᶜ, Hᶜᶠ)) + + # In a non-parallel grid we calculate only the interior + kernel_size = augmented_kernel_size(grid) + kernel_offsets = augmented_kernel_offsets(grid) + + kernel_parameters = KernelParameters(kernel_size, kernel_offsets) + + return SplitExplicitAuxiliaryFields(Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ, kernel_parameters) +end + +"""Integrate z at locations `location` and set! `height`` with the result""" +@inline function calculate_column_height!(height, location) + dz = GridMetricOperation(location, Δz, height.grid) + return sum!(height, dz) +end + +@inline function augmented_kernel_size(grid::DistributedGrid) + Nx, Ny, _ = size(grid) + Hx, Hy, _ = halo_size(grid) + + Tx, Ty, _ = topology(grid) + + Rx, Ry, _ = architecture(grid).ranks + + Ax = Rx == 1 ? Nx : (Tx == RightConnected || Tx == LeftConnected ? Nx + Hx - 1 : Nx + 2Hx - 2) + Ay = Ry == 1 ? Ny : (Ty == RightConnected || Ty == LeftConnected ? Ny + Hy - 1 : Ny + 2Hy - 2) + + return (Ax, Ay) +end + +@inline function augmented_kernel_offsets(grid::DistributedGrid) + Hx, Hy, _ = halo_size(grid) + Tx, Ty, _ = topology(grid) + + Rx, Ry, _ = architecture(grid).ranks + + Ax = Rx == 1 || Tx == RightConnected ? 0 : - Hx + 1 + Ay = Ry == 1 || Ty == RightConnected ? 0 : - Hy + 1 + + return (Ax, Ay) +end + +# Internal function for HydrostaticFreeSurfaceModel +function materialize_free_surface(free_surface::SplitExplicitFreeSurface, velocities, grid::DistributedGrid) + + settings = free_surface.settings + + old_halos = halo_size(grid) + Nsubsteps = length(settings.substepping.averaging_weights) + + extended_halos = distributed_split_explicit_halos(old_halos, Nsubsteps+1, grid) + extended_grid = with_halo(extended_halos, grid) + + Nze = size(extended_grid, 3) + η = ZFaceField(extended_grid, indices = (:, :, Nze+1)) + + return SplitExplicitFreeSurface(η, + SplitExplicitState(extended_grid, settings.timestepper), + SplitExplicitAuxiliaryFields(extended_grid), + free_surface.gravitational_acceleration, + free_surface.settings) +end + +@inline function distributed_split_explicit_halos(old_halos, step_halo, grid::DistributedGrid) + + Rx, Ry, _ = architecture(grid).ranks + + Ax = Rx == 1 ? old_halos[1] : max(step_halo, old_halos[1]) + Ay = Ry == 1 ? old_halos[2] : max(step_halo, old_halos[2]) + + return (Ax, Ay, old_halos[3]) +end + +const DistributedSplitExplicit = SplitExplicitFreeSurface{<:DistributedField} + +wait_free_surface_communication!(::DistributedSplitExplicit, ::SynchronizedDistributed) = nothing + +function wait_free_surface_communication!(free_surface::DistributedSplitExplicit, arch) + + state = free_surface.state + + for field in (state.U̅, state.V̅) + synchronize_communication!(field) + end + + auxiliary = free_surface.auxiliary + + for field in (auxiliary.Gᵁ, auxiliary.Gⱽ) + synchronize_communication!(field) + end + + return nothing +end diff --git a/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl index e788a1f555..56fd8d0049 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl @@ -50,26 +50,8 @@ end ##### Time stepping ##### -step_free_surface!(free_surface::ExplicitFreeSurface, model, timestepper::QuasiAdamsBashforth2TimeStepper, Δt) = - @apply_regionally explicit_ab2_step_free_surface!(free_surface, model, Δt, timestepper.χ) - -step_free_surface!(free_surface::ExplicitFreeSurface, model, timestepper::SplitRungeKutta3TimeStepper, Δt) = - @apply_regionally explicit_rk3_step_free_surface!(free_surface, model, Δt, timestepper) - -@inline rk3_coeffs(ts, ::Val{1}) = (1, 0) -@inline rk3_coeffs(ts, ::Val{2}) = (ts.γ², ts.ζ²) -@inline rk3_coeffs(ts, ::Val{3}) = (ts.γ³, ts.ζ³) - -function explicit_rk3_step_free_surface!(free_surface, model, Δt, timestepper) - - γⁿ, ζⁿ = rk3_coeffs(timestepper, Val(model.clock.stage)) - - launch!(model.architecture, model.grid, :xy, - _explicit_rk3_step_free_surface!, free_surface.η, Δt, γⁿ, ζⁿ, - model.timestepper.Gⁿ.η, model.timestepper.Ψ⁻.η, size(model.grid, 3)) - - return nothing -end +ab2_step_free_surface!(free_surface::ExplicitFreeSurface, model, Δt, χ) = + @apply_regionally explicit_ab2_step_free_surface!(free_surface, model, Δt, χ) explicit_ab2_step_free_surface!(free_surface, model, Δt, χ) = launch!(model.architecture, model.grid, :xy, @@ -77,86 +59,13 @@ explicit_ab2_step_free_surface!(free_surface, model, Δt, χ) = model.timestepper.Gⁿ.η, model.timestepper.G⁻.η, size(model.grid, 3)) ##### -##### Kernels +##### Kernel ##### -@kernel function _explicit_rk3_step_free_surface!(η, Δt, γⁿ, ζⁿ, Gⁿ, η⁻, Nz) - i, j = @index(Global, NTuple) - @inbounds η[i, j, Nz+1] += ζⁿ * η⁻[i, j, k] + γⁿ * (η[i, j, k] + convert(FT, Δt) * Gⁿ[i, j, k]) -end - @kernel function _explicit_ab2_step_free_surface!(η, Δt, χ::FT, Gηⁿ, Gη⁻, Nz) where FT i, j = @index(Global, NTuple) - @inbounds η[i, j, Nz+1] += Δt * ((FT(1.5) + χ) * Gηⁿ[i, j, Nz+1] - (FT(0.5) + χ) * Gη⁻[i, j, Nz+1]) -end - -##### -##### Tendency calculators for an explicit free surface -##### - -""" Calculate the right-hand-side of the free surface displacement (``η``) equation. """ -@kernel function compute_hydrostatic_free_surface_Gη!(Gη, grid, args) - i, j = @index(Global, NTuple) - @inbounds Gη[i, j, grid.Nz+1] = free_surface_tendency(i, j, grid, args...) -end - -""" - free_surface_tendency(i, j, grid, - velocities, - free_surface, - tracers, - auxiliary_fields, - forcings, - clock) -Return the tendency for an explicit free surface at horizontal grid point `i, j`. - -The tendency is called ``G_η`` and defined via - -```math -∂_t η = G_η -``` -""" -@inline function free_surface_tendency(i, j, grid, - velocities, - free_surface, - tracers, - auxiliary_fields, - forcings, - clock) - - k_top = grid.Nz + 1 - model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields) - - return @inbounds ( velocities.w[i, j, k_top] - + forcings.η(i, j, k_top, grid, clock, model_fields)) + @inbounds begin + η[i, j, Nz+1] += Δt * ((FT(1.5) + χ) * Gηⁿ[i, j, Nz+1] - (FT(0.5) + χ) * Gη⁻[i, j, Nz+1]) + end end - -compute_free_surface_tendency!(grid, model, ::ExplicitFreeSurface) = - @apply_regionally compute_explicit_free_surface_tendency!(grid, model) - -# Compute free surface tendency -function compute_explicit_free_surface_tendency!(grid, model) - - arch = architecture(grid) - - args = tuple(model.velocities, - model.free_surface, - model.tracers, - model.auxiliary_fields, - model.forcing, - model.clock) - - launch!(arch, grid, :xy, - compute_hydrostatic_free_surface_Gη!, model.timestepper.Gⁿ.η, - grid, args) - - args = (model.clock, - fields(model), - model.closure, - model.buoyancy) - - apply_flux_bcs!(model.timestepper.Gⁿ.η, displacement(model.free_surface), arch, args) - - return nothing -end \ No newline at end of file diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl index 6d53944ebd..c76ac3b525 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl @@ -9,16 +9,18 @@ import Oceananigans.TimeSteppers: ab2_step! ##### Step everything ##### -function ab2_step!(model::HydrostaticFreeSurfaceModel, Δt) +setup_free_surface!(model, free_surface, χ) = nothing - compute_free_surface_tendency!(model.grid, model, model.free_surface) +function ab2_step!(model::HydrostaticFreeSurfaceModel, Δt) χ = model.timestepper.χ + setup_free_surface!(model, model.free_surface, χ) # Step locally velocity and tracers @apply_regionally local_ab2_step!(model, Δt, χ) - step_free_surface!(model.free_surface, model, model.timestepper, Δt) + # blocking step for implicit free surface, non blocking for explicit + ab2_step_free_surface!(model.free_surface, model, Δt, χ) return nothing end @@ -26,11 +28,9 @@ end function local_ab2_step!(model, Δt, χ) ab2_step_velocities!(model.velocities, model, Δt, χ) ab2_step_tracers!(model.tracers, model, Δt, χ) - - return nothing + return nothing end - ##### ##### Step velocities ##### @@ -45,6 +45,9 @@ function ab2_step_velocities!(velocities, model, Δt, χ) launch!(model.architecture, model.grid, :xyz, ab2_step_field!, velocity_field, Δt, χ, Gⁿ, G⁻) + # TODO: let next implicit solve depend on previous solve + explicit velocity step + # Need to distinguish between solver events and tendency calculation events. + # Note that BatchedTridiagonalSolver has a hard `wait`; this must be solved first. implicit_step!(velocity_field, model.timestepper.implicit_solver, model.closure, @@ -58,7 +61,7 @@ function ab2_step_velocities!(velocities, model, Δt, χ) end ##### -##### Step Tracers +##### Step velocities ##### const EmptyNamedTuple = NamedTuple{(),Tuple{}} diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl index 5cc82b5f37..92f787bfd2 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl @@ -1,6 +1,5 @@ using Oceananigans.Grids: Center, Face using Oceananigans.Fields: XFaceField, YFaceField, ZFaceField, TracerFields -using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper, SplitRungeKutta3TimeStepper function HydrostaticFreeSurfaceVelocityFields(::Nothing, grid, clock, bcs=NamedTuple()) u = XFaceField(grid, boundary_conditions=bcs.u) @@ -10,34 +9,9 @@ function HydrostaticFreeSurfaceVelocityFields(::Nothing, grid, clock, bcs=NamedT end function HydrostaticFreeSurfaceTendencyFields(velocities, free_surface, grid, tracer_names) - u = XFaceField(grid) - v = YFaceField(grid) - tracers = TracerFields(tracer_names, grid) - return merge((u=u, v=v), tracers) -end - -function HydrostaticFreeSurfaceTendencyFields(velocities, free_surface::ExplicitFreeSurface, grid, tracer_names) u = XFaceField(grid) v = YFaceField(grid) η = free_surface_displacement_field(velocities, free_surface, grid) tracers = TracerFields(tracer_names, grid) return merge((u=u, v=v, η=η), tracers) end - -function HydrostaticFreeSurfaceTendencyFields(velocities, free_surface::SplitExplicitFreeSurface, grid, tracer_names) - u = XFaceField(grid) - v = YFaceField(grid) - U = deepcopy(free_surface.barotropic_velocities.U) - V = deepcopy(free_surface.barotropic_velocities.V) - tracers = TracerFields(tracer_names, grid) - return merge((u=u, v=v, U=U, V=V), tracers) -end - -PreviousHydrostaticTendencyFields(::Val{:QuasiAdamsBashforth2}, args...) = HydrostaticFreeSurfaceTendencyFields(args...) -PreviousHydrostaticTendencyFields(::Val{:SplitRungeKutta3}, args...) = nothing - -function PreviousHydrostaticTendencyFields(::Val{:SplitRungeKutta3}, velocities, free_surface::SplitExplicitFreeSurface, args...) - U = deepcopy(free_surface.barotropic_velocities.U) - V = deepcopy(free_surface.barotropic_velocities.V) - return (; U=U, V=V) -end diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl index 270ca8b42b..6b92127487 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl @@ -3,7 +3,7 @@ using OrderedCollections: OrderedDict using Oceananigans.DistributedComputations using Oceananigans.Architectures: AbstractArchitecture -using Oceananigans.Advection: AbstractAdvectionScheme, Centered, VectorInvariant, adapt_advection_order +using Oceananigans.Advection: AbstractAdvectionScheme, CenteredSecondOrder, VectorInvariant, adapt_advection_order using Oceananigans.BuoyancyModels: validate_buoyancy, regularize_buoyancy, SeawaterBuoyancy, g_Earth using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions using Oceananigans.Biogeochemistry: validate_biogeochemistry, AbstractBiogeochemistry, biogeochemical_auxiliary_fields @@ -12,7 +12,7 @@ using Oceananigans.Forcings: model_forcing using Oceananigans.Grids: AbstractCurvilinearGrid, AbstractHorizontallyCurvilinearGrid, architecture, halo_size using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid using Oceananigans.Models: AbstractModel, validate_model_halo, NaNChecker, validate_tracer_advection, extract_boundary_conditions -using Oceananigans.TimeSteppers: Clock, TimeStepper, update_state!, AbstractLagrangianParticles, SplitRungeKutta3TimeStepper +using Oceananigans.TimeSteppers: Clock, TimeStepper, update_state!, AbstractLagrangianParticles using Oceananigans.TurbulenceClosures: validate_closure, with_tracers, DiffusivityFields, add_closure_specific_boundary_conditions using Oceananigans.TurbulenceClosures: time_discretization, implicit_diffusion_solver using Oceananigans.Utils: tupleit @@ -56,21 +56,22 @@ default_free_surface(grid; gravitational_acceleration=g_Earth) = """ HydrostaticFreeSurfaceModel(; grid, clock = Clock{eltype(grid)}(time = 0), - momentum_advection = VectorInvariant(), - tracer_advection = Centered(), - buoyancy = SeawaterBuoyancy(eltype(grid)), - coriolis = nothing, - free_surface = default_free_surface(grid, gravitational_acceleration=g_Earth), - forcing::NamedTuple = NamedTuple(), - closure = nothing, - boundary_conditions::NamedTuple = NamedTuple(), - tracers = (:T, :S), - particles::ParticlesOrNothing = nothing, - biogeochemistry::AbstractBGCOrNothing = nothing, - velocities = nothing, - pressure = nothing, - diffusivity_fields = nothing, - auxiliary_fields = NamedTuple()) + momentum_advection = VectorInvariant(), + tracer_advection = CenteredSecondOrder(), + buoyancy = SeawaterBuoyancy(eltype(grid)), + coriolis = nothing, + free_surface = default_free_surface(grid, gravitational_acceleration=g_Earth), + forcing::NamedTuple = NamedTuple(), + closure = nothing, + boundary_conditions::NamedTuple = NamedTuple(), + tracers = (:T, :S), + particles::ParticlesOrNothing = nothing, +biogeochemistry::AbstractBGCOrNothing = nothing, + velocities = nothing, + pressure = nothing, + diffusivity_fields = nothing, + auxiliary_fields = NamedTuple(), + ) Construct a hydrostatic model with a free surface on `grid`. @@ -102,43 +103,38 @@ Keyword arguments - `auxiliary_fields`: `NamedTuple` of auxiliary fields. Default: `nothing`. """ function HydrostaticFreeSurfaceModel(; grid, - clock = Clock{eltype(grid)}(time = 0), - momentum_advection = VectorInvariant(), - tracer_advection = Centered(), - buoyancy = nothing, - coriolis = nothing, - free_surface = default_free_surface(grid, gravitational_acceleration=g_Earth), - tracers = nothing, - forcing::NamedTuple = NamedTuple(), - closure = nothing, - timestepper = :QuasiAdamsBashforth2, - boundary_conditions::NamedTuple = NamedTuple(), - particles::ParticlesOrNothing = nothing, - biogeochemistry::AbstractBGCOrNothing = nothing, + clock = Clock{eltype(grid)}(time = 0), + momentum_advection = VectorInvariant(), + tracer_advection = CenteredSecondOrder(), + buoyancy = nothing, + coriolis = nothing, + free_surface = default_free_surface(grid, gravitational_acceleration=g_Earth), + tracers = nothing, + forcing::NamedTuple = NamedTuple(), + closure = nothing, + boundary_conditions::NamedTuple = NamedTuple(), + particles::ParticlesOrNothing = nothing, + biogeochemistry::AbstractBGCOrNothing = nothing, velocities = nothing, - pressure = nothing, - diffusivity_fields = nothing, - auxiliary_fields = NamedTuple()) + pressure = nothing, + diffusivity_fields = nothing, + auxiliary_fields = NamedTuple() + ) # Check halos and throw an error if the grid's halo is too small @apply_regionally validate_model_halo(grid, momentum_advection, tracer_advection, closure) - # Validate biogeochemistry (add biogeochemical tracers automagically) - tracers = tupleit(tracers) # supports tracers=:c keyword argument (for example) - biogeochemical_fields = merge(auxiliary_fields, biogeochemical_auxiliary_fields(biogeochemistry)) - tracers, auxiliary_fields = validate_biogeochemistry(tracers, biogeochemical_fields, biogeochemistry, grid, clock) + arch = architecture(grid) - # Reduce the advection order in directions that do not have enough grid points @apply_regionally momentum_advection = validate_momentum_advection(momentum_advection, grid) - default_tracer_advection, tracer_advection = validate_tracer_advection(tracer_advection, grid) - default_generator(name, tracer_advection) = default_tracer_advection - # Generate tracer advection scheme for each tracer - tracer_advection_tuple = with_tracers(tracernames(tracers), tracer_advection, default_generator, with_velocities=false) - momentum_advection_tuple = (; momentum = momentum_advection) - advection = merge(momentum_advection_tuple, tracer_advection_tuple) - advection = NamedTuple(name => adapt_advection_order(scheme, grid) for (name, scheme) in pairs(advection)) + tracers = tupleit(tracers) # supports tracers=:c keyword argument (for example) + + # Reduce the advection order in directions that do not have enough grid points + momentum_advection = adapt_advection_order(momentum_advection, grid) + tracer_advection = adapt_advection_order(tracer_advection, grid) + tracers, auxiliary_fields = validate_biogeochemistry(tracers, merge(auxiliary_fields, biogeochemical_auxiliary_fields(biogeochemistry)), biogeochemistry, grid, clock) validate_buoyancy(buoyancy, tracernames(tracers)) buoyancy = regularize_buoyancy(buoyancy) @@ -156,8 +152,7 @@ function HydrostaticFreeSurfaceModel(; grid, # Next, we form a list of default boundary conditions: prognostic_field_names = (:u, :v, :w, tracernames(tracers)..., :η, keys(auxiliary_fields)...) - default_boundary_conditions = NamedTuple{prognostic_field_names}(Tuple(FieldBoundaryConditions() - for name in prognostic_field_names)) + default_boundary_conditions = NamedTuple{prognostic_field_names}(Tuple(FieldBoundaryConditions() for name in prognostic_field_names)) # Then we merge specified, embedded, and default boundary conditions. Specified boundary conditions # have precedence, followed by embedded, followed by default. @@ -166,11 +161,7 @@ function HydrostaticFreeSurfaceModel(; grid, # Finally, we ensure that closure-specific boundary conditions, such as # those required by CATKEVerticalDiffusivity, are enforced: - boundary_conditions = add_closure_specific_boundary_conditions(closure, - boundary_conditions, - grid, - tracernames(tracers), - buoyancy) + boundary_conditions = add_closure_specific_boundary_conditions(closure, boundary_conditions, grid, tracernames(tracers), buoyancy) # Ensure `closure` describes all tracers closure = with_tracers(tracernames(tracers), closure) @@ -186,24 +177,30 @@ function HydrostaticFreeSurfaceModel(; grid, @apply_regionally validate_velocity_boundary_conditions(grid, velocities) - arch = architecture(grid) free_surface = validate_free_surface(arch, free_surface) free_surface = materialize_free_surface(free_surface, velocities, grid) # Instantiate timestepper if not already instantiated - implicit_solver = implicit_diffusion_solver(time_discretization(closure), grid) - prognostic_fields = hydrostatic_prognostic_fields(velocities, free_surface, tracers) - - timestepper = TimeStepper(timestepper, grid, tracernames(tracers); + implicit_solver = implicit_diffusion_solver(time_discretization(closure), grid) + timestepper = TimeStepper(:QuasiAdamsBashforth2, grid, tracernames(tracers); implicit_solver = implicit_solver, Gⁿ = HydrostaticFreeSurfaceTendencyFields(velocities, free_surface, grid, tracernames(tracers)), - G⁻ = PreviousHydrostaticTendencyFields(Val(timestepper), velocities, free_surface, grid, tracernames(tracers)), - Ψ⁻ = deepcopy(prognostic_fields)) + G⁻ = HydrostaticFreeSurfaceTendencyFields(velocities, free_surface, grid, tracernames(tracers))) # Regularize forcing for model tracer and velocity fields. - model_fields = merge(prognostic_fields, auxiliary_fields) + model_fields = merge(hydrostatic_prognostic_fields(velocities, free_surface, tracers), auxiliary_fields) forcing = model_forcing(model_fields; forcing...) - + + default_tracer_advection, tracer_advection = validate_tracer_advection(tracer_advection, grid) + + # Advection schemes + tracer_advection_tuple = with_tracers(tracernames(tracers), + tracer_advection, + (name, tracer_advection) -> default_tracer_advection, + with_velocities=false) + + advection = merge((momentum=momentum_advection,), tracer_advection_tuple) + model = HydrostaticFreeSurfaceModel(arch, grid, clock, advection, buoyancy, coriolis, free_surface, forcing, closure, particles, biogeochemistry, velocities, tracers, pressure, diffusivity_fields, timestepper, auxiliary_fields) @@ -234,6 +231,7 @@ validate_momentum_advection(momentum_advection::VectorInvariant, grid::Orthogona validate_momentum_advection(momentum_advection, grid::OrthogonalSphericalShellGrid) = error("$(typeof(momentum_advection)) is not supported with $(typeof(grid))") initialize!(model::HydrostaticFreeSurfaceModel) = initialize_free_surface!(model.free_surface, model.grid, model.velocities) +initialize_free_surface!(free_surface, grid, velocities) = nothing # return the total advective velocities @inline total_velocities(model::HydrostaticFreeSurfaceModel) = model.velocities diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk3_step.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk3_step.jl deleted file mode 100644 index 0a02235f89..0000000000 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk3_step.jl +++ /dev/null @@ -1,122 +0,0 @@ -using Oceananigans.Fields: location -using Oceananigans.TurbulenceClosures: implicit_step! -using Oceananigans.ImmersedBoundaries: retrieve_interior_active_cells_map, retrieve_surface_active_cells_map - -import Oceananigans.TimeSteppers: split_rk3_substep!, _split_rk3_substep_field! - -function split_rk3_substep!(model::HydrostaticFreeSurfaceModel, Δt, γⁿ, ζⁿ) - - grid = model.grid - timestepper = model.timestepper - free_surface = model.free_surface - - compute_free_surface_tendency!(grid, model, free_surface) - - rk3_substep_velocities!(model.velocities, model, Δt, γⁿ, ζⁿ) - rk3_substep_tracers!(model.tracers, model, Δt, γⁿ, ζⁿ) - - # Full step for Implicit and Split-Explicit, substep for Explicit - step_free_surface!(free_surface, model, timestepper, Δt) - - # Average free surface variables - # in the second stage - if model.clock.stage == 2 - rk3_average_free_surface!(free_surface, grid, timestepper, γⁿ, ζⁿ) - end - - return nothing -end - -rk3_average_free_surface!(free_surface, args...) = nothing - -function rk3_average_free_surface!(free_surface::ImplicitFreeSurface, grid, timestepper, γⁿ, ζⁿ) - arch = architecture(grid) - - ηⁿ⁻¹ = timestepper.Ψ⁻.η - ηⁿ = free_surface.η - - launch!(arch, grid, :xy, _rk3_average_free_surface!, ηⁿ, grid, ηⁿ⁻¹, γⁿ, ζⁿ) - - return nothing -end - -function rk3_average_free_surface!(free_surface::SplitExplicitFreeSurface, grid, timestepper, γⁿ, ζⁿ) - - arch = architecture(grid) - - Uⁿ⁻¹ = timestepper.Ψ⁻.U - Vⁿ⁻¹ = timestepper.Ψ⁻.V - Uⁿ = free_surface.barotropic_velocities.U - Vⁿ = free_surface.barotropic_velocities.V - - launch!(arch, grid, :xy, _rk3_average_free_surface!, Uⁿ, grid, Uⁿ⁻¹, γⁿ, ζⁿ) - launch!(arch, grid, :xy, _rk3_average_free_surface!, Vⁿ, grid, Vⁿ⁻¹, γⁿ, ζⁿ) - - return nothing -end - -@kernel function _rk3_average_free_surface!(η, grid, η⁻, γⁿ, ζⁿ) - i, j = @index(Global, NTuple) - k = grid.Nz + 1 - @inbounds η[i, j, k] = ζⁿ * η⁻[i, j, k] + γⁿ * η[i, j, k] -end - -##### -##### Time stepping in each substep -##### - -function rk3_substep_velocities!(velocities, model, Δt, γⁿ, ζⁿ) - - for name in (:u, :v) - Gⁿ = model.timestepper.Gⁿ[name] - Ψ⁻ = model.timestepper.Ψ⁻[name] - velocity_field = velocities[name] - - launch!(model.architecture, model.grid, :xyz, - _split_rk3_substep_field!, velocity_field, Δt, γⁿ, ζⁿ, Gⁿ, Ψ⁻) - - implicit_step!(velocity_field, - model.timestepper.implicit_solver, - model.closure, - model.diffusivity_fields, - nothing, - model.clock, - Δt) - end - - return nothing -end - -##### -##### Step Tracers -##### - -rk3_substep_tracers!(::EmptyNamedTuple, model, Δt, γⁿ, ζⁿ) = nothing - -function rk3_substep_tracers!(tracers, model, Δt, γⁿ, ζⁿ) - - closure = model.closure - grid = model.grid - - # Tracer update kernels - for (tracer_index, tracer_name) in enumerate(propertynames(tracers)) - - Gⁿ = model.timestepper.Gⁿ[tracer_name] - Ψ⁻ = model.timestepper.Ψ⁻[tracer_name] - tracer_field = tracers[tracer_name] - closure = model.closure - - launch!(architecture(grid), grid, :xyz, - _split_rk3_substep_field!, tracer_field, Δt, γⁿ, ζⁿ, Gⁿ, Ψ⁻) - - implicit_step!(tracer_field, - model.timestepper.implicit_solver, - closure, - model.diffusivity_fields, - Val(tracer_index), - model.clock, - Δt) - end - - return nothing -end \ No newline at end of file diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl index 605849dbee..6de795e2b6 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl @@ -38,8 +38,8 @@ implicitly during time-stepping. diffusivities, hydrostatic_pressure_anomaly, auxiliary_fields, - clock, - forcing) + forcings, + clock) model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields) @@ -49,7 +49,7 @@ implicitly during time-stepping. - ∂xᶠᶜᶜ(i, j, k, grid, hydrostatic_pressure_anomaly) - ∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy) - immersed_∂ⱼ_τ₁ⱼ(i, j, k, grid, velocities, u_immersed_bc, closure, diffusivities, clock, model_fields) - + forcing(i, j, k, grid, clock, hydrostatic_prognostic_fields(velocities, free_surface, tracers))) + + forcings.u(i, j, k, grid, clock, hydrostatic_prognostic_fields(velocities, free_surface, tracers))) end """ @@ -77,8 +77,8 @@ implicitly during time-stepping. diffusivities, hydrostatic_pressure_anomaly, auxiliary_fields, - clock, - forcing) + forcings, + clock) model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields) @@ -88,7 +88,7 @@ implicitly during time-stepping. - ∂yᶜᶠᶜ(i, j, k, grid, hydrostatic_pressure_anomaly) - ∂ⱼ_τ₂ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy) - immersed_∂ⱼ_τ₂ⱼ(i, j, k, grid, velocities, v_immersed_bc, closure, diffusivities, clock, model_fields) - + forcing(i, j, k, grid, clock, model_fields)) + + forcings.v(i, j, k, grid, clock, model_fields)) end """ @@ -116,8 +116,8 @@ where `c = C[tracer_index]`. tracers, diffusivities, auxiliary_fields, - clock, - forcing) where tracer_index + forcing, + clock) where tracer_index @inbounds c = tracers[tracer_index] model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields) @@ -135,4 +135,29 @@ where `c = C[tracer_index]`. - immersed_∇_dot_qᶜ(i, j, k, grid, c, c_immersed_bc, closure, diffusivities, val_tracer_index, clock, model_fields) + biogeochemical_transition(i, j, k, grid, biogeochemistry, val_tracer_name, clock, model_fields) + forcing(i, j, k, grid, clock, model_fields)) -end \ No newline at end of file +end + +""" +Return the tendency for an explicit free surface at horizontal grid point `i, j`. + +The tendency is called ``G_η`` and defined via + +``` +∂_t η = G_η +``` +""" +@inline function free_surface_tendency(i, j, grid, + velocities, + free_surface, + tracers, + auxiliary_fields, + forcings, + clock) + + k_top = grid.Nz + 1 + model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields) + + return @inbounds ( velocities.w[i, j, k_top] + + forcings.η(i, j, k_top, grid, clock, model_fields)) +end + diff --git a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl index 7fab830011..e54ad5c49c 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl @@ -123,7 +123,10 @@ build_implicit_step_solver(::Val{:Default}, grid, settings, gravitational_accele """ Implicitly step forward η. """ -function step_free_surface!(free_surface::ImplicitFreeSurface, model, timestepper, Δt) +ab2_step_free_surface!(free_surface::ImplicitFreeSurface, model, Δt, χ) = + implicit_free_surface_step!(free_surface::ImplicitFreeSurface, model, Δt, χ) + +function implicit_free_surface_step!(free_surface::ImplicitFreeSurface, model, Δt, χ) η = free_surface.η g = free_surface.gravitational_acceleration rhs = free_surface.implicit_step_solver.right_hand_side diff --git a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl index a78b837017..94403dc62c 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl @@ -74,7 +74,7 @@ end function HydrostaticFreeSurfaceTendencyFields(::PrescribedVelocityFields, free_surface, grid, tracer_names) tracer_tendencies = TracerFields(tracer_names, grid) - momentum_tendencies = (u = nothing, v = nothing) + momentum_tendencies = (u = nothing, v = nothing, η = nothing) return merge(momentum_tendencies, tracer_tendencies) end @@ -89,8 +89,7 @@ end @inline datatuple(obj::PrescribedVelocityFields) = (; u = datatuple(obj.u), v = datatuple(obj.v), w = datatuple(obj.w)) ab2_step_velocities!(::PrescribedVelocityFields, args...) = nothing -rk3_substep_velocities!(::PrescribedVelocityFields, args...) = nothing -step_free_surface!(::Nothing, model, timestepper, Δt) = nothing +ab2_step_free_surface!(::Nothing, model, Δt, χ) = nothing compute_w_from_continuity!(::PrescribedVelocityFields, args...; kwargs...) = nothing validate_velocity_boundary_conditions(grid, ::PrescribedVelocityFields) = nothing diff --git a/src/Models/HydrostaticFreeSurfaceModels/nothing_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/rigid_lid.jl similarity index 100% rename from src/Models/HydrostaticFreeSurfaceModels/nothing_free_surface.jl rename to src/Models/HydrostaticFreeSurfaceModels/rigid_lid.jl diff --git a/src/Models/HydrostaticFreeSurfaceModels/show_hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/show_hydrostatic_free_surface_model.jl index bf461a43f1..88cd4ab915 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/show_hydrostatic_free_surface_model.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/show_hydrostatic_free_surface_model.jl @@ -27,7 +27,7 @@ function Base.show(io::IO, model::HydrostaticFreeSurfaceModel) end if typeof(model.free_surface).name.wrapper == SplitExplicitFreeSurface - print(io, "│ └── substepping: $(summary(model.free_surface.substepping))", "\n") + print(io, "│ └── substepping: $(summary(model.free_surface.settings.substepping))", "\n") end end diff --git a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl index dee6e4d278..d4d3dec8d8 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl @@ -43,7 +43,7 @@ end validate_velocity_boundary_conditions(::SingleColumnGrid, velocities) = nothing validate_velocity_boundary_conditions(::SingleColumnGrid, ::PrescribedVelocityFields) = nothing validate_momentum_advection(momentum_advection, ::SingleColumnGrid) = nothing -validate_tracer_advection(tracer_advection_tuple::NamedTuple, ::SingleColumnGrid) = Centered(), tracer_advection_tuple +validate_tracer_advection(tracer_advection_tuple::NamedTuple, ::SingleColumnGrid) = CenteredSecondOrder(), tracer_advection_tuple validate_tracer_advection(tracer_advection::AbstractAdvectionScheme, ::SingleColumnGrid) = tracer_advection, NamedTuple() compute_w_from_continuity!(velocities, arch, ::SingleColumnGrid; kwargs...) = nothing @@ -53,9 +53,11 @@ compute_w_from_continuity!(::PrescribedVelocityFields, arch, ::SingleColumnGrid; ##### Time-step optimizations ##### +compute_free_surface_tendency!(::SingleColumnGrid, args...) = nothing + # Disambiguation -compute_free_surface_tendency!(::SingleColumnGrid, model, ::ExplicitFreeSurface) = nothing -compute_free_surface_tendency!(::SingleColumnGrid, model, ::SplitExplicitFreeSurface) = nothing +compute_free_surface_tendency!(::SingleColumnGrid, ::ImplicitFreeSurfaceHFSM , args...) = nothing +compute_free_surface_tendency!(::SingleColumnGrid, ::SplitExplicitFreeSurfaceHFSM, args...) = nothing # Fast state update and halo filling diff --git a/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface.jl new file mode 100644 index 0000000000..5d9a00f2ed --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface.jl @@ -0,0 +1,424 @@ +using Oceananigans +using Oceananigans.Architectures +using Oceananigans.Fields +using Oceananigans.Grids +using Oceananigans.Grids: AbstractGrid +using Oceananigans.AbstractOperations: Δz, GridMetricOperation + +using Adapt +using Base +using KernelAbstractions: @index, @kernel + +import Oceananigans.TimeSteppers: reset! + +""" + struct SplitExplicitFreeSurface + +The split-explicit free surface solver. + +$(FIELDS) +""" +struct SplitExplicitFreeSurface{𝒩, 𝒮, ℱ, 𝒫 ,ℰ} <: AbstractFreeSurface{𝒩, 𝒫} + "The instantaneous free surface (`ReducedField`)" + η :: 𝒩 + "The entire state for the split-explicit solver (`SplitExplicitState`)" + state :: 𝒮 + "Parameters for timestepping split-explicit solver (`NamedTuple`)" + auxiliary :: ℱ + "Gravitational acceleration" + gravitational_acceleration :: 𝒫 + "Settings for the split-explicit scheme" + settings :: ℰ +end + +""" + SplitExplicitFreeSurface(grid = nothing; + gravitational_acceleration = g_Earth, + substeps = nothing, + cfl = nothing, + fixed_Δt = nothing, + averaging_kernel = averaging_shape_function, + timestepper = ForwardBackwardScheme()) + +Return a `SplitExplicitFreeSurface` representing an explicit time discretization +of a free surface dynamics with `gravitational_acceleration`. + +Keyword Arguments +================= + +- `gravitational_acceleration`: the gravitational acceleration (default: `g_Earth`) + +- `substeps`: The number of substeps that divide the range `(t, t + 2Δt)`, where `Δt` is the baroclinic + timestep. Note that some averaging functions do not require substepping until `2Δt`. + The number of substeps is reduced automatically to the last index of `averaging_kernel` + for which `averaging_kernel > 0`. + +- `cfl`: If set then the number of `substeps` are computed based on the advective timescale imposed from + the barotropic gravity-wave speed that corresponds to depth `grid.Lz`. If `fixed_Δt` is provided, + then the number of `substeps` adapts to maintain an exact `cfl`. If not, the effective cfl will + always be lower than the specified `cfl` provided that the baroclinic time step satisfies + `Δt_baroclinic < fixed_Δt`. + +!!! info "Needed keyword arguments" + Either `substeps` _or_ `cfl` need to be prescribed. + + When `cfl` is prescribed then `grid` is also required as a positional argument. + +- `fixed_Δt`: The maximum baroclinic timestep allowed. If `fixed_Δt` is a `nothing` and a cfl is provided, + then the number of substeps will be computed on the fly from the baroclinic time step to + maintain a constant cfl. + +- `averaging_kernel`: A function of `τ` used to average the barotropic transport `U` and the free surface + `η` within the barotropic advancement. `τ` is the fractional substep going from 0 to 2 + with the baroclinic time step `t + Δt` located at `τ = 1`. The `averaging_kernel` + function should be centered at `τ = 1`, that is, ``∑ (aₘ m / M) = 1``, where the + the summation occurs for ``m = 1, ..., M_*``. Here, ``m = 0`` and ``m = M`` correspond + to the two consecutive baroclinic timesteps between which the barotropic timestepping + occurs and ``M_*`` corresponds to the last barotropic time step for which the + `averaging_kernel > 0`. By default, the averaging kernel described by [Shchepetkin2005](@citet) + is used. + +- `timestepper`: Time stepping scheme used for the barotropic advancement. Choose one of: + * `ForwardBackwardScheme()` (default): `η = f(U)` then `U = f(η)`, + * `AdamsBashforth3Scheme()`: `η = f(U, Uᵐ⁻¹, Uᵐ⁻²)` then `U = f(η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻²)`. + +References +========== + +Shchepetkin, A. F., & McWilliams, J. C. (2005). The regional oceanic modeling system (ROMS): a split-explicit, free-surface, topography-following-coordinate oceanic model. Ocean Modelling, 9(4), 347-404. +""" +function SplitExplicitFreeSurface(grid = nothing; + gravitational_acceleration = g_Earth, + substeps = nothing, + cfl = nothing, + fixed_Δt = nothing, + averaging_kernel = averaging_shape_function, + timestepper = ForwardBackwardScheme()) + + settings = SplitExplicitSettings(grid; + gravitational_acceleration, + substeps, + cfl, + fixed_Δt, + averaging_kernel, + timestepper) + + return SplitExplicitFreeSurface(nothing, + nothing, + nothing, + gravitational_acceleration, + settings) +end + +# Internal function for HydrostaticFreeSurfaceModel +function materialize_free_surface(free_surface::SplitExplicitFreeSurface, velocities, grid) + settings = SplitExplicitSettings(grid; free_surface.settings.settings_kwargs...) + + η = free_surface_displacement_field(velocities, free_surface, grid) + + gravitational_acceleration = convert(eltype(grid), free_surface.gravitational_acceleration) + + return SplitExplicitFreeSurface(η, + SplitExplicitState(grid, settings.timestepper), + SplitExplicitAuxiliaryFields(grid), + gravitational_acceleration, + settings) +end + + +""" + struct SplitExplicitState + +A type containing the state fields for the split-explicit free surface. + +$(FIELDS) +""" +Base.@kwdef struct SplitExplicitState{CC, ACC, FC, AFC, CF, ACF} + "The free surface at time `m`. (`ReducedField` over ``z``)" + ηᵐ :: ACC + "The free surface at time `m-1`. (`ReducedField` over ``z``)" + ηᵐ⁻¹ :: ACC + "The free surface at time `m-2`. (`ReducedField` over ``z``)" + ηᵐ⁻² :: ACC + "The barotropic zonal velocity at time `m`. (`ReducedField` over ``z``)" + U :: FC + "The barotropic zonal velocity at time `m-1`. (`ReducedField` over ``z``)" + Uᵐ⁻¹ :: AFC + "The barotropic zonal velocity at time `m-2`. (`ReducedField` over ``z``)" + Uᵐ⁻² :: AFC + "The barotropic meridional velocity at time `m`. (`ReducedField` over ``z``)" + V :: CF + "The barotropic meridional velocity at time `m-1`. (`ReducedField` over ``z``)" + Vᵐ⁻¹ :: ACF + "The barotropic meridional velocity at time `m-2`. (`ReducedField` over ``z``)" + Vᵐ⁻² :: ACF + "The time-filtered free surface. (`ReducedField` over ``z``)" + η̅ :: CC + "The time-filtered barotropic zonal velocity. (`ReducedField` over ``z``)" + U̅ :: FC + "The time-filtered barotropic meridional velocity. (`ReducedField` over ``z``)" + V̅ :: CF +end + +""" + SplitExplicitState(grid, timestepper) + +Return the split-explicit state for `grid`. + +Note that `η̅` is solely used for setting the `η` at the next substep iteration -- it essentially +acts as a filter for `η`. Values with superscripts `m-1` and `m-2` correspond to previous stored +time steps to allow using a higher-order time stepping scheme, e.g., `AdamsBashforth3Scheme`. +""" +function SplitExplicitState(grid::AbstractGrid, timestepper) + + Nz = size(grid, 3) + + η̅ = ZFaceField(grid, indices = (:, :, Nz+1)) + + ηᵐ = auxiliary_free_surface_field(grid, timestepper) + ηᵐ⁻¹ = auxiliary_free_surface_field(grid, timestepper) + ηᵐ⁻² = auxiliary_free_surface_field(grid, timestepper) + + U = XFaceField(grid, indices = (:, :, Nz)) + V = YFaceField(grid, indices = (:, :, Nz)) + + Uᵐ⁻¹ = auxiliary_barotropic_U_field(grid, timestepper) + Vᵐ⁻¹ = auxiliary_barotropic_V_field(grid, timestepper) + Uᵐ⁻² = auxiliary_barotropic_U_field(grid, timestepper) + Vᵐ⁻² = auxiliary_barotropic_V_field(grid, timestepper) + + U̅ = XFaceField(grid, indices = (:, :, Nz)) + V̅ = YFaceField(grid, indices = (:, :, Nz)) + + return SplitExplicitState(; ηᵐ, ηᵐ⁻¹, ηᵐ⁻², U, Uᵐ⁻¹, Uᵐ⁻², V, Vᵐ⁻¹, Vᵐ⁻², η̅, U̅, V̅) +end + +""" + struct SplitExplicitAuxiliaryFields + +A type containing auxiliary fields for the split-explicit free surface. + +The barotropic time stepping is launched on a grid `(kernel_size[1], kernel_size[2])` +large (or `:xy` in case of a serial computation), and start computing from +`(i - kernel_offsets[1], j - kernel_offsets[2])`. + +$(FIELDS) +""" +Base.@kwdef struct SplitExplicitAuxiliaryFields{𝒞ℱ, ℱ𝒞, 𝒦} + "Vertically-integrated slow barotropic forcing function for `U` (`ReducedField` over ``z``)" + Gᵁ :: ℱ𝒞 + "Vertically-integrated slow barotropic forcing function for `V` (`ReducedField` over ``z``)" + Gⱽ :: 𝒞ℱ + "Depth at `(Face, Center)` (`ReducedField` over ``z``)" + Hᶠᶜ :: ℱ𝒞 + "Depth at `(Center, Face)` (`ReducedField` over ``z``)" + Hᶜᶠ :: 𝒞ℱ + "kernel size for barotropic time stepping" + kernel_parameters :: 𝒦 +end + +""" + SplitExplicitAuxiliaryFields(grid) + +Return the `SplitExplicitAuxiliaryFields` for `grid`. +""" +function SplitExplicitAuxiliaryFields(grid::AbstractGrid) + + Gᵁ = Field((Face, Center, Nothing), grid) + Gⱽ = Field((Center, Face, Nothing), grid) + + Hᶠᶜ = Field((Face, Center, Nothing), grid) + Hᶜᶠ = Field((Center, Face, Nothing), grid) + + dz = GridMetricOperation((Face, Center, Center), Δz, grid) + sum!(Hᶠᶜ, dz) + + dz = GridMetricOperation((Center, Face, Center), Δz, grid) + sum!(Hᶜᶠ, dz) + + fill_halo_regions!((Hᶠᶜ, Hᶜᶠ)) + + kernel_parameters = :xy + + return SplitExplicitAuxiliaryFields(Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ, kernel_parameters) +end + +""" + struct SplitExplicitSettings + +A type containing settings for the split-explicit free surface. + +$(FIELDS) +""" +struct SplitExplicitSettings{𝒩, 𝒮} + substepping :: 𝒩 # Either `FixedSubstepNumber` or `FixedTimeStepSize`" + timestepper :: 𝒮 # time-stepping scheme + settings_kwargs :: NamedTuple # kwargs to reproduce current settings +end + +struct AdamsBashforth3Scheme end +struct ForwardBackwardScheme end + + +auxiliary_free_surface_field(grid, ::AdamsBashforth3Scheme) = ZFaceField(grid, indices = (:, :, size(grid, 3)+1)) +auxiliary_free_surface_field(grid, ::ForwardBackwardScheme) = nothing + +auxiliary_barotropic_U_field(grid, ::AdamsBashforth3Scheme) = XFaceField(grid, indices = (:, :, size(grid, 3))) +auxiliary_barotropic_U_field(grid, ::ForwardBackwardScheme) = nothing +auxiliary_barotropic_V_field(grid, ::AdamsBashforth3Scheme) = YFaceField(grid, indices = (:, :, size(grid, 3))) +auxiliary_barotropic_V_field(grid, ::ForwardBackwardScheme) = nothing + +# (p = 2, q = 4, r = 0.18927) minimize dispersion error from Shchepetkin and McWilliams (2005): https://doi.org/10.1016/j.ocemod.2004.08.002 +@inline function averaging_shape_function(τ::FT; p = 2, q = 4, r = FT(0.18927)) where FT + τ₀ = (p + 2) * (p + q + 2) / (p + 1) / (p + q + 1) + + return (τ / τ₀)^p * (1 - (τ / τ₀)^q) - r * (τ / τ₀) +end + +@inline cosine_averaging_kernel(τ::FT) where FT = τ ≥ 0.5 && τ ≤ 1.5 ? convert(FT, 1 + cos(2π * (τ - 1))) : zero(FT) +@inline constant_averaging_kernel(τ::FT) where FT = convert(FT, 1) + +""" An internal type for the `SplitExplicitFreeSurface` that allows substepping with +a fixed `Δt_barotropic` based on a CFL condition """ +struct FixedTimeStepSize{B, F} + Δt_barotropic :: B + averaging_kernel :: F +end + +""" An internal type for the `SplitExplicitFreeSurface` that allows substepping with +a fixed number of substeps with time step size of `fractional_step_size * Δt_baroclinic` """ +struct FixedSubstepNumber{B, F} + fractional_step_size :: B + averaging_weights :: F +end + +function FixedTimeStepSize(grid; + cfl = 0.7, + averaging_kernel = averaging_shape_function, + gravitational_acceleration = g_Earth) + + FT = eltype(grid) + + Δx⁻² = topology(grid)[1] == Flat ? 0 : 1 / minimum_xspacing(grid)^2 + Δy⁻² = topology(grid)[2] == Flat ? 0 : 1 / minimum_yspacing(grid)^2 + Δs = sqrt(1 / (Δx⁻² + Δy⁻²)) + + wave_speed = sqrt(gravitational_acceleration * grid.Lz) + + Δt_barotropic = convert(FT, cfl * Δs / wave_speed) + + return FixedTimeStepSize(Δt_barotropic, averaging_kernel) +end + +@inline function weights_from_substeps(FT, substeps, averaging_kernel) + + τᶠ = range(FT(0), FT(2), length = substeps+1) + Δτ = τᶠ[2] - τᶠ[1] + + averaging_weights = map(averaging_kernel, τᶠ[2:end]) + idx = searchsortedlast(averaging_weights, 0, rev=true) + substeps = idx + + averaging_weights = averaging_weights[1:idx] + averaging_weights ./= sum(averaging_weights) + + return Δτ, tuple(averaging_weights...) +end + +function SplitExplicitSettings(grid = nothing; + gravitational_acceleration = g_Earth, + substeps = nothing, + cfl = nothing, + fixed_Δt = nothing, + averaging_kernel = averaging_shape_function, + timestepper = ForwardBackwardScheme()) + + settings_kwargs = (; gravitational_acceleration, + substeps, + cfl, + fixed_Δt, + averaging_kernel, + timestepper) + + if !isnothing(grid) + FT = eltype(grid) + else + # this is a fallback and only used via the outer constructor, + # in case no grid is provided; when afterwards the free surfade + # is materialized via materialize_free_surface + # FT becomes eltype(grid) + FT = Float64 + end + + if (!isnothing(substeps) && !isnothing(cfl)) || (isnothing(substeps) && isnothing(cfl)) + throw(ArgumentError("either specify a cfl or a number of substeps")) + end + + if !isnothing(cfl) + if isnothing(grid) + throw(ArgumentError(string("Need to provide the grid to calculate the barotropic substeps from the cfl. ", + "For example, SplitExplicitFreeSurface(grid, cfl=0.7, ...)"))) + end + substepping = FixedTimeStepSize(grid; cfl, gravitational_acceleration, averaging_kernel) + if isnothing(fixed_Δt) + return SplitExplicitSettings(substepping, timestepper, settings_kwargs) + else + substeps = ceil(Int, 2 * fixed_Δt / substepping.Δt_barotropic) + end + end + + fractional_step_size, averaging_weights = weights_from_substeps(FT, substeps, averaging_kernel) + substepping = FixedSubstepNumber(fractional_step_size, averaging_weights) + + return SplitExplicitSettings(substepping, timestepper, settings_kwargs) +end + +# Convenience Functions for grabbing free surface +free_surface(free_surface::SplitExplicitFreeSurface) = free_surface.η + +# extend +@inline explicit_barotropic_pressure_x_gradient(i, j, k, grid, ::SplitExplicitFreeSurface) = zero(grid) +@inline explicit_barotropic_pressure_y_gradient(i, j, k, grid, ::SplitExplicitFreeSurface) = zero(grid) + +# convenience functor +(sefs::SplitExplicitFreeSurface)(settings::SplitExplicitSettings) = + SplitExplicitFreeSurface(sefs.η, sefs.state, sefs.auxiliary, sefs.gravitational_acceleration, settings) + +Base.summary(s::FixedTimeStepSize) = string("Barotropic time step equal to $(prettytime(s.Δt_barotropic))") +Base.summary(s::FixedSubstepNumber) = string("Barotropic fractional step equal to $(s.fractional_step_size) times the baroclinic step") + +Base.summary(sefs::SplitExplicitFreeSurface) = string("SplitExplicitFreeSurface with $(summary(sefs.settings.substepping))") + +Base.show(io::IO, sefs::SplitExplicitFreeSurface) = print(io, "$(summary(sefs))\n") + +function reset!(sefs::SplitExplicitFreeSurface) + for name in propertynames(sefs.state) + var = getproperty(sefs.state, name) + fill!(var, 0) + end + + fill!(sefs.auxiliary.Gᵁ, 0) + fill!(sefs.auxiliary.Gⱽ, 0) + + return nothing +end + +# Adapt +Adapt.adapt_structure(to, free_surface::SplitExplicitFreeSurface) = + SplitExplicitFreeSurface(Adapt.adapt(to, free_surface.η), nothing, nothing, + free_surface.gravitational_acceleration, nothing) + +for Type in (:SplitExplicitFreeSurface, + :SplitExplicitSettings, + :SplitExplicitState, + :SplitExplicitAuxiliaryFields, + :FixedTimeStepSize, + :FixedSubstepNumber) + + @eval begin + function on_architecture(to, fs::$Type) + args = Tuple(on_architecture(to, prop) for prop in propertynames(fs)) + return $Type(args...) + end + end +end diff --git a/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl b/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl new file mode 100644 index 0000000000..91ffed889f --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl @@ -0,0 +1,482 @@ +using Oceananigans.Grids: topology +using Oceananigans.Utils +using Oceananigans.AbstractOperations: Δz +using Oceananigans.BoundaryConditions +using Oceananigans.Operators +using Oceananigans.Architectures: convert_args +using Oceananigans.ImmersedBoundaries: peripheral_node, immersed_inactive_node, GFBIBG +using Oceananigans.ImmersedBoundaries: inactive_node, IBG, c, f +using Oceananigans.ImmersedBoundaries: mask_immersed_field!, retrieve_surface_active_cells_map, retrieve_interior_active_cells_map +using Oceananigans.ImmersedBoundaries: active_linear_index_to_tuple, ActiveCellsIBG, ActiveZColumnsIBG +using Oceananigans.DistributedComputations: child_architecture +using Oceananigans.DistributedComputations: Distributed + +using Printf +using KernelAbstractions: @index, @kernel +using KernelAbstractions.Extras.LoopInfo: @unroll + +# constants for AB3 time stepping scheme (from https://doi.org/10.1016/j.ocemod.2004.08.002) +const β = 0.281105 +const α = 1.5 + β +const θ = - 0.5 - 2β +const γ = 0.088 +const δ = 0.614 +const ϵ = 0.013 +const μ = 1 - δ - γ - ϵ + +# Evolution Kernels +# +# ∂t(η) = -∇⋅U +# ∂t(U) = - gH∇η + f +# +# the free surface field η and its average η̄ are located on `Face`s at the surface (grid.Nz +1). All other intermediate variables +# (U, V, Ū, V̄) are barotropic fields (`ReducedField`) for which a k index is not defined + +# Operators specific to the advancement of the Free surface and the Barotropic velocity. In particular, the base operators follow +# these rules: +# +# `δxᶠᵃᵃ_η` : Hardcodes Noflux or Periodic boundary conditions for the free surface η in x direction +# `δyᵃᶠᵃ_η` : Hardcodes Noflux or Periodic boundary conditions for the free surface η in y direction +# +# `δxᶜᵃᵃ_U` : Hardcodes NoPenetration or Periodic boundary conditions for the zonal barotropic velocity U in x direction +# `δyᵃᶜᵃ_V` : Hardcodes NoPenetration or Periodic boundary conditions for the meridional barotropic velocity V in y direction +# +# The functions `η★` `U★` and `V★` represent the value of free surface, barotropic zonal and meridional velocity at time step m+1/2 +@inline δxᶠᵃᵃ_η(i, j, k, grid, T, η★::Function, args...) = δxᶠᵃᵃ(i, j, k, grid, η★, args...) +@inline δyᵃᶠᵃ_η(i, j, k, grid, T, η★::Function, args...) = δyᵃᶠᵃ(i, j, k, grid, η★, args...) +@inline δxᶜᵃᵃ_U(i, j, k, grid, T, U★::Function, args...) = δxᶜᵃᵃ(i, j, k, grid, U★, args...) +@inline δyᵃᶜᵃ_V(i, j, k, grid, T, V★::Function, args...) = δyᵃᶜᵃ(i, j, k, grid, V★, args...) + +@inline δxᶠᵃᵃ_η(i, j, k, grid, ::Type{Periodic}, η★::Function, args...) = ifelse(i == 1, η★(1, j, k, grid, args...) - η★(grid.Nx, j, k, grid, args...), δxᶠᵃᵃ(i, j, k, grid, η★, args...)) +@inline δyᵃᶠᵃ_η(i, j, k, grid, ::Type{Periodic}, η★::Function, args...) = ifelse(j == 1, η★(i, 1, k, grid, args...) - η★(i, grid.Ny, k, grid, args...), δyᵃᶠᵃ(i, j, k, grid, η★, args...)) + +@inline δxᶜᵃᵃ_U(i, j, k, grid, ::Type{Periodic}, U★::Function, args...) = ifelse(i == grid.Nx, U★(1, j, k, grid, args...) - U★(grid.Nx, j, k, grid, args...), δxᶜᵃᵃ(i, j, k, grid, U★, args...)) +@inline δyᵃᶜᵃ_V(i, j, k, grid, ::Type{Periodic}, V★::Function, args...) = ifelse(j == grid.Ny, V★(i, 1, k, grid, args...) - V★(i, grid.Ny, k, grid, args...), δyᵃᶜᵃ(i, j, k, grid, V★, args...)) + +# Enforce NoFlux conditions for `η★` + +@inline δxᶠᵃᵃ_η(i, j, k, grid, ::Type{Bounded}, η★::Function, args...) = ifelse(i == 1, zero(grid), δxᶠᵃᵃ(i, j, k, grid, η★, args...)) +@inline δyᵃᶠᵃ_η(i, j, k, grid, ::Type{Bounded}, η★::Function, args...) = ifelse(j == 1, zero(grid), δyᵃᶠᵃ(i, j, k, grid, η★, args...)) +@inline δxᶠᵃᵃ_η(i, j, k, grid, ::Type{RightConnected}, η★::Function, args...) = ifelse(i == 1, zero(grid), δxᶠᵃᵃ(i, j, k, grid, η★, args...)) +@inline δyᵃᶠᵃ_η(i, j, k, grid, ::Type{RightConnected}, η★::Function, args...) = ifelse(j == 1, zero(grid), δyᵃᶠᵃ(i, j, k, grid, η★, args...)) + +# Enforce Impenetrability conditions for `U★` and `V★` + +@inline δxᶜᵃᵃ_U(i, j, k, grid, ::Type{Bounded}, U★::Function, args...) = ifelse(i == grid.Nx, - U★(i, j, k, grid, args...), + ifelse(i == 1, U★(2, j, k, grid, args...), δxᶜᵃᵃ(i, j, k, grid, U★, args...))) +@inline δyᵃᶜᵃ_V(i, j, k, grid, ::Type{Bounded}, V★::Function, args...) = ifelse(j == grid.Ny, - V★(i, j, k, grid, args...), + ifelse(j == 1, V★(i, 2, k, grid, args...), δyᵃᶜᵃ(i, j, k, grid, V★, args...))) + +@inline δxᶜᵃᵃ_U(i, j, k, grid, ::Type{LeftConnected}, U★::Function, args...) = ifelse(i == grid.Nx, - U★(i, j, k, grid, args...), δxᶜᵃᵃ(i, j, k, grid, U★, args...)) +@inline δyᵃᶜᵃ_V(i, j, k, grid, ::Type{LeftConnected}, V★::Function, args...) = ifelse(j == grid.Ny, - V★(i, j, k, grid, args...), δyᵃᶜᵃ(i, j, k, grid, V★, args...)) + +@inline δxᶜᵃᵃ_U(i, j, k, grid, ::Type{RightConnected}, U★::Function, args...) = ifelse(i == 1, U★(2, j, k, grid, args...), δxᶜᵃᵃ(i, j, k, grid, U★, args...)) +@inline δyᵃᶜᵃ_V(i, j, k, grid, ::Type{RightConnected}, V★::Function, args...) = ifelse(j == 1, V★(i, 2, k, grid, args...), δyᵃᶜᵃ(i, j, k, grid, V★, args...)) + +# Derivative Operators + +@inline ∂xᶠᶜᶠ_η(i, j, k, grid, T, η★::Function, args...) = δxᶠᵃᵃ_η(i, j, k, grid, T, η★, args...) / Δxᶠᶜᶠ(i, j, k, grid) +@inline ∂yᶜᶠᶠ_η(i, j, k, grid, T, η★::Function, args...) = δyᵃᶠᵃ_η(i, j, k, grid, T, η★, args...) / Δyᶜᶠᶠ(i, j, k, grid) + +@inline div_xᶜᶜᶠ_U(i, j, k, grid, TX, U★, args...) = 1 / Azᶜᶜᶠ(i, j, k, grid) * δxᶜᵃᵃ_U(i, j, k, grid, TX, Δy_qᶠᶜᶠ, U★, args...) +@inline div_yᶜᶜᶠ_V(i, j, k, grid, TY, V★, args...) = 1 / Azᶜᶜᶠ(i, j, k, grid) * δyᵃᶜᵃ_V(i, j, k, grid, TY, Δx_qᶜᶠᶠ, V★, args...) + +# Immersed Boundary Operators (Velocities are `0` on `peripheral_node`s and the free surface should ensure no-flux on `inactive_node`s) + +@inline conditional_U_fcc(i, j, k, grid, ibg::IBG, U★::Function, args...) = ifelse(peripheral_node(i, j, k, ibg, f, c, c), zero(ibg), U★(i, j, k, grid, args...)) +@inline conditional_V_cfc(i, j, k, grid, ibg::IBG, V★::Function, args...) = ifelse(peripheral_node(i, j, k, ibg, c, f, c), zero(ibg), V★(i, j, k, grid, args...)) + +@inline conditional_∂xᶠᶜᶠ_η(i, j, k, ibg::IBG, args...) = ifelse(inactive_node(i, j, k, ibg, c, c, f) | inactive_node(i-1, j, k, ibg, c, c, f), zero(ibg), ∂xᶠᶜᶠ_η(i, j, k, ibg.underlying_grid, args...)) +@inline conditional_∂yᶜᶠᶠ_η(i, j, k, ibg::IBG, args...) = ifelse(inactive_node(i, j, k, ibg, c, c, f) | inactive_node(i, j-1, k, ibg, c, c, f), zero(ibg), ∂yᶜᶠᶠ_η(i, j, k, ibg.underlying_grid, args...)) + +@inline δxᶜᵃᵃ_U(i, j, k, ibg::IBG, T, U★::Function, args...) = δxᶜᵃᵃ_U(i, j, k, ibg.underlying_grid, T, conditional_U_fcc, ibg, U★, args...) +@inline δyᵃᶜᵃ_V(i, j, k, ibg::IBG, T, V★::Function, args...) = δyᵃᶜᵃ_V(i, j, k, ibg.underlying_grid, T, conditional_V_cfc, ibg, V★, args...) +@inline ∂xᶠᶜᶠ_η(i, j, k, ibg::IBG, T, η★::Function, args...) = conditional_∂xᶠᶜᶠ_η(i, j, k, ibg, T, η★, args...) +@inline ∂yᶜᶠᶠ_η(i, j, k, ibg::IBG, T, η★::Function, args...) = conditional_∂yᶜᶠᶠ_η(i, j, k, ibg, T, η★, args...) + +# Disambiguation +for Topo in [:Periodic, :Bounded, :RightConnected, :LeftConnected] + @eval begin + @inline δxᶜᵃᵃ_U(i, j, k, ibg::IBG, T::Type{$Topo}, U★::Function, args...) = δxᶜᵃᵃ_U(i, j, k, ibg.underlying_grid, T, conditional_U_fcc, ibg, U★, args...) + @inline δyᵃᶜᵃ_V(i, j, k, ibg::IBG, T::Type{$Topo}, V★::Function, args...) = δyᵃᶜᵃ_V(i, j, k, ibg.underlying_grid, T, conditional_V_cfc, ibg, V★, args...) + end +end + +# Time stepping extrapolation U★, and η★ + +# AB3 step +@inline function U★(i, j, k, grid, ::AdamsBashforth3Scheme, Uᵐ, Uᵐ⁻¹, Uᵐ⁻²) + FT = eltype(grid) + return @inbounds FT(α) * Uᵐ[i, j, k] + FT(θ) * Uᵐ⁻¹[i, j, k] + FT(β) * Uᵐ⁻²[i, j, k] +end + +@inline function η★(i, j, k, grid, ::AdamsBashforth3Scheme, ηᵐ⁺¹, ηᵐ, ηᵐ⁻¹, ηᵐ⁻²) + FT = eltype(grid) + return @inbounds FT(δ) * ηᵐ⁺¹[i, j, k] + FT(μ) * ηᵐ[i, j, k] + FT(γ) * ηᵐ⁻¹[i, j, k] + FT(ϵ) * ηᵐ⁻²[i, j, k] +end + +# Forward Backward Step +@inline U★(i, j, k, grid, ::ForwardBackwardScheme, U, args...) = @inbounds U[i, j, k] +@inline η★(i, j, k, grid, ::ForwardBackwardScheme, η, args...) = @inbounds η[i, j, k] + +@inline advance_previous_velocity!(i, j, k, ::ForwardBackwardScheme, U, Uᵐ⁻¹, Uᵐ⁻²) = nothing + +@inline function advance_previous_velocity!(i, j, k, ::AdamsBashforth3Scheme, U, Uᵐ⁻¹, Uᵐ⁻²) + @inbounds Uᵐ⁻²[i, j, k] = Uᵐ⁻¹[i, j, k] + @inbounds Uᵐ⁻¹[i, j, k] = U[i, j, k] + + return nothing +end + +@inline advance_previous_free_surface!(i, j, k, ::ForwardBackwardScheme, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻²) = nothing + +@inline function advance_previous_free_surface!(i, j, k, ::AdamsBashforth3Scheme, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻²) + @inbounds ηᵐ⁻²[i, j, k] = ηᵐ⁻¹[i, j, k] + @inbounds ηᵐ⁻¹[i, j, k] = ηᵐ[i, j, k] + @inbounds ηᵐ[i, j, k] = η[i, j, k] + + return nothing +end + +@kernel function _split_explicit_free_surface!(grid, Δτ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², U, V, Uᵐ⁻¹, Uᵐ⁻², Vᵐ⁻¹, Vᵐ⁻², timestepper) + i, j = @index(Global, NTuple) + free_surface_evolution!(i, j, grid, Δτ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², U, V, Uᵐ⁻¹, Uᵐ⁻², Vᵐ⁻¹, Vᵐ⁻², timestepper) +end + +@inline function free_surface_evolution!(i, j, grid, Δτ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², U, V, Uᵐ⁻¹, Uᵐ⁻², Vᵐ⁻¹, Vᵐ⁻², timestepper) + k_top = grid.Nz+1 + TX, TY, _ = topology(grid) + + @inbounds begin + advance_previous_free_surface!(i, j, k_top, timestepper, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻²) + + η[i, j, k_top] -= Δτ * (div_xᶜᶜᶠ_U(i, j, k_top-1, grid, TX, U★, timestepper, U, Uᵐ⁻¹, Uᵐ⁻²) + + div_yᶜᶜᶠ_V(i, j, k_top-1, grid, TY, U★, timestepper, V, Vᵐ⁻¹, Vᵐ⁻²)) + end + + return nothing +end + +@kernel function _split_explicit_barotropic_velocity!(averaging_weight, grid, Δτ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², + U, Uᵐ⁻¹, Uᵐ⁻², V, Vᵐ⁻¹, Vᵐ⁻², + η̅, U̅, V̅, Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ, g, + timestepper) + i, j = @index(Global, NTuple) + velocity_evolution!(i, j, grid, Δτ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², + U, Uᵐ⁻¹, Uᵐ⁻², V, Vᵐ⁻¹, Vᵐ⁻², + η̅, U̅, V̅, averaging_weight, + Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ, g, + timestepper) +end + +@inline function velocity_evolution!(i, j, grid, Δτ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², + U, Uᵐ⁻¹, Uᵐ⁻², V, Vᵐ⁻¹, Vᵐ⁻², + η̅, U̅, V̅, averaging_weight, + Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ, g, + timestepper) + k_top = grid.Nz+1 + + TX, TY, _ = topology(grid) + + @inbounds begin + advance_previous_velocity!(i, j, k_top-1, timestepper, U, Uᵐ⁻¹, Uᵐ⁻²) + advance_previous_velocity!(i, j, k_top-1, timestepper, V, Vᵐ⁻¹, Vᵐ⁻²) + + # ∂τ(U) = - ∇η + G + U[i, j, k_top-1] += Δτ * (- g * Hᶠᶜ[i, j, 1] * ∂xᶠᶜᶠ_η(i, j, k_top, grid, TX, η★, timestepper, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻²) + Gᵁ[i, j, k_top-1]) + V[i, j, k_top-1] += Δτ * (- g * Hᶜᶠ[i, j, 1] * ∂yᶜᶠᶠ_η(i, j, k_top, grid, TY, η★, timestepper, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻²) + Gⱽ[i, j, k_top-1]) + + # time-averaging + η̅[i, j, k_top] += averaging_weight * η[i, j, k_top] + U̅[i, j, k_top-1] += averaging_weight * U[i, j, k_top-1] + V̅[i, j, k_top-1] += averaging_weight * V[i, j, k_top-1] + end +end + +# Barotropic Model Kernels +# u_Δz = u * Δz +@kernel function _barotropic_mode_kernel!(U, V, grid, ::Nothing, u, v) + i, j = @index(Global, NTuple) + k_top = grid.Nz+1 + + @inbounds U[i, j, k_top-1] = Δzᶠᶜᶜ(i, j, 1, grid) * u[i, j, 1] + @inbounds V[i, j, k_top-1] = Δzᶜᶠᶜ(i, j, 1, grid) * v[i, j, 1] + + for k in 2:grid.Nz + @inbounds U[i, j, k_top-1] += Δzᶠᶜᶜ(i, j, k, grid) * u[i, j, k] + @inbounds V[i, j, k_top-1] += Δzᶜᶠᶜ(i, j, k, grid) * v[i, j, k] + end +end + +# Barotropic Model Kernels +# u_Δz = u * Δz +@kernel function _barotropic_mode_kernel!(U, V, grid, active_cells_map, u, v) + idx = @index(Global, Linear) + i, j = active_linear_index_to_tuple(idx, active_cells_map) + k_top = grid.Nz+1 + + @inbounds U[i, j, k_top-1] = Δzᶠᶜᶜ(i, j, 1, grid) * u[i, j, 1] + @inbounds V[i, j, k_top-1] = Δzᶜᶠᶜ(i, j, 1, grid) * v[i, j, 1] + + for k in 2:grid.Nz + @inbounds U[i, j, k_top-1] += Δzᶠᶜᶜ(i, j, k, grid) * u[i, j, k] + @inbounds V[i, j, k_top-1] += Δzᶜᶠᶜ(i, j, k, grid) * v[i, j, k] + end +end + +@inline function compute_barotropic_mode!(U, V, grid, u, v) + active_cells_map = retrieve_surface_active_cells_map(grid) + + launch!(architecture(grid), grid, :xy, _barotropic_mode_kernel!, U, V, grid, active_cells_map, u, v; active_cells_map) + + return nothing +end + +function initialize_free_surface_state!(state, η, timestepper) + + parent(state.U) .= parent(state.U̅) + parent(state.V) .= parent(state.V̅) + + initialize_auxiliary_state!(state, η, timestepper) + + fill!(state.η̅, 0) + fill!(state.U̅, 0) + fill!(state.V̅, 0) + + return nothing +end + +initialize_auxiliary_state!(state, η, ::ForwardBackwardScheme) = nothing + +function initialize_auxiliary_state!(state, η, timestepper) + parent(state.Uᵐ⁻¹) .= parent(state.U̅) + parent(state.Vᵐ⁻¹) .= parent(state.V̅) + + parent(state.Uᵐ⁻²) .= parent(state.U̅) + parent(state.Vᵐ⁻²) .= parent(state.V̅) + + parent(state.ηᵐ) .= parent(η) + parent(state.ηᵐ⁻¹) .= parent(η) + parent(state.ηᵐ⁻²) .= parent(η) + + return nothing +end + +@kernel function _barotropic_split_explicit_corrector!(u, v, U̅, V̅, U, V, Hᶠᶜ, Hᶜᶠ, grid) + i, j, k = @index(Global, NTuple) + k_top = grid.Nz+1 + + @inbounds begin + u[i, j, k] = u[i, j, k] + (U̅[i, j, k_top-1] - U[i, j, k_top-1]) / Hᶠᶜ[i, j, 1] + v[i, j, k] = v[i, j, k] + (V̅[i, j, k_top-1] - V[i, j, k_top-1]) / Hᶜᶠ[i, j, 1] + end +end + +function barotropic_split_explicit_corrector!(u, v, free_surface, grid) + sefs = free_surface.state + U, V, U̅, V̅ = sefs.U, sefs.V, sefs.U̅, sefs.V̅ + Hᶠᶜ, Hᶜᶠ = free_surface.auxiliary.Hᶠᶜ, free_surface.auxiliary.Hᶜᶠ + arch = architecture(grid) + + + # take out "bad" barotropic mode, + # !!!! reusing U and V for this storage since last timestep doesn't matter + compute_barotropic_mode!(U, V, grid, u, v) + # add in "good" barotropic mode + launch!(arch, grid, :xyz, _barotropic_split_explicit_corrector!, + u, v, U̅, V̅, U, V, Hᶠᶜ, Hᶜᶠ, grid) + + return nothing +end + +""" +Explicitly step forward η in substeps. +""" +ab2_step_free_surface!(free_surface::SplitExplicitFreeSurface, model, Δt, χ) = + split_explicit_free_surface_step!(free_surface, model, Δt, χ) + +function initialize_free_surface!(sefs::SplitExplicitFreeSurface, grid, velocities) + @apply_regionally compute_barotropic_mode!(sefs.state.U̅, sefs.state.V̅, grid, velocities.u, velocities.v) + fill_halo_regions!((sefs.state.U̅, sefs.state.V̅, sefs.η)) +end + +function split_explicit_free_surface_step!(free_surface::SplitExplicitFreeSurface, model, Δt, χ) + + # Note: free_surface.η.grid != model.grid for DistributedSplitExplicitFreeSurface + # since halo_size(free_surface.η.grid) != halo_size(model.grid) + free_surface_grid = free_surface.η.grid + + # Wait for previous set up + wait_free_surface_communication!(free_surface, architecture(free_surface_grid)) + + # Calculate the substepping parameterers + settings = free_surface.settings + Nsubsteps = calculate_substeps(settings.substepping, Δt) + + # barotropic time step as fraction of baroclinic step and averaging weights + fractional_Δt, weights = calculate_adaptive_settings(settings.substepping, Nsubsteps) + Nsubsteps = length(weights) + + # barotropic time step in seconds + Δτᴮ = fractional_Δt * Δt + + # reset free surface averages + @apply_regionally begin + initialize_free_surface_state!(free_surface.state, free_surface.η, settings.timestepper) + + # Solve for the free surface at tⁿ⁺¹ + iterate_split_explicit!(free_surface, free_surface_grid, Δτᴮ, weights, Val(Nsubsteps)) + + # Reset eta for the next timestep + set!(free_surface.η, free_surface.state.η̅) + end + + fields_to_fill = (free_surface.state.U̅, free_surface.state.V̅) + fill_halo_regions!(fields_to_fill; async = true) + + # Preparing velocities for the barotropic correction + @apply_regionally begin + mask_immersed_field!(model.velocities.u) + mask_immersed_field!(model.velocities.v) + end + + return nothing +end + +# Change name +const FNS = FixedSubstepNumber +const FTS = FixedTimeStepSize + +# since weights can be negative in the first few substeps (as in the default averaging kernel), +# we set a minimum number of substeps to execute to avoid numerical issues +const MINIMUM_SUBSTEPS = 5 + +@inline calculate_substeps(substepping::FNS, Δt=nothing) = length(substepping.averaging_weights) +@inline calculate_substeps(substepping::FTS, Δt) = max(MINIMUM_SUBSTEPS, ceil(Int, 2 * Δt / substepping.Δt_barotropic)) + +@inline calculate_adaptive_settings(substepping::FNS, substeps) = substepping.fractional_step_size, substepping.averaging_weights +@inline calculate_adaptive_settings(substepping::FTS, substeps) = weights_from_substeps(eltype(substepping.Δt_barotropic), + substeps, substepping.averaging_kernel) + +const FixedSubstepsSetting{N} = SplitExplicitSettings{<:FixedSubstepNumber{<:Any, <:NTuple{N, <:Any}}} where N +const FixedSubstepsSplitExplicit{F} = SplitExplicitFreeSurface{<:Any, <:Any, <:Any, <:Any, <:FixedSubstepsSetting{N}} where N + +function iterate_split_explicit!(free_surface, grid, Δτᴮ, weights, ::Val{Nsubsteps}) where Nsubsteps + arch = architecture(grid) + + η = free_surface.η + state = free_surface.state + auxiliary = free_surface.auxiliary + settings = free_surface.settings + g = free_surface.gravitational_acceleration + + # unpack state quantities, parameters and forcing terms + U, V = state.U, state.V + Uᵐ⁻¹, Uᵐ⁻² = state.Uᵐ⁻¹, state.Uᵐ⁻² + Vᵐ⁻¹, Vᵐ⁻² = state.Vᵐ⁻¹, state.Vᵐ⁻² + ηᵐ, ηᵐ⁻¹, ηᵐ⁻² = state.ηᵐ, state.ηᵐ⁻¹, state.ηᵐ⁻² + η̅, U̅, V̅ = state.η̅, state.U̅, state.V̅ + Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ = auxiliary.Gᵁ, auxiliary.Gⱽ, auxiliary.Hᶠᶜ, auxiliary.Hᶜᶠ + + timestepper = settings.timestepper + + parameters = auxiliary.kernel_parameters + + free_surface_kernel!, _ = configure_kernel(arch, grid, parameters, _split_explicit_free_surface!) + barotropic_velocity_kernel!, _ = configure_kernel(arch, grid, parameters, _split_explicit_barotropic_velocity!) + + η_args = (grid, Δτᴮ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², + U, V, Uᵐ⁻¹, Uᵐ⁻², Vᵐ⁻¹, Vᵐ⁻², + timestepper) + + U_args = (grid, Δτᴮ, η, ηᵐ, ηᵐ⁻¹, ηᵐ⁻², + U, Uᵐ⁻¹, Uᵐ⁻², V, Vᵐ⁻¹, Vᵐ⁻², + η̅, U̅, V̅, Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ, g, + timestepper) + + GC.@preserve η_args U_args begin + + # We need to perform ~50 time-steps which means + # launching ~100 very small kernels: we are limited by + # latency of argument conversion to GPU-compatible values. + # To alleviate this penalty we convert first and then we substep! + converted_η_args = convert_args(arch, η_args) + converted_U_args = convert_args(arch, U_args) + + @unroll for substep in 1:Nsubsteps + Base.@_inline_meta + averaging_weight = weights[substep] + free_surface_kernel!(converted_η_args...) + barotropic_velocity_kernel!(averaging_weight, converted_U_args...) + end + end + + return nothing +end + +# Calculate RHS for the barotropic time step. +@kernel function _compute_integrated_ab2_tendencies!(Gᵁ, Gⱽ, grid, ::Nothing, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ) + i, j = @index(Global, NTuple) + k_top = grid.Nz + 1 + + @inbounds Gᵁ[i, j, k_top-1] = Δzᶠᶜᶜ(i, j, 1, grid) * ab2_step_Gu(i, j, 1, grid, Gu⁻, Guⁿ, χ) + @inbounds Gⱽ[i, j, k_top-1] = Δzᶜᶠᶜ(i, j, 1, grid) * ab2_step_Gv(i, j, 1, grid, Gv⁻, Gvⁿ, χ) + + for k in 2:grid.Nz + @inbounds Gᵁ[i, j, k_top-1] += Δzᶠᶜᶜ(i, j, k, grid) * ab2_step_Gu(i, j, k, grid, Gu⁻, Guⁿ, χ) + @inbounds Gⱽ[i, j, k_top-1] += Δzᶜᶠᶜ(i, j, k, grid) * ab2_step_Gv(i, j, k, grid, Gv⁻, Gvⁿ, χ) + end +end + +# Calculate RHS for the barotropic time step.q +@kernel function _compute_integrated_ab2_tendencies!(Gᵁ, Gⱽ, grid, active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ) + idx = @index(Global, Linear) + i, j = active_linear_index_to_tuple(idx, active_cells_map) + k_top = grid.Nz+1 + + @inbounds Gᵁ[i, j, k_top-1] = Δzᶠᶜᶜ(i, j, 1, grid) * ab2_step_Gu(i, j, 1, grid, Gu⁻, Guⁿ, χ) + @inbounds Gⱽ[i, j, k_top-1] = Δzᶜᶠᶜ(i, j, 1, grid) * ab2_step_Gv(i, j, 1, grid, Gv⁻, Gvⁿ, χ) + + for k in 2:grid.Nz + @inbounds Gᵁ[i, j, k_top-1] += Δzᶠᶜᶜ(i, j, k, grid) * ab2_step_Gu(i, j, k, grid, Gu⁻, Guⁿ, χ) + @inbounds Gⱽ[i, j, k_top-1] += Δzᶜᶠᶜ(i, j, k, grid) * ab2_step_Gv(i, j, k, grid, Gv⁻, Gvⁿ, χ) + end +end + +@inline ab2_step_Gu(i, j, k, grid, G⁻, Gⁿ, χ::FT) where FT = + @inbounds ifelse(peripheral_node(i, j, k, grid, f, c, c), zero(grid), (convert(FT, 1.5) + χ) * Gⁿ[i, j, k] - G⁻[i, j, k] * (convert(FT, 0.5) + χ)) + +@inline ab2_step_Gv(i, j, k, grid, G⁻, Gⁿ, χ::FT) where FT = + @inbounds ifelse(peripheral_node(i, j, k, grid, c, f, c), zero(grid), (convert(FT, 1.5) + χ) * Gⁿ[i, j, k] - G⁻[i, j, k] * (convert(FT, 0.5) + χ)) + +# Setting up the RHS for the barotropic step (tendencies of the barotropic velocity components) +# This function is called after `calculate_tendency` and before `ab2_step_velocities!` +function setup_free_surface!(model, free_surface::SplitExplicitFreeSurface, χ) + + # we start the time integration of η from the average ηⁿ + Gu⁻ = model.timestepper.G⁻.u + Gv⁻ = model.timestepper.G⁻.v + Guⁿ = model.timestepper.Gⁿ.u + Gvⁿ = model.timestepper.Gⁿ.v + + auxiliary = free_surface.auxiliary + + @apply_regionally setup_split_explicit_tendency!(auxiliary, model.grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ) + + fields_to_fill = (auxiliary.Gᵁ, auxiliary.Gⱽ) + fill_halo_regions!(fields_to_fill; async = true) + + return nothing +end + +@inline function setup_split_explicit_tendency!(auxiliary, grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ) + active_cells_map = retrieve_surface_active_cells_map(grid) + + launch!(architecture(grid), grid, :xy, _compute_integrated_ab2_tendencies!, auxiliary.Gᵁ, auxiliary.Gⱽ, grid, + active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ; active_cells_map) + + return nothing +end + +wait_free_surface_communication!(free_surface, arch) = nothing + diff --git a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl index 89aa0d6ab8..cf6b20eade 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl @@ -39,7 +39,6 @@ function update_state!(model::HydrostaticFreeSurfaceModel, grid, callbacks; comp @apply_regionally update_boundary_condition!(fields(model), model) fill_halo_regions!(prognostic_fields(model), model.clock, fields(model); async = true) - @apply_regionally replace_horizontal_vector_halos!(model.velocities, model.grid) @apply_regionally compute_auxiliaries!(model) diff --git a/src/Models/LagrangianParticleTracking/LagrangianParticleTracking.jl b/src/Models/LagrangianParticleTracking/LagrangianParticleTracking.jl index d9c6090d3a..43c39ee5e1 100644 --- a/src/Models/LagrangianParticleTracking/LagrangianParticleTracking.jl +++ b/src/Models/LagrangianParticleTracking/LagrangianParticleTracking.jl @@ -17,7 +17,7 @@ using Oceananigans.Grids: XYFlatGrid, YZFlatGrid, XZFlatGrid using Oceananigans.ImmersedBoundaries: immersed_cell using Oceananigans.Architectures: device, architecture using Oceananigans.Fields: interpolate, datatuple, compute!, location -using Oceananigans.Fields: fractional_indices +using Oceananigans.Fields: fractional_indices, truncate_fractional_indices using Oceananigans.TimeSteppers: AbstractLagrangianParticles using Oceananigans.Utils: prettysummary, launch!, SumOfArrays diff --git a/src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl b/src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl index 3445e24ac1..0bb693c171 100644 --- a/src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl +++ b/src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl @@ -1,6 +1,5 @@ -using Oceananigans.Utils: instantiate, KernelParameters +using Oceananigans.Utils: instantiate using Oceananigans.Models: total_velocities -using Oceananigans.Fields: interpolator ##### ##### Boundary conditions for Lagrangian particles @@ -55,35 +54,30 @@ bouncing the particle off the immersed boundary with a coefficient or `restituti @inline function bounce_immersed_particle((x, y, z), ibg, restitution, previous_particle_indices) X = flattened_node((x, y, z), ibg) - # Determine current particle cell from the interfaces - fi, fj, fk = fractional_indices(X, ibg.underlying_grid, f, f, f) - - i, i⁺, _ = interpolator(fi) - j, j⁺, _ = interpolator(fj) - k, k⁺, _ = interpolator(fk) + # Determine current particle cell + fi, fj, fk = fractional_indices(X, ibg.underlying_grid, c, c, c) + i, j, k = truncate_fractional_indices(fi, fj, fk) # Determine whether particle was _previously_ in a non-immersed cell i⁻, j⁻, k⁻ = previous_particle_indices - tx, ty, tz = map(immersed_boundary_topology, topology(ibg)) + # Left bounds of the previous cell + xᴿ = ξnode(i⁻ + 1, j⁻ + 1, k⁻ + 1, ibg, f, f, f) + yᴿ = ηnode(i⁻ + 1, j⁻ + 1, k⁻ + 1, ibg, f, f, f) + zᴿ = rnode(i⁻ + 1, j⁻ + 1, k⁻ + 1, ibg, f, f, f) # Right bounds of the previous cell - xᴿ = ξnode(i⁺, j, k, ibg, f, f, f) - yᴿ = ηnode(i, j⁺, k, ibg, f, f, f) - zᴿ = rnode(i, j, k⁺, ibg, f, f, f) - - # Left bounds of the previous cell xᴸ = ξnode(i⁻, j⁻, k⁻, ibg, f, f, f) yᴸ = ηnode(i⁻, j⁻, k⁻, ibg, f, f, f) zᴸ = rnode(i⁻, j⁻, k⁻, ibg, f, f, f) Cʳ = restitution - + tx, ty, tz = map(immersed_boundary_topology, topology(ibg)) xb⁺ = enforce_boundary_conditions(tx, x, xᴸ, xᴿ, Cʳ) yb⁺ = enforce_boundary_conditions(ty, y, yᴸ, yᴿ, Cʳ) zb⁺ = enforce_boundary_conditions(tz, z, zᴸ, zᴿ, Cʳ) - immersed = immersed_cell(i⁺, j⁺, k⁺, ibg) + immersed = immersed_cell(i, j, k, ibg) x⁺ = ifelse(immersed, xb⁺, x) y⁺ = ifelse(immersed, yb⁺, y) z⁺ = ifelse(immersed, zb⁺, z) @@ -96,7 +90,7 @@ end Return the index of the rightmost cell interface for a grid with `topology` and `N` cells. """ -rightmost_interface_index(::Bounded, N) = N + 1 +rightmost_interface_index(::Bounded, N) = N + 1 rightmost_interface_index(::Periodic, N) = N + 1 rightmost_interface_index(::Flat, N) = N @@ -109,12 +103,9 @@ given `velocities`, time-step `Δt, and coefficient of `restitution`. @inline function advect_particle((x, y, z), p, restitution, grid, Δt, velocities) X = flattened_node((x, y, z), grid) - # Obtain current particle indices, looking at the interfaces - fi, fj, fk = fractional_indices(X, grid, f, f, f) - - i, i⁺, _ = interpolator(fi) - j, j⁺, _ = interpolator(fj) - k, k⁺, _ = interpolator(fk) + # Obtain current particle indices + fi, fj, fk = fractional_indices(X, grid, c, c, c) + i, j, k = truncate_fractional_indices(fi, fj, fk) current_particle_indices = (i, j, k) @@ -146,16 +137,15 @@ given `velocities`, time-step `Δt, and coefficient of `restitution`. yᴸ = ηnode(i, 1, k, grid, f, f, f) zᴸ = rnode(i, j, 1, grid, f, f, f) - xᴿ = ξnode(iᴿ, j, k, grid, f, f, f) - yᴿ = ηnode(i, jᴿ, k, grid, f, f, f) - zᴿ = rnode(i, j, kᴿ, grid, f, f, f) + xᴿ = ξnode(iᴿ, j, k, grid, f, f, f) + yᴿ = ηnode(i, jᴿ, k, grid, f, f, f) + zᴿ = rnode(i, j, kᴿ, grid, f, f, f) # Enforce boundary conditions for particles. Cʳ = restitution x⁺ = enforce_boundary_conditions(tx, x⁺, xᴸ, xᴿ, Cʳ) y⁺ = enforce_boundary_conditions(ty, y⁺, yᴸ, yᴿ, Cʳ) z⁺ = enforce_boundary_conditions(tz, z⁺, zᴸ, zᴿ, Cʳ) - if grid isa ImmersedBoundaryGrid previous_particle_indices = current_particle_indices # particle has been advected (x⁺, y⁺, z⁺) = bounce_immersed_particle((x⁺, y⁺, z⁺), grid, Cʳ, previous_particle_indices) @@ -196,11 +186,11 @@ end function advect_lagrangian_particles!(particles, model, Δt) grid = model.grid arch = architecture(grid) - parameters = KernelParameters(1:length(particles)) + workgroup = min(length(particles), 256) + worksize = length(particles) - launch!(arch, grid, parameters, - _advect_particles!, - particles.properties, particles.restitution, model.grid, Δt, total_velocities(model)) + advect_particles_kernel! = _advect_particles!(device(arch), workgroup, worksize) + advect_particles_kernel!(particles.properties, particles.restitution, model.grid, Δt, total_velocities(model)) return nothing end diff --git a/src/Models/Models.jl b/src/Models/Models.jl index 4981986dd5..9399a5b503 100644 --- a/src/Models/Models.jl +++ b/src/Models/Models.jl @@ -11,7 +11,7 @@ export using Oceananigans: AbstractModel, fields, prognostic_fields using Oceananigans.AbstractOperations: AbstractOperation -using Oceananigans.Advection: AbstractAdvectionScheme, Centered, VectorInvariant +using Oceananigans.Advection: AbstractAdvectionScheme, CenteredSecondOrder, VectorInvariant using Oceananigans.Fields: AbstractField, Field, flattened_unique_values, boundary_conditions using Oceananigans.Grids: AbstractGrid, halo_size, inflate_halo_size using Oceananigans.OutputReaders: update_field_time_series!, extract_field_time_series @@ -80,7 +80,7 @@ extract_boundary_conditions(field::Field) = field.boundary_conditions """ Returns a default_tracer_advection, tracer_advection `tuple`. """ validate_tracer_advection(invalid_tracer_advection, grid) = error("$invalid_tracer_advection is invalid tracer_advection!") -validate_tracer_advection(tracer_advection_tuple::NamedTuple, grid) = Centered(), tracer_advection_tuple +validate_tracer_advection(tracer_advection_tuple::NamedTuple, grid) = CenteredSecondOrder(), tracer_advection_tuple validate_tracer_advection(tracer_advection::AbstractAdvectionScheme, grid) = tracer_advection, NamedTuple() validate_tracer_advection(tracer_advection::Nothing, grid) = nothing, NamedTuple() diff --git a/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl b/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl index cb0947b4be..ffa7865bb2 100644 --- a/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl +++ b/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl @@ -91,15 +91,15 @@ function compute_interior_tendency_contributions!(model, kernel_parameters; acti u_kernel_args = tuple(start_momentum_kernel_args..., u_immersed_bc, end_momentum_kernel_args..., - hydrostatic_pressure, clock, forcings.u) + forcings, hydrostatic_pressure, clock) v_kernel_args = tuple(start_momentum_kernel_args..., v_immersed_bc, end_momentum_kernel_args..., - hydrostatic_pressure, clock, forcings.v) + forcings, hydrostatic_pressure, clock) w_kernel_args = tuple(start_momentum_kernel_args..., w_immersed_bc, end_momentum_kernel_args..., - hydrostatic_pressure, clock, forcings.w) + forcings, hydrostatic_pressure, clock) exclude_periphery = true launch!(arch, grid, kernel_parameters, compute_Gu!, @@ -128,7 +128,7 @@ function compute_interior_tendency_contributions!(model, kernel_parameters; acti start_tracer_kernel_args..., c_immersed_bc, end_tracer_kernel_args..., - clock, forcing) + forcing, clock) launch!(arch, grid, kernel_parameters, compute_Gc!, c_tendency, grid, active_cells_map, args; diff --git a/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl b/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl index 91034fc4b3..43b41d5715 100644 --- a/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl +++ b/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl @@ -3,7 +3,7 @@ using OrderedCollections: OrderedDict using Oceananigans.Architectures: AbstractArchitecture using Oceananigans.DistributedComputations: Distributed -using Oceananigans.Advection: Centered, adapt_advection_order +using Oceananigans.Advection: CenteredSecondOrder, adapt_advection_order using Oceananigans.BuoyancyModels: validate_buoyancy, regularize_buoyancy, SeawaterBuoyancy using Oceananigans.Biogeochemistry: validate_biogeochemistry, AbstractBiogeochemistry, biogeochemical_auxiliary_fields using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions @@ -56,7 +56,7 @@ end """ NonhydrostaticModel(; grid, clock = Clock{eltype(grid)}(time = 0), - advection = Centered(), + advection = CenteredSecondOrder(), buoyancy = nothing, coriolis = nothing, stokes_drift = nothing, @@ -113,7 +113,7 @@ Keyword arguments """ function NonhydrostaticModel(; grid, clock = Clock{eltype(grid)}(time = 0), - advection = Centered(), + advection = CenteredSecondOrder(), buoyancy = nothing, coriolis = nothing, stokes_drift = nothing, diff --git a/src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl b/src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl index 061daed61f..cefb5eadc3 100644 --- a/src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl +++ b/src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl @@ -56,9 +56,9 @@ pressure anomaly. tracers, auxiliary_fields, diffusivities, + forcings, hydrostatic_pressure, - clock, - forcing) + clock) model_fields = merge(velocities, tracers, auxiliary_fields) @@ -66,7 +66,7 @@ pressure anomaly. v = SumOfArrays{2}(velocities.v, background_fields.velocities.v), w = SumOfArrays{2}(velocities.w, background_fields.velocities.w)) - total_velocities = with_advective_forcing(forcing, total_velocities) + total_velocities = with_advective_forcing(forcings.u, total_velocities) return ( - div_𝐯u(i, j, k, grid, advection, total_velocities, velocities.u) - div_𝐯u(i, j, k, grid, advection, velocities, background_fields.velocities.u) @@ -77,7 +77,7 @@ pressure anomaly. - immersed_∂ⱼ_τ₁ⱼ(i, j, k, grid, velocities, u_immersed_bc, closure, diffusivities, clock, model_fields) + x_curl_Uˢ_cross_U(i, j, k, grid, stokes_drift, velocities, clock.time) + ∂t_uˢ(i, j, k, grid, stokes_drift, clock.time) - + forcing(i, j, k, grid, clock, model_fields)) + + forcings.u(i, j, k, grid, clock, model_fields)) end """ @@ -119,9 +119,9 @@ pressure anomaly. tracers, auxiliary_fields, diffusivities, + forcings, hydrostatic_pressure, - clock, - forcing) + clock) model_fields = merge(velocities, tracers, auxiliary_fields) @@ -129,7 +129,7 @@ pressure anomaly. v = SumOfArrays{2}(velocities.v, background_fields.velocities.v), w = SumOfArrays{2}(velocities.w, background_fields.velocities.w)) - total_velocities = with_advective_forcing(forcing, total_velocities) + total_velocities = with_advective_forcing(forcings.v, total_velocities) return ( - div_𝐯v(i, j, k, grid, advection, total_velocities, velocities.v) - div_𝐯v(i, j, k, grid, advection, velocities, background_fields.velocities.v) @@ -140,7 +140,7 @@ pressure anomaly. - immersed_∂ⱼ_τ₂ⱼ(i, j, k, grid, velocities, v_immersed_bc, closure, diffusivities, clock, model_fields) + y_curl_Uˢ_cross_U(i, j, k, grid, stokes_drift, velocities, clock.time) + ∂t_vˢ(i, j, k, grid, stokes_drift, clock.time) - + forcing(i, j, k, grid, clock, model_fields)) + + forcings.v(i, j, k, grid, clock, model_fields)) end # Only add buoyancy if the hydrostatic pressure isa Nothing @@ -185,9 +185,9 @@ velocity components, tracer fields, and precalculated diffusivities where applic tracers, auxiliary_fields, diffusivities, + forcings, hydrostatic_pressure, - clock, - forcing) + clock) model_fields = merge(velocities, tracers, auxiliary_fields) @@ -195,7 +195,7 @@ velocity components, tracer fields, and precalculated diffusivities where applic v = SumOfArrays{2}(velocities.v, background_fields.velocities.v), w = SumOfArrays{2}(velocities.w, background_fields.velocities.w)) - total_velocities = with_advective_forcing(forcing, total_velocities) + total_velocities = with_advective_forcing(forcings.w, total_velocities) return ( - div_𝐯w(i, j, k, grid, advection, total_velocities, velocities.w) - div_𝐯w(i, j, k, grid, advection, velocities, background_fields.velocities.w) @@ -205,7 +205,7 @@ velocity components, tracer fields, and precalculated diffusivities where applic - immersed_∂ⱼ_τ₃ⱼ(i, j, k, grid, velocities, w_immersed_bc, closure, diffusivities, clock, model_fields) + z_curl_Uˢ_cross_U(i, j, k, grid, stokes_drift, velocities, clock.time) + ∂t_wˢ(i, j, k, grid, stokes_drift, clock.time) - + forcing(i, j, k, grid, clock, model_fields)) + + forcings.w(i, j, k, grid, clock, model_fields)) end """ @@ -247,8 +247,8 @@ velocity components, tracer fields, and precalculated diffusivities where applic tracers, auxiliary_fields, diffusivities, - clock, - forcing) where tracer_index + forcing, + clock) where tracer_index @inbounds c = tracers[tracer_index] @inbounds background_fields_c = background_fields.tracers[tracer_index] diff --git a/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl b/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl index 705ca507ac..671253be84 100644 --- a/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl +++ b/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl @@ -66,4 +66,4 @@ function compute_auxiliaries!(model::NonhydrostaticModel; p_parameters = tuple(p update_hydrostatic_pressure!(model; parameters = ppar) end return nothing -end +end \ No newline at end of file diff --git a/src/Models/ShallowWaterModels/compute_shallow_water_tendencies.jl b/src/Models/ShallowWaterModels/compute_shallow_water_tendencies.jl index d56995a570..fa5a7708ff 100644 --- a/src/Models/ShallowWaterModels/compute_shallow_water_tendencies.jl +++ b/src/Models/ShallowWaterModels/compute_shallow_water_tendencies.jl @@ -75,14 +75,14 @@ function compute_interior_tendency_contributions!(tendencies, formulation) transport_args = (grid, gravitational_acceleration, advection.momentum, velocities, coriolis, closure, - bathymetry, solution, tracers, diffusivities, clock, formulation) + bathymetry, solution, tracers, diffusivities, forcings, clock, formulation) h_args = (grid, gravitational_acceleration, advection.mass, coriolis, closure, - solution, tracers, diffusivities, clock, formulation) + solution, tracers, diffusivities, forcings, clock, formulation) - launch!(arch, grid, :xyz, compute_Guh!, tendencies[1], transport_args..., forcings[1]; exclude_periphery=true) - launch!(arch, grid, :xyz, compute_Gvh!, tendencies[2], transport_args..., forcings[2]; exclude_periphery=true) - launch!(arch, grid, :xyz, compute_Gh!, tendencies[3], h_args..., forcings.h) + launch!(arch, grid, :xyz, compute_Guh!, tendencies[1], transport_args...; exclude_periphery=true) + launch!(arch, grid, :xyz, compute_Gvh!, tendencies[2], transport_args...; exclude_periphery=true) + launch!(arch, grid, :xyz, compute_Gh!, tendencies[3], h_args...) for (tracer_index, tracer_name) in enumerate(propertynames(tracers)) @inbounds Gc = tendencies[tracer_index+3] @@ -90,7 +90,7 @@ function compute_interior_tendency_contributions!(tendencies, @inbounds c_advection = advection[tracer_name] launch!(arch, grid, :xyz, compute_Gc!, Gc, grid, Val(tracer_index), - c_advection, closure, solution, tracers, diffusivities, clock, formulation, forcing) + c_advection, closure, solution, tracers, diffusivities, forcing, clock, formulation) end return nothing diff --git a/src/Models/ShallowWaterModels/shallow_water_model.jl b/src/Models/ShallowWaterModels/shallow_water_model.jl index 549dc8a6fa..7e51210703 100644 --- a/src/Models/ShallowWaterModels/shallow_water_model.jl +++ b/src/Models/ShallowWaterModels/shallow_water_model.jl @@ -3,7 +3,7 @@ using Oceananigans: AbstractModel, AbstractOutputWriter, AbstractDiagnostic using Oceananigans.Architectures: AbstractArchitecture, CPU using Oceananigans.AbstractOperations: @at, KernelFunctionOperation using Oceananigans.DistributedComputations -using Oceananigans.Advection: Centered, VectorInvariant +using Oceananigans.Advection: CenteredSecondOrder, VectorInvariant using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions using Oceananigans.Fields: Field, tracernames, TracerFields, XFaceField, YFaceField, CenterField, compute! using Oceananigans.Forcings: model_forcing @@ -62,7 +62,7 @@ struct VectorInvariantFormulation end ShallowWaterModel(; grid, gravitational_acceleration, clock = Clock{eltype(grid)}(time = 0), - momentum_advection = UpwindBiased(order=5), + momentum_advection = UpwindBiasedFifthOrder(), tracer_advection = WENO(), mass_advection = WENO(), coriolis = nothing, @@ -86,7 +86,7 @@ Keyword arguments - `gravitational_acceleration`: (required) The gravitational acceleration constant. - `clock`: The `clock` for the model. - `momentum_advection`: The scheme that advects velocities. See `Oceananigans.Advection`. - Default: `UpwindBiased(order=5)`. + Default: `UpwindBiasedFifthOrder()`. - `tracer_advection`: The scheme that advects tracers. See `Oceananigans.Advection`. Default: `WENO()`. - `mass_advection`: The scheme that advects the mass equation. See `Oceananigans.Advection`. Default: `WENO()`. @@ -113,7 +113,7 @@ function ShallowWaterModel(; grid, gravitational_acceleration, clock = Clock{eltype(grid)}(time=0), - momentum_advection = UpwindBiased(order=5), + momentum_advection = UpwindBiasedFifthOrder(), tracer_advection = WENO(), mass_advection = WENO(), coriolis = nothing, diff --git a/src/Models/ShallowWaterModels/solution_and_tracer_tendencies.jl b/src/Models/ShallowWaterModels/solution_and_tracer_tendencies.jl index cc3a6c4636..02d5ac2e55 100644 --- a/src/Models/ShallowWaterModels/solution_and_tracer_tendencies.jl +++ b/src/Models/ShallowWaterModels/solution_and_tracer_tendencies.jl @@ -32,9 +32,9 @@ Compute the tendency for the x-directional transport, uh solution, tracers, diffusivities, + forcings, clock, - formulation, - forcing) + formulation) g = gravitational_acceleration @@ -45,7 +45,7 @@ Compute the tendency for the x-directional transport, uh - x_f_cross_U(i, j, k, grid, coriolis, solution) - bathymetry_contribution_x(i, j, k, grid, g, solution.h, bathymetry, formulation) - sw_∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, formulation) - + forcing(i, j, k, grid, clock, merge(solution, tracers))) + + forcings[1](i, j, k, grid, clock, merge(solution, tracers))) end """ @@ -61,9 +61,9 @@ Compute the tendency for the y-directional transport, vh. solution, tracers, diffusivities, + forcings, clock, - formulation, - forcing) + formulation) g = gravitational_acceleration @@ -74,7 +74,7 @@ Compute the tendency for the y-directional transport, vh. - y_f_cross_U(i, j, k, grid, coriolis, solution) - bathymetry_contribution_y(i, j, k, grid, g, solution.h, bathymetry, formulation) - sw_∂ⱼ_τ₂ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, formulation) - + forcing(i, j, k, grid, clock, merge(solution, tracers))) + + forcings[2](i, j, k, grid, clock, merge(solution, tracers))) end """ @@ -88,12 +88,12 @@ Compute the tendency for the height, h. solution, tracers, diffusivities, + forcings, clock, - formulation, - forcing) + formulation) return ( - div_Uh(i, j, k, grid, advection, solution, formulation) - + forcing(i, j, k, grid, clock, merge(solution, tracers))) + + forcings.h(i, j, k, grid, clock, merge(solution, tracers))) end @inline function tracer_tendency(i, j, k, grid, @@ -103,9 +103,9 @@ end solution, tracers, diffusivities, + forcing, clock, - formulation, - forcing) where tracer_index + formulation) where tracer_index @inbounds c = tracers[tracer_index] diff --git a/src/Models/seawater_density.jl b/src/Models/seawater_density.jl index c5e723f101..6dbfdefb69 100644 --- a/src/Models/seawater_density.jl +++ b/src/Models/seawater_density.jl @@ -82,7 +82,7 @@ julia> model = NonhydrostaticModel(; grid, buoyancy, tracers) NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0) ├── grid: 1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo ├── timestepper: RungeKutta3TimeStepper -├── advection scheme: Centered(order=2) +├── advection scheme: Centered reconstruction order 2 ├── tracers: (T, S) ├── closure: Nothing ├── buoyancy: SeawaterBuoyancy with g=9.80665 and BoussinesqEquationOfState{Float64} with ĝ = NegativeZDirection() diff --git a/src/MultiRegion/cubed_sphere_grid.jl b/src/MultiRegion/cubed_sphere_grid.jl index 0cfa3db30e..f4ae660b1e 100644 --- a/src/MultiRegion/cubed_sphere_grid.jl +++ b/src/MultiRegion/cubed_sphere_grid.jl @@ -136,13 +136,11 @@ Example ```jldoctest cubedspheregrid julia> using Oceananigans -julia> using Oceananigans.MultiRegion: ConformalCubedSphereGrid - julia> grid = ConformalCubedSphereGrid(panel_size=(12, 12, 1), z=(-1, 0), radius=1) -ConformalCubedSphereGrid{Float64, Oceananigans.Grids.FullyConnected, Oceananigans.Grids.FullyConnected, Bounded} partitioned on CPU(): -├── grids: 12×12×1 OrthogonalSphericalShellGrid{Float64, Oceananigans.Grids.FullyConnected, Oceananigans.Grids.FullyConnected, Bounded} on CPU with 3×3×3 halo and with precomputed metrics -├── partitioning: CubedSpherePartition with (1 region in each panel) -├── connectivity: CubedSphereConnectivity +ConformalCubedSphereGrid{Float64, FullyConnected, FullyConnected, Bounded} partitioned on CPU(): +├── grids: 12×12×1 OrthogonalSphericalShellGrid{Float64, FullyConnected, FullyConnected, Bounded} on CPU with 3×3×3 halo and with precomputed metrics +├── partitioning: CubedSpherePartition with (1 region in each panel) +├── connectivity: CubedSphereConnectivity └── devices: (CPU(), CPU(), CPU(), CPU(), CPU(), CPU()) ``` diff --git a/src/MultiRegion/multi_region_boundary_conditions.jl b/src/MultiRegion/multi_region_boundary_conditions.jl index 92342a1106..4c21a6d8f7 100644 --- a/src/MultiRegion/multi_region_boundary_conditions.jl +++ b/src/MultiRegion/multi_region_boundary_conditions.jl @@ -101,28 +101,28 @@ end function fill_halo_regions!(c::MultiRegionObject, bcs, indices, loc, mrg::MultiRegionGrid, buffers, args...; fill_boundary_normal_velocities = true, kwargs...) arch = architecture(mrg) - - if fill_boundary_normal_velocities - apply_regionally!(fill_open_boundary_regions!, c, bcs, indices, loc, mrg, args...) - end - - @apply_regionally fill_halos!, permuted_bcs = multi_region_permute_boundary_conditions(bcs) - + @apply_regionally fill_halos!, bcs = multi_region_permute_boundary_conditions(bcs) + # The number of tasks is fixed to 3 (see `multi_region_permute_boundary_conditions`). # When we want to allow asynchronous communication, we will might need to split the halos sides # and the number of tasks might increase. for task in 1:3 @apply_regionally begin - bcs_side = getindex(permuted_bcs, task) + bcs_side = getindex(bcs, task) fill_halo_side! = getindex(fill_halos!, task) fill_multiregion_send_buffers!(c, buffers, mrg, bcs_side) end buff = Reference(buffers.regional_objects) + if fill_boundary_normal_velocities + apply_regionally!(fill_open_boundary_regions!, c, bcs_side, indices, loc, mrg, args...) + end + apply_regionally!(fill_halo_event!, c, fill_halo_side!, bcs_side, indices, loc, arch, mrg, buff, args...; kwargs...) + end return nothing diff --git a/src/MultiRegion/multi_region_field.jl b/src/MultiRegion/multi_region_field.jl index 6abbee7ce2..cb6d06cda8 100644 --- a/src/MultiRegion/multi_region_field.jl +++ b/src/MultiRegion/multi_region_field.jl @@ -70,8 +70,6 @@ Base.size(f::GriddedMultiRegionField) = size(getregion(f.grid, 1)) Reconstruct a global field from `mrf::MultiRegionField` on the `CPU`. """ function reconstruct_global_field(mrf::MultiRegionField) - - # TODO: Is this correct? Shall we reconstruct a global field on the architecture of the grid? global_grid = on_architecture(CPU(), reconstruct_global_grid(mrf.grid)) indices = reconstruct_global_indices(mrf.indices, mrf.grid.partition, size(global_grid)) global_field = Field(location(mrf), global_grid; indices) diff --git a/src/MultiRegion/multi_region_grid.jl b/src/MultiRegion/multi_region_grid.jl index 264b6a29be..0c67d75a49 100644 --- a/src/MultiRegion/multi_region_grid.jl +++ b/src/MultiRegion/multi_region_grid.jl @@ -3,10 +3,9 @@ using Oceananigans.ImmersedBoundaries: GridFittedBottom, PartialCellBottom, Grid import Oceananigans.Grids: architecture, size, new_data, halo_size import Oceananigans.Grids: with_halo, on_architecture -import Oceananigans.Grids: minimum_spacing, destantiate -import Oceananigans.Grids: minimum_xspacing, minimum_yspacing, minimum_zspacing import Oceananigans.Models.HydrostaticFreeSurfaceModels: default_free_surface import Oceananigans.DistributedComputations: reconstruct_global_grid +import Oceananigans.Grids: minimum_spacing, destantiate struct MultiRegionGrid{FT, TX, TY, TZ, P, C, G, D, Arch} <: AbstractMultiRegionGrid{FT, TX, TY, TZ, Arch} architecture :: Arch @@ -20,7 +19,7 @@ struct MultiRegionGrid{FT, TX, TY, TZ, P, C, G, D, Arch} <: AbstractMultiRegionG new{FT, TX, TY, TZ, P, C, G, D, A}(arch, partition, connectivity, region_grids, devices) end -const ImmersedMultiRegionGrid = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:MultiRegionGrid} +const ImmersedMultiRegionGrid = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:MultiRegionGrid} const MultiRegionGrids = Union{MultiRegionGrid, ImmersedMultiRegionGrid} @@ -42,15 +41,6 @@ number_of_regions(mrg::MultiRegionGrids) = lastindex(mrg) minimum_spacing(dir, grid::MultiRegionGrid, ℓx, ℓy, ℓz) = minimum(minimum_spacing(dir, grid[r], ℓx, ℓy, ℓz) for r in 1:number_of_regions(grid)) -minimum_xspacing(grid::MultiRegionGrid, ℓx, ℓy, ℓz) = - minimum(minimum_xspacing(grid[r], ℓx, ℓy, ℓz) for r in 1:number_of_regions(grid)) - -minimum_yspacing(grid::MultiRegionGrid, ℓx, ℓy, ℓz) = - minimum(minimum_yspacing(grid[r], ℓx, ℓy, ℓz) for r in 1:number_of_regions(grid)) - -minimum_zspacing(grid::MultiRegionGrid, ℓx, ℓy, ℓz) = - minimum(minimum_zspacing(grid[r], ℓx, ℓy, ℓz) for r in 1:number_of_regions(grid)) - @inline getdevice(mrg::ImmersedMultiRegionGrid, i) = getdevice(mrg.underlying_grid.region_grids, i) @inline switch_device!(mrg::ImmersedMultiRegionGrid, i) = switch_device!(getdevice(mrg.underlying_grid, i)) @inline devices(mrg::ImmersedMultiRegionGrid) = devices(mrg.underlying_grid.region_grids) @@ -78,7 +68,7 @@ Positional Arguments Keyword Arguments ================= -- `partition`: the partitioning required. The implemented partitioning are `XPartition` +- `partition`: the partitioning required. The implemented partitioning are `XPartition` (division along the ``x`` direction) and `YPartition` (division along the ``y`` direction). @@ -92,14 +82,23 @@ Keyword Arguments Example ======= -```@example multiregion +```jldoctest; filter = r".*@ Oceananigans.MultiRegion.*" julia> using Oceananigans -julia> using Oceananigans.MultiRegion: MultiRegionGrid, XPartition - -julia> grid = RectilinearGrid(size=(12, 12), extent=(1, 1), topology=(Bounded, Bounded, Flat)); +julia> grid = RectilinearGrid(size=(12, 12), extent=(1, 1), topology=(Bounded, Bounded, Flat)) +12×12×1 RectilinearGrid{Float64, Bounded, Bounded, Flat} on CPU with 3×3×0 halo +├── Bounded x ∈ [0.0, 1.0] regularly spaced with Δx=0.0833333 +├── Bounded y ∈ [0.0, 1.0] regularly spaced with Δy=0.0833333 +└── Flat z julia> multi_region_grid = MultiRegionGrid(grid, partition = XPartition(4)) +┌ Warning: MultiRegion functionalities are experimental: help the development by reporting bugs or non-implemented features! +└ @ Oceananigans.MultiRegion ~/Research/OC11.jl/src/MultiRegion/multi_region_grid.jl:108 +MultiRegionGrid{Float64, Bounded, Bounded, Flat} partitioned on CPU(): +├── grids: 3×12×1 RectilinearGrid{Float64, RightConnected, Bounded, Flat} on CPU with 3×3×0 halo +├── partitioning: Equal partitioning in X with (4 regions) +├── connectivity: MultiRegionObject{Tuple{@NamedTuple{west::Nothing, east::Oceananigans.MultiRegion.RegionalConnectivity{Oceananigans.MultiRegion.East, Oceananigans.MultiRegion.West}, north::Nothing, south::Nothing}, @NamedTuple{west::Oceananigans.MultiRegion.RegionalConnectivity{Oceananigans.MultiRegion.West, Oceananigans.MultiRegion.East}, east::Oceananigans.MultiRegion.RegionalConnectivity{Oceananigans.MultiRegion.East, Oceananigans.MultiRegion.West}, north::Nothing, south::Nothing}, @NamedTuple{west::Oceananigans.MultiRegion.RegionalConnectivity{Oceananigans.MultiRegion.West, Oceananigans.MultiRegion.East}, east::Oceananigans.MultiRegion.RegionalConnectivity{Oceananigans.MultiRegion.East, Oceananigans.MultiRegion.West}, north::Nothing, south::Nothing}, @NamedTuple{west::Oceananigans.MultiRegion.RegionalConnectivity{Oceananigans.MultiRegion.West, Oceananigans.MultiRegion.East}, east::Nothing, north::Nothing, south::Nothing}}, NTuple{4, CPU}} +└── devices: (CPU(), CPU(), CPU(), CPU()) ``` """ function MultiRegionGrid(global_grid; partition = XPartition(2), @@ -107,7 +106,7 @@ function MultiRegionGrid(global_grid; partition = XPartition(2), validate = true) @warn "MultiRegion functionalities are experimental: help the development by reporting bugs or non-implemented features!" - + if length(partition) == 1 return global_grid end @@ -182,17 +181,16 @@ end ##### `ImmersedMultiRegionGrid` functionalities ##### -function reconstruct_global_grid(mrg::ImmersedMultiRegionGrid) +function reconstruct_global_grid(mrg::ImmersedMultiRegionGrid) global_grid = reconstruct_global_grid(mrg.underlying_grid) - global_immersed_boundary = reconstruct_global_immersed_boundary(mrg.immersed_boundary) - global_immersed_boundary = on_architecture(architecture(mrg), global_immersed_boundary) + global_boundary = reconstruct_global_boundary(mrg.immersed_boundary) - return ImmersedBoundaryGrid(global_grid, global_immersed_boundary) + return ImmersedBoundaryGrid(global_grid, global_boundary) end -reconstruct_global_immersed_boundary(g::GridFittedBottom{<:Field}) = GridFittedBottom(reconstruct_global_field(g.bottom_height), g.immersed_condition) -reconstruct_global_immersed_boundary(g::PartialCellBottom{<:Field}) = PartialCellBottom(reconstruct_global_field(g.bottom_height), g.minimum_fractional_cell_height) -reconstruct_global_immersed_boundary(g::GridFittedBoundary{<:Field}) = GridFittedBoundary(reconstruct_global_field(g.mask)) +reconstruct_global_boundary(g::GridFittedBottom{<:Field}) = GridFittedBottom(reconstruct_global_field(g.bottom_height), g.immersed_condition) +reconstruct_global_boundary(g::PartialCellBottom{<:Field}) = PartialCellBottom(reconstruct_global_field(g.bottom_height), g.minimum_fractional_cell_height) +reconstruct_global_boundary(g::GridFittedBoundary{<:Field}) = GridFittedBoundary(reconstruct_global_field(g.mask)) @inline getregion(mrg::ImmersedMultiRegionGrid{FT, TX, TY, TZ}, r) where {FT, TX, TY, TZ} = ImmersedBoundaryGrid{TX, TY, TZ}(_getregion(mrg.underlying_grid, r), _getregion(mrg.immersed_boundary, r)) @inline _getregion(mrg::ImmersedMultiRegionGrid{FT, TX, TY, TZ}, r) where {FT, TX, TY, TZ} = ImmersedBoundaryGrid{TX, TY, TZ}( getregion(mrg.underlying_grid, r), getregion(mrg.immersed_boundary, r)) @@ -213,14 +211,14 @@ end # Fallback! multi_region_object_from_array(a::AbstractArray, grid) = on_architecture(architecture(grid), a) -#### +#### #### Utilities for MultiRegionGrid #### new_data(FT::DataType, mrg::MultiRegionGrids, args...) = construct_regionally(new_data, FT, mrg, args...) # This is kind of annoying but it is necessary to have compatible MultiRegion and Distributed -function with_halo(new_halo, mrg::MultiRegionGrid) +function with_halo(new_halo, mrg::MultiRegionGrid) devices = mrg.devices partition = mrg.partition cpu_mrg = on_architecture(CPU(), mrg) @@ -234,11 +232,11 @@ end function on_architecture(::CPU, mrg::MultiRegionGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} new_grids = construct_regionally(on_architecture, CPU(), mrg) - devices = Tuple(CPU() for i in 1:length(mrg)) + devices = Tuple(CPU() for i in 1:length(mrg)) return MultiRegionGrid{FT, TX, TY, TZ}(CPU(), mrg.partition, mrg.connectivity, new_grids, devices) end -Base.summary(mrg::MultiRegionGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = +Base.summary(mrg::MultiRegionGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = "MultiRegionGrid{$FT, $TX, $TY, $TZ} with $(summary(mrg.partition)) on $(string(typeof(mrg.region_grids[1]).name.wrapper))" Base.show(io::IO, mrg::MultiRegionGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = @@ -253,13 +251,13 @@ function Base.:(==)(mrg₁::MultiRegionGrid, mrg₂::MultiRegionGrid) vals = construct_regionally(Base.:(==), mrg₁, mrg₂) return all(vals.regional_objects) end - + #### #### This works only for homogenous partitioning #### -size(mrg::MultiRegionGrids) = size(getregion(mrg, 1)) -halo_size(mrg::MultiRegionGrids) = halo_size(getregion(mrg, 1)) +size(mrg::MultiRegionGrids) = size(getregion(mrg, 1)) +halo_size(mrg::MultiRegionGrids) = halo_size(getregion(mrg, 1)) #### #### Get property for `MultiRegionGrid` (gets the properties of region 1) diff --git a/src/MultiRegion/multi_region_models.jl b/src/MultiRegion/multi_region_models.jl index 21f22f70da..95f677e434 100644 --- a/src/MultiRegion/multi_region_models.jl +++ b/src/MultiRegion/multi_region_models.jl @@ -35,6 +35,8 @@ Types = (:HydrostaticFreeSurfaceModel, :ImplicitFreeSurface, :ExplicitFreeSurface, :QuasiAdamsBashforth2TimeStepper, + :SplitExplicitAuxiliaryFields, + :SplitExplicitState, :SplitExplicitFreeSurface, :PrescribedVelocityFields, :ConjugateGradientSolver, diff --git a/src/MultiRegion/multi_region_split_explicit_free_surface.jl b/src/MultiRegion/multi_region_split_explicit_free_surface.jl index 0d6e4dc086..fb36463b39 100644 --- a/src/MultiRegion/multi_region_split_explicit_free_surface.jl +++ b/src/MultiRegion/multi_region_split_explicit_free_surface.jl @@ -1,47 +1,25 @@ using Oceananigans.Utils using Oceananigans.AbstractOperations: GridMetricOperation, Δz -using Oceananigans.Models.HydrostaticFreeSurfaceModels: free_surface_displacement_field -using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces -using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces: calculate_substeps, - barotropic_bc, - materialize_timestepper +using Oceananigans.Models.HydrostaticFreeSurfaceModels: SplitExplicitFreeSurface, + SplitExplicitSettings, + SplitExplicitState, + FixedSubstepNumber, FixedTimeStepSize, + calculate_substeps -import Oceananigans.Models.HydrostaticFreeSurfaceModels: materialize_free_surface +import Oceananigans.Models.HydrostaticFreeSurfaceModels: materialize_free_surface, SplitExplicitAuxiliaryFields -# Internal function for HydrostaticFreeSurfaceModel -function materialize_free_surface(free_surface::SplitExplicitFreeSurface, velocities, grid::MultiRegionGrids) - - free_surface.substepping isa FixedTimeStepSize && - throw(ArgumentError("SplitExplicitFreeSurface on MultiRegionGrids only suports FixedSubstepNumber; re-initialize SplitExplicitFreeSurface using substeps kwarg")) - - switch_device!(grid.devices[1]) - - old_halos = halo_size(getregion(grid, 1)) - Nsubsteps = calculate_substeps(free_surface.substepping) +function SplitExplicitAuxiliaryFields(grid::MultiRegionGrids) - extended_halos = multiregion_split_explicit_halos(old_halos, Nsubsteps+1, grid.partition) - extended_grid = with_halo(extended_halos, grid) + Gᵁ = Field((Face, Center, Nothing), grid) + Gⱽ = Field((Center, Face, Nothing), grid) - η = free_surface_displacement_field(velocities, free_surface, extended_grid) - η̅ = free_surface_displacement_field(velocities, free_surface, extended_grid) + Hᶠᶜ = Field((Face, Center, Nothing), grid) + Hᶜᶠ = Field((Center, Face, Nothing), grid) - u_baroclinic = velocities.u - v_baroclinic = velocities.v + @apply_regionally calculate_column_height!(Hᶠᶜ, (Face, Center, Center)) + @apply_regionally calculate_column_height!(Hᶜᶠ, (Center, Face, Center)) - @apply_regionally u_bc = barotropic_bc(u_baroclinic) - @apply_regionally v_bc = barotropic_bc(v_baroclinic) - - U = Field{Center, Center, Nothing}(extended_grid, boundary_conditions = u_bc) - V = Field{Center, Center, Nothing}(extended_grid, boundary_conditions = v_bc) - - U̅ = Field{Center, Center, Nothing}(extended_grid, boundary_conditions = u_bc) - V̅ = Field{Center, Center, Nothing}(extended_grid, boundary_conditions = v_bc) - - filtered_state = (η = η̅, U = U̅, V = V̅) - barotropic_velocities = (U = U, V = V) - - gravitational_acceleration = convert(eltype(extended_grid), free_surface.gravitational_acceleration) - timestepper = materialize_timestepper(free_surface.timestepper, extended_grid, free_surface, velocities, u_bc, v_bc) + fill_halo_regions!((Hᶠᶜ, Hᶜᶠ)) # In a non-parallel grid we calculate only the interior @apply_regionally kernel_size = augmented_kernel_size(grid, grid.partition) @@ -49,17 +27,14 @@ function materialize_free_surface(free_surface::SplitExplicitFreeSurface, veloci @apply_regionally kernel_parameters = KernelParameters(kernel_size, kernel_offsets) - return SplitExplicitFreeSurface(η, - barotropic_velocities, - filtered_state, - gravitational_acceleration, - kernel_parameters, - free_surface.substepping, - timestepper) + return SplitExplicitAuxiliaryFields(Gᵁ, Gⱽ, Hᶠᶜ, Hᶜᶠ, kernel_parameters) end -@inline multiregion_split_explicit_halos(old_halos, step_halo, ::XPartition) = (max(step_halo, old_halos[1]), old_halos[2], old_halos[3]) -@inline multiregion_split_explicit_halos(old_halos, step_halo, ::YPartition) = (old_halos[1], max(step_halo, old_halo[2]), old_halos[3]) +@inline function calculate_column_height!(height, location) + dz = GridMetricOperation(location, Δz, height.grid) + sum!(height, dz) + return nothing +end @inline augmented_kernel_size(grid, ::XPartition) = (size(grid, 1) + 2halo_size(grid)[1]-2, size(grid, 2)) @inline augmented_kernel_size(grid, ::YPartition) = (size(grid, 1), size(grid, 2) + 2halo_size(grid)[2]-2) @@ -69,3 +44,29 @@ end @inline augmented_kernel_offsets(grid, ::YPartition) = (0, - halo_size(grid)[2] + 1) @inline augmented_kernel_offsets(grid, ::CubedSpherePartition) = (- halo_size(grid)[2] + 1, - halo_size(grid)[2] + 1) +# Internal function for HydrostaticFreeSurfaceModel +function materialize_free_surface(free_surface::SplitExplicitFreeSurface, velocities, grid::MultiRegionGrids) + settings = SplitExplicitSettings(grid; free_surface.settings.settings_kwargs...) + + settings.substepping isa FixedTimeStepSize && + throw(ArgumentError("SplitExplicitFreeSurface on MultiRegionGrids only suports FixedSubstepNumber; re-initialize SplitExplicitFreeSurface using substeps kwarg")) + + switch_device!(grid.devices[1]) + + old_halos = halo_size(getregion(grid, 1)) + Nsubsteps = calculate_substeps(settings.substepping) + + new_halos = multiregion_split_explicit_halos(old_halos, Nsubsteps+1, grid.partition) + new_grid = with_halo(new_halos, grid) + + η = ZFaceField(new_grid, indices = (:, :, size(new_grid, 3)+1)) + + return SplitExplicitFreeSurface(η, + SplitExplicitState(new_grid, free_surface.settings.timestepper), + SplitExplicitAuxiliaryFields(new_grid), + free_surface.gravitational_acceleration, + free_surface.settings) +end + +@inline multiregion_split_explicit_halos(old_halos, step_halo, ::XPartition) = (max(step_halo, old_halos[1]), old_halos[2], old_halos[3]) +@inline multiregion_split_explicit_halos(old_halos, step_halo, ::YPartition) = (old_halos[1], max(step_halo, old_halo[2]), old_halos[3]) diff --git a/src/Oceananigans.jl b/src/Oceananigans.jl index 39701328c5..db93e518ef 100644 --- a/src/Oceananigans.jl +++ b/src/Oceananigans.jl @@ -6,27 +6,35 @@ module Oceananigans export # Architectures - CPU, GPU, + CPU, GPU, + + # Logging + OceananigansLogger, # Grids Center, Face, - Periodic, Bounded, Flat, - RectilinearGrid, LatitudeLongitudeGrid, OrthogonalSphericalShellGrid, - nodes, xnodes, ynodes, znodes, λnodes, φnodes, - xspacings, yspacings, zspacings, λspacings, φspacings, + Periodic, Bounded, Flat, + FullyConnected, LeftConnected, RightConnected, + RectilinearGrid, + LatitudeLongitudeGrid, + OrthogonalSphericalShellGrid, + xnodes, ynodes, znodes, nodes, + λnodes, φnodes, + xspacings, yspacings, zspacings, minimum_xspacing, minimum_yspacing, minimum_zspacing, # Immersed boundaries - ImmersedBoundaryGrid, - GridFittedBoundary, GridFittedBottom, PartialCellBottom, - ImmersedBoundaryCondition, + ImmersedBoundaryGrid, GridFittedBoundary, GridFittedBottom, ImmersedBoundaryCondition, # Distributed Distributed, Partition, # Advection schemes - Centered, UpwindBiased, WENO, - VectorInvariant, WENOVectorInvariant, FluxFormAdvection, + Centered, CenteredSecondOrder, CenteredFourthOrder, + UpwindBiased, UpwindBiasedFirstOrder, UpwindBiasedThirdOrder, UpwindBiasedFifthOrder, + WENO, WENOThirdOrder, WENOFifthOrder, + VectorInvariant, WENOVectorInvariant, EnergyConserving, EnstrophyConserving, + FluxFormAdvection, # Boundary conditions BoundaryCondition, @@ -36,7 +44,7 @@ export # Fields and field manipulation Field, CenterField, XFaceField, YFaceField, ZFaceField, Average, Integral, CumulativeIntegral, Reduction, Accumulation, BackgroundField, - interior, set!, compute!, regrid!, + interior, set!, compute!, regrid!, location, # Forcing functions Forcing, Relaxation, LinearTarget, GaussianMask, AdvectiveForcing, @@ -60,13 +68,11 @@ export HorizontalScalarBiharmonicDiffusivity, ScalarBiharmonicDiffusivity, SmagorinskyLilly, - Smagorinsky, - LillyCoefficient, - DynamicCoefficient, AnisotropicMinimumDissipation, ConvectiveAdjustmentVerticalDiffusivity, - CATKEVerticalDiffusivity, RiBasedVerticalDiffusivity, + IsopycnalSkewSymmetricDiffusivity, + FluxTapering, VerticallyImplicitTimeDiscretization, viscosity, diffusivity, @@ -74,28 +80,32 @@ export LagrangianParticles, # Models - NonhydrostaticModel, HydrostaticFreeSurfaceModel, ShallowWaterModel, - ConservativeFormulation, VectorInvariantFormulation, - PressureField, fields, + NonhydrostaticModel, + HydrostaticFreeSurfaceModel, + ShallowWaterModel, ConservativeFormulation, VectorInvariantFormulation, + PressureField, + fields, # Hydrostatic free surface model stuff VectorInvariant, ExplicitFreeSurface, ImplicitFreeSurface, SplitExplicitFreeSurface, - HydrostaticSphericalCoriolis, PrescribedVelocityFields, + HydrostaticSphericalCoriolis, + PrescribedVelocityFields, # Time stepping Clock, TimeStepWizard, conjure_time_step_wizard!, time_step!, # Simulations - Simulation, run!, Callback, add_callback!, iteration, + Simulation, run!, Callback, add_callback!, iteration, stopwatch, iteration_limit_exceeded, stop_time_exceeded, wall_time_limit_exceeded, + TimeStepCallsite, TendencyCallsite, UpdateStateCallsite, # Diagnostics - CFL, AdvectiveCFL, DiffusiveCFL, + StateChecker, CFL, AdvectiveCFL, DiffusiveCFL, # Output writers NetCDFOutputWriter, JLD2OutputWriter, Checkpointer, - TimeInterval, IterationInterval, WallTimeInterval, AveragedTimeInterval, - SpecifiedTimes, FileSizeLimit, AndSchedule, OrSchedule, written_names, + TimeInterval, IterationInterval, AveragedTimeInterval, SpecifiedTimes, + FileSizeLimit, AndSchedule, OrSchedule, written_names, # Output readers FieldTimeSeries, FieldDataset, InMemory, OnDisk, @@ -103,32 +113,38 @@ export # Abstract operations ∂x, ∂y, ∂z, @at, KernelFunctionOperation, - # Utils - prettytime + # MultiRegion and Cubed sphere + MultiRegionGrid, MultiRegionField, + XPartition, YPartition, + CubedSpherePartition, ConformalCubedSphereGrid, CubedSphereField, + # Utils + prettytime, apply_regionally!, construct_regionally, @apply_regionally, MultiRegionObject, + + # Units + Time + +using Printf +using Logging +using Statistics +using LinearAlgebra using CUDA +using Adapt using DocStringExtensions +using OffsetArrays using FFTW - -function __init__() - threads = Threads.nthreads() - if threads > 1 - @info "Oceananigans will use $threads threads" - - # See: https://github.com/CliMA/Oceananigans.jl/issues/1113 - FFTW.set_num_threads(4threads) - end - - if CUDA.has_cuda() - @debug "CUDA-enabled GPU(s) detected:" - for (gpu, dev) in enumerate(CUDA.devices()) - @debug "$dev: $(CUDA.name(dev))" - end - - CUDA.allowscalar(false) - end -end - +using JLD2 + +using Base: @propagate_inbounds +using Statistics: mean + +import Base: + +, -, *, /, + size, length, eltype, + iterate, similar, show, + getindex, lastindex, setindex!, + push! + ##### ##### Abstract types ##### @@ -246,4 +262,23 @@ using .Simulations using .AbstractOperations using .MultiRegion +function __init__() + threads = Threads.nthreads() + if threads > 1 + @info "Oceananigans will use $threads threads" + + # See: https://github.com/CliMA/Oceananigans.jl/issues/1113 + FFTW.set_num_threads(4threads) + end + + if CUDA.has_cuda() + @debug "CUDA-enabled GPU(s) detected:" + for (gpu, dev) in enumerate(CUDA.devices()) + @debug "$dev: $(CUDA.name(dev))" + end + + CUDA.allowscalar(false) + end +end + end # module diff --git a/src/Operators/Operators.jl b/src/Operators/Operators.jl index a033f62d9e..570d43dd6b 100644 --- a/src/Operators/Operators.jl +++ b/src/Operators/Operators.jl @@ -10,10 +10,6 @@ export Axᶠᶠᶠ, Axᶠᶠᶜ, Axᶠᶜᶠ, Axᶠᶜᶜ, Axᶜᶠᶠ, Axᶜᶠ export Ayᶠᶠᶠ, Ayᶠᶠᶜ, Ayᶠᶜᶠ, Ayᶠᶜᶜ, Ayᶜᶠᶠ, Ayᶜᶠᶜ, Ayᶜᶜᶠ, Ayᶜᶜᶜ export Azᶠᶠᶠ, Azᶠᶠᶜ, Azᶠᶜᶠ, Azᶠᶜᶜ, Azᶜᶠᶠ, Azᶜᶠᶜ, Azᶜᶜᶠ, Azᶜᶜᶜ -export Axᵃᶜᶜ, Axᵃᶠᶠ, Axᶜᵃᶜ, Axᶠᵃᶠ, Axᶜᶜᵃ, Axᶠᶠᵃ -export Ayᵃᶜᶜ, Ayᵃᶠᶠ, Ayᶜᵃᶜ, Ayᶠᵃᶠ, Ayᶜᶜᵃ, Ayᶠᶠᵃ -export Azᵃᶜᶜ, Azᵃᶠᶠ, Azᶜᵃᶜ, Azᶠᵃᶠ, Azᶜᶜᵃ, Azᶠᶠᵃ - # Volumes export Vᶠᶠᶠ, Vᶠᶠᶜ, Vᶠᶜᶠ, Vᶠᶜᶜ, Vᶜᶠᶠ, Vᶜᶠᶜ, Vᶜᶜᶠ, Vᶜᶜᶜ @@ -34,8 +30,6 @@ export δyᶠᶠᶠ, δyᶠᶠᶜ, δyᶠᶜᶠ, δyᶠᶜᶜ, δyᶜᶠᶠ, δy export δzᶠᶠᶠ, δzᶠᶠᶜ, δzᶠᶜᶠ, δzᶠᶜᶜ, δzᶜᶠᶠ, δzᶜᶠᶜ, δzᶜᶜᶠ, δzᶜᶜᶜ # Derivatives -export ∂xᶜᵃᵃ, ∂xᶠᵃᵃ, ∂yᵃᶜᵃ, ∂yᵃᶠᵃ, ∂zᵃᵃᶜ, ∂zᵃᵃᶠ - export ∂xᶠᶠᶠ, ∂xᶠᶠᶜ, ∂xᶠᶜᶠ, ∂xᶠᶜᶜ, ∂xᶜᶠᶠ, ∂xᶜᶠᶜ, ∂xᶜᶜᶠ, ∂xᶜᶜᶜ export ∂yᶠᶠᶠ, ∂yᶠᶠᶜ, ∂yᶠᶜᶠ, ∂yᶠᶜᶜ, ∂yᶜᶠᶠ, ∂yᶜᶠᶜ, ∂yᶜᶜᶠ, ∂yᶜᶜᶜ export ∂zᶠᶠᶠ, ∂zᶠᶠᶜ, ∂zᶠᶜᶠ, ∂zᶠᶜᶜ, ∂zᶜᶠᶠ, ∂zᶜᶠᶜ, ∂zᶜᶜᶠ, ∂zᶜᶜᶜ @@ -66,10 +60,6 @@ export ℑxᶜᵃᵃ, ℑxᶠᵃᵃ, ℑyᵃᶜᵃ, ℑyᵃᶠᵃ, ℑzᵃᵃᶜ export ℑxyᶜᶜᵃ, ℑxyᶠᶜᵃ, ℑxyᶠᶠᵃ, ℑxyᶜᶠᵃ, ℑxzᶜᵃᶜ, ℑxzᶠᵃᶜ, ℑxzᶠᵃᶠ, ℑxzᶜᵃᶠ, ℑyzᵃᶜᶜ, ℑyzᵃᶠᶜ, ℑyzᵃᶠᶠ, ℑyzᵃᶜᶠ export ℑxyzᶜᶜᶠ, ℑxyzᶜᶠᶜ, ℑxyzᶠᶜᶜ, ℑxyzᶜᶠᶠ, ℑxyzᶠᶜᶠ, ℑxyzᶠᶠᶜ, ℑxyzᶜᶜᶜ, ℑxyzᶠᶠᶠ -# Topology-aware operators -export δxTᶠᵃᵃ, δyTᵃᶠᵃ, δxTᶜᵃᵃ, δyTᵃᶜᵃ -export ∂xTᶠᶜᶠ, ∂yTᶜᶠᶠ - # Reference frame conversion export intrinsic_vector, extrinsic_vector @@ -81,30 +71,12 @@ import Oceananigans.Grids: xspacing, yspacing, zspacing ##### Convenient aliases ##### -const AG = AbstractGrid +const AG = AbstractGrid const Δx = xspacing const Δy = yspacing const Δz = zspacing -const RG = RectilinearGrid -const RGX = XRegularRG -const RGY = YRegularRG -const RGZ = ZRegularRG - -const OSSG = OrthogonalSphericalShellGrid -const OSSGZ = ZRegOrthogonalSphericalShellGrid - -const LLG = LatitudeLongitudeGrid -const LLGX = XRegularLLG -const LLGY = YRegularLLG -const LLGZ = ZRegularLLG - -# On the fly calculations of metrics -const LLGF = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Nothing} -const LLGFX = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Nothing, <:Any, <:Number} -const LLGFY = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Nothing, <:Any, <:Any, <:Number} - include("difference_operators.jl") include("interpolation_operators.jl") include("interpolation_utils.jl") @@ -114,7 +86,6 @@ include("products_between_fields_and_grid_metrics.jl") include("derivative_operators.jl") include("divergence_operators.jl") -include("topology_aware_operators.jl") include("vorticity_operators.jl") include("laplacian_operators.jl") diff --git a/src/Operators/derivative_operators.jl b/src/Operators/derivative_operators.jl index a78dfef074..67f3287df3 100644 --- a/src/Operators/derivative_operators.jl +++ b/src/Operators/derivative_operators.jl @@ -2,8 +2,8 @@ ##### First derivative operators ##### -for LX in (:ᶜ, :ᶠ, :ᵃ), LY in (:ᶜ, :ᶠ, :ᵃ), LZ in (:ᶜ, :ᶠ, :ᵃ) - +for LX in (:ᶜ, :ᶠ), LY in (:ᶜ, :ᶠ), LZ in (:ᶜ, :ᶠ) + x_derivative = Symbol(:∂x, LX, LY, LZ) x_spacing = Symbol(:Δx, LX, LY, LZ) x_difference = Symbol(:δx, LX, LY, LZ) @@ -35,7 +35,7 @@ end ##### Second, Third, and Fourth derivatives ##### -@inline insert_symbol(dir, L, L1, L2) = +@inline insert_symbol(dir, L, L1, L2) = dir == :x ? (L, L1, L2) : dir == :y ? @@ -43,7 +43,7 @@ end (L1, L2, L) -for dir in (:x, :y, :z), L1 in (:ᶜ, :ᶠ, :ᵃ), L2 in (:ᶜ, :ᶠ, :ᵃ) +for dir in (:x, :y, :z), L1 in (:ᶜ, :ᶠ), L2 in (:ᶜ, :ᶠ) first_order_face = Symbol(:∂, dir, insert_symbol(dir, :ᶠ, L1, L2)...) second_order_face = Symbol(:∂², dir, insert_symbol(dir, :ᶠ, L1, L2)...) @@ -78,8 +78,8 @@ end ##### Operators of the form A*∂(q) where A is an area and q is some quantity. ##### -for dir in (:x, :y, :z), LX in (:ᶜ, :ᶠ, :ᵃ), LY in (:ᶜ, :ᶠ, :ᵃ), LZ in (:ᶜ, :ᶠ, :ᵃ) - +for dir in (:x, :y, :z), LX in (:ᶜ, :ᶠ), LY in (:ᶜ, :ᶠ), LZ in (:ᶜ, :ᶠ) + operator = Symbol(:A, dir, :_∂, dir, LX, LY, LZ) area = Symbol(:A, dir, LX, LY, LZ) derivative = Symbol(:∂, dir, LX, LY, LZ) diff --git a/src/Operators/interpolation_utils.jl b/src/Operators/interpolation_utils.jl index a2ff28e4d5..748eacf242 100644 --- a/src/Operators/interpolation_utils.jl +++ b/src/Operators/interpolation_utils.jl @@ -8,10 +8,10 @@ interpolation_code(from, to) = interpolation_code(to) interpolation_code(::Type{Face}) = :ᶠ interpolation_code(::Type{Center}) = :ᶜ -interpolation_code(::Type{Nothing}) = :ᵃ +interpolation_code(::Type{Nothing}) = :ᶜ interpolation_code(::Face) = :ᶠ interpolation_code(::Center) = :ᶜ -interpolation_code(::Nothing) = :ᵃ +interpolation_code(::Nothing) = :ᶜ # Intercept non-interpolations interpolation_code(from::L, to::L) where L = :ᵃ @@ -41,7 +41,7 @@ for i = 1:number_of_identities @inline $identity(i, j, k, grid, F::TF, args...) where TF<:Function = F(i, j, k, grid, args...) end end - + torus(x, lower, upper) = lower + rem(x - lower, upper - lower, RoundDown) identify_an_identity(number) = Symbol(:identity, torus(number, 1, number_of_identities)) identity_counter = 0 @@ -118,10 +118,10 @@ for LX in (:Center, :Face), LY in (:Center, :Face), LZ in (:Center, :Face) to = (eval(IX), eval(IY), eval(IZ)) interp_func = Symbol(interpolation_operator(from, to)) @eval begin - @inline ℑxyz(i, j, k, grid, from::F, to::T, c) where {F<:Tuple{<:$LX, <:$LY, <:$LZ}, T<:Tuple{<:$IX, <:$IY, <:$IZ}} = + @inline ℑxyz(i, j, k, grid, from::F, to::T, c) where {F<:Tuple{<:$LX, <:$LY, <:$LZ}, T<:Tuple{<:$IX, <:$IY, <:$IZ}} = $interp_func(i, j, k, grid, c) - - @inline ℑxyz(i, j, k, grid, from::F, to::T, f, args...) where {F<:Tuple{<:$LX, <:$LY, <:$LZ}, T<:Tuple{<:$IX, <:$IY, <:$IZ}} = + + @inline ℑxyz(i, j, k, grid, from::F, to::T, f, args...) where {F<:Tuple{<:$LX, <:$LY, <:$LZ}, T<:Tuple{<:$IX, <:$IY, <:$IZ}} = $interp_func(i, j, k, grid, f, args...) end end diff --git a/src/Operators/spacings_and_areas_and_volumes.jl b/src/Operators/spacings_and_areas_and_volumes.jl index 649be89975..5735444c63 100644 --- a/src/Operators/spacings_and_areas_and_volumes.jl +++ b/src/Operators/spacings_and_areas_and_volumes.jl @@ -1,5 +1,23 @@ using Oceananigans.Grids: Center, Face +const RG = RectilinearGrid +const RGX = XRegularRG +const RGY = YRegularRG +const RGZ = ZRegularRG + +const OSSG = OrthogonalSphericalShellGrid +const OSSGZ = ZRegOrthogonalSphericalShellGrid + +const LLG = LatitudeLongitudeGrid +const LLGX = XRegularLLG +const LLGY = YRegularLLG +const LLGZ = ZRegularLLG + +# On the fly calculations of metrics +const LLGF = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Nothing} +const LLGFX = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Nothing, <:Any, <:Number} +const LLGFY = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Nothing, <:Any, <:Any, <:Number} + @inline hack_cosd(φ) = cos(π * φ / 180) @inline hack_sind(φ) = sin(π * φ / 180) @@ -35,8 +53,27 @@ The operators in this file fall into three categories: ##### ##### +@inline Δxᶠᵃᵃ(i, j, k, grid) = nothing +@inline Δxᶜᵃᵃ(i, j, k, grid) = nothing +@inline Δyᵃᶠᵃ(i, j, k, grid) = nothing +@inline Δyᵃᶜᵃ(i, j, k, grid) = nothing + +const ZRG = Union{LLGZ, RGZ} + +@inline Δzᵃᵃᶠ(i, j, k, grid) = @inbounds grid.Δzᵃᵃᶠ[k] +@inline Δzᵃᵃᶜ(i, j, k, grid) = @inbounds grid.Δzᵃᵃᶜ[k] + +@inline Δzᵃᵃᶠ(i, j, k, grid::ZRG) = grid.Δzᵃᵃᶠ +@inline Δzᵃᵃᶜ(i, j, k, grid::ZRG) = grid.Δzᵃᵃᶜ + +@inline Δzᵃᵃᶜ(i, j, k, grid::OSSG) = @inbounds grid.Δzᵃᵃᶜ[k] +@inline Δzᵃᵃᶠ(i, j, k, grid::OSSG) = @inbounds grid.Δzᵃᵃᶠ[k] + +@inline Δzᵃᵃᶜ(i, j, k, grid::OSSGZ) = grid.Δzᵃᵃᶜ +@inline Δzᵃᵃᶠ(i, j, k, grid::OSSGZ) = grid.Δzᵃᵃᶠ + # Convenience Functions for all grids -for LX in (:ᶜ, :ᶠ, :ᵃ), LY in (:ᶜ, :ᶠ, :ᵃ) +for LX in (:ᶜ, :ᶠ), LY in (:ᶜ, :ᶠ) x_spacing_1D = Symbol(:Δx, LX, :ᵃ, :ᵃ) x_spacing_2D = Symbol(:Δx, LX, LY, :ᵃ) @@ -68,106 +105,51 @@ end ##### Rectilinear Grids (Flat grids already have Δ = 1) ##### -@inline Δxᶠᵃᵃ(i, j, k, grid::RG) = @inbounds grid.Δxᶠᵃᵃ[i] -@inline Δxᶜᵃᵃ(i, j, k, grid::RG) = @inbounds grid.Δxᶜᵃᵃ[i] -@inline Δxᶜᵃᶜ(i, j, k, grid::RG) = @inbounds grid.Δxᶜᵃᵃ[i] -@inline Δxᶠᵃᶜ(i, j, k, grid::RG) = @inbounds grid.Δxᶠᵃᵃ[i] -@inline Δxᶜᵃᶠ(i, j, k, grid::RG) = @inbounds grid.Δxᶜᵃᵃ[i] - -@inline Δyᵃᶠᵃ(i, j, k, grid::RG) = @inbounds grid.Δyᵃᶠᵃ[j] -@inline Δyᵃᶜᵃ(i, j, k, grid::RG) = @inbounds grid.Δyᵃᶜᵃ[j] -@inline Δyᶜᵃᶜ(i, j, k, grid::RG) = @inbounds grid.Δyᵃᶜᵃ[j] -@inline Δyᶠᵃᶜ(i, j, k, grid::RG) = @inbounds grid.Δyᵃᶜᵃ[j] -@inline Δyᶜᵃᶠ(i, j, k, grid::RG) = @inbounds grid.Δyᵃᶜᵃ[j] - -@inline Δzᵃᵃᶠ(i, j, k, grid::RG) = @inbounds grid.Δzᵃᵃᶠ[k] -@inline Δzᵃᵃᶜ(i, j, k, grid::RG) = @inbounds grid.Δzᵃᵃᶜ[k] -@inline Δzᶜᵃᶜ(i, j, k, grid::RG) = @inbounds grid.Δzᵃᵃᶜ[k] -@inline Δzᶠᵃᶜ(i, j, k, grid::RG) = @inbounds grid.Δzᵃᵃᶜ[k] -@inline Δzᶜᵃᶠ(i, j, k, grid::RG) = @inbounds grid.Δzᵃᵃᶠ[k] - -## XRegularRG +@inline Δxᶠᵃᵃ(i, j, k, grid::RG) = @inbounds grid.Δxᶠᵃᵃ[i] +@inline Δxᶜᵃᵃ(i, j, k, grid::RG) = @inbounds grid.Δxᶜᵃᵃ[i] +@inline Δyᵃᶠᵃ(i, j, k, grid::RG) = @inbounds grid.Δyᵃᶠᵃ[j] +@inline Δyᵃᶜᵃ(i, j, k, grid::RG) = @inbounds grid.Δyᵃᶜᵃ[j] @inline Δxᶠᵃᵃ(i, j, k, grid::RGX) = grid.Δxᶠᵃᵃ @inline Δxᶜᵃᵃ(i, j, k, grid::RGX) = grid.Δxᶜᵃᵃ - -@inline Δxᶜᵃᶜ(i, j, k, grid::RGX) = grid.Δxᶜᵃᵃ -@inline Δxᶠᵃᶜ(i, j, k, grid::RGX) = grid.Δxᶠᵃᵃ -@inline Δxᶜᵃᶠ(i, j, k, grid::RGX) = grid.Δxᶜᵃᵃ - -## YRegularRG - @inline Δyᵃᶠᵃ(i, j, k, grid::RGY) = grid.Δyᵃᶠᵃ @inline Δyᵃᶜᵃ(i, j, k, grid::RGY) = grid.Δyᵃᶜᵃ -@inline Δyᶜᵃᶜ(i, j, k, grid::RGY) = grid.Δyᵃᶜᵃ -@inline Δyᶠᵃᶜ(i, j, k, grid::RGY) = grid.Δyᵃᶜᵃ -@inline Δyᶜᵃᶠ(i, j, k, grid::RGY) = grid.Δyᵃᶜᵃ - -## ZRegularRG - -@inline Δzᵃᵃᶠ(i, j, k, grid::RGZ) = grid.Δzᵃᵃᶠ -@inline Δzᵃᵃᶜ(i, j, k, grid::RGZ) = grid.Δzᵃᵃᶜ - -@inline Δzᶜᵃᶜ(i, j, k, grid::RGZ) = grid.Δzᵃᵃᶜ -@inline Δzᶠᵃᶜ(i, j, k, grid::RGZ) = grid.Δzᵃᵃᶜ -@inline Δzᶜᵃᶠ(i, j, k, grid::RGZ) = grid.Δzᵃᵃᶠ - ##### ##### LatitudeLongitudeGrid ##### ## Pre computed metrics -@inline Δxᶜᶠᵃ(i, j, k, grid::LLG) = @inbounds grid.Δxᶜᶠᵃ[i, j] -@inline Δxᶠᶜᵃ(i, j, k, grid::LLG) = @inbounds grid.Δxᶠᶜᵃ[i, j] -@inline Δxᶠᶠᵃ(i, j, k, grid::LLG) = @inbounds grid.Δxᶠᶠᵃ[i, j] -@inline Δxᶜᶜᵃ(i, j, k, grid::LLG) = @inbounds grid.Δxᶜᶜᵃ[i, j] - -@inline Δyᶜᶠᵃ(i, j, k, grid::LLG) = @inbounds grid.Δyᶜᶠᵃ[j] -@inline Δyᶠᶜᵃ(i, j, k, grid::LLG) = @inbounds grid.Δyᶠᶜᵃ[j] -@inline Δyᶜᶜᵃ(i, j, k, grid::LLG) = Δyᶠᶜᵃ(i, j, k, grid) -@inline Δyᶠᶠᵃ(i, j, k, grid::LLG) = Δyᶜᶠᵃ(i, j, k, grid) - -@inline Δzᵃᵃᶠ(i, j, k, grid::LLG) = @inbounds grid.Δzᵃᵃᶠ[k] -@inline Δzᵃᵃᶜ(i, j, k, grid::LLG) = @inbounds grid.Δzᵃᵃᶜ[k] - -### XRegularLLG with pre-computed metrics - +@inline Δxᶜᶠᵃ(i, j, k, grid::LLG) = @inbounds grid.Δxᶜᶠᵃ[i, j] +@inline Δxᶠᶜᵃ(i, j, k, grid::LLG) = @inbounds grid.Δxᶠᶜᵃ[i, j] +@inline Δxᶠᶠᵃ(i, j, k, grid::LLG) = @inbounds grid.Δxᶠᶠᵃ[i, j] +@inline Δxᶜᶜᵃ(i, j, k, grid::LLG) = @inbounds grid.Δxᶜᶜᵃ[i, j] @inline Δxᶠᶜᵃ(i, j, k, grid::LLGX) = @inbounds grid.Δxᶠᶜᵃ[j] @inline Δxᶜᶠᵃ(i, j, k, grid::LLGX) = @inbounds grid.Δxᶜᶠᵃ[j] @inline Δxᶠᶠᵃ(i, j, k, grid::LLGX) = @inbounds grid.Δxᶠᶠᵃ[j] @inline Δxᶜᶜᵃ(i, j, k, grid::LLGX) = @inbounds grid.Δxᶜᶜᵃ[j] -### YRegularLLG with pre-computed metrics - +@inline Δyᶜᶠᵃ(i, j, k, grid::LLG) = @inbounds grid.Δyᶜᶠᵃ[j] +@inline Δyᶠᶜᵃ(i, j, k, grid::LLG) = @inbounds grid.Δyᶠᶜᵃ[j] @inline Δyᶜᶠᵃ(i, j, k, grid::LLGY) = grid.Δyᶜᶠᵃ @inline Δyᶠᶜᵃ(i, j, k, grid::LLGY) = grid.Δyᶠᶜᵃ - -### ZRegularLLG with pre-computed metrics - -@inline Δzᵃᵃᶠ(i, j, k, grid::LLGZ) = grid.Δzᵃᵃᶠ -@inline Δzᵃᵃᶜ(i, j, k, grid::LLGZ) = grid.Δzᵃᵃᶜ +@inline Δyᶜᶜᵃ(i, j, k, grid::LLG) = Δyᶠᶜᵃ(i, j, k, grid) +@inline Δyᶠᶠᵃ(i, j, k, grid::LLG) = Δyᶜᶠᵃ(i, j, k, grid) ## On the fly metrics -@inline Δxᶠᶜᵃ(i, j, k, grid::LLGF) = @inbounds grid.radius * deg2rad(grid.Δλᶠᵃᵃ[i]) * hack_cosd(grid.φᵃᶜᵃ[j]) -@inline Δxᶜᶠᵃ(i, j, k, grid::LLGF) = @inbounds grid.radius * deg2rad(grid.Δλᶜᵃᵃ[i]) * hack_cosd(grid.φᵃᶠᵃ[j]) -@inline Δxᶠᶠᵃ(i, j, k, grid::LLGF) = @inbounds grid.radius * deg2rad(grid.Δλᶠᵃᵃ[i]) * hack_cosd(grid.φᵃᶠᵃ[j]) -@inline Δxᶜᶜᵃ(i, j, k, grid::LLGF) = @inbounds grid.radius * deg2rad(grid.Δλᶜᵃᵃ[i]) * hack_cosd(grid.φᵃᶜᵃ[j]) +@inline Δxᶠᶜᵃ(i, j, k, grid::LLGF) = @inbounds grid.radius * deg2rad(grid.Δλᶠᵃᵃ[i]) * hack_cosd(grid.φᵃᶜᵃ[j]) +@inline Δxᶜᶠᵃ(i, j, k, grid::LLGF) = @inbounds grid.radius * deg2rad(grid.Δλᶜᵃᵃ[i]) * hack_cosd(grid.φᵃᶠᵃ[j]) +@inline Δxᶠᶠᵃ(i, j, k, grid::LLGF) = @inbounds grid.radius * deg2rad(grid.Δλᶠᵃᵃ[i]) * hack_cosd(grid.φᵃᶠᵃ[j]) +@inline Δxᶜᶜᵃ(i, j, k, grid::LLGF) = @inbounds grid.radius * deg2rad(grid.Δλᶜᵃᵃ[i]) * hack_cosd(grid.φᵃᶜᵃ[j]) +@inline Δxᶠᶜᵃ(i, j, k, grid::LLGFX) = @inbounds grid.radius * deg2rad(grid.Δλᶠᵃᵃ) * hack_cosd(grid.φᵃᶜᵃ[j]) +@inline Δxᶜᶠᵃ(i, j, k, grid::LLGFX) = @inbounds grid.radius * deg2rad(grid.Δλᶜᵃᵃ) * hack_cosd(grid.φᵃᶠᵃ[j]) +@inline Δxᶠᶠᵃ(i, j, k, grid::LLGFX) = @inbounds grid.radius * deg2rad(grid.Δλᶠᵃᵃ) * hack_cosd(grid.φᵃᶠᵃ[j]) +@inline Δxᶜᶜᵃ(i, j, k, grid::LLGFX) = @inbounds grid.radius * deg2rad(grid.Δλᶜᵃᵃ) * hack_cosd(grid.φᵃᶜᵃ[j]) @inline Δyᶜᶠᵃ(i, j, k, grid::LLGF) = @inbounds grid.radius * deg2rad(grid.Δφᵃᶠᵃ[j]) @inline Δyᶠᶜᵃ(i, j, k, grid::LLGF) = @inbounds grid.radius * deg2rad(grid.Δφᵃᶜᵃ[j]) - -### XRegularLLG with on-the-fly metrics - -@inline Δxᶠᶜᵃ(i, j, k, grid::LLGFX) = @inbounds grid.radius * deg2rad(grid.Δλᶠᵃᵃ) * hack_cosd(grid.φᵃᶜᵃ[j]) -@inline Δxᶜᶠᵃ(i, j, k, grid::LLGFX) = @inbounds grid.radius * deg2rad(grid.Δλᶜᵃᵃ) * hack_cosd(grid.φᵃᶠᵃ[j]) -@inline Δxᶠᶠᵃ(i, j, k, grid::LLGFX) = @inbounds grid.radius * deg2rad(grid.Δλᶠᵃᵃ) * hack_cosd(grid.φᵃᶠᵃ[j]) -@inline Δxᶜᶜᵃ(i, j, k, grid::LLGFX) = @inbounds grid.radius * deg2rad(grid.Δλᶜᵃᵃ) * hack_cosd(grid.φᵃᶜᵃ[j]) - -### YRegularLLG with on-the-fly metrics - @inline Δyᶜᶠᵃ(i, j, k, grid::LLGFY) = grid.radius * deg2rad(grid.Δφᵃᶠᵃ) @inline Δyᶠᶜᵃ(i, j, k, grid::LLGFY) = grid.radius * deg2rad(grid.Δφᵃᶜᵃ) @@ -185,19 +167,13 @@ end @inline Δyᶜᶠᵃ(i, j, k, grid::OSSG) = @inbounds grid.Δyᶜᶠᵃ[i, j] @inline Δyᶠᶠᵃ(i, j, k, grid::OSSG) = @inbounds grid.Δyᶠᶠᵃ[i, j] -@inline Δzᵃᵃᶜ(i, j, k, grid::OSSG) = @inbounds grid.Δzᵃᵃᶜ[k] -@inline Δzᵃᵃᶠ(i, j, k, grid::OSSG) = @inbounds grid.Δzᵃᵃᶠ[k] - -@inline Δzᵃᵃᶜ(i, j, k, grid::OSSGZ) = grid.Δzᵃᵃᶜ -@inline Δzᵃᵃᶠ(i, j, k, grid::OSSGZ) = grid.Δzᵃᵃᶠ - ##### ##### ##### Areas!! ##### ##### -for LX in (:ᶜ, :ᶠ, :ᵃ), LY in (:ᶜ, :ᶠ, :ᵃ) +for LX in (:ᶜ, :ᶠ), LY in (:ᶜ, :ᶠ) x_spacing_2D = Symbol(:Δx, LX, LY, :ᵃ) y_spacing_2D = Symbol(:Δy, LX, LY, :ᵃ) @@ -253,16 +229,14 @@ end ##### ##### -for LX in (:ᶠ, :ᶜ), LY in (:ᶠ, :ᶜ), LZ in (:ᶠ, :ᶜ) - - volume = Symbol(:V, LX, LY, LZ) - z_area = Symbol(:Az, LX, LY, LZ) - z_spacing = Symbol(:Δz, LX, LY, LZ) - - @eval begin - @inline $volume(i, j, k, grid) = $z_area(i, j, k, grid) * $z_spacing(i, j, k, grid) - end -end +@inline Vᶜᶜᶜ(i, j, k, grid) = Azᶜᶜᶜ(i, j, k, grid) * Δzᶜᶜᶜ(i, j, k, grid) +@inline Vᶠᶜᶜ(i, j, k, grid) = Azᶠᶜᶜ(i, j, k, grid) * Δzᶠᶜᶜ(i, j, k, grid) +@inline Vᶜᶠᶜ(i, j, k, grid) = Azᶜᶠᶜ(i, j, k, grid) * Δzᶜᶠᶜ(i, j, k, grid) +@inline Vᶜᶜᶠ(i, j, k, grid) = Azᶜᶜᶠ(i, j, k, grid) * Δzᶜᶜᶠ(i, j, k, grid) +@inline Vᶠᶠᶜ(i, j, k, grid) = Azᶠᶠᶜ(i, j, k, grid) * Δzᶠᶠᶜ(i, j, k, grid) +@inline Vᶠᶜᶠ(i, j, k, grid) = Azᶠᶜᶠ(i, j, k, grid) * Δzᶠᶜᶠ(i, j, k, grid) +@inline Vᶜᶠᶠ(i, j, k, grid) = Azᶜᶠᶠ(i, j, k, grid) * Δzᶜᶠᶠ(i, j, k, grid) +@inline Vᶠᶠᶠ(i, j, k, grid) = Azᶠᶠᶠ(i, j, k, grid) * Δzᶠᶠᶠ(i, j, k, grid) ##### ##### Generic functions for specified locations @@ -274,9 +248,9 @@ end location_code(LX, LY, LZ) = Symbol(interpolation_code(LX), interpolation_code(LY), interpolation_code(LZ)) -for LX in (:Center, :Face, :Nothing) - for LY in (:Center, :Face, :Nothing) - for LZ in (:Center, :Face, :Nothing) +for LX in (:Center, :Face) + for LY in (:Center, :Face) + for LZ in (:Center, :Face) LXe = @eval $LX LYe = @eval $LY LZe = @eval $LZ @@ -297,3 +271,4 @@ for LX in (:Center, :Face, :Nothing) end end end + diff --git a/src/Operators/topology_aware_operators.jl b/src/Operators/topology_aware_operators.jl deleted file mode 100644 index cab68c3b6e..0000000000 --- a/src/Operators/topology_aware_operators.jl +++ /dev/null @@ -1,67 +0,0 @@ -using Oceananigans.Grids: AbstractUnderlyingGrid - -const AGXB = AbstractUnderlyingGrid{FT, Bounded} where FT -const AGXP = AbstractUnderlyingGrid{FT, Periodic} where FT -const AGXR = AbstractUnderlyingGrid{FT, RightConnected} where FT -const AGXL = AbstractUnderlyingGrid{FT, LeftConnected} where FT - -const AGYB = AbstractUnderlyingGrid{FT, <:Any, Bounded} where FT -const AGYP = AbstractUnderlyingGrid{FT, <:Any, Periodic} where FT -const AGYR = AbstractUnderlyingGrid{FT, <:Any, RightConnected} where FT -const AGYL = AbstractUnderlyingGrid{FT, <:Any, LeftConnected} where FT - -# Topology-aware Operators with the following convention: -# -# `δxTᶠᵃᵃ` : Hardcodes `Noflux` or `Periodic` boundary conditions for a (Center, Center, Center) function `f` in the x-direction. -# `δyTᵃᶠᵃ` : Hardcodes `Noflux` or `Periodic` boundary conditions for a (Center, Center, Center) function `f` in the y-direction -# -# `δxTᶜᵃᵃ` : Hardcodes `NoPenetration` or `Periodic` boundary conditions for a (Face, Center, Center) function `U` in x direction -# `δyTᵃᶜᵃ` : Hardcodes `NoPenetration` or `Periodic` boundary conditions for a (Center, Face, Center) function `V` in y direction -# -# Note: The naming convention is that `T` denotes a topology-aware operator. So `δxTᶠᵃᵃ` is the topology-aware version of `δxᶠᵃᵃ`. - -# Fallback - -@inline δxTᶠᵃᵃ(i, j, k, grid, f::Function, args...) = δxᶠᵃᵃ(i, j, k, grid, f, args...) -@inline δyTᵃᶠᵃ(i, j, k, grid, f::Function, args...) = δyᵃᶠᵃ(i, j, k, grid, f, args...) -@inline δxTᶜᵃᵃ(i, j, k, grid, f::Function, args...) = δxᶜᵃᵃ(i, j, k, grid, f, args...) -@inline δyTᵃᶜᵃ(i, j, k, grid, f::Function, args...) = δyᵃᶜᵃ(i, j, k, grid, f, args...) - -# Enforce Periodic conditions - -@inline δxTᶠᵃᵃ(i, j, k, grid::AGXP, f::Function, args...) = ifelse(i == 1, f(1, j, k, grid, args...) - f(grid.Nx, j, k, grid, args...), δxᶠᵃᵃ(i, j, k, grid, f, args...)) -@inline δyTᵃᶠᵃ(i, j, k, grid::AGYP, f::Function, args...) = ifelse(j == 1, f(i, 1, k, grid, args...) - f(i, grid.Ny, k, grid, args...), δyᵃᶠᵃ(i, j, k, grid, f, args...)) - -@inline δxTᶜᵃᵃ(i, j, k, grid::AGXP, f::Function, args...) = ifelse(i == grid.Nx, f(1, j, k, grid, args...) - f(grid.Nx, j, k, grid, args...), δxᶜᵃᵃ(i, j, k, grid, f, args...)) -@inline δyTᵃᶜᵃ(i, j, k, grid::AGYP, f::Function, args...) = ifelse(j == grid.Ny, f(i, 1, k, grid, args...) - f(i, grid.Ny, k, grid, args...), δyᵃᶜᵃ(i, j, k, grid, f, args...)) - -# Enforce NoFlux conditions - -@inline δxTᶠᵃᵃ(i, j, k, grid::AGXB{FT}, f::Function, args...) where FT = ifelse(i == 1, zero(FT), δxᶠᵃᵃ(i, j, k, grid, f, args...)) -@inline δyTᵃᶠᵃ(i, j, k, grid::AGYB{FT}, f::Function, args...) where FT = ifelse(j == 1, zero(FT), δyᵃᶠᵃ(i, j, k, grid, f, args...)) - -@inline δxTᶠᵃᵃ(i, j, k, grid::AGXR{FT}, f::Function, args...) where FT = ifelse(i == 1, zero(FT), δxᶠᵃᵃ(i, j, k, grid, f, args...)) -@inline δyTᵃᶠᵃ(i, j, k, grid::AGYR{FT}, f::Function, args...) where FT = ifelse(j == 1, zero(FT), δyᵃᶠᵃ(i, j, k, grid, f, args...)) - -# Enforce Impenetrability conditions - -@inline δxTᶜᵃᵃ(i, j, k, grid::AGXB, f::Function, args...) = - ifelse(i == grid.Nx, - f(i, j, k, grid, args...), - ifelse(i == 1, f(2, j, k, grid, args...), - δxᶜᵃᵃ(i, j, k, grid, f, args...))) - -@inline δyTᵃᶜᵃ(i, j, k, grid::AGYB, f::Function, args...) = - ifelse(j == grid.Ny, - f(i, j, k, grid, args...), - ifelse(j == 1, f(i, 2, k, grid, args...), - δyᵃᶜᵃ(i, j, k, grid, f, args...))) - -@inline δxTᶜᵃᵃ(i, j, k, grid::AGXL, f::Function, args...) = ifelse(i == grid.Nx, - f(i, j, k, grid, args...), δxᶜᵃᵃ(i, j, k, grid, f, args...)) -@inline δyTᵃᶜᵃ(i, j, k, grid::AGYL, f::Function, args...) = ifelse(j == grid.Ny, - f(i, j, k, grid, args...), δyᵃᶜᵃ(i, j, k, grid, f, args...)) - -@inline δxTᶜᵃᵃ(i, j, k, grid::AGXR, f::Function, args...) = ifelse(i == 1, f(2, j, k, grid, args...), δxᶜᵃᵃ(i, j, k, grid, f, args...)) -@inline δyTᵃᶜᵃ(i, j, k, grid::AGYR, f::Function, args...) = ifelse(j == 1, f(i, 2, k, grid, args...), δyᵃᶜᵃ(i, j, k, grid, f, args...)) - -# Derivative operators - -@inline ∂xTᶠᶜᶠ(i, j, k, grid, f::Function, args...) = δxTᶠᵃᵃ(i, j, k, grid, f, args...) / Δxᶠᶜᶠ(i, j, k, grid) -@inline ∂yTᶜᶠᶠ(i, j, k, grid, f::Function, args...) = δyTᵃᶠᵃ(i, j, k, grid, f, args...) / Δyᶜᶠᶠ(i, j, k, grid) diff --git a/src/Operators/vector_rotation_operators.jl b/src/Operators/vector_rotation_operators.jl index f42dc70466..7ac9da502c 100644 --- a/src/Operators/vector_rotation_operators.jl +++ b/src/Operators/vector_rotation_operators.jl @@ -57,31 +57,30 @@ _intrinsic_ coordinate systems are equivalent. However, for other grids (e.g., f # 2D vectors @inline function intrinsic_vector(i, j, k, grid::OrthogonalSphericalShellGrid, uₑ, vₑ) - φᶠᶠᵃ⁺⁺ = φnode(i+1, j+1, 1, grid, Face(), Face(), Center()) - φᶠᶠᵃ⁺⁻ = φnode(i+1, j, 1, grid, Face(), Face(), Center()) - φᶠᶠᵃ⁻⁺ = φnode(i, j+1, 1, grid, Face(), Face(), Center()) - φᶠᶠᵃ⁻⁻ = φnode(i, j, 1, grid, Face(), Face(), Center()) - - Δyᶠᶜᵃ⁺ = Δyᶠᶜᶜ(i+1, j, 1, grid) - Δyᶠᶜᵃ⁻ = Δyᶠᶜᶜ(i, j, 1, grid) - Δxᶜᶠᵃ⁺ = Δxᶜᶠᶜ(i, j+1, 1, grid) - Δxᶜᶠᵃ⁻ = Δxᶜᶠᶜ(i, j, 1, grid) + φᶜᶠᵃ₊ = φnode(i, j+1, 1, grid, Center(), Face(), Center()) + φᶜᶠᵃ₋ = φnode(i, j, 1, grid, Center(), Face(), Center()) + Δyᶜᶜᵃ = Δyᶜᶜᶜ(i, j, 1, grid) # θᵢ is the rotation angle between intrinsic and extrinsic reference frame - Rcosθ = (deg2rad(φᶠᶠᵃ⁺⁺ - φᶠᶠᵃ⁺⁻) / Δyᶠᶜᵃ⁺ + deg2rad(φᶠᶠᵃ⁻⁺ - φᶠᶠᵃ⁻⁻) / Δyᶠᶜᵃ⁻) / 2 - Rsinθ = - (deg2rad(φᶠᶠᵃ⁺⁺ - φᶠᶠᵃ⁻⁺) / Δxᶜᶠᵃ⁺ + deg2rad(φᶠᶠᵃ⁺⁻ - φᶠᶠᵃ⁻⁻) / Δxᶜᶠᵃ⁻) / 2 + Rcosθᵢ = deg2rad(φᶜᶠᵃ₊ - φᶜᶠᵃ₋) / Δyᶜᶜᵃ + + φᶠᶜᵃ₊ = φnode(i+1, j, 1, grid, Face(), Center(), Center()) + φᶠᶜᵃ₋ = φnode(i, j, 1, grid, Face(), Center(), Center()) + Δxᶜᶜᵃ = Δxᶜᶜᶜ(i, j, 1, grid) + + Rsinθᵢ = - deg2rad(φᶠᶜᵃ₊ - φᶠᶜᵃ₋) / Δxᶜᶜᵃ # Normalization for the rotation angles - R = sqrt(Rcosθ^2 + Rsinθ^2) + Rᵢ = sqrt(Rcosθᵢ^2 + Rsinθᵢ^2) u = getvalue(uₑ, i, j, k, grid) v = getvalue(vₑ, i, j, k, grid) - cosθ = Rcosθ / R - sinθ = Rsinθ / R + cosθᵢ = Rcosθᵢ / Rᵢ + sinθᵢ = Rsinθᵢ / Rᵢ - uᵢ = u * cosθ + v * sinθ - vᵢ = - u * sinθ + v * cosθ + uᵢ = u * cosθᵢ + v * sinθᵢ + vᵢ = - u * sinθᵢ + v * cosθᵢ return uᵢ, vᵢ end @@ -98,31 +97,30 @@ end # 2D vectors @inline function extrinsic_vector(i, j, k, grid::OrthogonalSphericalShellGrid, uᵢ, vᵢ) - φᶠᶠᵃ⁺⁺ = φnode(i+1, j+1, 1, grid, Face(), Face(), Center()) - φᶠᶠᵃ⁺⁻ = φnode(i+1, j, 1, grid, Face(), Face(), Center()) - φᶠᶠᵃ⁻⁺ = φnode(i, j+1, 1, grid, Face(), Face(), Center()) - φᶠᶠᵃ⁻⁻ = φnode(i, j, 1, grid, Face(), Face(), Center()) + φᶜᶠᵃ₊ = φnode(i, j+1, 1, grid, Center(), Face(), Center()) + φᶜᶠᵃ₋ = φnode(i, j, 1, grid, Center(), Face(), Center()) + Δyᶜᶜᵃ = Δyᶜᶜᶜ(i, j, 1, grid) - Δyᶠᶜᵃ⁺ = Δyᶠᶜᶜ(i+1, j, 1, grid) - Δyᶠᶜᵃ⁻ = Δyᶠᶜᶜ(i, j, 1, grid) - Δxᶜᶠᵃ⁺ = Δxᶜᶠᶜ(i, j+1, 1, grid) - Δxᶜᶠᵃ⁻ = Δxᶜᶠᶜ(i, j, 1, grid) + # θₑ is the rotation angle between intrinsic and extrinsic reference frame + Rcosθₑ = deg2rad(φᶜᶠᵃ₊ - φᶜᶠᵃ₋) / Δyᶜᶜᵃ - # θᵢ is the rotation angle between intrinsic and extrinsic reference frame - Rcosθ = (deg2rad(φᶠᶠᵃ⁺⁺ - φᶠᶠᵃ⁺⁻) / Δyᶠᶜᵃ⁺ + deg2rad(φᶠᶠᵃ⁻⁺ - φᶠᶠᵃ⁻⁻) / Δyᶠᶜᵃ⁻) / 2 - Rsinθ = - (deg2rad(φᶠᶠᵃ⁺⁺ - φᶠᶠᵃ⁻⁺) / Δxᶜᶠᵃ⁺ + deg2rad(φᶠᶠᵃ⁺⁻ - φᶠᶠᵃ⁻⁻) / Δxᶜᶠᵃ⁻) / 2 + φᶠᶜᵃ₊ = φnode(i+1, j, 1, grid, Face(), Center(), Center()) + φᶠᶜᵃ₋ = φnode(i, j, 1, grid, Face(), Center(), Center()) + Δxᶜᶜᵃ = Δxᶜᶜᶜ(i, j, 1, grid) + + Rsinθₑ = - deg2rad(φᶠᶜᵃ₊ - φᶠᶜᵃ₋) / Δxᶜᶜᵃ # Normalization for the rotation angles - R = sqrt(Rcosθ^2 + Rsinθ^2) + Rₑ = sqrt(Rcosθₑ^2 + Rsinθₑ^2) u = getvalue(uᵢ, i, j, k, grid) v = getvalue(vᵢ, i, j, k, grid) - cosθ = Rcosθ / R - sinθ = Rsinθ / R + cosθₑ = Rcosθₑ / Rₑ + sinθₑ = Rsinθₑ / Rₑ - uₑ = u * cosθ - v * sinθ - vₑ = u * sinθ + v * cosθ + uₑ = u * cosθₑ - v * sinθₑ + vₑ = u * sinθₑ + v * cosθₑ return uₑ, vₑ end diff --git a/src/Operators/vorticity_operators.jl b/src/Operators/vorticity_operators.jl index bf1f95908d..170a013a28 100644 --- a/src/Operators/vorticity_operators.jl +++ b/src/Operators/vorticity_operators.jl @@ -10,12 +10,6 @@ The vertical vorticity associated with horizontal velocities ``u`` and ``v``. """ @inline ζ₃ᶠᶠᶜ(i, j, k, grid, u, v) = Γᶠᶠᶜ(i, j, k, grid, u, v) / Azᶠᶠᶜ(i, j, k, grid) -# South-west, south-east, north-west, north-east corners -@inline on_south_west_corner(i, j, grid) = (i == 1) & (j == 1) -@inline on_south_east_corner(i, j, grid) = (i == grid.Nx+1) & (j == 1) -@inline on_north_east_corner(i, j, grid) = (i == grid.Nx+1) & (j == grid.Ny+1) -@inline on_north_west_corner(i, j, grid) = (i == 1) & (j == grid.Ny+1) - ##### ##### Vertical circulation at the corners of the cubed sphere needs to treated in a special manner. ##### See: https://github.com/CliMA/Oceananigans.jl/issues/1584 @@ -27,13 +21,24 @@ The vertical vorticity associated with horizontal velocities ``u`` and ``v``. The vertical circulation associated with horizontal velocities ``u`` and ``v``. """ @inline function Γᶠᶠᶜ(i, j, k, grid::ConformalCubedSpherePanel, u, v) - Hx, Hy = grid.Hx, grid.Hy - Γ = ifelse(on_south_west_corner(i, j, grid) | on_north_west_corner(i, j, grid), - Δy_qᶜᶠᶜ(i, j, k, grid, v) - Δx_qᶠᶜᶜ(i, j, k, grid, u) + Δx_qᶠᶜᶜ(i, j-1, k, grid, u), - ifelse(on_south_east_corner(i, j, grid) | on_north_east_corner(i, j, grid), - - Δy_qᶜᶠᶜ(i-1, j, k, grid, v) + Δx_qᶠᶜᶜ(i, j-1, k, grid, u) - Δx_qᶠᶜᶜ(i, j, k, grid, u), - δxᶠᶠᶜ(i, j, k, grid, Δy_qᶜᶠᶜ, v) - δyᶠᶠᶜ(i, j, k, grid, Δx_qᶠᶜᶜ, u) - ) - ) - return Γ + # South-west corner + if i == 1 && j == 1 + return Δy_qᶜᶠᶜ(i, j, k, grid, v) - Δx_qᶠᶜᶜ(i, j, k, grid, u) + Δx_qᶠᶜᶜ(i, j-1, k, grid, u) + + # South-east corner + elseif i == grid.Nx+1 && j == 1 + return - Δy_qᶜᶠᶜ(i-1, j, k, grid, v) - Δx_qᶠᶜᶜ(i, j, k, grid, u) + Δx_qᶠᶜᶜ(i, j-1, k, grid, u) + + # North-west corner + elseif i == 1 && j == grid.Ny+1 + return Δy_qᶜᶠᶜ(i, j, k, grid, v) - Δx_qᶠᶜᶜ(i, j, k, grid, u) + Δx_qᶠᶜᶜ(i, j-1, k, grid, u) + + # North-east corner + elseif i == grid.Nx+1 && j == grid.Ny+1 + return - Δy_qᶜᶠᶜ(i-1, j, k, grid, v) - Δx_qᶠᶜᶜ(i, j, k, grid, u) + Δx_qᶠᶜᶜ(i, j-1, k, grid, u) + + # Not a corner + else + return δxᶠᵃᵃ(i, j, k, grid, Δy_qᶜᶠᶜ, v) - δyᵃᶠᵃ(i, j, k, grid, Δx_qᶠᶜᶜ, u) + end end diff --git a/src/OutputReaders/field_dataset.jl b/src/OutputReaders/field_dataset.jl index 9dbe0433f4..dc6072eb6d 100644 --- a/src/OutputReaders/field_dataset.jl +++ b/src/OutputReaders/field_dataset.jl @@ -1,8 +1,7 @@ -struct FieldDataset{F, M, P, KW} - fields :: F - metadata :: M - filepath :: P - reader_kw :: KW +struct FieldDataset{F, M, P} + fields :: F + metadata :: M + filepath :: P end """ @@ -23,24 +22,17 @@ linearly. `file["metadata"]`. - `grid`: May be specified to override the grid used in the JLD2 file. - -- `reader_kw`: A named tuple or dictionary of keyword arguments to pass to the reader - (currently only JLD2) to be used when opening files. """ function FieldDataset(filepath; - architecture = CPU(), - grid = nothing, - backend = InMemory(), - metadata_paths = ["metadata"], - reader_kw = NamedTuple()) + architecture=CPU(), grid=nothing, backend=InMemory(), metadata_paths=["metadata"]) - file = jldopen(filepath; reader_kw...) + file = jldopen(filepath) field_names = keys(file["timeseries"]) filter!(k -> k != "t", field_names) # Time is not a field. ds = Dict{String, FieldTimeSeries}( - name => FieldTimeSeries(filepath, name; architecture, backend, grid, reader_kw) + name => FieldTimeSeries(filepath, name; architecture, backend, grid) for name in field_names ) @@ -52,7 +44,7 @@ function FieldDataset(filepath; close(file) - return FieldDataset(ds, metadata, abspath(filepath), reader_kw) + return FieldDataset(ds, metadata, abspath(filepath)) end Base.getindex(fds::FieldDataset, inds...) = Base.getindex(fds.fields, inds...) diff --git a/src/OutputReaders/field_time_series.jl b/src/OutputReaders/field_time_series.jl index 19a1bdb2c8..7841daf5fb 100644 --- a/src/OutputReaders/field_time_series.jl +++ b/src/OutputReaders/field_time_series.jl @@ -85,7 +85,7 @@ period = t[end] - t[1] + Δt """ struct Cyclical{FT} period :: FT -end +end Cyclical() = Cyclical(nothing) @@ -164,7 +164,7 @@ Nt = 5 backend = InMemory(4, 3) # so we have (4, 5, 1) n = 1 # so, the right answer is m̃ = 3 m = 1 - (4 - 1) # = -2 -m̃ = mod1(-2, 5) # = 3 ✓ +m̃ = mod1(-2, 5) # = 3 ✓ ``` # Another shifting + wrapping example @@ -213,28 +213,26 @@ Base.length(backend::PartlyInMemory) = backend.length ##### FieldTimeSeries ##### -mutable struct FieldTimeSeries{LX, LY, LZ, TI, K, I, D, G, ET, B, χ, P, N, KW} <: AbstractField{LX, LY, LZ, G, ET, 4} - data :: D - grid :: G - backend :: K +mutable struct FieldTimeSeries{LX, LY, LZ, TI, K, I, D, G, ET, B, χ, P, N} <: AbstractField{LX, LY, LZ, G, ET, 4} + data :: D + grid :: G + backend :: K boundary_conditions :: B - indices :: I - times :: χ - path :: P - name :: N - time_indexing :: TI - reader_kw :: KW - + indices :: I + times :: χ + path :: P + name :: N + time_indexing :: TI + function FieldTimeSeries{LX, LY, LZ}(data::D, grid::G, backend::K, bcs::B, - indices::I, + indices::I, times, path, name, - time_indexing, - reader_kw) where {LX, LY, LZ, K, D, G, B, I} + time_indexing) where {LX, LY, LZ, K, D, G, B, I} ET = eltype(data) @@ -252,7 +250,7 @@ mutable struct FieldTimeSeries{LX, LY, LZ, TI, K, I, D, G, ET, B, χ, P, N, KW} times = on_architecture(architecture(grid), times) end - + if time_indexing isa Cyclical{Nothing} # we have to infer the period Δt = @allowscalar times[end] - times[end-1] period = @allowscalar times[end] - times[1] + Δt @@ -263,25 +261,23 @@ mutable struct FieldTimeSeries{LX, LY, LZ, TI, K, I, D, G, ET, B, χ, P, N, KW} TI = typeof(time_indexing) P = typeof(path) N = typeof(name) - KW = typeof(reader_kw) - return new{LX, LY, LZ, TI, K, I, D, G, ET, B, χ, P, N, KW}(data, grid, backend, bcs, - indices, times, path, name, - time_indexing, reader_kw) + return new{LX, LY, LZ, TI, K, I, D, G, ET, B, χ, P, N}(data, grid, backend, bcs, + indices, times, path, name, + time_indexing) end end -on_architecture(to, fts::FieldTimeSeries{LX, LY, LZ}) where {LX, LY, LZ} = +on_architecture(to, fts::FieldTimeSeries{LX, LY, LZ}) where {LX, LY, LZ} = FieldTimeSeries{LX, LY, LZ}(on_architecture(to, fts.data), on_architecture(to, fts.grid), on_architecture(to, fts.backend), on_architecture(to, fts.bcs), - on_architecture(to, fts.indices), + on_architecture(to, fts.indices), on_architecture(to, fts.times), on_architecture(to, fts.path), on_architecture(to, fts.name), - on_architecture(to, fts.time_indexing), - on_architecture(to, fts.reader_kw)) + on_architecture(to, fts.time_indexing)) ##### ##### Minimal implementation of FieldTimeSeries for use in GPU kernels @@ -294,7 +290,7 @@ struct GPUAdaptedFieldTimeSeries{LX, LY, LZ, TI, K, ET, D, χ} <: AbstractField{ times :: χ backend :: K time_indexing :: TI - + function GPUAdaptedFieldTimeSeries{LX, LY, LZ}(data::D, times::χ, backend::K, @@ -317,7 +313,7 @@ const FTS{LX, LY, LZ, TI, K} = FieldTimeSeries{LX, LY, LZ, TI, K} w const GPUFTS{LX, LY, LZ, TI, K} = GPUAdaptedFieldTimeSeries{LX, LY, LZ, TI, K} where {LX, LY, LZ, TI, K} const FlavorOfFTS{LX, LY, LZ, TI, K} = Union{GPUFTS{LX, LY, LZ, TI, K}, - FTS{LX, LY, LZ, TI, K}} where {LX, LY, LZ, TI, K} + FTS{LX, LY, LZ, TI, K}} where {LX, LY, LZ, TI, K} const InMemoryFTS = FlavorOfFTS{<:Any, <:Any, <:Any, <:Any, <:AbstractInMemoryBackend} const OnDiskFTS = FlavorOfFTS{<:Any, <:Any, <:Any, <:Any, <:OnDisk} @@ -349,7 +345,7 @@ instantiate(T::Type) = T() new_data(FT, grid, loc, indices, ::Nothing) = nothing # Apparently, not explicitly specifying Int64 in here makes this function -# fail on x86 processors where `Int` is implied to be `Int32` +# fail on x86 processors where `Int` is implied to be `Int32` # see ClimaOcean commit 3c47d887659d81e0caed6c9df41b7438e1f1cd52 at https://github.com/CliMA/ClimaOcean.jl/actions/runs/8804916198/job/24166354095) function new_data(FT, grid, loc, indices, Nt::Union{Int, Int64}) space_size = total_size(grid, loc, indices) @@ -364,13 +360,12 @@ time_indices_length(backend::PartlyInMemory, times) = length(backend) time_indices_length(::OnDisk, times) = nothing function FieldTimeSeries(loc, grid, times=(); - indices = (:, :, :), + indices = (:, :, :), backend = InMemory(), - path = nothing, + path = nothing, name = nothing, time_indexing = Linear(), - boundary_conditions = nothing, - reader_kw = NamedTuple()) + boundary_conditions = nothing) LX, LY, LZ = loc @@ -378,12 +373,15 @@ function FieldTimeSeries(loc, grid, times=(); data = new_data(eltype(grid), grid, loc, indices, Nt) if backend isa OnDisk + isnothing(name) && isnothing(name) && + error(ArgumentError("Must provide the keyword arguments `path` and `name` when `backend=OnDisk()`.")) + isnothing(path) && error(ArgumentError("Must provide the keyword argument `path` when `backend=OnDisk()`.")) isnothing(name) && error(ArgumentError("Must provide the keyword argument `name` when `backend=OnDisk()`.")) end - - return FieldTimeSeries{LX, LY, LZ}(data, grid, backend, boundary_conditions, indices, - times, path, name, time_indexing, reader_kw) + + return FieldTimeSeries{LX, LY, LZ}(data, grid, backend, boundary_conditions, + indices, times, path, name, time_indexing) end """ @@ -410,16 +408,10 @@ end struct UnspecifiedBoundaryConditions end """ - FieldTimeSeries(path, name; - backend = InMemory(), - architecture = nothing, + FieldTimeSeries(path, name, backend = InMemory(); grid = nothing, - location = nothing, - boundary_conditions = UnspecifiedBoundaryConditions(), - time_indexing = Linear(), iterations = nothing, - times = nothing, - reader_kw = Dict{Symbol, Any}()) + times = nothing) Return a `FieldTimeSeries` containing a time-series of the field `name` load from JLD2 output located at `path`. @@ -438,9 +430,6 @@ Keyword arguments - `times`: Save times to load, as determined through an approximate floating point comparison to recorded save times. Defaults to times associated with `iterations`. Takes precedence over `iterations` if `times` is specified. - -- `reader_kw`: A named tuple or dictionary of keyword arguments to pass to the reader - (currently only JLD2) to be used when opening files. """ function FieldTimeSeries(path::String, name::String; backend = InMemory(), @@ -450,16 +439,19 @@ function FieldTimeSeries(path::String, name::String; boundary_conditions = UnspecifiedBoundaryConditions(), time_indexing = Linear(), iterations = nothing, - times = nothing, - reader_kw = NamedTuple()) + times = nothing) - file = jldopen(path; reader_kw...) + file = jldopen(path) # Defaults isnothing(iterations) && (iterations = parse.(Int, keys(file["timeseries/t"]))) isnothing(times) && (times = [file["timeseries/t/$i"] for i in iterations]) isnothing(location) && (Location = file["timeseries/$name/serialized/location"]) + if boundary_conditions isa UnspecifiedBoundaryConditions + boundary_conditions = file["timeseries/$name/serialized/boundary_conditions"] + end + indices = try file["timeseries/$name/serialized/indices"] catch @@ -476,12 +468,6 @@ function FieldTimeSeries(path::String, name::String; end end - if boundary_conditions isa UnspecifiedBoundaryConditions - boundary_conditions = file["timeseries/$name/serialized/boundary_conditions"] - boundary_conditions = on_architecture(architecture, boundary_conditions) - end - - # This should be removed eventually... (4/5/2022) grid = try on_architecture(architecture, grid) @@ -537,8 +523,8 @@ function FieldTimeSeries(path::String, name::String; Nt = time_indices_length(backend, times) data = new_data(eltype(grid), grid, loc, indices, Nt) - time_series = FieldTimeSeries{LX, LY, LZ}(data, grid, backend, boundary_conditions, indices, - times, path, name, time_indexing, reader_kw) + time_series = FieldTimeSeries{LX, LY, LZ}(data, grid, backend, boundary_conditions, + indices, times, path, name, time_indexing) set!(time_series, path, name) @@ -550,8 +536,7 @@ end grid = nothing, architecture = nothing, indices = (:, :, :), - boundary_conditions = nothing, - reader_kw = NamedTuple()) + boundary_conditions = nothing) Load a field called `name` saved in a JLD2 file at `path` at `iter`ation. Unless specified, the `grid` is loaded from `path`. @@ -560,8 +545,7 @@ function Field(location, path::String, name::String, iter; grid = nothing, architecture = nothing, indices = (:, :, :), - boundary_conditions = nothing, - reader_kw = NamedTuple()) + boundary_conditions = nothing) # Default to CPU if neither architecture nor grid is specified if isnothing(architecture) @@ -571,9 +555,9 @@ function Field(location, path::String, name::String, iter; architecture = Architectures.architecture(grid) end end - + # Load the grid and data from file - file = jldopen(path; reader_kw...) + file = jldopen(path) isnothing(grid) && (grid = file["serialized/grid"]) raw_data = file["timeseries/$name/$iter"] @@ -584,7 +568,7 @@ function Field(location, path::String, name::String, iter; grid = on_architecture(architecture, grid) raw_data = on_architecture(architecture, raw_data) data = offset_data(raw_data, grid, location, indices) - + return Field(location, grid; boundary_conditions, indices, data) end @@ -644,3 +628,4 @@ function fill_halo_regions!(fts::InMemoryFTS) return nothing end + diff --git a/src/OutputReaders/field_time_series_indexing.jl b/src/OutputReaders/field_time_series_indexing.jl index d849192f67..6a4683f9d0 100644 --- a/src/OutputReaders/field_time_series_indexing.jl +++ b/src/OutputReaders/field_time_series_indexing.jl @@ -14,7 +14,7 @@ import Oceananigans.Fields: interpolate # Cyclical implementation if out-of-bounds (wrap around the time-series) @inline function interpolating_time_indices(ti::Cyclical, times, t) Nt = length(times) - t¹ = first(times) + t¹ = first(times) tᴺ = last(times) T = ti.period @@ -32,14 +32,14 @@ import Oceananigans.Fields: interpolate uncycled_indices = (ñ, n₁, n₂) return ifelse(cycling, cycled_indices, uncycled_indices) -end +end # Clamp mode if out-of-bounds, i.e get the neareast neighbor @inline function interpolating_time_indices(::Clamp, times, t) n, n₁, n₂ = time_index_binary_search(times, t) beyond_indices = (0, n₂, n₂) # Beyond the last time: return n₂ - before_indices = (0, n₁, n₁) # Before the first time: return n₁ + before_indices = (0, n₁, n₁) # Before the first time: return n₁ unclamped_indices = (n, n₁, n₂) # Business as usual Nt = length(times) @@ -53,13 +53,13 @@ end @inline function time_index_binary_search(times, t) Nt = length(times) - # n₁ and n₂ are the index to interpolate inbetween and + # n₁ and n₂ are the index to interpolate inbetween and # n is a fractional index where 0 ≤ n ≤ 1 n₁, n₂ = index_binary_search(times, t, Nt) @inbounds begin - t₁ = times[n₁] - t₂ = times[n₂] + t₁ = times[n₁] + t₂ = times[n₂] end # "Fractional index" ñ ∈ (0, 1) @@ -79,7 +79,7 @@ import Base: getindex function getindex(fts::OnDiskFTS, n::Int) # Load data arch = architecture(fts) - file = jldopen(fts.path; fts.reader_kw...) + file = jldopen(fts.path) iter = keys(file["timeseries/t"])[n] raw_data = on_architecture(arch, file["timeseries/$(fts.name)/$iter"]) close(file) @@ -117,7 +117,7 @@ const YZFTS = FlavorOfFTS{Nothing, <:Any, <:Any, <:Any, <:Any} @inline function interpolating_getindex(fts, i, j, k, time_index) ñ, n₁, n₂ = interpolating_time_indices(fts.time_indexing, fts.times, time_index.time) - + @inbounds begin ψ₁ = getindex(fts, i, j, k, n₁) ψ₂ = getindex(fts, i, j, k, n₂) @@ -168,17 +168,13 @@ end to_time = to_time_index.time - if topology(from_grid) === (Flat, Flat, Flat) - ix = iy = iz = (1, 1, 0) - else - # Build space interpolators - to_node = flatten_node(to_node...) - ii, jj, kk = fractional_indices(to_node, from_grid, from_loc...) + # Build space interpolators + to_node = flatten_node(to_node...) + ii, jj, kk = fractional_indices(to_node, from_grid, from_loc...) - ix = interpolator(ii) - iy = interpolator(jj) - iz = interpolator(kk) - end + ix = interpolator(ii) + iy = interpolator(jj) + iz = interpolator(kk) ñ, n₁, n₂ = interpolating_time_indices(time_indexing, times, to_time) @@ -233,14 +229,14 @@ end ##### FieldTimeSeries updating ##### -# Let's make sure `times` is available on the CPU. This is always the case -# for ranges. if `times` is a vector that resides on the GPU, it has to be moved to the CPU for safe indexing. +# Let's make sure `times` is available on the CPU. This is always the case +# for ranges. if `times` is a vector that resides on the GPU, it has to be moved to the CPU for safe indexing. # TODO: Copying the whole array is a bit unclean, maybe find a way that avoids the penalty of allocating and copying memory. # This would require refactoring `FieldTimeSeries` to include a cpu-allocated times array cpu_interpolating_time_indices(::CPU, times, time_indexing, t, arch) = interpolating_time_indices(time_indexing, times, t) cpu_interpolating_time_indices(::CPU, times::AbstractVector, time_indexing, t) = interpolating_time_indices(time_indexing, times, t) -function cpu_interpolating_time_indices(::GPU, times::AbstractVector, time_indexing, t) +function cpu_interpolating_time_indices(::GPU, times::AbstractVector, time_indexing, t) cpu_times = on_architecture(CPU(), times) return interpolating_time_indices(time_indexing, cpu_times, t) end @@ -283,3 +279,4 @@ function getindex(fts::InMemoryFTS, n::Int) return Field(location(fts), fts.grid; data, fts.boundary_conditions, fts.indices) end + diff --git a/src/OutputReaders/set_field_time_series.jl b/src/OutputReaders/set_field_time_series.jl index 577082782e..d450926b3e 100644 --- a/src/OutputReaders/set_field_time_series.jl +++ b/src/OutputReaders/set_field_time_series.jl @@ -11,7 +11,7 @@ find_time_index(time::Number, file_times) = findfirst(t -> t ≈ time, fil find_time_index(time::AbstractTime, file_times) = findfirst(t -> t == time, file_times) function set!(fts::InMemoryFTS, path::String=fts.path, name::String=fts.name) - file = jldopen(path; fts.reader_kw...) + file = jldopen(path) file_iterations = iterations_from_file(file) file_times = [file["timeseries/t/$i"] for i in file_iterations] close(file) @@ -33,7 +33,7 @@ function set!(fts::InMemoryFTS, path::String=fts.path, name::String=fts.name) end file_iter = file_iterations[file_index] - + # Note: use the CPU for this step field_n = Field(location(fts), path, name, file_iter, architecture = cpu_architecture(arch), @@ -51,7 +51,7 @@ set!(fts::InMemoryFTS, value, n::Int) = set!(fts[n], value) function set!(fts::InMemoryFTS, fields_vector::AbstractVector{<:AbstractField}) raw_data = parent(fts) - file = jldopen(path; fts.reader_kw...) + file = jldopen(path) for (n, field) in enumerate(fields_vector) nth_raw_data = view(raw_data, :, :, :, n) @@ -68,7 +68,7 @@ end function maybe_write_property!(file, property, data) try test = file[property] - catch + catch file[property] = data end end @@ -101,3 +101,4 @@ function initialize_file!(file, name, fts) end set!(fts::OnDiskFTS, path::String, name::String) = nothing + diff --git a/src/OutputWriters/jld2_output_writer.jl b/src/OutputWriters/jld2_output_writer.jl index b8f4bd4a59..50a0a31937 100644 --- a/src/OutputWriters/jld2_output_writer.jl +++ b/src/OutputWriters/jld2_output_writer.jl @@ -30,7 +30,7 @@ ext(::Type{JLD2OutputWriter}) = ".jld2" JLD2OutputWriter(model, outputs; filename, schedule, dir = ".", indices = (:, :, :), - with_halos = true, + with_halos = false, array_type = Array{Float64}, file_splitting = NoFileSplitting(), overwrite_existing = false, @@ -70,8 +70,10 @@ Keyword arguments halo regions are removed from `indices`. For example, `indices = (:, :, 1)` will save xy-slices of the bottom-most index. -- `with_halos` (Bool): Whether or not to slice off or keep halo regions from fields before writing output. - Preserving halo region data can be useful for postprocessing. Default: true. +- `with_halos` (Bool): Whether or not to slice halo regions from fields before writing output. + Note, that to postprocess saved output (e.g., compute derivatives, etc) + information about the boundary conditions is often crucial. In that case + you might need to set `with_halos = true`. - `array_type`: The array type to which output arrays are converted to prior to saving. Default: `Array{Float64}`. @@ -110,7 +112,7 @@ Example Write out 3D fields for ``u``, ``v``, ``w``, and a tracer ``c``, along with a horizontal average: -```@example +``` using Oceananigans using Oceananigans.Utils: hour, minute @@ -131,21 +133,39 @@ simulation.output_writers[:velocities] = JLD2OutputWriter(model, model.velocitie filename = "some_data.jld2", schedule = TimeInterval(20minute), init = init_save_some_metadata!) + +# output +JLD2OutputWriter scheduled on TimeInterval(20 minutes): +├── filepath: ./some_data.jld2 +├── 3 outputs: (u, v, w) +├── array type: Array{Float64} +├── including: [:grid, :coriolis, :buoyancy, :closure] +├── file_splitting: NoFileSplitting +└── file size: 28.0 KiB ``` and a time- and horizontal-average of tracer ``c`` every 20 minutes of simulation time to a file called `some_averaged_data.jld2` -```@example +``` simulation.output_writers[:avg_c] = JLD2OutputWriter(model, (; c=c_avg), filename = "some_averaged_data.jld2", schedule = AveragedTimeInterval(20minute, window=5minute)) + +# output +JLD2OutputWriter scheduled on TimeInterval(20 minutes): +├── filepath: ./some_averaged_data.jld2 +├── 1 outputs: c averaged on AveragedTimeInterval(window=5 minutes, stride=1, interval=20 minutes) +├── array type: Array{Float64} +├── including: [:grid, :coriolis, :buoyancy, :closure] +├── file_splitting: NoFileSplitting +└── file size: 17.8 KiB ``` """ function JLD2OutputWriter(model, outputs; filename, schedule, dir = ".", indices = (:, :, :), - with_halos = true, + with_halos = false, array_type = Array{Float64}, file_splitting = NoFileSplitting(), overwrite_existing = false, @@ -157,8 +177,7 @@ function JLD2OutputWriter(model, outputs; filename, schedule, mkpath(dir) filename = auto_extension(filename, ".jld2") - filepath = abspath(joinpath(dir, filename)) - + filepath = joinpath(dir, filename) initialize!(file_splitting, model) update_file_splitting_schedule!(file_splitting, filepath) overwrite_existing && isfile(filepath) && rm(filepath, force=true) @@ -328,7 +347,7 @@ function Base.show(io::IO, ow::JLD2OutputWriter) Noutputs = length(ow.outputs) print(io, "JLD2OutputWriter scheduled on $(summary(ow.schedule)):", "\n", - "├── filepath: ", relpath(ow.filepath), "\n", + "├── filepath: $(ow.filepath)", "\n", "├── $Noutputs outputs: ", prettykeys(ow.outputs), show_averaging_schedule(averaging_schedule), "\n", "├── array type: ", show_array_type(ow.array_type), "\n", "├── including: ", ow.including, "\n", diff --git a/src/OutputWriters/netcdf_output_writer.jl b/src/OutputWriters/netcdf_output_writer.jl index 0660019d23..60ef6cc512 100644 --- a/src/OutputWriters/netcdf_output_writer.jl +++ b/src/OutputWriters/netcdf_output_writer.jl @@ -265,7 +265,7 @@ Examples Saving the ``u`` velocity field and temperature fields, the full 3D fields and surface 2D slices to separate NetCDF files: -```@example netcdf1 +```jldoctest netcdf1 using Oceananigans grid = RectilinearGrid(size=(16, 16, 16), extent=(1, 1, 1)) @@ -278,26 +278,53 @@ fields = Dict("u" => model.velocities.u, "c" => model.tracers.c) simulation.output_writers[:field_writer] = NetCDFOutputWriter(model, fields, filename="fields.nc", schedule=TimeInterval(60)) + +# output +NetCDFOutputWriter scheduled on TimeInterval(1 minute): +├── filepath: ./fields.nc +├── dimensions: zC(16), zF(17), xC(16), yF(16), xF(16), yC(16), time(0) +├── 2 outputs: (c, u) +└── array type: Array{Float64} +├── file_splitting: NoFileSplitting +└── file size: 14.9 KiB ``` -```@example netcdf1 +```jldoctest netcdf1 simulation.output_writers[:surface_slice_writer] = NetCDFOutputWriter(model, fields, filename="surface_xy_slice.nc", schedule=TimeInterval(60), indices=(:, :, grid.Nz)) + +# output +NetCDFOutputWriter scheduled on TimeInterval(1 minute): +├── filepath: ./surface_xy_slice.nc +├── dimensions: zC(1), zF(1), xC(16), yF(16), xF(16), yC(16), time(0) +├── 2 outputs: (c, u) +└── array type: Array{Float64} +├── file_splitting: NoFileSplitting +└── file size: 14.9 KiB ``` -```@example netcdf1 +```jldoctest netcdf1 simulation.output_writers[:averaged_profile_writer] = NetCDFOutputWriter(model, fields, filename = "averaged_z_profile.nc", schedule = AveragedTimeInterval(60, window=20), indices = (1, 1, :)) + +# output +NetCDFOutputWriter scheduled on TimeInterval(1 minute): +├── filepath: ./averaged_z_profile.nc +├── dimensions: zC(16), zF(17), xC(1), yF(1), xF(1), yC(1), time(0) +├── 2 outputs: (c, u) averaged on AveragedTimeInterval(window=20 seconds, stride=1, interval=1 minute) +└── array type: Array{Float64} +├── file_splitting: NoFileSplitting +└── file size: 17.6 KiB ``` `NetCDFOutputWriter` also accepts output functions that write scalars and arrays to disk, provided that their `dimensions` are provided: -```@example +```jldoctest using Oceananigans Nx, Ny, Nz = 16, 16, 16 @@ -334,13 +361,22 @@ simulation.output_writers[:things] = NetCDFOutputWriter(model, outputs, schedule=IterationInterval(1), filename="things.nc", dimensions=dims, verbose=true, global_attributes=global_attributes, output_attributes=output_attributes) + +# output +NetCDFOutputWriter scheduled on IterationInterval(1): +├── filepath: ./things.nc +├── dimensions: zC(16), zF(17), xC(16), yF(16), xF(16), yC(16), time(0) +├── 3 outputs: (profile, slice, scalar) +└── array type: Array{Float64} +├── file_splitting: NoFileSplitting +└── file size: 17.8 KiB ``` `NetCDFOutputWriter` can also be configured for `outputs` that are interpolated or regridded to a different grid than `model.grid`. To use this functionality, include the keyword argument `grid = output_grid`. -```@example +```jldoctest using Oceananigans using Oceananigans.Fields: interpolate! @@ -357,6 +393,15 @@ output_writer = NetCDFOutputWriter(model, outputs; grid = coarse_grid, filename = "coarse_u.nc", schedule = IterationInterval(1)) + +# output +NetCDFOutputWriter scheduled on IterationInterval(1): +├── filepath: ./coarse_u.nc +├── dimensions: zC(4), zF(5), xC(1), yF(1), xF(1), yC(1), time(0) +├── 1 outputs: u +└── array type: Array{Float64} +├── file_splitting: NoFileSplitting +└── file size: 14.6 KiB ``` """ function NetCDFOutputWriter(model, outputs; @@ -377,7 +422,7 @@ function NetCDFOutputWriter(model, outputs; verbose = false) mkpath(dir) filename = auto_extension(filename, ".nc") - filepath = abspath(joinpath(dir, filename)) + filepath = joinpath(dir, filename) initialize!(file_splitting, model) update_file_splitting_schedule!(file_splitting, filepath) @@ -389,6 +434,7 @@ function NetCDFOutputWriter(model, outputs; overwrite_existing = true end else + if isfile(filepath) && !overwrite_existing @warn "$filepath already exists and `overwrite_existing = false`. Mode will be set to append to existing file. " * "You might experience errors when writing output if the existing file belonged to a different simulation!" @@ -592,7 +638,7 @@ function Base.show(io::IO, ow::NetCDFOutputWriter) Noutputs = length(ow.outputs) print(io, "NetCDFOutputWriter scheduled on $(summary(ow.schedule)):", "\n", - "├── filepath: ", relpath(ow.filepath), "\n", + "├── filepath: ", ow.filepath, "\n", "├── dimensions: $dims", "\n", "├── $Noutputs outputs: ", prettykeys(ow.outputs), show_averaging_schedule(averaging_schedule), "\n", "└── array type: ", show_array_type(ow.array_type), "\n", diff --git a/src/OutputWriters/output_writer_utils.jl b/src/OutputWriters/output_writer_utils.jl index 7b79755b3b..af6c6f6834 100644 --- a/src/OutputWriters/output_writer_utils.jl +++ b/src/OutputWriters/output_writer_utils.jl @@ -98,7 +98,7 @@ function saveproperty!(file, address, bcs::FieldBoundaryConditions) if bc.condition isa Function || bc.condition isa ContinuousBoundaryFunction file[address * "/$boundary/condition"] = missing else - file[address * "/$boundary/condition"] = on_architecture(CPU(), bc.condition) + file[address * "/$boundary/condition"] = bc.condition end end end @@ -130,13 +130,13 @@ function serializeproperty!(file, address, grid::DistributedGrid) file[address] = on_architecture(cpu_arch, grid) end -function serializeproperty!(file, address, fbcs::FieldBoundaryConditions) +function serializeproperty!(file, address, p::FieldBoundaryConditions) # TODO: it'd be better to "filter" `FieldBoundaryCondition` and then serialize # rather than punting with `missing` instead. - if has_reference(Function, fbcs) + if has_reference(Function, p) file[address] = missing else - file[address] = on_architecture(CPU(), fbcs) + file[address] = p end end diff --git a/src/OutputWriters/windowed_time_average.jl b/src/OutputWriters/windowed_time_average.jl index 11f8be5682..c36a7cf528 100644 --- a/src/OutputWriters/windowed_time_average.jl +++ b/src/OutputWriters/windowed_time_average.jl @@ -59,7 +59,7 @@ AveragedTimeInterval(window=2 days, stride=1, interval=4 days) An `AveragedTimeInterval` schedule directs an output writer to time-average its outputs before writing them to disk: -```@example averaged_time_interval +```jldoctest averaged_time_interval using Oceananigans using Oceananigans.Units @@ -70,6 +70,15 @@ simulation = Simulation(model, Δt=10minutes, stop_time=30days) simulation.output_writers[:velocities] = JLD2OutputWriter(model, model.velocities, filename= "averaged_velocity_data.jld2", schedule = AveragedTimeInterval(4days, window=2days, stride=2)) + +# output +JLD2OutputWriter scheduled on TimeInterval(4 days): +├── filepath: ./averaged_velocity_data.jld2 +├── 3 outputs: (u, v, w) averaged on AveragedTimeInterval(window=2 days, stride=2, interval=4 days) +├── array type: Array{Float64} +├── including: [:grid, :coriolis, :buoyancy, :closure] +├── file_splitting: NoFileSplitting +└── file size: 27.6 KiB ``` """ function AveragedTimeInterval(interval; window=interval, stride=1) diff --git a/src/Simulations/Simulations.jl b/src/Simulations/Simulations.jl index 8c4da49e9d..9dc720b815 100644 --- a/src/Simulations/Simulations.jl +++ b/src/Simulations/Simulations.jl @@ -5,6 +5,7 @@ export Simulation export run! export Callback, add_callback! export iteration +export stopwatch using Oceananigans.Models using Oceananigans.Diagnostics diff --git a/src/Simulations/simulation.jl b/src/Simulations/simulation.jl index bd6a0fbbe1..3b73d85702 100644 --- a/src/Simulations/simulation.jl +++ b/src/Simulations/simulation.jl @@ -98,7 +98,7 @@ function Base.show(io::IO, s::Simulation) "├── Elapsed wall time: $(prettytime(s.run_wall_time))", "\n", "├── Wall time per iteration: $(prettytime(s.run_wall_time / iteration(s)))", "\n", "├── Stop time: $(prettytime(s.stop_time))", "\n", - "├── Stop iteration: $(s.stop_iteration)", "\n", + "├── Stop iteration : $(s.stop_iteration)", "\n", "├── Wall time limit: $(s.wall_time_limit)", "\n", "├── Callbacks: $(ordered_dict_show(s.callbacks, "│"))", "\n", "├── Output writers: $(ordered_dict_show(s.output_writers, "│"))", "\n", diff --git a/src/Solvers/Solvers.jl b/src/Solvers/Solvers.jl index c89a56e40e..bbad459ea9 100644 --- a/src/Solvers/Solvers.jl +++ b/src/Solvers/Solvers.jl @@ -53,6 +53,6 @@ const GridWithFourierTridiagonalSolver = Union{XYRegularRG, XZRegularRG, YZRegul fft_poisson_solver(grid::XYZRegularRG) = FFTBasedPoissonSolver(grid) fft_poisson_solver(grid::GridWithFourierTridiagonalSolver) = - FourierTridiagonalPoissonSolver(grid) + FourierTridiagonalPoissonSolver(grid.underlying_grid) end # module diff --git a/src/Solvers/conjugate_gradient_poisson_solver.jl b/src/Solvers/conjugate_gradient_poisson_solver.jl index eb9ec9db21..7a4ba2f117 100644 --- a/src/Solvers/conjugate_gradient_poisson_solver.jl +++ b/src/Solvers/conjugate_gradient_poisson_solver.jl @@ -1,4 +1,4 @@ -using Oceananigans.Operators +using Oceananigans.Operators: divᶜᶜᶜ, ∇²ᶜᶜᶜ using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid using Statistics: mean @@ -107,7 +107,7 @@ function compute_preconditioner_rhs!(solver::FourierTridiagonalPoissonSolver, rh arch = architecture(grid) tridiagonal_dir = solver.batched_tridiagonal_solver.tridiagonal_direction launch!(arch, grid, :xyz, fourier_tridiagonal_preconditioner_rhs!, - solver.storage, tridiagonal_dir, grid, rhs) + solver.storage, tridiagonal_dir, rhs) return nothing end @@ -115,7 +115,7 @@ const FFTBasedPreconditioner = Union{FFTBasedPoissonSolver, FourierTridiagonalPo function precondition!(p, preconditioner::FFTBasedPreconditioner, r, args...) compute_preconditioner_rhs!(preconditioner, r) - solve!(p, preconditioner) + p = solve!(p, preconditioner) mean_p = mean(p) grid = p.grid diff --git a/src/TimeSteppers/TimeSteppers.jl b/src/TimeSteppers/TimeSteppers.jl index e0e50c8970..2955ca6999 100644 --- a/src/TimeSteppers/TimeSteppers.jl +++ b/src/TimeSteppers/TimeSteppers.jl @@ -59,6 +59,5 @@ include("clock.jl") include("store_tendencies.jl") include("quasi_adams_bashforth_2.jl") include("runge_kutta_3.jl") -include("split_hydrostatic_runge_kutta_3.jl") end # module diff --git a/src/TimeSteppers/quasi_adams_bashforth_2.jl b/src/TimeSteppers/quasi_adams_bashforth_2.jl index 6d8d7ba24a..ed9f55750b 100644 --- a/src/TimeSteppers/quasi_adams_bashforth_2.jl +++ b/src/TimeSteppers/quasi_adams_bashforth_2.jl @@ -41,8 +41,7 @@ function QuasiAdamsBashforth2TimeStepper(grid, tracers, χ = 0.1; implicit_solver::IT = nothing, Gⁿ = TendencyFields(grid, tracers), - G⁻ = TendencyFields(grid, tracers), - kw...) where IT + G⁻ = TendencyFields(grid, tracers)) where IT FT = eltype(grid) GT = typeof(Gⁿ) diff --git a/src/TimeSteppers/split_hydrostatic_runge_kutta_3.jl b/src/TimeSteppers/split_hydrostatic_runge_kutta_3.jl deleted file mode 100644 index 2dbf8d2d1f..0000000000 --- a/src/TimeSteppers/split_hydrostatic_runge_kutta_3.jl +++ /dev/null @@ -1,175 +0,0 @@ -using Oceananigans.Architectures: architecture -using Oceananigans: fields - -""" - SplitRungeKutta3TimeStepper{FT, TG} <: AbstractTimeStepper - -Holds parameters and tendency fields for a low storage, third-order Runge-Kutta-Wray -time-stepping scheme described by [Lan2022](@citet). -""" -struct SplitRungeKutta3TimeStepper{FT, TG, TE, PF, TI} <: AbstractTimeStepper - γ² :: FT - γ³ :: FT - ζ² :: FT - ζ³ :: FT - Gⁿ :: TG - G⁻ :: TE # only used as storage when needed - Ψ⁻ :: PF # prognostic state at the previous timestep - implicit_solver :: TI -end - -""" - SplitRungeKutta3TimeStepper(grid, tracers; - implicit_solver = nothing, - Gⁿ = TendencyFields(grid, tracers), - G⁻ = TendencyFields(grid, tracers), - Ψ⁻ = TendencyFields(grid, tracers)) - -Return a 3rd-order `SplitRungeKutta3TimeStepper` on `grid` and with `tracers`. -The tendency fields `Gⁿ` and `G⁻`, as well as the previous prognostic state Ψ⁻ can be specified via optional `kwargs`. - -The scheme described by [Lan2022](@citet). In a nutshel, the 3rd-order -Runge Kutta timestepper steps forward the state `Uⁿ` by `Δt` via 3 substeps. A baroptropic velocity correction -step is applied after at each substep. - -The state `U` after each substep `m` is - -```julia -Uᵐ⁺¹ = ζᵐ * Uⁿ + γᵐ * (Uᵐ + Δt * Gᵐ) -``` - -where `Uᵐ` is the state at the ``m``-th substep, `Uⁿ` is the state at the ``n``-th timestep, `Gᵐ` is the tendency -at the ``m``-th substep, and constants ``γ¹ = 1`, ``γ² = 1/4``, ``γ³ = 1/3``, -``ζ¹ = 0``, ``ζ² = 3/4``, ``ζ³ = 1/3``. - -The state at the first substep is taken to be the one that corresponds to the ``n``-th timestep, -`U¹ = Uⁿ`, and the state after the third substep is then the state at the `Uⁿ⁺¹ = U³`. -""" -function SplitRungeKutta3TimeStepper(grid, tracers; - implicit_solver::TI = nothing, - Gⁿ::TG = TendencyFields(grid, tracers), - G⁻::TE = TendencyFields(grid, tracers), - Ψ⁻::PF = TendencyFields(grid, tracers)) where {TI, TG, TE, PF} - - - @warn("Split barotropic-baroclinic time stepping with SplitRungeKutta3TimeStepper is not tested and experimental.\n" * - "Use at own risk, and report any issues encountered.") - - !isnothing(implicit_solver) && - @warn("Implicit-explicit time-stepping with SplitRungeKutta3TimeStepper is not tested. " * - "\n implicit_solver: $(typeof(implicit_solver))") - - γ² = 1 // 4 - γ³ = 2 // 3 - - ζ² = 3 // 4 - ζ³ = 1 // 3 - - FT = eltype(grid) - - return SplitRungeKutta3TimeStepper{FT, TG, TE, PF, TI}(γ², γ³, ζ², ζ³, Gⁿ, G⁻, Ψ⁻, implicit_solver) -end - - -function time_step!(model::AbstractModel{<:SplitRungeKutta3TimeStepper}, Δt; callbacks=[]) - Δt == 0 && @warn "Δt == 0 may cause model blowup!" - - # Be paranoid and update state at iteration 0, in case run! is not used: - model.clock.iteration == 0 && update_state!(model, callbacks; compute_tendencies = true) - - γ² = model.timestepper.γ² - γ³ = model.timestepper.γ³ - - ζ² = model.timestepper.ζ² - ζ³ = model.timestepper.ζ³ - - store_fields!(model) - - #### - #### First stage - #### - - model.clock.stage = 1 - - split_rk3_substep!(model, Δt, nothing, nothing) - calculate_pressure_correction!(model, Δt) - pressure_correct_velocities!(model, Δt) - update_state!(model, callbacks; compute_tendencies = true) - - #### - #### Second stage - #### - - model.clock.stage = 2 - - split_rk3_substep!(model, Δt, γ², ζ²) - calculate_pressure_correction!(model, Δt) - pressure_correct_velocities!(model, Δt) - update_state!(model, callbacks; compute_tendencies = true) - - #### - #### Third stage - #### - - model.clock.stage = 3 - - split_rk3_substep!(model, Δt, γ³, ζ³) - calculate_pressure_correction!(model, Δt) - pressure_correct_velocities!(model, Δt) - update_state!(model, callbacks; compute_tendencies = true) - - step_lagrangian_particles!(model, Δt) - - tick!(model.clock, Δt) - model.clock.last_Δt = Δt - - return nothing -end - -@kernel function _split_rk3_substep_field!(cⁿ, Δt, γⁿ::FT, ζⁿ, Gⁿ, cⁿ⁻¹) where FT - i, j, k = @index(Global, NTuple) - cⁿ[i, j, k] = ζⁿ * cⁿ⁻¹[i, j, k] + γⁿ * (cⁿ[i, j, k] + convert(FT, Δt) * Gⁿ[i, j, k]) -end - -@kernel function _split_rk3_substep_field!(cⁿ, Δt, ::Nothing, ::Nothing, Gⁿ, cⁿ⁻¹) - i, j, k = @index(Global, NTuple) - cⁿ[i, j, k] = cⁿ[i, j, k] + Δt * Gⁿ[i, j, k] -end - -function split_rk3_substep!(model, Δt, γⁿ, ζⁿ) - - grid = model.grid - arch = architecture(grid) - model_fields = prognostic_fields(model) - - for (i, field) in enumerate(model_fields) - kernel_args = (field, Δt, γⁿ, ζⁿ, model.timestepper.Gⁿ[i], model.timestepper.Ψ⁻[i]) - launch!(arch, grid, :xyz, rk3_substep_field!, kernel_args...; exclude_periphery=true) - - # TODO: function tracer_index(model, field_index) = field_index - 3, etc... - tracer_index = Val(i - 3) # assumption - - implicit_step!(field, - model.timestepper.implicit_solver, - model.closure, - model.diffusivity_fields, - tracer_index, - model.clock, - stage_Δt(Δt, γⁿ, ζⁿ)) - end -end - -function store_fields!(model) - - previous_fields = model.timestepper.Ψ⁻ - model_fields = prognostic_fields(model) - - for name in keys(previous_fields) - if !isnothing(previous_fields[name]) - parent(previous_fields[name]) .= parent(model_fields[name]) # Storing also the halos - end - end - - return nothing -end - diff --git a/src/TurbulenceClosures/TurbulenceClosures.jl b/src/TurbulenceClosures/TurbulenceClosures.jl index 2e988816a4..5b0aabcac3 100644 --- a/src/TurbulenceClosures/TurbulenceClosures.jl +++ b/src/TurbulenceClosures/TurbulenceClosures.jl @@ -12,9 +12,6 @@ export ScalarBiharmonicDiffusivity, TwoDimensionalLeith, SmagorinskyLilly, - Smagorinsky, - LillyCoefficient, - DynamicCoefficient, AnisotropicMinimumDissipation, ConvectiveAdjustmentVerticalDiffusivity, RiBasedVerticalDiffusivity, @@ -53,8 +50,7 @@ using Oceananigans.Utils using Oceananigans.Architectures: AbstractArchitecture, device using Oceananigans.Fields: FunctionField -using Oceananigans.ImmersedBoundaries -using Oceananigans.ImmersedBoundaries: AbstractGridFittedBottom +using Oceananigans.ImmersedBoundaries: z_bottom import Oceananigans.Grids: required_halo_size_x, required_halo_size_y, required_halo_size_z import Oceananigans.Architectures: on_architecture @@ -82,7 +78,7 @@ compute_diffusivities!(K, closure::AbstractTurbulenceClosure, args...; kwargs... # be calculated from local information, still `B = 1`, because we need at least one additional # point at each side to calculate viscous fluxes at the edge of the domain. # If diffusivity itself requires one halo to be computed (e.g. κ = ℑxᶠᵃᵃ(i, j, k, grid, ℑxᶜᵃᵃ, T), -# or `AnisotropicMinimumDissipation` and `Smagorinsky`) then B = 2 +# or `AnisotropicMinimumDissipation` and `SmagorinskyLilly`) then B = 2 @inline required_halo_size_x(::AbstractTurbulenceClosure{TD, B}) where {TD, B} = B @inline required_halo_size_y(::AbstractTurbulenceClosure{TD, B}) where {TD, B} = B @inline required_halo_size_z(::AbstractTurbulenceClosure{TD, B}) where {TD, B} = B @@ -124,21 +120,14 @@ end @inline clip(x) = max(zero(x), x) -##### -##### Height, Depth and Bottom interfaces -##### - const c = Center() const f = Face() -const AGFBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractGridFittedBottom} - @inline z_top(i, j, grid) = znode(i, j, grid.Nz+1, grid, c, c, f) -@inline z_bottom(i, j, grid) = znode(i, j, 1, grid, c, c, f) -@inline z_bottom(i, j, ibg::AGFBIBG) = @inbounds ibg.immersed_boundary.bottom_height[i, j, 1] -@inline depthᶜᶜᶠ(i, j, k, grid) = clip(z_top(i, j, grid) - znode(i, j, k, grid, c, c, f)) -@inline depthᶜᶜᶜ(i, j, k, grid) = clip(z_top(i, j, grid) - znode(i, j, k, grid, c, c, c)) +@inline depthᶜᶜᶠ(i, j, k, grid) = clip(z_top(i, j, grid) - znode(i, j, k, grid, c, c, f)) +@inline depthᶜᶜᶜ(i, j, k, grid) = clip(z_top(i, j, grid) - znode(i, j, k, grid, c, c, c)) +@inline total_depthᶜᶜᵃ(i, j, grid) = clip(z_top(i, j, grid) - z_bottom(i, j, grid)) @inline function height_above_bottomᶜᶜᶠ(i, j, k, grid) h = znode(i, j, k, grid, c, c, f) - z_bottom(i, j, grid) @@ -177,8 +166,8 @@ include("turbulence_closure_implementations/nothing_closure.jl") # AbstractScalarDiffusivity closures: include("turbulence_closure_implementations/scalar_diffusivity.jl") include("turbulence_closure_implementations/scalar_biharmonic_diffusivity.jl") +include("turbulence_closure_implementations/smagorinsky_lilly.jl") include("turbulence_closure_implementations/anisotropic_minimum_dissipation.jl") -include("turbulence_closure_implementations/Smagorinskys/Smagorinskys.jl") include("turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl") include("turbulence_closure_implementations/TKEBasedVerticalDiffusivities/TKEBasedVerticalDiffusivities.jl") include("turbulence_closure_implementations/ri_based_vertical_diffusivity.jl") @@ -189,8 +178,6 @@ include("turbulence_closure_implementations/isopycnal_skew_symmetric_diffusivity include("turbulence_closure_implementations/leith_enstrophy_diffusivity.jl") using .TKEBasedVerticalDiffusivities: CATKEVerticalDiffusivity, TKEDissipationVerticalDiffusivity -using .Smagorinskys: Smagorinsky, DynamicSmagorinsky, SmagorinskyLilly -using .Smagorinskys: LillyCoefficient, DynamicCoefficient, LagrangianAveraging # Miscellaneous utilities include("diffusivity_fields.jl") diff --git a/src/TurbulenceClosures/turbulence_closure_diagnostics.jl b/src/TurbulenceClosures/turbulence_closure_diagnostics.jl index a2a82ba65b..839f6d4456 100644 --- a/src/TurbulenceClosures/turbulence_closure_diagnostics.jl +++ b/src/TurbulenceClosures/turbulence_closure_diagnostics.jl @@ -44,7 +44,7 @@ function cell_diffusion_timescale(closure::ScalarBiharmonicDiffusivity{Dir}, dif return min(Δ^4/ max_ν, Δ^4 / max_κ) end -function cell_diffusion_timescale(closure::Smagorinsky, diffusivities, grid) +function cell_diffusion_timescale(closure::SmagorinskyLilly, diffusivities, grid) Δ = min_Δxyz(grid, formulation(closure)) min_Pr = closure.Pr isa NamedTuple{()} ? 1 : minimum(closure.Pr) # Innocuous value is there's no tracers max_νκ = maximum(diffusivities.νₑ.data.parent) * max(1, 1/min_Pr) diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/Smagorinskys.jl b/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/Smagorinskys.jl deleted file mode 100644 index a7822fcd11..0000000000 --- a/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/Smagorinskys.jl +++ /dev/null @@ -1,15 +0,0 @@ -module Smagorinskys - -using Oceananigans.Operators -using Oceananigans.Fields - -using Oceananigans.Grids: AbstractGrid - -using KernelAbstractions: @kernel, @index - -include("smagorinsky.jl") -include("dynamic_coefficient.jl") -include("lilly_coefficient.jl") -include("scale_invariant_operators.jl") - -end diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/dynamic_coefficient.jl b/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/dynamic_coefficient.jl deleted file mode 100644 index e99db2a606..0000000000 --- a/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/dynamic_coefficient.jl +++ /dev/null @@ -1,266 +0,0 @@ -using Oceananigans.Architectures: architecture -using Oceananigans.Fields: interpolate -using Statistics - -struct DynamicCoefficient{A, FT, S} - averaging :: A - minimum_numerator :: FT - schedule :: S -end - -const DynamicSmagorinsky = Smagorinsky{<:Any, <:DynamicCoefficient} - -function DynamicSmagorinsky(time_discretization=ExplicitTimeDiscretization(), FT=Float64; averaging, - Pr=1.0, schedule=IterationInterval(1), minimum_numerator=1e-32) - coefficient = DynamicCoefficient(FT; averaging, schedule, minimum_numerator) - TD = typeof(time_discretization) - Pr = convert_diffusivity(FT, Pr; discrete_form=false) - return Smagorinsky{TD}(coefficient, Pr) -end - -DynamicSmagorinsky(FT::DataType; kwargs...) = DynamicSmagorinsky(ExplicitTimeDiscretization(), FT; kwargs...) - -Adapt.adapt_structure(to, dc::DynamicCoefficient) = DynamicCoefficient(dc.averaging, dc.minimum_numerator, nothing) - -const DirectionallyAveragedCoefficient{N} = DynamicCoefficient{<:Union{NTuple{N, Int}, Int, Colon}} where N -const DirectionallyAveragedDynamicSmagorinsky{N} = Smagorinsky{<:Any, <:DirectionallyAveragedCoefficient{N}} where N - -struct LagrangianAveraging end -const LagrangianAveragedCoefficient = DynamicCoefficient{<:LagrangianAveraging} -const LagrangianAveragedDynamicSmagorinsky = Smagorinsky{<:Any, <:LagrangianAveragedCoefficient} - -""" - DynamicCoefficient([FT=Float64;] averaging, schedule=IterationInterval(1), minimum_numerator=1e-32) - -When used with `Smagorinsky`, it calculates the Smagorinsky coefficient dynamically from the flow -according to the scale invariant procedure in [BouZeid05](@citet). - -`DynamicCoefficient` requires an `averaging` procedure, which can be a `LagrangianAveraging` (which -averages fluid parcels along their Lagrangian trajectory) or a tuple of integers indicating -a directional averaging procedure along chosen dimensions (e.g. `averaging=(1,2)` uses averages -in the `x` and `y` directions). - -`DynamicCoefficient` is updated according to `schedule`, and `minimum_numerator` defines the minimum -value that is acceptable in the denominator of the final calculation. -""" -function DynamicCoefficient(FT=Float64; averaging, schedule=IterationInterval(1), minimum_numerator=1e-32) - minimum_numerator = convert(FT, minimum_numerator) - return DynamicCoefficient(averaging, minimum_numerator, schedule) -end - -Base.summary(dc::DynamicCoefficient) = string("DynamicCoefficient(averaging = $(dc.averaging), schedule = $(dc.schedule))") -Base.show(io::IO, dc::DynamicCoefficient) = print(io, "DynamicCoefficient with\n", - "├── averaging = ", dc.averaging, "\n", - "├── schedule = ", dc.schedule, "\n", - "└── minimum_numerator = ", dc.minimum_numerator) - -##### -##### Some common utilities independent of averaging -##### - -@inline function square_smagorinsky_coefficient(i, j, k, grid, closure::DynamicSmagorinsky, diffusivity_fields, args...) - 𝒥ᴸᴹ = diffusivity_fields.𝒥ᴸᴹ - 𝒥ᴹᴹ = diffusivity_fields.𝒥ᴹᴹ - 𝒥ᴸᴹ_min = closure.coefficient.minimum_numerator - - @inbounds begin - 𝒥ᴸᴹ_ijk = max(𝒥ᴸᴹ[i, j, k], 𝒥ᴸᴹ_min) - 𝒥ᴹᴹ_ijk = 𝒥ᴹᴹ[i, j, k] - end - - return ifelse(𝒥ᴹᴹ_ijk == 0, zero(grid), 𝒥ᴸᴹ_ijk / 𝒥ᴹᴹ_ijk) -end - -@kernel function _compute_Σ_Σ̄!(Σ, Σ̄, grid, u, v, w) - i, j, k = @index(Global, NTuple) - - @inbounds begin - Σ[i, j, k] = √(ΣᵢⱼΣᵢⱼᶜᶜᶜ(i, j, k, grid, u, v, w)) - Σ̄[i, j, k] = √(Σ̄ᵢⱼΣ̄ᵢⱼᶜᶜᶜ(i, j, k, grid, u, v, w)) - end -end - -@kernel function _compute_LM_MM!(LM, MM, Σ, Σ̄, grid, u, v, w) - i, j, k = @index(Global, NTuple) - LM_ijk, MM_ijk = LM_and_MM(i, j, k, grid, Σ, Σ̄, u, v, w) - @inbounds begin - LM[i, j, k] = LM_ijk - MM[i, j, k] = MM_ijk - end -end - -@inline function LM_and_MM(i, j, k, grid, Σ, Σ̄, u, v, w) - L₁₁ = L₁₁ᶜᶜᶜ(i, j, k, grid, u, v, w) - L₂₂ = L₂₂ᶜᶜᶜ(i, j, k, grid, u, v, w) - L₃₃ = L₃₃ᶜᶜᶜ(i, j, k, grid, u, v, w) - L₁₂ = L₁₂ᶜᶜᶜ(i, j, k, grid, u, v, w) - L₁₃ = L₁₃ᶜᶜᶜ(i, j, k, grid, u, v, w) - L₂₃ = L₂₃ᶜᶜᶜ(i, j, k, grid, u, v, w) - - M₁₁ = M₁₁ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ, Σ̄) - M₂₂ = M₂₂ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ, Σ̄) - M₃₃ = M₃₃ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ, Σ̄) - M₁₂ = M₁₂ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ, Σ̄) - M₁₃ = M₁₃ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ, Σ̄) - M₂₃ = M₂₃ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ, Σ̄) - - LM_ijk = L₁₁ * M₁₁ + L₂₂ * M₂₂ + L₃₃ * M₃₃ + 2L₁₂ * M₁₂ + 2L₁₃ * M₁₃ + 2L₂₃ * M₂₃ - MM_ijk = M₁₁ * M₁₁ + M₂₂ * M₂₂ + M₃₃ * M₃₃ + 2M₁₂ * M₁₂ + 2M₁₃ * M₁₃ + 2M₂₃ * M₂₃ - - return LM_ijk, MM_ijk -end - -##### -##### Directionally-averaged functionality -##### - -function compute_coefficient_fields!(diffusivity_fields, closure::DirectionallyAveragedDynamicSmagorinsky, model; parameters) - grid = model.grid - arch = architecture(grid) - velocities = model.velocities - cˢ = closure.coefficient - - if cˢ.schedule(model) - Σ = diffusivity_fields.Σ - Σ̄ = diffusivity_fields.Σ̄ - launch!(arch, grid, :xyz, _compute_Σ_Σ̄!, Σ, Σ̄, grid, velocities...) - - LM = diffusivity_fields.LM - MM = diffusivity_fields.MM - launch!(arch, grid, :xyz, _compute_LM_MM!, LM, MM, Σ, Σ̄, grid, velocities...) - - 𝒥ᴸᴹ = diffusivity_fields.𝒥ᴸᴹ - 𝒥ᴹᴹ = diffusivity_fields.𝒥ᴹᴹ - compute!(𝒥ᴸᴹ) - compute!(𝒥ᴹᴹ) - end - - return nothing -end - -function allocate_coefficient_fields(closure::DirectionallyAveragedDynamicSmagorinsky, grid) - LM = CenterField(grid) - MM = CenterField(grid) - - Σ = CenterField(grid) - Σ̄ = CenterField(grid) - - 𝒥ᴸᴹ = Field(Average(LM, dims=closure.coefficient.averaging)) - 𝒥ᴹᴹ = Field(Average(MM, dims=closure.coefficient.averaging)) - - return (; Σ, Σ̄, LM, MM, 𝒥ᴸᴹ, 𝒥ᴹᴹ) -end - -##### -##### Lagrangian-averaged functionality -##### - -const c = Center() - -@kernel function _lagrangian_average_LM_MM!(𝒥ᴸᴹ, 𝒥ᴹᴹ, 𝒥ᴸᴹ⁻, 𝒥ᴹᴹ⁻, 𝒥ᴸᴹ_min, Σ, Σ̄, grid, Δt, u, v, w) - i, j, k = @index(Global, NTuple) - LM, MM = LM_and_MM(i, j, k, grid, Σ, Σ̄, u, v, w) - FT = eltype(grid) - - @inbounds begin - 𝒥ᴸᴹ⁻ᵢⱼₖ = max(𝒥ᴸᴹ⁻[i, j, k], 𝒥ᴸᴹ_min) - 𝒥ᴹᴹ⁻ᵢⱼₖ = 𝒥ᴹᴹ⁻[i, j, k] - - # Compute time scale - 𝒥ᴸᴹ𝒥ᴹᴹ = 𝒥ᴸᴹ⁻ᵢⱼₖ * 𝒥ᴹᴹ⁻ᵢⱼₖ - - T⁻ = convert(FT, 1.5) * Δᶠ(i, j, k, grid) / ∜(∜(𝒥ᴸᴹ𝒥ᴹᴹ)) - τ = Δt / T⁻ - ϵ = τ / (1 + τ) - - # Compute interpolation - x = xnode(i, j, k, grid, c, c, c) - y = ynode(i, j, k, grid, c, c, c) - z = znode(i, j, k, grid, c, c, c) - - # Displacements - δx = u[i, j, k] * Δt - δy = v[i, j, k] * Δt - δz = w[i, j, k] * Δt - - # Prevent displacements from getting too big? - Δx = Δxᶜᶜᶜ(i, j, k, grid) - Δy = Δyᶜᶜᶜ(i, j, k, grid) - Δz = Δzᶜᶜᶜ(i, j, k, grid) - - δx = clamp(δx, -Δx, Δx) - δy = clamp(δy, -Δy, Δy) - δz = clamp(δz, -Δz, Δz) - - # Previous locations - x⁻ = x - δx - y⁻ = y - δy - z⁻ = z - δz - X⁻ = (x⁻, y⁻, z⁻) - - itp_𝒥ᴹᴹ⁻ = interpolate(X⁻, 𝒥ᴹᴹ⁻, (c, c, c), grid) - itp_𝒥ᴸᴹ⁻ = interpolate(X⁻, 𝒥ᴸᴹ⁻, (c, c, c), grid) - - # Take time-step - 𝒥ᴹᴹ[i, j, k] = ϵ * MM + (1 - ϵ) * itp_𝒥ᴹᴹ⁻ - - 𝒥ᴸᴹ★ = ϵ * LM + (1 - ϵ) * max(itp_𝒥ᴸᴹ⁻, 𝒥ᴸᴹ_min) - 𝒥ᴸᴹ[i, j, k] = max(𝒥ᴸᴹ★, 𝒥ᴸᴹ_min) - end -end - -function compute_coefficient_fields!(diffusivity_fields, closure::LagrangianAveragedDynamicSmagorinsky, model; parameters) - grid = model.grid - arch = architecture(grid) - clock = model.clock - cˢ = closure.coefficient - t⁻ = diffusivity_fields.previous_compute_time - u, v, w = model.velocities - - Δt = clock.time - t⁻[] - t⁻[] = model.clock.time - - if cˢ.schedule(model) - Σ = diffusivity_fields.Σ - Σ̄ = diffusivity_fields.Σ̄ - launch!(arch, grid, :xyz, _compute_Σ_Σ̄!, Σ, Σ̄, grid, u, v, w) - - parent(diffusivity_fields.𝒥ᴸᴹ⁻) .= parent(diffusivity_fields.𝒥ᴸᴹ) - parent(diffusivity_fields.𝒥ᴹᴹ⁻) .= parent(diffusivity_fields.𝒥ᴹᴹ) - - 𝒥ᴸᴹ⁻ = diffusivity_fields.𝒥ᴸᴹ⁻ - 𝒥ᴹᴹ⁻ = diffusivity_fields.𝒥ᴹᴹ⁻ - 𝒥ᴸᴹ = diffusivity_fields.𝒥ᴸᴹ - 𝒥ᴹᴹ = diffusivity_fields.𝒥ᴹᴹ - 𝒥ᴸᴹ_min = cˢ.minimum_numerator - - if !isfinite(clock.last_Δt) || Δt == 0 # first time-step - launch!(arch, grid, :xyz, _compute_LM_MM!, 𝒥ᴸᴹ, 𝒥ᴹᴹ, Σ, Σ̄, grid, u, v, w) - parent(𝒥ᴸᴹ) .= max(mean(𝒥ᴸᴹ), 𝒥ᴸᴹ_min) - parent(𝒥ᴹᴹ) .= mean(𝒥ᴹᴹ) - else - launch!(arch, grid, :xyz, - _lagrangian_average_LM_MM!, 𝒥ᴸᴹ, 𝒥ᴹᴹ, 𝒥ᴸᴹ⁻, 𝒥ᴹᴹ⁻, 𝒥ᴸᴹ_min, Σ, Σ̄, grid, Δt, u, v, w) - - end - end - - return nothing -end - -function allocate_coefficient_fields(closure::LagrangianAveragedDynamicSmagorinsky, grid) - 𝒥ᴸᴹ⁻ = CenterField(grid) - 𝒥ᴹᴹ⁻ = CenterField(grid) - - 𝒥ᴸᴹ = CenterField(grid) - 𝒥ᴹᴹ = CenterField(grid) - - Σ = CenterField(grid) - Σ̄ = CenterField(grid) - - previous_compute_time = Ref(zero(grid)) - - return (; Σ, Σ̄, 𝒥ᴸᴹ, 𝒥ᴹᴹ, 𝒥ᴸᴹ⁻, 𝒥ᴹᴹ⁻, previous_compute_time) -end - - diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/lilly_coefficient.jl b/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/lilly_coefficient.jl deleted file mode 100644 index bf748cb2ce..0000000000 --- a/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/lilly_coefficient.jl +++ /dev/null @@ -1,146 +0,0 @@ -using Oceananigans.BuoyancyModels: ∂z_b - -struct LillyCoefficient{FT} - smagorinsky :: FT - reduction_factor :: FT -end - -""" - LillyCoefficient([FT=Float64;] smagorinsky=0.16, reduction_factor=1) - -When used with `Smagorinsky`, it calculates the Smagorinsky coefficient according to closure -proposed by [Lilly62](@citet), and [Lilly66](@citet), which has an eddy viscosity of the form - -``` -νₑ = (Cˢ * Δᶠ)² * √(2Σ²) * √(1 - Cb * N² / Σ²) -``` - -where `C²` is the Smagorinsky coefficient, `Δᶠ` is the filter width, `Σ² = ΣᵢⱼΣᵢⱼ` is the double dot -product of the strain tensor `Σᵢⱼ`, `N²` is the total buoyancy gradient, and `Cb` is a constant the -multiplies the Richardson number modification to the eddy viscosity. - -Arguments -========= - -* `FT`: Float type; default `Float64`. - -Keyword arguments -================= - -* `smagorinsky`: Smagorinsky coefficient `Cˢ`. Default value is 0.16 as obtained by Lilly (1966). - -* `reduction_factor`: Buoyancy term multipler `Cb` based on Lilly (1962) (`reduction_factor = 0` - turns it off, `reduction_factor ≠ 0` turns it on. - Typically, and according to the original work by Lilly (1962), `Cb = 1 / Pr`.) - -References -========== - -Lilly, D. K. "On the numerical simulation of buoyant convection." Tellus (1962) - -Lilly, D. K. "The representation of small-scale turbulence in numerical simulation experiments." - NCAR Manuscript No. 281, 0, (1966) -""" -LillyCoefficient(FT=Float64; smagorinsky=0.16, reduction_factor=1) = - LillyCoefficient(convert(FT, smagorinsky), convert(FT, reduction_factor)) - -const SmagorinskyLilly = Smagorinsky{<:Any, <:LillyCoefficient} - -""" - SmagorinskyLilly([time_discretization::TD = ExplicitTimeDiscretization(), FT=Float64;] C=0.16, Cb=1.0, Pr=1.0) - -Return a `SmagorinskyLilly` type associated with the turbulence closure proposed by -[Lilly62](@citet), [Smagorinsky1958](@citet), [Smagorinsky1963](@citet), and [Lilly66](@citet), -which has an eddy viscosity of the form - -``` -νₑ = (Cˢ * Δᶠ)² * √(2Σ²) * √(1 - Cb * N² / Σ²) -``` - -and an eddy diffusivity of the form - -``` -κₑ = νₑ / Pr -``` - -where `Δᶠ` is the filter width, `Σ² = ΣᵢⱼΣᵢⱼ` is the double dot product of -the strain tensor `Σᵢⱼ`, `Pr` is the turbulent Prandtl number, `N²` is the -total buoyancy gradient, and `Cb` is a constant the multiplies the Richardson -number modification to the eddy viscosity. - -Arguments -========= - -* `time_discretization`: Either `ExplicitTimeDiscretization()` or `VerticallyImplicitTimeDiscretization()`, - which integrates the terms involving only ``z``-derivatives in the - viscous and diffusive fluxes with an implicit time discretization. - Default `ExplicitTimeDiscretization()`. - -* `FT`: Float type; default `Float64`. - -Keyword arguments -================= - -* `Cˢ`: Smagorinsky coefficient. Default value is 0.16 as obtained by Lilly (1966). - -* `Cb`: Buoyancy term multipler based on Lilly (1962) (`Cb = 0` turns it off, `Cb ≠ 0` turns it on. - Typically, and according to the original work by Lilly (1962), `Cb = 1 / Pr`.) - -* `Pr`: Turbulent Prandtl numbers for each tracer. Either a constant applied to every - tracer, or a `NamedTuple` with fields for each tracer individually. - -References -========== - -Smagorinsky, J. "On the numerical integration of the primitive equations of motion for - baroclinic flow in a closed region." Monthly Weather Review (1958) - -Lilly, D. K. "On the numerical simulation of buoyant convection." Tellus (1962) - -Smagorinsky, J. "General circulation experiments with the primitive equations: I. - The basic experiment." Monthly Weather Review (1963) - -Lilly, D. K. "The representation of small-scale turbulence in numerical simulation experiments." - NCAR Manuscript No. 281, 0, (1966) -""" -function SmagorinskyLilly(time_discretization=ExplicitTimeDiscretization(), FT=Float64; C=0.16, Cb=1, Pr=1) - coefficient = LillyCoefficient(FT, smagorinsky=C, reduction_factor=Cb) - TD = typeof(time_discretization) - Pr = convert_diffusivity(FT, Pr; discrete_form=false) - return Smagorinsky{TD}(coefficient, Pr) -end - -SmagorinskyLilly(FT::DataType; kwargs...) = SmagorinskyLilly(ExplicitTimeDiscretization(), FT; kwargs...) - -""" - stability(N², Σ², Cb) - -Return the stability function - -```math - \\sqrt(1 - Cb N^2 / Σ^2 ) -``` - -when ``N^2 > 0``, and 1 otherwise. -""" -@inline function stability(N²::FT, Σ²::FT, cᵇ::FT) where FT - N²⁺ = max(zero(FT), N²) # clip - ς² = one(FT) - min(one(FT), cᵇ * N²⁺ / Σ²) - return ifelse(Σ²==0, zero(FT), sqrt(ς²)) -end - -@inline function square_smagorinsky_coefficient(i, j, k, grid, closure::SmagorinskyLilly, - diffusivity_fields, Σ², buoyancy, tracers) - N² = ℑzᵃᵃᶜ(i, j, k, grid, ∂z_b, buoyancy, tracers) - c₀ = closure.coefficient.smagorinsky - cᵇ = closure.coefficient.reduction_factor - ς = stability(N², Σ², cᵇ) # Use unity Prandtl number. - return ς * c₀^2 -end - -Base.summary(dc::LillyCoefficient) = string("LillyCoefficient(smagorinsky = $(dc.smagorinsky), reduction_factor = $(dc.reduction_factor))") -Base.show(io::IO, dc::LillyCoefficient) = print(io, "LillyCoefficient with\n", - "├── Smagorinsky coefficient = ", dc.smagorinsky, "\n", - "└── reduction_factor = ", dc.reduction_factor) - - diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/scale_invariant_operators.jl b/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/scale_invariant_operators.jl deleted file mode 100644 index 0a84c9c6cc..0000000000 --- a/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/scale_invariant_operators.jl +++ /dev/null @@ -1,181 +0,0 @@ -using Oceananigans.TurbulenceClosures: Σ₁₁, Σ₂₂, Σ₃₃, Σ₁₂, Σ₁₃, Σ₂₃ -using Oceananigans.TurbulenceClosures: tr_Σ², Σ₁₂², Σ₁₃², Σ₂₃² -using Oceananigans.Operators: volume - -##### -##### Double dot product of strain on cell edges (currently unused) -##### - -"Return the double dot product of strain at `ccc`." -@inline ΣᵢⱼΣᵢⱼᶜᶜᶜ(i, j, k, grid, u, v, w) = tr_Σ²(i, j, k, grid, u, v, w) + - 2 * ℑxyᶜᶜᵃ(i, j, k, grid, Σ₁₂², u, v, w) + - 2 * ℑxzᶜᵃᶜ(i, j, k, grid, Σ₁₃², u, v, w) + - 2 * ℑyzᵃᶜᶜ(i, j, k, grid, Σ₂₃², u, v, w) - -"Return the double dot product of strain at `ffc`." -@inline ΣᵢⱼΣᵢⱼᶠᶠᶜ(i, j, k, grid, u, v, w) = ℑxyᶠᶠᵃ(i, j, k, grid, tr_Σ², u, v, w) + - 2 * Σ₁₂²(i, j, k, grid, u, v, w) + - 2 * ℑyzᵃᶠᶜ(i, j, k, grid, Σ₁₃², u, v, w) + - 2 * ℑxzᶠᵃᶜ(i, j, k, grid, Σ₂₃², u, v, w) - -"Return the double dot product of strain at `fcf`." -@inline ΣᵢⱼΣᵢⱼᶠᶜᶠ(i, j, k, grid, u, v, w) = ℑxzᶠᵃᶠ(i, j, k, grid, tr_Σ², u, v, w) + - 2 * ℑyzᵃᶜᶠ(i, j, k, grid, Σ₁₂², u, v, w) + - 2 * Σ₁₃²(i, j, k, grid, u, v, w) + - 2 * ℑxyᶠᶜᵃ(i, j, k, grid, Σ₂₃², u, v, w) - -"Return the double dot product of strain at `cff`." -@inline ΣᵢⱼΣᵢⱼᶜᶠᶠ(i, j, k, grid, u, v, w) = ℑyzᵃᶠᶠ(i, j, k, grid, tr_Σ², u, v, w) + - 2 * ℑxzᶜᵃᶠ(i, j, k, grid, Σ₁₂², u, v, w) + - 2 * ℑxyᶜᶠᵃ(i, j, k, grid, Σ₁₃², u, v, w) + - 2 * Σ₂₃²(i, j, k, grid, u, v, w) - -"Return the double dot product of strain at `ccf`." -@inline ΣᵢⱼΣᵢⱼᶜᶜᶠ(i, j, k, grid, u, v, w) = ℑzᵃᵃᶠ(i, j, k, grid, tr_Σ², u, v, w) + - 2 * ℑxyzᶜᶜᶠ(i, j, k, grid, Σ₁₂², u, v, w) + - 2 * ℑxᶜᵃᵃ(i, j, k, grid, Σ₁₃², u, v, w) + - 2 * ℑyᵃᶜᵃ(i, j, k, grid, Σ₂₃², u, v, w) - - -##### -##### Filtering -##### - -@inline ℱ₂ᶠᶜᶜ(i, j, k, grid, f, args...) = ℑxyzᶠᶜᶜ(i, j, k, grid, ℑxyzᶜᶠᶠ, f, args...) -@inline ℱ₂ᶜᶠᶜ(i, j, k, grid, f, args...) = ℑxyzᶜᶠᶜ(i, j, k, grid, ℑxyzᶠᶜᶠ, f, args...) -@inline ℱ₂ᶜᶜᶠ(i, j, k, grid, f, args...) = ℑxyzᶜᶜᶠ(i, j, k, grid, ℑxyzᶠᶠᶜ, f, args...) -@inline ℱ₂ᶜᶜᶜ(i, j, k, grid, f, args...) = ℑxyzᶜᶜᶜ(i, j, k, grid, ℑxyzᶠᶠᶠ, f, args...) - -##### -##### Velocity gradients -##### - -# Diagonal -@inline ∂x_ū(i, j, k, grid, u) = ∂xᶜᶜᶜ(i, j, k, grid, ℱ₂ᶠᶜᶜ, u) -@inline ∂y_v̄(i, j, k, grid, v) = ∂yᶜᶜᶜ(i, j, k, grid, ℱ₂ᶜᶠᶜ, v) -@inline ∂z_w̄(i, j, k, grid, w) = ∂zᶜᶜᶜ(i, j, k, grid, ℱ₂ᶜᶜᶠ, w) - -# Off-diagonal -@inline ∂x_v̄(i, j, k, grid, v) = ∂xᶠᶠᶜ(i, j, k, grid, ℱ₂ᶜᶠᶜ, v) -@inline ∂x_w̄(i, j, k, grid, w) = ∂xᶠᶜᶠ(i, j, k, grid, ℱ₂ᶜᶜᶠ, w) - -@inline ∂y_ū(i, j, k, grid, u) = ∂yᶠᶠᶜ(i, j, k, grid, ℱ₂ᶠᶜᶜ, u) -@inline ∂y_w̄(i, j, k, grid, w) = ∂yᶜᶠᶠ(i, j, k, grid, ℱ₂ᶜᶜᶠ, w) - -@inline ∂z_ū(i, j, k, grid, u) = ∂zᶠᶜᶠ(i, j, k, grid, ℱ₂ᶠᶜᶜ, u) -@inline ∂z_v̄(i, j, k, grid, v) = ∂zᶜᶠᶠ(i, j, k, grid, ℱ₂ᶜᶠᶜ, v) - -##### -##### Strain components -##### - -const AG = AbstractGrid - -# ccc strain components -@inline Σ̄₁₁(i, j, k, grid, u) = ∂x_ū(i, j, k, grid, u) -@inline Σ̄₂₂(i, j, k, grid, v) = ∂y_v̄(i, j, k, grid, v) -@inline Σ̄₃₃(i, j, k, grid, w) = ∂z_w̄(i, j, k, grid, w) - -@inline tr_Σ̄(i, j, k, grid, u, v, w) = Σ̄₁₁(i, j, k, grid, u) + Σ̄₂₂(i, j, k, grid, v) + Σ̄₃₃(i, j, k, grid, w) -@inline tr_Σ̄²(ijk...) = Σ̄₁₁(ijk...)^2 + Σ̄₂₂(ijk...)^2 + Σ̄₃₃(ijk...)^2 - -# ffc -@inline Σ̄₁₂(i, j, k, grid::AG{FT}, u, v) where FT = FT(0.5) * (∂y_ū(i, j, k, grid, u) + ∂x_v̄(i, j, k, grid, v)) -@inline Σ̄₁₂²(i, j, k, grid, u, v) = Σ̄₁₂(i, j, k, grid, u, v)^2 - -# fcf -@inline Σ̄₁₃(i, j, k, grid::AG{FT}, u, w) where FT = FT(0.5) * (∂z_ū(i, j, k, grid, u) + ∂x_w̄(i, j, k, grid, w)) -@inline Σ̄₁₃²(i, j, k, grid, u, w) = Σ̄₁₃(i, j, k, grid, u, w)^2 - -# cff -@inline Σ̄₂₃(i, j, k, grid::AG{FT}, v, w) where FT = FT(0.5) * (∂z_v̄(i, j, k, grid, v) + ∂y_w̄(i, j, k, grid, w)) -@inline Σ̄₂₃²(i, j, k, grid, v, w) = Σ̄₂₃(i, j, k, grid, v, w)^2 - -@inline Σ̄₁₁(i, j, k, grid, u, v, w) = Σ̄₁₁(i, j, k, grid, u) -@inline Σ̄₂₂(i, j, k, grid, u, v, w) = Σ̄₂₂(i, j, k, grid, v) -@inline Σ̄₃₃(i, j, k, grid, u, v, w) = Σ̄₃₃(i, j, k, grid, w) - -@inline Σ̄₁₂(i, j, k, grid, u, v, w) = Σ̄₁₂(i, j, k, grid, u, v) -@inline Σ̄₁₃(i, j, k, grid, u, v, w) = Σ̄₁₃(i, j, k, grid, u, w) -@inline Σ̄₂₃(i, j, k, grid, u, v, w) = Σ̄₂₃(i, j, k, grid, v, w) - -@inline Σ̄₁₂²(i, j, k, grid, u, v, w) = Σ̄₁₂²(i, j, k, grid, u, v) -@inline Σ̄₁₃²(i, j, k, grid, u, v, w) = Σ̄₁₃²(i, j, k, grid, u, w) -@inline Σ̄₂₃²(i, j, k, grid, u, v, w) = Σ̄₂₃²(i, j, k, grid, v, w) - -##### -##### Double dot product of strain on cell edges -##### - -"Return the double dot product of strain at `ccc` on a 2δ test grid." -@inline Σ̄ᵢⱼΣ̄ᵢⱼᶜᶜᶜ(i, j, k, grid, u, v, w) = (tr_Σ̄²(i, j, k, grid, u, v, w) - + 2 * ℑxyᶜᶜᵃ(i, j, k, grid, Σ̄₁₂², u, v, w) - + 2 * ℑxzᶜᵃᶜ(i, j, k, grid, Σ̄₁₃², u, v, w) - + 2 * ℑyzᵃᶜᶜ(i, j, k, grid, Σ̄₂₃², u, v, w)) - -# Here the notation ⟨A⟩ is equivalent to Ā: a filter of size 2Δᶠ, where Δᶠ is the grid scale. - -@inline ΣΣ₁₁ᶜᶜᶜ(i, j, k, grid, u, v, w, Σᶜᶜᶜ) = @inbounds Σᶜᶜᶜ[i, j, k] * Σ₁₁(i, j, k, grid, u, v, w) -@inline ΣΣ₂₂ᶜᶜᶜ(i, j, k, grid, u, v, w, Σᶜᶜᶜ) = @inbounds Σᶜᶜᶜ[i, j, k] * Σ₂₂(i, j, k, grid, u, v, w) -@inline ΣΣ₃₃ᶜᶜᶜ(i, j, k, grid, u, v, w, Σᶜᶜᶜ) = @inbounds Σᶜᶜᶜ[i, j, k] * Σ₃₃(i, j, k, grid, u, v, w) - -@inline ΣΣ₁₂ᶜᶜᶜ(i, j, k, grid, u, v, w, Σᶜᶜᶜ) = @inbounds Σᶜᶜᶜ[i, j, k] * ℑxyᶜᶜᵃ(i, j, k, grid, Σ₁₂, u, v, w) -@inline ΣΣ₁₃ᶜᶜᶜ(i, j, k, grid, u, v, w, Σᶜᶜᶜ) = @inbounds Σᶜᶜᶜ[i, j, k] * ℑxzᶜᵃᶜ(i, j, k, grid, Σ₁₃, u, v, w) -@inline ΣΣ₂₃ᶜᶜᶜ(i, j, k, grid, u, v, w, Σᶜᶜᶜ) = @inbounds Σᶜᶜᶜ[i, j, k] * ℑyzᵃᶜᶜ(i, j, k, grid, Σ₂₃, u, v, w) - -@inline var"⟨ΣΣ₁₁⟩ᶜᶜᶜ"(i, j, k, grid, u, v, w, Σᶜᶜᶜ) = ℱ₂ᶜᶜᶜ(i, j, k, grid, ΣΣ₁₁ᶜᶜᶜ, u, v, w, Σᶜᶜᶜ) -@inline var"⟨ΣΣ₂₂⟩ᶜᶜᶜ"(i, j, k, grid, u, v, w, Σᶜᶜᶜ) = ℱ₂ᶜᶜᶜ(i, j, k, grid, ΣΣ₂₂ᶜᶜᶜ, u, v, w, Σᶜᶜᶜ) -@inline var"⟨ΣΣ₃₃⟩ᶜᶜᶜ"(i, j, k, grid, u, v, w, Σᶜᶜᶜ) = ℱ₂ᶜᶜᶜ(i, j, k, grid, ΣΣ₃₃ᶜᶜᶜ, u, v, w, Σᶜᶜᶜ) - -@inline var"⟨ΣΣ₁₂⟩ᶜᶜᶜ"(i, j, k, grid, u, v, w, Σᶜᶜᶜ) = ℱ₂ᶜᶜᶜ(i, j, k, grid, ΣΣ₁₂ᶜᶜᶜ, u, v, w, Σᶜᶜᶜ) -@inline var"⟨ΣΣ₁₃⟩ᶜᶜᶜ"(i, j, k, grid, u, v, w, Σᶜᶜᶜ) = ℱ₂ᶜᶜᶜ(i, j, k, grid, ΣΣ₁₃ᶜᶜᶜ, u, v, w, Σᶜᶜᶜ) -@inline var"⟨ΣΣ₂₃⟩ᶜᶜᶜ"(i, j, k, grid, u, v, w, Σᶜᶜᶜ) = ℱ₂ᶜᶜᶜ(i, j, k, grid, ΣΣ₂₃ᶜᶜᶜ, u, v, w, Σᶜᶜᶜ) - -@inline Σ̄Σ̄₁₁ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ̄ᶜᶜᶜ) = @inbounds Σ̄ᶜᶜᶜ[i, j, k] * Σ̄₁₁(i, j, k, grid, u, v, w) -@inline Σ̄Σ̄₂₂ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ̄ᶜᶜᶜ) = @inbounds Σ̄ᶜᶜᶜ[i, j, k] * Σ̄₂₂(i, j, k, grid, u, v, w) -@inline Σ̄Σ̄₃₃ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ̄ᶜᶜᶜ) = @inbounds Σ̄ᶜᶜᶜ[i, j, k] * Σ̄₃₃(i, j, k, grid, u, v, w) - -@inline Σ̄Σ̄₁₂ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ̄ᶜᶜᶜ) = @inbounds Σ̄ᶜᶜᶜ[i, j, k] * ℑxyᶜᶜᵃ(i, j, k, grid, Σ̄₁₂, u, v, w) -@inline Σ̄Σ̄₁₃ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ̄ᶜᶜᶜ) = @inbounds Σ̄ᶜᶜᶜ[i, j, k] * ℑxzᶜᵃᶜ(i, j, k, grid, Σ̄₁₃, u, v, w) -@inline Σ̄Σ̄₂₃ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ̄ᶜᶜᶜ) = @inbounds Σ̄ᶜᶜᶜ[i, j, k] * ℑyzᵃᶜᶜ(i, j, k, grid, Σ̄₂₃, u, v, w) - -const ᾱ² = 4 -const β = 1 -@inline Δᶠ(i, j, k, grid) = ∛volume(i, j, k, grid, Center(), Center(), Center()) -@inline M₁₁ᶜᶜᶜ(i, j, k, grid, u, v, w, Σᶜᶜᶜ, Σ̄ᶜᶜᶜ) = 2*Δᶠ(i, j, k, grid)^2 * (var"⟨ΣΣ₁₁⟩ᶜᶜᶜ"(i, j, k, grid, u, v, w, Σᶜᶜᶜ) - ᾱ² * β * Σ̄Σ̄₁₁ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ̄ᶜᶜᶜ)) -@inline M₂₂ᶜᶜᶜ(i, j, k, grid, u, v, w, Σᶜᶜᶜ, Σ̄ᶜᶜᶜ) = 2*Δᶠ(i, j, k, grid)^2 * (var"⟨ΣΣ₂₂⟩ᶜᶜᶜ"(i, j, k, grid, u, v, w, Σᶜᶜᶜ) - ᾱ² * β * Σ̄Σ̄₂₂ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ̄ᶜᶜᶜ)) -@inline M₃₃ᶜᶜᶜ(i, j, k, grid, u, v, w, Σᶜᶜᶜ, Σ̄ᶜᶜᶜ) = 2*Δᶠ(i, j, k, grid)^2 * (var"⟨ΣΣ₃₃⟩ᶜᶜᶜ"(i, j, k, grid, u, v, w, Σᶜᶜᶜ) - ᾱ² * β * Σ̄Σ̄₃₃ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ̄ᶜᶜᶜ)) - -@inline M₁₂ᶜᶜᶜ(i, j, k, grid, u, v, w, Σᶜᶜᶜ, Σ̄ᶜᶜᶜ) = 2*Δᶠ(i, j, k, grid)^2 * (var"⟨ΣΣ₁₂⟩ᶜᶜᶜ"(i, j, k, grid, u, v, w, Σᶜᶜᶜ) - ᾱ² * β * Σ̄Σ̄₁₂ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ̄ᶜᶜᶜ)) -@inline M₁₃ᶜᶜᶜ(i, j, k, grid, u, v, w, Σᶜᶜᶜ, Σ̄ᶜᶜᶜ) = 2*Δᶠ(i, j, k, grid)^2 * (var"⟨ΣΣ₁₃⟩ᶜᶜᶜ"(i, j, k, grid, u, v, w, Σᶜᶜᶜ) - ᾱ² * β * Σ̄Σ̄₁₃ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ̄ᶜᶜᶜ)) -@inline M₂₃ᶜᶜᶜ(i, j, k, grid, u, v, w, Σᶜᶜᶜ, Σ̄ᶜᶜᶜ) = 2*Δᶠ(i, j, k, grid)^2 * (var"⟨ΣΣ₂₃⟩ᶜᶜᶜ"(i, j, k, grid, u, v, w, Σᶜᶜᶜ) - ᾱ² * β * Σ̄Σ̄₂₃ᶜᶜᶜ(i, j, k, grid, u, v, w, Σ̄ᶜᶜᶜ)) - -@inline uᵢ²(i, j, k, grid, uᵢ) = @inbounds uᵢ[i, j, k]^2 -@inline u₁u₁ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℑxᶜᵃᵃ(i, j, k, grid, uᵢ², u) -@inline u₂u₂ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℑyᵃᶜᵃ(i, j, k, grid, uᵢ², v) -@inline u₃u₃ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℑzᵃᵃᶜ(i, j, k, grid, uᵢ², w) - -@inline ū²(i, j, k, grid, u) = ℱ₂ᶠᶜᶜ(i, j, k, grid, u)^2 -@inline v̄²(i, j, k, grid, v) = ℱ₂ᶜᶠᶜ(i, j, k, grid, v)^2 -@inline w̄²(i, j, k, grid, w) = ℱ₂ᶜᶜᶠ(i, j, k, grid, w)^2 - -@inline ū₁ū₁ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℑxᶜᵃᵃ(i, j, k, grid, ū², u) -@inline ū₂ū₂ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℑyᵃᶜᵃ(i, j, k, grid, v̄², v) -@inline ū₃ū₃ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℑzᵃᵃᶜ(i, j, k, grid, w̄², w) - -@inline u₁u₂ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℑxᶜᵃᵃ(i, j, k, grid, u) * ℑyᵃᶜᵃ(i, j, k, grid, v) -@inline u₁u₃ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℑxᶜᵃᵃ(i, j, k, grid, u) * ℑzᵃᵃᶜ(i, j, k, grid, w) -@inline u₂u₃ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℑyᵃᶜᵃ(i, j, k, grid, v) * ℑzᵃᵃᶜ(i, j, k, grid, w) - -@inline ū₁ū₂ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℑxᶜᵃᵃ(i, j, k, grid, ℱ₂ᶠᶜᶜ, u) * ℑyᵃᶜᵃ(i, j, k, grid, ℱ₂ᶜᶠᶜ, v) -@inline ū₁ū₃ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℑxᶜᵃᵃ(i, j, k, grid, ℱ₂ᶠᶜᶜ, u) * ℑzᵃᵃᶜ(i, j, k, grid, ℱ₂ᶜᶜᶠ, w) -@inline ū₂ū₃ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℑyᵃᶜᵃ(i, j, k, grid, ℱ₂ᶜᶠᶜ, v) * ℑzᵃᵃᶜ(i, j, k, grid, ℱ₂ᶜᶜᶠ, w) - -@inline L₁₁ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℱ₂ᶜᶜᶜ(i, j, k, grid, u₁u₁ᶜᶜᶜ, u, v, w) - ū₁ū₁ᶜᶜᶜ(i, j, k, grid, u, v, w) -@inline L₂₂ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℱ₂ᶜᶜᶜ(i, j, k, grid, u₂u₂ᶜᶜᶜ, u, v, w) - ū₂ū₂ᶜᶜᶜ(i, j, k, grid, u, v, w) -@inline L₃₃ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℱ₂ᶜᶜᶜ(i, j, k, grid, u₃u₃ᶜᶜᶜ, u, v, w) - ū₃ū₃ᶜᶜᶜ(i, j, k, grid, u, v, w) - -@inline L₁₂ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℱ₂ᶜᶜᶜ(i, j, k, grid, u₁u₂ᶜᶜᶜ, u, v, w) - ū₁ū₂ᶜᶜᶜ(i, j, k, grid, u, v, w) -@inline L₁₃ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℱ₂ᶜᶜᶜ(i, j, k, grid, u₁u₃ᶜᶜᶜ, u, v, w) - ū₁ū₃ᶜᶜᶜ(i, j, k, grid, u, v, w) -@inline L₂₃ᶜᶜᶜ(i, j, k, grid, u, v, w) = ℱ₂ᶜᶜᶜ(i, j, k, grid, u₂u₃ᶜᶜᶜ, u, v, w) - ū₂ū₃ᶜᶜᶜ(i, j, k, grid, u, v, w) - diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/smagorinsky.jl b/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/smagorinsky.jl deleted file mode 100644 index d7968017b0..0000000000 --- a/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/smagorinsky.jl +++ /dev/null @@ -1,147 +0,0 @@ -using Oceananigans.AbstractOperations: Average -using Oceananigans.Fields: FieldBoundaryConditions -using Oceananigans.Utils: launch!, IterationInterval - -using Adapt - -using ..TurbulenceClosures: - AbstractScalarDiffusivity, - ThreeDimensionalFormulation, - ExplicitTimeDiscretization, - convert_diffusivity - -import Oceananigans.Utils: with_tracers - -import ..TurbulenceClosures: - viscosity, - diffusivity, - κᶠᶜᶜ, - κᶜᶠᶜ, - κᶜᶜᶠ, - compute_diffusivities!, - DiffusivityFields, - tracer_diffusivities - -##### -##### The turbulence closure proposed by Smagorinsky and Lilly. -##### - -# struct LagrangianAveragedCoefficient end - -struct Smagorinsky{TD, C, P} <: AbstractScalarDiffusivity{TD, ThreeDimensionalFormulation, 2} - coefficient :: C - Pr :: P - - function Smagorinsky{TD}(coefficient, Pr) where TD - P = typeof(Pr) - C = typeof(coefficient) - return new{TD, C, P}(coefficient, Pr) - end -end - -@inline viscosity(::Smagorinsky, K) = K.νₑ -@inline diffusivity(closure::Smagorinsky, K, ::Val{id}) where id = K.νₑ / closure.Pr[id] - -const ConstantSmagorinsky = Smagorinsky{<:Any, <:Number} - -""" - Smagorinsky([time_discretization::TD = ExplicitTimeDiscretization(), FT=Float64;] - coefficient = 0.16, - Pr = 1.0) - -Return a `Smagorinsky` type associated with the turbulence closure proposed by -[Smagorinsky1958](@citet) and [Smagorinsky1963](@citet) -which has an eddy viscosity of the form - -``` -νₑ = (Cˢ * Δᶠ)² * √(2Σ²) -``` - -and an eddy diffusivity of the form - -``` -κₑ = νₑ / Pr. -``` - -where `Δᶠ` is the filter width, `Σ² = ΣᵢⱼΣᵢⱼ` is the double dot product of -the strain tensor `Σᵢⱼ`, `Pr` is the turbulent Prandtl number, `N²` is the -total buoyancy gradient, and `Cb` is a constant the multiplies the Richardson -number modification to the eddy viscosity. - -`Cˢ` is the Smagorinsky coefficient and the default value is 0.16, according -to the analysis by [Lilly66](@citet). For other options, see `LillyCoefficient` -and `DynamicCoefficient`. -""" -function Smagorinsky(time_discretization::TD = ExplicitTimeDiscretization(), FT=Float64; - coefficient = 0.16, Pr = 1.0) where TD - Pr = convert_diffusivity(FT, Pr; discrete_form=false) - return Smagorinsky{TD}(coefficient, Pr) -end - -Smagorinsky(FT::DataType; kwargs...) = Smagorinsky(ExplicitTimeDiscretization(), FT; kwargs...) - -function with_tracers(tracers, closure::Smagorinsky{TD}) where TD - Pr = tracer_diffusivities(tracers, closure.Pr) - return Smagorinsky{TD}(closure.coefficient, Pr) -end - -@kernel function _compute_smagorinsky_viscosity!(diffusivity_fields, grid, closure, buoyancy, velocities, tracers) - i, j, k = @index(Global, NTuple) - - # Strain tensor dot product - Σ² = ΣᵢⱼΣᵢⱼᶜᶜᶜ(i, j, k, grid, velocities.u, velocities.v, velocities.w) - - # Filter width - Δ³ = Δxᶜᶜᶜ(i, j, k, grid) * Δyᶜᶜᶜ(i, j, k, grid) * Δzᶜᶜᶜ(i, j, k, grid) - Δᶠ = cbrt(Δ³) - cˢ² = square_smagorinsky_coefficient(i, j, k, grid, closure, diffusivity_fields, Σ², buoyancy, tracers) - - νₑ = diffusivity_fields.νₑ - - @inbounds νₑ[i, j, k] = cˢ² * Δᶠ^2 * sqrt(2Σ²) -end - -@inline square_smagorinsky_coefficient(i, j, k, grid, c::ConstantSmagorinsky, args...) = c.coefficient^2 - -compute_coefficient_fields!(diffusivity_fields, closure, model; parameters) = nothing - -function compute_diffusivities!(diffusivity_fields, closure::Smagorinsky, model; parameters = :xyz) - arch = model.architecture - grid = model.grid - buoyancy = model.buoyancy - velocities = model.velocities - tracers = model.tracers - - compute_coefficient_fields!(diffusivity_fields, closure, model; parameters) - - launch!(arch, grid, parameters, _compute_smagorinsky_viscosity!, - diffusivity_fields, grid, closure, buoyancy, velocities, tracers) - - return nothing -end - -allocate_coefficient_fields(closure, grid) = NamedTuple() - -function DiffusivityFields(grid, tracer_names, bcs, closure::Smagorinsky) - coefficient_fields = allocate_coefficient_fields(closure, grid) - - default_eddy_viscosity_bcs = (; νₑ = FieldBoundaryConditions(grid, (Center, Center, Center))) - bcs = merge(default_eddy_viscosity_bcs, bcs) - νₑ = CenterField(grid, boundary_conditions=bcs.νₑ) - viscosity_nt = (; νₑ) - - return merge(viscosity_nt, coefficient_fields) -end - -@inline κᶠᶜᶜ(i, j, k, grid, c::Smagorinsky, K, ::Val{id}, args...) where id = ℑxᶠᵃᵃ(i, j, k, grid, K.νₑ) / c.Pr[id] -@inline κᶜᶠᶜ(i, j, k, grid, c::Smagorinsky, K, ::Val{id}, args...) where id = ℑyᵃᶠᵃ(i, j, k, grid, K.νₑ) / c.Pr[id] -@inline κᶜᶜᶠ(i, j, k, grid, c::Smagorinsky, K, ::Val{id}, args...) where id = ℑzᵃᵃᶠ(i, j, k, grid, K.νₑ) / c.Pr[id] - -Base.summary(closure::Smagorinsky) = string("Smagorinsky with coefficient = ", summary(closure.coefficient), ", Pr=$(closure.Pr)") -function Base.show(io::IO, closure::Smagorinsky) - coefficient_summary = closure.coefficient isa Number ? closure.coefficient : summary(closure.coefficient) - print(io, "Smagorinsky closure with\n", - "├── coefficient = ", coefficient_summary, "\n", - "└── Pr = ", closure.Pr) -end - diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_equation.jl b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_equation.jl index bf2aaa7bfd..a08fb4b969 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_equation.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_equation.jl @@ -58,7 +58,7 @@ end ℓ★ = ifelse(isnan(ℓ★), zero(grid), ℓ★) ℓᴰ = max(ℓ★, ℓʰ) - H = static_column_depthᶜᶜᵃ(i, j, grid) + H = total_depthᶜᶜᵃ(i, j, grid) return min(H, ℓᴰ) end diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_mixing_length.jl b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_mixing_length.jl index dce68696b6..8b2620f6b0 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_mixing_length.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_mixing_length.jl @@ -5,7 +5,7 @@ using ..TurbulenceClosures: height_above_bottomᶜᶜᶠ, depthᶜᶜᶜ, height_above_bottomᶜᶜᶜ, - static_column_depthᶜᶜᵃ + total_depthᶜᶜᵃ """ struct CATKEMixingLength{FT} @@ -16,7 +16,7 @@ Base.@kwdef struct CATKEMixingLength{FT} Cˢ :: FT = 1.131 # Surface distance coefficient for shear length scale Cᵇ :: FT = Inf # Bottom distance coefficient for shear length scale Cˢᵖ :: FT = 0.505 # Sheared convective plume coefficient - CRiᵟ :: FT = 1.02 # Stability function width + CRiᵟ :: FT = 1.02 # Stability function width CRi⁰ :: FT = 0.254 # Stability function lower Ri Cʰⁱu :: FT = 0.242 # Shear mixing length coefficient for momentum at high Ri Cˡᵒu :: FT = 0.361 # Shear mixing length coefficient for momentum at low Ri @@ -131,7 +131,7 @@ end ϵˢᵖ = 1 - Cˢᵖ * Riᶠ # ϵ = Sheared convection factor ℓᵉ = clip(ϵˢᵖ * ℓᵉ) =# - + # Figure out which mixing length applies convecting = (Jᵇ > Jᵇᵋ) & (N² < 0) entraining = (Jᵇ > Jᵇᵋ) & (N² > 0) & (N²_above < 0) @@ -232,7 +232,7 @@ end ℓ★ = ifelse(isnan(ℓ★), zero(grid), ℓ★) ℓu = max(ℓ★, ℓʰ) - H = static_column_depthᶜᶜᵃ(i, j, grid) + H = total_depthᶜᶜᵃ(i, j, grid) return min(H, ℓu) end @@ -252,7 +252,7 @@ end ℓ★ = ifelse(isnan(ℓ★), zero(grid), ℓ★) ℓc = max(ℓ★, ℓʰ) - H = static_column_depthᶜᶜᵃ(i, j, grid) + H = total_depthᶜᶜᵃ(i, j, grid) return min(H, ℓc) end @@ -272,7 +272,7 @@ end ℓ★ = ifelse(isnan(ℓ★), zero(grid), ℓ★) ℓe = max(ℓ★, ℓʰ) - H = static_column_depthᶜᶜᵃ(i, j, grid) + H = total_depthᶜᶜᵃ(i, j, grid) return min(H, ℓe) end @@ -299,3 +299,4 @@ Base.show(io::IO, ml::CATKEMixingLength) = " ├── Cˢᵖ: ", ml.Cˢᵖ, '\n', " ├── CRiᵟ: ", ml.CRiᵟ, '\n', " └── CRi⁰: ", ml.CRi⁰) + diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl index 4847a3802a..2adf895c82 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl @@ -1,3 +1,4 @@ + struct CATKEVerticalDiffusivity{TD, CL, FT, DT, TKE} <: AbstractScalarDiffusivity{TD, VerticalFormulation, 2} mixing_length :: CL turbulent_kinetic_energy_equation :: TKE @@ -17,7 +18,7 @@ function CATKEVerticalDiffusivity{TD}(mixing_length::CL, maximum_viscosity::FT, minimum_tke::FT, minimum_convective_buoyancy_flux::FT, - negative_tke_damping_time_scale::FT, + negative_tke_damping_time_scale::FT, tke_time_step::DT) where {TD, CL, FT, DT, TKE} return CATKEVerticalDiffusivity{TD, CL, FT, DT, TKE}(mixing_length, @@ -172,7 +173,7 @@ function DiffusivityFields(grid, tracer_names, bcs, closure::FlavorOfCATKE) return (; κu, κc, κe, Le, Jᵇ, previous_compute_time, previous_velocities, _tupled_tracer_diffusivities, _tupled_implicit_linear_coefficients) -end +end @inline viscosity_location(::FlavorOfCATKE) = (c, c, f) @inline diffusivity_location(::FlavorOfCATKE) = (c, c, f) @@ -227,7 +228,7 @@ end Jᵇᵋ = closure.minimum_convective_buoyancy_flux Jᵇᵢⱼ = @inbounds Jᵇ[i, j, 1] Jᵇ⁺ = max(Jᵇᵋ, Jᵇᵢⱼ, Jᵇ★) # selects fastest (dominant) time-scale - t★ = cbrt(ℓᴰ^2 / Jᵇ⁺) + t★ = (ℓᴰ^2 / Jᵇ⁺)^(1/3) ϵ = Δt / t★ @inbounds Jᵇ[i, j, 1] = (Jᵇᵢⱼ + ϵ * Jᵇ★) / (1 + ϵ) @@ -262,9 +263,7 @@ end ℓu = momentum_mixing_lengthᶜᶜᶠ(i, j, k, grid, closure, velocities, tracers, buoyancy, surface_buoyancy_flux) κu = ℓu * w★ κu_max = closure.maximum_viscosity - κu★ = min(κu, κu_max) - FT = eltype(grid) - return κu★::FT + return min(κu, κu_max) end @inline function κcᶜᶜᶠ(i, j, k, grid, closure, velocities, tracers, buoyancy, surface_buoyancy_flux) @@ -272,9 +271,7 @@ end ℓc = tracer_mixing_lengthᶜᶜᶠ(i, j, k, grid, closure, velocities, tracers, buoyancy, surface_buoyancy_flux) κc = ℓc * w★ κc_max = closure.maximum_tracer_diffusivity - κc★ = min(κc, κc_max) - FT = eltype(grid) - return κc★::FT + return min(κc, κc_max) end @inline function κeᶜᶜᶠ(i, j, k, grid, closure, velocities, tracers, buoyancy, surface_buoyancy_flux) @@ -282,14 +279,12 @@ end ℓe = TKE_mixing_lengthᶜᶜᶠ(i, j, k, grid, closure, velocities, tracers, buoyancy, surface_buoyancy_flux) κe = ℓe * w★ κe_max = closure.maximum_tke_diffusivity - κe★ = min(κe, κe_max) - FT = eltype(grid) - return κe★::FT + return min(κe, κe_max) end @inline viscosity(::FlavorOfCATKE, diffusivities) = diffusivities.κu @inline diffusivity(::FlavorOfCATKE, diffusivities, ::Val{id}) where id = diffusivities._tupled_tracer_diffusivities[id] - + ##### ##### Show ##### @@ -339,3 +334,4 @@ function Base.show(io::IO, clo::CATKEVD) " ├── CᵂwΔ: ", prettysummary(clo.turbulent_kinetic_energy_equation.CᵂwΔ), '\n', " └── Cᵂϵ: ", prettysummary(clo.turbulent_kinetic_energy_equation.Cᵂϵ)) end + diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/time_step_catke_equation.jl b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/time_step_catke_equation.jl index 892a8c041c..6658bd56f1 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/time_step_catke_equation.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/time_step_catke_equation.jl @@ -117,7 +117,7 @@ end # Then the contribution of Jᵉ to the implicit flux is # # Lᵂ = - Cᵂϵ * √e / Δz. - + on_bottom = !inactive_cell(i, j, k, grid) & inactive_cell(i, j, k-1, grid) Δz = Δzᶜᶜᶜ(i, j, k, grid) Cᵂϵ = closure_ij.turbulent_kinetic_energy_equation.Cᵂϵ @@ -190,7 +190,7 @@ function tracer_tendency_kernel_function(model::HFSM, ::Val{:e}, closures::Tuple else catke_closure = closures[catke_index] catke_diffusivity_fields = diffusivity_fields[catke_index] - return compute_hydrostatic_free_surface_Ge!, catke_closure, catke_diffusivity_fields + return compute_hydrostatic_free_surface_Ge!, catke_closure, catke_diffusivity_fields end end @@ -202,3 +202,4 @@ end return NamedTuple{tuple(names...)}(tuple(values...)) end =# + diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/smagorinsky_lilly.jl b/src/TurbulenceClosures/turbulence_closure_implementations/smagorinsky_lilly.jl new file mode 100644 index 0000000000..3652ce0d38 --- /dev/null +++ b/src/TurbulenceClosures/turbulence_closure_implementations/smagorinsky_lilly.jl @@ -0,0 +1,213 @@ +##### +##### The turbulence closure proposed by Smagorinsky and Lilly. +##### We also call this 'Constant Smagorinsky'. +##### + +struct SmagorinskyLilly{TD, FT, P} <: AbstractScalarDiffusivity{TD, ThreeDimensionalFormulation, 2} + C :: FT + Cb :: FT + Pr :: P + + function SmagorinskyLilly{TD, FT}(C, Cb, Pr) where {TD, FT} + Pr = convert_diffusivity(FT, Pr; discrete_form=false) + P = typeof(Pr) + return new{TD, FT, P}(C, Cb, Pr) + end +end + +@inline viscosity(::SmagorinskyLilly, K) = K.νₑ +@inline diffusivity(closure::SmagorinskyLilly, K, ::Val{id}) where id = K.νₑ / closure.Pr[id] + +""" + SmagorinskyLilly([time_discretization::TD = ExplicitTimeDiscretization(), FT=Float64;] C=0.16, Cb=1.0, Pr=1.0) + +Return a `SmagorinskyLilly` type associated with the turbulence closure proposed by +[Lilly62](@citet), [Smagorinsky1958](@citet), [Smagorinsky1963](@citet), and [Lilly66](@citet), +which has an eddy viscosity of the form + +``` +νₑ = (C * Δᶠ)² * √(2Σ²) * √(1 - Cb * N² / Σ²) +``` + +and an eddy diffusivity of the form + +``` +κₑ = νₑ / Pr +``` + +where `Δᶠ` is the filter width, `Σ² = ΣᵢⱼΣᵢⱼ` is the double dot product of +the strain tensor `Σᵢⱼ`, `Pr` is the turbulent Prandtl number, `N²` is the +total buoyancy gradient, and `Cb` is a constant the multiplies the Richardson +number modification to the eddy viscosity. + +Arguments +========= + +* `time_discretization`: Either `ExplicitTimeDiscretization()` or `VerticallyImplicitTimeDiscretization()`, + which integrates the terms involving only ``z``-derivatives in the + viscous and diffusive fluxes with an implicit time discretization. + Default `ExplicitTimeDiscretization()`. + +* `FT`: Float type; default `Float64`. + +Keyword arguments +================= + +* `C`: Smagorinsky constant. Default value is 0.16 as obtained by Lilly (1966). + +* `Cb`: Buoyancy term multipler based on Lilly (1962) (`Cb = 0` turns it off, `Cb ≠ 0` turns it on. + Typically, and according to the original work by Lilly (1962), `Cb = 1 / Pr`.) + +* `Pr`: Turbulent Prandtl numbers for each tracer. Either a constant applied to every + tracer, or a `NamedTuple` with fields for each tracer individually. + +References +========== + +Smagorinsky, J. "On the numerical integration of the primitive equations of motion for + baroclinic flow in a closed region." Monthly Weather Review (1958) + +Lilly, D. K. "On the numerical simulation of buoyant convection." Tellus (1962) + +Smagorinsky, J. "General circulation experiments with the primitive equations: I. + The basic experiment." Monthly Weather Review (1963) + +Lilly, D. K. "The representation of small-scale turbulence in numerical simulation experiments." + NCAR Manuscript No. 281, 0, (1966) +""" +SmagorinskyLilly(time_discretization::TD = ExplicitTimeDiscretization(), FT=Float64; C=0.16, Cb=1.0, Pr=1.0) where TD = + SmagorinskyLilly{TD, FT}(C, Cb, Pr) + +SmagorinskyLilly(FT::DataType; kwargs...) = SmagorinskyLilly(ExplicitTimeDiscretization(), FT; kwargs...) + +function with_tracers(tracers, closure::SmagorinskyLilly{TD, FT}) where {TD, FT} + Pr = tracer_diffusivities(tracers, closure.Pr) + return SmagorinskyLilly{TD, FT}(closure.C, closure.Cb, Pr) +end + +""" + stability(N², Σ², Cb) + +Return the stability function + +```math + \\sqrt(1 - Cb N^2 / Σ^2 ) +``` + +when ``N^2 > 0``, and 1 otherwise. +""" +@inline function stability(N²::FT, Σ²::FT, Cb::FT) where FT + N²⁺ = max(zero(FT), N²) # clip + ς² = one(FT) - min(one(FT), Cb * N²⁺ / Σ²) + return ifelse(Σ²==0, zero(FT), sqrt(ς²)) +end + +@kernel function _compute_smagorinsky_viscosity!(νₑ, grid, closure, buoyancy, velocities, tracers) + i, j, k = @index(Global, NTuple) + + # Strain tensor dot product + Σ² = ΣᵢⱼΣᵢⱼᶜᶜᶜ(i, j, k, grid, velocities.u, velocities.v, velocities.w) + + # Stability function + N² = ℑzᵃᵃᶜ(i, j, k, grid, ∂z_b, buoyancy, tracers) + ς = stability(N², Σ², closure.Cb) # Use unity Prandtl number. + + # Filter width + Δ³ = Δxᶜᶜᶜ(i, j, k, grid) * Δyᶜᶜᶜ(i, j, k, grid) * Δzᶜᶜᶜ(i, j, k, grid) + Δᶠ = cbrt(Δ³) + C = closure.C # free parameter + + @inbounds νₑ[i, j, k] = ς * (C * Δᶠ)^2 * sqrt(2Σ²) +end + +function compute_diffusivities!(diffusivity_fields, closure::SmagorinskyLilly, model; parameters = :xyz) + arch = model.architecture + grid = model.grid + buoyancy = model.buoyancy + velocities = model.velocities + tracers = model.tracers + + launch!(arch, grid, parameters, _compute_smagorinsky_viscosity!, + diffusivity_fields.νₑ, grid, closure, buoyancy, velocities, tracers) + + return nothing +end + +@inline κᶠᶜᶜ(i, j, k, grid, closure::SmagorinskyLilly, K, ::Val{id}, args...) where id = ℑxᶠᵃᵃ(i, j, k, grid, K.νₑ) / closure.Pr[id] +@inline κᶜᶠᶜ(i, j, k, grid, closure::SmagorinskyLilly, K, ::Val{id}, args...) where id = ℑyᵃᶠᵃ(i, j, k, grid, K.νₑ) / closure.Pr[id] +@inline κᶜᶜᶠ(i, j, k, grid, closure::SmagorinskyLilly, K, ::Val{id}, args...) where id = ℑzᵃᵃᶠ(i, j, k, grid, K.νₑ) / closure.Pr[id] + +##### +##### Double dot product of strain on cell edges (currently unused) +##### + +# tr_Σ² : ccc +# Σ₁₂ : ffc +# Σ₁₃ : fcf +# Σ₂₃ : cff + +"Return the double dot product of strain at `ccc`." +@inline function ΣᵢⱼΣᵢⱼᶜᶜᶜ(i, j, k, grid, u, v, w) + return ( + tr_Σ²(i, j, k, grid, u, v, w) + + 2 * ℑxyᶜᶜᵃ(i, j, k, grid, Σ₁₂², u, v, w) + + 2 * ℑxzᶜᵃᶜ(i, j, k, grid, Σ₁₃², u, v, w) + + 2 * ℑyzᵃᶜᶜ(i, j, k, grid, Σ₂₃², u, v, w) + ) +end + +"Return the double dot product of strain at `ffc`." +@inline function ΣᵢⱼΣᵢⱼᶠᶠᶜ(i, j, k, grid, u, v, w) + return ( + ℑxyᶠᶠᵃ(i, j, k, grid, tr_Σ², u, v, w) + + 2 * Σ₁₂²(i, j, k, grid, u, v, w) + + 2 * ℑyzᵃᶠᶜ(i, j, k, grid, Σ₁₃², u, v, w) + + 2 * ℑxzᶠᵃᶜ(i, j, k, grid, Σ₂₃², u, v, w) + ) +end + +"Return the double dot product of strain at `fcf`." +@inline function ΣᵢⱼΣᵢⱼᶠᶜᶠ(i, j, k, grid, u, v, w) + return ( + ℑxzᶠᵃᶠ(i, j, k, grid, tr_Σ², u, v, w) + + 2 * ℑyzᵃᶜᶠ(i, j, k, grid, Σ₁₂², u, v, w) + + 2 * Σ₁₃²(i, j, k, grid, u, v, w) + + 2 * ℑxyᶠᶜᵃ(i, j, k, grid, Σ₂₃², u, v, w) + ) +end + +"Return the double dot product of strain at `cff`." +@inline function ΣᵢⱼΣᵢⱼᶜᶠᶠ(i, j, k, grid, u, v, w) + return ( + ℑyzᵃᶠᶠ(i, j, k, grid, tr_Σ², u, v, w) + + 2 * ℑxzᶜᵃᶠ(i, j, k, grid, Σ₁₂², u, v, w) + + 2 * ℑxyᶜᶠᵃ(i, j, k, grid, Σ₁₃², u, v, w) + + 2 * Σ₂₃²(i, j, k, grid, u, v, w) + ) +end + +"Return the double dot product of strain at `ccf`." +@inline function ΣᵢⱼΣᵢⱼᶜᶜᶠ(i, j, k, grid, u, v, w) + return ( + ℑzᵃᵃᶠ(i, j, k, grid, tr_Σ², u, v, w) + + 2 * ℑxyzᶜᶜᶠ(i, j, k, grid, Σ₁₂², u, v, w) + + 2 * ℑxᶜᵃᵃ(i, j, k, grid, Σ₁₃², u, v, w) + + 2 * ℑyᵃᶜᵃ(i, j, k, grid, Σ₂₃², u, v, w) + ) +end + +Base.summary(closure::SmagorinskyLilly) = string("SmagorinskyLilly: C=$(closure.C), Cb=$(closure.Cb), Pr=$(closure.Pr)") +Base.show(io::IO, closure::SmagorinskyLilly) = print(io, summary(closure)) + +##### +##### For closures that only require an eddy viscosity νₑ field. +##### + +function DiffusivityFields(grid, tracer_names, bcs, closure::SmagorinskyLilly) + + default_eddy_viscosity_bcs = (; νₑ = FieldBoundaryConditions(grid, (Center, Center, Center))) + bcs = merge(default_eddy_viscosity_bcs, bcs) + νₑ = CenterField(grid, boundary_conditions=bcs.νₑ) + + return (; νₑ) +end diff --git a/src/Utils/kernel_launching.jl b/src/Utils/kernel_launching.jl index 097ce8fe00..72e4bcb96f 100644 --- a/src/Utils/kernel_launching.jl +++ b/src/Utils/kernel_launching.jl @@ -2,7 +2,6 @@ ##### Utilities for launching kernels ##### -using Oceananigans: location using Oceananigans.Architectures using Oceananigans.Grids using Oceananigans.Grids: AbstractGrid @@ -80,9 +79,6 @@ end contiguousrange(range::NTuple{N, Int}, offset::NTuple{N, Int}) where N = Tuple(1+o:r+o for (r, o) in zip(range, offset)) flatten_reduced_dimensions(worksize, dims) = Tuple(d ∈ dims ? 1 : worksize[d] for d = 1:3) -# Support for 1D -heuristic_workgroup(Wx) = min(Wx, 256) - # This supports 2D, 3D and 4D work sizes (but the 3rd and 4th dimension are discarded) function heuristic_workgroup(Wx, Wy, Wz=nothing, Wt=nothing) @@ -272,9 +268,12 @@ end @inline function _launch!(arch, grid, workspec, kernel!, first_kernel_arg, other_kernel_args...; exclude_periphery = false, reduced_dimensions = (), - active_cells_map = nothing) + active_cells_map = nothing, + # TODO: these two kwargs do nothing: + only_local_halos = false, + async = false) - location = Oceananigans.location(first_kernel_arg) + location = Oceananigans.Grids.location(first_kernel_arg) loop!, worksize = configure_kernel(arch, grid, workspec, kernel!; location, @@ -283,15 +282,7 @@ end active_cells_map) # Don't launch kernels with no size - haswork = if worksize isa OffsetStaticSize - length(worksize) > 0 - elseif worksize isa Number - worksize > 0 - else - true - end - - if haswork + if worksize != 0 loop!(first_kernel_arg, other_kernel_args...) end @@ -333,16 +324,16 @@ end @pure get(::OffsetStaticSize{S}) where {S} = S @pure Base.getindex(::OffsetStaticSize{S}, i::Int) where {S} = i <= length(S) ? S[i] : 1 @pure Base.ndims(::OffsetStaticSize{S}) where {S} = length(S) -@pure Base.length(::OffsetStaticSize{S}) where {S} = prod(map(worksize, S)) +@pure Base.length(::OffsetStaticSize{S}) where {S} = prod(worksize.(S)) @inline getrange(::OffsetStaticSize{S}) where {S} = worksize(S), offsets(S) @inline getrange(::Type{OffsetStaticSize{S}}) where {S} = worksize(S), offsets(S) @inline offsets(ranges::Tuple{Vararg{UnitRange}}) = Tuple(r.start - 1 for r in ranges) -@inline worksize(t::Tuple) = map(worksize, t) -@inline worksize(sz::Int) = sz -@inline worksize(r::UnitRange) = length(r) +@inline worksize(i::Tuple) = worksize.(i) +@inline worksize(i::Int) = i +@inline worksize(i::UnitRange) = length(i) """a type used to store offsets in `NDRange` types""" struct KernelOffsets{O} diff --git a/test/dependencies_for_runtests.jl b/test/dependencies_for_runtests.jl index 5d9b4e0680..3e15e0a3d8 100644 --- a/test/dependencies_for_runtests.jl +++ b/test/dependencies_for_runtests.jl @@ -45,7 +45,7 @@ using Oceananigans.Logger using Oceananigans.Units using Oceananigans.Utils -using Oceananigans: Clock, location +using Oceananigans: Clock using Oceananigans.Architectures: device, array_type # to resolve conflict with CUDA.device using Oceananigans.Architectures: on_architecture using Oceananigans.AbstractOperations: UnaryOperation, Derivative, BinaryOperation, MultiaryOperation @@ -71,10 +71,7 @@ closures = ( :ScalarDiffusivity, :ScalarBiharmonicDiffusivity, :TwoDimensionalLeith, - :ConstantSmagorinsky, :SmagorinskyLilly, - :LagrangianAveragedDynamicSmagorinsky, - :DirectionallyAveragedDynamicSmagorinsky, :AnisotropicMinimumDissipation, :ConvectiveAdjustmentVerticalDiffusivity, ) diff --git a/test/regression_tests/ocean_large_eddy_simulation_regression_test.jl b/test/regression_tests/ocean_large_eddy_simulation_regression_test.jl index a247c89730..8a9ac391e9 100644 --- a/test/regression_tests/ocean_large_eddy_simulation_regression_test.jl +++ b/test/regression_tests/ocean_large_eddy_simulation_regression_test.jl @@ -3,14 +3,7 @@ using Oceananigans.TimeSteppers: update_state! using Oceananigans.DistributedComputations: cpu_architecture, partition function run_ocean_large_eddy_simulation_regression_test(arch, grid_type, closure) - if first(closure) isa SmagorinskyLilly - name = "ocean_large_eddy_simulation_SmagorinskyLilly" - else - firstclosure = first(closure) - closurename = typeof(firstclosure).name.wrapper - closurestr = string(closurename) - name = "ocean_large_eddy_simulation_$closurestr" - end + name = "ocean_large_eddy_simulation_" * string(typeof(first(closure)).name.wrapper) spinup_steps = 10000 test_steps = 10 @@ -47,7 +40,7 @@ function run_ocean_large_eddy_simulation_regression_test(arch, grid_type, closur boundary_conditions = (u=u_bcs, T=T_bcs, S=S_bcs)) # The type of the underlying data, not the offset array. - ArrayType = typeof(parent(model.velocities.u)) + ArrayType = typeof(model.velocities.u.data.parent) nx, ny, nz = size(model.tracers.T) u, v, w = model.velocities diff --git a/test/runtests.jl b/test/runtests.jl index a47f0e3d3d..c78eeefbb6 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -46,9 +46,7 @@ CUDA.allowscalar() do if group == :unit || group == :all @testset "Unit tests" begin include("test_grids.jl") - include("test_immersed_boundary_grid.jl") include("test_operators.jl") - include("test_vector_rotation_operators.jl") include("test_boundary_conditions.jl") include("test_field.jl") include("test_regrid.jl") @@ -80,7 +78,6 @@ CUDA.allowscalar() do if group == :poisson_solvers_2 || group == :all @testset "Poisson Solvers 2" begin include("test_poisson_solvers_stretched_grids.jl") - include("test_conjugate_gradient_poisson_solver.jl") end end @@ -174,24 +171,18 @@ CUDA.allowscalar() do if group == :distributed || group == :all MPI.Initialized() || MPI.Init() - # In case CUDA is not found, we reset CUDA and restart the julia session - reset_cuda_if_necessary() archs = test_architectures() include("test_distributed_models.jl") end if group == :distributed_solvers || group == :all MPI.Initialized() || MPI.Init() - # In case CUDA is not found, we reset CUDA and restart the julia session - reset_cuda_if_necessary() include("test_distributed_transpose.jl") include("test_distributed_poisson_solvers.jl") end if group == :distributed_hydrostatic_model || group == :all MPI.Initialized() || MPI.Init() - # In case CUDA is not found, we reset CUDA and restart the julia session - reset_cuda_if_necessary() archs = test_architectures() include("test_hydrostatic_regression.jl") include("test_distributed_hydrostatic_model.jl") @@ -199,8 +190,6 @@ CUDA.allowscalar() do if group == :distributed_nonhydrostatic_regression || group == :all MPI.Initialized() || MPI.Init() - # In case CUDA is not found, we reset CUDA and restart the julia session - reset_cuda_if_necessary() archs = nonhydrostatic_regression_test_architectures() include("test_nonhydrostatic_regression.jl") end diff --git a/test/test_checkpointer.jl b/test/test_checkpointer.jl index a5555cd2e5..c32fad63b3 100644 --- a/test/test_checkpointer.jl +++ b/test/test_checkpointer.jl @@ -14,10 +14,8 @@ function test_model_equality(test_model, true_model) for name in field_names @test all(test_model_fields[name].data .≈ true_model_fields[name].data) - if name ∈ keys(test_model.timestepper.Gⁿ) - @test all(test_model.timestepper.Gⁿ[name].data .≈ true_model.timestepper.Gⁿ[name].data) - @test all(test_model.timestepper.G⁻[name].data .≈ true_model.timestepper.G⁻[name].data) - end + @test all(test_model.timestepper.Gⁿ[name].data .≈ true_model.timestepper.Gⁿ[name].data) + @test all(test_model.timestepper.G⁻[name].data .≈ true_model.timestepper.G⁻[name].data) end end diff --git a/test/test_conjugate_gradient_poisson_solver.jl b/test/test_conjugate_gradient_poisson_solver.jl deleted file mode 100644 index 5e92a8f241..0000000000 --- a/test/test_conjugate_gradient_poisson_solver.jl +++ /dev/null @@ -1,28 +0,0 @@ -include("dependencies_for_runtests.jl") -using Oceananigans.Solvers: fft_poisson_solver, ConjugateGradientPoissonSolver - -@testset "Conjugate gradient Poisson solver" begin - @info "Testing Conjugate gradient poisson solver..." - for arch in archs - @testset "Conjugate gradient Poisson solver unit tests [$arch]" begin - @info "Unit testing Conjugate gradient poisson solver..." - - # Test the generic fft_poisson_solver constructor - x = y = (0, 1) - z = (0, 1) - grid = RectilinearGrid(arch, size=(2, 2, 2); x, y, z) - solver = ConjugateGradientPoissonSolver(grid, preconditioner=fft_poisson_solver(grid)) - pressure = CenterField(grid) - solve!(pressure, solver.conjugate_gradient_solver, solver.right_hand_side) - @test solver isa ConjugateGradientPoissonSolver - - z = [0, 0.2, 1] - grid = RectilinearGrid(arch, size=(2, 2, 2); x, y, z) - solver = ConjugateGradientPoissonSolver(grid, preconditioner=fft_poisson_solver(grid)) - pressure = CenterField(grid) - solve!(pressure, solver.conjugate_gradient_solver, solver.right_hand_side) - @test solver isa ConjugateGradientPoissonSolver - end - end -end - diff --git a/test/test_cubed_spheres.jl b/test/test_cubed_spheres.jl new file mode 100644 index 0000000000..d8984cb0a1 --- /dev/null +++ b/test/test_cubed_spheres.jl @@ -0,0 +1,232 @@ +include("dependencies_for_runtests.jl") +include("data_dependencies.jl") + +using Statistics: mean +using Oceananigans.Operators +using Oceananigans.CubedSpheres +using Oceananigans.Models.HydrostaticFreeSurfaceModels +using Oceananigans.Models.HydrostaticFreeSurfaceModels: VerticalVorticityField + +using OrthogonalSphericalShellGrids + +# To be used in the test below as `KernelFunctionOperation`s +@inline intrinsic_vector_x_component(i, j, k, grid, uₑ, vₑ) = + @inbounds intrinsic_vector(i, j, k, grid, uₑ, vₑ)[1] + +@inline intrinsic_vector_y_component(i, j, k, grid, uₑ, vₑ) = + @inbounds intrinsic_vector(i, j, k, grid, uₑ, vₑ)[2] + +@inline extrinsic_vector_x_component(i, j, k, grid, uₑ, vₑ) = + @inbounds intrinsic_vector(i, j, k, grid, uₑ, vₑ)[1] + +@inline extrinsic_vector_y_component(i, j, k, grid, uₑ, vₑ) = + @inbounds intrinsic_vector(i, j, k, grid, uₑ, vₑ)[2] + +function kinetic_energy(u, v) + ke = Field(0.5 * (u * u + v * v)) + return compute!(ke) +end + +function test_vector_rotation(grid) + u = XFaceField(grid) + v = YFaceField(grid) + + # Purely longitudinal flow in the extrinsic coordinate system + set!(u, 1) + set!(v, 0) + + # Convert it to an "Instrinsic" reference frame + uᵢ = KernelFunctionOperation{Face, Center, Center}(intrinsic_vector_x_component, grid, u, v) + vᵢ = KernelFunctionOperation{Center, Face, Center}(intrinsic_vector_y_component, grid, u, v) + + uᵢ = compute!(Field(uᵢ)) + vᵢ = compute!(Field(vᵢ)) + + # The extrema of u and v, as well as their mean value should + # be equivalent on an "intrinsic" frame + @test maximum(uᵢ) ≈ maximum(vᵢ) + @test minimum(uᵢ) ≈ minimum(vᵢ) + @test mean(uᵢ) ≈ mean(vᵢ) + @test mean(uᵢ) > 0 # The mean value should be positive + + # Kinetic energy should remain the same + KE = kinetic_energy(uᵢ, vᵢ) + @test all(on_architecture(CPU(), interior(KE)) .≈ 0.5) + + # Convert it back to a purely zonal velocity (vₑ == 0) + uₑ = KernelFunctionOperation{Face, Center, Center}(extrinsic_vector_x_component, grid, uᵢ, vᵢ) + vₑ = KernelFunctionOperation{Center, Face, Center}(extrinsic_vector_y_component, grid, uᵢ, vᵢ) + + uₑ = compute!(Field(uₑ)) + vₑ = compute!(Field(vₑ)) + + # Make sure that the flow was converted back to a + # purely zonal flow in the extrensic frame (v ≈ 0) + @test all(on_architecture(CPU(), interior(vₑ)) .≈ 0) + @test all(on_architecture(CPU(), interior(uₑ)) .≈ 1) + + # Purely meridional flow in the extrinsic coordinate system + set!(u, 0) + set!(v, 1) + + # Convert it to an "Instrinsic" reference frame + uᵢ = KernelFunctionOperation{Face, Center, Center}(intrinsic_vector_x_component, grid, u, v) + vᵢ = KernelFunctionOperation{Center, Face, Center}(intrinsic_vector_y_component, grid, u, v) + + uᵢ = compute!(Field(uᵢ)) + vᵢ = compute!(Field(vᵢ)) + + # The extrema of u and v, as well as their mean value should + # be equivalent on an "intrinsic" frame + @test maximum(uᵢ) ≈ maximum(vᵢ) + @test minimum(uᵢ) ≈ minimum(vᵢ) + @test mean(uᵢ) ≈ mean(vᵢ) + @test mean(vᵢ) > 0 # The mean value should be positive + + # Kinetic energy should remain the same + KE = kinetic_energy(uᵢ, vᵢ) + @test all(on_architecture(CPU(), interior(KE)) .≈ 0.5) + + # Convert it back to a purely zonal velocity (vₑ == 0) + uₑ = KernelFunctionOperation{Face, Center, Center}(extrinsic_vector_x_component, grid, uᵢ, vᵢ) + vₑ = KernelFunctionOperation{Center, Face, Center}(extrinsic_vector_y_component, grid, uᵢ, vᵢ) + + uₑ = compute!(Field(uₑ)) + vₑ = compute!(Field(vₑ)) + + # Make sure that the flow was converted back to a + # purely zonal flow in the extrensic frame (v ≈ 0) + @test all(on_architecture(CPU(), interior(vₑ)) .≈ 1) + @test all(on_architecture(CPU(), interior(uₑ)) .≈ 0) + + # Mixed zonal and meridional flow. + set!(u, 0.5) + set!(v, 0.5) + + # Convert it to an "Instrinsic" reference frame + uᵢ = KernelFunctionOperation{Face, Center, Center}(intrinsic_vector_x_component, grid, u, v) + vᵢ = KernelFunctionOperation{Center, Face, Center}(intrinsic_vector_y_component, grid, u, v) + + uᵢ = compute!(Field(uᵢ)) + vᵢ = compute!(Field(vᵢ)) + + # The extrema of u and v, as well as their mean value should + # be equivalent on an "intrinsic" frame + @test maximum(uᵢ) ≈ maximum(vᵢ) + @test minimum(uᵢ) ≈ minimum(vᵢ) + @test mean(uᵢ) ≈ mean(vᵢ) + @test mean(vᵢ) > 0 # The mean value should be positive + + # Kinetic energy should remain the same + KE = kinetic_energy(uᵢ, vᵢ) + @test all(on_architecture(CPU(), interior(KE)) .≈ 0.25) + + # Convert it back to a purely zonal velocity (vₑ == 0) + uₑ = KernelFunctionOperation{Face, Center, Center}(extrinsic_vector_x_component, grid, uᵢ, vᵢ) + vₑ = KernelFunctionOperation{Center, Face, Center}(extrinsic_vector_y_component, grid, uᵢ, vᵢ) + + uₑ = compute!(Field(uₑ)) + vₑ = compute!(Field(vₑ)) + + # Make sure that the flow was converted back to a + # purely zonal flow in the extrensic frame (v ≈ 0) + @test all(on_architecture(CPU(), interior(vₑ)) .≈ 0.5) + @test all(on_architecture(CPU(), interior(uₑ)) .≈ 0.5) +end + + +@testset "Cubed spheres" begin + + @testset "Conformal cubed sphere grid" begin + @info " Testing conformal cubed sphere grid..." + + grid = ConformalCubedSphereGrid(panel_size=(10, 10, 1), z=(-1, 0)) + @test try show(grid); println(); true; catch; false; end + end + + for arch in archs + + @info " Constructing a ConformalCubedSphereGrid from file [$(typeof(arch))]..." + + # These tests cause an undefined `Bound Access Error` on GPU's CI with the new CUDA version. + # The error is not reproducible neither on Tartarus nor on Sverdrup. + # These are excised for the moment (PR #2253) as Cubed sphere will be reworked + if !(arch isa GPU) + # Prototype grid and model for subsequent tests + cs32_filepath = datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2" + grid = OldConformalCubedSphereGrid(cs32_filepath, arch, Nz=1, z=(-1, 0)) + + @info " Constructing a HydrostaticFreeSurfaceModel on a ConformalCubedSphereGrid [$(typeof(arch))]..." + + free_surface = ExplicitFreeSurface(gravitational_acceleration=0.1) + model = HydrostaticFreeSurfaceModel(; grid, free_surface, + momentum_advection = VectorInvariant(), + coriolis = nothing, + closure = nothing, + tracers = :c, + buoyancy = nothing) + + @testset "Constructing a grid from file [$(typeof(arch))]" begin + @test grid isa ConformalCubedSphereGrid + end + + @testset "Conversion from Intrinsic to Extrinsic reference frame [$(typeof(arch))]" begin + @info " Testing the conversion of a vector between the Intrinsic and Extrinsic reference frame" + trg_grid = TripolarGrid(arch, size = (20, 20, 1), z = (0, 1)) + + test_vector_rotation(grid) + test_vector_rotation(trg_grid) + end + + @testset "CubedSphereData and CubedSphereFields [$(typeof(arch))]" begin + @info " Testing CubedSphereData and CubedSphereFields [$(typeof(arch))]..." + c = model.tracers.c + η = model.free_surface.η + + set!(c, 0) + set!(η, 0) + + CUDA.allowscalar(true) + @test all(all(face_c .== 0) for face_c in faces(c)) + @test all(all(face_η .== 0) for face_η in faces(η)) + CUDA.allowscalar(false) + + @test maximum(abs, c) == 0 + @test minimum(abs, c) == 0 + @test mean(c) == 0 + + @test maximum(abs, η) == 0 + @test minimum(abs, η) == 0 + @test mean(η) == 0 + end + + @testset "Constructing a HydrostaticFreeSurfaceModel [$(typeof(arch))]" begin + @test model isa HydrostaticFreeSurfaceModel + end + + @testset "Time stepping a HydrostaticFreeSurfaceModel [$(typeof(arch))]" begin + @info " Time-stepping HydrostaticFreeSurfaceModel on a ConformalCubedSphereGrid [$(typeof(arch))]..." + time_step!(model, 1) + @test try time_step!(model, 1); true; catch; false; end + end + + @testset "VerticalVorticityField on ConformalCubedSphereGrid [$(typeof(arch))]" begin + @info " Testing VerticalVorticityField on a ConformalCubedSphereGrid [$(typeof(arch))]..." + ζ = VerticalVorticityField(model) + + @test ζ isa Field + + set!(model, u = (x, y, z) -> rand()) + + @test try + compute!(ζ) + true + catch err + println(sprint(showerror, err)) + false + end + @test maximum(abs, ζ) > 0 # fingers crossed + end + end + end +end diff --git a/test/test_distributed_poisson_solvers.jl b/test/test_distributed_poisson_solvers.jl index 5ab7c90a7e..dca406d8c3 100644 --- a/test/test_distributed_poisson_solvers.jl +++ b/test/test_distributed_poisson_solvers.jl @@ -121,7 +121,9 @@ function divergence_free_poisson_tridiagonal_solution(grid_points, ranks, stretc return Array(interior(∇²ϕ)) ≈ Array(R) end -@testset "Distributed FFT-based Poisson solver" begin +@testset "Distributed FFT-based Poisson solver" begin + child_arch = test_child_arch() + for topology in ((Periodic, Periodic, Periodic), (Periodic, Periodic, Bounded), (Periodic, Bounded, Bounded), diff --git a/test/test_distributed_transpose.jl b/test/test_distributed_transpose.jl index 2eff450b27..5bc0fe6292 100644 --- a/test/test_distributed_transpose.jl +++ b/test/test_distributed_transpose.jl @@ -38,6 +38,8 @@ function test_transpose(grid_points, ranks, topo, child_arch) end @testset "Distributed Transpose" begin + child_arch = test_child_arch() + for topology in ((Periodic, Periodic, Periodic), (Periodic, Periodic, Bounded), (Periodic, Bounded, Bounded), diff --git a/test/test_enzyme.jl b/test/test_enzyme.jl index cc86656889..45faab48ff 100644 --- a/test/test_enzyme.jl +++ b/test/test_enzyme.jl @@ -157,8 +157,7 @@ end end end - -@testset "Enzyme for advection and diffusion with various boundary conditions" begin +@testset "Enzyme on advection and diffusion" begin Nx = Ny = 64 Nz = 8 @@ -184,17 +183,16 @@ end fill_halo_regions!(u) fill_halo_regions!(v) - @inline function tracer_flux(i, j, grid, clock, model_fields, p) + @inline function tracer_flux(x, y, t, c, p) c₀ = p.surface_tracer_concentration u★ = p.piston_velocity - return - u★ * (c₀ - model_fields.c[i, j, p.level]) + return - u★ * (c₀ - c) end parameters = (surface_tracer_concentration = 1, - piston_velocity = 0.1, - level = Nz) + piston_velocity = 0.1) - top_c_bc = FluxBoundaryCondition(tracer_flux; discrete_form=true, parameters) + top_c_bc = FluxBoundaryCondition(tracer_flux, field_dependencies=:c; parameters) c_bcs = FieldBoundaryConditions(top=top_c_bc) # TODO: @@ -202,51 +200,37 @@ end # 2. Add surface fluxes # 3. Do a problem where we invert for the tracer fluxes (maybe with CATKE) - model_no_bc = HydrostaticFreeSurfaceModel(; grid, - tracer_advection = WENO(), - tracers = :c, - velocities = PrescribedVelocityFields(; u, v), - closure = diffusion) - - model_bc = HydrostaticFreeSurfaceModel(; grid, - tracer_advection = WENO(), - tracers = :c, - velocities = PrescribedVelocityFields(; u, v), - boundary_conditions = (; c=c_bcs), - closure = diffusion) - - models = [model_no_bc, model_bc] - - @show "Advection-diffusion results, first without then with flux BC" - - for i in 1:2 - # Compute derivative by hand - κ₁, κ₂ = 0.9, 1.1 - c²₁ = stable_diffusion!(models[i], 1, κ₁) - c²₂ = stable_diffusion!(models[i], 1, κ₂) - dc²_dκ_fd = (c²₂ - c²₁) / (κ₂ - κ₁) - - # Now for real - amplitude = 1.0 - κ = 1.0 - dmodel = Enzyme.make_zero(models[i]) - set_diffusivity!(dmodel, 0) - - dc²_dκ = autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse), - stable_diffusion!, - Duplicated(models[i], dmodel), - Const(amplitude), - Active(κ)) - - @info """ \n - Advection-diffusion: + model = HydrostaticFreeSurfaceModel(; grid, + tracer_advection = WENO(), + tracers = :c, + velocities = PrescribedVelocityFields(; u, v), + closure = diffusion) + + # Compute derivative by hand + κ₁, κ₂ = 0.9, 1.1 + c²₁ = stable_diffusion!(model, 1, κ₁) + c²₂ = stable_diffusion!(model, 1, κ₂) + dc²_dκ_fd = (c²₂ - c²₁) / (κ₂ - κ₁) + + # Now for real + amplitude = 1.0 + κ = 1.0 + dmodel = Enzyme.make_zero(model) + set_diffusivity!(dmodel, 0) + + dc²_dκ = autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse), + stable_diffusion!, + Duplicated(model, dmodel), + Const(amplitude), + Active(κ)) + + @info """ \n Enzyme computed $dc²_dκ Finite differences computed $dc²_dκ_fd - """ + """ - tol = 0.01 - rel_error = abs(dc²_dκ[1][3] - dc²_dκ_fd) / abs(dc²_dκ_fd) - @test rel_error < tol - end + tol = 0.01 + rel_error = abs(dc²_dκ[1][3] - dc²_dκ_fd) / abs(dc²_dκ_fd) + @test rel_error < tol end diff --git a/test/test_field.jl b/test/test_field.jl index 6759eb5870..65c8787b94 100644 --- a/test/test_field.jl +++ b/test/test_field.jl @@ -6,11 +6,6 @@ using Oceananigans.Grids: total_length using Oceananigans.Fields: ReducedField, has_velocities using Oceananigans.Fields: VelocityFields, TracerFields, interpolate, interpolate! using Oceananigans.Fields: reduced_location -using Oceananigans.Fields: fractional_indices, interpolator -using Oceananigans.Grids: ξnode, ηnode, rnode - -using Random -using CUDA: @allowscalar """ correct_field_size(grid, FieldType, Tx, Ty, Tz) @@ -409,46 +404,6 @@ end end end - @testset "Unit interpolation" begin - for arch in archs - hu = (-1, 1) - hs = range(-1, 1, length=21) - zu = (-100, 0) - zs = range(-100, 0, length=33) - - for latitude in (hu, hs), longitude in (hu, hs), z in (zu, zs), loc in (Center(), Face()) - @info " Testing interpolation for $(latitude) latitude and longitude, $(z) z on $(typeof(loc))s..." - grid = LatitudeLongitudeGrid(arch; size = (20, 20, 32), longitude, latitude, z, halo = (5, 5, 5)) - - # Test random positions, - # set seed for reproducibility - Random.seed!(1234) - Xs = [(2rand()-1, 2rand()-1, -100rand()) for p in 1:20] - - for X in Xs - (x, y, z) = X - fi, fj, fk = @allowscalar fractional_indices(X, grid, loc, loc, loc) - - i⁻, i⁺, _ = interpolator(fi) - j⁻, j⁺, _ = interpolator(fj) - k⁻, k⁺, _ = interpolator(fk) - - x⁻ = @allowscalar ξnode(i⁻, j⁻, k⁻, grid, loc, loc, loc) - y⁻ = @allowscalar ηnode(i⁻, j⁻, k⁻, grid, loc, loc, loc) - z⁻ = @allowscalar rnode(i⁻, j⁻, k⁻, grid, loc, loc, loc) - - x⁺ = @allowscalar ξnode(i⁺, j⁺, k⁺, grid, loc, loc, loc) - y⁺ = @allowscalar ηnode(i⁺, j⁺, k⁺, grid, loc, loc, loc) - z⁺ = @allowscalar rnode(i⁺, j⁺, k⁺, grid, loc, loc, loc) - - @test x⁻ ≤ x ≤ x⁺ - @test y⁻ ≤ y ≤ y⁺ - @test z⁻ ≤ z ≤ z⁺ - end - end - end - end - @testset "Field interpolation" begin @info " Testing field interpolation..." diff --git a/test/test_field_scans.jl b/test/test_field_scans.jl index 781c208536..8fb864639a 100644 --- a/test/test_field_scans.jl +++ b/test/test_field_scans.jl @@ -92,7 +92,7 @@ interior_array(a, i, j, k) = Array(interior(a, i, j, k)) for T′ in (Tx, Txy) @test T′.operand.operand === T end - + for w′ in (wx, wxy) @test w′.operand.operand === w end @@ -255,7 +255,7 @@ interior_array(a, i, j, k) = Array(interior(a, i, j, k)) push!(results, all(interior(C) .== 1)) end - @test mean(results) == 1.0 + @test mean(results) == 1.0 end @testset "Allocating reductions [$arch_str]" begin @@ -331,7 +331,7 @@ interior_array(a, i, j, k) = Array(interior(a, i, j, k)) @testset "Immersed Fields reduction [$(typeof(arch))]" begin @info " Testing reductions of immersed Fields [$(typeof(arch))]" underlying_grid = RectilinearGrid(arch, size=(3, 3, 3), extent=(1, 1, 1)) - + grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom((x, y) -> y < 0.5 ? - 0.6 : 0)) c = Field((Center, Center, Nothing), grid) @@ -363,23 +363,6 @@ interior_array(a, i, j, k) = Array(interior(a, i, j, k)) compute!(bottom_half_average_field) bottom_half_average_array = Array(interior(bottom_half_average_field)) @test bottom_half_average_array[1, 1, 1] == bottom_half_average_manual - - # See: https://github.com/CliMA/Oceananigans.jl/issues/3948 - underlying_grid = RectilinearGrid( - topology=(Periodic, Periodic, Periodic), - size=(3, 3, 3), - x=(0, 1), y=(0, 1), z=(0, 1) - ) - - grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom((x, y) -> x + y)) - - c = CenterField(grid) - set!(c, (x, y, z) -> x + y + z) - - max_c² = Field(Reduction(maximum, c^2, dims=3)) - ∫max_c² = Integral(max_c², dims=(1, 2)) - compute!(∫max_c²) - @test ∫max_c² isa Reduction end end end diff --git a/test/test_forcings.jl b/test/test_forcings.jl index c4056ea7e0..e4e4e48294 100644 --- a/test/test_forcings.jl +++ b/test/test_forcings.jl @@ -118,7 +118,7 @@ end function time_step_with_field_time_series_forcing(arch) grid = RectilinearGrid(arch, size=(1, 1, 1), extent=(1, 1, 1)) - + u_forcing = FieldTimeSeries{Face, Center, Center}(grid, 0:1:3) for (t, time) in enumerate(u_forcing.times) @@ -134,14 +134,14 @@ function time_step_with_field_time_series_forcing(arch) model = NonhydrostaticModel(; grid, forcing=(; u=u_forcing)) time_step!(model, 2) time_step!(model, 2) - + @test u_forcing.backend.start == 4 return true end function relaxed_time_stepping(arch) - x_relax = Relaxation(rate = 1/60, mask = GaussianMask{:x}(center=0.5, width=0.1), + x_relax = Relaxation(rate = 1/60, mask = GaussianMask{:x}(center=0.5, width=0.1), target = LinearTarget{:x}(intercept=π, gradient=ℯ)) y_relax = Relaxation(rate = 1/60, mask = GaussianMask{:y}(center=0.5, width=0.1), @@ -197,7 +197,7 @@ end function two_forcings(arch) grid = RectilinearGrid(arch, size=(4, 5, 6), extent=(1, 1, 1), halo=(4, 4, 4)) - + forcing1 = Relaxation(rate=1) forcing2 = Relaxation(rate=2) @@ -221,7 +221,7 @@ function seven_forcings(arch) peculiar_forcing(x, y, z, t) = 2t / z eccentric_forcing(x, y, z, t) = x + y + z + t unconventional_forcing(x, y, z, t) = 10x * y - + F1 = Forcing(weird_forcing) F2 = Forcing(wonky_forcing) F3 = Forcing(strange_forcing) @@ -269,7 +269,7 @@ end @test time_step_with_multiple_field_dependent_forcing(arch) @test time_step_with_parameterized_field_dependent_forcing(arch) - end + end @testset "Relaxation forcing functions [$A]" begin @info " Testing relaxation forcing functions [$A]..." diff --git a/test/test_grids.jl b/test/test_grids.jl index b30b48c1bb..4ed536c732 100644 --- a/test/test_grids.jl +++ b/test/test_grids.jl @@ -2,9 +2,9 @@ include("dependencies_for_runtests.jl") include("data_dependencies.jl") using Oceananigans.Grids: total_extent, - xspacings, yspacings, zspacings, + xspacings, yspacings, zspacings, xnode, ynode, znode, λnode, φnode, - λspacing, φspacing, λspacings, φspacings + λspacings, φspacings, λspacing, φspacing using Oceananigans.Operators: Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δxᶜᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, Azᶜᶜᵃ @@ -182,12 +182,8 @@ function test_regular_rectilinear_xnode_ynode_znode_and_spacings(arch, FT) variably_spaced_grid = RectilinearGrid(arch, FT; size, topology, x=domain, y=domain, z=domain) - ibg_regular_spaced_grid = ImmersedBoundaryGrid(regular_spaced_grid, GridFittedBottom((x, y) -> 0)) - - ibg_variably_spaced_grid = ImmersedBoundaryGrid(variably_spaced_grid, GridFittedBottom((x, y) -> 0)) - - grids_types = ["regularly spaced", "variably spaced", "IBG regularly spaced", "IBG variably spaced"] - grids = [regular_spaced_grid, variably_spaced_grid, ibg_regular_spaced_grid, ibg_variably_spaced_grid] + grids_types = ["regularly spaced", "variably spaced"] + grids = [regular_spaced_grid, variably_spaced_grid] for (grid_type, grid) in zip(grids_types, grids) @info " Testing grid utils on $grid_type grid...." @@ -212,9 +208,9 @@ function test_regular_rectilinear_xnode_ynode_znode_and_spacings(arch, FT) @test all(y ≈ FT(π/N) for y in yspacings(grid, Face())) @test all(z ≈ FT(π/N) for z in zspacings(grid, Face())) - @test all(xspacings(grid, Face()) .== xspacings(grid, Face(), Center(), Center())) - @test all(yspacings(grid, Face()) .== yspacings(grid, Center(), Face(), Center())) - @test all(zspacings(grid, Face()) .== zspacings(grid, Center(), Center(), Face())) + @test xspacings(grid, Face()) == xspacings(grid, Face(), Center(), Center()) + @test yspacings(grid, Face()) == yspacings(grid, Center(), Face(), Center()) + @test zspacings(grid, Face()) == zspacings(grid, Center(), Center(), Face()) @test xspacing(1, 1, 1, grid, Face(), Center(), Center()) ≈ FT(π/N) @test yspacing(1, 1, 1, grid, Center(), Face(), Center()) ≈ FT(π/N) @@ -339,6 +335,7 @@ end function test_grid_equality_over_architectures() grid_cpu = RectilinearGrid(CPU(), topology=(Periodic, Periodic, Bounded), size=(3, 7, 9), x=(0, 1), y=(-1, 1), z=0:9) grid_gpu = RectilinearGrid(GPU(), topology=(Periodic, Periodic, Bounded), size=(3, 7, 9), x=(0, 1), y=(-1, 1), z=0:9) + return grid_cpu == grid_gpu end @@ -411,9 +408,8 @@ function test_rectilinear_grid_correct_spacings(FT, N) @test all(isapprox.( grid.zᵃᵃᶜ[1:N], zᵃᵃᶜ.(1:N) )) @test all(isapprox.( grid.Δzᵃᵃᶜ[1:N], Δzᵃᵃᶜ.(1:N) )) - @test all(isapprox.(zspacings(grid, Face()), reshape(grid.Δzᵃᵃᶠ[1:N+1], 1, 1, N+1))) - @test all(isapprox.(zspacings(grid, Center()), reshape(grid.Δzᵃᵃᶜ[1:N], 1, 1, N))) - + @test all(isapprox.(zspacings(grid, Face(), with_halos=true), grid.Δzᵃᵃᶠ)) + @test all(isapprox.(zspacings(grid, Center(), with_halos=true), grid.Δzᵃᵃᶜ)) @test zspacing(1, 1, 2, grid, Center(), Center(), Face()) == grid.Δzᵃᵃᶠ[2] @test minimum_zspacing(grid, Center(), Center(), Center()) ≈ minimum(grid.Δzᵃᵃᶜ[1:grid.Nz]) @@ -544,21 +540,21 @@ function test_basic_lat_lon_general_grid(FT) @test typeof(grid_reg.Δzᵃᵃᶜ) == typeof(grid_reg.Δzᵃᵃᶠ) == FT - @test all(xspacings(grid_reg, Center(), Center()) .== reshape(grid_reg.Δxᶜᶜᵃ[1:Nφ], 1, Nφ, 1)) - @test all(xspacings(grid_reg, Center(), Face() ) .== reshape(grid_reg.Δxᶜᶠᵃ[1:Nφ+1], 1, Nφ+1, 1)) - @test all(xspacings(grid_reg, Face(), Center()) .== reshape(grid_reg.Δxᶠᶜᵃ[1:Nφ], 1, Nφ, 1)) - @test all(xspacings(grid_reg, Face(), Face()) .== reshape(grid_reg.Δxᶠᶠᵃ[1:Nφ+1], 1, Nφ+1, 1)) - @test all(yspacings(grid_reg, Center(), Face()) .== grid_reg.Δyᶜᶠᵃ) - @test all(yspacings(grid_reg, Face(), Center()) .== grid_reg.Δyᶠᶜᵃ) - @test all(zspacings(grid_reg, Center()) .== grid_reg.Δzᵃᵃᶜ) - @test all(zspacings(grid_reg, Face()) .== grid_reg.Δzᵃᵃᶠ) - - @test all(xspacings(grid_reg, Center(), Center(), Center()) .== xspacings(grid_reg, Center(), Center())) - @test all(xspacings(grid_reg, Face(), Face(), Center()) .== xspacings(grid_reg, Face(), Face())) - @test all(yspacings(grid_reg, Center(), Face(), Center()) .== yspacings(grid_reg, Center(), Face())) - @test all(yspacings(grid_reg, Face(), Center(), Center()) .== yspacings(grid_reg, Face(), Center())) - @test all(zspacings(grid_reg, Face(), Center(), Center()) .== zspacings(grid_reg, Center())) - @test all(zspacings(grid_reg, Face(), Center(), Face() ) .== zspacings(grid_reg, Face())) + @test xspacings(grid_reg, Center(), Center(), with_halos=true) == grid_reg.Δxᶜᶜᵃ + @test xspacings(grid_reg, Center(), Face(), with_halos=true) == grid_reg.Δxᶜᶠᵃ + @test xspacings(grid_reg, Face(), Center(), with_halos=true) == grid_reg.Δxᶠᶜᵃ + @test xspacings(grid_reg, Face(), Face(), with_halos=true) == grid_reg.Δxᶠᶠᵃ + @test yspacings(grid_reg, Center(), Face(), with_halos=true) == grid_reg.Δyᶜᶠᵃ + @test yspacings(grid_reg, Face(), Center(), with_halos=true) == grid_reg.Δyᶠᶜᵃ + @test zspacings(grid_reg, Center(), with_halos=true) == grid_reg.Δzᵃᵃᶜ + @test zspacings(grid_reg, Face(), with_halos=true) == grid_reg.Δzᵃᵃᶠ + + @test xspacings(grid_reg, Center(), Center(), Center()) == xspacings(grid_reg, Center(), Center()) + @test xspacings(grid_reg, Face(), Face(), Center()) == xspacings(grid_reg, Face(), Face()) + @test yspacings(grid_reg, Center(), Face(), Center()) == yspacings(grid_reg, Center(), Face()) + @test yspacings(grid_reg, Face(), Center(), Center()) == yspacings(grid_reg, Face(), Center()) + @test zspacings(grid_reg, Face(), Face(), Center()) == zspacings(grid_reg, Center()) + @test zspacings(grid_reg, Face(), Center(), Face() ) == zspacings(grid_reg, Face()) @test xspacing(1, 2, 3, grid_reg, Center(), Center(), Center()) == grid_reg.Δxᶜᶜᵃ[2] @test xspacing(1, 2, 3, grid_reg, Center(), Face(), Center()) == grid_reg.Δxᶜᶠᵃ[2] @@ -567,10 +563,10 @@ function test_basic_lat_lon_general_grid(FT) @test zspacing(1, 2, 3, grid_reg, Center(), Center(), Face() ) == grid_reg.Δzᵃᵃᶠ @test zspacing(1, 2, 3, grid_reg, Center(), Center(), Center()) == grid_reg.Δzᵃᵃᶜ - @test all(λspacings(grid_reg, Center()) .== grid_reg.Δλᶜᵃᵃ) - @test all(λspacings(grid_reg, Face()) .== grid_reg.Δλᶠᵃᵃ) - @test all(φspacings(grid_reg, Center()) .== grid_reg.Δφᵃᶜᵃ) - @test all(φspacings(grid_reg, Face()) .== grid_reg.Δφᵃᶠᵃ) + @test λspacings(grid_reg, Center(), with_halos=true) == grid_reg.Δλᶜᵃᵃ + @test λspacings(grid_reg, Face(), with_halos=true) == grid_reg.Δλᶠᵃᵃ + @test φspacings(grid_reg, Center(), with_halos=true) == grid_reg.Δφᵃᶜᵃ + @test φspacings(grid_reg, Face(), with_halos=true) == grid_reg.Δφᵃᶠᵃ @test λspacing(1, 2, 3, grid_reg, Face(), Center(), Face()) == grid_reg.Δλᶠᵃᵃ @test φspacing(1, 2, 3, grid_reg, Center(), Face(), Center()) == grid_reg.Δφᵃᶠᵃ @@ -585,13 +581,13 @@ function test_basic_lat_lon_general_grid(FT) @test length(grid_str.λᶠᵃᵃ) == length(grid_reg.λᶠᵃᵃ) == Nλ + 2Hλ @test length(grid_str.λᶜᵃᵃ) == length(grid_reg.λᶜᵃᵃ) == Nλ + 2Hλ - + @test length(grid_str.φᵃᶠᵃ) == length(grid_reg.φᵃᶠᵃ) == Nφ + 2Hφ + 1 @test length(grid_str.φᵃᶜᵃ) == length(grid_reg.φᵃᶜᵃ) == Nφ + 2Hφ - + @test length(grid_str.zᵃᵃᶠ) == length(grid_reg.zᵃᵃᶠ) == Nz + 2Hz + 1 @test length(grid_str.zᵃᵃᶜ) == length(grid_reg.zᵃᵃᶜ) == Nz + 2Hz - + @test length(grid_str.Δzᵃᵃᶠ) == Nz + 2Hz + 1 @test length(grid_str.Δzᵃᵃᶜ) == Nz + 2Hz @@ -605,19 +601,22 @@ function test_basic_lat_lon_general_grid(FT) @test sum(grid_str.Δzᵃᵃᶜ) == grid_reg.Δzᵃᵃᶜ * length(grid_str.Δzᵃᵃᶜ) @test sum(grid_str.Δzᵃᵃᶠ) == grid_reg.Δzᵃᵃᶠ * length(grid_str.Δzᵃᵃᶠ) - @test all(xspacings(grid_str, Center(), Center()) .== reshape(grid_str.Δxᶜᶜᵃ[1:Nλ, 1:Nφ], Nλ, Nφ, 1)) - @test all(xspacings(grid_str, Center(), Face()) .== reshape(grid_str.Δxᶜᶠᵃ[1:Nλ, 1:Nφ+1], Nλ, Nφ+1, 1)) - @test all(xspacings(grid_str, Face(), Center()) .== reshape(grid_str.Δxᶠᶜᵃ[1:Nλ, 1:Nφ], Nλ, Nφ, 1)) - @test all(xspacings(grid_str, Face(), Face()) .== reshape(grid_str.Δxᶠᶠᵃ[1:Nλ, 1:Nφ+1], Nλ, Nφ+1, 1)) - - @test all(yspacings(grid_str, Center(), Face()) .== grid_str.Δyᶜᶠᵃ) - @test all(yspacings(grid_str, Face(), Center()) .== grid_str.Δyᶠᶜᵃ) + @test xspacings(grid_str, Center(), Center(), with_halos=true) == grid_str.Δxᶜᶜᵃ + @test xspacings(grid_str, Center(), Face(), with_halos=true) == grid_str.Δxᶜᶠᵃ + @test xspacings(grid_str, Face(), Center(), with_halos=true) == grid_str.Δxᶠᶜᵃ + @test xspacings(grid_str, Face(), Face(), with_halos=true) == grid_str.Δxᶠᶠᵃ + @test yspacings(grid_str, Center(), Face(), with_halos=true) == grid_str.Δyᶜᶠᵃ + @test yspacings(grid_str, Face(), Center(), with_halos=true) == grid_str.Δyᶠᶜᵃ + @test zspacings(grid_str, Center(), with_halos=true) == grid_str.Δzᵃᵃᶜ + @test zspacings(grid_str, Face(), with_halos=true) == grid_str.Δzᵃᵃᶠ - @test all(zspacings(grid_str, Center()) .== reshape(grid_str.Δzᵃᵃᶜ[1:Nz], 1, 1, Nz)) - @test all(zspacings(grid_str, Face()) .== reshape(grid_str.Δzᵃᵃᶠ[1:Nz+1], 1, 1, Nz+1)) + @test xspacings(grid_str, Center(), Center()) == grid_str.Δxᶜᶜᵃ[1:grid_str.Nx, 1:grid_str.Ny] + @test xspacings(grid_str, Center(), Face()) == grid_str.Δxᶜᶠᵃ[1:grid_str.Nx, 1:grid_str.Ny+1] + @test zspacings(grid_str, Center()) == grid_str.Δzᵃᵃᶜ[1:grid_str.Nz] + @test zspacings(grid_str, Face()) == grid_str.Δzᵃᵃᶠ[1:grid_str.Nz+1] - @test all(zspacings(grid_str, Center()) .== zspacings(grid_str, Center(), Center(), Center())) - @test all(zspacings(grid_str, Face()) .== zspacings(grid_str, Face(), Center(), Face())) + @test zspacings(grid_str, Face(), Face(), Center()) == zspacings(grid_str, Center()) + @test zspacings(grid_str, Face(), Center(), Face() ) == zspacings(grid_str, Face()) return nothing end @@ -691,7 +690,7 @@ function test_lat_lon_precomputed_metrics(FT, arch) println("$lat, $lon, $z") grid_pre = LatitudeLongitudeGrid(arch, FT, size=N, halo=H, latitude=lat, longitude=lon, z=z, precompute_metrics=true) grid_fly = LatitudeLongitudeGrid(arch, FT, size=N, halo=H, latitude=lat, longitude=lon, z=z) - + @test all(Array([all(Array([Δxᶠᶜᵃ(i, j, 1, grid_pre) ≈ Δxᶠᶜᵃ(i, j, 1, grid_fly) for i in 1-Hλ+1:Nλ+Hλ-1])) for j in 1-Hφ+1:Nφ+Hφ-1])) @test all(Array([all(Array([Δxᶜᶠᵃ(i, j, 1, grid_pre) ≈ Δxᶜᶠᵃ(i, j, 1, grid_fly) for i in 1-Hλ+1:Nλ+Hλ-1])) for j in 1-Hφ+1:Nφ+Hφ-1])) @test all(Array([all(Array([Δxᶠᶠᵃ(i, j, 1, grid_pre) ≈ Δxᶠᶠᵃ(i, j, 1, grid_fly) for i in 1-Hλ+1:Nλ+Hλ-1])) for j in 1-Hφ+1:Nφ+Hφ-1])) @@ -701,7 +700,7 @@ function test_lat_lon_precomputed_metrics(FT, arch) @test all(Array([all(Array([Azᶜᶠᵃ(i, j, 1, grid_pre) ≈ Azᶜᶠᵃ(i, j, 1, grid_fly) for i in 1-Hλ+1:Nλ+Hλ-1])) for j in 1-Hφ+1:Nφ+Hφ-1])) @test all(Array([all(Array([Azᶠᶠᵃ(i, j, 1, grid_pre) ≈ Azᶠᶠᵃ(i, j, 1, grid_fly) for i in 1-Hλ+1:Nλ+Hλ-1])) for j in 1-Hφ+1:Nφ+Hφ-1])) @test all(Array([all(Array([Azᶜᶜᵃ(i, j, 1, grid_pre) ≈ Azᶜᶜᵃ(i, j, 1, grid_fly) for i in 1-Hλ+1:Nλ+Hλ-1])) for j in 1-Hφ+1:Nφ+Hφ-1])) - end + end end end @@ -739,18 +738,18 @@ function test_orthogonal_shell_grid_array_sizes_and_spacings(FT) @test size(grid.φᶜᶠᵃ) == (Nx + 2Hx, Ny + 2Hy + 1) @test size(grid.φᶠᶠᵃ) == (Nx + 2Hx + 1, Ny + 2Hy + 1) - @test all(xspacings(grid, Center(), Center(), Face()) .== xspacings(grid, Center(), Center()) .== grid.Δxᶜᶜᵃ[1:Nx, 1:Ny]) - @test all(xspacings(grid, Center(), Face(), Face()) .== xspacings(grid, Center(), Face() ) .== grid.Δxᶜᶠᵃ[1:Nx, 1:Ny+1]) - @test all(xspacings(grid, Face(), Center(), Face()) .== xspacings(grid, Face(), Center()) .== grid.Δxᶠᶜᵃ[1:Nx+1, 1:Ny]) - @test all(xspacings(grid, Face(), Face(), Face()) .== xspacings(grid, Face(), Face() ) .== grid.Δxᶠᶠᵃ[1:Nx+1, 1:Ny+1]) + @test xspacings(grid, Center(), Center(), Face(), with_halos=true) == xspacings(grid, Center(), Center(), with_halos=true) == grid.Δxᶜᶜᵃ + @test xspacings(grid, Center(), Face(), Face(), with_halos=true) == xspacings(grid, Center(), Face(), with_halos=true) == grid.Δxᶜᶠᵃ + @test xspacings(grid, Face(), Center(), Face()) == xspacings(grid, Face(), Center()) == grid.Δxᶠᶜᵃ[1:grid.Nx+1, 1:grid.Ny] + @test xspacings(grid, Face(), Face(), Face()) == xspacings(grid, Face(), Face()) == grid.Δxᶠᶠᵃ[1:grid.Nx+1, 1:grid.Ny+1] - @test all(yspacings(grid, Center(), Center(), Face()) .== yspacings(grid, Center(), Center()) .== grid.Δyᶜᶜᵃ[1:Nx, 1:Ny]) - @test all(yspacings(grid, Center(), Face(), Face()) .== yspacings(grid, Center(), Face() ) .== grid.Δyᶜᶠᵃ[1:Nx, 1:Ny+1]) - @test all(yspacings(grid, Face(), Center(), Face()) .== yspacings(grid, Face(), Center()) .== grid.Δyᶠᶜᵃ[1:Nx+1, 1:Ny]) - @test all(yspacings(grid, Face(), Face(), Face()) .== yspacings(grid, Face(), Face() ) .== grid.Δyᶠᶠᵃ[1:Nx+1, 1:Ny+1]) + @test yspacings(grid, Center(), Center(), Face(), with_halos=true) == yspacings(grid, Center(), Center(), with_halos=true) == grid.Δyᶜᶜᵃ + @test yspacings(grid, Center(), Face(), Face(), with_halos=true) == yspacings(grid, Center(), Face(), with_halos=true) == grid.Δyᶜᶠᵃ + @test yspacings(grid, Face(), Center(), Face()) == yspacings(grid, Face(), Center()) == grid.Δyᶠᶜᵃ[1:grid.Nx+1, 1:grid.Ny] + @test yspacings(grid, Face(), Face(), Face()) == yspacings(grid, Face(), Face()) == grid.Δyᶠᶠᵃ[1:grid.Nx+1, 1:grid.Ny+1] - @test all(zspacings(grid, Center(), Center(), Face() ) .== zspacings(grid, Face() ) .== grid.Δzᵃᵃᶠ) - @test all(zspacings(grid, Center(), Center(), Center()) .== zspacings(grid, Center()) .== grid.Δzᵃᵃᶜ) + @test zspacings(grid, Center(), Face(), Face(), with_halos=true) == zspacings(grid, Face(), with_halos=true) == grid.Δzᵃᵃᶠ + @test zspacings(grid, Center(), Face(), Center()) == zspacings(grid, Center()) == grid.Δzᵃᵃᶜ return nothing end @@ -808,7 +807,7 @@ end @testset "Grid equality" begin @info " Testing grid equality operator (==)..." - + for arch in archs test_grid_equality(arch) end @@ -820,7 +819,7 @@ end # Testing show function topo = (Periodic, Periodic, Periodic) - + grid = RectilinearGrid(CPU(), topology=topo, size=(3, 7, 9), x=(0, 1), y=(-π, π), z=(0, 2π)) @test try @@ -831,7 +830,7 @@ end println(sprint(showerror, err)) false end - + @test grid isa RectilinearGrid end @@ -861,7 +860,7 @@ end # Testing show function Nz = 20 grid = RectilinearGrid(arch, size=(1, 1, Nz), x=(0, 1), y=(0, 1), z=collect(0:Nz).^2) - + @test try show(grid); println() true @@ -870,21 +869,11 @@ end println(sprint(showerror, err)) false end - - @test grid isa RectilinearGrid - end - - for arch in archs - @info " Testing on_architecture for RectilinearGrid..." - cpu_grid = RectilinearGrid(CPU(), size=(1, 1, 4), x=(0, 1), y=(0, 1), z=collect(0:4).^2) - grid = on_architecture(arch, cpu_grid) + @test grid isa RectilinearGrid - @test architecture(grid) == arch - cpu_grid_again = on_architecture(CPU(), grid) - @test cpu_grid_again == cpu_grid end end - + @testset "Latitude-longitude grid" begin @info " Testing general latitude-longitude grid..." @@ -895,7 +884,7 @@ end test_lat_lon_areas(FT) end - @info " Testing precomputed metrics on LatitudeLongitudeGrid..." + @info " Testing precomputed metrics on latitude-longitude grid..." for arch in archs, FT in float_types test_lat_lon_precomputed_metrics(FT, arch) test_lat_lon_xyzλφ_node_nodes(FT, arch) @@ -903,7 +892,7 @@ end # Testing show function for regular grid grid = LatitudeLongitudeGrid(CPU(), size=(36, 32, 1), longitude=(-180, 180), latitude=(-80, 80), z=(0, 1)) - + @test try show(grid); println() true @@ -915,40 +904,19 @@ end @test grid isa LatitudeLongitudeGrid - for arch in archs - @info " Testing show for vertically-stretched LatitudeLongitudeGrid..." - grid = LatitudeLongitudeGrid(arch, - size = (36, 32, 10), - longitude = (-180, 180), - latitude = (-80, 80), - z = collect(0:10)) - - @test try - show(grid); println() - true - catch err - println("error in show(::LatitudeLongitudeGrid)") - println(sprint(showerror, err)) - false - end + # Testing show function for stretched grid + grid = LatitudeLongitudeGrid(CPU(), size=(36, 32, 10), longitude=(-180, 180), latitude=(-80, 80), z=collect(0:10)) - @test grid isa LatitudeLongitudeGrid + @test try + show(grid); println() + true + catch err + println("error in show(::LatitudeLongitudeGrid)") + println(sprint(showerror, err)) + false end - for arch in archs - @info " Testing on_architecture for LatitudeLongitudeGrid..." - cpu_grid = LatitudeLongitudeGrid(CPU(), - size = (36, 32, 10), - longitude = (-180, 180), - latitude = (-80, 80), - z = collect(0:10)) - grid = on_architecture(arch, cpu_grid) - @test grid isa LatitudeLongitudeGrid - @test architecture(grid) == arch - - cpu_grid_again = on_architecture(CPU(), grid) - @test cpu_grid_again == cpu_grid - end + @test grid isa LatitudeLongitudeGrid end @testset "Single column grids" begin @@ -992,7 +960,7 @@ end end end end - + @testset "Conformal cubed sphere face grid" begin @info " Testing OrthogonalSphericalShellGrid grid..." @@ -1002,7 +970,7 @@ end # Testing show function grid = conformal_cubed_sphere_panel(CPU(), size=(10, 10, 1), z=(0, 1)) - + @test try show(grid); println() true diff --git a/test/test_hydrostatic_free_surface_immersed_boundaries_implicit_solve.jl b/test/test_hydrostatic_free_surface_immersed_boundaries_implicit_solve.jl index f7e5c4142f..9cf8116bbb 100644 --- a/test/test_hydrostatic_free_surface_immersed_boundaries_implicit_solve.jl +++ b/test/test_hydrostatic_free_surface_immersed_boundaries_implicit_solve.jl @@ -5,7 +5,7 @@ using Oceananigans.Architectures: on_architecture using Oceananigans.TurbulenceClosures using Oceananigans.Models.HydrostaticFreeSurfaceModels: compute_vertically_integrated_volume_flux!, compute_implicit_free_surface_right_hand_side!, - step_free_surface!, + implicit_free_surface_step!, pressure_correct_velocities! @testset "Immersed boundaries test divergent flow solve with hydrostatic free surface models" begin @@ -55,7 +55,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels: compute_vertically_integ v[imm1, jmm1, 1:Nz] .= 1 v[imm1, jmp1, 1:Nz] .= -1 - step_free_surface!(model.free_surface, model, model.timestepper, 1.0) + implicit_free_surface_step!(model.free_surface, model, 1.0, 1.5) sol = (sol..., model.free_surface.η) f = (f..., model.free_surface) diff --git a/test/test_hydrostatic_free_surface_models.jl b/test/test_hydrostatic_free_surface_models.jl index fcfa61a99a..d303a3d186 100644 --- a/test/test_hydrostatic_free_surface_models.jl +++ b/test/test_hydrostatic_free_surface_models.jl @@ -3,7 +3,7 @@ include("dependencies_for_runtests.jl") using Oceananigans.Models.HydrostaticFreeSurfaceModels: VectorInvariant, PrescribedVelocityFields using Oceananigans.Models.HydrostaticFreeSurfaceModels: ExplicitFreeSurface, ImplicitFreeSurface using Oceananigans.Models.HydrostaticFreeSurfaceModels: SingleColumnGrid -using Oceananigans.Advection: EnergyConserving, EnstrophyConserving, FluxFormAdvection +using Oceananigans.Advection: EnergyConserving, EnstrophyConserving using Oceananigans.TurbulenceClosures using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity @@ -11,16 +11,15 @@ function time_step_hydrostatic_model_works(grid; coriolis = nothing, free_surface = ExplicitFreeSurface(), momentum_advection = nothing, - tracers = [:b], - tracer_advection = nothing, closure = nothing, velocities = nothing) + tracers = [:b] buoyancy = BuoyancyTracer() closure isa CATKEVerticalDiffusivity && push!(tracers, :e) model = HydrostaticFreeSurfaceModel(; grid, coriolis, tracers, velocities, buoyancy, - momentum_advection, tracer_advection, free_surface, closure) + momentum_advection, free_surface, closure) simulation = Simulation(model, Δt=1.0, stop_iteration=1) @@ -54,31 +53,6 @@ function hydrostatic_free_surface_model_tracers_and_forcings_work(arch) return nothing end -function time_step_hydrostatic_model_with_catke_works(arch, FT) - grid = LatitudeLongitudeGrid( - arch, - FT, - topology = (Bounded, Bounded, Bounded), - size = (8, 8, 8), - longitude = (0, 1), - latitude = (0, 1), - z = (-100, 0) - ) - - model = HydrostaticFreeSurfaceModel(; - grid, - buoyancy = BuoyancyTracer(), - tracers = (:b, :e), - closure = CATKEVerticalDiffusivity(FT) - ) - - simulation = Simulation(model, Δt=1.0, stop_iteration=1) - - run!(simulation) - - return model.clock.iteration == 1 -end - topo_1d = (Flat, Flat, Bounded) topos_2d = ((Periodic, Flat, Bounded), @@ -91,7 +65,7 @@ topos_3d = ((Periodic, Periodic, Bounded), @testset "Hydrostatic free surface Models" begin @info "Testing hydrostatic free surface models..." - + @testset "$topo_1d model construction" begin @info " Testing $topo_1d model construction..." for arch in archs, FT in [Float64] #float_types @@ -105,7 +79,7 @@ topos_3d = ((Periodic, Periodic, Bounded), @test !(:η ∈ keys(fields(model))) # doesn't include free surface end end - + for topo in topos_2d @testset "$topo model construction" begin @info " Testing $topo model construction..." @@ -117,7 +91,7 @@ topos_3d = ((Periodic, Periodic, Bounded), end end end - + for topo in topos_3d @testset "$topo model construction" begin @info " Testing $topo model construction..." @@ -134,10 +108,10 @@ topos_3d = ((Periodic, Periodic, Bounded), grid = RectilinearGrid(topology=topo, size=(1, 1, 1), extent=(1, 2, 3), halo=(1, 1, 1)) hcabd_closure = ScalarBiharmonicDiffusivity() - @test_throws ArgumentError HydrostaticFreeSurfaceModel(grid=grid, tracer_advection=Centered(order=4)) - @test_throws ArgumentError HydrostaticFreeSurfaceModel(grid=grid, tracer_advection=UpwindBiased(order=3)) - @test_throws ArgumentError HydrostaticFreeSurfaceModel(grid=grid, tracer_advection=UpwindBiased(order=5)) - @test_throws ArgumentError HydrostaticFreeSurfaceModel(grid=grid, momentum_advection=UpwindBiased(order=5)) + @test_throws ArgumentError HydrostaticFreeSurfaceModel(grid=grid, tracer_advection=CenteredFourthOrder()) + @test_throws ArgumentError HydrostaticFreeSurfaceModel(grid=grid, tracer_advection=UpwindBiasedThirdOrder()) + @test_throws ArgumentError HydrostaticFreeSurfaceModel(grid=grid, tracer_advection=UpwindBiasedFifthOrder()) + @test_throws ArgumentError HydrostaticFreeSurfaceModel(grid=grid, momentum_advection=UpwindBiasedFifthOrder()) @test_throws ArgumentError HydrostaticFreeSurfaceModel(grid=grid, closure=hcabd_closure) # Big enough @@ -146,13 +120,13 @@ topos_3d = ((Periodic, Periodic, Bounded), model = HydrostaticFreeSurfaceModel(grid=bigger_grid, closure=hcabd_closure) @test model isa HydrostaticFreeSurfaceModel - model = HydrostaticFreeSurfaceModel(grid=bigger_grid, momentum_advection=UpwindBiased(order=5)) + model = HydrostaticFreeSurfaceModel(grid=bigger_grid, momentum_advection=UpwindBiasedFifthOrder()) @test model isa HydrostaticFreeSurfaceModel model = HydrostaticFreeSurfaceModel(grid=bigger_grid, closure=hcabd_closure) @test model isa HydrostaticFreeSurfaceModel - model = HydrostaticFreeSurfaceModel(grid=bigger_grid, tracer_advection=UpwindBiased(order=5)) + model = HydrostaticFreeSurfaceModel(grid=bigger_grid, tracer_advection=UpwindBiasedFifthOrder()) @test model isa HydrostaticFreeSurfaceModel end end @@ -185,7 +159,6 @@ topos_3d = ((Periodic, Periodic, Bounded), end for arch in archs - for topo in topos_3d grid = RectilinearGrid(arch, size=(1, 1, 1), extent=(1, 1, 1), topology=topo) @@ -197,18 +170,14 @@ topos_3d = ((Periodic, Periodic, Bounded), z_face_generator(; Nz=1, p=1, H=1) = k -> -H + (k / (Nz+1))^p # returns a generating function - H = 7 - halo = (7, 7, 7) - rectilinear_grid = RectilinearGrid(arch; size=(H, H, 1), extent=(1, 1, 1), halo) - vertically_stretched_grid = RectilinearGrid(arch; size=(H, H, 1), x=(0, 1), y=(0, 1), z=z_face_generator(), halo=(H, H, H)) + rectilinear_grid = RectilinearGrid(arch, size=(3, 3, 1), extent=(1, 1, 1), halo=(3, 3, 3)) + vertically_stretched_grid = RectilinearGrid(arch, size=(3, 3, 1), x=(0, 1), y=(0, 1), z=z_face_generator(), halo=(3, 3, 3)) - precompute_metrics = true - lat_lon_sector_grid = LatitudeLongitudeGrid(arch; size=(H, H, H), longitude=(0, 60), latitude=(15, 75), z=(-1, 0), precompute_metrics, halo) - lat_lon_strip_grid = LatitudeLongitudeGrid(arch; size=(H, H, H), longitude=(-180, 180), latitude=(15, 75), z=(-1, 0), precompute_metrics, halo) - - z = z_face_generator() - lat_lon_sector_grid_stretched = LatitudeLongitudeGrid(arch; size=(H, H, H), longitude=(0, 60), latitude=(15, 75), z, precompute_metrics, halo) - lat_lon_strip_grid_stretched = LatitudeLongitudeGrid(arch; size=(H, H, H), longitude=(-180, 180), latitude=(15, 75), z, precompute_metrics, halo) + lat_lon_sector_grid = LatitudeLongitudeGrid(arch, size=(3, 3, 3), longitude=(0, 60), latitude=(15, 75), z=(-1, 0), precompute_metrics=true) + lat_lon_strip_grid = LatitudeLongitudeGrid(arch, size=(3, 3, 3), longitude=(-180, 180), latitude=(15, 75), z=(-1, 0), precompute_metrics=true) + + lat_lon_sector_grid_stretched = LatitudeLongitudeGrid(arch, size=(3, 3, 3), longitude=(0, 60), latitude=(15, 75), z=z_face_generator(), precompute_metrics=true) + lat_lon_strip_grid_stretched = LatitudeLongitudeGrid(arch, size=(3, 3, 3), longitude=(-180, 180), latitude=(15, 75), z=z_face_generator(), precompute_metrics=true) grids = (rectilinear_grid, vertically_stretched_grid, lat_lon_sector_grid, lat_lon_strip_grid, @@ -221,7 +190,7 @@ topos_3d = ((Periodic, Periodic, Bounded), topo = topology(grid) grid_type = typeof(grid).name.wrapper free_surface_type = typeof(free_surface).name.wrapper - test_label = "[$arch, $grid_type, $topo, $free_surface_type]" + test_label = "[$arch, $grid_type, $topo, $free_surface_type]" @testset "Time-stepping HydrostaticFreeSurfaceModels with various grids $test_label" begin @info " Testing time-stepping HydrostaticFreeSurfaceModels with various grids $test_label..." @test time_step_hydrostatic_model_works(grid; free_surface) @@ -241,34 +210,22 @@ topos_3d = ((Periodic, Periodic, Bounded), HydrostaticSphericalCoriolis(scheme=EnstrophyConserving())) @testset "Time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(coriolis))]" begin - @test time_step_hydrostatic_model_works(lat_lon_sector_grid; coriolis) - @test time_step_hydrostatic_model_works(lat_lon_strip_grid; coriolis) + @test time_step_hydrostatic_model_works(lat_lon_sector_grid, coriolis=coriolis) + @test time_step_hydrostatic_model_works(lat_lon_strip_grid, coriolis=coriolis) end end - for momentum_advection in (VectorInvariant(), WENOVectorInvariant(), Centered(), WENO()) + for momentum_advection in (VectorInvariant(), CenteredSecondOrder(), WENO()) @testset "Time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(momentum_advection))]" begin @info " Testing time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(momentum_advection))]..." - @test time_step_hydrostatic_model_works(rectilinear_grid; momentum_advection) + @test time_step_hydrostatic_model_works(rectilinear_grid, momentum_advection=momentum_advection) end end - for momentum_advection in (VectorInvariant(), WENOVectorInvariant()) - @testset "Time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(momentum_advection))]" begin - @info " Testing time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(momentum_advection))]..." - @test time_step_hydrostatic_model_works(lat_lon_sector_grid; momentum_advection) - end - end - - for tracer_advection in [WENO(), - FluxFormAdvection(WENO(), WENO(), Centered()), - (b=WENO(), c=nothing)] - - T = typeof(tracer_advection) - @testset "Time-stepping HydrostaticFreeSurfaceModels with tracer advection [$arch, $T]" begin - @info " Testing time-stepping HydrostaticFreeSurfaceModels with tracer advection [$arch, $T]..." - @test time_step_hydrostatic_model_works(rectilinear_grid; tracer_advection, tracers=[:b, :c]) - end + momentum_advection = VectorInvariant() + @testset "Time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(momentum_advection))]" begin + @info " Testing time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(momentum_advection))]..." + @test time_step_hydrostatic_model_works(lat_lon_sector_grid, momentum_advection=momentum_advection) end for closure in (ScalarDiffusivity(), @@ -281,8 +238,8 @@ topos_3d = ((Periodic, Periodic, Bounded), @testset "Time-stepping Curvilinear HydrostaticFreeSurfaceModels [$arch, $(typeof(closure).name.wrapper)]" begin @info " Testing time-stepping Curvilinear HydrostaticFreeSurfaceModels [$arch, $(typeof(closure).name.wrapper)]..." @test_skip time_step_hydrostatic_model_works(arch, vertically_stretched_grid, closure=closure) - @test time_step_hydrostatic_model_works(lat_lon_sector_grid; closure) - @test time_step_hydrostatic_model_works(lat_lon_strip_grid; closure) + @test time_step_hydrostatic_model_works(lat_lon_sector_grid, closure=closure) + @test time_step_hydrostatic_model_works(lat_lon_strip_grid, closure=closure) end end @@ -303,7 +260,7 @@ topos_3d = ((Periodic, Periodic, Bounded), @test time_step_hydrostatic_model_works(rectilinear_grid, momentum_advection = nothing, velocities = velocities) @test time_step_hydrostatic_model_works(lat_lon_sector_grid, momentum_advection = nothing, velocities = velocities) - + parameters = (U=1, m=0.1, W=0.001) u(x, y, z, t, p) = p.U v(x, y, z, t, p) = exp(p.m * z) @@ -319,11 +276,5 @@ topos_3d = ((Periodic, Periodic, Bounded), @info " Testing HydrostaticFreeSurfaceModel with tracers and forcings [$arch]..." hydrostatic_free_surface_model_tracers_and_forcings_work(arch) end - - # See: https://github.com/CliMA/Oceananigans.jl/issues/3870 - @testset "HydrostaticFreeSurfaceModel with Float32 CATKE [$arch]" begin - @info " Testing HydrostaticFreeSurfaceModel with Float32 CATKE [$arch]..." - @test time_step_hydrostatic_model_with_catke_works(arch, Float32) - end end end diff --git a/test/test_hydrostatic_regression.jl b/test/test_hydrostatic_regression.jl index d989ba6c27..f7937d2eaa 100644 --- a/test/test_hydrostatic_regression.jl +++ b/test/test_hydrostatic_regression.jl @@ -3,7 +3,7 @@ include("data_dependencies.jl") using Oceananigans.Grids: topology, XRegularLLG, YRegularLLG, ZRegularLLG -function show_hydrostatic_test(grid, free_surface, precompute_metrics) +function show_hydrostatic_test(grid, free_surface, precompute_metrics) typeof(grid) <: XRegularLLG ? gx = :regular : gx = :stretched typeof(grid) <: YRegularLLG ? gy = :regular : gy = :stretched @@ -80,10 +80,10 @@ include("regression_tests/hydrostatic_free_turbulence_regression_test.jl") substeps = 5) for free_surface in [explicit_free_surface, implicit_free_surface, split_explicit_free_surface] - + # GPU + ImplicitFreeSurface + precompute metrics cannot be tested on sverdrup at the moment - # because "uses too much parameter space (maximum 0x1100 bytes)" error - if !(precompute_metrics && free_surface isa ImplicitFreeSurface && arch isa GPU) && + # because "uses too much parameter space (maximum 0x1100 bytes)" error + if !(precompute_metrics && free_surface isa ImplicitFreeSurface && arch isa GPU) && !(free_surface isa ImplicitFreeSurface && arch isa Distributed) # Also no implicit free surface on distributed testset_str, info_str = show_hydrostatic_test(grid, free_surface, precompute_metrics) diff --git a/test/test_immersed_advection.jl b/test/test_immersed_advection.jl index 3136987208..dc36776b6a 100644 --- a/test/test_immersed_advection.jl +++ b/test/test_immersed_advection.jl @@ -12,8 +12,7 @@ using Oceananigans.Advection: _biased_interpolate_yᵃᶠᵃ, FluxFormAdvection -linear_advection_schemes = [Centered, UpwindBiased] -advection_schemes = [linear_advection_schemes... WENO] +advection_schemes = [Centered, UpwindBiased, WENO] @inline advective_order(buffer, ::Type{Centered}) = buffer * 2 @inline advective_order(buffer, AdvectionType) = buffer * 2 - 1 @@ -104,19 +103,12 @@ for arch in archs mask_immersed_field!(c) fill_halo_regions!(c) - for adv in linear_advection_schemes, buffer in [1, 2, 3, 4, 5] + for adv in advection_schemes, buffer in [1, 2, 3, 4, 5] scheme = adv(order = advective_order(buffer, adv)) @info " Testing immersed tracer reconstruction [$(typeof(arch)), $(summary(scheme))]" run_tracer_interpolation_test(c, ibg, scheme) end - - for buffer in [2, 3, 4, 5], bounds in (nothing, (0, 1)) - scheme = WENO(; order = advective_order(buffer, WENO), bounds) - - @info " Testing immersed tracer reconstruction [$(typeof(arch)), $(summary(scheme))]" - run_tracer_interpolation_test(c, ibg, scheme) - end end @testset "Immersed tracer conservation" begin diff --git a/test/test_immersed_boundary_grid.jl b/test/test_immersed_boundary_grid.jl index b84763f16c..5be526b3be 100644 --- a/test/test_immersed_boundary_grid.jl +++ b/test/test_immersed_boundary_grid.jl @@ -1,20 +1,16 @@ include("dependencies_for_runtests.jl") -grid = RectilinearGrid(; size=(100, 100, 100), extent = (1, 1, 1)) +grid = RectilinearGrid(; size=(2, 2, 2), extent = (1, 1, 1)) @testset "Testing Immersed Boundaries" begin @info "Testing the immersed boundary construction..." - bottom(x, y) = rand() - + + bottom(x, y) = -1 + 0.5 * exp(-x^2 - y^2) ibg = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom)) - # Test that the bottom is at the same position - bottom_height = interior(ibg.immersed_boundary.bottom_height) + # Unit test (bottom is at the right position) - zfaces = znodes(ibg, Face()) + @info "Testing stably stratified initial conditions..." - for i in 1:size(ibg, 1), j in 1:size(ibg, 2) - @test bottom_height[i, j, 1] ∈ zfaces - end end diff --git a/test/test_implicit_free_surface_solver.jl b/test/test_implicit_free_surface_solver.jl index c4ef0f2a19..d56c47562a 100644 --- a/test/test_implicit_free_surface_solver.jl +++ b/test/test_implicit_free_surface_solver.jl @@ -10,7 +10,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels: PCGImplicitFreeSurfaceSolver, MatrixImplicitFreeSurfaceSolver, compute_vertically_integrated_lateral_areas!, - step_free_surface!, + implicit_free_surface_step!, implicit_free_surface_linear_operation! using Oceananigans.Grids: with_halo @@ -52,7 +52,7 @@ function run_implicit_free_surface_solver_tests(arch, grid, free_surface) free_surface) set_simple_divergent_velocity!(model) - step_free_surface!(model.free_surface, model, model.timestepper, Δt) + implicit_free_surface_step!(model.free_surface, model, Δt, 1.5) acronym = free_surface.solver_method == :HeptadiagonalIterativeSolver ? "Matrix" : "PCG" @@ -167,9 +167,9 @@ end for m in (mat_model, pcg_model, fft_model) set_simple_divergent_velocity!(m) - step_free_surface!(m.free_surface, m, m.timestepper, Δt₁) - step_free_surface!(m.free_surface, m, m.timestepper, Δt₁) - step_free_surface!(m.free_surface, m, m.timestepper, Δt₂) + implicit_free_surface_step!(m.free_surface, m, Δt₁, 1.5) + implicit_free_surface_step!(m.free_surface, m, Δt₁, 1.5) + implicit_free_surface_step!(m.free_surface, m, Δt₂, 1.5) end mat_η = mat_model.free_surface.η diff --git a/test/test_jld2_output_writer.jl b/test/test_jld2_output_writer.jl index b36bd0f472..46e512bdcd 100644 --- a/test/test_jld2_output_writer.jl +++ b/test/test_jld2_output_writer.jl @@ -49,7 +49,6 @@ function test_jld2_size_file_splitting(arch) function fake_bc_init(file, model) file["boundary_conditions/fake"] = π end - ow = JLD2OutputWriter(model, (; u=model.velocities.u); dir = ".", filename = "test.jld2", @@ -157,7 +156,6 @@ function test_jld2_time_averaging_of_horizontal_averages(model) simulation.output_writers[:fluxes] = JLD2OutputWriter(model, average_fluxes, schedule = AveragedTimeInterval(4Δt, window=2Δt), dir = ".", - with_halos = false, filename = "jld2_time_averaging_test.jld2", overwrite_existing = true) @@ -166,7 +164,7 @@ function test_jld2_time_averaging_of_horizontal_averages(model) test_file_name = "jld2_time_averaging_test.jld2" file = jldopen(test_file_name) - # Data is saved without halos + # Data is saved without halos by default wu = file["timeseries/wu/4"][1, 1, 3] uv = file["timeseries/uv/4"][1, 1, 3] wT = file["timeseries/wT/4"][1, 1, 3] diff --git a/test/test_lagrangian_particle_tracking.jl b/test/test_lagrangian_particle_tracking.jl index 6ee6492a3f..fb60c23beb 100644 --- a/test/test_lagrangian_particle_tracking.jl +++ b/test/test_lagrangian_particle_tracking.jl @@ -324,22 +324,4 @@ lagrangian_particle_test_curvilinear_grid(arch, z) = grid = lagrangian_particle_test_curvilinear_grid(arch, z) run_simple_particle_tracking_tests(grid) end - - for arch in archs - @info " Testing Lagrangian particle tracking [$(typeof(arch))] with 0 particles ..." - xp = Array{Float64}(undef, 0) - yp = Array{Float64}(undef, 0) - zp = Array{Float64}(undef, 0) - - xp = on_architecture(arch, xp) - yp = on_architecture(arch, yp) - zp = on_architecture(arch, zp) - - grid = RectilinearGrid(arch, size=(1, 1, 1), extent=(1, 1, 1)) - particles = LagrangianParticles(x=xp, y=yp, z=zp) - model = NonhydrostaticModel(; grid, particles) - time_step!(model, 1) - @test model.particles isa LagrangianParticles - end end - diff --git a/test/test_nonhydrostatic_models.jl b/test/test_nonhydrostatic_models.jl index e49d4a47db..bd10653bfd 100644 --- a/test/test_nonhydrostatic_models.jl +++ b/test/test_nonhydrostatic_models.jl @@ -45,7 +45,7 @@ using Oceananigans.Grids: required_halo_size_x, required_halo_size_y, required_h @test model.grid.Hx == 1 && model.grid.Hy == 3 && model.grid.Hz == 4 # Model ensures that halos are at least of size 2 - for scheme in (Centered(order=4), UpwindBiased(order=3)) + for scheme in (CenteredFourthOrder(), UpwindBiasedThirdOrder()) model = NonhydrostaticModel(advection=scheme, grid=minimal_grid) @test model.grid.Hx == 2 && model.grid.Hy == 2 && model.grid.Hz == 2 @@ -54,7 +54,7 @@ using Oceananigans.Grids: required_halo_size_x, required_halo_size_y, required_h end # Model ensures that halos are at least of size 3 - for scheme in (WENO(), UpwindBiased(order=5)) + for scheme in (WENO(), UpwindBiasedFifthOrder()) model = NonhydrostaticModel(advection=scheme, grid=minimal_grid) @test model.grid.Hx == 3 && model.grid.Hy == 3 && model.grid.Hz == 3 diff --git a/test/test_output_readers.jl b/test/test_output_readers.jl index 064de64bad..fd51e1ca7a 100644 --- a/test/test_output_readers.jl +++ b/test/test_output_readers.jl @@ -221,40 +221,6 @@ end @test v1[2] isa Field end end - - if arch isa GPU - @testset "FieldTimeSeries with CuArray boundary conditions [$(typeof(arch))]" begin - @info " Testing FieldTimeSeries with CuArray boundary conditions..." - - x = y = z = (0, 1) - grid = RectilinearGrid(GPU(); size=(1, 1, 1), x, y, z) - - τx = CuArray(zeros(size(grid)...)) - τy = Field{Center, Face, Nothing}(grid) - u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(τx)) - v_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(τy)) - model = NonhydrostaticModel(; grid, boundary_conditions = (; u=u_bcs, v=v_bcs)) - simulation = Simulation(model; Δt=1, stop_iteration=1) - - simulation.output_writers[:jld2] = JLD2OutputWriter(model, model.velocities, - filename = "test_cuarray_bc.jld2", - schedule=IterationInterval(1), - overwrite_existing = true) - - run!(simulation) - - ut = FieldTimeSeries("test_cuarray_bc.jld2", "u") - vt = FieldTimeSeries("test_cuarray_bc.jld2", "v") - @test ut.boundary_conditions.top.classification isa Flux - @test ut.boundary_conditions.top.condition isa Array - - τy_ow = vt.boundary_conditions.top.condition - @test τy_ow isa Field{Center, Face, Nothing} - @test architecture(τy_ow) isa CPU - @test parent(τy_ow) isa Array - rm("test_cuarray_bc.jld2") - end - end end for arch in archs @@ -324,8 +290,8 @@ end @test t[1, 1, 1] == 3.8 end - @testset "Test chunked abstraction" begin - @info " Testing Chunked abstraction..." + @testset "Test chunked abstraction" begin + @info " Testing Chunked abstraction..." filepath = "testfile.jld2" fts = FieldTimeSeries(filepath, "c") fts_chunked = FieldTimeSeries(filepath, "c"; backend = InMemory(2), time_indexing = Cyclical()) @@ -376,8 +342,8 @@ end end for Backend in [InMemory, OnDisk] - @testset "FieldDataset{$Backend} indexing" begin - @info " Testing FieldDataset{$Backend} indexing..." + @testset "FieldDataset{$Backend}" begin + @info " Testing FieldDataset{$Backend}..." ds = FieldDataset(filepath3d, backend=Backend()) @@ -388,7 +354,7 @@ end @test ds[var_str] isa FieldTimeSeries @test ds[var_str][1] isa Field end - + for var_sym in (:u, :v, :w, :T, :S, :b, :ζ, :ke) @test ds[var_sym] isa FieldTimeSeries @test ds[var_sym][2] isa Field @@ -405,36 +371,6 @@ end end end - for Backend in [InMemory, OnDisk] - @testset "FieldTimeSeries{$Backend} parallel reading" begin - @info " Testing FieldTimeSeries{$Backend} parallel reading..." - - reader_kw = Dict(:parallel_read => true) - u3 = FieldTimeSeries(filepath3d, "u"; backend=Backend(), reader_kw) - b3 = FieldTimeSeries(filepath3d, "b"; backend=Backend(), reader_kw) - - @test u3 isa FieldTimeSeries - @test b3 isa FieldTimeSeries - @test u3[1] isa Field - @test b3[1] isa Field - end - end - - for Backend in [InMemory, OnDisk] - @testset "FieldDataset{$Backend} parallel reading" begin - @info " Testing FieldDataset{$Backend} parallel reading..." - - reader_kw = (; parallel_read = true) - ds = FieldDataset(filepath3d; backend=Backend(), reader_kw) - - @test ds isa FieldDataset - @test ds.u isa FieldTimeSeries - @test ds.b isa FieldTimeSeries - @test ds.u[1] isa Field - @test ds.b[1] isa Field - end - end - rm(filepath1d) rm(filepath2d) rm(filepath3d) diff --git a/test/test_poisson_solvers.jl b/test/test_poisson_solvers.jl index ac12c92dbd..7e37253a3e 100644 --- a/test/test_poisson_solvers.jl +++ b/test/test_poisson_solvers.jl @@ -1,8 +1,6 @@ include("dependencies_for_runtests.jl") include("dependencies_for_poisson_solvers.jl") -using Oceananigans.Solvers: fft_poisson_solver - ##### ##### Run pressure solver tests 1 ##### @@ -40,18 +38,6 @@ two_dimensional_topologies = [(Flat, Bounded, Bounded), @test poisson_solver_instantiates(grid, FFTW.ESTIMATE) @test poisson_solver_instantiates(grid, FFTW.MEASURE) end - - # Test the generic fft_poisson_solver constructor - x = y = (0, 1) - z = (0, 1) - regular_grid = RectilinearGrid(arch, size=(2, 2, 2); x, y, z) - fft_based_solver = fft_poisson_solver(regular_grid) - @test fft_based_solver isa FFTBasedPoissonSolver - - z = [0, 0.2, 1] - vertically_stretched_grid = RectilinearGrid(arch, size=(2, 2, 2); x, y, z) - fourier_tridiagonal_solver = fft_poisson_solver(vertically_stretched_grid) - @test fourier_tridiagonal_solver isa FourierTridiagonalPoissonSolver end end diff --git a/test/test_shallow_water_models.jl b/test/test_shallow_water_models.jl index e6304d8839..26c8bf85fb 100644 --- a/test/test_shallow_water_models.jl +++ b/test/test_shallow_water_models.jl @@ -197,7 +197,7 @@ end end # Advection = nothing is broken as halo does not have a maximum - for advection in (nothing, Centered(), WENO()) + for advection in (nothing, CenteredSecondOrder(), WENO()) @testset "Time-stepping ShallowWaterModels [$arch, $(typeof(advection))]" begin @info " Testing time-stepping ShallowWaterModels [$arch, $(typeof(advection))]..." @test time_stepping_shallow_water_model_works(arch, topos[1], nothing, advection) diff --git a/test/test_split_explicit_free_surface_solver.jl b/test/test_split_explicit_free_surface_solver.jl index f186dd72f9..2786bf70c9 100644 --- a/test/test_split_explicit_free_surface_solver.jl +++ b/test/test_split_explicit_free_surface_solver.jl @@ -1,13 +1,15 @@ include("dependencies_for_runtests.jl") -using Oceananigans.Fields: VelocityFields using Oceananigans.Models.HydrostaticFreeSurfaceModels -using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces: calculate_substeps, - calculate_adaptive_settings, - constant_averaging_kernel, - materialize_free_surface, - SplitExplicitFreeSurface, - iterate_split_explicit! +using Oceananigans.Models.HydrostaticFreeSurfaceModels: calculate_substeps, + calculate_adaptive_settings, + constant_averaging_kernel, + materialize_free_surface, + SplitExplicitFreeSurface, + SplitExplicitState, + SplitExplicitAuxiliaryFields, + SplitExplicitSettings, + iterate_split_explicit! @testset "Split-Explicit Dynamics" begin @@ -24,19 +26,14 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces x = (0, Lx), y = (0, Ly), z = (-Lz, 0), halo = (1, 1, 1)) - velocities = VelocityFields(grid) - sefs = SplitExplicitFreeSurface(substeps = 200, averaging_kernel = constant_averaging_kernel) - sefs = materialize_free_surface(sefs, velocities, grid) + sefs = materialize_free_surface(sefs, nothing, grid) sefs.η .= 0 - GU = Field{Face, Center, Nothing}(grid) - GV = Field{Center, Face, Nothing}(grid) @testset " One timestep test " begin - state = sefs.filtered_state - U, V = sefs.barotropic_velocities - η̅, U̅, V̅ = state.η, state.U, state.V + state = sefs.state + U, V, η̅, U̅, V̅ = state.U, state.V, state.η̅, state.U̅, state.V̅ η = sefs.η Δτ = 1.0 @@ -44,10 +41,10 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces η₀(x, y, z) = sin(x) set!(η, η₀) - Nsubsteps = calculate_substeps(sefs.substepping, 1) - fractional_Δt, weights = calculate_adaptive_settings(sefs.substepping, Nsubsteps) # barotropic time step in fraction of baroclinic step and averaging weights + Nsubsteps = calculate_substeps(sefs.settings.substepping, 1) + fractional_Δt, weights = calculate_adaptive_settings(sefs.settings.substepping, Nsubsteps) # barotropic time step in fraction of baroclinic step and averaging weights - iterate_split_explicit!(sefs, grid, GU, GV, Δτ, weights, Val(1)) + iterate_split_explicit!(sefs, grid, Δτ, weights, Val(1)) U_computed = Array(U.data.parent)[2:Nx+1, 2:Ny+1] U_exact = (reshape(-cos.(grid.xᶠᵃᵃ), (length(grid.xᶜᵃᵃ), 1)).+reshape(0 * grid.yᵃᶜᵃ, (1, length(grid.yᵃᶜᵃ))))[2:Nx+1, 2:Ny+1] @@ -56,9 +53,10 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces end @testset "Multi-timestep test " begin - state = sefs.filtered_state - U, V = sefs.barotropic_velocities - η̅, U̅, V̅ = state.η, state.U, state.V + state = sefs.state + auxiliary = sefs.auxiliary + U, V, η̅, U̅, V̅ = state.U, state.V, state.η̅, state.U̅, state.V̅ + Gᵁ, Gⱽ = auxiliary.Gᵁ, auxiliary.Gⱽ η = sefs.η T = 2π @@ -67,31 +65,33 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces Δτ_end = T - Nt * Δτ sefs = SplitExplicitFreeSurface(substeps = Nt, averaging_kernel = constant_averaging_kernel) - sefs = materialize_free_surface(sefs, velocities, grid) + sefs = materialize_free_surface(sefs, nothing, grid) # set!(η, f(x, y)) η₀(x, y, z) = sin(x) set!(η, η₀) - set!(U, 0) - set!(V, 0) + U₀(x, y, z) = 0 + set!(U, U₀) + V₀(x, y, z) = 0 + set!(V, V₀) η̅ .= 0 U̅ .= 0 V̅ .= 0 - GU .= 0 - GV .= 0 + Gᵁ .= 0 + Gⱽ .= 0 - weights = sefs.substepping.averaging_weights + weights = sefs.settings.substepping.averaging_weights for _ in 1:Nt - iterate_split_explicit!(sefs, grid, GU, GV, Δτ, weights, Val(1)) + iterate_split_explicit!(sefs, grid, Δτ, weights, Val(1)) end - iterate_split_explicit!(sefs, grid, GU, GV, Δτ_end, weights, Val(1)) + iterate_split_explicit!(sefs, grid, Δτ_end, weights, Val(1)) U_computed = Array(deepcopy(interior(U))) η_computed = Array(deepcopy(interior(η))) set!(η, η₀) - set!(U, 0) + set!(U, U₀) U_exact = Array(deepcopy(interior(U))) η_exact = Array(deepcopy(interior(η))) @@ -100,16 +100,18 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces end sefs = SplitExplicitFreeSurface(substeps = 200, averaging_kernel = constant_averaging_kernel) - sefs = materialize_free_surface(sefs, velocities, grid) + sefs = materialize_free_surface(sefs, nothing, grid) sefs.η .= 0 @testset "Averaging / Do Nothing test " begin - state = sefs.filtered_state - U, V = sefs.barotropic_velocities - η̅, U̅, V̅ = state.η, state.U, state.V - η = sefs.η + state = sefs.state + auxiliary = sefs.auxiliary + U, V, η̅, U̅, V̅ = state.U, state.V, state.η̅, state.U̅, state.V̅ + Gᵁ, Gⱽ = auxiliary.Gᵁ, auxiliary.Gⱽ + g = sefs.gravitational_acceleration + η = sefs.η Δτ = 2π / maximum([Nx, Ny]) * 1e-2 # the last factor is essentially the order of accuracy @@ -124,15 +126,16 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces fill!(η̅ , 0) fill!(U̅ , 0) fill!(V̅ , 0) - fill!(GU, 0) - fill!(GV, 0) + fill!(Gᵁ, 0) + fill!(Gⱽ, 0) + settings = sefs.settings - Nsubsteps = calculate_substeps(sefs.substepping, 1) - fractional_Δt, weights = calculate_adaptive_settings(sefs.substepping, Nsubsteps) # barotropic time step in fraction of baroclinic step and averaging weights + Nsubsteps = calculate_substeps(settings.substepping, 1) + fractional_Δt, weights = calculate_adaptive_settings(settings.substepping, Nsubsteps) # barotropic time step in fraction of baroclinic step and averaging weights for step in 1:Nsubsteps - iterate_split_explicit!(sefs, grid, GU, GV, Δτ, weights, Val(1)) + iterate_split_explicit!(sefs, grid, Δτ, weights, Val(1)) end U_computed = Array(deepcopy(interior(U))) @@ -166,12 +169,13 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces Nt = floor(Int, T / Δτ) Δτ_end = T - Nt * Δτ - sefs = SplitExplicitFreeSurface(grid; substeps = Nt + 1, averaging_kernel = constant_averaging_kernel) - sefs = materialize_free_surface(sefs, velocities, grid) + sefs = SplitExplicitFreeSurface(substeps=200) + sefs = materialize_free_surface(sefs, nothing, grid) - state = sefs.filtered_state - U, V = sefs.barotropic_velocities - η̅, U̅, V̅ = state.η, state.U, state.V + state = sefs.state + auxiliary = sefs.auxiliary + U, V, η̅, U̅, V̅ = state.U, state.V, state.η̅, state.U̅, state.V̅ + Gᵁ, Gⱽ = auxiliary.Gᵁ, auxiliary.Gⱽ η = sefs.η g = sefs.gravitational_acceleration @@ -188,14 +192,17 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces η̅ .= 0 U̅ .= 0 V̅ .= 0 - GU .= gu_c - GV .= gv_c + Gᵁ .= gu_c + Gⱽ .= gv_c + + settings = SplitExplicitSettings(grid; substeps = Nt + 1, averaging_kernel = constant_averaging_kernel) + sefs = sefs(settings) - weights = sefs.substepping.averaging_weights + weights = settings.substepping.averaging_weights for i in 1:Nt - iterate_split_explicit!(sefs, grid, GU, GV, Δτ, weights, Val(1)) + iterate_split_explicit!(sefs, grid, Δτ, weights, Val(1)) end - iterate_split_explicit!(sefs, grid, GU, GV, Δτ_end, weights, Val(1)) + iterate_split_explicit!(sefs, grid, Δτ_end, weights, Val(1)) η_mean_after = mean(Array(interior(η))) @@ -215,11 +222,11 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces # ∂ₜₜ(η) = Δη η_exact = cos(ω * T) * (Array(interior(η, :, 1, 1)) .- 1) .+ 1 - U₀(x, y) = kx * cos(kx * x) * sin(ky * y) # ∂ₜU = - ∂x(η), since we know η + U₀(x, y, z) = kx * cos(kx * x) * sin(ky * y) # ∂ₜU = - ∂x(η), since we know η set!(U, U₀) U_exact = -(sin(ω * T) * 1 / ω) .* Array(interior(U, :, 1, 1)) .+ gu_c * T - V₀(x, y) = ky * sin(kx * x) * cos(ky * y) # ∂ₜV = - ∂y(η), since we know η + V₀(x, y, z) = ky * sin(kx * x) * cos(ky * y) # ∂ₜV = - ∂y(η), since we know η set!(V, V₀) V_exact = -(sin(ω * T) * 1 / ω) .* Array(interior(V, :, 1, 1)) .+ gv_c * T diff --git a/test/test_split_explicit_vertical_integrals.jl b/test/test_split_explicit_vertical_integrals.jl index 406f31c057..ca7713ad2b 100644 --- a/test/test_split_explicit_vertical_integrals.jl +++ b/test/test_split_explicit_vertical_integrals.jl @@ -1,12 +1,13 @@ include("dependencies_for_runtests.jl") -using Oceananigans.Fields: VelocityFields using Oceananigans.Models.HydrostaticFreeSurfaceModels -using Oceananigans.Models.HydrostaticFreeSurfaceModels: materialize_free_surface using Oceananigans.Models.HydrostaticFreeSurfaceModels: SplitExplicitFreeSurface -using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces: compute_barotropic_mode!, - barotropic_split_explicit_corrector!, - initialize_free_surface_state! +using Oceananigans.Models.HydrostaticFreeSurfaceModels: SplitExplicitState, + SplitExplicitAuxiliaryFields, + SplitExplicitSettings +using Oceananigans.Models.HydrostaticFreeSurfaceModels: compute_barotropic_mode!, + barotropic_split_explicit_corrector!, + initialize_free_surface_state! @testset "Barotropic Kernels" begin @@ -17,28 +18,28 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces Lx = Ly = Lz = 2π grid = RectilinearGrid(arch, topology = topology, size = (Nx, Ny, Nz), x = (0, Lx), y = (0, Ly), z = (-Lz, 0)) - - velocities = VelocityFields(grid) + + tmp = SplitExplicitFreeSurface(substeps = 200) + + sefs = SplitExplicitState(grid, tmp.settings.timestepper) + sefs = SplitExplicitAuxiliaryFields(grid) sefs = SplitExplicitFreeSurface(substeps = 200) - sefs = materialize_free_surface(sefs, velocities, grid) + sefs = materialize_free_surface(sefs, nothing, grid) - state = sefs.filtered_state - barotropic_velocities = sefs.barotropic_velocities - η̅, U̅, V̅ = state.η, state.U, state.V - U, V = barotropic_velocities + state = sefs.state + auxiliary = sefs.auxiliary + U, V, η̅, U̅, V̅ = state.U, state.V, state.η̅, state.U̅, state.V̅ + Gᵁ, Gⱽ = auxiliary.Gᵁ, auxiliary.Gⱽ u = Field{Face, Center, Center}(grid) v = Field{Center, Face, Center}(grid) - GU = Field{Face, Center, Nothing}(grid) - GV = Field{Center, Face, Nothing}(grid) - @testset "Average to zero" begin # set equal to something else η̅ .= U̅ .= V̅ .= 1.0 # now set equal to zero - initialize_free_surface_state!(sefs, sefs.timestepper, sefs.timestepper, Val(1)) + initialize_free_surface_state!(sefs.state, sefs.η, sefs.settings.timestepper) # don't forget the halo points fill_halo_regions!(η̅) @@ -57,7 +58,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces Δz .= grid.Δzᵃᵃᶠ set_u_check(x, y, z) = cos((π / 2) * z / Lz) - set_U_check(x, y) = (sin(0) - (-2 * Lz / (π))) + set_U_check(x, y, z) = (sin(0) - (-2 * Lz / (π))) set!(u, set_u_check) exact_U = similar(U) set!(exact_U, set_U_check) @@ -66,7 +67,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces @test all((Array(interior(U) .- interior(exact_U))) .< tolerance) set_v_check(x, y, z) = sin(x * y) * cos((π / 2) * z / Lz) - set_V_check(x, y) = sin(x * y) * (sin(0) - (-2 * Lz / (π))) + set_V_check(x, y, z) = sin(x * y) * (sin(0) - (-2 * Lz / (π))) set!(v, set_v_check) exact_V = similar(V) set!(exact_V, set_V_check) @@ -89,7 +90,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces @test all(Array(interior(U)) .≈ Lz) set_u_check(x, y, z) = sin(x) - set_U_check(x, y) = sin(x) * Lz + set_U_check(x, y, z) = sin(x) * Lz set!(u, set_u_check) exact_U = similar(U) set!(exact_U, set_U_check) @@ -97,7 +98,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces @test all(Array(interior(U)) .≈ Array(interior(exact_U))) set_v_check(x, y, z) = sin(x) * z * cos(y) - set_V_check(x, y) = -sin(x) * Lz^2 / 2.0 * cos(y) + set_V_check(x, y, z) = -sin(x) * Lz^2 / 2.0 * cos(y) set!(v, set_v_check) exact_V = similar(V) set!(exact_V, set_V_check) @@ -112,32 +113,38 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces Nx, Ny, Nz = 128, 64, 32 Lx = Ly = Lz = 2π - grid = RectilinearGrid(arch, topology = topology, size = (Nx, Ny, Nz), x = (0, Lx), y = (0, Ly), z = (-Lz, 0)) - velocities = VelocityFields(grid) + grid = RectilinearGrid(arch, topology = topology, size = (Nx, Ny, Nz), x = (0, Lx), y = (0, Ly), z = (-Lz, 0)) sefs = SplitExplicitFreeSurface(grid, cfl=0.7) - sefs = materialize_free_surface(sefs, velocities, grid) + sefs = materialize_free_surface(sefs, nothing, grid) + + state = sefs.state + auxiliary = sefs.auxiliary + U, V, η̅, U̅, V̅ = state.U, state.V, state.η̅, state.U̅, state.V̅ + Gᵁ, Gⱽ = auxiliary.Gᵁ, auxiliary.Gⱽ - U, V = sefs.barotropic_velocities - u = velocities.u - v = velocities.v + u = Field{Face, Center, Center}(grid) + v = Field{Center, Face, Center}(grid) u_corrected = similar(u) v_corrected = similar(v) set_u(x, y, z) = z + Lz / 2 + sin(x) - set_U̅(x, y) = cos(x) * Lz + set_U̅(x, y, z) = cos(x) * Lz set_u_corrected(x, y, z) = z + Lz / 2 + cos(x) set!(u, set_u) - set!(U, set_U̅) + set!(U̅, set_U̅) set!(u_corrected, set_u_corrected) set_v(x, y, z) = (z + Lz / 2) * sin(y) + sin(x) - set_V̅(x, y) = (cos(x) + x) * Lz + set_V̅(x, y, z) = (cos(x) + x) * Lz set_v_corrected(x, y, z) = (z + Lz / 2) * sin(y) + cos(x) + x set!(v, set_v) - set!(V, set_V̅) + set!(V̅, set_V̅) set!(v_corrected, set_v_corrected) + Δz = zeros(Nz) + Δz .= grid.Δzᵃᵃᶜ + barotropic_split_explicit_corrector!(u, v, sefs, grid) @test all(Array((interior(u) .- interior(u_corrected))) .< 1e-14) @test all(Array((interior(v) .- interior(v_corrected))) .< 1e-14) diff --git a/test/test_time_stepping.jl b/test/test_time_stepping.jl index 033ded8574..a3fbc45bd3 100644 --- a/test/test_time_stepping.jl +++ b/test/test_time_stepping.jl @@ -5,7 +5,6 @@ using Oceananigans.Grids: topological_tuple_length, total_size using Oceananigans.Fields: BackgroundField using Oceananigans.TimeSteppers: Clock using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity -using Oceananigans.TurbulenceClosures.Smagorinskys: LagrangianAveraging, DynamicSmagorinsky function time_stepping_works_with_flat_dimensions(arch, topology) size = Tuple(1 for i = 1:topological_tuple_length(topology...)) @@ -239,32 +238,23 @@ end Planes = (FPlane, ConstantCartesianCoriolis, BetaPlane, NonTraditionalBetaPlane) -BuoyancyModifiedAnisotropicMinimumDissipation(FT=Float64) = AnisotropicMinimumDissipation(FT, Cb=1.0) - -ConstantSmagorinsky(FT=Float64) = Smagorinsky(FT, coefficient=0.16) -DirectionallyAveragedDynamicSmagorinsky(FT=Float64) = - Smagorinsky(FT, coefficient=DynamicCoefficient(averaging=(1, 2))) -LagrangianAveragedDynamicSmagorinsky(FT=Float64) = - Smagorinsky(FT, coefficient=DynamicCoefficient(averaging=LagrangianAveraging())) +BuoyancyModifiedAnisotropicMinimumDissipation(FT) = AnisotropicMinimumDissipation(FT, Cb=1.0) Closures = (ScalarDiffusivity, ScalarBiharmonicDiffusivity, TwoDimensionalLeith, IsopycnalSkewSymmetricDiffusivity, - ConstantSmagorinsky, SmagorinskyLilly, - DirectionallyAveragedDynamicSmagorinsky, - LagrangianAveragedDynamicSmagorinsky, AnisotropicMinimumDissipation, BuoyancyModifiedAnisotropicMinimumDissipation, CATKEVerticalDiffusivity) advection_schemes = (nothing, - UpwindBiased(order=1), - Centered(order=2), - UpwindBiased(order=3), - Centered(order=4), - UpwindBiased(order=5), + UpwindBiasedFirstOrder(), + CenteredSecondOrder(), + UpwindBiasedThirdOrder(), + CenteredFourthOrder(), + UpwindBiasedFifthOrder(), WENO()) @inline ∂t_uˢ_uniform(z, t, h) = exp(z / h) * cos(t) @@ -329,7 +319,7 @@ timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3) end end - @testset "Flat dimensions" begin + @testset "Flat dimensions" begin for arch in archs for topology in ((Flat, Periodic, Periodic), (Periodic, Flat, Periodic), @@ -392,8 +382,6 @@ timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3) elseif Closure === CATKEVerticalDiffusivity # CATKE isn't supported with NonhydrostaticModel yet @test_skip time_stepping_works_with_closure(arch, FT, Closure) - elseif Closure() isa DynamicSmagorinsky - @test_skip time_stepping_works_with_closure(arch, FT, Closure) else @test time_stepping_works_with_closure(arch, FT, Closure) end diff --git a/test/test_turbulence_closures.jl b/test/test_turbulence_closures.jl index 28c8d226b2..00b8e34bef 100644 --- a/test/test_turbulence_closures.jl +++ b/test/test_turbulence_closures.jl @@ -10,19 +10,11 @@ using Oceananigans.TurbulenceClosures: diffusive_flux_x, diffusive_flux_y, diffu viscous_flux_vx, viscous_flux_vy, viscous_flux_vz, viscous_flux_wx, viscous_flux_wy, viscous_flux_wz -using Oceananigans.TurbulenceClosures: ScalarDiffusivity, - ScalarBiharmonicDiffusivity, - TwoDimensionalLeith, - ConvectiveAdjustmentVerticalDiffusivity, - Smagorinsky, - DynamicSmagorinsky, - SmagorinskyLilly, - LagrangianAveraging, - AnisotropicMinimumDissipation - -ConstantSmagorinsky(FT=Float64) = Smagorinsky(FT, coefficient=0.16) -DirectionallyAveragedDynamicSmagorinsky(FT=Float64) = DynamicSmagorinsky(FT, averaging=(1, 2)) -LagrangianAveragedDynamicSmagorinsky(FT=Float64) = DynamicSmagorinsky(FT, averaging=LagrangianAveraging()) +for closure in closures + @eval begin + using Oceananigans.TurbulenceClosures: $closure + end +end function tracer_specific_horizontal_diffusivity(T=Float64; νh=T(0.3), κh=T(0.7)) closure = HorizontalScalarDiffusivity(κ=(T=κh, S=κh), ν=νh) @@ -156,45 +148,6 @@ function time_step_with_variable_discrete_diffusivity(arch) return true end -function diffusivity_fields_sizes_are_correct(arch) - grid = RectilinearGrid(arch, size=(2, 3, 4), extent=(1, 2, 3)) - - closure = Smagorinsky(coefficient=DynamicCoefficient(averaging=1)) - model = NonhydrostaticModel(; grid, closure) - @test size(model.diffusivity_fields.𝒥ᴸᴹ) == (1, grid.Ny, grid.Nz) - @test size(model.diffusivity_fields.𝒥ᴹᴹ) == (1, grid.Ny, grid.Nz) - @test size(model.diffusivity_fields.LM) == size(grid) - @test size(model.diffusivity_fields.MM) == size(grid) - @test size(model.diffusivity_fields.Σ) == size(grid) - @test size(model.diffusivity_fields.Σ̄) == size(grid) - - closure = DynamicSmagorinsky(averaging=(1, 2)) - model = NonhydrostaticModel(; grid, closure) - @test size(model.diffusivity_fields.𝒥ᴸᴹ) == (1, 1, grid.Nz) - @test size(model.diffusivity_fields.𝒥ᴹᴹ) == (1, 1, grid.Nz) - - closure = DynamicSmagorinsky(averaging=Colon()) - model = NonhydrostaticModel(; grid, closure) - @test size(model.diffusivity_fields.𝒥ᴸᴹ) == (1, 1, 1) - @test size(model.diffusivity_fields.𝒥ᴹᴹ) == (1, 1, 1) - - closure = DynamicSmagorinsky(averaging=(2, 3)) - model = NonhydrostaticModel(; grid, closure) - @test size(model.diffusivity_fields.𝒥ᴸᴹ) == (grid.Nx, 1, 1) - @test size(model.diffusivity_fields.𝒥ᴹᴹ) == (grid.Nx, 1, 1) - - closure = DynamicSmagorinsky(averaging=LagrangianAveraging()) - model = NonhydrostaticModel(; grid, closure) - @test size(model.diffusivity_fields.𝒥ᴸᴹ) == size(grid) - @test size(model.diffusivity_fields.𝒥ᴹᴹ) == size(grid) - @test size(model.diffusivity_fields.𝒥ᴸᴹ⁻) == size(grid) - @test size(model.diffusivity_fields.𝒥ᴹᴹ⁻) == size(grid) - @test size(model.diffusivity_fields.Σ) == size(grid) - @test size(model.diffusivity_fields.Σ̄) == size(grid) - - return true -end - function time_step_with_tupled_closure(FT, arch) closure_tuple = (AnisotropicMinimumDissipation(FT), ScalarDiffusivity(FT)) @@ -259,12 +212,11 @@ end @testset "Closure instantiation" begin @info " Testing closure instantiation..." for closurename in closures - closure = @eval $closurename() + closure = getproperty(TurbulenceClosures, closurename)() @test closure isa TurbulenceClosures.AbstractTurbulenceClosure - closure isa DynamicSmagorinsky && continue # `DynamicSmagorinsky`s `_compute_LM_MM!()` kernel isn't compiling on buildkite grid = RectilinearGrid(CPU(), size=(2, 2, 2), extent=(1, 2, 3)) - model = NonhydrostaticModel(; grid, closure, tracers=:c) + model = NonhydrostaticModel(grid=grid, closure=closure, tracers=:c) c = model.tracers.c u = model.velocities.u κ = diffusivity(model.closure, model.diffusivity_fields, Val(:c)) @@ -336,14 +288,6 @@ end end end - @testset "Dynamic Smagorinsky closures" begin - @info " Testing that dynamic Smagorinsky closures produce diffusivity fields of correct sizes..." - for arch in archs - # `DynamicSmagorinsky`s `_compute_LM_MM!()` kernel isn't compiling on buildkite - @test_skip diffusivity_fields_sizes_are_correct(arch) - end - end - @testset "Time-stepping with CATKE closure" begin @info " Testing time-stepping with CATKE closure and closure tuples with CATKE..." for arch in archs @@ -382,8 +326,7 @@ end @testset "Diagnostics" begin @info " Testing turbulence closure diagnostics..." for closurename in closures - closure = @eval $closurename() - closure isa DynamicSmagorinsky && continue # `DynamicSmagorinsky`s `_compute_LM_MM!()` kernel isn't compiling on buildkite + closure = getproperty(TurbulenceClosures, closurename)() compute_closure_specific_diffusive_cfl(closure) end diff --git a/test/test_vector_rotation_operators.jl b/test/test_vector_rotation_operators.jl deleted file mode 100644 index eeb61a8bb9..0000000000 --- a/test/test_vector_rotation_operators.jl +++ /dev/null @@ -1,178 +0,0 @@ -include("dependencies_for_runtests.jl") - -using Statistics: mean -using Oceananigans.Operators - -# To be used in the test below as `KernelFunctionOperation`s -@inline intrinsic_vector_x_component(i, j, k, grid, uₑ, vₑ) = - @inbounds intrinsic_vector(i, j, k, grid, uₑ, vₑ)[1] - -@inline intrinsic_vector_y_component(i, j, k, grid, uₑ, vₑ) = - @inbounds intrinsic_vector(i, j, k, grid, uₑ, vₑ)[2] - -@inline extrinsic_vector_x_component(i, j, k, grid, uᵢ, vᵢ) = - @inbounds extrinsic_vector(i, j, k, grid, uᵢ, vᵢ)[1] - -@inline extrinsic_vector_y_component(i, j, k, grid, uᵢ, vᵢ) = - @inbounds extrinsic_vector(i, j, k, grid, uᵢ, vᵢ)[2] - -@inline function kinetic_energyᶜᶜᶜ(i, j, k, grid, uᶜᶜᶜ, vᶜᶜᶜ) - @inbounds u² = uᶜᶜᶜ[i, j, k]^2 - @inbounds v² = vᶜᶜᶜ[i, j, k]^2 - return (u² + v²) / 2 -end - -function kinetic_energy(u, v) - grid = u.grid - ke_op = KernelFunctionOperation{Center, Center, Center}(kinetic_energyᶜᶜᶜ, grid, u, v) - ke = Field(ke_op) - return compute!(ke) -end - -function pointwise_approximate_equal(field, val) - CPU_field = on_architecture(CPU(), field) - @test all(interior(CPU_field) .≈ val) -end - -# A purely zonal flow with an west-east velocity > 0 -# on a cubed sphere in an intrinsic coordinate system -# has the following properties: -function test_purely_zonal_flow(uᵢ, vᵢ, grid) - c1 = maximum(uᵢ) ≈ - minimum(vᵢ) - c2 = minimum(uᵢ) ≈ - maximum(vᵢ) - c3 = mean(uᵢ) ≈ - mean(vᵢ) - c4 = mean(uᵢ) > 0 # The mean value should be positive) - - return c1 & c2 & c3 & c4 -end - -# A purely meridional flow with a south-north velocity > 0 -# on a cubed sphere in an intrinsic coordinate system -# has the following properties: -function test_purely_meridional_flow(uᵢ, vᵢ, grid) - c1 = maximum(uᵢ) ≈ maximum(vᵢ) - c2 = minimum(uᵢ) ≈ minimum(vᵢ) - c3 = mean(uᵢ) ≈ mean(vᵢ) - c4 = mean(vᵢ) > 0 # The mean value should be positive - - return c1 & c2 & c3 & c4 -end - -function test_vector_rotation(grid) - u = CenterField(grid) - v = CenterField(grid) - - # Purely longitudinal flow in the extrinsic coordinate system - fill!(u, 1) - fill!(v, 0) - - # Convert it to an "Instrinsic" reference frame - uᵢ = KernelFunctionOperation{Center, Center, Center}(intrinsic_vector_x_component, grid, u, v) - vᵢ = KernelFunctionOperation{Center, Center, Center}(intrinsic_vector_y_component, grid, u, v) - - uᵢ = compute!(Field(uᵢ)) - vᵢ = compute!(Field(vᵢ)) - - # The extrema of u and v, as well as their mean value should - # be equivalent on an "intrinsic" frame - @test test_purely_zonal_flow(uᵢ, vᵢ, grid) - - # Kinetic energy should remain the same - KE = kinetic_energy(uᵢ, vᵢ) - @apply_regionally pointwise_approximate_equal(KE, 0.5) - - # Convert it back to a purely zonal velocity (vₑ == 0) - uₑ = KernelFunctionOperation{Center, Center, Center}(extrinsic_vector_x_component, grid, uᵢ, vᵢ) - vₑ = KernelFunctionOperation{Center, Center, Center}(extrinsic_vector_y_component, grid, uᵢ, vᵢ) - - uₑ = compute!(Field(uₑ)) - vₑ = compute!(Field(vₑ)) - - # Make sure that the flow was converted back to a - # purely zonal flow in the extrensic frame (v ≈ 0) - if architecture(grid) isa CPU - # Note that on the GPU, there are (apparently?) larger numerical errors - # which lead to -1e-17 < vₑ < 1e-17 for which this test fails. - @apply_regionally pointwise_approximate_equal(vₑ, 0) - end - - @apply_regionally pointwise_approximate_equal(uₑ, 1) - - # Purely meridional flow in the extrinsic coordinate system - fill!(u, 0) - fill!(v, 1) - - # Convert it to an "Instrinsic" reference frame - uᵢ = KernelFunctionOperation{Center, Center, Center}(intrinsic_vector_x_component, grid, u, v) - vᵢ = KernelFunctionOperation{Center, Center, Center}(intrinsic_vector_y_component, grid, u, v) - - uᵢ = compute!(Field(uᵢ)) - vᵢ = compute!(Field(vᵢ)) - - # The extrema of u and v, as well as their mean value should - # be equivalent on an "intrinsic" frame - @test test_purely_meridional_flow(uᵢ, vᵢ, grid) - - # Kinetic energy should remain the same - KE = kinetic_energy(uᵢ, vᵢ) - @apply_regionally pointwise_approximate_equal(KE, 0.5) - - # Convert it back to a purely zonal velocity (vₑ == 0) - uₑ = KernelFunctionOperation{Center, Center, Center}(extrinsic_vector_x_component, grid, uᵢ, vᵢ) - vₑ = KernelFunctionOperation{Center, Center, Center}(extrinsic_vector_y_component, grid, uᵢ, vᵢ) - - uₑ = compute!(Field(uₑ)) - vₑ = compute!(Field(vₑ)) - - # Make sure that the flow was converted back to a - # purely zonal flow in the extrensic frame (v ≈ 0) - @apply_regionally pointwise_approximate_equal(vₑ, 1) - - if architecture(grid) isa CPU - # Note that on the GPU, there are (apparently?) larger numerical errors - # which lead to - 4e-17 < uₑ < 4e-17 for which this test fails. - @apply_regionally pointwise_approximate_equal(uₑ, 0) - end - - # Mixed zonal and meridional flow. - fill!(u, 0.5) - fill!(v, 0.5) - - # Convert it to an "Instrinsic" reference frame - uᵢ = KernelFunctionOperation{Center, Center, Center}(intrinsic_vector_x_component, grid, u, v) - vᵢ = KernelFunctionOperation{Center, Center, Center}(intrinsic_vector_y_component, grid, u, v) - - uᵢ = compute!(Field(uᵢ)) - vᵢ = compute!(Field(vᵢ)) - - # The extrema of u and v, should be equivalent on an "intrinsic" frame - # when u == v on an extrinsic frame - @test maximum(uᵢ) ≈ maximum(vᵢ) - @test minimum(uᵢ) ≈ minimum(vᵢ) - - # Kinetic energy should remain the same - KE = kinetic_energy(uᵢ, vᵢ) - @apply_regionally pointwise_approximate_equal(KE, 0.25) - - # Convert it back to a purely zonal velocity (vₑ == 0) - uₑ = KernelFunctionOperation{Center, Center, Center}(extrinsic_vector_x_component, grid, uᵢ, vᵢ) - vₑ = KernelFunctionOperation{Center, Center, Center}(extrinsic_vector_y_component, grid, uᵢ, vᵢ) - - uₑ = compute!(Field(uₑ)) - vₑ = compute!(Field(vₑ)) - - # Make sure that the flow was converted back to a - # purely zonal flow in the extrensic frame (v ≈ 0) - @apply_regionally pointwise_approximate_equal(vₑ, 0.5) - @apply_regionally pointwise_approximate_equal(uₑ, 0.5) -end - -@testset "Vector rotation" begin - for arch in archs - @testset "Conversion from Intrinsic to Extrinsic reference frame [$(typeof(arch))]" begin - @info " Testing the conversion of a vector between the Intrinsic and Extrinsic reference frame" - grid = ConformalCubedSphereGrid(arch; panel_size=(10, 10, 1), z=(-1, 0)) - test_vector_rotation(grid) - end - end -end diff --git a/test/utils_for_runtests.jl b/test/utils_for_runtests.jl index 11dc6ecd5e..7bc06e528e 100644 --- a/test/utils_for_runtests.jl +++ b/test/utils_for_runtests.jl @@ -3,55 +3,21 @@ using Oceananigans.DistributedComputations: Distributed, Partition, child_archit import Oceananigans.Fields: interior -# Are the test running on the GPUs? -# Are the test running in parallel? -child_arch = get(ENV, "GPU_TEST", nothing) == "true" ? GPU() : CPU() -mpi_test = get(ENV, "MPI_TEST", nothing) == "true" - -# Sometimes when running tests in parallel, the CUDA.jl package is not loaded correctly. -# This function is a failsafe to re-load CUDA.jl using the suggested cach compilation from -# https://github.com/JuliaGPU/CUDA.jl/blob/a085bbb3d7856dfa929e6cdae04a146a259a2044/src/initialization.jl#L105 -# To make sure Julia restarts, an error is thrown. -function reset_cuda_if_necessary() - - # Do nothing if we are on the CPU - if child_arch isa CPU - return - end - - try - c = CUDA.zeros(10) # This will fail if CUDA is not available - catch err +test_child_arch() = CUDA.has_cuda() ? GPU() : CPU() - # Avoid race conditions and precompile on rank 0 only - if MPI.Comm_rank(MPI.COMM_WORLD) == 0 - pkg = Base.PkgId(Base.UUID("76a88914-d11a-5bdc-97e0-2f5a05c973a2"), "CUDA_Runtime_jll") - Base.compilecache(pkg) - @info "CUDA.jl was not correctly loaded. Re-loading CUDA.jl and re-starting Julia." - end - - MPI.Barrier(MPI.COMM_WORLD) +function test_architectures() + child_arch = test_child_arch() - # re-start Julia and re-load CUDA.jl - throw(err) - end -end - -function test_architectures() # If MPI is initialized with MPI.Comm_size > 0, we are running in parallel. - # We test several different configurations: `Partition(x = 4)`, `Partition(y = 4)`, + # We test several different configurations: `Partition(x = 4)`, `Partition(y = 4)`, # `Partition(x = 2, y = 2)`, and different fractional subdivisions in x, y and xy - if mpi_test - if MPI.Initialized() && MPI.Comm_size(MPI.COMM_WORLD) == 4 - return (Distributed(child_arch; partition = Partition(4)), - Distributed(child_arch; partition = Partition(1, 4)), - Distributed(child_arch; partition = Partition(2, 2)), - Distributed(child_arch; partition = Partition(x = Fractional(1, 2, 3, 4))), - Distributed(child_arch; partition = Partition(y = Fractional(1, 2, 3, 4))), - Distributed(child_arch; partition = Partition(x = Fractional(1, 2), y = Equal()))) - else - return throw("The MPI partitioning is not correctly configured.") - end + if MPI.Initialized() && MPI.Comm_size(MPI.COMM_WORLD) == 4 + return (Distributed(child_arch; partition = Partition(4)), + Distributed(child_arch; partition = Partition(1, 4)), + Distributed(child_arch; partition = Partition(2, 2)), + Distributed(child_arch; partition = Partition(x = Fractional(1, 2, 3, 4))), + Distributed(child_arch; partition = Partition(y = Fractional(1, 2, 3, 4))), + Distributed(child_arch; partition = Partition(x = Fractional(1, 2), y = Equal()))) else return tuple(child_arch) end @@ -59,18 +25,16 @@ end # For nonhydrostatic simulations we cannot use `Fractional` at the moment (requirements # for the tranpose are more stringent than for hydrostatic simulations). -function nonhydrostatic_regression_test_architectures() +function nonhydrostatic_regression_test_architectures() + child_arch = test_child_arch() + # If MPI is initialized with MPI.Comm_size > 0, we are running in parallel. - # We test 3 different configurations: `Partition(x = 4)`, `Partition(y = 4)` + # We test 3 different configurations: `Partition(x = 4)`, `Partition(y = 4)` # and `Partition(x = 2, y = 2)` - if mpi_test - if MPI.Initialized() && MPI.Comm_size(MPI.COMM_WORLD) == 4 - return (Distributed(child_arch; partition = Partition(4)), - Distributed(child_arch; partition = Partition(1, 4)), - Distributed(child_arch; partition = Partition(2, 2))) - else - return throw("The MPI partitioning is not correctly configured.") - end + if MPI.Initialized() && MPI.Comm_size(MPI.COMM_WORLD) == 4 + return (Distributed(child_arch; partition = Partition(4)), + Distributed(child_arch; partition = Partition(1, 4)), + Distributed(child_arch; partition = Partition(2, 2))) else return tuple(child_arch) end @@ -98,8 +62,8 @@ end # TODO: docstring? function center_clustered_coord(N, L, x₀) - Δz(k) = k < N / 2 + 1 ? 2 / (N - 1) * (k - 1) + 1 : - 2 / (N - 1) * (k - N) + 1 - z_faces = zeros(N+1) + Δz(k) = k < N / 2 + 1 ? 2 / (N - 1) * (k - 1) + 1 : - 2 / (N - 1) * (k - N) + 1 + z_faces = zeros(N+1) for k = 2:N+1 z_faces[k] = z_faces[k-1] + 3 - Δz(k-1) end @@ -109,12 +73,12 @@ end # TODO: docstring? function boundary_clustered_coord(N, L, x₀) - Δz(k) = k < N / 2 + 1 ? 2 / (N - 1) * (k - 1) + 1 : - 2 / (N - 1) * (k - N) + 1 - z_faces = zeros(N+1) + Δz(k) = k < N / 2 + 1 ? 2 / (N - 1) * (k - 1) + 1 : - 2 / (N - 1) * (k - N) + 1 + z_faces = zeros(N+1) for k = 2:N+1 z_faces[k] = z_faces[k-1] + Δz(k-1) end - z_faces = z_faces ./ z_faces[end] .* L .+ x₀ + z_faces = z_faces ./ z_faces[end] .* L .+ x₀ return z_faces end diff --git a/validation/advection/plot_rates_convergence_advection.jl b/validation/advection/plot_rates_convergence_advection.jl index 948550ed06..932ba53edf 100644 --- a/validation/advection/plot_rates_convergence_advection.jl +++ b/validation/advection/plot_rates_convergence_advection.jl @@ -56,13 +56,13 @@ pnorm = 1 uh(x, y, z) = U * h(x, y, z) schemes = ( - Centered(order=4), - UpwindBiased(order=5), - WENO(order=3), - WENO(order=5), - WENO(order=7), - WENO(order=9), - WENO(order=11) + CenteredFourthOrder(), + UpwindBiasedFifthOrder(), + WENO(order = 3), + WENO(order = 5), + WENO(order = 7), + WENO(order = 9), + WENO(order = 11) )s error = Dict() diff --git a/validation/biogeochemistry/simple_plankton_continuous_form_biogeochemistry.jl b/validation/biogeochemistry/simple_plankton_continuous_form_biogeochemistry.jl index 7d94dcca9b..28eca3083c 100644 --- a/validation/biogeochemistry/simple_plankton_continuous_form_biogeochemistry.jl +++ b/validation/biogeochemistry/simple_plankton_continuous_form_biogeochemistry.jl @@ -5,7 +5,7 @@ using Oceananigans.Grids: znode using Oceananigans.Forcings: maybe_constant_field using Oceananigans.Architectures: device, architecture using Oceananigans.Utils: launch! -using Oceananigans.Advection: Centered +using Oceananigans.Advection: CenteredSecondOrder using Oceananigans.Fields: Field, TracerFields, CenterField import Oceananigans.Biogeochemistry: @@ -43,7 +43,7 @@ function SimplePlanktonGrowthDeath(FT=Float64; grid, advection_scheme = nothing) if sinking_velocity != 0 - advection_scheme = Centered() + advection_scheme = CenteredSecondOrder() end u, v, w = maybe_constant_field.((0, 0, - sinking_velocity)) diff --git a/validation/convergence_tests/src/OneDimensionalCosineAdvectionDiffusion.jl b/validation/convergence_tests/src/OneDimensionalCosineAdvectionDiffusion.jl index ce17239da3..1793c63903 100644 --- a/validation/convergence_tests/src/OneDimensionalCosineAdvectionDiffusion.jl +++ b/validation/convergence_tests/src/OneDimensionalCosineAdvectionDiffusion.jl @@ -14,7 +14,7 @@ using ConvergenceTests: compute_error c(x, y, z, t, U, κ) = exp(-κ * t) * cos(x - U * t) function run_test(; Nx, Δt, stop_iteration, U = 1, κ = 1e-4, - architecture = CPU(), topo = (Periodic, Periodic, Periodic), advection = Centered()) + architecture = CPU(), topo = (Periodic, Periodic, Periodic), advection = CenteredSecondOrder()) ##### ##### Test advection-diffusion in the x-direction diff --git a/validation/convergence_tests/src/OneDimensionalGaussianAdvectionDiffusion.jl b/validation/convergence_tests/src/OneDimensionalGaussianAdvectionDiffusion.jl index 828226f4e1..71865859b7 100644 --- a/validation/convergence_tests/src/OneDimensionalGaussianAdvectionDiffusion.jl +++ b/validation/convergence_tests/src/OneDimensionalGaussianAdvectionDiffusion.jl @@ -15,7 +15,7 @@ using ConvergenceTests: compute_error c(x, y, z, t, U, κ, t₀) = 1 / √(4π * κ * (t + t₀)) * exp(-(x - U * t)^2 / σ(t, κ, t₀)) function run_test(; Nx, Δt, stop_iteration, U = 1, κ = 1e-4, width = 0.05, - architecture = CPU(), topo = (Periodic, Periodic, Periodic), advection = Centered()) + architecture = CPU(), topo = (Periodic, Periodic, Periodic), advection = CenteredSecondOrder()) t₀ = width^2 / 4κ diff --git a/validation/convergence_tests/src/TwoDimensionalGaussianAdvectionDiffusion.jl b/validation/convergence_tests/src/TwoDimensionalGaussianAdvectionDiffusion.jl index 12e6140824..d48829c097 100644 --- a/validation/convergence_tests/src/TwoDimensionalGaussianAdvectionDiffusion.jl +++ b/validation/convergence_tests/src/TwoDimensionalGaussianAdvectionDiffusion.jl @@ -16,7 +16,7 @@ using ConvergenceTests: compute_error c(x, y, z, t, U, κ, t₀) = 1 / √(4π * κ * (t + t₀)) * exp(-((x - U * t)^2 + (y - U * t)^2) / σ(t, κ, t₀)) function run_test(; Nx, Δt, stop_iteration, U = 1, κ = 1e-4, width = 0.05, - architecture = CPU(), topo = (Periodic, Periodic, Periodic), advection = Centered()) + architecture = CPU(), topo = (Periodic, Periodic, Periodic), advection = CenteredSecondOrder()) t₀ = width^2 / 4κ diff --git a/validation/distributed_simulations/distributed_shallow_water_turbulence.jl b/validation/distributed_simulations/distributed_shallow_water_turbulence.jl index a5279fd8b8..343671a312 100644 --- a/validation/distributed_simulations/distributed_shallow_water_turbulence.jl +++ b/validation/distributed_simulations/distributed_shallow_water_turbulence.jl @@ -19,7 +19,7 @@ local_rank = MPI.Comm_rank(arch.communicator) model = ShallowWaterModel(grid = grid, timestepper = :RungeKutta3, - momentum_advection = UpwindBiased(order=5), + momentum_advection = UpwindBiasedFifthOrder(), gravitational_acceleration = 1) set!(model, h=1) diff --git a/validation/immersed_boundaries/2D_rough_rayleighbenard.jl b/validation/immersed_boundaries/2D_rough_rayleighbenard.jl index 160bb04d6b..d620e0bb22 100644 --- a/validation/immersed_boundaries/2D_rough_rayleighbenard.jl +++ b/validation/immersed_boundaries/2D_rough_rayleighbenard.jl @@ -77,7 +77,7 @@ function run_simulation(solver, preconditioner; Nr, Ra, Nz, Pr=1, IPS_reltol=1e- if solver == "FFT" model = NonhydrostaticModel(; grid, - # advection = Centered(), + # advection = CenteredSecondOrder(), advection = WENO(order=7), tracers = (:b), buoyancy = BuoyancyTracer(), @@ -87,7 +87,7 @@ function run_simulation(solver, preconditioner; Nr, Ra, Nz, Pr=1, IPS_reltol=1e- else model = NonhydrostaticModel(; grid, pressure_solver = ImmersedPoissonSolver(grid, preconditioner=preconditioner, reltol=IPS_reltol), - # advection = Centered(), + # advection = CenteredSecondOrder(), advection = WENO(order=7), tracers = (:b), buoyancy = BuoyancyTracer(), diff --git a/validation/immersed_boundaries/cylinder_flow_with_tracer.jl b/validation/immersed_boundaries/cylinder_flow_with_tracer.jl index 42ef6af4bf..358b7ec1c7 100644 --- a/validation/immersed_boundaries/cylinder_flow_with_tracer.jl +++ b/validation/immersed_boundaries/cylinder_flow_with_tracer.jl @@ -27,14 +27,14 @@ const nu = 2*R/Re # viscosity """ run_cylinder_steadyflow(output_time_interval = 1, stop_time = 100, arch = CPU(), Nh = 250, ν = 1/20, - momentum_advection = UpwindBiased(order=5), radius = R) + momentum_advection = UpwindBiasedFifthOrder(), radius = R) Run the steady state cylinder validation experiment until `stop_time` using `momentum_advection` scheme or formulation, with horizontal resolution `Nh`, viscosity `ν`, and cylinder radius `R`, on `arch`itecture. """ function run_cylinder_steadystate(; output_time_interval = 1, stop_time = 100, arch = CPU(), Nh = 250, ν = nu, - advection = UpwindBiased(order=5), radius = R) + advection = UpwindBiasedFifthOrder(), radius = R) inside_cylinder(x, y, z) = (x^2 + y^2) <= radius # immersed solid @@ -380,7 +380,7 @@ function analyze_cylinder_steadystate(experiment_name) Plots.savefig(analysisplot, "analysis_" * experiment_name * ".png") end -advection = Centered() +advection = CenteredSecondOrder() experiment_name = run_cylinder_steadystate(Nh = 350, advection = advection, radius = R, stop_time = 100, ν = nu) visualize_cylinder_steadystate(experiment_name) analyze_cylinder_steadystate(experiment_name) diff --git a/validation/immersed_boundaries/internal_tide.jl b/validation/immersed_boundaries/internal_tide.jl index 2085cbab4c..466372c0fa 100644 --- a/validation/immersed_boundaries/internal_tide.jl +++ b/validation/immersed_boundaries/internal_tide.jl @@ -45,7 +45,7 @@ tidal_forcing(x, y, z, t) = 1e-4 * cos(t) for free_surface in (ExplicitFreeSurface, ) model = HydrostaticFreeSurfaceModel(grid = grid_with_bump, - momentum_advection = Centered(), + momentum_advection = CenteredSecondOrder(), free_surface = free_surface(gravitational_acceleration=10), closure = ScalarDiffusivity(VerticallyImplicitTimeDiscretization(), ν=1e-2, κ=1e-2), tracers = :b, diff --git a/validation/immersed_boundaries/nonlinear_topography.jl b/validation/immersed_boundaries/nonlinear_topography.jl index fa330ac43d..5628f8fd17 100644 --- a/validation/immersed_boundaries/nonlinear_topography.jl +++ b/validation/immersed_boundaries/nonlinear_topography.jl @@ -57,7 +57,7 @@ function run_simulation(solver, preconditioner) if solver == "FFT" model = NonhydrostaticModel(; grid, # advection = WENO(), - advection = Centered(), + advection = CenteredSecondOrder(), tracers = :b, buoyancy = BuoyancyTracer(), # timestepper = :RungeKutta3, @@ -67,7 +67,7 @@ function run_simulation(solver, preconditioner) model = NonhydrostaticModel(; grid, pressure_solver = ImmersedPoissonSolver(grid, preconditioner=preconditioner, reltol=1e-8), # advection = WENO(), - advection = Centered(), + advection = CenteredSecondOrder(), tracers = :b, buoyancy = BuoyancyTracer(), # timestepper = :RungeKutta3, diff --git a/validation/immersed_boundaries/resting_stratified_bumpy_ocean.jl b/validation/immersed_boundaries/resting_stratified_bumpy_ocean.jl index f6114b4b1c..892b942951 100644 --- a/validation/immersed_boundaries/resting_stratified_bumpy_ocean.jl +++ b/validation/immersed_boundaries/resting_stratified_bumpy_ocean.jl @@ -5,8 +5,8 @@ using Printf using GLMakie arch = CPU() -tracer_advection = Centered() -momentum_advection = Centered() +tracer_advection = CenteredSecondOrder() +momentum_advection = CenteredSecondOrder() underlying_grid = RectilinearGrid(arch, size=(128, 64), halo=(3, 3), diff --git a/validation/immersed_boundaries/tracer_advection_over_bump.jl b/validation/immersed_boundaries/tracer_advection_over_bump.jl index d22a148677..c619dce4b3 100644 --- a/validation/immersed_boundaries/tracer_advection_over_bump.jl +++ b/validation/immersed_boundaries/tracer_advection_over_bump.jl @@ -6,7 +6,7 @@ using Oceananigans.Operators: ℑxᶠᵃᵃ, ℑyᵃᶠᵃ, ℑzᵃᵃᶠ using Printf arch = CPU() -tracer_advection = Centered() +tracer_advection = CenteredSecondOrder() underlying_grid = RectilinearGrid(arch, size=(128, 64), halo=(3, 3), diff --git a/validation/isotropic_mixing_closures/3d_turbulence_smagorinsky.jl b/validation/isotropic_mixing_closures/3d_turbulence_smagorinsky.jl deleted file mode 100644 index b142949349..0000000000 --- a/validation/isotropic_mixing_closures/3d_turbulence_smagorinsky.jl +++ /dev/null @@ -1,102 +0,0 @@ -using Oceananigans -using Oceananigans.Fields: interpolate! -using Statistics -using Oceananigans.TurbulenceClosures: DirectionallyAveragedCoefficient -using Printf: @printf - -N = 64 -arch = GPU() -grid = RectilinearGrid(arch, size=(N, N, N), extent=(2π, 2π, 2π), topology=(Periodic, Periodic, Periodic)) -coarse_grid = RectilinearGrid(arch, size=(N÷4, N÷4, N÷4), extent=(2π, 2π, 2π), topology=(Periodic, Periodic, Periodic)) - - -function run_3d_turbulence(closure; grid = grid, coarse_grid = coarse_grid) - model = NonhydrostaticModel(; grid, timestepper = :RungeKutta3, - advection = WENO(order=5), - closure = closure) - - random_c = CenterField(coarse_grid) # Technically this shouldn't be a CenterField, but oh well - noise(x, y, z) = rand() - set!(random_c, noise) - - u, v, w = model.velocities - interpolate!(u, random_c) - interpolate!(v, random_c) - interpolate!(w, random_c) - u .-= mean(u) - v .-= mean(v) - w .-= mean(w) - - simulation = Simulation(model, Δt=0.5minimum_zspacing(grid)/maximum(u), stop_time=80) - - wizard = TimeStepWizard(cfl=0.7, max_change=1.1, max_Δt=0.5) - add_callback!(simulation, wizard, IterationInterval(10)) - - start_time = time_ns() # so we can print the total elapsed wall time - progress_message(sim) = @printf("Iteration: %04d, time: %s, Δt: %s, max|u|: %.2e m/s, wall time: %s\n", - iteration(sim), prettytime(time(sim)), prettytime(sim.Δt), maximum(abs, sim.model.velocities.u), - prettytime((time_ns() - start_time) * 1e-9)) - add_callback!(simulation, Callback(progress_message, IterationInterval(100))) - - S² = KernelFunctionOperation{Center, Center, Center}(Oceananigans.TurbulenceClosures.ΣᵢⱼΣᵢⱼᶜᶜᶜ, model.grid, u, v, w) - - if closure.coefficient isa DirectionallyAveragedCoefficient - c²ₛ = model.diffusivity_fields.LM_avg / model.diffusivity_fields.MM_avg - outputs = (; S², c²ₛ) - else - outputs = (; S²) - end - - filename = "3d_turbulence_" * string(nameof(typeof(closure.coefficient))) - simulation.output_writers[:fields] = JLD2OutputWriter(model, outputs, - schedule = TimeInterval(0.6), - filename = filename * ".jld2", - overwrite_existing = true) - run!(simulation) -end - -closures = [SmagorinskyLilly(coefficient=0.16), SmagorinskyLilly(coefficient=DirectionallyAveragedCoefficient(:))] -for closure in closures - @info "Running" closure - run_3d_turbulence(closure) -end - - -@info "Start plotting" -using CairoMakie -set_theme!(Theme(fontsize = 18)) -fig = Figure(size = (800, 800)) -axis_kwargs = (xlabel = "x", ylabel = "y", limits = ((0, 2π), (0, 2π)), aspect = AxisAspect(1)) -n = Observable(1) - -for (i, closure) in enumerate(closures) - closure_name = string(nameof(typeof(closure.coefficient))) - local filename = "3d_turbulence_" * closure_name - @info "Plotting from " * filename - local S²_timeseries = FieldTimeSeries(filename * ".jld2", "S²") - local S² = @lift interior(S²_timeseries[$n], :, :, 1) - - local ax = Axis(fig[2, i]; title = "ΣᵢⱼΣᵢⱼ; $closure_name", axis_kwargs...) - local xc, yc, zc = nodes(S²_timeseries) - heatmap!(ax, xc, yc, S²; colormap = :speed, colorrange = (0, 2)) - - global times = S²_timeseries.times - if closure.coefficient isa DirectionallyAveragedCoefficient - c²ₛ_timeseries = FieldTimeSeries(filename * ".jld2", "c²ₛ") - c²ₛ = interior(c²ₛ_timeseries, 1, 1, 1, :) - global cₛ = sqrt.(max.(c²ₛ, 0)) - local ax_cₛ = Axis(fig[3, 1:length(closures)]; title = "Smagorinsky coefficient", xlabel = "Time", limits = ((0, nothing), (0, 0.25))) - lines!(ax_cₛ, times, cₛ, color=:black, label="Scale Invariant Smagorinsky") - hlines!(ax_cₛ, [0.16], linestyle=:dash, color=:blue) - timesₙ = @lift times[$n] - vlines!(ax_cₛ, timesₙ, linestyle=:dash, color=:black) - end -end - -title = @lift "t = " * string(round(times[$n], digits=2)) * ", cₛ = " * string(round(cₛ[$n], digits=4)) -Label(fig[1, 1:2], title, fontsize=24, tellwidth=false) -frames = 1:length(times) -@info "Making a neat animation of vorticity and speed..." -record(fig, "3d_turbulence_smagorinsky.mp4", frames, framerate=24) do i - n[] = i -end diff --git a/validation/isotropic_mixing_closures/compare_closures.jl b/validation/isotropic_mixing_closures/compare_closures.jl deleted file mode 100644 index b4255aff60..0000000000 --- a/validation/isotropic_mixing_closures/compare_closures.jl +++ /dev/null @@ -1,143 +0,0 @@ -using Oceananigans -using Oceananigans.Units -using Oceananigans.TurbulenceClosures: Smagorinsky, DynamicCoefficient, LagrangianAveraging -using Printf - -stopwatch = Ref(time_ns()) -function progress(sim) - elapsed = 1e-9 * (time_ns() - stopwatch[]) - - msg = @sprintf("Iter: %d, time: %s, Δt: %s, wall time: %s", - iteration(sim), prettytime(sim), prettytime(sim.Δt), prettytime(elapsed)) - - u, v, w = sim.model.velocities - msg *= @sprintf(", max|u|: (%.4f, %.4f, %.4f) m s⁻¹", - maximum(abs, interior(u)), - maximum(abs, interior(v)), - maximum(abs, interior(w))) - - @info msg - - stopwatch[] = time_ns() - - return nothing -end - -function wind_driven_turbulence_simulation(grid, advection, closure; stop_time=9hours, τx=-1e-4, f=1e-4, N²=1e-5) - @info "Running closure $closure" - coriolis = FPlane(; f) - u_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(τx)) - model = NonhydrostaticModel(; grid, closure, coriolis, advection, - boundary_conditions = (; u=u_bcs), - tracers = :b, - buoyancy = BuoyancyTracer()) - - Δz = minimum_zspacing(grid) - δb = N² * Δz - u★ = sqrt(abs(τx)) - uᵢ(x, y, z) = 1e-2 * u★ * (2rand() - 1) - bᵢ(x, y, z) = N² * z + 1e-2 * δb - set!(model, u=uᵢ, v=uᵢ, w=uᵢ, b=bᵢ) - - Δt = 1e-1 * Δz / u★ - simulation = Simulation(model; Δt, stop_time) - conjure_time_step_wizard!(simulation, cfl=0.7) - add_callback!(simulation, progress, IterationInterval(100)) - - return simulation -end - -arch = GPU() -Nx = Ny = Nz = 128 -x = y = (0, 128) -z = (-64, 0) -grid = RectilinearGrid(arch; size=(Nx, Ny, Nz), halo=(5, 5, 5), x, y, z, topology=(Periodic, Periodic, Bounded)) -Δz = @show 10 * round(Int, - z[1] / Nz) -save_interval = 1hour - -schedule = TimeInterval(save_interval) -filename = "wind_driven_WENO9_$Δz" -advection = WENO(order=9) -closure = nothing -simulation = wind_driven_turbulence_simulation(grid, advection, closure) -outputs = merge(simulation.model.velocities, simulation.model.tracers) -output_writer = JLD2OutputWriter(simulation.model, outputs; filename, schedule, overwrite_existing=true) -simulation.output_writers[:jld2] = output_writer -run!(simulation) - -schedule = TimeInterval(save_interval) -filename = "wind_driven_WENO5_$Δz" -advection = WENO(order=5) -closure = nothing -simulation = wind_driven_turbulence_simulation(grid, advection, closure) -outputs = merge(simulation.model.velocities, simulation.model.tracers) -output_writer = JLD2OutputWriter(simulation.model, outputs; filename, schedule, overwrite_existing=true) -simulation.output_writers[:jld2] = output_writer -run!(simulation) - -schedule = TimeInterval(save_interval) -filename = "wind_driven_AMD_$Δz" -advection = Centered(order=2) -closure = AnisotropicMinimumDissipation() -simulation = wind_driven_turbulence_simulation(grid, advection, closure) -outputs = merge(simulation.model.velocities, simulation.model.tracers) -νₑ = simulation.model.diffusivity_fields.νₑ -κₑ = simulation.model.diffusivity_fields.κₑ.b -outputs = merge(outputs, (; νₑ, κₑ)) -output_writer = JLD2OutputWriter(simulation.model, outputs; filename, schedule, overwrite_existing=true) -simulation.output_writers[:jld2] = output_writer -run!(simulation) - -schedule = TimeInterval(save_interval) -filename = "wind_driven_smagorinsky_lilly_$Δz" -advection = Centered(order=2) -closure = SmagorinskyLilly() -simulation = wind_driven_turbulence_simulation(grid, advection, closure) -outputs = merge(simulation.model.velocities, simulation.model.tracers) -νₑ = simulation.model.diffusivity_fields.νₑ -outputs = merge(outputs, (; νₑ)) -output_writer = JLD2OutputWriter(simulation.model, outputs; filename, schedule, overwrite_existing=true) -simulation.output_writers[:jld2] = output_writer -run!(simulation) - -schedule = TimeInterval(save_interval) -filename = "wind_driven_constant_smagorinsky_$Δz" -advection = Centered(order=2) -closure = Smagorinsky(coefficient=0.16) -simulation = wind_driven_turbulence_simulation(grid, advection, closure) -outputs = merge(simulation.model.velocities, simulation.model.tracers) -νₑ = simulation.model.diffusivity_fields.νₑ -outputs = merge(outputs, (; νₑ)) -output_writer = JLD2OutputWriter(simulation.model, outputs; filename, schedule, overwrite_existing=true) -simulation.output_writers[:jld2] = output_writer -run!(simulation) - -schedule = TimeInterval(save_interval) -filename = "wind_driven_directional_smagorinsky_$Δz" -advection = Centered(order=2) -closure = Smagorinsky(coefficient=DynamicCoefficient(averaging=(1, 2))) -simulation = wind_driven_turbulence_simulation(grid, advection, closure) -outputs = merge(simulation.model.velocities, simulation.model.tracers) -𝒥ᴸᴹ = simulation.model.diffusivity_fields.𝒥ᴸᴹ -𝒥ᴹᴹ = simulation.model.diffusivity_fields.𝒥ᴹᴹ -νₑ = simulation.model.diffusivity_fields.νₑ -outputs = merge(outputs, (; 𝒥ᴸᴹ, 𝒥ᴹᴹ, νₑ)) -output_writer = JLD2OutputWriter(simulation.model, outputs; filename, schedule, overwrite_existing=true) -simulation.output_writers[:jld2] = output_writer -run!(simulation) - -schedule = TimeInterval(save_interval) -filename = "wind_driven_lagrangian_smagorinsky_$Δz" -advection = Centered(order=2) -closure = Smagorinsky(coefficient=DynamicCoefficient(averaging=LagrangianAveraging())) -simulation = wind_driven_turbulence_simulation(grid, advection, closure) -outputs = merge(simulation.model.velocities, simulation.model.tracers) -𝒥ᴸᴹ = simulation.model.diffusivity_fields.𝒥ᴸᴹ -𝒥ᴹᴹ = simulation.model.diffusivity_fields.𝒥ᴹᴹ -𝒥ᴸᴹ⁻ = simulation.model.diffusivity_fields.𝒥ᴸᴹ⁻ -𝒥ᴹᴹ⁻ = simulation.model.diffusivity_fields.𝒥ᴹᴹ⁻ -νₑ = simulation.model.diffusivity_fields.νₑ -outputs = merge(outputs, (; 𝒥ᴸᴹ, 𝒥ᴹᴹ, 𝒥ᴸᴹ⁻, 𝒥ᴹᴹ⁻, νₑ)) -output_writer = JLD2OutputWriter(simulation.model, outputs; filename, schedule, overwrite_existing=true) -simulation.output_writers[:jld2] = output_writer -run!(simulation) diff --git a/validation/isotropic_mixing_closures/visualize_closure_comparison.jl b/validation/isotropic_mixing_closures/visualize_closure_comparison.jl deleted file mode 100644 index cc1e261d75..0000000000 --- a/validation/isotropic_mixing_closures/visualize_closure_comparison.jl +++ /dev/null @@ -1,102 +0,0 @@ -using Oceananigans -using GLMakie - -r = "20" - -filenames = [ - "wind_driven_AMD_$r.jld2", - "wind_driven_WENO5_$r.jld2", - "wind_driven_WENO9_$r.jld2", - "wind_driven_constant_smagorinsky_$r.jld2", - "wind_driven_smagorinsky_lilly_$r.jld2", - "wind_driven_directional_smagorinsky_$r.jld2", - "wind_driven_lagrangian_smagorinsky_$r.jld2", -] - -labels = [ - "AMD", - "WENO(order=5)", - "WENO(order=9)", - "Constant Smagorinsky", - "Smagorinsky-Lilly", - "Dynamic Smagorinsky (directional average)", - "Dynamic Smagorinsky (Lagrangian average)", -] - -hasnu = [ - true, - false, - false, - false, - true, - true, - true, -] - -Bs = [] -Us = [] -w²s = [] -νs = [] -for (hasν, name) in zip(hasnu, filenames) - bt = FieldTimeSeries(name, "b", backend=OnDisk()) - ut = FieldTimeSeries(name, "u", backend=OnDisk()) - wt = FieldTimeSeries(name, "w", backend=OnDisk()) - - bn = bt[end] - un = ut[end] - wn = wt[end] - - Bn = compute!(Field(Average(bn, dims=(1, 2)))) - Un = compute!(Field(Average(un, dims=(1, 2)))) - w²n = compute!(Field(Average(wn^2, dims=(1, 2)))) - - push!(Bs, Bn) - push!(Us, Un) - push!(w²s, w²n) - - if hasν - νt = FieldTimeSeries(name, "νₑ", backend=OnDisk()) - νn = νt[end] - Nun = compute!(Field(Average(νn, dims=(1, 2)))) - push!(νs, Nun) - else - push!(νs, nothing) - end -end - -fig = Figure(size=(1400, 800)) - -axb = Axis(fig[1, 1], ylabel="z (m)", xlabel="Buoyancy (m s⁻²)") -axu = Axis(fig[1, 2], ylabel="z (m)", xlabel="x-velocity (m s⁻¹)") -axw = Axis(fig[1, 3], ylabel="z (m)", xlabel="Vertical velocity variance, w² (m² s⁻²)") -axν = Axis(fig[1, 4], ylabel="z (m)", xlabel="Eddy viscosity (m² s⁻¹)") - -for (label, Bn) in zip(labels, Bs) - lines!(axb, Bn; label) -end - -for Un in Us - lines!(axu, Un, label="u") -end - -for w²n in w²s - lines!(axw, w²n) -end - -for i = 1:length(hasnu) - hasν = hasnu[i] - νn = νs[i] - label = labels[i] - if hasν - lines!(axν, νn; label) - else - lines!(axν, [NaN], [NaN]) - end -end - -axislegend(axb, position=:lt) -axislegend(axν, position=:rb) - -display(fig) - -save("closure_comparison_$r.png", fig) diff --git a/validation/isotropic_mixing_closures/wall_flow.jl b/validation/isotropic_mixing_closures/wall_flow.jl deleted file mode 100644 index c3677a08d0..0000000000 --- a/validation/isotropic_mixing_closures/wall_flow.jl +++ /dev/null @@ -1,136 +0,0 @@ -using Oceananigans -using Oceananigans.Units -using Oceananigans.Grids: znode -using Oceananigans.TurbulenceClosures: DirectionallyAveragedCoefficient -using Printf: @printf -simname = "wall_flow" - -const κ = 0.4 - -function run_wall_flow(closure; arch=CPU(), H=1, L=2π*H, N=32, u★=1, z₀ = 1e-4*H, stop_time=30) - grid = RectilinearGrid(arch, size=(N, N, N÷2), topology=(Periodic, Periodic, Bounded), - x=(0, L), y=(0, L), z=(0, H)) - @show grid - - z₁ = first(znodes(grid, Center())) - cᴰ = (κ / log(z₁ / z₀))^2 - - @inline drag_u(x, y, t, u, v, p) = - p.cᴰ * √(u^2 + v^2) * u - @inline drag_v(x, y, t, u, v, p) = - p.cᴰ * √(u^2 + v^2) * v - - u_bcs = FieldBoundaryConditions(bottom=FluxBoundaryCondition(drag_u, field_dependencies = (:u, :v), parameters = (; cᴰ))) - v_bcs = FieldBoundaryConditions(bottom=FluxBoundaryCondition(drag_v, field_dependencies = (:u, :v), parameters = (; cᴰ))) - - @inline x_pressure_gradient(x, y, z, t, p) = p.u★^2 / p.H - u_forcing = Forcing(x_pressure_gradient, parameters=(; u★, H)) - - model = NonhydrostaticModel(; grid, timestepper = :RungeKutta3, - advection = WENO(grid, order=5), - boundary_conditions = (; u=u_bcs, v=v_bcs), - forcing = (; u = u_forcing), - closure = closure) - @show model - - noise(x, y, z) = u★ * randn() * exp(-(z/0.1)^2) - u₀(x, y, z) = (u★/κ) * log(z/z₀) - set!(model, u=u₀, v=noise, w=noise) - - Δt₀ = 1e-4 * (H / u★) / N - simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) - - wizard = TimeStepWizard(max_change=1.1, cfl=0.9) - simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(2)) - - start_time = time_ns() # so we can print the total elapsed wall time - progress_message(sim) = @printf("Iteration: %04d, time: %s, Δt: %s, max|u|: %.2e m/s, wall time: %s\n", - iteration(sim), prettytime(time(sim)), prettytime(sim.Δt), maximum(abs, sim.model.velocities.u), - prettytime((time_ns() - start_time) * 1e-9)) - add_callback!(simulation, Callback(progress_message, IterationInterval(100))) - - closure_name = string(nameof(typeof(closure.coefficient))) - - u, v, w = model.velocities - - U = Average(u, dims=(1,2)) - z = @show KernelFunctionOperation{Nothing, Nothing, Face}(znode, model.grid, Center(), Center(), Face()) - ϕ = @show @at (Nothing, Nothing, Center) κ * z / u★ * ∂z(Field(U)) - - if closure.coefficient isa DirectionallyAveragedCoefficient - cₛ² = model.diffusivity_fields.LM_avg / model.diffusivity_fields.MM_avg - else - cₛ² = Field{Nothing, Nothing, Center}(grid) - cₛ² .= model.closure.coefficient^2 - end - outputs = (; u, w, U, ϕ, cₛ²) - - simulation.output_writers[:fields] = NetCDFOutputWriter(model, outputs; - filename = joinpath(@__DIR__, simname *"_"* closure_name *".nc"), - schedule = TimeInterval(1), - global_attributes = (; u★, z₀, H, L), - overwrite_existing = true) - run!(simulation) - -end - -closures = [SmagorinskyLilly(coefficient=0.16), SmagorinskyLilly(coefficient=DirectionallyAveragedCoefficient(1,2))] -for closure in closures - @info "Running" closure - run_wall_flow(closure, N=64, stop_time=60, arch=GPU()) -end - - - -@info "Start plotting" -using CairoMakie -using NCDatasets -using Statistics: mean -set_theme!(Theme(fontsize = 18)) -fig = Figure(size = (800, 500)) -ax1 = Axis(fig[2, 1]; xlabel = "cₛ", ylabel = "z", limits = ((0, 0.3), (0, 0.8))) -ax2 = Axis(fig[2, 2]; xlabel = "z", ylabel = "U", limits = ((1e-3, 1), (10, 30)), xscale = log10) -ax3 = Axis(fig[2:3, 3]; xlabel = "ϕ = κ z ∂z(U) / u★", ylabel = "z", limits = ((0.0, 2.0), (0, 0.8))) -ax4 = Axis(fig[3, 1]; xlabel = "x (SmagLilly)", ylabel = "y") -ax5 = Axis(fig[3, 2]; xlabel = "x (ScaleInv)", ylabel = "y") -n = Observable(1) - -colors = [:red, :blue] -Δt = 10 -for (i, closure) in enumerate(closures) - closure_name = string(nameof(typeof(closure.coefficient))) - local filename = simname * "_" * closure_name - @info "Plotting from " * filename - ds = NCDataset(filename * ".nc", "r") - - xc, yc, zc = ds["xC"], ds["yC"], ds["zC"] - - #cₛ² = @lift sqrt.(max.(ds["cₛ²"], 0))[:, $n] - cₛ² = mean(sqrt.(max.(ds["cₛ²"], 0))[:, end-Δt:end], dims=2)[:, 1] - scatterlines!(ax1, cₛ², zc, color=colors[i], markercolor=colors[i], label=closure_name) - - #U = @lift ds["U"][:, $n] - U = mean(ds["U"][:, end-Δt:end], dims=2)[:, 1] - scatterlines!(ax2, zc, U, color=colors[i], markercolor=colors[i]) - - #ϕ = @lift ds["ϕ"][:, $n] - ϕ = mean(ds["ϕ"][:, end-Δt:end], dims=2)[:, 1] - scatterlines!(ax3, ϕ, zc, color=colors[i], markercolor=colors[i]) - - if i == 1 - û = (ds.attrib["u★"] / κ) * log.(zc / ds.attrib["z₀"]) - lines!(ax2, zc, û, color=:black) - lines!(ax3, ones(length(zc)), zc, color=:black, linestyle=:dash) - end - - u = @lift ds["u"][:, :, 4, $n] - heatmap!(fig[3, i], xc, yc, u,) - - global times = ds["time"] -end - -axislegend(ax1, labelsize=10) -title = @lift "t = " * string(round(times[$n], digits=2)) -Label(fig[1, 1:2], title, fontsize=24, tellwidth=false) -frames = 1:length(times) -record(fig, simname * ".mp4", frames, framerate=8) do i - n[] = i -end diff --git a/validation/lagrangian_particles/particles_in_convection.jl b/validation/lagrangian_particles/particles_in_convection.jl index 091a65047a..4be9fef74a 100644 --- a/validation/lagrangian_particles/particles_in_convection.jl +++ b/validation/lagrangian_particles/particles_in_convection.jl @@ -41,7 +41,7 @@ particles = LagrangianParticles(x=x₀, y=y₀, z=z₀, restitution=0) b_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(1e-8)) model = NonhydrostaticModel(; grid, particles, - advection = UpwindBiased(order=5), + advection = UpwindBiasedFifthOrder(), timestepper = :RungeKutta3, tracers = :b, buoyancy = BuoyancyTracer(), diff --git a/validation/mesoscale_turbulence/zonally_averaged_abernathey_channel.jl b/validation/mesoscale_turbulence/zonally_averaged_abernathey_channel.jl index 8a7303f77d..6863a0e7b5 100644 --- a/validation/mesoscale_turbulence/zonally_averaged_abernathey_channel.jl +++ b/validation/mesoscale_turbulence/zonally_averaged_abernathey_channel.jl @@ -142,7 +142,7 @@ closure = (horizontal_closure, vertical_closure) model = NonhydrostaticModel(architecture, grid = grid, - advection = UpwindBiased(order=5), + advection = UpwindBiasedFifthOrder(), buoyancy = BuoyancyTracer(), coriolis = coriolis, closure = closure, diff --git a/validation/multi_region/multi_region_internal_tide.jl b/validation/multi_region/multi_region_internal_tide.jl index 7f0312e659..99fcae1e33 100644 --- a/validation/multi_region/multi_region_internal_tide.jl +++ b/validation/multi_region/multi_region_internal_tide.jl @@ -49,7 +49,7 @@ mrg_with_bump = MultiRegionGrid(grid_with_bump, partition=XPartition(2), device tidal_forcing(x, y, z, t) = 1e-4 * cos(t) model = HydrostaticFreeSurfaceModel(grid = mrg_with_bump, - momentum_advection = Centered(), + momentum_advection = CenteredSecondOrder(), free_surface = ExplicitFreeSurface(gravitational_acceleration=10), closure = ScalarDiffusivity(VerticallyImplicitTimeDiscretization(), ν=1e-2, κ=1e-2), tracers = :b, @@ -80,7 +80,7 @@ run!(simulation) """ model_ref = HydrostaticFreeSurfaceModel(grid = grid_with_bump, - momentum_advection = Centered(), + momentum_advection = CenteredSecondOrder(), free_surface = ExplicitFreeSurface(gravitational_acceleration=10), closure = ScalarDiffusivity(VerticallyImplicitTimeDiscretization(), ν=1e-2, κ=1e-2), tracers = :b, diff --git a/validation/open_boundaries/oscillating_flow.jl b/validation/open_boundaries/oscillating_flow.jl index 966b0d35b5..6bf48337ea 100644 --- a/validation/open_boundaries/oscillating_flow.jl +++ b/validation/open_boundaries/oscillating_flow.jl @@ -1,11 +1,12 @@ -# This validation script shows open boundaries working in a simple case where a flow past a 2D -# cylinder oscillates in two directions. All boundaries have the same -# `FlatExtrapolationOpenBoundaryCondition`s. This is similar to a more realistic case where we know -# some arbitary external conditions. First we test an xy flow and then we test an xz flow (the -# forcings and boundary conditions originally designed for `v` aere then used for `w` without -# modification). -# -# This case also has a stretched grid to validate the matching scheme on a stretched grid. +# This validation script shows open boundaries working in a simple case where the +# oscillates sinusoidally so changes sign across two open boundaries. This is similar +# to a more realistic case where we know some arbitary external conditions. +# This necessitates using a combination allowing information to exit the domain, in +# this case by setting the wall normal velocity gradient to zero, as well as forcing +# to the external value in this example by relaxing to it. + +# This case also has a stretched grid to validate the zero wall normal velocity +# gradient matching scheme on a stretched grid. using Oceananigans, CairoMakie using Oceananigans.BoundaryConditions: FlatExtrapolationOpenBoundaryCondition @@ -21,125 +22,133 @@ end architecture = CPU() # model parameters +Re = 200 U = 1 -D = 1.0 -T = 50 +D = 1. +resolution = D / 10 + +# add extra downstream distance to see if the solution near the cylinder changes +extra_downstream = 0 cylinder = Cylinder(; D) -L = 10 -Nx = Ny = 40 +x = (-5, 5 + extra_downstream) .* D -β = 0.2 -x_faces(i) = L/2 * (β * ((2 * (i - 1)) / Nx - 1)^3 + (2 * (i - 1)) / Nx - 1) / (β + 1) -y = (-L/2, +L/2) .* D +Ny = Int(10 / resolution) -xygrid = RectilinearGrid(architecture; topology = (Bounded, Bounded, Flat), size = (Nx, Ny), x = x_faces, y) -xzgrid = RectilinearGrid(architecture; topology = (Bounded, Flat, Bounded), size = (Nx, Ny), x = x_faces, z = y) +Nx = Int((10 + extra_downstream) / resolution) -Δt = .5 * minimum_xspacing(xygrid) / abs(U) +function Δy(j) + if Ny/2 - 2/resolution < j < Ny/2 + 2/resolution + return resolution + elseif j <= Ny/2 - 2/resolution + return resolution * (1 + (Ny/2 - 2/resolution - j) / (Ny/2 - 2/resolution)) + elseif j >= Ny/2 + 2/resolution + return resolution * (1 + (j - Ny/2 - 2/resolution) / (Ny/2 - 2/resolution)) + else + Throw(ArgumentError("$j not in range")) + end +end -@inline u∞(y, t, p) = p.U * cos(t * 2π / p.T) * (1 + 0.01 * randn()) -@inline v∞(x, t, p) = p.U * sin(t * 2π / p.T) * (1 + 0.01 * randn()) +y(j) = sum(Δy.([1:j;])) - sum(Δy.([1:Ny;]))/2 -function run_cylinder(grid, boundary_conditions; plot=true, stop_time = 50, simname = "") - @info "Testing $simname with grid" grid +ν = U * D / Re - cylinder_forcing = Relaxation(; rate = 1 / (2 * Δt), mask = cylinder) +closure = ScalarDiffusivity(;ν, κ = ν) - global model = NonhydrostaticModel(; grid, - advection = UpwindBiased(order=5), - forcing = (u = cylinder_forcing, v = cylinder_forcing, w = cylinder_forcing), - boundary_conditions) +grid = RectilinearGrid(architecture; topology = (Bounded, Bounded, Flat), size = (Nx, Ny), x = y, y = x) - @info "Constructed model" +T = 20 / U - # initial noise to induce turbulance faster - set!(model, u = U) +@inline u(t, p) = p.U * sin(t * 2π / p.T) +@inline u(y, t, p) = u(t, p) - @info "Set initial conditions" - simulation = Simulation(model; Δt = Δt, stop_time = stop_time) +relaxation_timescale = 0.15 - # Callbacks - wizard = TimeStepWizard(cfl = 0.1) - simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(100)) +u_boundaries = FieldBoundaryConditions(east = FlatExtrapolationOpenBoundaryCondition(u; relaxation_timescale, parameters = (; U, T)), + west = FlatExtrapolationOpenBoundaryCondition(u; relaxation_timescale, parameters = (; U, T)), + south = GradientBoundaryCondition(0), + north = GradientBoundaryCondition(0)) - progress(sim) = @info "$(time(sim)) with Δt = $(prettytime(sim.Δt)) in $(prettytime(sim.run_wall_time))" - simulation.callbacks[:progress] = Callback(progress, IterationInterval(1000)) +v_boundaries = FieldBoundaryConditions(east = GradientBoundaryCondition(0), + west = GradientBoundaryCondition(0), + south = FlatExtrapolationOpenBoundaryCondition(0; relaxation_timescale), + north = FlatExtrapolationOpenBoundaryCondition(0; relaxation_timescale)) - u, v, w = model.velocities - filename = "cylinder_$(simname)_Nx_$Nx.jld2" +Δt = .3 * resolution / U - if grid isa Oceananigans.Grids.ZFlatGrid - outputs = (; model.velocities..., ζ = (@at (Center, Center, Center) ∂x(v) - ∂y(u))) - elseif grid isa Oceananigans.Grids.YFlatGrid - outputs = (; model.velocities..., ζ = (@at (Center, Center, Center) ∂x(w) - ∂z(u))) - end - simulation.output_writers[:velocity] = JLD2OutputWriter(model, outputs, - overwrite_existing = true, - filename = filename, - schedule = TimeInterval(0.5), - with_halos = true) - run!(simulation) - - if plot - # load the results - ζ_ts = FieldTimeSeries(filename, "ζ") - u_ts = FieldTimeSeries(filename, "u") - @info "Loaded results" - - xζ, yζ, zζ = nodes(ζ_ts, with_halos=true) - xu, yu, zu = nodes(u_ts, with_halos=true) - - # plot the results - fig = Figure(size = (600, 600)) - n = Observable(1) - - if grid isa Oceananigans.Grids.ZFlatGrid - ζ_plt = @lift ζ_ts[:, :, 1, $n].parent - ax = Axis(fig[1, 1], aspect = DataAspect(), xlabel = "x", ylabel = "y", width = 400, height = 400, title = "ζ") - heatmap!(ax, collect(xζ), collect(yζ), ζ_plt, colorrange = (-2, 2), colormap = :curl) - - u_plt = @lift u_ts[:, :, 1, $n].parent - ax = Axis(fig[1, 2], aspect = DataAspect(), xlabel = "x", ylabel = "y", width = 400, height = 400, title = "u") - heatmap!(ax, collect(xu), collect(yu), u_plt, colorrange = (-2, 2), colormap = :curl) - - elseif grid isa Oceananigans.Grids.YFlatGrid - ζ_plt = @lift ζ_ts[:, 1, :, $n].parent - ax = Axis(fig[1, 1], aspect = DataAspect(), xlabel = "x", ylabel = "z", width = 400, height = 400, title = "ζ") - heatmap!(ax, collect(xζ), collect(zζ), ζ_plt, colorrange = (-2, 2), colormap = :curl) - - u_plt = @lift u_ts[:, 1, :, $n].parent - ax = Axis(fig[1, 2], aspect = DataAspect(), xlabel = "x", ylabel = "z", width = 400, height = 400, title = "u") - heatmap!(ax, collect(xu), collect(zu), u_plt, colorrange = (-2, 2), colormap = :curl) - - end - resize_to_layout!(fig) - record(fig, "ζ_$filename.mp4", 1:length(ζ_ts.times), framerate = 16) do i; - n[] = i - i % 10 == 0 && @info "$(n.val) of $(length(ζ_ts.times))" - end - end -end +@show Δt -matching_scheme_name(obc) = string(nameof(typeof(obc.classification.matching_scheme))) -for grid in (xygrid, xzgrid) +u_forcing = Relaxation(; rate = 1 / (2 * Δt), mask = cylinder) +v_forcing = Relaxation(; rate = 1 / (2 * Δt), mask = cylinder) - u_fe = FlatExtrapolationOpenBoundaryCondition(u∞, parameters = (; U, T), relaxation_timescale = 1) - v_fe = FlatExtrapolationOpenBoundaryCondition(v∞, parameters = (; U, T), relaxation_timescale = 1) - w_fe = FlatExtrapolationOpenBoundaryCondition(v∞, parameters = (; U, T), relaxation_timescale = 1) +model = NonhydrostaticModel(; grid, + closure, + forcing = (u = u_forcing, v = v_forcing), + boundary_conditions = (u = u_boundaries, v = v_boundaries)) - u_boundaries_fe = FieldBoundaryConditions(west = u_fe, east = u_fe) - v_boundaries_fe = FieldBoundaryConditions(south = v_fe, north = v_fe) - w_boundaries_fe = FieldBoundaryConditions(bottom = w_fe, top = w_fe) +@info "Constructed model" - if grid isa Oceananigans.Grids.ZFlatGrid - boundary_conditions = (u = u_boundaries_fe, v = v_boundaries_fe) - simname = "xy_" * matching_scheme_name(u_boundaries_fe.east) - elseif grid isa Oceananigans.Grids.YFlatGrid - boundary_conditions = (u = u_boundaries_fe, w = w_boundaries_fe) - simname = "xz_" * matching_scheme_name(u_boundaries_fe.east) - end - run_cylinder(grid, boundary_conditions, simname = simname, stop_time = T) +# initial noise to induce turbulance faster +set!(model, u = (x, y) -> randn() * U * 0.01, v = (x, y) -> randn() * U * 0.01) + +@info "Set initial conditions" + +simulation = Simulation(model; Δt = Δt, stop_time = 300) + +wizard = TimeStepWizard(cfl = 0.3, max_Δt = Δt) + +simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(100)) + +progress(sim) = @info "$(time(sim)) with Δt = $(prettytime(sim.Δt)) in $(prettytime(sim.run_wall_time))" + +simulation.callbacks[:progress] = Callback(progress, IterationInterval(1000)) + +simulation.output_writers[:velocity] = JLD2OutputWriter(model, model.velocities, + overwrite_existing = true, + filename = "oscillating_cylinder_$(extra_downstream)_Re_$Re.jld2", + schedule = TimeInterval(1), + with_halos = true) + +run!(simulation) + +# load the results + +u_ts = FieldTimeSeries("oscillating_cylinder_$(extra_downstream)_Re_$Re.jld2", "u") +v_ts = FieldTimeSeries("oscillating_cylinder_$(extra_downstream)_Re_$Re.jld2", "v") + +u′, v′, w′ = Oceananigans.Fields.VelocityFields(u_ts.grid) + +ζ = Field((@at (Center, Center, Center) ∂x(v′)) - (@at (Center, Center, Center) ∂y(u′))) + +# there is probably a more memory efficient way todo this + +ζ_ts = zeros(size(grid, 1), size(grid, 2), length(u_ts.times)) # u_ts.grid so its always on cpu + +for n in 1:length(u_ts.times) + set!(u′, u_ts[n]) + set!(v′, v_ts[n]) + compute!(ζ) + ζ_ts[:, :, n] = interior(ζ, :, :, 1) end +@info "Loaded results" + +# plot the results + +fig = Figure(size = (600, 600)) + +xc, yc, zc = nodes(ζ) + +ax = Axis(fig[1, 1], aspect = DataAspect(), limits = (minimum(xc), maximum(xc), minimum(yc), maximum(yc))) + +n = Observable(1) + +ζ_plt = @lift ζ_ts[:, :, $n] + +contour!(ax, xc, yc, ζ_plt, levels = [-1, 1], colorrange = (-1, 1), colormap = :roma) + +record(fig, "oscillating_ζ_Re_$(Re)_no_exterior.mp4", 1:length(u_ts.times), framerate = 5) do i; + n[] = i + i % 10 == 0 && @info "$(n.val) of $(length(u_ts.times))" +end \ No newline at end of file diff --git a/validation/periodic_advection/periodic_advection.jl b/validation/periodic_advection/periodic_advection.jl index 900fe24a74..954900e7c3 100644 --- a/validation/periodic_advection/periodic_advection.jl +++ b/validation/periodic_advection/periodic_advection.jl @@ -99,7 +99,7 @@ end L = 1 ϕs = (ϕ_Gaussian, ϕ_Square) time_steppers = (:RungeKutta3,) -advection_schemes = (Centered(order=2), Centered(order=4), UpwindBiased(order=3), UpwindBiased(order=5), WENO()) +advection_schemes = (CenteredSecondOrder(), CenteredFourthOrder(), UpwindBiasedThirdOrder(), UpwindBiasedFifthOrder(), WENO()) Ns = [16, 64] CFLs = (0.5, 1.7) Us = [+1, -1] diff --git a/validation/stommel_gyre/plot_stommel_gyre_advection.py b/validation/stommel_gyre/plot_stommel_gyre_advection.py index 92303c2f1e..9d479396ba 100644 --- a/validation/stommel_gyre/plot_stommel_gyre_advection.py +++ b/validation/stommel_gyre/plot_stommel_gyre_advection.py @@ -6,7 +6,7 @@ import ffmpeg shapes = ("Gaussian", "Square") -schemes = ("Centered", "CenteredFourthOrder", "WENO") +schemes = ("CenteredSecondOrder", "CenteredFourthOrder", "WENO") Ns = (32, 256) CFLs = (0.05, 0.30) diff --git a/validation/stommel_gyre/stommel_gyre_advection.jl b/validation/stommel_gyre/stommel_gyre_advection.jl index 95bdae46f0..d381c69ca8 100644 --- a/validation/stommel_gyre/stommel_gyre_advection.jl +++ b/validation/stommel_gyre/stommel_gyre_advection.jl @@ -100,7 +100,7 @@ ic_name(::typeof(ϕ_λ_Gaussian)) = ic_name(ϕ_Gaussian) ic_name(::typeof(ϕ_λ_Square)) = ic_name(ϕ_Square) ϕ_λs = (ϕ_λ_Gaussian, ϕ_λ_Square) -schemes = (Centered(), CenteredFourthOrder(), WENO()) +schemes = (CenteredSecondOrder(), CenteredFourthOrder(), WENO()) Ns = (32, 256) CFLs = (0.05,) diff --git a/validation/thermal_bubble/plot_thermal_bubble_advection.py b/validation/thermal_bubble/plot_thermal_bubble_advection.py index cbbb332bdb..0f2a94dcdd 100644 --- a/validation/thermal_bubble/plot_thermal_bubble_advection.py +++ b/validation/thermal_bubble/plot_thermal_bubble_advection.py @@ -5,7 +5,7 @@ import cmocean import ffmpeg -# schemes = ("Centered", "CenteredFourthOrder", "WENO") +# schemes = ("CenteredSecondOrder", "CenteredFourthOrder", "WENO") schemes = ("WENO",) Ns = (32, 128) diff --git a/validation/thermal_bubble/thermal_bubble.jl b/validation/thermal_bubble/thermal_bubble.jl index 5358786ae0..eb23bf1391 100644 --- a/validation/thermal_bubble/thermal_bubble.jl +++ b/validation/thermal_bubble/thermal_bubble.jl @@ -69,7 +69,7 @@ function print_progress(simulation) progress, i, t, u_max, w_max, T_min, T_max, CFL) end -schemes = (WENO(), Centered(order=4)) +schemes = (WENO(), CenteredFourthOrder()) Ns = (32, 128) for scheme in schemes, N in Ns diff --git a/validation/vertical_mixing_closures/gpu_tkevd_ensemble.jl b/validation/vertical_mixing_closures/gpu_tkevd_ensemble.jl index 149f86a4ab..64cf261388 100644 --- a/validation/vertical_mixing_closures/gpu_tkevd_ensemble.jl +++ b/validation/vertical_mixing_closures/gpu_tkevd_ensemble.jl @@ -17,13 +17,13 @@ grid = RectilinearGrid(size=sz, halo=halo, z=(-128, 0), topology=(Flat, Flat, Bo closure = CuArray([CATKEVerticalDiffusivity() for i=1:Ex, j=1:Ey]) -Jᵇ = CuArray([+1e-8 for i=1:Ex, j=1:Ey]) -τx = CuArray([-1e-4 for i=1:Ex, j=1:Ey]) -τy = CuArray([0.0 for i=1:Ex, j=1:Ey]) +Qᵇ = CuArray([+1e-8 for i=1:Ex, j=1:Ey]) +Qᵘ = CuArray([-1e-4 for i=1:Ex, j=1:Ey]) +Qᵛ = CuArray([0.0 for i=1:Ex, j=1:Ey]) -u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(τx)) -v_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(τy)) -b_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Jᵇ)) +u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵘ)) +v_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵛ)) +b_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵇ)) # Half rotating, half not f_ij(i, j) = j < Ey/2 ? 1e-4 : 0.0 diff --git a/validation/vertical_mixing_closures/heterogeneous_windy_convection.jl b/validation/vertical_mixing_closures/heterogeneous_windy_convection.jl index 267be16537..6faf2950d0 100644 --- a/validation/vertical_mixing_closures/heterogeneous_windy_convection.jl +++ b/validation/vertical_mixing_closures/heterogeneous_windy_convection.jl @@ -34,15 +34,15 @@ grid = RectilinearGrid(size = (Nx, Ny, Nz), z = z, topology=(Periodic, Bounded, Bounded)) -bottom_height(x, y) = - Lz * (1 - (2y / Ly)^2) -grid = ImmersedBoundaryGrid(grid, PartialCellBottom(bottom_height, minimum_fractional_cell_height=0.1)) +z_bottom(x, y) = - Lz * (1 - (2y / Ly)^2) +grid = ImmersedBoundaryGrid(grid, PartialCellBottom(z_bottom, minimum_fractional_cell_height=0.1)) @show grid -@inline Jᵇ(x, y, t) = 1e-7 -@inline τx(x, y, t) = -1e-3 * cos(π * y / Ly) +@inline Qᵇ(x, y, t) = 1e-7 +@inline Qᵘ(x, y, t) = -1e-3 * cos(π * y / Ly) -b_top_bc = FluxBoundaryCondition(Jᵇ) -u_top_bc = FluxBoundaryCondition(τx) +b_top_bc = FluxBoundaryCondition(Qᵇ) +u_top_bc = FluxBoundaryCondition(Qᵘ) b_bcs = FieldBoundaryConditions(top=b_top_bc) u_bcs = FieldBoundaryConditions(top=u_top_bc) From ebf663972f45e481d647798d76efca070c025fd0 Mon Sep 17 00:00:00 2001 From: Xin Kai Lee Date: Tue, 3 Dec 2024 07:04:50 -0500 Subject: [PATCH 114/114] fix script changes for modewater after equilibration --- ...BLRifirstzone510_doublegyre_model_modewater.jl | 15 ++++++++++----- baseclosure_doublegyre_model_modewater.jl | 13 +++++++------ physicalclosure_doublegyre_model_modewater.jl | 15 ++++++++++----- 3 files changed, 27 insertions(+), 16 deletions(-) diff --git a/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl b/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl index 1f107c5cb9..9fbf9cf1fd 100644 --- a/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl +++ b/NNclosure_Ri_nof_BBLRifirstzone510_doublegyre_model_modewater.jl @@ -21,9 +21,11 @@ using SeawaterPolynomials:TEOS10 using ColorSchemes #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_NN_closure_NDE5_Ri_BBLRifirztzone510_temp" +const Qᵀ_mode = 3.5e-4 +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_NN_closure_NDE5_Ri_BBLRifirztzone510_temp_QT$(Qᵀ_mode)" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +@info "$(FILE_DIR)" mkpath(FILE_DIR) # Architecture @@ -75,7 +77,6 @@ const μ_T = 1/30days const X₀ = -Lx/2 + 800kilometers const Y₀ = -Ly/2 + 1500kilometers const R₀ = 700kilometers -const Qᵀ_mode = 4.5e-4 const σ_mode = 20kilometers ##### @@ -107,7 +108,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), @inline T_ref(y) = T_mid - ΔT / Ly * y -@inline Qᵀ_winter(t) = max(0, -Qᵀ_mode * sin(2π * t / 360days)) +@inline Qᵀ_winter(t) = ifelse(t < 10800days, 0, max(0, -Qᵀ_mode * sin(2π * t / 360days))) @inline Qᵀ_subpolar(x, y, t) = ifelse((x - X₀)^2 + (y - Y₀)^2 <= R₀^2, Qᵀ_winter(t), exp(-(sqrt((x - X₀)^2 + (y - Y₀)^2) - R₀)^2 / (2 * σ_mode^2)) * Qᵀ_winter(t)) @@ -179,7 +180,7 @@ update_state!(model) ##### Simulation building ##### Δt₀ = 5minutes -stop_time = 10800days +stop_time = 12600days simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) @@ -394,7 +395,11 @@ simulation.output_writers[:BBL] = JLD2OutputWriter(model, (; first_index, last_i simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), filename = "$(FILE_DIR)/averaged_fields_streamfunction", - schedule = AveragedTimeInterval(1825days, window=1825days)) + schedule = AveragedTimeInterval(1800days, window=1800days)) + +simulation.output_writers[:complete_fields] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields", + schedule = TimeInterval(1800days)) simulation.output_writers[:checkpointer] = Checkpointer(model, schedule = TimeInterval(730days), diff --git a/baseclosure_doublegyre_model_modewater.jl b/baseclosure_doublegyre_model_modewater.jl index 11328e9ede..40da3d8fd1 100644 --- a/baseclosure_doublegyre_model_modewater.jl +++ b/baseclosure_doublegyre_model_modewater.jl @@ -20,9 +20,11 @@ using SeawaterPolynomials:TEOS10 using ColorSchemes #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_newbaseclosure" +const Qᵀ_mode = 3.5e-4 +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_newbaseclosure_QT$(Qᵀ_mode)" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +@info "$(FILE_DIR)" mkpath(FILE_DIR) # Architecture @@ -73,7 +75,6 @@ const μ_T = 1/30days const X₀ = -Lx/2 + 800kilometers const Y₀ = -Ly/2 + 1500kilometers const R₀ = 700kilometers -const Qᵀ_mode = 4.5e-4 const σ_mode = 20kilometers ##### @@ -105,7 +106,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), @inline T_ref(y) = T_mid - ΔT / Ly * y -@inline Qᵀ_winter(t) = max(0, -Qᵀ_mode * sin(2π * t / 360days)) +@inline Qᵀ_winter(t) = ifelse(t < 10800days, 0, max(0, -Qᵀ_mode * sin(2π * t / 360days))) @inline Qᵀ_subpolar(x, y, t) = ifelse((x - X₀)^2 + (y - Y₀)^2 <= R₀^2, Qᵀ_winter(t), exp(-(sqrt((x - X₀)^2 + (y - Y₀)^2) - R₀)^2 / (2 * σ_mode^2)) * Qᵀ_winter(t)) @@ -177,7 +178,7 @@ update_state!(model) ##### Simulation building ##### Δt₀ = 5minutes -stop_time = 10800days +stop_time = 12600days simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) @@ -316,11 +317,11 @@ simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_output simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), filename = "$(FILE_DIR)/averaged_fields_streamfunction", - schedule = AveragedTimeInterval(1825days, window=1825days)) + schedule = AveragedTimeInterval(1800days, window=1800days)) simulation.output_writers[:complete_fields] = JLD2OutputWriter(model, outputs, filename = "$(FILE_DIR)/instantaneous_fields", - schedule = TimeInterval(1825days)) + schedule = TimeInterval(1800days)) simulation.output_writers[:checkpointer] = Checkpointer(model, schedule = TimeInterval(730days), diff --git a/physicalclosure_doublegyre_model_modewater.jl b/physicalclosure_doublegyre_model_modewater.jl index 6b80b049fb..9a5777b97a 100644 --- a/physicalclosure_doublegyre_model_modewater.jl +++ b/physicalclosure_doublegyre_model_modewater.jl @@ -27,9 +27,11 @@ using ColorSchemes using Glob #%% -filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_CATKEVerticalDiffusivity" +const Qᵀ_mode = 3.5e-4 +filename = "doublegyre_30Cwarmflushbottom10_relaxation_30days_modewater_zWENO5_CATKEVerticalDiffusivity_QT$(Qᵀ_mode)" FILE_DIR = "./Output/$(filename)" # FILE_DIR = "/storage6/xinkai/NN_Oceananigans/$(filename)" +@info "$(FILE_DIR)" mkpath(FILE_DIR) # Architecture @@ -85,7 +87,6 @@ const μ_T = 1/30days const X₀ = -Lx/2 + 800kilometers const Y₀ = -Ly/2 + 1500kilometers const R₀ = 700kilometers -const Qᵀ_mode = 4.5e-4 const σ_mode = 20kilometers ##### @@ -116,7 +117,7 @@ v_bcs = FieldBoundaryConditions( top = FluxBoundaryCondition(0), @inline T_ref(y) = T_mid - ΔT / Ly * y -@inline Qᵀ_winter(t) = max(0, -Qᵀ_mode * sin(2π * t / 360days)) +@inline Qᵀ_winter(t) = ifelse(t < 10800days, 0, max(0, -Qᵀ_mode * sin(2π * t / 360days))) @inline Qᵀ_subpolar(x, y, t) = ifelse((x - X₀)^2 + (y - Y₀)^2 <= R₀^2, Qᵀ_winter(t), exp(-(sqrt((x - X₀)^2 + (y - Y₀)^2) - R₀)^2 / (2 * σ_mode^2)) * Qᵀ_winter(t)) @@ -188,7 +189,7 @@ update_state!(model) ##### Simulation building ##### Δt₀ = 5minutes -stop_time = 10800days +stop_time = 12600days simulation = Simulation(model, Δt = Δt₀, stop_time = stop_time) @@ -332,7 +333,11 @@ simulation.output_writers[:zonal_average] = JLD2OutputWriter(model, zonal_output simulation.output_writers[:streamfunction] = JLD2OutputWriter(model, (; Ψ=Ψ,), filename = "$(FILE_DIR)/averaged_fields_streamfunction", - schedule = AveragedTimeInterval(1825days, window=1825days)) + schedule = AveragedTimeInterval(1800days, window=1800days)) + +simulation.output_writers[:complete_fields] = JLD2OutputWriter(model, outputs, + filename = "$(FILE_DIR)/instantaneous_fields", + schedule = TimeInterval(1800days)) simulation.output_writers[:checkpointer] = Checkpointer(model, schedule = TimeInterval(730days),