A NN parameterization to test in the neverworld framework

.gitignore

@@ -6,4 +6,5 @@
 *.svg
 *.out
 *.sh
-*core*
+*core*
+Manifest.toml

Project.toml

@@ -9,25 +9,27 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterTools = "35a29f4d-8980-5a13-9543-d66fff28ecb8"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
 Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 
 [compat]
 Adapt = "^3"
+CUDA = "4, 5"
 Documenter = "1"
 DocumenterTools = "0.1"
 FFTW = "^1"
 JLD2 = "0.4"
+KernelAbstractions = "^0.9"
 MPI = "0.20"
 MPIPreferences = "0.1"
-CUDA = "4, 5"
-Printf = "1.9"
-KernelAbstractions = "^0.9"
 Oceananigans = "^0.90"
+Printf = "1.9"
 julia = "1.9"
 
 [extras]

examples/run.jl

@@ -2,7 +2,7 @@ using WenoNeverworld
 using Oceananigans
 using Oceananigans.Units
 using Oceananigans.Grids: φnodes, λnodes, znodes, on_architecture
-using CairoMakie 
+using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities: CATKEVerticalDiffusivity, MixingLength
 
 output_dir    = joinpath(@__DIR__, "./")
 @show output_prefix = output_dir * "/neverworld_quarter_resolution"
@@ -38,50 +38,20 @@ buoyancy_relaxation = BuoyancyRelaxationBoundaryCondition(ΔB = 0.06, λ = 7days
 # @inline seasonal_cosine_scaling(y, t) = cos(π * y / 70) * sin(2π * t / 1year)
 # buoyancy_relaxation = BuoyancyRelaxationBoundaryCondition(seasonal_cosine_scaling; ΔB = 0.06, λ = 7days)    
 
+mixing_length = MixingLength(; Cᵇ = 0.01)
+vertical_diffusivity = CATKEVerticalDiffusivity(; mixing_length)
+
 # Construct the neverworld simulation
 simulation = weno_neverworld_simulation(grid; Δt, stop_time,
                                               wind_stress,
                                               buoyancy_relaxation,
                                               interp_init,
-                                              init_file)
+                                              vertical_diffusivity,
+                                              init_file,
+                                              tracers = (:b, :e))
 
 model = simulation.model
 
-# Let's visualize our boundary conditions!
-cpu_grid = on_architecture(CPU(), grid)
-φ = φnodes(cpu_grid.underlying_grid, Center(), Center(), Center())
-τ_bcs = Array(model.velocities.u.boundary_conditions.top.condition.func.stress)
-b_bcs = zeros(length(τ_bcs))
-b = Array(interior(model.tracers.b))
-for j in 1:grid.Ny
-    b_bcs[j] = buoyancy_relaxation(grid.Nx÷2, j, grid, model.clock, (; b))
-end
-
-fig = Figure()
-ax  = Axis(fig[1, 1], title = L"\text{Wind stress profile}")
-lines!(ax, φ, - τ_bcs .* 1000, linewidth = 5) # (we re-convert the wind stress form kg/m² to Nm)
-
-ax  = Axis(fig[1, 2], title = L"\text{Restoring buoyancy flux}")
-lines!(ax, φ, - b_bcs, linewidth = 5)
-
-CairoMakie.save("boundary_conditions.png", fig)
-
-# Let's plot the initial conditions to make sure they are reasonable
-λ = λnodes(cpu_grid.underlying_grid, Center(), Center(), Center())
-z = znodes(cpu_grid.underlying_grid, Center(), Center(), Center())
-
-fig = Figure(resolution = (800, 2000))
-ax  = Axis(fig[1:4, 1], title = "Surface initial buoyancy")
-hm  = heatmap!(ax, λ, φ, b[:, :, grid.Nz], colormap = :thermometer)
-cb  = Colorbar(fig[1:4, 2], hm) 
-
-ax  = Axis(fig[5, 1], title = "Mid longitude buoyancy")
-hm  = heatmap!(ax, φ, z, b[grid.Nx ÷ 2, :, :], colormap = :thermometer)
-ct  = contour!(ax, φ, z, b[grid.Nx ÷ 2, :, :], levels = range(0, 0.06, length = 10), color = :black)
-cb  = Colorbar(fig[5, 2], hm) 
-
-CairoMakie.save("initial_conditions.png", fig)
-
 # Add outputs (check other outputs to attach in `src/neverworld_outputs.jl`)
 checkpoint_outputs!(simulation, output_prefix)
 

