diff --git a/Project.toml b/Project.toml
index 167892e..9fd69d9 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "OceanTransportMatrixBuilder"
 uuid = "c2b4a04e-6049-4fc4-aa6a-5508a29a1e1c"
 authors = ["Benoit Pasquier <briochemc@gmail.com> and contributors"]
-version = "0.2.8"
+version = "0.3.0"
 
 [deps]
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
@@ -18,6 +18,7 @@ julia = "1.10"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
+Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
 GibbsSeaWater = "9a22fb26-0b63-4589-b28e-8f9d0b5c3d05"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -32,4 +33,4 @@ YAXArrays = "c21b50f5-aa40-41ea-b809-c0f5e47bfa5c"
 Zarr = "0a941bbe-ad1d-11e8-39d9-ab76183a1d99"
 
 [targets]
-test = ["DimensionalData", "GibbsSeaWater", "GLMakie", "LinearAlgebra", "Makie", "NaNStatistics", "NetCDF", "Test", "TestItemRunner", "TestItems", "Unitful", "YAXArrays", "CSV", "DataFrames", "Zarr"]
+test = ["Downloads", "DimensionalData", "GibbsSeaWater", "GLMakie", "LinearAlgebra", "Makie", "NaNStatistics", "NetCDF", "Test", "TestItemRunner", "TestItems", "Unitful", "YAXArrays", "CSV", "DataFrames", "Zarr"]
diff --git a/src/OceanTransportMatrixBuilder.jl b/src/OceanTransportMatrixBuilder.jl
index 8a3825c..1e4eaa8 100644
--- a/src/OceanTransportMatrixBuilder.jl
+++ b/src/OceanTransportMatrixBuilder.jl
@@ -9,8 +9,8 @@ using Distances: haversine
 
 include("preprocessing.jl")
 include("matrixbuilding.jl")
-include("grids.jl")
-include("topology.jl")
+include("gridcellgeometry.jl")
+include("gridtopology.jl")
 include("derivatives.jl")
 include("extratools.jl") # <- I think this should be in a separate "base" repo
 
diff --git a/src/derivatives.jl b/src/derivatives.jl
index a582c3a..6e98fcd 100644
--- a/src/derivatives.jl
+++ b/src/derivatives.jl
@@ -12,11 +12,11 @@ Returns the discrete difference
 # forward derivative in the i direction for centered data (A-grid)
 function ∂ᵢ₊(χ, modelgrid)
     χ = χ |> Array
-    (; lon, lat, gridtype) = modelgrid
+    (; lon, lat, arakawagrid) = modelgrid
     ∂ᵢ₊χ = zeros(size(χ))
     for I in CartesianIndices(lon)
         P = (lon[I], lat[I])
-        Iᵢ₊₁ = i₊₁(I, gridtype)
+        Iᵢ₊₁ = i₊₁(I, arakawagrid)
         Pᵢ₊₁ = (lon[Iᵢ₊₁], lat[Iᵢ₊₁])
         d = haversine(P, Pᵢ₊₁)
         ∂ᵢ₊χ[I, :] = (χ[Iᵢ₊₁, :] - χ[I, :]) / d
@@ -26,11 +26,11 @@ end
 # backward derivative in the i direction for centered data (A-grid)
 function ∂ᵢ₋(χ, modelgrid)
     χ = χ |> Array
-    (; lon, lat, gridtype) = modelgrid
+    (; lon, lat, arakawagrid) = modelgrid
     ∂ᵢ₋χ = zeros(size(χ))
     for I in CartesianIndices(lon)
         P = (lon[I], lat[I])
-        Iᵢ₋₁ = i₋₁(I, gridtype)
+        Iᵢ₋₁ = i₋₁(I, arakawagrid)
         Pᵢ₋₁ = (lon[Iᵢ₋₁], lat[Iᵢ₋₁])
         d = haversine(P, Pᵢ₋₁)
         ∂ᵢ₋χ[I, :] = (χ[I, :] - χ[Iᵢ₋₁, :]) / d
@@ -40,11 +40,11 @@ end
 # forward derivative in the j direction
 function ∂ⱼ₊(χ, modelgrid)
     χ = χ |> Array
-    (; lon, lat, gridtype) = modelgrid
+    (; lon, lat, arakawagrid) = modelgrid
     ∂ⱼ₊χ = zeros(size(χ))
     for I in CartesianIndices(lon)
         P = (lon[I], lat[I])
-        Iᵢ₊₁ = j₊₁(I, gridtype)
+        Iᵢ₊₁ = j₊₁(I, arakawagrid)
         Pᵢ₊₁ = (lon[Iᵢ₊₁], lat[Iᵢ₊₁])
         d = haversine(P, Pᵢ₊₁)
         ∂ⱼ₊χ[I, :] = (χ[Iᵢ₊₁, :] - χ[I, :]) / d
@@ -54,11 +54,11 @@ end
 # backward derivative in the j direction
 function ∂ⱼ₋(χ, modelgrid)
     χ = χ |> Array
-    (; lon, lat, gridtype) = modelgrid
+    (; lon, lat, arakawagrid) = modelgrid
     ∂ⱼ₋χ = zeros(size(χ))
     for I in CartesianIndices(lon)
         P = (lon[I], lat[I])
-        Iᵢ₋₁ = j₋₁(I, gridtype)
+        Iᵢ₋₁ = j₋₁(I, arakawagrid)
         Pᵢ₋₁ = (lon[Iᵢ₋₁], lat[Iᵢ₋₁])
         d = haversine(P, Pᵢ₋₁)
         ∂ⱼ₋χ[I, :] = (χ[I, :] - χ[Iᵢ₋₁, :]) / d
@@ -68,12 +68,12 @@ end
 # derivative in the k direction
 function ∂ₖ₊(χ, modelgrid)
     χ = χ |> Array
-    (; zt, DZT3d, gridtype) = modelgrid
+    (; zt, thkcello, arakawagrid) = modelgrid
     ∂ₖ₊χ = zeros(size(χ))
     for I in CartesianIndices(χ)
-        Iᵢ₊₁ = k₊₁(I, gridtype)
+        Iᵢ₊₁ = k₊₁(I, arakawagrid)
         isnothing(Iᵢ₊₁) && continue
-        h = (DZT3d[I] + DZT3d[Iᵢ₊₁]) / 2
+        h = (thkcello[I] + thkcello[Iᵢ₊₁]) / 2
         ∂ₖ₊χ[I] = (χ[Iᵢ₊₁] - χ[I]) / h
     end
     return ∂ₖ₊χ
@@ -81,12 +81,12 @@ end
 # backward derivative in the k direction
 function ∂ₖ₋(χ, modelgrid)
     χ = χ |> Array
-    (; zt, DZT3d, gridtype) = modelgrid
+    (; zt, thkcello, arakawagrid) = modelgrid
     ∂ₖ₋χ = zeros(size(χ))
     for I in CartesianIndices(χ)
-        Iᵢ₋₁ = k₋₁(I, gridtype)
+        Iᵢ₋₁ = k₋₁(I, arakawagrid)
         isnothing(Iᵢ₋₁) && continue
-        h = (DZT3d[I] + DZT3d[Iᵢ₋₁]) / 2
+        h = (thkcello[I] + thkcello[Iᵢ₋₁]) / 2
         ∂ₖ₋χ[I] = (χ[I] - χ[Iᵢ₋₁]) / h
     end
     return ∂ₖ₋χ
@@ -94,49 +94,56 @@ end
 
 # interpolation in the i direction
 function itpₖ₊(χ, modelgrid)
-    (; gridtype) = modelgrid
+    (; arakawagrid) = modelgrid
     χ = χ |> Array
     itpₖ₊χ = zeros(size(χ))
     for I in CartesianIndices(χ)
-        Iᵢ₊₁ = k₊₁(I, gridtype)
+        Iᵢ₊₁ = k₊₁(I, arakawagrid)
         isnothing(Iᵢ₊₁) && continue
         itpₖ₊χ[I] = (χ[Iᵢ₊₁] + χ[I]) / 2
     end
     return itpₖ₊χ
 end
 function itpₖ₋(χ, modelgrid)
-    (; gridtype) = modelgrid
+    (; arakawagrid) = modelgrid
     χ = χ |> Array
     itpₖ₋χ = zeros(size(χ))
     for I in CartesianIndices(χ)
