Fix local tests

Author: briochemc

.gitignore

@@ -2,4 +2,5 @@
 
 .vscode/
 
-test/plots/
+test/plots/
+plots/

Project.toml

@@ -1,7 +1,7 @@
 name = "OceanTransportMatrixBuilder"
 uuid = "c2b4a04e-6049-4fc4-aa6a-5508a29a1e1c"
 authors = ["Benoit Pasquier <[email protected]> and contributors"]
-version = "0.6.0"
+version = "0.6.1"
 
 [deps]
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"

README.md

@@ -30,7 +30,7 @@ using NetCDF
 using YAXArrays
 using OceanTransportMatrixBuilder
 
-inputdir = "/Users/benoitpasquier/Data/TMIP/data/ACCESS-ESM1-5/historical/r1i1p1f1/Jan1990-Dec1999"
+inputdir = "$(ENV["HOME"])/Data/TMIP/data/ACCESS-ESM1-5/historical/r1i1p1f1/Jan1990-Dec1999"
 # The path above is for my local machine. See test/onlinebuild.jl for an example using data from the cloud.
 
 # Load datasets

src/extratools.jl

@@ -31,7 +31,7 @@ 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, volume; f=Returns(true), di=2, dj=2, dk=1)
+function lump_and_spray(wet3D, vol; f=Returns(true), di=2, dj=2, dk=1)
 
     # extend the grid to avoid lumping cells outside of bounds
     nxyz = size(wet3D)
@@ -54,7 +54,7 @@ function lump_and_spray(wet3D, volume; f=Returns(true), di=2, dj=2, dk=1)
             c += 1
         end
     end
-    LUMP = sparse(LUMPidx[:], 1:length(LUMPidx), 1)
+    LUMP = sparse(LUMPidx[C][:], 1:length(C), 1)
 
     # Find wet points in coarsened grid
     wet = wet3D[:]
@@ -66,15 +66,15 @@ function lump_and_spray(wet3D, volume; f=Returns(true), di=2, dj=2, dk=1)
     # Make the LUMP operator volume-conserving
     # by volume integrating on the right and dividing by the coarse
     # volume on the left
