Merge branch 'main' into ss/typo-in-generate-fluxes

.buildkite/examples_build.yml

@@ -4,8 +4,6 @@ agents:
   modules: climacommon/2024_05_27
   slurm_time: 48:00:00
 
-timeout_in_minutes: 1440
-
 env:
   JULIA_LOAD_PATH: "${JULIA_LOAD_PATH}:${BUILDKITE_BUILD_CHECKOUT_PATH}/.buildkite"
   OPENBLAS_NUM_THREADS: 1
@@ -45,7 +43,7 @@ steps:
     env:
       JULIA_DEBUG: "Documenter"
       # This environment variable is needed to avoid SSL verification errors when Downloads.jl 
-      # tries to download the bathymetry data. It should not be required so we need to fix our certificates.
+      # tries to download the bathymetry data. It should not be required so we need to fix our certificates
       # and remove this environment variable. ref: https://github.com/JuliaLang/Downloads.jl/issues/97
       JULIA_SSL_NO_VERIFY: "**" 
 

docs/Project.toml

@@ -11,4 +11,5 @@ OrthogonalSphericalShellGrids = "c2be9673-fb75-4747-82dc-aa2bb9f4aed0"
 CairoMakie = "0.10.12, 0.11, 0.12"
 DataDeps = "0.7"
 Documenter = "1"
-Oceananigans = "0.94"
+Oceananigans = "0.94.3"
+OrthogonalSphericalShellGrids = "0.1.8"

examples/near_global_ocean_simulation.jl

@@ -4,7 +4,7 @@
 # ClimaOcean.jl. The simulation covers latitudes from 75°S to 75°N with a horizontal
 # resolution of 1/4 degree and 40 vertical levels.
 #
-# The simulation's results are visualized using the CairoMakie.jl package.
+# The simulation's results are visualized with the CairoMakie.jl package.
 #
 # ## Initial setup with package imports
 #
@@ -13,14 +13,13 @@
 # These packages provide the foundational tools for setting up the simulation environment,
 # including grid setup, physical processes modeling, and data visualization.
 
-using Printf
+using ClimaOcean
 using Oceananigans
 using Oceananigans.Units
-using ClimaOcean
 using CairoMakie
-
 using CFTime
 using Dates
+using Printf
 
 # ### Grid configuration 
 #
@@ -58,7 +57,7 @@ bottom_height = regrid_bathymetry(grid;
                                   interpolation_passes = 5,
                                   major_basins = 3)
 
-grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height))
+grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true)
 
 # Let's see what the bathymetry looks like:
 

examples/single_column_os_papa_simulation.jl

@@ -1,4 +1,4 @@
-# # Single column ocean simulation forced by JRA55 re-analysis
+# # Single-column ocean simulation forced by JRA55 re-analysis
 #
 # In this example, we simulate the evolution of an ocean water column 
 # forced by an atmosphere derived from the JRA55 re-analysis.
@@ -15,18 +15,15 @@
 # ```
 
 using ClimaOcean
-
 using Oceananigans
 using Oceananigans.Units
 using Oceananigans.BuoyancyModels: buoyancy_frequency
 using Oceananigans.Units: Time
-
-using CairoMakie
 using Printf
 
 # # Construct the grid
 #
-# First, we construct a single column grid with 2 meter spacing
+# First, we construct a single-column grid with 2 meter spacing
 # located at ocean station Papa.
 
 # Ocean station papa location