-        Iᵢ₋₁ = k₋₁(I, gridtype)
+        Iᵢ₋₁ = k₋₁(I, arakawagrid)
         isnothing(Iᵢ₋₁) && continue
         itpₖ₋χ[I] = (χ[Iᵢ₋₁] + χ[I]) / 2
     end
     return itpₖ₋χ
 end
 function itpᵢ₊(χ, modelgrid)
-    (; gridtype) = modelgrid
+    (; arakawagrid) = modelgrid
     χ = χ |> Array
     itpᵢ₊χ = zeros(size(χ))
     for I in CartesianIndices(χ)
-        Iᵢ₊₁ = i₊₁(I, gridtype)
+        Iᵢ₊₁ = i₊₁(I, arakawagrid)
         itpᵢ₊χ[I] = (χ[Iᵢ₊₁] + χ[I]) / 2
     end
     return itpᵢ₊χ
 end
 function itpⱼ₊(χ, modelgrid)
-    (; gridtype) = modelgrid
+    (; arakawagrid) = modelgrid
     χ = χ |> Array
     itpⱼ₊χ = zeros(size(χ))
     for I in CartesianIndices(χ)
-        Iᵢ₊₁ = j₊₁(I, gridtype)
+        Iᵢ₊₁ = j₊₁(I, arakawagrid)
         itpⱼ₊χ[I] = (χ[Iᵢ₊₁] + χ[I]) / 2
     end
     return itpⱼ₊χ
 end
 
+"""
+    bolus_GM_velocity(σ, modelgrid; κGM = 600, maxslope = 0.01)
+
+Returns the bolus velocity field due to the Gent-McWilliams parameterization,
+computed from the neutral density field `σ` (or potential density, ρθ in kg/m³).
 
+Note: This is experimental at this stage.
+"""
 function bolus_GM_velocity(σ, modelgrid; κGM = 600, maxslope = 0.01)
     # σ is neutral density (or potential density, ρθ in kg/m³)
     σ = replace(σ, missing => 0, NaN => 0)
@@ -193,13 +200,13 @@ end
 
 
 # function ∂ᵢ(χ, modelgrid, shift=0,  d=:f)
-#     (;lon, lat, gridtype) = modelgrid
+#     (;lon, lat, arakawagrid) = modelgrid
 #     m = 1
 #     ∂ᵢχ = zeros(size(χ))
 #     # Draw a line along coordinate i and compute the distances
 #     d = dᵢ₊(lon, lat)
 #     for I in CartesianIndices(lon)
-#         movingI = SVector(ishift(I, gridtype, xxx))
+#         movingI = SVector(ishift(I, arakawagrid, xxx))
 #         i, j = I.I
 #         x = SVector(0, d[i])
 #         w = stencil(x, 0, χ)
diff --git a/src/grids.jl b/src/gridcellgeometry.jl
similarity index 55%
rename from src/grids.jl
rename to src/gridcellgeometry.jl
index c0c76ee..01aff6f 100644
--- a/src/grids.jl
+++ b/src/gridcellgeometry.jl
@@ -1,19 +1,25 @@
 
 
-abstract type ArakawaGrid end
+abstract type ArakawaGridCell end
 
-struct AGrid <: ArakawaGrid
+struct AGridCell <: ArakawaGridCell
+    u_pos::Symbol
+    v_pos::Symbol
 end
 
-struct BGrid <: ArakawaGrid
+struct BGridCell <: ArakawaGridCell
+    u_pos::Symbol
+    v_pos::Symbol
 end
 
-struct CGrid <: ArakawaGrid
+struct CGridCell <: ArakawaGridCell
+    u_pos::Symbol
+    v_pos::Symbol
 end
 
 
 """
-    arakawa, uo_pos, vo_pos, uo_relerr, vo_relerr = gridtype(uo_lon, uo_lat, vo_lon, vo_lat, modelgrid)
+    arakawa = getarakawagrid(u_lon, u_lat, v_lon, v_lat, modelgrid)
 
 Returns the type of the grid (A, B, or C), the grid position of the velocity points,
 and the error of that position relative to the perimeter of the cell.
@@ -43,14 +49,14 @@ The different Arakawa grids recongnized here are:
 where each grid is centered in C.
 Also returns the distances from the velocity points to the grid points.
 """
-function gridtype(uo_lon, uo_lat, vo_lon, vo_lat, modelgrid)
+function getarakawagrid(u_lon, u_lat, v_lon, v_lat, modelgrid)
 
     # Unpack modelgrid
     (; lon, lat, lon_vertices, lat_vertices) = modelgrid
 
     i = j = 1
-    uo_point = (uo_lon[i, j], uo_lat[i, j])
-    vo_point = (vo_lon[i, j], vo_lat[i, j])
+    u_point = (u_lon[i, j], u_lat[i, j])
+    v_point = (v_lon[i, j], v_lat[i, j])
 
     C = (lon[i, j], lat[i, j])
     SW = (lon_vertices[1, i, j], lat_vertices[1, i, j])
@@ -64,75 +70,74 @@ function gridtype(uo_lon, uo_lat, vo_lon, vo_lat, modelgrid)
 
     cell = (; C, SW, SE, NE, NW, S, N, W, E)
 
-    uo_distances = (; (k => haversine(P, uo_point) for (k,P) in pairs(cell))...)
-    vo_distances = (; (k => haversine(P, vo_point) for (k,P) in pairs(cell))...)
+    u_distances = (; (k => haversine(P, u_point) for (k,P) in pairs(cell))...)
+    v_distances = (; (k => haversine(P, v_point) for (k,P) in pairs(cell))...)
 
-    uo_distance, uo_pos = findmin(uo_distances)
-    vo_distance, vo_pos = findmin(vo_distances)
+    u_distance, u_pos = findmin(u_distances)
+    v_distance, v_pos = findmin(v_distances)
 
     # Arakawa grid type
-    if uo_pos == vo_pos == :C
-        arakawa = AGrid()
-    elseif uo_pos == vo_pos && uo_pos ∈ (:NE, :NW, :SE, :SW)
-        arakawa = BGrid()
-    elseif uo_pos ∈ (:E, :W) && vo_pos ∈ (:N, :S)
-        arakawa = CGrid()
+    if u_pos == v_pos == :C
+        arakawa = AGridCell(u_pos, v_pos)
+    elseif u_pos == v_pos && u_pos ∈ (:NE, :NW, :SE, :SW)
+        arakawa = BGridCell(u_pos, v_pos)
+    elseif u_pos ∈ (:E, :W) && v_pos ∈ (:N, :S)
+        arakawa = CGridCell(u_pos, v_pos)
     else
         error("Unknown Arakawa grid type")
     end
 
     cellperimeter = haversine(SW, SE) + haversine(SE, NE) + haversine(NE, NW) + haversine(NW, SW)
-    uo_relerr = uo_distance / cellperimeter
-    vo_relerr = vo_distance / cellperimeter
+    relerr = (u_distance + v_distance) / cellperimeter
 
-    return arakawa, uo_pos, vo_pos, uo_relerr, vo_relerr
+    relerr > 0.01 && warn("Relative error in grid positions in $arakawa is $relerr")
+
+    return arakawa
 
 end
 
 """
-    uo, uo_lon, uo_lat, vo, vo_lon, vo_lat = interpolateontodefaultCgrid(uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, modelgrid)
+    u, u_lon, u_lat, v, v_lon, v_lat = interpolateontodefaultCgrid(u, u_lon, u_lat, v, v_lon, v_lat, modelgrid)
 
-Interpolates the velocity fields `uo` and `vo` from B- or C-grid
+Interpolates the velocity fields `u` and `v` from B- or C-grid
 onto the default C-grid (centered on the cell faces).
 """
