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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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/108] 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)