WIP new smart coarsening with Graphs.jl: Don't lump disconnected cells!

Author: briochemc

Project.toml

@@ -1,24 +1,26 @@
 name = "OceanTransportMatrixBuilder"
 uuid = "c2b4a04e-6049-4fc4-aa6a-5508a29a1e1c"
+version = "0.8.0"
 authors = ["Benoit Pasquier <briochemc@gmail.com> and contributors"]
-version = "0.7.2"
 
 [deps]
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
+Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
 Distances = "0.10"
+Graphs = "1.13.1"
 LinearAlgebra = "1"
 SparseArrays = "1"
 StaticArrays = "1"
 julia = "1.10"
 
 [extras]
-CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
 Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
@@ -35,4 +37,4 @@ YAXArrays = "c21b50f5-aa40-41ea-b809-c0f5e47bfa5c"
 Zarr = "0a941bbe-ad1d-11e8-39d9-ab76183a1d99"
 
 [targets]
-test = ["CairoMakie", "CSV", "DataFrames", "DimensionalData", "Downloads", "GibbsSeaWater", "LinearAlgebra", "Makie", "NaNStatistics", "NetCDF", "Test", "TestItemRunner", "TestItems", "Unitful", "YAXArrays", "Zarr"]
+test = ["CairoMakie", "CSV", "DataFrames", "DimensionalData", "Downloads", "GibbsSeaWater", "LinearAlgebra", "Makie", "NaNStatistics", "NetCDF", "Test", "TestItemRunner", "TestItems", "Unitful", "YAXArrays", "Zarr"]

src/OceanTransportMatrixBuilder.jl

@@ -7,6 +7,7 @@ using LinearAlgebra
 # deps
 using Distances: haversine
 using StaticArrays: FieldVector
+using Graphs: SimpleGraph, connected_components
 
 # Main functions to be exposed
 include("matrixbuilding.jl")

src/extratools.jl

