Skip to content

Commit

Permalink
Fix allocations + add interactive profiler
Browse files Browse the repository at this point in the history
  • Loading branch information
briochemc committed Oct 30, 2024
1 parent 5af59c4 commit ded800d
Show file tree
Hide file tree
Showing 6 changed files with 167 additions and 41 deletions.
2 changes: 1 addition & 1 deletion Project.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "OceanTransportMatrixBuilder"
uuid = "c2b4a04e-6049-4fc4-aa6a-5508a29a1e1c"
authors = ["Benoit Pasquier <[email protected]> and contributors"]
version = "0.5.0"
version = "0.5.1"

[deps]
Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
Expand Down
9 changes: 6 additions & 3 deletions src/gridcellgeometry.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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

2 changes: 1 addition & 1 deletion src/gridtopology.jl
Original file line number Diff line number Diff line change
Expand Up @@ -115,4 +115,4 @@ idx₊₁(::Jcoord) = j₊₁
idx₊₁(::Kcoord) = k₊₁
idx₋₁(::Icoord) = i₋₁
idx₋₁(::Jcoord) = j₋₁
idx₋₁(::Kcoord) = k₋₁
idx₋₁(::Kcoord) = k₋₁
72 changes: 41 additions & 31 deletions src/matrixbuilding.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.")
Expand Down Expand Up @@ -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.")
Expand All @@ -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.")
Expand All @@ -129,23 +129,36 @@ 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





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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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𝑖]
Expand Down
10 changes: 5 additions & 5 deletions src/velocities.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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))
Expand All @@ -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
Expand Down
113 changes: 113 additions & 0 deletions test/interactive.jl
Original file line number Diff line number Diff line change
@@ -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

2 comments on commit ded800d

@briochemc
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register

Release notes:

  • "Fix" allocations by using sizehint!. TODO: pre-allocate the sparse matrices themselves, or some other process to sum over indices/values
  • Add small interactive profiler script in test/

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/118354

Tagging

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.5.1 -m "<description of version>" ded800d2c365aa03887a356d60229be45a3a7420
git push origin v0.5.1

Please sign in to comment.