examples/run_nn.jl

@@ -0,0 +1,64 @@
+using WenoNeverworld
+using WenoNeverworld.Parameterizations
+using WenoNeverworld.Auxiliaries
+using Oceananigans
+using Oceananigans.Units
+
+output_dir    = joinpath(@__DIR__, "./")
+@show output_prefix = output_dir * "/neverworld_quarter_resolution"
+
+arch = GPU()
+
+# The resolution in degrees
+degree_resolution = 1/4
+
+grid = NeverworldGrid(degree_resolution; arch)
+
+# Do we need to interpolate? (interp_init) If `true` from which file?
+interp_init = false # If interpolating from a different grid: `interp_init = true`
+init_file   = nothing # To restart from a file: `init_file = /path/to/restart`
+
+# Simulation parameters
+Δt        = 1minutes
+stop_time = 200years
+
+# Latitudinal wind stress acting on the zonal velocity
+# a piecewise-cubic profile interpolated between
+# x = φs (latitude) and y = τs (stress)
+φs = (-70.0, -45.0, -15.0,  0.0,  15.0, 45.0, 70.0)
+τs = (  0.0,   0.2,  -0.1, -0.02, -0.1,  0.1,  0.0)
+wind_stress = WindStressBoundaryCondition(; φs, τs)
+
+# Buoyancy relaxation profile:
+# a parabolic profile between 0, at the poles, and ΔB = 0.06 at the equator
+# the restoring time is λ = 7days
+buoyancy_relaxation = BuoyancyRelaxationBoundaryCondition(ΔB = 0.06, λ = 7days)
+
+# Here we use test `NNbackscatteringClosure` closure
+horizontal_closure = NNbackscatteringClosure(eltype(grid); 
+                                             architecture = arch, 
+                                             weight_path = "model_weights.jld2")
+
+# Construct the neverworld simulation
+simulation = weno_neverworld_simulation(grid; Δt, stop_time,
+                                              wind_stress,
+                                              buoyancy_relaxation,
+                                              interp_init,
+                                              horizontal_closure,
+                                              init_file)
+
+# Add outputs (check other outputs to attach in `src/neverworld_outputs.jl`)
+checkpoint_outputs!(simulation, output_prefix)
+
+# Initialize with a small time step and increase it after the 
+# inital spin up has completed
+increase_simulation_Δt!(simulation; cutoff_time = 20days,  new_Δt =  2minutes)
+increase_simulation_Δt!(simulation; cutoff_time = 50days,  new_Δt =  4minutes)
+increase_simulation_Δt!(simulation; cutoff_time = 100days, new_Δt =  6minutes)
+increase_simulation_Δt!(simulation; cutoff_time = 150days, new_Δt =  8minutes)
+increase_simulation_Δt!(simulation; cutoff_time = 200days, new_Δt = 10minutes)
+
+# initializing the time for wall_time calculation
+@info "Running with Δt = $(prettytime(simulation.Δt))"
+run_simulation!(simulation; interp_init, init_file)
+

src/Auxiliaries/Auxiliaries.jl

@@ -9,7 +9,7 @@ using Oceananigans
 using Oceananigans.Utils
 using Oceananigans.Fields: interpolate
 using Oceananigans.Grids: λnode, φnode, halo_size, on_architecture
-using Oceananigans.Architectures: arch_array, architecture
+using Oceananigans.Architectures: on_architecture, architecture
 using Oceananigans.Utils: instantiate
 using Oceananigans.BoundaryConditions
 using JLD2

src/Diagnostics/spurious_mixing.jl

@@ -1,6 +1,6 @@
 using Oceananigans.AbstractOperations: GridMetricOperation
 using Oceananigans.Grids: architecture, znode
-using Oceananigans.Architectures: device, arch_array
+using Oceananigans.Architectures: device, on_architecture
 
 function calculate_z★_diagnostics(b::FieldTimeSeries)
 