@@ -42,7 +39,7 @@ grid = RectilinearGrid(size = 200,
 # # An "ocean simulation"
 #
 # Next, we use ClimaOcean's ocean_simulation constructor to build a realistic
-# ocean simulation on the single column grid,
+# ocean simulation on the single-column grid,
 
 ocean = ocean_simulation(grid; Δt=10minutes, coriolis=FPlane(latitude = φ★))
 
@@ -56,7 +53,7 @@ set!(ocean.model, T=ECCOMetadata(:temperature), S=ECCOMetadata(:salinity))
 
 # # A prescribed atmosphere based on JRA55 re-analysis
 #
-# We build a PrescribedAtmosphere at the same location as the single colunm grid
+# We build a PrescribedAtmosphere at the same location as the single-colunm grid
 # which is based on the JRA55 reanalysis.
 
 simulation_days = 31
@@ -79,6 +76,10 @@ Ta = interior(atmosphere.tracers.T, 1, 1, 1, :)
 qa = interior(atmosphere.tracers.q, 1, 1, 1, :)
 t_days = atmosphere.times / days
 
+using CairoMakie
+
+set_theme!(Theme(linewidth=3, fontsize=24))
+
 fig = Figure(size=(800, 600))
 axu = Axis(fig[2, 1], xlabel="Days since Jan 1 1990", ylabel="Atmosphere \n velocity (m s⁻¹)")
 axT = Axis(fig[3, 1], xlabel="Days since Jan 1 1990", ylabel="Atmosphere \n temperature (K)")
@@ -93,7 +94,9 @@ axislegend(axu, framevisible=false, nbanks=2, position=:lb)
 lines!(axT, t_days, Ta)
 lines!(axq, t_days, qa)
 
-display(fig)
+current_figure()
+
+# We continue constructing a simulation.
 
 radiation = Radiation()
 coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation)
@@ -216,29 +219,26 @@ for n = 1:Nt
     Pt[n]   =  Pr[1, 1, 1, Time(t)] + Ps[1, 1, 1, Time(t)]
 end
 
-set_theme!(Theme(linewidth=3))
-
-fig = Figure(size=(2400, 1800))
+fig = Figure(size=(1800, 1800))
 
-axτ = Axis(fig[1, 1:2], xlabel="Days since Oct 1 1992", ylabel="Wind stress (N m⁻²)")
-axu = Axis(fig[2, 1:2], xlabel="Days since Oct 1 1992", ylabel="Velocities (m s⁻¹)")
-axQ = Axis(fig[1, 3:4], xlabel="Days since Oct 1 1992", ylabel="Heat flux (W m⁻²)")
-axT = Axis(fig[2, 3:4], xlabel="Days since Oct 1 1992", ylabel="Surface temperature (ᵒC)")
-axF = Axis(fig[1, 5:6], xlabel="Days since Oct 1 1992", ylabel="Freshwater volume flux (m s⁻¹)")
-axS = Axis(fig[2, 5:6], xlabel="Days since Oct 1 1992", ylabel="Surface salinity (g kg⁻¹)")
+axτ = Axis(fig[1, 1:3], xlabel="Days since Oct 1 1992", ylabel="Wind stress (N m⁻²)")
+axQ = Axis(fig[1, 4:6], xlabel="Days since Oct 1 1992", ylabel="Heat flux (W m⁻²)")
+axu = Axis(fig[2, 1:3], xlabel="Days since Oct 1 1992", ylabel="Velocities (m s⁻¹)")
+axT = Axis(fig[2, 4:6], xlabel="Days since Oct 1 1992", ylabel="Surface temperature (ᵒC)")
+axF = Axis(fig[3, 1:3], xlabel="Days since Oct 1 1992", ylabel="Freshwater volume flux (m s⁻¹)")
+axS = Axis(fig[3, 4:6], xlabel="Days since Oct 1 1992", ylabel="Surface salinity (g kg⁻¹)")
 
-axuz = Axis(fig[3, 1], xlabel="Velocities (m s⁻¹)",                ylabel="z (m)")
-axTz = Axis(fig[3, 2], xlabel="Temperature (ᵒC)",                  ylabel="z (m)")
-axSz = Axis(fig[3, 3], xlabel="Salinity (g kg⁻¹)",                 ylabel="z (m)")
-axNz = Axis(fig[3, 4], xlabel="Buoyancy frequency (s⁻²)",          ylabel="z (m)")
-axκz = Axis(fig[3, 5], xlabel="Eddy diffusivity (m² s⁻¹)",         ylabel="z (m)", xscale=log10)
-axez = Axis(fig[3, 6], xlabel="Turbulent kinetic energy (m² s⁻²)", ylabel="z (m)", xscale=log10)
+axuz = Axis(fig[4:5, 1:2], xlabel="Velocities (m s⁻¹)",                ylabel="z (m)")
+axTz = Axis(fig[4:5, 3:4], xlabel="Temperature (ᵒC)",                  ylabel="z (m)")
+axSz = Axis(fig[4:5, 5:6], xlabel="Salinity (g kg⁻¹)",                 ylabel="z (m)")
+axNz = Axis(fig[6:7, 1:2], xlabel="Buoyancy frequency (s⁻²)",          ylabel="z (m)")
+axκz = Axis(fig[6:7, 3:4], xlabel="Eddy diffusivity (m² s⁻¹)",         ylabel="z (m)", xscale=log10)
+axez = Axis(fig[6:7, 5:6], xlabel="Turbulent kinetic energy (m² s⁻²)", ylabel="z (m)", xscale=log10)
 
