document smart lump and spray function

Author: briochemc

src/extratools.jl

@@ -1,8 +1,12 @@
 """
-    LUMP, SPRAY, v_c = lump_and_spray(wet3D, volume; f=Returns(true),di=2, dj=2, dk=1)
+    LUMP, SPRAY, v_c = lump_and_spray(wet3D, vol, T, mask=trues(size(wet3D)); di=2, dj=2, dk=1)
 
 returns the LUMP and SPRAY matrices, and the coarsened volume vector.
 
+The `vol` vector of volumes is used to make the LUMP operator volume-conserving.
+The `T` matrix is used to determine connectivity of the grid cells, so as to avoid lumping disconnected cells.
+The `mask` boolean array (same size as `wet3D`) can be used to restrict the lumping to a certain region.
+
 To get the coarsened vector from fine vector x, use
 
     LUMP * x
@@ -108,20 +112,6 @@ function lump_and_spray(wet3D, vol, T, 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
-
-
 
 """
 as2D(x, wet3D)

test/interactive.jl

@@ -122,19 +122,16 @@ end
 # @profview_allocs transportmatrix(; ϕ, mlotst, gridmetrics, indices, ρ, κH, κVML, κVdeep) sample_rate=0.9
 
 
+# Test extreme coarsening and plot it
 # 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

test/local_full.jl

@@ -137,11 +137,12 @@ end
 
     # 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))
+    colormap = cgrad(:okabe_ito, categorical=true)
+    sprayedrand = SPRAY * rand(1:length(colormap), 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,
+        colormap,
         shading = false
     )
     outputfile = joinpath(outputdir, "LUMP_and_SPRAY_coarsening.png")