From ded800d2c365aa03887a356d60229be45a3a7420 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Wed, 30 Oct 2024 13:25:08 +1100 Subject: [PATCH] Fix allocations + add interactive profiler --- Project.toml | 2 +- src/gridcellgeometry.jl | 9 ++-- src/gridtopology.jl | 2 +- src/matrixbuilding.jl | 72 ++++++++++++++----------- src/velocities.jl | 10 ++-- test/interactive.jl | 113 ++++++++++++++++++++++++++++++++++++++++ 6 files changed, 167 insertions(+), 41 deletions(-) create mode 100644 test/interactive.jl diff --git a/Project.toml b/Project.toml index ae635e8..af5c00a 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "OceanTransportMatrixBuilder" uuid = "c2b4a04e-6049-4fc4-aa6a-5508a29a1e1c" authors = ["Benoit Pasquier and contributors"] -version = "0.5.0" +version = "0.5.1" [deps] Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7" diff --git a/src/gridcellgeometry.jl b/src/gridcellgeometry.jl index 1d1cc8f..929cb76 100644 --- a/src/gridcellgeometry.jl +++ b/src/gridcellgeometry.jl @@ -193,6 +193,7 @@ function horizontaldistance(lon, lat, I, J) PJ = (getindexornan(lon, J), getindexornan(lat, J)) return haversine(PI, PJ) end +horizontaldistance(lon, lat, I, ::Nothing) = NaN function verticaldistance(Z::Vector, I, J) I = verticalindex(I) J = verticalindex(J) @@ -304,12 +305,14 @@ function makegridmetrics(; areacello, volcello, lon, lat, lev, lon_vertices, lat C = CartesianIndices(size(lon)) + gridtopology = getgridtopology(lon_vertices, lat_vertices, zt) + dirs = (:south, :east, :north, :west) + 𝑗s = (j₋₁, i₊₁, j₊₁, i₋₁) 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) + distance_to_neighbour_2D = Dict(d=>[horizontaldistance(lon, lat, 𝑖, 𝑗(𝑖, gridtopology)) for 𝑖 in C] for (d,𝑗) in zip(dirs,𝑗s)) - gridtopology = getgridtopology(lon_vertices, lat_vertices, zt) - - return (; area2D, v3D, thkcello, lon_vertices, lat_vertices, lon, lat, Z3D, zt, edge_length_2D, distance_to_edge_2D, gridtopology) + return (; area2D, v3D, thkcello, lon_vertices, lat_vertices, lon, lat, Z3D, zt, edge_length_2D, distance_to_edge_2D, distance_to_neighbour_2D, gridtopology) end diff --git a/src/gridtopology.jl b/src/gridtopology.jl index 347667b..1b308e6 100644 --- a/src/gridtopology.jl +++ b/src/gridtopology.jl @@ -115,4 +115,4 @@ idx₊₁(::Jcoord) = j₊₁ idx₊₁(::Kcoord) = k₊₁ idx₋₁(::Icoord) = i₋₁ idx₋₁(::Jcoord) = j₋₁ -idx₋₁(::Kcoord) = k₋₁ \ No newline at end of file +idx₋₁(::Kcoord) = k₋₁ diff --git a/src/matrixbuilding.jl b/src/matrixbuilding.jl index b1cbb7e..8c91766 100644 --- a/src/matrixbuilding.jl +++ b/src/matrixbuilding.jl @@ -32,7 +32,7 @@ Build the advection operator Tadv. function buildTadv(; ϕ, gridmetrics, indices, ρ) # default ρ = 1035 kg/m^3 is the value originally used by Chamberlain et al. (2019) - @info "Building Tadv" + @debug "Building Tadv" 𝑖s, 𝑗s, Tvals = upwind_advection_operator_sparse_entries(ϕ, gridmetrics, indices, ρ) N = indices.N @@ -56,7 +56,7 @@ function buildTκH(; gridmetrics, indices, ρ, κH) # Wet mask for horizontal diffusivity ΩH = trues(N) - @info "Building TκH" + @debug "Building TκH" 𝑖s, 𝑗s, Tvals = horizontal_diffusion_operator_sparse_entries(; gridmetrics, indices, κH, ΩH) any(isnan.(Tvals)) && error("TκH contains NaNs.") @@ -85,7 +85,7 @@ function buildTκVML(; mlotst, gridmetrics, indices, κVML) # Wet mask for mixed layer diffusivity Ω = replace(reshape(zt, 1, 1, length(zt)) .< mlotst, missing=>false)[Lwet] - @info "Building TκVML " + @debug "Building TκVML " 𝑖s, 𝑗s, Tvals = vertical_diffusion_operator_sparse_entries(; gridmetrics, indices, κV = κVML, Ω) any(isnan.(Tvals)) && error("TκVML contains NaNs.") @@ -108,7 +108,7 @@ function buildTκVdeep(; mlotst, gridmetrics, indices, κVdeep) # Deep mask for vertical diffusivity Ω = trues(N) # TODO (maybe): make Ωdeep not overlap with ΩML at MLD? - @info "Building TκVdeep" + @debug "Building TκVdeep" 𝑖s, 𝑗s, Tvals = vertical_diffusion_operator_sparse_entries(; gridmetrics, indices, κV = κVdeep, Ω) any(isnan.(Tvals)) && error("TκVdeep contains NaNs.") @@ -129,15 +129,20 @@ function transportmatrix(; ϕ, mlotst, gridmetrics, indices, ρ, κH = 500.0, # m^2/s, κVML = 0.1, # m^2/s, κVdeep = 1e-5, # m^2/s, - Tadv = buildTadv(; ϕ, gridmetrics, indices, ρ), - TκH = buildTκH(; gridmetrics, indices, ρ, κH), - TκVML = buildTκVML(; mlotst, gridmetrics, indices, κVML), - TκVdeep = buildTκVdeep(; mlotst, gridmetrics, indices, κVdeep), + Tadv = nothing, + TκH = nothing, + TκVML = nothing, + TκVdeep = nothing, ) - @info "Building T" + isnothing(Tadv) && (Tadv = buildTadv(; ϕ, gridmetrics, indices, ρ)) + isnothing(TκH) && (TκH = buildTκH(; gridmetrics, indices, ρ, κH)) + isnothing(TκVML) && (TκVML = buildTκVML(; mlotst, gridmetrics, indices, κVML)) + isnothing(TκVdeep) && (TκVdeep = buildTκVdeep(; mlotst, gridmetrics, indices, κVdeep)) - @time T = Tadv + TκH + TκVML + TκVdeep + @debug "Building T" + + T = Tadv + TκH + TκVML + TκVdeep return (; T, Tadv, TκH, TκVML, TκVdeep) end @@ -145,7 +150,15 @@ end - +function preallocate_sparse_entries(sizehint) + 𝑖s = Int64[] + 𝑗s = Int64[] + Tvals = Float64[] + sizehint!(𝑖s, sizehint) + sizehint!(𝑗s, sizehint) + sizehint!(Tvals, sizehint) + return 𝑖s, 𝑗s, Tvals +end # Some personal notes # ϕ are water mass transports, in kg/s @@ -184,13 +197,13 @@ function upwind_advection_operator_sparse_entries(ϕ, gridmetrics, indices, ρ) # Unpack model grid (; v3D, gridtopology) = gridmetrics # Unpack indices - (; Lwet, Lwet3D, C) = indices + (; Lwet, Lwet3D, C, N) = indices - any(isnan.(ρ[Lwet])) && error("ρ contains NaNs") + any(isnan, ρ[Lwet]) && error("ρ contains NaNs") - 𝑖s, 𝑗s, Tvals = Int[], Int[], Float64[] + 𝑖s, 𝑗s, Tvals = preallocate_sparse_entries(6N) # 6 directions for upwind advection - @time for 𝑖 in eachindex(Lwet) + for 𝑖 in eachindex(Lwet) L𝑖 = Lwet[𝑖] C𝑖 = C[L𝑖] i, j, k = C𝑖.I @@ -315,18 +328,19 @@ Return the sparse (i, j, v) for the horizontal diffusion operator TκH. function horizontal_diffusion_operator_sparse_entries(; gridmetrics, indices, κH, ΩH) # Unpack model grid - (; v3D, edge_length_2D, lon, lat, thkcello, gridtopology) = gridmetrics + (; v3D, edge_length_2D, lon, lat, thkcello, gridtopology, distance_to_neighbour_2D) = gridmetrics # Unpack indices - (; wet3D, Lwet, Lwet3D, C) = indices + (; wet3D, Lwet, Lwet3D, C, N) = indices - 𝑖s, 𝑗s, Tvals = Int[], Int[], Float64[] + 𝑖s, 𝑗s, Tvals = preallocate_sparse_entries(8N) # 2 × 4 directions for horizontal diffusion ny = size(wet3D, 2) # Should not be needed once oppdir is dealt by topology functions - @time for 𝑖 in eachindex(Lwet) + for 𝑖 in eachindex(Lwet) ΩH[𝑖] || continue # only continue if inside ΩH L𝑖 = Lwet[𝑖] C𝑖 = C[L𝑖] + C𝑖srf = horizontalindex(C𝑖) i, j, k = C𝑖.I V = v3D[C𝑖] # From West @@ -340,8 +354,7 @@ function horizontal_diffusion_operator_sparse_entries(; gridmetrics, indices, κ 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) + d = distance_to_neighbour_2D[:west][C𝑖srf] pushTmixingvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗W, κH, a, d, V) end end @@ -355,7 +368,7 @@ function horizontal_diffusion_operator_sparse_entries(; gridmetrics, indices, κ 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) + d = distance_to_neighbour_2D[:east][C𝑖srf] pushTmixingvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗E, κH, a, d, V) end end @@ -369,7 +382,7 @@ function horizontal_diffusion_operator_sparse_entries(; gridmetrics, indices, κ 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) + d = distance_to_neighbour_2D[:south][C𝑖srf] pushTmixingvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗S, κH, a, d, V) end end @@ -386,7 +399,7 @@ function horizontal_diffusion_operator_sparse_entries(; gridmetrics, indices, κ 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) + d = distance_to_neighbour_2D[:north][C𝑖srf] pushTmixingvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗N, κH, a, d, V) end end @@ -403,10 +416,6 @@ Pushes the sparse indices and values into (𝑖s, 𝑗s, Tvals) corresponding to """ function pushTmixingvalues!(𝑖s, 𝑗s, Tvals, 𝑖, 𝑗, κ, a, d, V) Tval = κ * a / (d * V) - if d == 0 || V == 0 || isnan(Tval) - @show 𝑖, 𝑗, κ, a, d, V - error() - end push!(𝑖s, 𝑖) push!(𝑗s, 𝑖) push!(Tvals, Tval) @@ -424,13 +433,14 @@ function vertical_diffusion_operator_sparse_entries(; gridmetrics, indices, κV, # Unpack model grid (; v3D, area2D, zt, gridtopology) = gridmetrics # Unpack indices - (; wet3D, Lwet, Lwet3D, C) = indices + (; wet3D, Lwet, Lwet3D, C, N) = indices + + 𝑖s, 𝑗s, Tvals = preallocate_sparse_entries(4 * N) # 2 × 2 directions for vertical diffusion - 𝑖s, 𝑗s, Tvals = Int[], Int[], Float64[] nxyz = size(wet3D) _, _, nz = nxyz - @time for 𝑖 in eachindex(Lwet) + for 𝑖 in eachindex(Lwet) Ω[𝑖] || continue # only continue if inside Ω L𝑖 = Lwet[𝑖] C𝑖 = C[L𝑖] diff --git a/src/velocities.jl b/src/velocities.jl index 314f18c..a5ea67c 100644 --- a/src/velocities.jl +++ b/src/velocities.jl @@ -161,9 +161,9 @@ function facefluxes(umo, vmo, gridmetrics, indices; FillValue) (; C) = indices (; gridtopology) = gridmetrics - @info "Making ϕeast" + @debug "Making ϕeast" ϕeast = replace(umo, NaN=>0.0, FillValue=>0.0) .|> Float64 - @info "Making ϕwest" + @debug "Making ϕwest" # ϕwest[i] is ϕeast[west of i] ϕwest = zeros(size(ϕeast)) for i in eachindex(ϕwest) @@ -172,10 +172,10 @@ function facefluxes(umo, vmo, gridmetrics, indices; FillValue) ϕwest[i] = ϕeast[W] end - @info "Making ϕnorth" + @debug "Making ϕnorth" # Check that south pole ϕnorth is zero (so that it causes no issues with circular shift) ϕnorth = replace(vmo, NaN=>0.0, FillValue=>0.0) .|> Float64 - @info "Making ϕsouth" + @debug "Making ϕsouth" # ϕsouth[i] is ϕnorth[south of i] # Note from BP: This might break if there is a seam at the South Pole ϕsouth = zeros(size(ϕnorth)) @@ -185,7 +185,7 @@ function facefluxes(umo, vmo, gridmetrics, indices; FillValue) ϕsouth[i] = ϕnorth[S] end - @info "Making ϕtop and ϕbottom" + @debug "Making ϕtop and ϕbottom" # Then build ϕtop and ϕbottom from bottom up # Mass conservation implies that # ϕwest + ϕsouth + ϕbottom - ϕeast - ϕnorth - ϕtop = 0 diff --git a/test/interactive.jl b/test/interactive.jl new file mode 100644 index 0000000..e4aeeb2 --- /dev/null +++ b/test/interactive.jl @@ -0,0 +1,113 @@ +using Pkg +Pkg.activate(".") +Pkg.instantiate() +using TestEnv +TestEnv.activate(); + + +using Test +using OceanTransportMatrixBuilder +using NetCDF +using YAXArrays +using GibbsSeaWater +using DimensionalData +using NaNStatistics + +# stdlib +using SparseArrays +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" + +# and for output files +@show version = "v$(pkgversion(OceanTransportMatrixBuilder))" +outputdir = joinpath("plots", version) +mkpath(outputdir) + +# Load datasets +umo_ds = open_dataset(joinpath(inputdir, "umo.nc")) +vmo_ds = open_dataset(joinpath(inputdir, "vmo.nc")) +uo_ds = open_dataset(joinpath(inputdir, "uo.nc")) +vo_ds = open_dataset(joinpath(inputdir, "vo.nc")) +mlotst_ds = open_dataset(joinpath(inputdir, "mlotst.nc")) +volcello_ds = open_dataset(joinpath(inputdir, "volcello.nc")) +areacello_ds = open_dataset(joinpath(inputdir, "areacello.nc")) + + +# Load variables in memory +umo = readcubedata(umo_ds.umo) +vmo = readcubedata(vmo_ds.vmo) +umo_lon = readcubedata(umo_ds.lon) +umo_lat = readcubedata(umo_ds.lat) +vmo_lon = readcubedata(vmo_ds.lon) +vmo_lat = readcubedata(vmo_ds.lat) +mlotst = readcubedata(mlotst_ds.mlotst) +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) # xmip issue: https://github.com/jbusecke/xMIP/issues/369 +lat_vertices = readcubedata(volcello_ds.lat_verticies) # xmip issue: https://github.com/jbusecke/xMIP/issues/369 + +# Make makegridmetrics +gridmetrics = makegridmetrics(; areacello, volcello, lon, lat, lev, lon_vertices, lat_vertices) +(; lon_vertices, lat_vertices, Z3D, v3D, zt, thkcello) = gridmetrics + +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) + +# Load thetato and so to compute density +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) +@test 0 < nanmean(thetao) < 20 +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) +# TODO: Check if this is correct usage of gsw functions! +# Alternatively use a fixed density: +# ρ = 1035.0 # kg/m^3 + +# Below is commented out but should eventually be teseted for neutral/potential density +# @show nanmean(ct) +ρθ = gsw_rho.(so, ct, 1000) +# @show nanmean(ρθ) +# from MATLAB GSW toolbox: +# gsw_rho.(so, ct, p) +# so = Absolute Salinity (g/kg) +# ct = Conservative Temperature (ITS-90) (°C) +# p = sea pressure (dbar) (here using 0 pressure to get potential density + +# Diffusivites +κH = 500.0 # m^2/s +κVML = 0.1 # m^2/s +κVdeep = 1e-5 # m^2/s + +# Make indices +indices = makeindices(gridmetrics.v3D) + +@test all(.!isnan.(ρ[indices.wet3D])) == true + +# Make fuxes from all directions +ϕ = facefluxesfrommasstransport(; umo, vmo, gridmetrics, indices) + +# Make fuxes from all directions from velocities +ϕ_bis = facefluxesfromvelocities(; uo, uo_lon, uo_lat, vo, vo_lon, vo_lat, gridmetrics, indices, ρ) + +# Make transport matrix +(; T, Tadv, TκH, TκVML, TκVdeep) = transportmatrix(; ϕ, mlotst, gridmetrics, indices, ρ, κH, κVML, κVdeep) + +@profview transportmatrix(; ϕ, mlotst, gridmetrics, indices, ρ, κH, κVML, κVdeep) + +@profview_allocs transportmatrix(; ϕ, mlotst, gridmetrics, indices, ρ, κH, κVML, κVdeep) sample_rate=0.9 \ No newline at end of file