ImmersedBoundaryCondition for drag, take 2 (#137)

* take 2 * forse restart build * Update src/OceanSimulations/OceanSimulations.jl Co-authored-by: Gregory L. Wagner <[email protected]> * bugfix * adapt to new syntax * update project * update Oceananigans * Update OceanSimulations.jl * add also generate fluxes * this should work now? * try adding manifest * remove the manifest --------- Co-authored-by: Gregory L. Wagner <[email protected]>
CliMA · Nov 10, 2024 · ff9366b · ff9366b
1 parent 2e3ce57
commit ff9366b
Show file tree

Hide file tree

Showing 7 changed files with 88 additions and 50 deletions.
diff --git a/docs/make.jl b/docs/make.jl
@@ -15,7 +15,7 @@ const OUTPUT_DIR   = joinpath(@__DIR__, "src/literated")
 to_be_literated = [
     "inspect_ecco_data.jl",
     "generate_bathymetry.jl",
-    # "generate_surface_fluxes.jl",
+    "generate_surface_fluxes.jl",
     "single_column_os_papa_simulation.jl",
     # "mediterranean_simulation_with_ecco_restoring.jl",
     "near_global_ocean_simulation.jl"
@@ -50,7 +50,7 @@ pages = [
     "Examples" => [
         "Inspect ECCO2 data" => "literated/inspect_ecco_data.md",
         "Generate bathymetry" => "literated/generate_bathymetry.md",
-        # "Surface fluxes" => "literated/generate_surface_fluxes.md",
+        "Surface fluxes" => "literated/generate_surface_fluxes.md",
         "Single column simulation" => "literated/single_column_os_papa_simulation.md",
         # "Mediterranean simulation with ECCO restoring" => "literated/mediterranean_simulation_with_ecco_restoring.md",
         "Near-global Ocean simulation" => "literated/near_global_ocean_simulation.md",

diff --git a/examples/generate_bathymetry.jl b/examples/generate_bathymetry.jl
@@ -48,9 +48,6 @@ h_one_basin = regrid_bathymetry(grid; major_basins = 1)
 nothing #hide
 
 # Finally, we visualize the generated bathymetry data for the Mediterranean Sea using CairoMakie.
-
-λ, φ, z = nodes(h_smooth)
-
 land_smooth = interior(h_smooth) .>= 0
 interior(h_smooth)[land_smooth] .= NaN
 
@@ -63,13 +60,13 @@ interior(h_one_basin)[land_one_basin] .= NaN
 fig = Figure(size=(850, 1150))
 
 ax = Axis(fig[1, 1], title = "Rough bathymetry", xlabel = "Longitude", ylabel = "Latitude")
-hm = heatmap!(ax, λ, φ, - interior(h_rough, :, :, 1), nan_color=:white, colormap = Reverse(:deep))
+hm = heatmap!(ax, h_rough, nan_color=:grey, colormap = Reverse(:deep))
 
 ax = Axis(fig[2, 1], title = "Smooth bathymetry", xlabel = "Longitude", ylabel = "Latitude")
-hm = heatmap!(ax, λ, φ, - interior(h_smooth, :, :, 1), nan_color=:white, colormap = Reverse(:deep))
+hm = heatmap!(ax, h_smooth, nan_color=:grey, colormap = Reverse(:deep))
 
-ax = Axis(fig[3, 1], title = "Bathymetry without only one basin", xlabel = "Longitude", ylabel = "Latitude")
-hm = heatmap!(ax, λ, φ, - interior(h_one_basin, :, :, 1), nan_color=:white, colormap = Reverse(:deep))
+ax = Axis(fig[3, 1], title = "Bathymetry with only one basin", xlabel = "Longitude", ylabel = "Latitude")
+hm = heatmap!(ax, h_one_basin, nan_color=:grey, colormap = Reverse(:deep))
 
 cb = Colorbar(fig[1:3, 2], hm, height = Relative(3/4), label = "Depth (m)")
 

diff --git a/examples/generate_surface_fluxes.jl b/examples/generate_surface_fluxes.jl
@@ -18,16 +18,13 @@ using CairoMakie
 
 # # Computing fluxes on the ECCO2 grid
 #
-# We start by building the ECCO2 grid, using `ECCO_mask`
-# to identify missing cells.
+# We start by building the ECCO2 grid, using `ECCO_bottom_height` to identify the bottom height.
 
-mask = ECCO_mask()
-grid = mask.grid
-grid = ImmersedBoundaryGrid(grid, GridFittedBoundary(mask))
+grid = ECCO_immersed_grid()
 
 fig = Figure()
 ax  = Axis(fig[1, 1])
-heatmap!(ax, interior(grid.immersed_boundary.mask, :, :, grid.Nz))
+heatmap!(ax, interior(grid.immersed_boundary.bottom_height, :, :, 1))
 save("ECCO_continents.png", fig) #hide
 
 # ![](ECCO_continents.png)
@@ -71,24 +68,25 @@ coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation=Radiation())
 # Now that the surface fluxes are computed, we can extract and visualize them.
 # The turbulent fluxes are stored in `coupled_model.fluxes.turbulent`.
 
-fluxes = coupled_model.fluxes.turbulent.fields
+fluxes  = coupled_model.fluxes.turbulent.fields
+λ, φ, z = nodes(fluxes.sensible_heat)
 
-fig = Figure(size = (800, 400))
+fig = Figure(size = (800, 800), fontsize = 15)
 
-ax = Axis(fig[1, 1], title = "Sensible heat flux (W m⁻²)")
-heatmap!(ax, fluxes.sensible_heat; colormap = :bwr)
+ax = Axis(fig[1, 1], title = "Sensible heat flux (W m⁻²)", ylabel = "Latitude")
+heatmap!(ax, λ, φ, interior(fluxes.sensible_heat, :, :, 1); colormap = :bwr)
 
 ax = Axis(fig[1, 2], title = "Latent heat flux (W m⁻²)")
-heatmap!(ax, fluxes.latent_heat; colormap = :bwr)
+heatmap!(ax, λ, φ, interior(fluxes.latent_heat, :, :, 1); colormap = :bwr)
 
-ax = Axis(fig[2, 1], title = "Zonal wind stress (N m)")
-heatmap!(ax, fluxes.x_momentum; colormap = :bwr)
+ax = Axis(fig[2, 1], title = "Zonal wind stress (N m)", ylabel = "Latitude")
+heatmap!(ax, λ, φ, interior(fluxes.x_momentum, :, :, 1); colormap = :bwr)
 
-ax = Axis(fig[2, 2], title = "Meridional wind stress (N m)")
-heatmap!(ax, fluxes.y_momentum; colormap = :bwr)
+ax = Axis(fig[2, 2], title = "Meridional wind stress (N m)", xlabel = "Longitude")
+heatmap!(ax, λ, φ, interior(fluxes.y_momentum, :, :, 1); colormap = :bwr)
 
-ax = Axis(fig[3, 1], title = "Water vapor flux (kg m⁻² s⁻¹)")
-heatmap!(ax, Mv; colormap = :bwr)
+ax = Axis(fig[3, 1], title = "Water vapor flux (kg m⁻² s⁻¹)", xlabel = "Longitude", ylabel = "Latitude")
+heatmap!(ax, λ, φ, interior(fluxes.water_vapor, :, :, 1); colormap = :bwr)
 
 save("fluxes.png", fig)
 # ![](fluxes.png)
diff --git a/examples/single_column_os_papa_simulation.jl b/examples/single_column_os_papa_simulation.jl
@@ -65,10 +65,7 @@ last_time = simulation_days * snapshots_per_day
 atmosphere = JRA55_prescribed_atmosphere(1:last_time;
                                          longitude = λ★,
                                          latitude = φ★,
-                                         #longitude = (λ★ - 1/4, λ★ + 1/4),
-                                         #latitude  = (φ★ - 1/4, φ★ + 1/4),
-                                         backend = InMemory(),
-                                         include_rivers_and_icebergs = false)
+                                         backend = InMemory())
 
 # This builds a representation of the atmosphere on the small grid
 

diff --git a/src/DataWrangling/ECCO/ECCO.jl b/src/DataWrangling/ECCO/ECCO.jl
@@ -1,6 +1,6 @@
 module ECCO
 
-export ECCOMetadata, ECCO_field, ECCO_mask, adjusted_ECCO_tracers, initialize!
+export ECCOMetadata, ECCO_field, ECCO_mask, ECCO_immersed_grid, adjusted_ECCO_tracers, initialize!
 export ECCO2Monthly, ECCO4Monthly, ECCO2Daily
 export ECCORestoring, LinearlyTaperedPolarMask
 
@@ -88,7 +88,7 @@ empty_ECCO_field(variable_name::Symbol; kw...) = empty_ECCO_field(ECCOMetadata(v
 
 function empty_ECCO_field(metadata::ECCOMetadata;
                           architecture = CPU(), 
-                          horizontal_halo = (3, 3))
+                          horizontal_halo = (7, 7))
 
     Nx, Ny, Nz, _ = size(metadata)
     loc = location(metadata)
@@ -128,7 +128,7 @@ in the x and y direction. The halo in the z-direction is one.
 """
 function ECCO_field(metadata::ECCOMetadata;
                     architecture = CPU(),
-                    horizontal_halo = (3, 3))
+                    horizontal_halo = (7, 7))
 
     download_dataset!(metadata)
     path = metadata_path(metadata)

diff --git a/src/DataWrangling/ECCO/ECCO_mask.jl b/src/DataWrangling/ECCO/ECCO_mask.jl
@@ -1,4 +1,5 @@
 using Oceananigans.Architectures: AbstractArchitecture
+using Oceananigans.Grids: znode
 import ClimaOcean: stateindex
 
 """
@@ -40,6 +41,40 @@ end
     @inbounds mask[i, j, k] = (Tᵢ[i, j, k] == 0) 
 end
 
+"""
+    ECCO_immersed_grid(metadata, architecture = CPU())
+
+Compute the `ImmersedBoundaryGrid` for `metadata` with a bottom height field that is defined 
+by the first non-missing value from the bottom up.
+"""
+function ECCO_immersed_grid(metadata, architecture = CPU())
+
+    mask = ECCO_mask(metadata, architecture)
+    grid = mask.grid
+    bottom = Field{Center, Center, Nothing}(grid)
+
+    # Set the mask with zeros where field is defined
+    launch!(architecture, grid, :xy, _set_height_from_mask!, bottom, grid, mask)
+
+    return ImmersedBoundaryGrid(grid, GridFittedBottom(bottom))
+end
+
+# Default
+ECCO_immersed_grid(arch::AbstractArchitecture=CPU()) = ECCO_immersed_grid(ECCOMetadata(:temperature), arch)
+
+@kernel function _set_height_from_mask!(bottom, grid, mask)
+    i, j = @index(Global, NTuple)
+
+    # Starting from the bottom
+    @inbounds bottom[i, j, 1] = znode(i, j, 1, grid, Center(), Center(), Face())
+
+    # Sweep up
+    for k in 1:grid.Nz
+        z⁺ = znode(i, j, k+1, grid, Center(), Center(), Face())
+        @inbounds bottom[i, j, k] = ifelse(mask[i, j, k], z⁺, bottom[i, j, k])
+    end
+end
+
 struct LinearlyTaperedPolarMask{N, S, Z} 
     northern :: N
     southern :: S
@@ -88,4 +123,4 @@ end
     LX, LY, LZ = loc 
     λ, φ, z = node(i, j, k, grid, LX(), LY(), LZ())
     return mask(φ, z)
-end
+end
diff --git a/src/OceanSimulations/OceanSimulations.jl b/src/OceanSimulations/OceanSimulations.jl
@@ -67,12 +67,9 @@ default_tracer_advection() = FluxFormAdvection(WENO(order=7),
 @inline u_quadratic_bottom_drag(i, j, grid, c, Φ, μ) = @inbounds - μ * Φ.u[i, j, 1] * spᶠᶜᶜ(i, j, 1, grid, Φ)
 @inline v_quadratic_bottom_drag(i, j, grid, c, Φ, μ) = @inbounds - μ * Φ.v[i, j, 1] * spᶜᶠᶜ(i, j, 1, grid, Φ)
 
-@inline is_immersed_drag_u(i, j, k, grid) = Int(immersed_peripheral_node(i, j, k-1, grid, Face(), Center(), Center()) & !inactive_node(i, j, k, grid, Face(), Center(), Center()))
-@inline is_immersed_drag_v(i, j, k, grid) = Int(immersed_peripheral_node(i, j, k-1, grid, Center(), Face(), Center()) & !inactive_node(i, j, k, grid, Center(), Face(), Center()))
-
 # Keep a constant linear drag parameter independent on vertical level
-@inline u_immersed_bottom_drag(i, j, k, grid, clock, fields, μ) = @inbounds - μ * fields.u[i, j, k] * is_immersed_drag_u(i, j, k, grid) * spᶠᶜᶜ(i, j, k, grid, fields) / Δzᶠᶜᶜ(i, j, k, grid)
-@inline v_immersed_bottom_drag(i, j, k, grid, clock, fields, μ) = @inbounds - μ * fields.v[i, j, k] * is_immersed_drag_v(i, j, k, grid) * spᶜᶠᶜ(i, j, k, grid, fields) / Δzᶜᶠᶜ(i, j, k, grid)
+@inline u_immersed_bottom_drag(i, j, k, grid, clock, fields, μ) = @inbounds - μ * fields.u[i, j, k] * spᶠᶜᶜ(i, j, k, grid, fields) 
+@inline v_immersed_bottom_drag(i, j, k, grid, clock, fields, μ) = @inbounds - μ * fields.v[i, j, k] * spᶜᶠᶜ(i, j, k, grid, fields) 
 
 # TODO: Specify the grid to a grid on the sphere; otherwise we can provide a different
 # function that requires latitude and longitude etc for computing coriolis=FPlane...
@@ -105,8 +102,18 @@ function ocean_simulation(grid;
         # Don't let users use advection in a single column model
         tracer_advection = nothing
         momentum_advection = nothing
+
+        # No immersed boundaries in a single column grid
+        u_immersed_bc = nothing
+        v_immersed_bc = nothing
     else
         bottom_drag_coefficient = default_or_override(bottom_drag_coefficient)
+
+        u_immersed_drag = FluxBoundaryCondition(u_immersed_bottom_drag, discrete_form=true, parameters=bottom_drag_coefficient)
+        v_immersed_drag = FluxBoundaryCondition(v_immersed_bottom_drag, discrete_form=true, parameters=bottom_drag_coefficient)
+
+        u_immersed_bc = ImmersedBoundaryCondition(bottom = u_immersed_drag)
+        v_immersed_bc = ImmersedBoundaryCondition(bottom = v_immersed_drag)
     end
 
     bottom_drag_coefficient = convert(FT, bottom_drag_coefficient)
@@ -117,19 +124,23 @@ function ocean_simulation(grid;
     top_ocean_heat_flux          = Jᵀ = Field{Center, Center, Nothing}(grid)
     top_salt_flux                = Jˢ = Field{Center, Center, Nothing}(grid)
 
+    # Construct ocean boundary conditions including surface forcing and bottom drag
+    u_top_bc = FluxBoundaryCondition(τx)
+    v_top_bc = FluxBoundaryCondition(τy)
+    T_top_bc = FluxBoundaryCondition(Jᵀ)
+    S_top_bc = FluxBoundaryCondition(Jˢ)
+
     u_bot_bc = FluxBoundaryCondition(u_quadratic_bottom_drag, discrete_form=true, parameters=bottom_drag_coefficient)
     v_bot_bc = FluxBoundaryCondition(v_quadratic_bottom_drag, discrete_form=true, parameters=bottom_drag_coefficient)
-
-    ocean_boundary_conditions = (u = FieldBoundaryConditions(top = FluxBoundaryCondition(τx), bottom = u_bot_bc),
-                                 v = FieldBoundaryConditions(top = FluxBoundaryCondition(τy), bottom = v_bot_bc),
-                                 T = FieldBoundaryConditions(top = FluxBoundaryCondition(Jᵀ)),
-                                 S = FieldBoundaryConditions(top = FluxBoundaryCondition(Jˢ)))
-
-    if grid isa ImmersedBoundaryGrid
-        Fu = Forcing(u_immersed_bottom_drag, discrete_form=true, parameters=bottom_drag_coefficient)
-        Fv = Forcing(v_immersed_bottom_drag, discrete_form=true, parameters=bottom_drag_coefficient)
-        forcing = merge(forcing, (u=Fu, v=Fv))
-    end
+
+    ocean_boundary_conditions = (u = FieldBoundaryConditions(top = u_top_bc, bottom = u_bot_bc, immersed = u_immersed_bc),
+                                 v = FieldBoundaryConditions(top = v_top_bc, bottom = v_bot_bc, immersed = v_immersed_bc),
+                                 T = FieldBoundaryConditions(top = T_top_bc),
+                                 S = FieldBoundaryConditions(top = S_top_bc))
+
+    # Use the TEOS10 equation of state
+    teos10 = TEOS10EquationOfState(; reference_density)
+    buoyancy = SeawaterBuoyancy(; gravitational_acceleration, equation_of_state=teos10)
 
     buoyancy = SeawaterBuoyancy(; gravitational_acceleration, equation_of_state)