-interpolateontodefaultCgrid(uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, modelgrid) = interpolateontodefaultCgrid(uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, modelgrid, gridtype(uo_lon, uo_lat, vo_lon, vo_lat, modelgrid))
-interpolateontodefaultCgrid(uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, modelgrid, ::CGrid) = uo, uo_lon, uo_lat, vo, vo_lon, vo_lat
-interpolateontodefaultCgrid(uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, modelgrid, ::AGrid) = error("Interpolation not implemented for this grid type")
-# TODO this is clumsy to have to pass the gridtype but then recompute it
-function interpolateontodefaultCgrid(uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, modelgrid, arakawa::BGrid)
-
-    arakawa, uo_pos, vo_pos, uo_relerr, vo_relerr = gridtype(uo_lon, uo_lat, vo_lon, vo_lat, modelgrid)
+interpolateontodefaultCgrid(u, u_lon, u_lat, v, v_lon, v_lat, modelgrid) = interpolateontodefaultCgrid(u, u_lon, u_lat, v, v_lon, v_lat, modelgrid, getarakawagrid(u_lon, u_lat, v_lon, v_lat, modelgrid))
+interpolateontodefaultCgrid(u, u_lon, u_lat, v, v_lon, v_lat, modelgrid, ::CGridCell) = u, u_lon, u_lat, v, v_lon, v_lat
+interpolateontodefaultCgrid(u, u_lon, u_lat, v, v_lon, v_lat, modelgrid, ::AGridCell) = error("Interpolation not implemented for A-grid type")
+function interpolateontodefaultCgrid(u, u_lon, u_lat, v, v_lon, v_lat, modelgrid, arakawa::BGridCell)
 
-    uo_pos == vo_pos == :NE || error("Interpolation not implemented for this B-grid type")
+    (; u_pos, v_pos) = arakawa
+    u_pos == v_pos == :NE || error("Interpolation not implemented for this B-grid($u_pos,$v_pos) type")
 
-    _FillValue = uo.properties["_FillValue"]
+    _FillValue = u.properties["_FillValue"]
 
-    # Make sure uo and vo are in memory
-    uo = uo |> Array
-    vo = vo |> Array
+    # Make sure u and v are in memory
+    u = u |> Array
+    v = v |> Array
 
-    nx, ny, nz = size(uo)
+    nx, ny, nz = size(u)
 
     # unpack modelgrid
     (; lon_vertices, lat_vertices) = modelgrid
 
-    # It seems that uo/vo is NaN on boundaries (for ACCESS-ESM1-5)
-    # and that umo/vmo were computed as if uo/vo were 0 on boundaries
+    # It seems that u/v is NaN on boundaries (for ACCESS-ESM1-5)
+    # and that umo/vmo were computed as if u/v were 0 on boundaries
     # so that's what we do here
-    uo2 = replace(uo, _FillValue => 0.0)
-    vo2 = replace(vo, _FillValue => 0.0)
-    uo2 = 0.5(uo2 + [fill(0.0, nx, 1, nz);; uo2[:, 1:end-1, :]])
-    vo2 = 0.5(vo2 + [fill(0.0, 1, ny, nz); vo2[1:end-1, :, :]])
+    u2 = replace(u, _FillValue => 0.0)
+    v2 = replace(v, _FillValue => 0.0)
+    u2 = 0.5(u2 + [fill(0.0, nx, 1, nz);; u2[:, 1:end-1, :]])
+    v2 = 0.5(v2 + [fill(0.0, 1, ny, nz); v2[1:end-1, :, :]])
     SE_points = [(lon, lat) for (lon, lat) in zip(lon_vertices[2, :, :], lat_vertices[2, :, :])]
     NE_points = [(lon, lat) for (lon, lat) in zip(lon_vertices[3, :, :], lat_vertices[3, :, :])]
     NW_points = [(lon, lat) for (lon, lat) in zip(lon_vertices[4, :, :], lat_vertices[4, :, :])]
-    uo2_points = [midpointonsphere(SE, NE) for (SE, NE) in zip(NE_points, SE_points)]
-    vo2_points = [midpointonsphere(NE, NW) for (NE, NW) in zip(NW_points, NE_points)]
-    uo2_lon = [P[1] for P in uo2_points]
-    uo2_lat = [P[2] for P in uo2_points]
-    vo2_lon = [P[1] for P in vo2_points]
-    vo2_lat = [P[2] for P in vo2_points]
-    return uo2, uo2_lon, uo2_lat, vo2, vo2_lon, vo2_lat
+    u2_points = [midpointonsphere(SE, NE) for (SE, NE) in zip(NE_points, SE_points)]
+    v2_points = [midpointonsphere(NE, NW) for (NE, NW) in zip(NW_points, NE_points)]
+    u2_lon = [P[1] for P in u2_points]
+    u2_lat = [P[2] for P in u2_points]
+    v2_lon = [P[1] for P in v2_points]
+    v2_lat = [P[2] for P in v2_points]
+    return u2, u2_lon, u2_lat, v2, v2_lon, v2_lat
 
 end
 
diff --git a/src/gridtopology.jl b/src/gridtopology.jl
new file mode 100644
index 0000000..fffcdae
--- /dev/null
+++ b/src/gridtopology.jl
@@ -0,0 +1,104 @@
+abstract type AbstractGridTopology end
+struct BipolarGridTopology <: AbstractGridTopology
+    nx::Int64
+    ny::Int64
+    nz::Int64
+end
+struct TripolarGridTopology <: AbstractGridTopology
+    nx::Int64
+    ny::Int64
+    nz::Int64
+end
+struct UnknownGridTopology <: AbstractGridTopology
+    nx::Int64
+    ny::Int64
+    nz::Int64
+end
+
+"""
+    getgridtopology(lon_vertices, lat_vertices, lev)
+
+Returns the type of grid based on how the vertices connect at the north pole.
+"""
+function getgridtopology(lon_vertices, lat_vertices, lev)
+    nx = size(lon_vertices, 2)
+    ny = size(lon_vertices, 3)
+    nz = length(lev)
+    # North pole vertices
+    NPlon = @view lon_vertices[3:4,:,end]
+    NPlat = @view lat_vertices[3:4,:,end]
+    # If all the latitudes of the "northmost" vertices are 90, then it's a "regular" grid with 2 poles
+    if all(NPlat .== 90)
+        return BipolarGridTopology(nx,ny,nz)
+    # Otherwise check if the north pole is split in two
+    elseif NPlon == rot180(NPlon) && NPlat == rot180(NPlat)
+        return TripolarGridTopology(nx,ny,nz)
+    else
+        return UnknownGridTopology(nx,ny,nz)
+    end
+end
+
+# Default behavior for grids (assumed bipolar)
+# wrap around in the i direction (longitude)
+i₊₁(C, g::AbstractGridTopology) = C.I[1] < g.nx ? C + CartesianIndex(1, 0, 0) : CartesianIndex(1, C.I[2], C.I[3])
+i₋₁(C, g::AbstractGridTopology) = C.I[1] > 1 ? C + CartesianIndex(-1, 0, 0) : CartesianIndex(g.nx, C.I[2], C.I[3])
+i₊₁(C::CartesianIndex{2}, g::AbstractGridTopology) = C.I[1] < g.nx ? C + CartesianIndex(1, 0) : CartesianIndex(1, C.I[2])
+i₋₁(C::CartesianIndex{2}, g::AbstractGridTopology) = C.I[1] > 1 ? C + CartesianIndex(-1, 0) : CartesianIndex(g.nx, C.I[2])
+# No connection by default in the j direction (latitude)
+j₊₁(C, g::AbstractGridTopology) = C.I[2] < g.ny ? C + CartesianIndex(0, 1, 0) : nothing
+j₋₁(C, ::AbstractGridTopology) = C.I[2] > 1 ? C + CartesianIndex(0, -1, 0) : nothing
+j₊₁(C::CartesianIndex{2}, g::AbstractGridTopology) = C.I[2] < g.ny ? C + CartesianIndex(0, 1) : nothing
+j₋₁(C::CartesianIndex{2}, ::AbstractGridTopology) = C.I[2] > 1 ? C + CartesianIndex(0, -1) : nothing
+# No connection by default in the k direction (depth)
+k₊₁(C, g::AbstractGridTopology) = C.I[3] < g.nz ? C + CartesianIndex(0, 0, 1) : nothing
+k₋₁(C, ::AbstractGridTopology) = C.I[3] > 1 ? C + CartesianIndex(0, 0, -1) : nothing
+
+
+
+function ishift(C, g::AbstractGridTopology, n=0)
+    i, j, k = C.I
+    CartesianIndex(mod1(i + n, g.nx), j, k)
+end
+function jshift(C, g::AbstractGridTopology, n=0)
+    i, j, k = C.I
+    j2 = j + n
+    (1 ≤ j2 ≤ g.ny) ? CartesianIndex(i, j2, k) : nothing
+end
+function kshift(C, g::AbstractGridTopology, n=0)
+    i, j, k = C.I
+    k2 = k + n
+    (1 ≤ k2 ≤ g.nz) ? CartesianIndex(i, j, k2) : nothing
+end
+
+# Special behavior for tripolar grids where the seam connects the top and "folds" it around
+# ┌─────┬─────┬─────┬─────┬─────┬─────┐
+# │  6  │  5  │  4  │  3  │  2  │  1  │
+# ├─────┼─────┼─────┼─────┼─────┼─────┤ ◄─ seam
+# │  1  │  2  │  3  │  4  │  5  │  6  │
+# └─────┴─────┴─────┴─────┴─────┴─────┘
+#            ─► increasing i ("east")
+j₊₁(C, g::TripolarGridTopology) = C.I[2] < g.ny ? C + CartesianIndex(0, 1, 0) : CartesianIndex(g.nx - C.I[1] + 1, g.ny, C.I[3])
+j₊₁(C::CartesianIndex{2}, g::TripolarGridTopology) = C.I[2] < g.ny ? C + CartesianIndex(0, 1) : CartesianIndex(g.nx - C.I[1] + 1, g.ny)
+
+function jshift(C, g::TripolarGridTopology, n=0)
+    i, j, k = C.I
+    j2 = j + n
+    if !(1 ≤ j2)
+        return nothing
+    elseif j2 ≤ g.ny
+        return CartesianIndex(i, j2, k)
+    else
+        i2 = g.nx - i + 1
+        return CartesianIndex(i2, j, k)
+    end
+end
+
+# error for unknown grids
+i₊₁(C, ::UnknownGridTopology) = error("Unknown grid type")
+i₋₁(C, ::UnknownGridTopology) = error("Unknown grid type")
+j₊₁(C, ::UnknownGridTopology) = error("Unknown grid type")
+j₋₁(C, ::UnknownGridTopology) = error("Unknown grid type")
+k₊₁(C, ::UnknownGridTopology) = error("Unknown grid type")
+k₋₁(C, ::UnknownGridTopology) = error("Unknown grid type")
+
+
diff --git a/src/matrixbuilding.jl b/src/matrixbuilding.jl
index 860f065..c42467e 100644
--- a/src/matrixbuilding.jl
+++ b/src/matrixbuilding.jl
@@ -8,10 +8,12 @@ function buildTadv(; ϕ, modelgrid, indices, ρ)
     # default ρ = 1035 kg/m^3 is the value originally used by Chamberlain et al. (2019)
 
 	@info "Building Tadv"
