Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement neural network parameterization of salty turbulent mixing in the upper ocean #3819

Draft
wants to merge 119 commits into
base: main
Choose a base branch
from
Draft
Changes from 1 commit
Commits
Show all changes
119 commits
Select commit Hold shift + click to select a range
13cd632
working version on CPU
xkykai Jun 24, 2024
bb70e5c
a working model for GPU!
xkykai Jun 24, 2024
e845dd5
it actually works on the gpu now!
xkykai Jun 25, 2024
b6552ab
build scaling from tuple values
xkykai Jun 25, 2024
c9a9c65
Update XinKaiLocalVerticalDiffusivity with 2 Pr values
xkykai Jun 26, 2024
9a07b82
fix function construct_scaling to construct_zeromeanunitvariance_scaling
xkykai Jun 26, 2024
74c8b61
update using KernelParameters
xkykai Jun 26, 2024
6f97c3f
fix nn closure bug
xkykai Jun 26, 2024
65f15e4
test script for nn closure
xkykai Jun 26, 2024
cb601ce
fix N2 average, tracer diffusivity expression for local diffusivity c…
xkykai Jun 27, 2024
dbe5759
2Pr version of local physical closure
xkykai Jun 28, 2024
e9ac435
nonlocal physical closure of vertical diffusivity
xkykai Jun 28, 2024
5e548da
2Pr nonlocal physical closure
xkykai Jun 28, 2024
772c235
working GPU version of NN closure with scaling and correction
xkykai Jun 28, 2024
9c6873f
validation script for oceananigans NN implementation
xkykai Jul 1, 2024
f6d508a
close file and record video of validation
xkykai Jul 1, 2024
1faa5a6
Update Project.toml with new package dependencies
xkykai Jul 2, 2024
930cc7e
Update NN closure model to use a larger neural network
xkykai Jul 2, 2024
fd95fc6
Remove unused import of StaticArrays
xkykai Jul 2, 2024
72aebc1
Add ComponentArrays dependency
xkykai Jul 2, 2024
c2d336e
add total_size and KernelParameters dependency
xkykai Jul 2, 2024
d4a1156
Coarsen LES data for NN closure model
xkykai Jul 17, 2024
ae7b563
rename file and compare LES with NN closure
xkykai Jul 17, 2024
4d118c2
run LES for hald sinusoid cooling
xkykai Jul 18, 2024
fc081cc
run 3D simulation with limited extent
xkykai Aug 21, 2024
a78109e
add sponge at bottom
xkykai Aug 21, 2024
ca8d78c
fix metres to meters
xkykai Aug 21, 2024
78e7a67
fix temperature flux
xkykai Aug 21, 2024
eaf2183
Calculate average velocities and tracers in 3D model LES simulation
xkykai Aug 21, 2024
7e144ff
reduce size of model
xkykai Aug 21, 2024
ba7f58d
run double gyre with physical closure
xkykai Aug 29, 2024
fcb466a
rename file
xkykai Aug 29, 2024
b63e7d6
fix S forcing
xkykai Aug 29, 2024
67e011c
use RiBasedVerticalDiffusivity as closure
xkykai Aug 29, 2024
55d635a
fix sign errors in tracer fluxes
xkykai Aug 29, 2024
8319ba8
fix boundary conditions, initialize trivial model first
xkykai Aug 29, 2024
229feeb
update boundary conditions, initial state
xkykai Aug 29, 2024
a853275
fix tyop_buoyancy_flux bug
xkykai Aug 29, 2024
477c6be
fix initial conditions, run for 10 years, set up checkpointing
xkykai Aug 29, 2024
232b67d
run double gyre with new physical closure
xkykai Aug 30, 2024
148a582
plotting barotropic streamfunction
xkykai Sep 4, 2024
e9f2f0e
fix plot units
xkykai Sep 4, 2024
7f17cf2
Merge branch 'main' into xk/embed-nn
xkykai Sep 9, 2024
ed8ec7d
run double gyre with CATKE
xkykai Sep 9, 2024
f4081bf
Merge branch 'main' into xk/embed-nn
xkykai Sep 9, 2024
893e529
add TKE tracer for CATKE
xkykai Sep 9, 2024
ba5491d
fix closure to use only CATKE
xkykai Sep 10, 2024
21a8356
update CATKE configuration
xkykai Sep 17, 2024
de190d2
use older dependencies for compatibility with neural networks
xkykai Sep 17, 2024
7e7a04f
add fields to be calculated for CATKE
xkykai Sep 19, 2024
1c31868
local diffusivity for 2step calibration
xkykai Sep 19, 2024
3609d6c
new NN closure with nof and base boundary layer criteria
xkykai Sep 19, 2024
8eaaaae
fix bug in NN closure implementation
xkykai Sep 20, 2024
6535be5
using Grids.total_size
xkykai Sep 23, 2024
812df88
add using KernelParameters
xkykai Sep 23, 2024
eb50cca
update NN model
xkykai Sep 24, 2024
b2cd2fa
use BBL integral metric to compute base of boundary layer
xkykai Sep 24, 2024
d3fa8b8
Update xin_kai_vertical_diffusivity.jl
xkykai Oct 4, 2024
2448003
remove type piracy TEOS10.s
xkykai Oct 6, 2024
cb5f864
NN closure using BBL zone below nonbabkground kappa
xkykai Oct 6, 2024
f443296
run double gyre with NDE BBLkappazonelast41
xkykai Oct 6, 2024
2f3eed9
change initialized state to 8day restoration forcing
xkykai Oct 6, 2024
987f1ed
run double gyre withj baseclosure initialized
xkykai Oct 6, 2024
c4351f7
NN closure for augmenting flux in a zone below MLD
xkykai Oct 6, 2024
19f71d3
8 day relaxation double gyre for baseclosure
xkykai Oct 8, 2024
3c4102a
uncomment CairoMakie
xkykai Oct 8, 2024
b122e7b
NN closure and CATKE with 8 day restoration and warm flush
xkykai Oct 8, 2024
e0bd65a
fix initial temperature issue
xkykai Oct 8, 2024
8dcb6d3
add zonal average calculations to double gyre simulation
xkykai Oct 9, 2024
132188a
run double gyre with seasonal forcing
xkykai Oct 10, 2024
57c3ac5
increase simulation run time
xkykai Oct 16, 2024
d5a320e
wall restoration to maintain strratification
xkykai Oct 17, 2024
0cd618a
baseclosure doublegyre with wallrestoration
xkykai Oct 17, 2024
4abf7e0
change filenames
xkykai Oct 17, 2024
7cf7090
using wider zone for NN closure
xkykai Oct 17, 2024
70dbe2f
fix variable error
xkykai Oct 17, 2024
fd57bf8
update NN configuration
xkykai Oct 18, 2024
b582447
NDE double gyre script for seasonal forcing and wall restoration
xkykai Oct 18, 2024
6235bc4
run NNclosure with Ri nof BBLkappazonelast55
xkykai Oct 20, 2024
3609c21
updated NN model with no Ri
xkykai Oct 22, 2024
823e8c7
run double gyre with NNclosure with no Ri input kappazonelast55
xkykai Oct 22, 2024
893b7bc
add fluxes calculations and zonal average
xkykai Oct 24, 2024
ea77669
Merge branch 'main' into xk/embed-nn
xkykai Oct 24, 2024
3428f5d
Merge branch 'main' into xk/embed-nn
xkykai Oct 24, 2024
4a5a3e1
using centered second order in z instead of WENO
xkykai Oct 24, 2024
0f6bf66
remove background vertical scalar diffusivity
xkykai Oct 24, 2024
8385229
diffusivity fields indexing for base closure
xkykai Oct 24, 2024
45eef0a
change temperature restoration to 30days
xkykai Oct 30, 2024
e29bfd9
change vertical advection scheme to WENO5
xkykai Oct 30, 2024
0fba127
update neural network to new weights
xkykai Oct 30, 2024
73afa68
fix file name change
xkykai Oct 30, 2024
1fab19f
run double gyre with centered second order and wall restoration
xkykai Oct 31, 2024
49d901f
fix dynamic function invocation that doesn't affect the baseclosure s…
xkykai Oct 31, 2024
534d21d
add y resolution variable
xkykai Oct 31, 2024
2551c38
run NN closure recording xz slices
xkykai Nov 1, 2024
ff272c4
recording and plotting xz yz slices and fluxes
xkykai Nov 1, 2024
93bf0bc
change default z advection scheme to weno 5
xkykai Nov 1, 2024
d36677f
run with linear ramp seasonal forcing
xkykai Nov 6, 2024
ce68cc7
fix T_seasonal
xkykai Nov 7, 2024
e1603ca
fix temperature restoration
xkykai Nov 8, 2024
a3ab871
implement different NNclosure with Ri zone
xkykai Nov 21, 2024
51acb9d
run double gyre on updated closure
xkykai Nov 21, 2024
33f2ea8
new NN closure witth Rifirstzone
xkykai Nov 22, 2024
36cecb6
fix NN closure
xkykai Nov 22, 2024
946d710
fix NN_closure
xkykai Nov 22, 2024
0c55c44
updated base closure with new calibration
xkykai Nov 26, 2024
dada2c5
run base closure and CATKE with mode waters
xkykai Nov 26, 2024
1d0ef1f
fix advection scheme and file name
xkykai Nov 26, 2024
4898be4
new NN closure including Ri
xkykai Nov 26, 2024
3e84aaf
actual new closure including Ri
xkykai Nov 26, 2024
88f7fa4
run double gyre with new NN closure
xkykai Nov 26, 2024
030c890
run CATKE for 10800days
xkykai Nov 26, 2024
f4f7383
Merge branch 'main' into xk/embed-nn
xkykai Nov 28, 2024
35dbad3
flush mode water after equilibration
xkykai Dec 2, 2024
88e3cf1
fix filename
xkykai Dec 2, 2024
5b34027
Revert "fix filename"
xkykai Dec 3, 2024
8e458e0
Revert "flush mode water after equilibration"
xkykai Dec 3, 2024
34f7437
Revert "Merge branch 'main' into xk/embed-nn"
xkykai Dec 3, 2024
ebf6639
fix script changes for modewater after equilibration
xkykai Dec 3, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
new NN closure witth Rifirstzone
  • Loading branch information
xkykai committed Nov 22, 2024
commit 33f2ea823a4abbb1584ad0aa779dcad8d36be27e
270 changes: 270 additions & 0 deletions NN_closure_global_nof_BBLRifirstzone510.jl
Original file line number Diff line number Diff line change
@@ -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]