diff --git a/.gitignore b/.gitignore index edb030721..bf2ce4e33 100644 --- a/.gitignore +++ b/.gitignore @@ -63,4 +63,5 @@ docs/src/literated/ .CondaPkg CondaPkg.toml +docs/plans/* !docs/CondaPkg.toml diff --git a/Project.toml b/Project.toml index 9faa65131..3c00b0707 100644 --- a/Project.toml +++ b/Project.toml @@ -79,7 +79,7 @@ SeawaterPolynomials = "0.3.5" SpeedyWeather = "0.17.4" StaticArrays = "1" Statistics = "<0.0.1, 1" -Thermodynamics = "0.15.3" +Thermodynamics = "0.15.3, 1.0" WorldOceanAtlasTools = "0.6" XESMF = "0.1.6" ZipFile = "0.10" diff --git a/src/Bathymetry/Bathymetry.jl b/src/Bathymetry/Bathymetry.jl index 40eed858d..6fbfc45e6 100644 --- a/src/Bathymetry/Bathymetry.jl +++ b/src/Bathymetry/Bathymetry.jl @@ -1,6 +1,10 @@ module Bathymetry export regrid_bathymetry, ORCAGrid +export Basin +export atlantic_ocean_basin, indian_ocean_basin, southern_ocean_basin, pacific_ocean_basin, arctic_ocean_basin +export label_ocean_basins +export meridional_barrier using Downloads using ImageMorphology @@ -12,7 +16,7 @@ using Oceananigans.BoundaryConditions using Oceananigans.DistributedComputations using Oceananigans.DistributedComputations: DistributedGrid, reconstruct_global_grid, all_reduce, @root using Oceananigans.Fields: interpolate! -using Oceananigans.Grids: x_domain, y_domain, topology +using Oceananigans.Grids: x_domain, y_domain, topology, AbstractGrid using Oceananigans.Utils: launch! using OffsetArrays using NCDatasets @@ -22,7 +26,9 @@ using Scratch using ..DataWrangling: Metadatum, native_grid, metadata_path, download_dataset using ..DataWrangling.ETOPO: ETOPO2022 +include("label_ocean_basins.jl") include("regrid_bathymetry.jl") +include("ocean_basin.jl") include("orca_grid.jl") end # module diff --git a/src/Bathymetry/label_ocean_basins.jl b/src/Bathymetry/label_ocean_basins.jl new file mode 100644 index 000000000..58d0bb385 --- /dev/null +++ b/src/Bathymetry/label_ocean_basins.jl @@ -0,0 +1,208 @@ +using Oceananigans.OrthogonalSphericalShellGrids: TripolarGridOfSomeKind +using Oceananigans.Fields: convert_to_0_360 +using ..DataWrangling: BoundingBox + +##### +##### Barrier application functions +##### + +# A barrier is a `BoundingBox` whose horizontal extent defines a rectangular +# region to be temporarily marked as land during connected-component labeling. +# The `z` field of the `BoundingBox` is ignored: the flood-fill operates on +# the 2D bathymetry slice, and barriers act at every depth. + +""" + meridional_barrier(longitude, south, north; width=2) + +Create a narrow meridional `BoundingBox` centered at `longitude` with meridional +extent `[south, north]` and zonal width `width` degrees. Useful for closing +straits like Cape Agulhas or Indonesian passages during basin labeling. +""" +meridional_barrier(longitude, south, north; width=2) = BoundingBox(longitude=(longitude - width/2, longitude + width/2), latitude=(south, north)) + +""" + apply_barrier!(zb, grid, barrier::BoundingBox) + +Mark all cells within the barrier's horizontal region as land (z = 0). +""" +apply_barrier!(zb, grid, barrier::BoundingBox) = launch!(architecture(grid), grid, :xy, _apply_barrier!, zb, grid, barrier) + +apply_barrier!(zb, grid, barriers::Nothing) = zb + +function apply_barrier!(zb, grid, barriers::AbstractVector) + for barrier in barriers + apply_barrier!(zb, grid, barrier) + end + return zb +end + +@kernel function _apply_barrier!(zb, grid, barrier::BoundingBox) + i, j = @index(Global, NTuple) + + in_lon = if isnothing(barrier.longitude) || (barrier.longitude[2] - barrier.longitude[1] >= 360) + true + else + bw = convert_to_0_360(barrier.longitude[1]) + be = convert_to_0_360(barrier.longitude[2]) + λ = λnode(i, j, 1, grid, Center(), Center(), Center()) + λ = convert_to_0_360(λ) + (bw <= λ <= be) + end + + in_lat = if isnothing(barrier.latitude) + true + else + φ = φnode(i, j, 1, grid, Center(), Center(), Center()) + barrier.latitude[1] <= φ <= barrier.latitude[2] + end + + @inbounds zb[i, j, 1] = ifelse(in_lon & in_lat, zero(grid), zb[i, j, 1]) +end + +# Since the strel algorithm in `label_components` does not recognize periodic boundaries, +# before labeling connected regions we copy half the longitude extent on each side so that +# water cells near the boundary are correctly identified as connected. +function copy_periodic_longitude(zb_cpu::Field, ::Periodic) + Nλ = size(zb_cpu, 1) + half = Nλ ÷ 2 + + # 2D slice of the bathymetry at the bottom level. + zb_data = zb_cpu.data[1:Nλ, :, 1] + zb_parent = parent(zb_data) + + # Concatenate a copy of the eastern half on the left and the western half on the right. + # This is an O(Nλ × Nφ) CPU allocation, acceptable for the serial labeling step. + concatenated_zb = vcat(zb_parent[half:Nλ, :], zb_parent, zb_parent[1:half, :]) + + # Offsets so that index 1 maps back to the original domain start + xoffset = - half - 1 + yoffset = zb_cpu.data.offsets[2] + + return OffsetArray(concatenated_zb, xoffset, yoffset) +end + +copy_periodic_longitude(zb_cpu::Field, tx) = interior(zb_cpu, :, :, 1) + +""" + label_ocean_basins(zb_field, TX, core_size) + +Creates a matrix with a unique label for each connected basin. Useful for inpainting the bathymetry and +computing the masks for oceanic basins. + +Handles periodic boundary extension internally and returns labels for the core region only. +""" +function label_ocean_basins(zb_field, TX, size) + zb = copy_periodic_longitude(zb_field, TX()) # Outputs a 2D AbstractArray + + water = zb .< 0 + + connectivity = ImageMorphology.strel(water) + labels = ImageMorphology.label_components(connectivity) + + Nx, Ny = size[1], size[2] + core_labels = labels[1:Nx, 1:Ny] + + # Enforce periodicity: merge labels that should be connected across + # the periodic boundaries. This handles cases where a barrier (e.g., blocking + # the Southern Ocean) prevents the extended domain from connecting basins + # that are actually periodic neighbors. + enforce_periodic_labels!(core_labels, TX()) + enforce_tripolar_labels!(core_labels, zb_field.grid) + + return core_labels +end + +""" + enforce_periodic_labels!(labels, ::Periodic) + +Merge labels that should be connected due to periodic boundary conditions. +For each latitude, if the westernmost and easternmost cells are both water +(non-zero labels), they are periodic neighbors and must have the same label. +""" +function enforce_periodic_labels!(labels, ::Periodic) + Nx, Ny = Base.size(labels) + + for j in 1:Ny + label_west = labels[1, j] + label_east = labels[Nx, j] + + # Both cells are water and have different labels: merge them + if label_west != 0 && label_east != 0 && label_west != label_east + # Replace all occurrences of label_east with label_west + replace!(labels, label_east => label_west) + end + end + + return labels +end + +# No-op for non-periodic domains +enforce_periodic_labels!(labels, tx) = labels + +""" + enforce_tripolar_labels!(labels, ::TripolarGridOfSomeKind) + +Merge labels that should be connected due to zipper boundary conditions. +For each longitude, if cells reflected across the fold (Nx÷2) are water +(non-zero labels), they are periodic neighbors and must have the same label. +""" +function enforce_tripolar_labels!(labels, ::TripolarGridOfSomeKind) + Nx, Ny = Base.size(labels) + + for i in 1:Nx÷2 + label_west = labels[i, Ny] + label_east = labels[Nx-i+1, Ny] + + # Both cells are water and have different labels: merge them + if label_west != 0 && label_east != 0 && label_west != label_east + # Replace all occurrences of label_east with label_west + replace!(labels, label_east => label_west) + end + end + + return labels +end + +# No-op for non-tripolar grids +enforce_tripolar_labels!(labels, grid) = labels + +# Utilities to label ocean basins passing only the grid +function label_ocean_basins(grid::AbstractGrid; barriers=nothing) + @warn "The grid is not immersed, there is only one ocean basin!" + Nx, Ny, Nz = size(grid) + return zeros(Int, Nx, Ny) +end + +""" + label_ocean_basins(grid::ImmersedBoundaryGrid; barriers=nothing) + +Label connected ocean basins in an ImmersedBoundaryGrid. + +Keyword Arguments +================= +- `barriers`: Collection of `BoundingBox`es applied before labeling. Each barrier + temporarily marks its horizontal rectangle as land, allowing separation + of connected ocean basins (e.g., separating Atlantic from Pacific via + the Southern Ocean). +""" +function label_ocean_basins(grid::ImmersedBoundaryGrid; barriers=nothing) + + # The labelling algorithm works only on CPUs + cpu_grid = on_architecture(CPU(), grid) + + TX = topology(cpu_grid, 1) + zb = cpu_grid.immersed_boundary.bottom_height + + # If barriers are specified, apply them to a copy of the bathymetry. + # Using `similar(zb)` preserves the field type, which allows for future + # bathymetry representations beyond plain 2D `GridFittedBottom` fields. + if !isnothing(barriers) + zb_modified = similar(zb) + parent(zb_modified) .= parent(zb) + apply_barrier!(zb_modified, cpu_grid, barriers) + else + zb_modified = zb + end + + return label_ocean_basins(zb_modified, TX, size(cpu_grid)) +end diff --git a/src/Bathymetry/ocean_basin.jl b/src/Bathymetry/ocean_basin.jl new file mode 100644 index 000000000..e4c9eca74 --- /dev/null +++ b/src/Bathymetry/ocean_basin.jl @@ -0,0 +1,395 @@ +using Oceananigans.Grids: λnode, φnode, on_architecture +using ImageMorphology + +##### +##### Basin struct +##### + +""" + Basin{M, G, S} + +A connected water region (ocean basin) identified on a grid, together with the +boolean mask that labels cells belonging to it and the seed points used to pick +the connected component. + +Fields +====== +- `mask`: a 2D `Field{Center, Center, Nothing}` with `Bool` values — `true` for + cells belonging to the basin, `false` otherwise. +- `grid`: the grid on which the basin is defined. +- `seed_points`: `(λ, φ)` coordinate pairs used to identify which connected + component corresponds to the desired basin. The algorithm finds + the label at each seed point and builds the mask from that label. +""" +struct Basin{M, G, S} + mask :: M + grid :: G + seed_points :: S +end + +Base.summary(basin::Basin) = "Basin" + +function Base.show(io::IO, basin::Basin) + print(io, summary(basin), " on ", summary(basin.grid)) +end + +# Forward getindex to the underlying mask +Base.getindex(basin::Basin, i, j, k) = basin.mask[i, j, k] + +##### +##### Basin union +##### + +""" + Base.:|(basin1::Basin, basin2::Basin) + +Combine two basins into a new one whose mask is the union of both. Cells that +belong to either input basin are `true` in the combined mask. This mirrors the +elementwise `.|` applied to the underlying boolean masks. +""" +function Base.:|(basin1::Basin, basin2::Basin) + grid = basin1.grid + combined_mask = Field{Center, Center, Nothing}(grid, Bool) + + # Union: a cell is in the combined mask if it is in either input mask + parent(combined_mask) .= parent(basin1.mask) .| parent(basin2.mask) + fill_halo_regions!(combined_mask) + + # Combine seed points + combined_seeds = (basin1.seed_points..., basin2.seed_points...) + + return Basin(combined_mask, grid, combined_seeds) +end + +##### +##### Connected component utilities +##### + +""" + find_label_at_point(labels, grid, λs, φs; radius = 2) + +Find the connected component label at the given longitude/latitude seed point. +Returns the label value, or 0 if the point is on land or outside the domain. +Checks in a circular cap of radius `radius` degrees around the seed point +""" +function find_label_at_point(labels, grid, λs, φs; radius = 2) + Nx, Ny, _ = size(grid) + + # Find grid cell containing the seed point + for j in 1:Ny + for i in 1:Nx + λ = λnode(i, j, 1, grid, Center(), Center(), Center()) + φ = φnode(i, j, 1, grid, Center(), Center(), Center()) + + λ = convert_to_0_360(λ) + + Δλ = if isnothing(λs) + zero(λ) + else + λs = convert_to_0_360(λs) + λ - λs + end + + # Check if this cell contains the seed point (within a certain radius) + if Δλ^2 + (φ - φs)^2 < radius^2 + return labels[i, j] + end + end + end + + return 0 # Seed point not found +end + +""" + create_basin_mask_from_label(grid, labels, basin_label) + +Create a mask field for all cells with the given label. +""" +function create_basin_mask_from_label(grid, labels, basin_label) + Nx, Ny, _ = size(grid) + + mask = Field{Center, Center, Nothing}(grid, Bool) + + launch!(CPU(), grid, :xy, _compute_basin_mask!, + mask, grid, labels, basin_label) + + fill_halo_regions!(mask) + + return mask +end + +@kernel function _compute_basin_mask!(mask, grid, labels, basin_label) + i, j = @index(Global, NTuple) + correct_basin = @inbounds labels[i, j] == basin_label + @inbounds mask[i, j, 1] = correct_basin +end + +##### +##### Some useful Basin seeds and barriers +##### + +const SOUTHERN_OCEAN_SEPARATION_BARRIER = BoundingBox(longitude=(-180, 180), latitude=(-56, -54)) + +const ATLANTIC_OCEAN_BARRIERS = [ + meridional_barrier(20, -90, -30), # Cape Agulhas + meridional_barrier(289, -90, -30), # Drake Passage +] + +const INDIAN_OCEAN_BARRIERS = [ + meridional_barrier(141, -90, -3), # Indonesian side + meridional_barrier(20, -90, -30), # Cape Agulhas + BoundingBox(longitude=(105, 141), latitude=(-4, -3)), # Indonesian/Asian seas (zonal barrier at 3.5ᵒ S) +] + +const SOUTHERN_OCEAN_BARRIERS = [SOUTHERN_OCEAN_SEPARATION_BARRIER] + +const PACIFIC_OCEAN_BARRIERS = [ + meridional_barrier(141, -90, -3), # Indonesian side + meridional_barrier(20, -90, -30), # Cape Agulhas + BoundingBox(longitude=(105, 141), latitude=(-4, -3)), # Indonesian/Asian seas (zonal barrier at 3.5ᵒ S) +] + +# Seed points for Atlantic Ocean (definitely in the Atlantic) +const ATLANTIC_SEED_POINTS = [ + (-30, 0), # Central equatorial Atlantic + (-40, 30), # North Atlantic + (-25, -20), # South Atlantic +] + +# Seed points for Indian Ocean +const INDIAN_SEED_POINTS = [ + (70, -10), # Central Indian Ocean + (60, 10), # Arabian Sea region + (90, -20), # Eastern Indian Ocean +] + +# Seed points for Southern Ocean +const SOUTHERN_SEED_POINTS = [ + (0, -60), # South Atlantic sector + (90, -60), # Indian Ocean sector + (180, -60), # Pacific sector (date line) + (-90, -60), # South Pacific sector +] + +# Seed points for Pacific Ocean +const PACIFIC_SEED_POINTS = [ + (180, 0), # Central equatorial Pacific (dateline) + (-150, 20), # North Pacific (Hawaii region) + (-120, -20), # South Pacific + # Same values but mapped to [0, 360) + (180, 0), # Central equatorial Pacific + (-150 + 360, 20), # North Pacific + (-120 + 360, -20), # South Pacific +] + +##### +##### Basin constructor +##### + +add_barrier(v::AbstractVector, b::BoundingBox) = [v..., b] +add_barrier(v::BoundingBox, b::BoundingBox) = [v, b] +add_barrier(::Nothing, b::BoundingBox) = b + +""" + Basin(grid; + south_boundary = nothing, + north_boundary = nothing, + seed_points = [(0, 0)], + barriers = nothing) + +Build a `Basin` — a single connected water region on `grid` together with its +boolean mask. + +The algorithm first labels every connected water region (using +`ImageMorphology.label_components`), optionally splitting regions with +`barriers`. It then retrieves the basin whose label matches the first +`seed_point` that falls on water. + +Multiple `seed_points` are tried in order as fallbacks (the first hit wins); +multiple `barriers` are applied simultaneously before labeling so that each +barrier independently blocks water connectivity. + +Arguments +========= +- `grid`: An `ImmersedBoundaryGrid` whose immersed boundary defines the coastlines. + +Keyword Arguments +================= +- `south_boundary`: Southern latitude limit — cells south of this become land. Default: `nothing`. +- `north_boundary`: Northern latitude limit — cells north of this become land. Default: `nothing`. +- `seed_points`: `(λ, φ)` pairs identifying the target basin. The first seed that lands + on water determines which connected component becomes `true` in the mask. + Multiple seeds are tried as fallbacks for grids where a single point + may fall on land. Default: `[(0, 0)]`. +- `barriers`: A `BoundingBox` (or a `Vector` of them) applied before labeling. + Each barrier temporarily marks its horizontal rectangle as land, + preventing the flood-fill from crossing it (e.g., closing Drake Passage + separates the Atlantic from the Pacific). The `z` field of the + `BoundingBox` is ignored. Default: `nothing`. +""" +function Basin(grid; + south_boundary = nothing, + north_boundary = nothing, + seed_points = [(0, 0)], + barriers = nothing) + + # The computations are 2D and require serial algorithms, so + # we perform the computation on the CPU then move the output + # to the GPU if the initial grid was a GPU grid + cpu_grid = Oceananigans.on_architecture(CPU(), grid) + + # Enforce north and south boundaries + if !isnothing(south_boundary) + barriers = add_barrier(barriers, BoundingBox(longitude=nothing, latitude=(-90, south_boundary))) + end + + if !isnothing(north_boundary) + barriers = add_barrier(barriers, BoundingBox(longitude=nothing, latitude=(north_boundary, 90))) + end + + # Compute connected component labels for all ocean cells + # Barriers are applied to temporarily separate connected basins + labels = label_ocean_basins(cpu_grid; barriers) + + # Find the basin label using seed points + basin_label = 0 + for (λs, φs) in seed_points + label = find_label_at_point(labels, cpu_grid, λs, φs) + if label > 0 + basin_label = label + break + end + end + + if basin_label == 0 + @warn "Could not find the basin in grid. Returning empty mask." + mask = Field{Center, Center, Nothing}(grid, Bool) + return Basin(mask, grid, seed_points) + end + + # Create mask from label with latitude bounds + mask = create_basin_mask_from_label(cpu_grid, labels, basin_label) + mask = Oceananigans.on_architecture(architecture(grid), mask) + + return Basin(mask, grid, seed_points) +end + +##### +##### Convenience functions for Earth's ocean basins +##### + +""" + atlantic_ocean_basin(grid; include_southern_ocean=false, kw...) + +Build a `Basin` for Earth's Atlantic Ocean with predefined barriers and seed points. + +Keyword Arguments +================= +- `include_southern_ocean`: If `true`, extends the Atlantic basin into the Southern Ocean + sector below the standard separation latitude (~55°S). Default: `false`. +- `south_boundary`: Southern latitude limit. Default: -50.0 (or -90.0 if `include_southern_ocean=true`) +- `north_boundary`: Northern latitude limit. Default: 65.0 +- Other keyword arguments are passed to `Basin`. +""" +function atlantic_ocean_basin(grid; + include_southern_ocean = true, + south_boundary = include_southern_ocean ? -90 : -50, + north_boundary = 65, + barriers = ATLANTIC_OCEAN_BARRIERS, + seed_points = ATLANTIC_SEED_POINTS, + kw...) + + if !include_southern_ocean + barriers = [barriers..., SOUTHERN_OCEAN_SEPARATION_BARRIER] + end + + return Basin(grid; south_boundary, north_boundary, barriers, seed_points, kw...) +end + +""" + indian_ocean_basin(grid; include_southern_ocean=false, kw...) + +Build a `Basin` for Earth's Indian Ocean with predefined barriers and seed points. + +Keyword Arguments +================= +- `include_southern_ocean`: If `true`, extends the Indian basin into the Southern Ocean + sector below the standard separation latitude (~55°S). Default: `false`. +- `south_boundary`: Southern latitude limit. Default: -50.0 (or -90.0 if `include_southern_ocean=true`) +- `north_boundary`: Northern latitude limit. Default: 30.0 +- Other keyword arguments are passed to `Basin`. +""" +function indian_ocean_basin(grid; + include_southern_ocean = true, + south_boundary = include_southern_ocean ? -90 : -50, + north_boundary = 30, + barriers = INDIAN_OCEAN_BARRIERS, + seed_points = INDIAN_SEED_POINTS, + kw...) + + if !include_southern_ocean + barriers = [barriers..., SOUTHERN_OCEAN_SEPARATION_BARRIER] + end + + return Basin(grid; south_boundary, north_boundary, barriers, seed_points, kw...) +end + +""" + southern_ocean_basin(grid; kw...) + +Build a `Basin` for Earth's Southern Ocean with predefined barriers and seed points. +Default boundaries: south=-90.0, north=-35.0 +""" +function southern_ocean_basin(grid; + south_boundary = -90, + north_boundary = -35, + barriers = SOUTHERN_OCEAN_BARRIERS, + seed_points = SOUTHERN_SEED_POINTS, + kw...) + return Basin(grid; south_boundary, north_boundary, barriers, seed_points, kw...) +end + +""" + pacific_ocean_basin(grid; include_southern_ocean=false, kw...) + +Build a `Basin` for Earth's Pacific Ocean with predefined barriers and seed points. + +Keyword Arguments +================= +- `include_southern_ocean`: If `true`, extends the Pacific basin into the Southern Ocean + sector below the standard separation latitude (~55°S). Default: `false`. +- `south_boundary`: Southern latitude limit. Default: -50.0 (or -90.0 if `include_southern_ocean=true`) +- `north_boundary`: Northern latitude limit. Default: 65.0 +- Other keyword arguments are passed to `Basin`. +""" +function pacific_ocean_basin(grid; + include_southern_ocean = true, + south_boundary = include_southern_ocean ? -90 : -50, + north_boundary = 65, + barriers = PACIFIC_OCEAN_BARRIERS, + seed_points = PACIFIC_SEED_POINTS, + kw...) + + if !include_southern_ocean + barriers = [barriers..., SOUTHERN_OCEAN_SEPARATION_BARRIER] + end + + return Basin(grid; south_boundary, north_boundary, barriers, seed_points, kw...) +end + +""" + arctic_ocean_basin(grid; kw...) + +Build a `Basin` for Earth's Arctic Ocean with predefined seed points. +Default boundaries: south=65.0, north=91.0 +""" +function arctic_ocean_basin(grid; + include_southern_ocean = true, + south_boundary = 65, + north_boundary = 91, + barriers = nothing, + seed_points = [(nothing, 90)], + kw...) + + return Basin(grid; south_boundary, north_boundary, barriers, seed_points, kw...) +end diff --git a/src/Bathymetry/regrid_bathymetry.jl b/src/Bathymetry/regrid_bathymetry.jl index 430cad426..6227b78f0 100644 --- a/src/Bathymetry/regrid_bathymetry.jl +++ b/src/Bathymetry/regrid_bathymetry.jl @@ -413,49 +413,6 @@ Arguments """ function remove_minor_basins!(zb::Field, keep_major_basins) - zb_cpu = on_architecture(CPU(), zb) - TX = topology(zb_cpu.grid, 1) - - core_size = Nx, Ny, _ = size(zb_cpu.grid) - zb_data = maybe_extend_longitude(zb_cpu, TX()) # Outputs a 2D AbstractArray - - # For periodic grids, `zb_data` may include an extended longitude domain so that - # ImageMorphology can identify connected components correctly across periodic - # boundaries. However, we want `keep_major_basins` to refer to the number of - # major basins in the *physical* model domain only. - # - # To achieve this, we rank basins by their area within the core region - # `1:Nx, 1:Ny` of `zb_data`, rather than within the whole extended array. - remove_minor_basins!(zb_data, keep_major_basins, core_size) - set!(zb, zb_data[1:Nx, 1:Ny]) - - return zb -end - -maybe_extend_longitude(zb_cpu, tx) = interior(zb_cpu, :, :, 1) - -# Since the strel algorithm in `remove_major_basins` does not recognize periodic boundaries, -# before removing connected regions, we extend the longitude direction if it is periodic. -# An extension of half the domain is enough. -function maybe_extend_longitude(zb_cpu, ::Periodic) - Nx = size(zb_cpu, 1) - nx = Nx ÷ 2 - - zb_data = zb_cpu.data[1:Nx, :, 1] - zb_parent = zb_data.parent - - # Add information on the LHS and to the RHS - zb_parent = vcat(zb_parent[nx:Nx, :], zb_parent, zb_parent[1:nx, :]) - - # Update offsets - yoffsets = zb_cpu.data.offsets[2] - xoffsets = - nx - - return OffsetArray(zb_parent, xoffsets, yoffsets) -end - -function remove_minor_basins!(zb, keep_major_basins, core_size) - if !isfinite(keep_major_basins) throw(ArgumentError("`keep_major_basins` must be a finite number!")) end @@ -464,56 +421,60 @@ function remove_minor_basins!(zb, keep_major_basins, core_size) throw(ArgumentError("keep_major_basins must be larger than 0.")) end - water = zb .< 0 + zb_cpu = on_architecture(CPU(), zb) + TX, TY, _ = topology(zb_cpu.grid) - connectivity = ImageMorphology.strel(water) - labels = ImageMorphology.label_components(connectivity) + Nx = Base.length(Center(), TX(), zb_cpu.grid.Nx) + Ny = Base.length(Center(), TY(), zb_cpu.grid.Ny) + # Get labels for the core region (extension is handled internally by label_ocean_basins) + labels = label_ocean_basins(zb_cpu, TX, (Nx, Ny)) nlabels = maximum(labels) - # Rank labels by the number of elements they occupy within the *core* region. - - Nx, Ny = core_size[1], core_size[2] - core_view = labels[1:Nx, 1:Ny] + if nlabels == 0 + return zb # No basins found + end + # Rank labels by the number of elements they occupy total_elements = zeros(nlabels) label_elements = zeros(Int, nlabels) for e in 1:nlabels - cnt = count(==(e), core_view) + cnt = count(==(e), labels) total_elements[e] = cnt label_elements[e] = e end - # Ignore labels that do not intersect the core at all. - # These correspond to basins that live only in the periodic extension. + # Find valid basins (those with at least one cell) valid = findall(>(0), total_elements) - mm_basins = [] # major basins indexes + major_basins = Int[] # indices of major basins to keep m = 1 - # We add basin indexes until we reach the specified number (m == keep_major_basins) or - # we run out of basins to keep -> isempty(valid) + # Keep the largest basins up to keep_major_basins while (m <= keep_major_basins) && !isempty(valid) - # Among the remaining valid labels, find the one with the largest core area. _, idx = findmax(total_elements[valid]) next_label = label_elements[valid[idx]] - push!(mm_basins, next_label) + push!(major_basins, next_label) deleteat!(valid, idx) m += 1 end - labels = map(Float64, labels) + # Modify the bathymetry: set minor basin cells to 0 (land) + # Work on interior view which directly modifies the underlying data + zb_data = interior(zb_cpu, :, :, 1) - for ℓ = 1:maximum(labels) - remove_basin = all(ℓ != m for m in mm_basins) - if remove_basin - labels[labels .== ℓ] .= NaN # Regions to remove + for j in 1:Ny, i in 1:Nx + label = labels[i, j] + if label > 0 && !(label in major_basins) + zb_data[i, j] = 0 # Flatten this cell (make it land) end end - # Flatten minor basins, corresponding to regions where `labels == NaN` - zb[isnan.(labels)] .= 0 + # If original field was on a different architecture, copy back + if zb !== zb_cpu + set!(zb, zb_cpu) + end - return nothing + return zb end diff --git a/src/EarthSystemModels/InterfaceComputations/sea_ice_ocean_fluxes.jl b/src/EarthSystemModels/InterfaceComputations/sea_ice_ocean_fluxes.jl index a715493d3..ce3df1880 100644 --- a/src/EarthSystemModels/InterfaceComputations/sea_ice_ocean_fluxes.jl +++ b/src/EarthSystemModels/InterfaceComputations/sea_ice_ocean_fluxes.jl @@ -198,8 +198,8 @@ end # ============================================= # Returns interfacial heat flux, melt rate qᵐ, and interface T, S 𝒬ⁱᵒ, qᵐ, Tᵦ, Sᵦ = compute_interface_heat_flux(flux_formulation, - ocean_surface_state, ice_state, - liquidus, ocean_properties, ℰ, u★) + ocean_surface_state, ice_state, + liquidus, ocean_properties, ℰ, u★) # Store interface values and heat flux @inbounds 𝒬ⁱⁿᵗ[i, j, 1] = 𝒬ⁱᵒ diff --git a/src/NumericalEarth.jl b/src/NumericalEarth.jl index 9dd134f19..47205051d 100644 --- a/src/NumericalEarth.jl +++ b/src/NumericalEarth.jl @@ -32,6 +32,14 @@ export OSPapaHourly, JRA55NetCDFBackend, regrid_bathymetry, + label_ocean_basins, + Basin, + meridional_barrier, + atlantic_ocean_basin, + indian_ocean_basin, + southern_ocean_basin, + pacific_ocean_basin, + arctic_ocean_basin, Metadata, Metadatum, ECCOMetadatum, diff --git a/src/Oceans/ocean_simulation.jl b/src/Oceans/ocean_simulation.jl index b2876b97c..ad716622a 100644 --- a/src/Oceans/ocean_simulation.jl +++ b/src/Oceans/ocean_simulation.jl @@ -191,6 +191,7 @@ defaults on a per-field basis. function ocean_simulation(grid; Δt = estimate_maximum_Δt(grid), closure = default_ocean_closure(), + clock = Clock(grid), tracers = (:T, :S), free_surface = default_free_surface(grid), reference_density = 1020, @@ -316,6 +317,7 @@ function ocean_simulation(grid; end ocean_model = HydrostaticFreeSurfaceModel(grid; + clock, buoyancy, closure, biogeochemistry, diff --git a/test/test_bathymetry.jl b/test/test_bathymetry.jl index e5edaecf7..b7ceac273 100644 --- a/test/test_bathymetry.jl +++ b/test/test_bathymetry.jl @@ -6,7 +6,14 @@ using NumericalEarth.Bathymetry: remove_minor_basins!, BathymetryRegridding, cache_filename, load_bathymetry_cache, - save_bathymetry_cache + save_bathymetry_cache, + label_ocean_basins, + find_label_at_point, + Basin, + atlantic_ocean_basin, + meridional_barrier, + ATLANTIC_OCEAN_BARRIERS +using NumericalEarth.DataWrangling: BoundingBox using NumericalEarth.DataWrangling.ETOPO using Statistics @@ -144,3 +151,97 @@ end result4 = regrid_bathymetry(grid; cache=false) @test parent(result1) == parent(result4) end + +@testset "Barrier geometry" begin + @info "Testing barrier geometry utilities..." + + # A barrier is a `BoundingBox` whose horizontal extent defines the rectangle + # to be treated as land during connected-component labeling. + barrier = BoundingBox(longitude=(-10, 10), latitude=(-5, 5)) + @test barrier.longitude == (-10, 10) + @test barrier.latitude == (-5, 5) + + # Test meridional_barrier constructor (longitude, south, north) + meridional = meridional_barrier(20, -36, -30) + @test meridional.longitude == (19, 21) # 20 ± 2/2 + @test meridional.latitude == (-36, -30) + + # Test meridional_barrier with custom width + meridional_wide = meridional_barrier(20, -36, -30; width=4) + @test meridional_wide.longitude == (18, 22) + + # Latitude band covering the full longitude range + band = BoundingBox(longitude=(-180, 180), latitude=(-60, -55)) + @test band.longitude == (-180, 180) + @test band.latitude == (-60, -55) +end + +@testset "Ocean basin labeling with barriers" begin + @info "Testing ocean basin labeling with barriers..." + + for arch in test_architectures + # Create a global grid + grid = LatitudeLongitudeGrid(arch; + size = (90, 45, 10), + longitude = (-180, 180), + latitude = (-90, 90), + z = (-6000, 0)) + + # Regrid real bathymetry onto this grid + bottom_height = regrid_bathymetry(grid) + ibg = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height)) + + # Test labeling without barriers - all oceans should have the same label + # (they're connected via the Southern Ocean) + labels_no_barrier = label_ocean_basins(ibg) + + # Find labels at Atlantic and Pacific seed points + atlantic_label = find_label_at_point(labels_no_barrier, ibg, -30.0, 0.0) + pacific_label = find_label_at_point(labels_no_barrier, ibg, -170.0, 0.0) + + # Without barriers, Atlantic and Pacific should have the same label + # (connected via Southern Ocean) + @test atlantic_label == pacific_label + @test atlantic_label > 0 # Should find a valid basin + + # Test labeling with barriers - oceans should be separated + labels_with_barrier = label_ocean_basins(ibg; barriers=ATLANTIC_OCEAN_BARRIERS) + + # With barriers, Atlantic should still be found + atlantic_label_with_barrier = find_label_at_point(labels_with_barrier, ibg, -30.0, 0.0) + @test atlantic_label_with_barrier > 0 + end +end + +@testset "Basin creation" begin + @info "Testing Basin creation..." + + for arch in test_architectures + # Create a global grid at 1° resolution (needed to properly resolve + # Central America and separate Atlantic from Pacific) + grid = LatitudeLongitudeGrid(arch; + size = (360, 180, 10), + longitude = (-180, 180), + latitude = (-90, 90), + z = (-6000, 0)) + + bottom_height = regrid_bathymetry(grid) + ibg = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height)) + + # Test atlantic_ocean_basin creation + atlantic = atlantic_ocean_basin(ibg) + @test atlantic isa Basin + @test sum(interior(atlantic.mask)) > 0 # Should have some ocean cells + + mask = on_architecture(CPU(), atlantic.mask) + + # Test that the mask is properly bounded + # Atlantic mask should not include cells in the Pacific + # (seed point at -170°, 0° should be 0) + pacific_point_i = findfirst(i -> -175 < i < -165, range(-180, 180, length=360)) + equator_j = 90 # equator for 180 latitude points + if !isnothing(pacific_point_i) + @test mask[pacific_point_i, equator_j, 1] == 0 + end + end +end