-	𝑖s, 𝑗s, Tvals = upwind_advection_operator_sparse_entries(; ϕ, modelgrid, indices, ρ)
+	𝑖s, 𝑗s, Tvals = upwind_advection_operator_sparse_entries(ϕ, modelgrid, indices, ρ)
 
     N = indices.N
 
+	any(isnan.(Tvals)) && error("Tadv contains NaNs.")
+
 	Tadv = sparse(𝑖s, 𝑗s, Tvals, N, N)
 
 	return Tadv
@@ -32,6 +34,8 @@ function buildTκH(; modelgrid, indices, ρ, κH)
 	@info "Building TκH"
 	𝑖s, 𝑗s, Tvals = horizontal_diffusion_operator_sparse_entries(; modelgrid, indices, κH, ΩH)
 
+	any(isnan.(Tvals)) && error("TκH contains NaNs.")
+
 	TκH = sparse(𝑖s, 𝑗s, Tvals, N, N)
 
 	return TκH
@@ -49,7 +53,7 @@ function buildTκVML(; mlotst, modelgrid, indices, κVML)
     (; zt, ) = modelgrid
 
     # Unpack indices
-    (; Lwet) = indices
+    (; Lwet, N) = indices
 
 	mlotst = mlotst |> Array # to prevent slow getindex for lazily loaded data?
 
@@ -59,7 +63,7 @@ function buildTκVML(; mlotst, modelgrid, indices, κVML)
 	@info "Building TκVML "
 	𝑖s, 𝑗s, Tvals = vertical_diffusion_operator_sparse_entries(; modelgrid, indices, κV = κVML, Ω)
 
-    N = indices.N
+	any(isnan.(Tvals)) && error("TκVML contains NaNs.")
 
 	TκVML = sparse(𝑖s, 𝑗s, Tvals, N, N)
 
@@ -82,6 +86,8 @@ function buildTκVdeep(; mlotst, modelgrid, indices, κVdeep)
 	@info "Building TκVdeep"
 	𝑖s, 𝑗s, Tvals = vertical_diffusion_operator_sparse_entries(; modelgrid, indices, κV = κVdeep, Ω)
 
+	any(isnan.(Tvals)) && error("TκVdeep contains NaNs.")
+
 	TκVdeep = sparse(𝑖s, 𝑗s, Tvals, N, N)
 
 	return TκVdeep
@@ -129,6 +135,15 @@ end
 # 	T[i,j] = -ϕ[j→i] / m[i]             units = kg s⁻¹ / kg = s⁻¹
 # and ϕ[j→i] / m[j] should be added to the diagonal T[j,j].
 
+# Is this mass conserving?
+# Only if
+#   v[i] * T[i,i] + v[j] * T[j,i] ≈ 0
+#   = vi ϕ[i→j] / m[i] + vj -ϕ[i→j] / m[j]
+#   = ϕ[i→j] / ρ[i] - ϕ[i→j] / ρ[j]
+# i.e., iff
+#   ρ[i] = ρ[j]
+# So I must use the mean density between facing cells.
+
 # There are three ways to index:
 # - Cartesian indices (i,j,k)
 # - Linear index L𝑖
@@ -139,74 +154,85 @@ end
 # For the last one, I make a 3D array filled with the wet linear indices
 
 """
-    upwind_advection_operator_sparse_entries(; ϕ, modelgrid, indices, ρ)
+    upwind_advection_operator_sparse_entries(ϕ, modelgrid, indices, ρ)
 
 Return the sparse (i, j, v) for the upwind advection operator Tadv.
 """
-function upwind_advection_operator_sparse_entries(; ϕ, modelgrid, indices, ρ)
+function upwind_advection_operator_sparse_entries(ϕ, modelgrid, indices, ρ::Number)
+	# If ρ is a scalar, broadcast it to the modelgrid size
+	ρ = fill(ρ, size(modelgrid.v3D))
+	return upwind_advection_operator_sparse_entries(ϕ, modelgrid, indices, ρ)
+end
+function upwind_advection_operator_sparse_entries(ϕ, modelgrid, indices, ρ)
 
     # Unpack model grid
-    (; v3D, gridtype) = modelgrid
+    (; v3D, arakawagrid) = modelgrid
     # Unpack indices
-    (; wet3D, Lwet, Lwet3D, C) = indices
+    (; Lwet, Lwet3D, C) = indices
 
+	any(isnan.(ρ[Lwet])) && error("ρ contains NaNs")
 
 	𝑖s, 𝑗s, Tvals = Int[], Int[], Float64[]
-	nxyz = size(wet3D)
-    nx, ny, _ = nxyz
 
     @time for 𝑖 in eachindex(Lwet)
 		L𝑖 = Lwet[𝑖]
 		C𝑖 = C[L𝑖]
 		i, j, k = C𝑖.I
-		m𝑖 = v3D[C𝑖] * ρ
+		v𝑖 = v3D[C𝑖]
+		ρ𝑖 = ρ[C𝑖]
 		# From West
 		ϕwest = ϕ.west[C𝑖]
 		if ϕwest > 0
-			C𝑗 = i₋₁(C𝑖, gridtype)
+			C𝑗 = i₋₁(C𝑖, arakawagrid)
 			𝑗 = Lwet3D[C𝑗]
-			m𝑗 = v3D[C𝑗] * ρ
-			pushTadvectionvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, ϕwest, m𝑖, m𝑗)
+			v𝑗 = v3D[C𝑗]
+			ρ𝑗 = ρ[C𝑗]
+			pushTadvectionvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, ϕwest, ρ𝑖, ρ𝑗, v𝑖, v𝑗)
 		end
 		# From East
 		ϕeast = ϕ.east[C𝑖]
 		if ϕeast < 0
-			C𝑗 = i₊₁(C𝑖, gridtype)
+			C𝑗 = i₊₁(C𝑖, arakawagrid)
 			𝑗 = Lwet3D[C𝑗]
