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,160 @@
+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
+
+fig = Figure(resolution = (1600, 1200))
+
+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))
+
+
+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))
+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)
+
+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}
+}

paper/paper.md

@@ -55,28 +55,28 @@ Most traditional BGC models also require CPU-based supercomputers to run quickly
 One area where novel ideas must be explored with BGC codes is assessing ocean carbon dioxide removal (OCDR) strategies. Quantifying the effectiveness and identifying the impacts of OCDR is challenging due to the aforementioned complexity of the ocean BGC system.
 Moreover, field trials of OCDR interventions are generally small-scale and targeted, while the intervention required to have a climate-scale impact is regional or global.
 This necessitates adaptable, easy-to-use, and verifiable BGC modelling tools which can be used to assess OCDR strategies at the fast pace with which they are being developed [@NationalAcademies2022].
-We have built ``OceanBioME.jl`` to meet these challenges by creating a tool that provides a modular interface to the different components, within the ocean modelling framework provided by ``Oceananigans.jl``.
+We have built ``OceanBioME.jl`` to meet these challenges by creating a tool that provides a modular interface to the different components, within the ocean modelling framework provided by ``Oceananigans.jl`` [@Oceananigans].
 Previously implementing biogeochemical models in ``Oceananigans.jl`` required the user to add forcing terms and boundary conditions to generic tracers [e.g. @SimoesSousa2022].
 ``OceanBioME.jl`` provides a suite of biogeochemical models ranging from simple idealized to full-complexity models and models for boundary fluxes (e.g. air-sea gas exchange).
-``Oceananignas.jl`` and ``OceanBioME.jl`` are built from the ground-up to exploit the power of graphical processor units (GPUs), while also retaining the ability to run on CPUs.
+Facilitated by ``KernelAbstractions.jl`` and ``CUDA.jl`` [@KA; @CUDA], ``Oceananignas.jl`` and ``OceanBioME.jl`` are built from the ground-up to exploit the power of graphical processor units (GPUs) while also retaining the ability to run on CPUs.
 The flexibility of the ``Oceananigans.jl`` framework allows ``OceanBioME.jl`` to be applied across a wide range of scales and use cases, including small-scale large-eddy simulations and regional and global models.
 
 # Summary
 
 ``OceanBioME.jl`` is a flexible modelling environment written in Julia [@julia] for simulating the coupled interactions between ocean biogeochemistry, carbonate chemistry, and physics.
-``OceanBioME.jl`` can be used as a stand-alone box model, or integrated into ``Oceananigans.jl`` [@Oceananigans] for coupled physical-biogeochemical simulations in one, two, or three dimensions.
+``OceanBioME.jl`` can be used as a stand-alone box model, or integrated into ``Oceananigans.jl`` for coupled physical-biogeochemical simulations in one, two, or three dimensions.
 As a result, ``OceanBioME.jl`` and ``Oceananigans.jl`` can be used to simulate the biogeochemical response across an enormous range of scales: from surface boundary layer turbulence at the sub-meter scale to eddying global ocean simulations at the planetary scale, and on computational systems ranging from laptops to supercomputers.
 An example of a problem involving small-scale flow features is showcased in \autoref{eady}, which shows a simulation of a sub-mesoscale eddy in a 1km x 1km horizontal domain with an intermediate complexity biogeochemical model and a kelp growth model solved along the trajectories of drifting buoys (a list of of examples shown in this paper and links to source code are given at the end of the paper).
 ``OceanBioME.jl`` leverages Julia's multiple dispatch and effective inline capabilities to fuse its computations directly into existing ``Oceananigans.jl`` kernels, thus maintaining ``Oceananigans.jl``'s bespoke performance, memory- and cost-efficiency on GPUs in ``OceanBioME.jl``-augmented simulations.
 
-![In this simulation of baroclinic instability in the Eady problem, a background buoyancy gradient and corresponding thermal wind generates a sub-mesoscale eddy, roughly following the setup of @taylor:2016.
+![In this simulation of baroclinic instability in the Eady problem, a background buoyancy gradient and corresponding thermal wind generates a sub-mesoscale eddy, roughly following the setup of @taylor:2016. A submesoscale front develops on the periphery of the eddy where intense three-dimensional turbulence develops with inherently non-hydrostatic dynamics.
 To this physical setup, we added a medium complexity (9 tracers) biogeochemical model, some components of which are shown above.
 On top of this, we added particles modelling the growth of sugar kelp, which are free-floating and advected by the flow, and carbon dioxide exchange from the air.
-Thanks to Julia's speed and efficiency the above model (1 km × 1 km × 100 m with 64 × 64 × 16 grid points) took about 30 minutes of computing time to simulate 10 days of evolution on an Nvidia P100 GPU. 
-Panel (a) shows the domain with the colour representing the concentration of various biogeochemical tracer fields: inorganic carbon, organic carbon (dissolved and particulate), phytoplankton, and nutrients. 
-The increase in organic carbon concentration in the centre of the eddy can be seen, as well as carbon being subducted (most visible in the xz face in the organic carbon).
+Thanks to Julia's speed and efficiency the above model (1 km × 1 km × 100 m with 512 × 512 × 64 grid points) took about 2 hours of computing time to simulate 10 days of evolution on an Nvidia A100 GPU. 
+The figure shows the domain with the colour representing the concentration of various biogeochemical tracer fields: nutrients (nitrate and ammonia), phytoplankton, organic carbon (dissolved and particulate), and inorganic carbon. 
+Darker colours represent higher values.
+The increase in organic carbon concentration in the centre of the eddy can be seen, as well as carbon being subducted (most visible in the xz face in the phytoplankton).
 Points on the surface represent the kelp particle positions, with the colour representing the range of frond size.
-Panel (b) shows the carbon stored in each kelp frond, highlighting the variability depending on the nutrition availability in each particle's location history.
 Figure made with ``Makie.jl`` [@makie]. \label{eady}](eady_example.png)
 
 ``OceanBioME.jl`` is built with a highly modular design that allows user control and customization.
@@ -140,7 +140,7 @@ Strong-Wright et al. (in prep.) uses the coupling of both the biogeochemistry an
 
 | Example| OceanBioME features utilised        | Code location |
 |--------------|--------------|---------------|
-| Sub-mesoscale eddy (\autoref{eady}) | LOBSTER biogeochemical model$^1$ with carbonate model active, CO$_2$ exchange with the air$^2$, Light attenuation$^3$, mass conserving negativity protection$^4$ | `examples/eady.jl` with resolution increased to 64x64x16|
+| Sub-mesoscale eddy (\autoref{eady}) | LOBSTER biogeochemical model$^1$ with carbonate model active, CO$_2$ exchange with the air$^2$, Light attenuation$^3$, mass conserving negativity protection$^4$ |  `paper/figures/eady.jl`, similar to `examples/eady.jl` |
 | Near-global (\autoref{global}) proof of concept | Light attenuation$^3$, NPZD model$^5$| [https://github.com/OceanBioME/GlobalOceanBioME.jl/releases/tag/v0.0.1](https://github.com/OceanBioME/GlobalOceanBioME.jl/releases/tag/v0.0.1) |
 | Idealised 1D model with kelp individuals (\autoref{column}) |  LOBSTER biogeochemical model$^1$ with carbonate model and variable Redfield ratio for organic components active, CO$_2$ exchange with the air$^2$, light attenuation$^3$, mass conserving negativity protection$^4$, and Saccharina Latissima (sugar kelp) model$^6$ | `paper/figures/column.jl`, similar to `examples/column.jl` and `examples/kelp.jl` |