-    volume_c = LUMP * volume
-    LUMP = sparse(Diagonal(1 ./ volume_c)) * LUMP * sparse(Diagonal(volume))
+    vol_c = LUMP * vol
+    LUMP = sparse(Diagonal(1 ./ vol_c)) * LUMP * sparse(Diagonal(vol))
 
     # The SPRAY operator just copies the values back
     # so it is sinply 1's with the transposed sparsity structure
     SPRAY = copy(LUMP')
     SPRAY.nzval .= 1
 
-    return LUMP, SPRAY, volume_c
+    return LUMP, SPRAY, vol_c
 end
 
 

test/LocalBuiltMatrix.jl

@@ -16,7 +16,7 @@
     # My local directory for input files
     model = "ACCESS-ESM1-5"
     member = "r1i1p1f1"
-    inputdir = "/Users/benoitpasquier/Data/TMIP/data/$model/historical/$member/Jan1990-Dec1999"
+    inputdir = "$(ENV["HOME"])/Data/TMIP/data/$model/historical/$member/Jan1990-Dec1999"
 
     # and for output files
     @show version = "v$(pkgversion(OceanTransportMatrixBuilder))"
@@ -64,9 +64,9 @@
     thetao_ds = open_dataset(joinpath(inputdir, "thetao.nc"))
     so_ds = open_dataset(joinpath(inputdir, "so.nc"))
     # Load variables in memory
-    thetao = readcubedata(thetao_ds.thetao)
+    thetao = readcubedata(thetao_ds.thetao) |> Array
     @test 0 < nanmean(thetao) < 20
-    so = readcubedata(so_ds.so)
+    so = readcubedata(so_ds.so) |> Array
     @show 30 < nanmean(so) < 40
     # Convert thetao and so to density
     ct = gsw_ct_from_pt.(so, thetao)

test/interactive.jl

@@ -4,6 +4,7 @@ Pkg.instantiate()
 using TestEnv
 TestEnv.activate();
 
+using Revise
 
 using Test
 using OceanTransportMatrixBuilder
@@ -20,7 +21,7 @@ using LinearAlgebra
 # My local directory for input files
 model = "ACCESS-ESM1-5"
 member = "r1i1p1f1"
-inputdir = "/Users/benoitpasquier/Data/TMIP/data/$model/historical/$member/Jan1990-Dec1999"
+inputdir = "$(ENV["HOME"])/Data/TMIP/data/$model/historical/$member/Jan1990-Dec1999"
 
 # and for output files
 @show version = "v$(pkgversion(OceanTransportMatrixBuilder))"
@@ -73,8 +74,8 @@ thetao = readcubedata(thetao_ds.thetao)
 so = readcubedata(so_ds.so)
 @show 30 < nanmean(so) < 40
 # Convert thetao and so to density
-ct = gsw_ct_from_pt.(so, thetao)
-ρ = gsw_rho.(so, ct, Z3D)
+ct = map(gsw_ct_from_pt, so, thetao)
+ρ = map(gsw_rho, so, ct, Z3D)
 # TODO: Check if this is correct usage of gsw functions!
 # Alternatively use a fixed density:
 # ρ = 1035.0    # kg/m^3

test/local_fast.jl

@@ -1,7 +1,7 @@
 
 
 
-@testitem "grid checks" tags = [:skipci] begin
+@testitem "grid checks" tags = [:skipci, :local_fast] begin
 
     using Test
     using OceanTransportMatrixBuilder
@@ -14,8 +14,10 @@
     using Unitful: m, km
 
     # TODO fix this so that I don't have to track all members
-    models = ("ACCESS-ESM1-5", "ACCESS-CM2", "ACCESS1-3", "ACCESS1-0")
-    members = ("r1i1p1f1", "r1i1p1f1", "r1i1p1", "r1i1p1")
+    # models = ("ACCESS-ESM1-5", "ACCESS-CM2", "ACCESS1-3", "ACCESS1-0")
+    # members = ("r1i1p1f1", "r1i1p1f1", "r1i1p1", "r1i1p1")
+    models = ("ACCESS-ESM1-5",)# "ACCESS-CM2", "ACCESS1-3", "ACCESS1-0")
+    members = ("r1i1p1f1",)# "r1i1p1f1", "r1i1p1", "r1i1p1")
 
     # models_untested = ("CMCC-CM2-HR4",)
     # members_untested = ("r1i1p1f1",)
@@ -39,7 +41,8 @@
     for (model, member) in zip(models, members)
 
         # My local directory for input files
-        inputdir = "/Users/benoitpasquier/Data/TMIP/data/$model/historical/$member/Jan1990-Dec1999"
+        inputdir = "$(ENV["HOME"])/Data/TMIP/data/$model/historical/$member/Jan1990-Dec1999"
+        @info "inputdir = $inputdir"
 
         # Load datasets
         volcello_ds = open_dataset(joinpath(inputdir, "volcello.nc"))

test/local_full.jl

@@ -2,7 +2,7 @@
 
 
 
-@testitem "velocities and mass transports" setup=[LocalBuiltMatrix] tags=[:skipci] begin
+@testitem "velocities and mass transports" setup=[LocalBuiltMatrix] tags=[:skipci, :local_full] begin
 
     using NaNStatistics
     using CairoMakie
@@ -15,7 +15,7 @@
 
     # plot for sanity check
     # plot zonal average density
-    ρ2D = dropdims(nansum(ρ.data .* v3D, dims = 1) ./ nansum(v3D, dims = 1), dims = 1)
+    ρ2D = dropdims(nansum(ρ .* v3D, dims = 1) ./ nansum(v3D, dims = 1), dims = 1)
     fig = Figure()
     ax = Axis(fig[1,1], xlabel = "latitude (°)", ylabel = "depth (m)")
     # levels = 25:0.1:30
@@ -32,10 +32,6 @@
     @info "Saving density zonal average as image file:\n $(joinpath("test", outputfile))"
     save(outputfile, fig)
 
-    κGM = 600 # m^2/s
-    maxslope = 0.01
-    uGM, vGM = OceanTransportMatrixBuilder.bolus_GM_velocity(ρ, gridmetrics, indices; κGM, maxslope)
-
     # Plot location of cell center for volcello, umo, vmo, uo, vo
     fig = Figure()
     ax = Axis(fig[1,1], xlabel = "lon", ylabel = "lat")
@@ -47,31 +43,37 @@
     text!(ax, vo_lon[1], vo_lat[1]; text="vo (i,j)", align = (:center, :bottom))
     fig
 
+
+    κGM = 600 # m^2/s
+    maxslope = 0.01
+    uGM, vGM = OceanTransportMatrixBuilder.bolus_GM_velocity(ρ, gridmetrics, indices; κGM, maxslope)
+
+
     uo2, uo2_lon, uo2_lat, vo2, vo2_lon, vo2_lat = OceanTransportMatrixBuilder.interpolateontodefaultCgrid(uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, gridmetrics)
     fig = Figure(size=(1500, 800))
     ax = Axis(fig[1,1], xlabel = "lon", ylabel = "lat", xgridvisible = false, ygridvisible = false)
     i = 85 .+ (-1:5); j = 2 .+ (-1:4); k = 1
     [poly!(ax, lon_vertices[:,i,j] |> vec, lat_vertices[:,i,j] |> vec, color = (:blue, 0.1)) for i in i for j in j if indices.wet3D[i,j,1]];
     [lines!(ax, lon_vertices[:,i,j] |> vec, lat_vertices[:,i,j] |> vec, color = (:black, 0.1)) for i in i for j in j];
     # lines!(ax, lon_vertices[1,i,j] |> Array, lat_vertices[1,i,j] |> Array, color = :black)
-    arrows!(ax, uo_lon[i,j,k] |> vec, uo_lat[i,j,k] |> vec, uo[i,j,k] |> vec, 0uo[i,j,k] |> vec, arrowsize = 10, lengthscale = 10, arrowcolor = :darkblue, linecolor = :darkblue)
-    arrows!(ax, vo_lon[i,j,k] |> vec, vo_lat[i,j,k] |> vec, 0vo[i,j,k] |> vec, vo[i,j,k] |> vec, arrowsize = 10, lengthscale = 10, arrowcolor = :darkred, linecolor = :darkred)
-    arrows!(ax, uo2_lon[i,j,k] |> vec, uo2_lat[i,j,k] |> vec, uo2[i,j,k] |> vec, 0uo2[i,j,k] |> vec, arrowsize = 10, lengthscale = 10, arrowcolor = :blue, linecolor = :blue)
-    arrows!(ax, vo2_lon[i,j,k] |> vec, vo2_lat[i,j,k] |> vec, 0vo2[i,j,k] |> vec, vo2[i,j,k] |> vec, arrowsize = 10, lengthscale = 10, arrowcolor = :red, linecolor = :red)
-    arrows!(ax, uo2_lon[i,j,k] |> vec, uo2_lat[i,j,k] .+ 0.02 |> vec, umo[i,j,k] |> vec, 0umo[i,j,k] |> vec, arrowsize = 10, lengthscale = 3e-8, arrowcolor = (:blue, 0.3), linecolor = (:blue, 0.3))
-    arrows!(ax, vo2_lon[i,j,k] .+ 0.02 |> vec, vo2_lat[i,j,k] |> vec, 0vmo[i,j,k] |> vec, vmo[i,j,k] |> vec, arrowsize = 10, lengthscale = 3e-8, arrowcolor = (:red, 0.3), linecolor = (:red, 0.3))
+    arrows2d!(ax, uo_lon[i,j,k] |> vec, uo_lat[i,j,k] |> vec, uo[i,j,k] |> vec, 0(uo[i,j,k] |> vec), lengthscale = 10, tipcolor = :darkblue, shaftcolor = :darkblue)
+    arrows2d!(ax, vo_lon[i,j,k] |> vec, vo_lat[i,j,k] |> vec, 0(vo[i,j,k] |> vec), vo[i,j,k] |> vec, lengthscale = 10, tipcolor = :darkred, shaftcolor = :darkred)
+    arrows2d!(ax, uo2_lon[i,j,k] |> vec, uo2_lat[i,j,k] |> vec, uo2[i,j,k] |> vec, 0(uo2[i,j,k] |> vec), lengthscale = 10, tipcolor = :blue, shaftcolor = :blue)
+    arrows2d!(ax, vo2_lon[i,j,k] |> vec, vo2_lat[i,j,k] |> vec, 0(vo2[i,j,k] |> vec), vo2[i,j,k] |> vec, lengthscale = 10, tipcolor = :red, shaftcolor = :red)
+    arrows2d!(ax, uo2_lon[i,j,k] |> vec, uo2_lat[i,j,k] .+ 0.02 |> vec, umo[i,j,k] |> vec, 0(umo[i,j,k] |> vec), lengthscale = 3e-8, tipcolor = (:blue, 0.3), shaftcolor = (:blue, 0.3))
+    arrows2d!(ax, vo2_lon[i,j,k] .+ 0.02 |> vec, vo2_lat[i,j,k] |> vec, 0(vmo[i,j,k] |> vec), vmo[i,j,k] |> vec, lengthscale = 3e-8, tipcolor = (:red, 0.3), shaftcolor = (:red, 0.3))
     fig
     ax = Axis(fig[1,2], xlabel = "lon", ylabel = "lat", xgridvisible = false, ygridvisible = false)
     i = 160 .+ (-3:3); j = 200 .+ (-2:2); k = 1
     [poly!(ax, lon_vertices[:,i,j] |> vec, lat_vertices[:,i,j] |> vec, color = (:blue, 0.1)) for i in i for j in j if indices.wet3D[i,j,1]];
     [lines!(ax, lon_vertices[:,i,j] |> vec, lat_vertices[:,i,j] |> vec, color = (:black, 0.1)) for i in i for j in j];
     # lines!(ax, lon_vertices[1,i,j] |> Array, lat_vertices[1,i,j] |> Array, color = :black)
-    arrows!(ax, uo_lon[i,j,k] |> vec, uo_lat[i,j,k] |> vec, uo[i,j,k] |> vec, 0uo[i,j,k] |> vec, arrowsize = 10, lengthscale = 10, arrowcolor = :darkblue, linecolor = :darkblue)
-    arrows!(ax, vo_lon[i,j,k] |> vec, vo_lat[i,j,k] |> vec, 0vo[i,j,k] |> vec, vo[i,j,k] |> vec, arrowsize = 10, lengthscale = 10, arrowcolor = :darkred, linecolor = :darkred)
-    arrows!(ax, uo2_lon[i,j,k] |> vec, uo2_lat[i,j,k] |> vec, uo2[i,j,k] |> vec, 0uo2[i,j,k] |> vec, arrowsize = 10, lengthscale = 10, arrowcolor = :blue, linecolor = :blue)
-    arrows!(ax, vo2_lon[i,j,k] |> vec, vo2_lat[i,j,k] |> vec, 0vo2[i,j,k] |> vec, vo2[i,j,k] |> vec, arrowsize = 10, lengthscale = 10, arrowcolor = :red, linecolor = :red)
-    arrows!(ax, uo2_lon[i,j,k] |> vec, uo2_lat[i,j,k] .+ 0.02 |> vec, umo[i,j,k] |> vec, 0umo[i,j,k] |> vec, arrowsize = 10, lengthscale = 1e-8, arrowcolor = (:blue, 0.3), linecolor = (:blue, 0.3))
-    arrows!(ax, vo2_lon[i,j,k] .+ 0.02 |> vec, vo2_lat[i,j,k] |> vec, 0vmo[i,j,k] |> vec, vmo[i,j,k] |> vec, arrowsize = 10, lengthscale = 1e-8, arrowcolor = (:red, 0.3), linecolor = (:red, 0.3))
+    arrows2d!(ax, uo_lon[i,j,k] |> vec, uo_lat[i,j,k] |> vec, uo[i,j,k] |> vec, 0(uo[i,j,k] |> vec), lengthscale = 10, tipcolor = :darkblue, shaftcolor = :darkblue)
+    arrows2d!(ax, vo_lon[i,j,k] |> vec, vo_lat[i,j,k] |> vec, 0(vo[i,j,k] |> vec), vo[i,j,k] |> vec, lengthscale = 10, tipcolor = :darkred, shaftcolor = :darkred)
+    arrows2d!(ax, uo2_lon[i,j,k] |> vec, uo2_lat[i,j,k] |> vec, uo2[i,j,k] |> vec, 0(uo2[i,j,k] |> vec), lengthscale = 10, tipcolor = :blue, shaftcolor = :blue)
+    arrows2d!(ax, vo2_lon[i,j,k] |> vec, vo2_lat[i,j,k] |> vec, 0(vo2[i,j,k] |> vec), vo2[i,j,k] |> vec, lengthscale = 10, tipcolor = :red, shaftcolor = :red)
+    arrows2d!(ax, uo2_lon[i,j,k] |> vec, uo2_lat[i,j,k] .+ 0.02 |> vec, umo[i,j,k] |> vec, 0(umo[i,j,k] |> vec), lengthscale = 1e-8, tipcolor = (:blue, 0.3), shaftcolor = (:blue, 0.3))
+    arrows2d!(ax, vo2_lon[i,j,k] .+ 0.02 |> vec, vo2_lat[i,j,k] |> vec, 0(vmo[i,j,k] |> vec), vmo[i,j,k] |> vec, lengthscale = 1e-8, tipcolor = (:red, 0.3), shaftcolor = (:red, 0.3))
     fig
 
     fig = Figure()
@@ -143,6 +145,9 @@ end
     M_c = sparse(Diagonal(issrf_c))
     sΓ_c = LUMP * sΓ
     @time "Solving ideal mean age" Γ_c = (T_c + M_c) \ sΓ_c
+    # Just some timings from machines I have used (ACCESS-ESM1-5 coarsened 2x2x1)
+    # time      machine  proc.   RAM
+    # 3min  MacBook Air  8 × M2  8GB
     Γ = SPRAY * Γ_c
     Γyr = ustrip.(yr, Γ .* s)
     Γ3D = OceanTransportMatrixBuilder.as3D(Γyr, wet3D)
@@ -220,6 +225,8 @@ end
     # Difference between the fluxes from velocities and mass transport
     (; uo, vo, umo, vmo, uo_lon, uo_lat, vo_lon, vo_lat, gridmetrics, ρ) = LocalBuiltMatrix
     umo_bis, vmo_bis = velocity2fluxes(uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, gridmetrics, ρ)
+    umo_arr = umo |> Array
+    vmo_arr = vmo |> Array
     colorrange = 1e9 .* (-1, 1)
     colormap = cgrad(:RdBu, rev=true)
     Δcolorrange = (-5, 5)
@@ -239,22 +246,22 @@ end
         hidedecorations!(ax)
 
         ax = Axis(fig[1,2], xlabel = "i", ylabel = "j")
-        heatmap!(ax, umo[:,:,k]; colormap, colorrange)
+        heatmap!(ax, umo_arr[:,:,k]; colormap, colorrange)
         text!(ax, 0, 1; text = "umo", align = (:left, :top), offset = (5, -5), space = :relative)
         hidedecorations!(ax)
 
         ax = Axis(fig[2,2], xlabel = "i", ylabel = "j")
-        hm = heatmap!(ax, vmo[:,:,k]; colormap, colorrange)
+        hm = heatmap!(ax, vmo_arr[:,:,k]; colormap, colorrange)
         text!(ax, 0, 1; text = "vmo", align = (:left, :top), offset = (5, -5), space = :relative)
         hidedecorations!(ax)
 
         ax = Axis(fig[1,3], xlabel = "i", ylabel = "j")
-        heatmap!(ax, 100((umo_bis - umo) ./ umo)[:,:,k]; colormap = Δcolormap, colorrange = Δcolorrange)
+        heatmap!(ax, 100((umo_bis - umo_arr) ./ umo_arr)[:,:,k]; colormap = Δcolormap, colorrange = Δcolorrange)
         text!(ax, 0, 1; text = "(umo_bis − umo) / umo", align = (:left, :top), offset = (5, -5), space = :relative)
         hidedecorations!(ax)
 
         ax = Axis(fig[2,3], xlabel = "i", ylabel = "j")
-        Δhm = heatmap!(ax, 100((vmo_bis - vmo) ./ vmo)[:,:,k]; colormap = Δcolormap, colorrange = Δcolorrange)
+        Δhm = heatmap!(ax, 100((vmo_bis - vmo_arr) ./ vmo_arr)[:,:,k]; colormap = Δcolormap, colorrange = Δcolorrange)
         text!(ax, 0, 1; text = "(vmo_bis − vmo) / vmo", align = (:left, :top), offset = (5, -5), space = :relative)
         hidedecorations!(ax)
 

test/runtests.jl

@@ -1,3 +1,3 @@
 using TestItemRunner
 
-@run_package_tests verbose=true filter=ti->!(:skipci in ti.tags)
+@run_package_tests verbose=true filter=test->(:skipci ∉ test.tags)