-			m𝑗 = v3D[C𝑗] * ρ
-			pushTadvectionvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, -ϕeast, m𝑖, m𝑗)
+			v𝑗 = v3D[C𝑗]
+			ρ𝑗 = ρ[C𝑗]
+			pushTadvectionvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, -ϕeast, ρ𝑖, ρ𝑗, v𝑖, v𝑗)
 		end
 		# From South
 		ϕsouth = ϕ.south[C𝑖]
 		if ϕsouth > 0
-			C𝑗 = j₋₁(C𝑖, gridtype)
+			C𝑗 = j₋₁(C𝑖, arakawagrid)
 			𝑗 = Lwet3D[C𝑗]
-			m𝑗 = v3D[C𝑗] * ρ
-			pushTadvectionvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, ϕsouth, m𝑖, m𝑗)
+			v𝑗 = v3D[C𝑗]
+			ρ𝑗 = ρ[C𝑗]
+			pushTadvectionvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, ϕsouth, ρ𝑖, ρ𝑗, v𝑖, v𝑗)
 		end
 		# From North (Special case with north bipole)
 		ϕnorth = ϕ.north[C𝑖]
 		if ϕnorth < 0
-			C𝑗 = j₊₁(C𝑖, gridtype)
+			C𝑗 = j₊₁(C𝑖, arakawagrid)
 			𝑗 = Lwet3D[C𝑗]
-			m𝑗 = v3D[C𝑗] * ρ
-			pushTadvectionvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, -ϕnorth, m𝑖, m𝑗)
+			v𝑗 = v3D[C𝑗]
+			ρ𝑗 = ρ[C𝑗]
+			pushTadvectionvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, -ϕnorth, ρ𝑖, ρ𝑗, v𝑖, v𝑗)
 		end
 		# From Bottom
 		ϕbottom = ϕ.bottom[C𝑖]
 		if ϕbottom > 0
-			C𝑗 = k₊₁(C𝑖, gridtype)
+			C𝑗 = k₊₁(C𝑖, arakawagrid)
 			𝑗 = Lwet3D[C𝑗]
-			m𝑗 = v3D[C𝑗] * ρ
-			pushTadvectionvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, ϕbottom, m𝑖, m𝑗)
+			v𝑗 = v3D[C𝑗]
+			ρ𝑗 = ρ[C𝑗]
+			pushTadvectionvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, ϕbottom, ρ𝑖, ρ𝑗, v𝑖, v𝑗)
 		end
 		# From Top
 		ϕtop = ϕ.top[C𝑖]
 		if ϕtop < 0 && k > 1 # Evaporation/precipitation -> no change to χ
-			C𝑗 = k₋₁(C𝑖, gridtype)
+			C𝑗 = k₋₁(C𝑖, arakawagrid)
 			𝑗 = Lwet3D[C𝑗]
-			m𝑗 = v3D[C𝑗] * ρ
-			pushTadvectionvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, -ϕtop, m𝑖, m𝑗)
+			v𝑗 = v3D[C𝑗]
+			ρ𝑗 = ρ[C𝑗]
+			pushTadvectionvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, -ϕtop, ρ𝑖, ρ𝑗, v𝑖, v𝑗)
 		end
 	end
 	return 𝑖s, 𝑗s, Tvals
@@ -214,11 +240,14 @@ end
 
 
 """
-    pushTadvectionvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, ϕ, m𝑖, m𝑗)
+    pushTadvectionvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, ϕ, ρ𝑖, ρ𝑗, v𝑖, v𝑗)
 
 Pushes the sparse indices and values into (𝑖s, 𝑗s, Tvals) corresponding to the j→i advection.
 """
-function pushTadvectionvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, ϕ, m𝑖, m𝑗)
+function pushTadvectionvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, ϕ, ρ𝑖, ρ𝑗, v𝑖, v𝑗)
+	ρ = (ρ𝑖 + ρ𝑗) / 2
+	m𝑖 = ρ * v𝑖
+	m𝑗 = ρ * v𝑗
 	push!(𝑖s, 𝑖)
 	push!(𝑗s, 𝑗)
 	push!(Tvals, -ϕ / m𝑖)
@@ -270,13 +299,13 @@ Return the sparse (i, j, v) for the horizontal diffusion operator TκH.
 function horizontal_diffusion_operator_sparse_entries(; modelgrid, indices, κH, ΩH)
 
     # Unpack model grid
-    (; v3D, edge_length_2D, lon, lat, DZT3d, gridtype) = modelgrid
+    (; v3D, edge_length_2D, lon, lat, thkcello, arakawagrid) = modelgrid
     # Unpack indices
     (; wet3D, Lwet, Lwet3D, C) = indices
 
 	𝑖s, 𝑗s, Tvals = Int[], Int[], Float64[]
-	nxyz = size(wet3D)
-    nx, ny, _ = nxyz
+
+    ny = size(wet3D, 2) # Should not be needed once oppdir is dealt by topology functions
 
     @time for 𝑖 in eachindex(Lwet)
         ΩH[𝑖] || continue # only continue if inside ΩH
@@ -285,15 +314,15 @@ function horizontal_diffusion_operator_sparse_entries(; modelgrid, indices, κH,
 		i, j, k = C𝑖.I
 		V = v3D[C𝑖]
 		# From West
-		C𝑗W = i₋₁(C𝑖, gridtype)
+		C𝑗W = i₋₁(C𝑖, arakawagrid)
         if !isnothing(C𝑗W)
 			𝑗W = Lwet3D[C𝑗W]
 			if !ismissing(𝑗W) && ΩH[𝑗W]
 				iW, jW, _ = C𝑗W.I
 				# (𝑖 == 𝑗W) && @show(i, j, iW, jW)
 				# I take the minimum area from both dirs (through which mixing goes through)
-				aij = verticalfacearea(edge_length_2D, DZT3d, i, j, k, :west)
-				aji = verticalfacearea(edge_length_2D, DZT3d, iW, jW, k, :east)
+				aij = verticalfacearea(edge_length_2D, thkcello, i, j, k, :west)
+				aji = verticalfacearea(edge_length_2D, thkcello, iW, jW, k, :east)
 				a = min(aij, aji)
 				# I take the mean distance from both dirs
 				d = horizontalcentroiddistance(lon, lat, i, j, iW, jW)
@@ -301,35 +330,35 @@ function horizontal_diffusion_operator_sparse_entries(; modelgrid, indices, κH,
 			end
 		end
         # From East
-		C𝑗E = i₊₁(C𝑖, gridtype)
+		C𝑗E = i₊₁(C𝑖, arakawagrid)
         if !isnothing(C𝑗E)
 			𝑗E = Lwet3D[C𝑗E]
 			if !ismissing(𝑗E) && ΩH[𝑗E]
 				iE, jE, _ = C𝑗E.I
 				# (𝑖 == 𝑗E) && @show(i, j, iE, jE)
-				aij = verticalfacearea(edge_length_2D, DZT3d, i, j, k, :east)
-				aji = verticalfacearea(edge_length_2D, DZT3d, iE, jE, k, :west)
+				aij = verticalfacearea(edge_length_2D, thkcello, i, j, k, :east)
+				aji = verticalfacearea(edge_length_2D, thkcello, iE, jE, k, :west)
 				a = min(aij, aji)
 				d = horizontalcentroiddistance(lon, lat, i, j, iE, jE)
 				pushTmixingvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗E, κH, a, d, V)
 			end
 		end
         # From South
-		C𝑗S = j₋₁(C𝑖, gridtype)
+		C𝑗S = j₋₁(C𝑖, arakawagrid)
 		if !isnothing(C𝑗S)
 			𝑗S = Lwet3D[C𝑗S]
 			if !ismissing(𝑗S) && ΩH[𝑗S]
 				iS, jS, _ = C𝑗S.I
 				# (𝑖 == 𝑗S) && @show(i, j, iS, jS)
-				aij = verticalfacearea(edge_length_2D, DZT3d, i, j, k, :south)
-				aji = verticalfacearea(edge_length_2D, DZT3d, iS, jS, k, :north)
+				aij = verticalfacearea(edge_length_2D, thkcello, i, j, k, :south)
+				aji = verticalfacearea(edge_length_2D, thkcello, iS, jS, k, :north)
 				a = min(aij, aji)
 				d = horizontalcentroiddistance(lon, lat, i, j, iS, jS)
 				pushTmixingvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗S, κH, a, d, V)
 			end
 		end
         # From North