-title = @sprintf("Single column simulation at %.2f, %.2f", φ★, λ★)
+title = @sprintf("Single-column simulation at %.2f, %.2f", φ★, λ★)
 Label(fig[0, 1:6], title)
 
-slider = Slider(fig[4, 1:6], range=1:Nt, startvalue=1)
-n = slider.value
+n = Observable(1)
 
 times = (times .- times[1]) ./days
 Nt = length(times)
@@ -302,23 +302,27 @@ scatterlines!(axκz, κn,  zf)
 
 axislegend(axuz)
 
+ulim = max(maximum(abs, u), maximum(abs, v))
+xlims!(axuz, -ulim, ulim)
+
 Tmax = maximum(interior(T))
 Tmin = minimum(interior(T))
 xlims!(axTz, Tmin - 0.1, Tmax + 0.1)
 
 Nmax = maximum(interior(N²))
-Nmin = minimum(interior(N²))
-xlims!(axNz, Nmin / 2, Nmin * 1.1)
+xlims!(axNz, -Nmax/10, Nmax * 1.05)
+
+κmax = maximum(interior(κ))
+xlims!(axκz, 1e-9, κmax * 1.1)
 
 emax = maximum(interior(e))
-xlims!(axez, 8e-7, emax * 1.1)
-xlims!(axκz, 1e-7, 10)
+xlims!(axez, 1e-11, emax * 1.1)
 
 Smax = maximum(interior(S))
 Smin = minimum(interior(S))
 xlims!(axSz, Smin - 0.2, Smax + 0.2)
 
-record(fig, "single_column_profiles.mp4", 1:2:Nt, framerate=24) do nn
+record(fig, "single_column_profiles.mp4", 1:Nt, framerate=24) do nn
     @info "Drawing frame $nn of $Nt..."
     n[] = nn
 end

src/distributed_utils.jl

@@ -6,66 +6,81 @@ using MPI
 #####
 
 # Utilities to make the macro work importing only ClimaOcean and not MPI
-mpi_initialized() = MPI.Initialized()
-mpi_rank()        = MPI.Comm_rank(MPI.COMM_WORLD)
-mpi_size()        = MPI.Comm_size(MPI.COMM_WORLD)
-global_barrier()  = mpi_initialized() ? MPI.Barrier(MPI.COMM_WORLD) : nothing
+mpi_initialized()    = MPI.Initialized()
+mpi_rank(comm)       = MPI.Comm_rank(comm)
+mpi_size(comm)       = MPI.Comm_size(comm)
+global_barrier(comm) = MPI.Barrier(comm) 
 
 """
-    @root exs...
+    @root communicator exs...
 
-Perform `exs` only on rank 0, otherwise know as "root" rank.
-Other ranks will wait for the root rank to finish before continuing
+Perform `exs` only on rank 0 in communicator, otherwise known as the "root" rank.
+Other ranks will wait for the root rank to finish before continuing.
+If `communicator` is not provided, `MPI.COMM_WORLD` is used.
 """
-macro root(exp)
+macro root(communicator, exp)
     command = quote
         if ClimaOcean.mpi_initialized()
-            rank = ClimaOcean.mpi_rank()
+            rank = ClimaOcean.mpi_rank($communicator)
             if rank == 0
                 $exp
             end
-            ClimaOcean.global_barrier()
+            ClimaOcean.global_barrier($communicator)
         else
             $exp
         end
     end
     return esc(command)
 end
 
+macro root(exp)
+    command = quote
+        @root MPI.COMM_WORLD $exp
+    end
+    return esc(command)
+end
+
 """
-    @onrank rank, exs...
+    @onrank communicator rank exs...
 
-Perform `exp` only on rank `rank`
-Other ranks will wait for the root rank to finish before continuing.
-The expression is run anyways if MPI in not initialized
+Perform `exp` only on rank `rank` (0-based index) in `communicator`.
+Other ranks will wait for rank `rank` to finish before continuing.
+The expression is run anyways if MPI in not initialized.
+If `communicator` is not provided, `MPI.COMM_WORLD` is used.
 """