@@ -31,32 +31,56 @@ where the lumping should occur.
 Outside of this region, no lumping.
 Default is `f=Returns(true)`, i.e. lump everywhere.
 """
-function lump_and_spray(wet3D, vol, mask=trues(size(wet3D)); di=2, dj=2, dk=1)
+function lump_and_spray(wet3D, vol, T, mask=trues(size(wet3D)); di=2, dj=2, dk=1)
 
     # extend the grid to avoid lumping cells outside of bounds
     nxyz = size(wet3D)
     C = CartesianIndices(nxyz)
     LUMPidx = zeros(Int, nxyz .+ (di - 1, dj - 1, dk - 1))
 
+    wet = wet3D[:]
+    i, j = findnz(T)
+    L = LinearIndices(nxyz .+ (di - 1, dj - 1, dk - 1))
+    wet3Dext = falses(nxyz .+ (di - 1, dj - 1, dk - 1))
+    wet3Dext[C[wet3D]] .= true
+    Lwet = L[wet3Dext] # careful here, as L is larger than wet3D
+
+    connectivitymatrix = sparse(Lwet[i], Lwet[j], true, length(LUMPidx), length(LUMPidx))
+
     # loop once over all grid cells and assign lumped indices
-    c = 1 # index / counter
+    c = 2 # index / counter (we start at 2 because 1 is reserved for dry cells)
     neighbours = CartesianIndices((0:di-1, 0:dj-1, 0:dk-1))
     for 𝑖 in eachindex(C)
         C𝑖 = C[𝑖]
-        LUMPidx[C𝑖] ≠ 0 && mask[C𝑖] && continue # skip if index already assigned
+        # skip if index already assigned and in mask
+        #(if assigned but not in mask, must reassign)
+        LUMPidx[C𝑖] > 0 && mask[C𝑖] && continue
         if mask[C𝑖] # if in region, assign all neighbours also in region
-            LUMPidx[C𝑖 .+ neighbours] .= c
-        else # to only this cell if outside region
+            # list of neighbours
+            L𝑖 = L[C𝑖 .+ neighbours][:]
+            # First, assign dry index to dry cells
+            localwet = wet3Dext[L𝑖]
+            dryidx = L𝑖[.!localwet]
+            LUMPidx[dryidx] .= 1
+            # Then build a simple connectivity graph of the remaining cells to be lumped
+            wetidx = L𝑖[localwet]
+            localconnectivitymatrix = view(connectivitymatrix, wetidx, wetidx)
+            g = SimpleGraph(localconnectivitymatrix)
+            # and split it in connected components
+            for comp in connected_components(g)
+                LUMPidx[wetidx[comp]] .= c
+                c += 1
+            end
+        else # if outside region, assign new index
             LUMPidx[C𝑖] = c
+            c += 1
         end
-        c += 1
     end
 
     # Remove ghost cells outside of original grid
     LUMP = sparse(LUMPidx[C][:], 1:length(C), 1)
 
     # Find wet points in coarsened grid
-    wet = wet3D[:]
     wet_c = LUMP * wet .> 0
 
     # Extract only indices of wet grd points
@@ -84,6 +108,18 @@ function lump_and_spray(wet3D, vol, mask=trues(size(wet3D)); di=2, dj=2, dk=1)
 end
 
 
+"""
+diconnected_cells(wet3D, T; di=2, dj=2, dk=1)
+
+"""
+function connected_cells(wet3D, T; di=2, dj=2, dk=1)
+    i, j = findnz(SPRAY * LUMP)
+    connected_by_grid = sparse(i, j, true, size(T)...)
+    i, j = findnz(T)
+    connected_by_T = sparse(i, j, true, size(T)...)
+    connected = connected_by_grid .& connected_by_T
+    connected = ((connected^di)^dj)^dk
+end
 
 
 

test/interactive.jl

@@ -120,3 +120,31 @@ end
 
 # @profview transportmatrix(; ϕ, mlotst, gridmetrics, indices, ρ, κH, κVML, κVdeep)
 # @profview_allocs transportmatrix(; ϕ, mlotst, gridmetrics, indices, ρ, κH, κVML, κVdeep) sample_rate=0.9
+
+
+# unpack model grid
+(; v3D, lat, lon, zt) = gridmetrics
+# unpack indices
+(; wet3D, N) = indices
+
+v = v3D[wet3D]
+
+@info "coarsening grid"
+SOmask = lat .< -35
+NAmask = @. (lat > 50) & ((lon < 100) | (250 < lon))
+mymask = repeat(.!SOmask .& .!NAmask, 1, 1, size(wet3D, 3))
+
+# Test extreme coarsening and plot it
+LUMP, SPRAY, v_c = OceanTransportMatrixBuilder.lump_and_spray(wet3D, v, T, mymask; di=10, dj=10, dk=1)
+# Plot the lumping
+using CairoMakie
+sprayedrand = SPRAY * rand(size(LUMP, 1))
+sprayedrand3D = OceanTransportMatrixBuilder.as3D(sprayedrand, wet3D)
+fig, ax, hm = surface(mod.(lon .- 80, 360) .+ 80, lat, zeros(size(lat));
+    color=sprayedrand3D[:,:,45],
+    colormap = :turbo,
+    shading = false
+)
+outputfile = joinpath(outputdir, "LUMP_and_SPRAY_coarsening.png")
+@info "Saving coarsening as image file:\n $(joinpath("test", outputfile))"
+save(outputfile, fig)

test/local_full.jl

@@ -119,6 +119,7 @@ end
     using Unitful: s, yr
     using NaNStatistics
     using CairoMakie
+    using Graphs
 
     (; gridmetrics, indices, T, model, member, outputdir) = LocalBuiltMatrix
 
@@ -133,8 +134,22 @@ end
     SOmask = lat .< -35
     NAmask = @. (lat > 50) & ((lon < 100) | (250 < lon))
     mymask = repeat(.!SOmask .& .!NAmask, 1, 1, size(wet3D, 3))
-    LUMP, SPRAY, v_c = OceanTransportMatrixBuilder.lump_and_spray(wet3D, v, mymask; di=2, dj=2, dk=1)
-    LUMP, SPRAY, v_c = OceanTransportMatrixBuilder.lump_and_spray(wet3D, v; di=2, dj=2, dk=1)
+
+    # Test extreme coarsening to show behavior across land and plot it
+    LUMP, SPRAY, v_c = OceanTransportMatrixBuilder.lump_and_spray(wet3D, v, T, mymask; di=10, dj=10, dk=1)
+    sprayedrand = SPRAY * rand(size(LUMP, 1))
+    sprayedrand3D = OceanTransportMatrixBuilder.as3D(sprayedrand, wet3D)
+    fig, ax, hm = surface(mod.(lon .- 80, 360) .+ 80, lat, zeros(size(lat));
+        color=sprayedrand3D[:,:,1],
+        colormap = :turbo,
+        shading = false
+    )
+    outputfile = joinpath(outputdir, "LUMP_and_SPRAY_coarsening.png")
+    @info "Saving coarsening as image file:\n $(joinpath("test", outputfile))"
+    save(outputfile, fig)
+
+    # Redo coarsening with less extreme coarsening
+    LUMP, SPRAY, v_c = OceanTransportMatrixBuilder.lump_and_spray(wet3D, v, T; di=2, dj=2, dk=1)
 
     # surface mask
     issrf3D = copy(wet3D)

test/online.jl

@@ -128,6 +128,6 @@
     SOmask = lat .< -35
     NAmask = @. (lat > 50) & ((lon < 100) | (250 < lon))
     mymask = repeat(.!SOmask .& .!NAmask, 1, 1, size(wet3D, 3))
-    LUMP, SPRAY, v_c = OceanTransportMatrixBuilder.lump_and_spray(wet3D, v, mymask; di=2, dj=2, dk=1)
+    LUMP, SPRAY, v_c = OceanTransportMatrixBuilder.lump_and_spray(wet3D, v, T, mymask; di=2, dj=2, dk=1)
 
 end