-		C𝑗N = j₊₁(C𝑖, gridtype)
+		C𝑗N = j₊₁(C𝑖, arakawagrid)
         if !isnothing(C𝑗N)
 			𝑗N = Lwet3D[C𝑗N]
 			if !ismissing(𝑗N) && ΩH[𝑗N]
@@ -338,8 +367,8 @@ function horizontal_diffusion_operator_sparse_entries(; modelgrid, indices, κH,
 				# Note that the opposite direction (oppdir) is still north at j == ny
 				# TODO: implement this into a topology.jl function
 				oppdir = (j == ny) ? :north : :south
-				aij = verticalfacearea(edge_length_2D, DZT3d, i, j, k, :north)
-				aji = verticalfacearea(edge_length_2D, DZT3d, iN, jN, k, oppdir)
+				aij = verticalfacearea(edge_length_2D, thkcello, i, j, k, :north)
+				aji = verticalfacearea(edge_length_2D, thkcello, iN, jN, k, oppdir)
 				a = min(aij, aji)
 				d = horizontalcentroiddistance(lon, lat, i, j, iN, jN)
 				pushTmixingvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗N, κH, a, d, V)
@@ -377,7 +406,7 @@ end
 function vertical_diffusion_operator_sparse_entries(; modelgrid, indices, κV, Ω)
 
     # Unpack model grid
-    (; v3D, area2D, zt, gridtype) = modelgrid
+    (; v3D, area2D, zt, arakawagrid) = modelgrid
     # Unpack indices
     (; wet3D, Lwet, Lwet3D, C) = indices
 
