From 8a880b56850a74165f14e19d7e4295e1b6d02173 Mon Sep 17 00:00:00 2001
From: Simone Silvestri <silvestri.simone0@gmail.com>
Date: Thu, 14 Nov 2024 10:47:04 +0100
Subject: [PATCH] Fix major basin removals on Periodic longitudes (#245)

* this works

* compat for OffsetArrays

* Z -> zb

* bugfix

---------

Co-authored-by: Navid C. Constantinou <navidcy@users.noreply.github.com>
---
 Project.toml      |  2 ++
 src/Bathymetry.jl | 54 +++++++++++++++++++++++++++++++++++------------
 2 files changed, 43 insertions(+), 13 deletions(-)

diff --git a/Project.toml b/Project.toml
index f0613a66..f2a08e74 100644
--- a/Project.toml
+++ b/Project.toml
@@ -19,6 +19,7 @@ KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 OrthogonalSphericalShellGrids = "c2be9673-fb75-4747-82dc-aa2bb9f4aed0"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Scratch = "6c6a2e73-6563-6170-7368-637461726353"
@@ -42,6 +43,7 @@ KernelAbstractions = "0.9"
 MPI = "0.20"
 NCDatasets = "0.12, 0.13, 0.14"
 Oceananigans = "0.93.1"
+OffsetArrays = "1.14"
 OrthogonalSphericalShellGrids = "0.1.2, 0.2"
 Scratch = "1"
 SeawaterPolynomials = "0.3.4"
diff --git a/src/Bathymetry.jl b/src/Bathymetry.jl
index d1ee8a47..970b49f1 100644
--- a/src/Bathymetry.jl
+++ b/src/Bathymetry.jl
@@ -17,6 +17,7 @@ using Oceananigans.BoundaryConditions
 using KernelAbstractions: @kernel, @index
 using JLD2
 
+using OffsetArrays
 using ClimaOcean
 
 using NCDatasets
@@ -176,8 +177,7 @@ function regrid_bathymetry(target_grid;
     set!(native_z, z_data)
 
     target_z = interpolate_bathymetry_in_passes(native_z, target_grid; 
-                                                passes = interpolation_passes,
-                                                minimum_depth)
+                                                passes = interpolation_passes)
 
     if minimum_depth > 0
         zi = interior(target_z, :, :, 1)
@@ -198,8 +198,7 @@ end
 
 # Here we can either use `regrid!` (three dimensional version) or `interpolate`
 function interpolate_bathymetry_in_passes(native_z, target_grid; 
-                                          passes = 10,
-                                          minimum_depth = 0)
+                                          passes = 10)
     Nλt, Nφt = Nt = size(target_grid)
     Nλn, Nφn = Nn = size(native_z)
 
@@ -269,16 +268,45 @@ Arguments
                        Default is `Inf`, which means all connected regions are kept.
 
 """
-function remove_minor_basins!(Z::Field, keep_major_basins)
-    Zi = interior(Z, :, :, 1)
-    Zi_cpu = on_architecture(CPU(), Zi)
-    remove_minor_basins!(Zi_cpu, keep_major_basins)
-    set!(Z, Zi_cpu)
+function remove_minor_basins!(zb::Field, keep_major_basins)
+    zb_cpu = on_architecture(CPU(), zb)
+    TX     = topology(zb_cpu.grid, 1)
+    
+    Nx, Ny, _ = size(zb_cpu.grid)
+    zb_data   = maybe_extend_longitude(zb_cpu, TX()) # Outputs a 2D AbstractArray
+
+    remove_minor_basins!(zb_data, keep_major_basins)
+    set!(zb, zb_data[1:Nx, 1:Ny])
+
+    return zb
+end
+
+maybe_extend_longitude(zb_cpu, tx) = interior(zb_cpu, :, :, 1)
 
-    return Z
+# 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!(Z, keep_major_basins)
+remove_major_basins!(zb::OffsetArray, keep_minor_basins) = 
+    remove_minor_basins!(zb.parent, keep_minor_basins)
+
+function remove_minor_basins!(zb, keep_major_basins)
 
     if !isfinite(keep_major_basins)
         throw(ArgumentError("`keep_major_basins` must be a finite number!"))
@@ -288,7 +316,7 @@ function remove_minor_basins!(Z, keep_major_basins)
         throw(ArgumentError("keep_major_basins must be larger than 0."))
     end
 
-    water = Z .< 0
+    water = zb .< 0
     
     connectivity = ImageMorphology.strel(water)
     labels = ImageMorphology.label_components(connectivity)
@@ -324,7 +352,7 @@ function remove_minor_basins!(Z, keep_major_basins)
     end
 
     # Flatten minor basins, corresponding to regions where `labels == NaN`
-    Z[isnan.(labels)] .= 0
+    zb[isnan.(labels)] .= 0
 
     return nothing
 end