src/NeverworldBoundaries/NeverworldBoundaries.jl

@@ -16,7 +16,7 @@ using Oceananigans.Units
 using Oceananigans.Operators
 using Oceananigans.BoundaryConditions
 using Oceananigans.Fields: interpolate
-using Oceananigans.Architectures: architecture, arch_array
+using Oceananigans.Architectures: architecture, on_architecture
 using Oceananigans.Grids: λnode, φnode, halo_size, on_architecture
 using Oceananigans.Utils: instantiate
 using Oceananigans.ImmersedBoundaries: ImmersedBoundaryCondition

src/NeverworldBoundaries/wind_stress_bc.jl

@@ -36,5 +36,5 @@ Adapt.adapt_structure(to, ws::WindStressBoundaryCondition) = WindStressBoundaryC
         stress[j] = cubic_interpolate(φ, x₁ = φ₁, x₂ = φ₂, y₁ = τ₁, y₂ = τ₂) / 1000.0
     end
 
-    return WindStressBoundaryCondition(bc.φs, bc.τs, arch_array(arch, - stress))
+    return WindStressBoundaryCondition(bc.φs, bc.τs, on_architecture(arch, - stress))
 end

src/NeverworldGrids/NeverworldGrids.jl

@@ -12,7 +12,7 @@ using Oceananigans.Operators
 using Oceananigans.BoundaryConditions
 using Oceananigans.Units
 using Oceananigans.Grids
-using Oceananigans.Architectures: arch_array, architecture
+using Oceananigans.Architectures: on_architecture, architecture
 using Oceananigans.Grids: on_architecture
 using Oceananigans.ImmersedBoundaries
 

src/Parameterizations/Parameterizations.jl

@@ -1,6 +1,9 @@
 module Parameterizations
 
-export QGLeith, EnergyBackScattering
+export QGLeith, 
+       EnergyBackScattering, 
+       NNbackscatteringClosure,
+       XinKaiVerticalDiffusivity
 
 using Oceananigans
 using KernelAbstractions: @index, @kernel
@@ -23,6 +26,7 @@ using Oceananigans.TurbulenceClosures:
         VerticallyImplicitTimeDiscretization,
         ExplicitTimeDiscretization,
         FluxTapering,
+        getclosure,
         isopycnal_rotation_tensor_xz_ccf,
         isopycnal_rotation_tensor_yz_ccf,
         isopycnal_rotation_tensor_zz_ccf
@@ -34,19 +38,31 @@ import Oceananigans.TurbulenceClosures:
         diffusivity,
         diffusive_flux_x,
         diffusive_flux_y, 
-        diffusive_flux_z
+        diffusive_flux_z,
+        top_buoyancy_flux
+
+import Oceananigans.TurbulenceClosures: 
+        ∂ⱼ_τ₁ⱼ, 
+        ∂ⱼ_τ₂ⱼ, 
+        ∂ⱼ_τ₃ⱼ,
+        ∇_dot_qᶜ
 
 using Oceananigans.Utils: launch!
 using Oceananigans.Coriolis: fᶠᶠᵃ
 using Oceananigans.Operators
+using Oceananigans.BoundaryConditions: fill_halo_regions!
 using Oceananigans.BuoyancyModels: ∂x_b, ∂y_b, ∂z_b 
 
 using Oceananigans.Operators: ℑxyzᶜᶜᶠ, ℑyzᵃᶜᶠ, ℑxzᶜᵃᶠ, Δxᶜᶜᶜ, Δyᶜᶜᶜ
 
+using Adapt
+
 "Return the filter width for an Horizontal closure on a general grid."
 @inline Δ²ᶜᶜᶜ(i, j, k, grid) =  2 * (1 / (1 / Δxᶜᶜᶜ(i, j, k, grid)^2 + 1 / Δyᶜᶜᶜ(i, j, k, grid)^2))
 
 include("quasi_geostrophic_leith.jl")
 include("energy_backscattering.jl")
+include("xin_kai_vertical_diffusivity.jl")
+include("guillaumin_zanna_parameterization.jl")
 
 end