@@ -393,7 +422,7 @@ function vertical_diffusion_operator_sparse_entries(; modelgrid, indices, κV, 
 		V = v3D[C𝑖]
         a = area2D[i,j]
 		# From Bottom
-		C𝑗B = k₊₁(C𝑖, gridtype)
+		C𝑗B = k₊₁(C𝑖, arakawagrid)
 		if !isnothing(C𝑗B)
 			𝑗B = Lwet3D[C𝑗B]
 			if !ismissing(𝑗B) && Ω[𝑗B] # only continue if inside Ω
@@ -403,7 +432,7 @@ function vertical_diffusion_operator_sparse_entries(; modelgrid, indices, κV, 
 			end
 		end
 		# From Top
-		C𝑗T = k₋₁(C𝑖, gridtype)
+		C𝑗T = k₋₁(C𝑖, arakawagrid)
 		if !isnothing(C𝑗T)
 			𝑗T = Lwet3D[C𝑗T]
 			if !ismissing(𝑗T) && Ω[𝑗T] # only continue if inside Ω
diff --git a/src/preprocessing.jl b/src/preprocessing.jl
index 7fd1e04..02c4503 100644
--- a/src/preprocessing.jl
+++ b/src/preprocessing.jl
@@ -1,29 +1,77 @@
 
 """
-    velocity2fluxes(uo_ds, vo_ds; modelgrid, ρ)
+    velocity2fluxes(u, u_lon, u_lat, v, v_lon, v_lat, modelgrid, ρ)
 
-Return the fluxes (integrated over cell faces) given veloticies `uo_ds` and `vo_ds`.
+Return the fluxes (integrated over cell faces) given veloticies `u` and `v` and their (lon,lat).
 
-The datasets are required to access the longitude and latitude of the velocity points.
+The fluxes are calculated from the density ρ (number or 3D array) and from the cell face areas between the cells.
+The mean density from the two cells that share the face is used.
+The minimum thickness of the two cells that share the face is used.
 """
-function velocity2fluxes(; uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, modelgrid, ρ)
+function velocity2fluxes(u, u_lon, u_lat, v, v_lon, v_lat, modelgrid, ρ)
 
     # Unpack modelgrid
-    (; DZT3d, edge_length_2D) = modelgrid
+    (; thkcello, edge_length_2D, v3D, lon_vertices, lat_vertices, zt) = modelgrid
+
+    # make indices
+    indices = makeindices(v3D)
 
     # Interpolate to C-grid
-    uo, _, _, vo, _, _ = interpolateontodefaultCgrid(uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, modelgrid)
+    u, _, _, v, _, _ = interpolateontodefaultCgrid(u, u_lon, u_lat, v, v_lon, v_lat, modelgrid)
+
+    # grid type
+    gridtopology = getgridtopology(lon_vertices, lat_vertices, zt)
+
+    ϕᵢ = zeros(size(u))
+    ϕⱼ = zeros(size(v))
 
     # Calculate fluxes from velocities
-    umo = @. uo * ρ * DZT3d * edge_length_2D[:east]
-    vmo = @. vo * ρ * DZT3d * edge_length_2D[:north]
+    for 𝑖 in indices.C
+        i, j = 𝑖.I # indices to access the 2D array
+        𝑗 = i₊₁(𝑖, gridtopology) # grid cell to the "east"
+        ϕᵢ[𝑖] = u[𝑖] * twocellnanmean(ρ, 𝑖, 𝑗) * twocellnanmin(thkcello, 𝑖, 𝑗) * edge_length_2D[:east][i,j]
+        𝑗 = j₊₁(𝑖, gridtopology) # grid cell to the "north"
+        ϕⱼ[𝑖] = v[𝑖] * twocellnanmean(ρ, 𝑖, 𝑗) * twocellnanmin(thkcello, 𝑖, 𝑗) * edge_length_2D[:north][i,j]
+    end
     # TODO check that this orientation is correct for all models?
     # I think this only applies to Arakawa C-grids...
 
-    return umo, vmo
+    return ϕᵢ, ϕⱼ
 end
 
+"""
+    twocellnanmean(x, 𝑖, 𝑗)
+
+Return the nanmean of `x` at the two cell indices `𝑖` and `𝑗`.
+"""
+twocellnanmean(x::Number, 𝑖, 𝑗) = x
+twocellnanmean(x, 𝑖, 𝑗) = nanmean2(x[𝑖], x[𝑗])
 
+"""
+    nanmean2(a, b)
+
+Return the nanmean of `a` and `b` (scalars).
+"""
+function nanmean2(a, b)
+    wa = !isnan(a)
+    wb = !isnan(b)
+    (wa * a + wb * b) / (wa + wb)
+end
+
+"""
+    twocellnanmin(x, 𝑖, 𝑗)
+
+Return the nanmin of `x` at the two cell indices `𝑖` and `𝑗`.
+"""
+twocellnanmin(x::Number, 𝑖, 𝑗) = x
+twocellnanmin(x, 𝑖, 𝑗) = nanmin2(x[𝑖], x[𝑗])
+
+"""
+    nanmin2(a, b)
+
+Return the nanmin of `a` and `b` (scalars).
+"""
+nanmin2(a, b) = isnan(a) ? b : isnan(b) ? a : min(a, b)
 
 """
     facefluxesfrommasstransport(; umo, vmo)
@@ -62,7 +110,7 @@ function facefluxesfromvelocities(; uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, mode
 
     # Convert to in-memory Array to avoid slow getindex
     # Convert to Float64 for double-precision mass conservation
-    umo, vmo = velocity2fluxes(; uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, modelgrid, ρ)
+    umo, vmo = velocity2fluxes(uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, modelgrid, ρ)
 
     return facefluxes(umo, vmo; FillValue)
 
@@ -138,7 +186,7 @@ function makemodelgrid(; areacello, volcello, lon, lat, lev, lon_vertices, lat_v
     area2D = replace(area2D, missing => NaN, 0 => NaN, FillValue => NaN)
 
 	# depth and cell height (3D)
-	DZT3d = v3D ./ area2D
+	thkcello = v3D ./ area2D
 	zt = lev |> Array
 
     lat = lat |> Array
@@ -159,9 +207,9 @@ function makemodelgrid(; areacello, volcello, lon, lat, lev, lon_vertices, lat_v
     edge_length_2D = Dict(d=>[verticalfacewidth(lon_vertices, lat_vertices, 𝑖.I[1], 𝑖.I[2], d) for 𝑖 in C] for d in dirs)
     distance_to_edge_2D = Dict(d=>[centroid2edgedistance(lon, lat, lon_vertices, lat_vertices, 𝑖.I[1], 𝑖.I[2], d) for 𝑖 in C] for d in dirs)
 
-    gridtype = gridtopology(lon_vertices, lat_vertices, zt)
+    arakawagrid = getgridtopology(lon_vertices, lat_vertices, zt)
 
-	return (; area2D, v3D, DZT3d, lon_vertices, lat_vertices, lon, lat, zt, edge_length_2D, distance_to_edge_2D, gridtype)
+	return (; area2D, v3D, thkcello, lon_vertices, lat_vertices, lon, lat, zt, edge_length_2D, distance_to_edge_2D, arakawagrid)
 end
 
 function makeindices(v3D)
diff --git a/src/topology.jl b/src/topology.jl
deleted file mode 100644
index 973a979..0000000
--- a/src/topology.jl
+++ /dev/null
@@ -1,104 +0,0 @@
-abstract type AbstractGrid end
-struct BipolarGrid <: AbstractGrid
-    nx::Int64
-    ny::Int64
-    nz::Int64
-end
-struct TripolarGrid <: AbstractGrid
-    nx::Int64
-    ny::Int64
-    nz::Int64
-end
-struct UnknownGrid <: AbstractGrid
-    nx::Int64
-    ny::Int64
-    nz::Int64
-end
-
-"""
-    gridtopology(lon_vertices, lat_vertices, lev)
-
-Returns the type of grid based on how the vertices connect at the north pole.
-"""
-function gridtopology(lon_vertices, lat_vertices, lev)
-    nx = size(lon_vertices, 2)
-    ny = size(lon_vertices, 3)
-    nz = length(lev)
-    # North pole vertices
-    NPlon = @view lon_vertices[3:4,:,end]
-    NPlat = @view lat_vertices[3:4,:,end]
-    # If all the latitudes of the "northmost" vertices are 90, then it's a "regular" grid with 2 poles
-    if all(NPlat .== 90)
-        return BipolarGrid(nx,ny,nz)
-    # Otherwise check if the north pole is split in two
-    elseif NPlon == rot180(NPlon) && NPlat == rot180(NPlat)
-        return TripolarGrid(nx,ny,nz)
-    else
-        return UnknownGrid(nx,ny,nz)
-    end
-end
-
-# Default behavior for grids (assumed bipolar)
-# wrap around in the i direction (longitude)
-i₊₁(C, g::AbstractGrid) = C.I[1] < g.nx ? C + CartesianIndex(1, 0, 0) : CartesianIndex(1, C.I[2], C.I[3])
-i₋₁(C, g::AbstractGrid) = C.I[1] > 1 ? C + CartesianIndex(-1, 0, 0) : CartesianIndex(g.nx, C.I[2], C.I[3])
-i₊₁(C::CartesianIndex{2}, g::AbstractGrid) = C.I[1] < g.nx ? C + CartesianIndex(1, 0) : CartesianIndex(1, C.I[2])
-i₋₁(C::CartesianIndex{2}, g::AbstractGrid) = C.I[1] > 1 ? C + CartesianIndex(-1, 0) : CartesianIndex(g.nx, C.I[2])
-# No connection by default in the j direction (latitude)
-j₊₁(C, g::AbstractGrid) = C.I[2] < g.ny ? C + CartesianIndex(0, 1, 0) : nothing
-j₋₁(C, ::AbstractGrid) = C.I[2] > 1 ? C + CartesianIndex(0, -1, 0) : nothing
-j₊₁(C::CartesianIndex{2}, g::AbstractGrid) = C.I[2] < g.ny ? C + CartesianIndex(0, 1) : nothing
-j₋₁(C::CartesianIndex{2}, ::AbstractGrid) = C.I[2] > 1 ? C + CartesianIndex(0, -1) : nothing
-# No connection by default in the k direction (depth)
-k₊₁(C, g::AbstractGrid) = C.I[3] < g.nz ? C + CartesianIndex(0, 0, 1) : nothing
-k₋₁(C, ::AbstractGrid) = C.I[3] > 1 ? C + CartesianIndex(0, 0, -1) : nothing
-
-
-
-function ishift(C, g::AbstractGrid, n=0)
-    i, j, k = C.I
-    CartesianIndex(mod1(i + n, g.nx), j, k)
-end
-function jshift(C, g::AbstractGrid, n=0)
-    i, j, k = C.I
-    j2 = j + n
-    (1 ≤ j2 ≤ g.ny) ? CartesianIndex(i, j2, k) : nothing
-end
-function kshift(C, g::AbstractGrid, n=0)
-    i, j, k = C.I
-    k2 = k + n
-    (1 ≤ k2 ≤ g.nz) ? CartesianIndex(i, j, k2) : nothing
-end
-
-# Special behavior for tripolar grids where the seam connects the top and "folds" it around
-# ┌─────┬─────┬─────┬─────┬─────┬─────┐
-# │  6  │  5  │  4  │  3  │  2  │  1  │
-# ├─────┼─────┼─────┼─────┼─────┼─────┤ ◄─ seam
-# │  1  │  2  │  3  │  4  │  5  │  6  │
-# └─────┴─────┴─────┴─────┴─────┴─────┘
-#            ─► increasing i ("east")
-j₊₁(C, g::TripolarGrid) = C.I[2] < g.ny ? C + CartesianIndex(0, 1, 0) : CartesianIndex(g.nx - C.I[1] + 1, g.ny, C.I[3])
-j₊₁(C::CartesianIndex{2}, g::TripolarGrid) = C.I[2] < g.ny ? C + CartesianIndex(0, 1) : CartesianIndex(g.nx - C.I[1] + 1, g.ny)
-
-function jshift(C, g::TripolarGrid, n=0)
-    i, j, k = C.I
-    j2 = j + n
-    if !(1 ≤ j2)
-        return nothing
-    elseif j2 ≤ g.ny
-        return CartesianIndex(i, j2, k)
-    else
-        i2 = g.nx - i + 1
-        return CartesianIndex(i2, j, k)
-    end
-end
-
-# error for unknown grids
-i₊₁(C, ::UnknownGrid) = error("Unknown grid type")
-i₋₁(C, ::UnknownGrid) = error("Unknown grid type")
-j₊₁(C, ::UnknownGrid) = error("Unknown grid type")
-j₋₁(C, ::UnknownGrid) = error("Unknown grid type")
-k₊₁(C, ::UnknownGrid) = error("Unknown grid type")
-k₋₁(C, ::UnknownGrid) = error("Unknown grid type")
-
-
diff --git a/test/local_fast.jl b/test/local_fast.jl
index 9066ab2..96f8bb6 100644
--- a/test/local_fast.jl
+++ b/test/local_fast.jl
@@ -63,7 +63,7 @@
 
         # Make makemodelgrid
         modelgrid = makemodelgrid(; areacello, volcello, lon, lat, lev, lon_vertices, lat_vertices)
-        (; lon_vertices, lat_vertices, edge_length_2D, distance_to_edge_2D, lon, lat, zt, area2D, v3D, DZT3d) = modelgrid
+        (; lon_vertices, lat_vertices, edge_length_2D, distance_to_edge_2D, lon, lat, zt, area2D, v3D, thkcello) = modelgrid
 
         # Make indices
         indices = makeindices(v3D)
@@ -222,7 +222,7 @@
 
 
 
-        # plot DZT3d every 1000m
+        # plot thkcello every 1000m
         fig = Figure(size=(1000, 750))
         zs = [0 1000; 2000 3000; 4000 5000]
         colorrange = (0, 500) # m
@@ -231,7 +231,7 @@
             i, j = Tuple(idx)
             ax = Axis(fig[i, j], xlabel="i", ylabel="j")
             k = findfirst(zt .> zs[i, j])
-            hm = heatmap!(ax, ustrip.(m, DZT3d[:, :, k] * m); colorrange)
+            hm = heatmap!(ax, ustrip.(m, thkcello[:, :, k] * m); colorrange)
             customdecorations!(ax; i, j, imax=size(zs, 1))
             text!(ax, 0.5, 1, text="$(zs[i,j])m", align=(:center, :top), offset=(0, -2), space=:relative)
         end
diff --git a/test/local_full.jl b/test/local_full.jl
index f280549..c0357c1 100644
--- a/test/local_full.jl
+++ b/test/local_full.jl
@@ -8,9 +8,8 @@
     using NetCDF
     using YAXArrays
     using GibbsSeaWater
-    using GLMakie
-    using NaNStatistics
     using DimensionalData
+    using NaNStatistics
 
     # stdlib
     using SparseArrays
@@ -54,7 +53,7 @@
 
     # Make makemodelgrid
     modelgrid = makemodelgrid(; areacello, volcello, lon, lat, lev, lon_vertices, lat_vertices)
-    (; lon_vertices, lat_vertices, v3D, zt) = modelgrid
+    (; lon_vertices, lat_vertices, v3D, zt, thkcello) = modelgrid
 
     uo = readcubedata(uo_ds.uo)
     vo = readcubedata(vo_ds.vo)
@@ -82,6 +81,52 @@
     # ct = Conservative Temperature (ITS-90) (°C)
     # p = sea pressure (dbar) (here using 0 pressure to get potential density
     # TODO: CHECK IF THIS IS CORRECT
+
+    # Some parameter values
+    ρ = 1035.0    # kg/m^3
+    # Alternatively, rebuild density from thetao, so, and depth as approximate pressure
+    ZBOT3D = cumsum(thkcello, dims = 3)
+    Z3D = ZBOT3D - 0.5 * thkcello
+    ρ = gsw_rho.(so, ct, Z3D)
+
+    # Diffusivites
+    κH = 500.0    # m^2/s
+    κVML = 0.1    # m^2/s
+    κVdeep = 1e-5 # m^2/s
+
+    # Make fuxes from all directions
+    ϕ = facefluxesfrommasstransport(; umo, vmo)
+
+    # Make fuxes from all directions from velocities
+    ϕ_bis = facefluxesfromvelocities(; uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, modelgrid, ρ)
+
+    # Make indices
+    indices = makeindices(modelgrid.v3D)
+
+    @test all(.!isnan.(ρ[indices.wet3D])) == true
+
+    # Make transport matrix
+    (; T, Tadv, TκH, TκVML, TκVdeep) = transportmatrix(; ϕ, mlotst, modelgrid, indices, ρ, κH, κVML, κVdeep)
+
+    Tsyms = (:T, :Tadv, :TκH, :TκVML, :TκVdeep)
+	for Ttest in (T, Tadv, TκH, TκVML, TκVdeep)
+        @test Ttest isa SparseMatrixCSC{Float64, Int}
+    end
+
+
+end
+
+@testitem "velocities and mass transports" setup=[LocalBuiltMatrix] tags=[:skipci] begin
+
+    using NaNStatistics
+    using GLMakie
+
+    (; modelgrid, indices, ρθ, v3D, lat, lon, zt, uo, vo, uo_lon, uo_lat, vo_lon, vo_lat,
+    lon_vertices, lat_vertices, indices,
+    umo, vmo, umo_lon, umo_lat, vmo_lon, vmo_lat, model, member, outputdir) = LocalBuiltMatrix
+
+
+
     # plot for sanity check
     begin # plot zonal average
         ρθ2D = dropdims(nansum(ρθ .* v3D, dims = 1) ./ nansum(v3D, dims = 1), dims = 1)
@@ -102,39 +147,8 @@
     save(outputfile, fig)
 
     κGM = 600 # m^2/s
-    uGM, vGM = OceanTransportMatrixBuilder.bolus_GM_velocity(ρθ, κGM, modelgrid)
-    # Turn uGM into a YAXArray by rebuilding from uo
-    uGM_YAXArray = rebuild(uo;
-        data = uGM,
-        dims = dims(uo),
-        metadata = Dict(
-            "origin" => "u GM bolus velocity computed from thetao and so",
-            "kGM" => κGM,
-            "units" => "m/s",
-        )
-    )
-    arrays = Dict(:uo_GM => uGM_YAXArray, :lat => uo_lat, :lon => uo_lon)
-    uGM_ds = Dataset(; uo_ds.properties, arrays...)
-    # Save to netCDF file
-    outputfile = joinpath(outputdir, "uo_GM.nc")
-    @info "uo_GM as netCDF file:\n  $(outputfile)"
-    savedataset(uGM_ds, path = outputfile, driver = :netcdf, overwrite = true)
-    # Turn vGM into a YAXArray by rebuilding from uo
-    vGM_YAXArray = rebuild(vo;
-        data = vGM,
-        dims = dims(vo),
-        metadata = Dict(
-            "origin" => "v GM bolus velocity computed from thetao and so",
-            "kGM" => κGM,
-            "units" => "m/s",
-        )
-    )
-    arrays = Dict(:vo_GM => vGM_YAXArray, :lat => vo_lat, :lon => vo_lon)
-    vGM_ds = Dataset(; vo_ds.properties, arrays...)
-    # Save to netCDF file
-    outputfile = joinpath(outputdir, "vo_GM.nc")
-    @info "Saving vo_GM as netCDF file:\n  $(outputfile)"
-    savedataset(vGM_ds, path = outputfile, driver = :netcdf, overwrite = true)
+    maxslope = 0.01
+    uGM, vGM = OceanTransportMatrixBuilder.bolus_GM_velocity(ρθ, modelgrid; κGM, maxslope)
 
     # Plot location of cell center for volcello, umo, vmo, uo, vo
     fig = Figure()
@@ -147,23 +161,6 @@
     text!(ax, vo_lon[1], vo_lat[1]; text="vo (i,j)", align = (:center, :bottom))
     fig
 
-
-    # Some parameter values
-    ρ = 1035.0    # kg/m^3
-    κH = 500.0    # m^2/s
-    κVML = 0.1    # m^2/s
-    κVdeep = 1e-5 # m^2/s
-
-
-    # Make fuxes from all directions
-    ϕ = facefluxesfrommasstransport(; umo, vmo)
-
-    # Make fuxes from all directions from velocities
-    ϕ_bis = facefluxesfromvelocities(; uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, modelgrid, ρ)
-
-    # Make indices
-    indices = makeindices(modelgrid.v3D)
-
     uo2, uo2_lon, uo2_lat, vo2, vo2_lon, vo2_lat = OceanTransportMatrixBuilder.interpolateontodefaultCgrid(uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, modelgrid)
     fig = Figure(size=(1500, 800))
     ax = Axis(fig[1,1], xlabel = "lon", ylabel = "lat", xgridvisible = false, ygridvisible = false)
@@ -196,21 +193,9 @@
     hm = heatmap!(ax, indices.wet3D[:,:,1])
     translate!(hm, 0, 0, -100) # move the plot behind the grid
     fig
-
-    # Make transport matrix
-    (; T, Tadv, TκH, TκVML, TκVdeep) = transportmatrix(; ϕ, mlotst, modelgrid, indices, ρ, κH, κVML, κVdeep)
-
-    Tsyms = (:T, :Tadv, :TκH, :TκVML, :TκVdeep)
-	for Ttest in (T, Tadv, TκH, TκVML, TκVdeep)
-        @test Ttest isa SparseMatrixCSC{Float64, Int}
-    end
-
-
 end
 
-
-
-@testitem "Timescales (divergence and mass conservation)" setup=[LocalBuiltMatrix] tags=[:skipci] begin
+@testitem "Divergence and mass conservation" setup=[LocalBuiltMatrix] tags=[:skipci] begin
 
     using Unitful
     using Unitful: s, Myr
@@ -238,7 +223,7 @@ end
 	end
 end
 
-@testitem "Test if flux divergence (not convergence)" setup=[LocalBuiltMatrix] tags=[:skipci] begin
+@testitem "Test flux divergence" setup=[LocalBuiltMatrix] tags=[:skipci] begin
 
     using SparseArrays
     using LinearAlgebra
@@ -364,7 +349,7 @@ end
 
     # Difference between the fluxes from velocities and mass transport
     (; uo, vo, umo, vmo, uo_lon, uo_lat, vo_lon, vo_lat, modelgrid, ρ) = LocalBuiltMatrix
-    umo_bis, vmo_bis = velocity2fluxes(; uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, modelgrid, ρ)
+    umo_bis, vmo_bis = velocity2fluxes(uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, modelgrid, ρ)
     colorrange = 1e9 .* (-1, 1)
     colormap = cgrad(:RdBu, rev=true)
     Δcolorrange = (-5, 5)
diff --git a/test/online.jl b/test/online.jl
index fb666a3..d0e23d4 100644
--- a/test/online.jl
+++ b/test/online.jl
@@ -1,6 +1,6 @@
 
 
-@testitem "building matrices" begin
+@testitem "ACCESS-ESM1-5" begin
 
     # using Test
     # using OceanTransportMatrixBuilder
@@ -11,6 +11,7 @@
     using CSV
     using Unitful
     using Unitful: s, Myr
+    using Downloads: download
 
     # stdlib
     using SparseArrays