-macro onrank(exp_with_rank)
-    on_rank = exp_with_rank.args[1]
-    exp  = exp_with_rank.args[2]
+macro onrank(communicator, on_rank, exp)
     command = quote
         mpi_initialized = ClimaOcean.mpi_initialized()
-        rank = ClimaOcean.mpi_rank()
         if !mpi_initialized
             $exp
         else
+            rank = ClimaOcean.mpi_rank($communicator)
             if rank == $on_rank
                 $exp
             end
-            ClimaOcean.global_barrier()
+            ClimaOcean.global_barrier($communicator)
         end
     end
 
     return esc(command)
 end
 
+macro onrank(rank, exp)
+    command = quote
+        @onrank MPI.COMM_WORLD $rank $exp
+    end
+    return esc(command)
+end
+
 """ 
-    @distribute for i in iterable
+    @distribute communicator for i in iterable
         ...
     end
 
-Distribute a `for` loop among different ranks
+Distribute a `for` loop among different ranks in `communicator`.
+If `communicator` is not provided, `MPI.COMM_WORLD` is used.
 """
-macro distribute(exp)
+macro distribute(communicator, exp)
     if exp.head != :for
         error("The `@distribute` macro expects a `for` loop")
     end
@@ -74,46 +89,67 @@ macro distribute(exp)
     variable = exp.args[1].args[1]
     forbody  = exp.args[2]
 
+    # Safety net if the iterable variable has the same name as the 
+    # reserved variable names (nprocs, counter, rank)
+    nprocs  = ifelse(variable == :nprocs,  :othernprocs,  :nprocs)
+    counter = ifelse(variable == :counter, :othercounter, :counter)
+    rank    = ifelse(variable == :rank,    :otherrank,    :rank)
+
     new_loop = quote
         mpi_initialized = ClimaOcean.mpi_initialized()
         if !mpi_initialized
             $exp
         else
-            rank   = ClimaOcean.mpi_rank()
-            nprocs = ClimaOcean.mpi_size()
-            for (counter, $variable) in enumerate($iterable)
-                if (counter - 1) % nprocs == rank
+            $rank   = ClimaOcean.mpi_rank($communicator)
+            $nprocs = ClimaOcean.mpi_size($communicator)
+            for ($counter, $variable) in enumerate($iterable)
+                if ($counter - 1) % $nprocs == $rank
                     $forbody
                 end
             end
-            ClimaOcean.global_barrier()
+            ClimaOcean.global_barrier($communicator)
         end
     end
 
     return esc(new_loop)
 end
 
+macro distribute(exp)
+    command = quote
+        @distribute MPI.COMM_WORLD $exp
+    end
+    return esc(command)
+end
+
 """
-    @handshake exs...
+    @handshake communicator exs...
 
-perform `exs` on all ranks, but only one rank at a time, where
-ranks `r2 > r1` wait for rank `r1` to finish before executing `exs`
+perform `exs` on all ranks in `communicator`, but only one rank at a time, where
+ranks `r2 > r1` wait for rank `r1` to finish before executing `exs`.
+If `communicator` is not provided, `MPI.COMM_WORLD` is used.
 """
-macro handshake(exp)
+macro handshake(communicator, exp)
     command = quote
         mpi_initialized = ClimaOcean.mpi_initialized()
         if !mpi_initialized
             $exp
         else
-            rank   = ClimaOcean.mpi_rank()
-            nprocs = ClimaOcean.mpi_size()
+            rank   = ClimaOcean.mpi_rank($communicator)
+            nprocs = ClimaOcean.mpi_size($communicator)
             for r in 0 : nprocs -1
                 if rank == r
                     $exp
                 end
-                ClimaOcean.global_barrier()
+                ClimaOcean.global_barrier($communicator)
             end
         end
     end
     return esc(command)
-end
+end
+
+macro handshake(exp)
+    command = quote
+        @handshake MPI.COMM_WORLD $exp
+    end
+    return esc(command)
+end