Improve Eady figure

paper/figures/eady.jl

@@ -0,0 +1,133 @@
+using OceanBioME, Printf, Oceananigans, Oceananigans.Units
+
+using Oceananigans.Architectures: arch_array
+
+using Oceananigans.Solvers: FFTBasedPoissonSolver 
+
+using Random
+
+Random.seed!(43)
+
+h(k) = (k - 1) / Nz
+
+ζ₀(k) = 1 + (h(k) - 1) / refinement
+
+Σ(k) = (1 - exp(-stretching * h(k))) / (1 - exp(-stretching))
+
+z_faces(k) = Lz * (ζ₀(k) * Σ(k) - 1)
+
+Nx, Ny, Nz = 512, 512, 64
+Lx, Ly, Lz = 1kilometer, 1kilometer, 140.0
+
+refinement = 1.8
+stretching = 3
+
+arch = GPU()
+
+grid = RectilinearGrid(arch; size = (Nx, Ny, Nz), x = (0, Lx), y = (0, Ly), z = z_faces)
+
+coriolis = FPlane(f = 1e-4) # [s⁻¹]
+
+background_state_parameters = ( M = 1e-4,       # s⁻¹, geostrophic shear
+                                f = coriolis.f, # s⁻¹, Coriolis parameter
+                                N = 1e-4,       # s⁻¹, buoyancy frequency
+                                H = grid.Lz )
+
+# We assume a background buoyancy ``B`` with a constant stratification and also a constant lateral
+# gradient (in the zonal direction). The background velocity components ``U`` and ``V`` are prescribed
+# so that the thermal wind relationship is satisfied, that is, ``f \partial_z U = - \partial_y B`` and
+# ``f \partial_z V = \partial_x B``.
+B(x, y, z, t, p) = p.M ^ 2 * x + p.N ^ 2 * (z + p.H)
+V(x, y, z, t, p) = p.M ^ 2 / p.f * (z + p.H)
+
+V_field = BackgroundField(V, parameters = background_state_parameters)
+B_field = BackgroundField(B, parameters = background_state_parameters)
+
+νᵥ = κᵥ = 1e-4 # [m² s⁻¹]
+closure = ScalarDiffusivity(ν = νᵥ, κ = κᵥ)
+
+@info "Setting up kelp particles"
+
+n = 36 # must be a square number
+
+x = arch_array(arch, [repeat([Lx / (sqrt(n) + 1) * n for n in 1:Int(sqrt(n))], 1, Int(sqrt(n)))...])
+y = arch_array(arch, [repeat([Ly / (sqrt(n) + 1) * n for n in 1:Int(sqrt(n))], 1, Int(sqrt(n)))'...])
+z = arch_array(arch, zeros(Float64, n))
+
+particles = SLatissima(; architecture = arch,
+                         x, y, z, 
+                         A = arch_array(arch, 5.0 .* ones(n)), N = arch_array(arch, 0.01 .* ones(n)), C = arch_array(arch, 0.18 .* ones(n)), 
+                         latitude = 43.3,
+                         scalefactor = 500.0,
+                         pescribed_temperature = (args...) -> 12.0)
+
+biogeochemistry = LOBSTER(; grid,
+                            carbonates = true,
+                            open_bottom = true,
+                            particles,
+                            scale_negatives = true,
+			                surface_phytosynthetically_active_radiation = (x, y, t) -> 100)
+
+DIC_bcs = FieldBoundaryConditions(top = GasExchange(; gas = :CO₂, temperature = (args...) -> 12, salinity = (args...) -> 35))
+
+# Model instantiation
+model = NonhydrostaticModel(; grid,
+                              biogeochemistry,
+                              boundary_conditions = (DIC = DIC_bcs, ),
+                              advection = WENO(grid),
+                              timestepper = :RungeKutta3,
+                              coriolis,
+                              tracers = :b,
+                              buoyancy = BuoyancyTracer(),
+                              background_fields = (b = B_field, v = V_field),
+                              closure)
+
+model.clock.time = 50days
+
+Ξ(z) = randn() * z / grid.Lz * (z / grid.Lz + 1)
+
+Ũ = 1e-3
+uᵢ(x, y, z) = Ũ * Ξ(z)
+vᵢ(x, y, z) = Ũ * Ξ(z)
+
+set!(model, u=uᵢ, v=vᵢ, P = 0.03, Z = 0.03, NO₃ = 4.0, NH₄ = 0.05, DIC = 2200.0, Alk = 2409.0)
+
+Δx = minimum_xspacing(grid, Center(), Center(), Center())
+Δy = minimum_yspacing(grid, Center(), Center(), Center())
+Δz = minimum_zspacing(grid, Center(), Center(), Center())
+
+Δt₀ = 0.6 * min(Δx, Δy, Δz) / V(0, 0, 0, 0, background_state_parameters)
+
+simulation = Simulation(model, Δt = Δt₀, stop_time = 60days)
+
+# Adapt the time step while keeping the CFL number fixed.
+wizard = TimeStepWizard(cfl = 0.6, diffusive_cfl = 0.6, max_Δt = 30minutes, min_change = 0.1, max_change = 2)
+simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(1))
+
+# Create a progress message.
+progress(sim) = @printf("i: % 6d, sim time: % 10s, wall time: % 10s, Δt: % 10s\n",
+                        sim.model.clock.iteration,
+                        prettytime(sim.model.clock.time),
+                        prettytime(sim.run_wall_time),
+                        prettytime(sim.Δt))
+
+simulation.callbacks[:progress] = Callback(progress, IterationInterval(10))
+
+u, v, w = model.velocities
+
+# Periodically save the velocities and vorticity to a file.
+simulation.output_writers[:fields] = JLD2OutputWriter(model, merge(model.tracers, (; u, v, w));
+                                                      schedule = TimeInterval(6hours),
+                                                      filename = "eady.jld2",
+                                                      overwrite_existing = true)
+
+simulation.output_writers[:particles] = JLD2OutputWriter(model, (; particles);
+                                                         schedule = TimeInterval(6hours),
+                                                         filename = "eady_particles.jld2",
+                                                         overwrite_existing = true)
+
+run!(simulation)
+
+simulation.output_writers[:checkpointer] = Checkpointer(model, schedule=IterationInterval(1), prefix="model_checkpoint")
+
+run!(simulation)

paper/figures/eady_plot.jl

@@ -0,0 +1,167 @@
+using Oceananigans, GLMakie, JLD2
+
+Nx = 512
+
+# load all the fields
+
+field_path = "eady.jld2"
+
+field_all = FieldTimeSeries(field_path, "P")
+
+times = field_all.times
+grid = field_all.grid
+
+P = [Array(interior(field_all[n], 1:Int(Nx / 4), :, :)) for n in eachindex(times)]
+
+field_all = FieldTimeSeries(field_path, "NO₃")
+
+N = [Array(interior(field_all[n], Int(Nx / 4 + 1):Int(2 * Nx / 4), :, :)) for n in eachindex(times)]
+
+field_all = FieldTimeSeries(field_path, "NH₄")
+
+for n in eachindex(times)
+    N[n] .+= Array(interior(field_all[n], Int(Nx / 4 + 1):Int(2 * Nx / 4), :, :))
+end
+
+field_all = FieldTimeSeries(field_path, "sPOM")
+
+OC = [Array(interior(field_all[n], Int(2 * Nx / 4 + 1):Int(3 * Nx / 4), :, :)) for n in eachindex(times)]
+
+field_all = FieldTimeSeries(field_path, "bPOM")
+
+for n in eachindex(times)
+    OC[n] .+= Array(interior(field_all[n], Int(Nx / 4 + 1):Int(2 * Nx / 4), :, :))
+end
+
+field_all = FieldTimeSeries(field_path, "DOM")
+
+for n in eachindex(times)
+    OC[n] .+= Array(interior(field_all[n], Int(Nx / 4 + 1):Int(2 * Nx / 4), :, :))
+end
+
+field_all = FieldTimeSeries(field_path, "DIC")
+
+DIC = [Array(interior(field_all[n], Int(3 * Nx / 4 + 1):Int(4 * Nx / 4), :, :)) for n in eachindex(times)]
+
+# load particles
+
+file = jldopen("eady_particles.jld2")
+
+iterations = keys(file["timeseries/t"])
+
+x, y, z, A, N_kelp, C = ntuple(n -> ones(36, length(iterations)) .* NaN, 6)
+
+for (idx, it) in enumerate(iterations)
+    particles = file["timeseries/particles/$it"]
+
+    x[:, idx] = particles.x
+    y[:, idx] = particles.y
+    z[:, idx] = particles.z
+    A[:, idx] = particles.A
+    N_kelp[:, idx] = particles.N
+    C[:, idx] = particles.C
+end
+
+close(file)
+
+##### 
+##### Animation
+#####
+
+# setup the observables
+
+n = Observable(1)
+
+N_plt   =   @lift N[$n]
+P_plt   =   @lift P[$n]
+OC_plt  =  @lift OC[$n]
+DIC_plt = @lift DIC[$n]
+
+x_plt = @lift x[:, $n]
+y_plt = @lift y[:, $n]
+z_plt = @lift z[:, $n]
+A_plt = @lift A[:, $n]
+
+
+fig = Figure(resolution = (1600, 1000))
+
+ax = Axis3(fig[1:4, 1:4], aspect = (1, 1, 0.28), xticks = [0, 1000], yticks = [0, 1000], zticks = [-140, 0],
+                                             xlabel = "x (m)", ylabel = "y (m)", zlabel = "z (m)",
+                                             xgridvisible = false, ygridvisible = false, zgridvisible = false,
+                                             xspinesvisible = false, yspinesvisible = false, zspinesvisible = false,
+                                             protrusions = (50, 30, 30, 30))
+
+vm1 = contour!(ax, xc[1:Int(Nx / 4)], yc, zc, N_plt, levels = 50, colormap = Reverse(:bamako))
+vm2 = contour!(ax, xc[Int(Nx / 4 + 1):Int(2 * Nx / 4)], yc, zc, P_plt, levels = 50, colormap = Reverse(:batlow))
+vm3 = contour!(ax, xc[Int(2 * Nx / 4 + 1):Int(3 * Nx / 4)], yc, zc, OC_plt, levels = 50, colormap = :lajolla)
+vm4 = contour!(ax, xc[Int(3 * Nx / 4 + 1):Int(4 * Nx / 4)], yc, zc, DIC_plt, levels = 50, colormap = Reverse(:devon))
+
+sc = scatter!(ax, x_plt, y_plt, z_plt, color = A_plt, colormap=:grayC)
+
+txt = text!(ax,
+            [Point3f(xc[Int(1 + (i - 1) * Nx / 4)], yc[1], zc[1]) for i in 1:4],
+            text = ["Nutrients", "Phytoplankton", "Organic carbon", "Inorganic carbon"],
+            rotation = [0 for i in 1:4],
+            align = (:left, :top),
+            fontsize = 35,
+            markerspace = :data)
+
+supertitle = Label(fig[0, 1:4], "t = 0.0, n = 1")
+
+record(fig, "all.mp4", eachindex(times)[2:end], framerate = 10) do i
+    n[] = i
+    supertitle.text = "t = $(prettytime(times[i])), n = $i"
+    println("$i of $(length(times))")
+end
+
+n[] = 39
+fig
+
+#####
+##### Static Figure
+#####
+n = 37
+
+lims[1] = (min(minimum(N_plt[:, :, end]), minimum(N_plt[1, :, :]), minimum(N_plt[:, 1, :])), max(maximum(N_plt[:, :, end]), maximum(N_plt[1, :, :]), maximum(N_plt[:, 1, :])))
+
+
+fig = Figure(resolution = (1600, 1000))
+
+ax = Axis3(fig[1:4, 1:4], aspect = (1, 1, 0.28), xticks = [0, 1000], yticks = [0, 1000], zticks = [-140, 0],
+                                             xlabel = "x (m)", ylabel = "y (m)", zlabel = "z (m)",
+                                             xgridvisible = false, ygridvisible = false, zgridvisible = false,
+                                             xspinesvisible = false, yspinesvisible = false, zspinesvisible = false,
+                                             protrusions = (50, 5, 5, 5))
+
+
+N_plt   =   N[n]
+P_plt   =   P[n]
+OC_plt  =  OC[n]
+DIC_plt = DIC[n]
+
+x_plt = x[:, n]
+y_plt = y[:, n]
+z_plt = z[:, n]
+A_plt = A[:, n]
+
+vm1 = contour!(ax, xc[1:Int(Nx / 4)], yc, zc, N_plt, levels = 50, colormap = Reverse(:bamako), colorrange = lims[1])
+vm2 = contour!(ax, xc[Int(Nx / 4 + 1):Int(2 * Nx / 4)], yc, zc, P_plt, levels = 50, colormap = Reverse(:batlow))
+vm3 = contour!(ax, xc[Int(2 * Nx / 4 + 1):Int(3 * Nx / 4)], yc, zc, OC_plt, levels = 50, colormap = :lajolla)
+vm4 = contour!(ax, xc[Int(3 * Nx / 4 + 1):Int(4 * Nx / 4)], yc, zc, DIC_plt, levels = 50, colormap = Reverse(:devon))
+
+sc = scatter!(ax, x_plt, y_plt, z_plt, color = A_plt, colormap=:grayC)
+
+txt = text!(ax,
+            [Point3f(xc[Int(1 + (i - 1) * Nx / 4)], yc[1], zc[1]) for i in 1:4],
+            text = ["| Nutrients", "| Phytoplankton", "| Organic carbon", "| Inorganic carbon"],
+            rotation = [0 for i in 1:4],
+            align = (:left, :top),
+            fontsize = 33,
+            markerspace = :data)
+
+Colorbar(fig[5, 1], limits = (minimum(N_plt), maximum(N_plt)), colormap = Reverse(:bamako), vertical = false, label = "Nutrients (mmol N / m³)")
+Colorbar(fig[5, 2], limits = (minimum(P_plt), maximum(P_plt)), colormap = Reverse(:batlow), vertical = false, label = "Phytoplankton (mmol N / m³)")
+Colorbar(fig[5, 3], limits = (minimum(OC_plt), maximum(OC_plt)), colormap = :lajolla, vertical = false, label = "Organic Carbon (mmol C / m³)")
+Colorbar(fig[5, 4], limits = (minimum(DIC_plt), maximum(DIC_plt)), colormap = Reverse(:devon), vertical = false, label = "Inorganic Carbon (mmol C / m³)")
+
+save("eady.png", fig)

paper/paper.bib

@@ -186,3 +186,21 @@ @article{SimoesSousa2022
   publisher={Frontiers},
   doi = {10.3389/fmars.2022.825027}
 }
+@ARTICLE{CUDA,
+   author={Besard, Tim and Foket, Christophe and De Sutter, Bjorn},
+   journal={IEEE Transactions on Parallel and Distributed Systems}, 
+   title={Effective Extensible Programming: Unleashing Julia on GPUs}, 
+   year={2019},
+   volume={30},
+   number={4},
+   pages={827-841},
+   doi={10.1109/TPDS.2018.2872064}
+}
+@article{KA,
+   title={JuliaGPU/KernelAbstractions.jl: v0.9.8}, 
+   DOI={10.5281/zenodo.8176369}, 
+   publisher={Zenodo}, 
+   author={Valentin Churavy and Dilum Aluthge and Anton Smirnov and Julian Samaroo and James Schloss and Lucas C Wilcox and Simon Byrne and Tim Besard and Maciej Waruszewski and Ali Ramadhan and et al.}, 
+   year={2023}, 
+   month={Jul}
+}