Add fluxes2velocity function

Author: briochemc

Project.toml

@@ -1,6 +1,6 @@
 name = "OceanTransportMatrixBuilder"
 uuid = "c2b4a04e-6049-4fc4-aa6a-5508a29a1e1c"
-version = "0.8.2"
+version = "0.8.3"
 authors = ["Benoit Pasquier <briochemc@gmail.com> and contributors"]
 
 [deps]

src/OceanTransportMatrixBuilder.jl

@@ -28,7 +28,7 @@ include("RediGM.jl")
 # Extras
 include("extratools.jl") # <- I think this should be in a separate "base" repo
 
-export velocity2fluxes
+export velocity2fluxes, fluxes2velocity
 export facefluxesfrommasstransport
 export facefluxesfromvelocities
 export makegridmetrics

src/velocities.jl

@@ -38,6 +38,41 @@ function velocity2fluxes(u, u_lon, u_lat, v, v_lon, v_lat, gridmetrics, ρ)
     return ϕᵢ, ϕⱼ
 end
 
+"""
+    fluxes2velocity(ϕᵢ, ϕⱼ, gridmetrics, ρ)
+
+Return the velocities `u` and `v` on the C-grid given the mass fluxes `ϕᵢ` and `ϕⱼ` (in kg/s).
+
+The velocities 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 fluxes2velocity(ϕᵢ, ϕⱼ, gridmetrics, ρ)
+
+    # Unpack gridmetrics
+    (; thkcello, edge_length_2D, v3D, lon_vertices, lat_vertices, zt) = gridmetrics
+
+    # make indices
+    indices = makeindices(v3D)
+
+    # grid type
+    gridtopology = getgridtopology(lon_vertices, lat_vertices, zt)
+
+    u = zeros(size(ϕᵢ))
+    v = zeros(size(ϕⱼ))
+
+    # Calculate velocities from fluxes
+    for 𝑖 in indices.C
+        i, j = 𝑖.I
+        𝑗 = 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
+
+    return u, v
+end
+
 """
     twocellnanmean(x, 𝑖, 𝑗)
 

test/local_full.jl

@@ -297,4 +297,10 @@ end
         save(outputfile, fig)
     end
 
+    # Round-trip test: velocity → fluxes → velocity should recover the C-grid velocities
+    u_cgrid, _, _, v_cgrid, _, _ = OceanTransportMatrixBuilder.interpolateontodefaultCgrid(uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, gridmetrics)
+    u_rec, v_rec = fluxes2velocity(umo_bis, vmo_bis, gridmetrics, ρ)
+    @test u_rec[wet3D] ≈ u_cgrid[wet3D]
+    @test v_rec[wet3D] ≈ v_cgrid[wet3D]
+
 end

test/test_fluxes2velocity.jl

@@ -0,0 +1,55 @@
+# Standalone round-trip test for fluxes2velocity
+# Run with: julia --project test/test_fluxes2velocity.jl
+
+using Test
+using OceanTransportMatrixBuilder
+using NetCDF
+using YAXArrays
+using GibbsSeaWater
+using DimensionalData
+using NaNStatistics
+
+model = "ACCESS-ESM1-5"
+member = "r1i1p1f1"
+inputdir = "$(ENV["HOME"])/Data/TMIP/data/$model/historical/$member/Jan1990-Dec1999"
+
+# Load datasets
+uo_ds = open_dataset(joinpath(inputdir, "uo.nc"))
+vo_ds = open_dataset(joinpath(inputdir, "vo.nc"))
+volcello_ds = open_dataset(joinpath(inputdir, "volcello.nc"))
+areacello_ds = open_dataset(joinpath(inputdir, "areacello.nc"))
+thetao_ds = open_dataset(joinpath(inputdir, "thetao.nc"))
+so_ds = open_dataset(joinpath(inputdir, "so.nc"))
+
+areacello = readcubedata(areacello_ds.areacello)
+volcello = readcubedata(volcello_ds.volcello)
+lon = readcubedata(volcello_ds.lon)
+lat = readcubedata(volcello_ds.lat)
+lev = volcello_ds.lev
+lon_vertices = readcubedata(volcello_ds.lon_verticies)
+lat_vertices = readcubedata(volcello_ds.lat_verticies)
+uo = readcubedata(uo_ds.uo)
+vo = readcubedata(vo_ds.vo)
+uo_lon = readcubedata(uo_ds.lon)
+uo_lat = readcubedata(uo_ds.lat)
+vo_lon = readcubedata(vo_ds.lon)
+vo_lat = readcubedata(vo_ds.lat)
+thetao = readcubedata(thetao_ds.thetao) |> Array
+so = readcubedata(so_ds.so) |> Array
+
+gridmetrics = makegridmetrics(; areacello, volcello, lon, lat, lev, lon_vertices, lat_vertices)
+(; Z3D, v3D) = gridmetrics
+ct = gsw_ct_from_pt.(so, thetao)
+ρ = gsw_rho.(so, ct, Z3D)
+indices = makeindices(v3D)
+(; wet3D) = indices
+
+# Round-trip: velocity → flux → velocity
+u_cgrid, _, _, v_cgrid, _, _ = OceanTransportMatrixBuilder.interpolateontodefaultCgrid(uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, gridmetrics)
+umo_bis, vmo_bis = velocity2fluxes(uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, gridmetrics, ρ)
+u_rec, v_rec = fluxes2velocity(umo_bis, vmo_bis, gridmetrics, ρ)
+
+@test u_rec[wet3D] ≈ u_cgrid[wet3D]
+@test v_rec[wet3D] ≈ v_cgrid[wet3D]
+
+println("All tests passed!")