How should I deal with post processing data after running on GPUs?

[anpa5279]

Hello!

Currently, I am trying to figure out what is the best way to plot data from .jld2 files using JLD2OutputWriter after running the below code on a HPC. I modified the Langmuir Turbulence example so it could run on GPUs on NCAR's Derecho:

using Pkg
using MPI
using Oceananigans
using Oceananigans.DistributedComputations
using Oceananigans.Units: minute, minutes, hours
using Oceananigans.BuoyancyFormulations: g_Earth

mutable struct Params
    Nx::Int
    Ny::Int     # number of points in each of horizontal directions
    Nz::Int     # number of points in the vertical direction
    Lx::Float64
    Ly::Float64     # (m) domain horizontal extents
    Lz::Float64     # (m) domain depth 
    amplitude::Float64 # m
    wavelength::Float64 # m
    τx::Float64 # m² s⁻², surface kinematic momentum flux
    Jᵇ::Float64 # m² s⁻³, surface buoyancy flux
    N²::Float64 # s⁻², initial and bottom buoyancy gradient
    initial_mixed_layer_depth::Float64 #m
end

params = Params(32, 32, 32, 128.0, 128.0, 64.0, 0.8, 60.0, -3.72e-5, 2.307e-8, 1.936e-5, 33.0)


arch = Distributed(GPU())
rank = arch.local_rank
Nranks = MPI.Comm_size(arch.communicator)
println("Hello from process $rank out of $Nranks")

grid = RectilinearGrid(arch; size=(params.Nx, params.Ny, params.Nz), extent=(params.Lx, params.Ly, params.Lz))
@show grid

const wavenumber = 2π / params.wavelength # m⁻¹
const frequency = sqrt(g_Earth * wavenumber) # s⁻¹
@show wavenumber, frequency

const vertical_scale = params.wavelength / 4π
@show vertical_scale

const Uˢ = params.amplitude^2 * wavenumber * frequency # m s⁻¹
@show Uˢ

@inline uˢ(z) = Uˢ * exp(z / vertical_scale)

@inline ∂z_uˢ(z, t) = 1 / vertical_scale * Uˢ * exp(z / vertical_scale)

u_boundary_conditions = FieldBoundaryConditions(top = FluxBoundaryCondition(params.τx))
@show u_boundary_conditions

b_boundary_conditions = FieldBoundaryConditions(top = FluxBoundaryCondition(params.Jᵇ),
                                                bottom = GradientBoundaryCondition(params.N²))
@show b_boundary_conditions

coriolis = FPlane(f=1e-4) # s⁻¹

model = NonhydrostaticModel(; grid, coriolis,
                            advection = WENO(),
                            timestepper = :RungeKutta3,
                            tracers = :b,
                            buoyancy = BuoyancyTracer(),
                            closure = AnisotropicMinimumDissipation(),
                            stokes_drift = UniformStokesDrift(∂z_uˢ=∂z_uˢ),
                            boundary_conditions = (u=u_boundary_conditions, b=b_boundary_conditions))
@show model

@inline Ξ(z) = randn() * exp(z / 4)

@inline stratification(z) = z < - params.initial_mixed_layer_depth ? params.N² * z : params.N² * (-params.initial_mixed_layer_depth)

@inline bᵢ(x, y, z) = stratification(z) + 1e-1 * Ξ(z) * params.N² * model.grid.Lz

u★ = sqrt(abs(params.τx))
@inline uᵢ(x, y, z) = u★ * 1e-1 * Ξ(z)
@inline wᵢ(x, y, z) = u★ * 1e-1 * Ξ(z)

set!(model, u=uᵢ, w=wᵢ, b=bᵢ)

simulation = Simulation(model, Δt=45.0, stop_time=4hours)
@show simulation

conjure_time_step_wizard!(simulation, cfl=1.0, max_Δt=1minute)

output_interval = 5minutes

fields_to_output = merge(model.velocities, model.tracers, (; νₑ=model.diffusivity_fields.νₑ))

simulation.output_writers[:fields] = JLD2OutputWriter(model, fields_to_output,
                                                      schedule = TimeInterval(output_interval),
                                                      filename = "langmuir_turbulence_fields_$rank.jld2",
                                                      overwrite_existing = true,
                                                      with_halos = true)

u, v, w = model.velocities
b = model.tracers.b

U = Average(u, dims=(1, 2))
V = Average(v, dims=(1, 2))
B = Average(b, dims=(1, 2))
wu = Average(w * u, dims=(1, 2))
wv = Average(w * v, dims=(1, 2))

simulation.output_writers[:averages] = JLD2OutputWriter(model, (; U, V, B, wu, wv),
                                                        schedule = AveragedTimeInterval(output_interval, window=2minutes),
                                                        filename = "langmuir_turbulence_averages_$rank.jld2",
                                                        overwrite_existing = true,
                                                        with_halos = true)

run!(simulation)

time_series = (;
    w = FieldTimeSeries("langmuir_turbulence_fields_$rank.jld2", "w"),
    u = FieldTimeSeries("langmuir_turbulence_fields_$rank.jld2", "u"),
    B = FieldTimeSeries("langmuir_turbulence_averages_$rank.jld2", "B"),
    U = FieldTimeSeries("langmuir_turbulence_averages_$rank.jld2", "U"),
    V = FieldTimeSeries("langmuir_turbulence_averages_$rank.jld2", "V"),
    wu = FieldTimeSeries("langmuir_turbulence_averages_$rank.jld2", "wu"),
    wv = FieldTimeSeries("langmuir_turbulence_averages_$rank.jld2", "wv"))
                        
times = time_series.w.times

n = Observable(1)

wxy_title = @lift string("w(x, y, t) at z=-8 m and t = ", prettytime(times[$n]))
wxz_title = @lift string("w(x, z, t) at y=0 m and t = ", prettytime(times[$n]))
uxz_title = @lift string("u(x, z, t) at y=0 m and t = ", prettytime(times[$n]))

fig = Figure(size = (850, 850))

ax_B = Axis(fig[1, 4];
            xlabel = "Buoyancy (m s⁻²)",
            ylabel = "z (m)")

ax_U = Axis(fig[2, 4];
            xlabel = "Velocities (m s⁻¹)",
            ylabel = "z (m)",
            limits = ((-0.07, 0.07), nothing))

ax_fluxes = Axis(fig[3, 4];
                xlabel = "Momentum fluxes (m² s⁻²)",
                ylabel = "z (m)",
                limits = ((-3.5e-5, 3.5e-5), nothing))

ax_wxy = Axis(fig[1, 1:2];
            xlabel = "x (m)",
            ylabel = "y (m)",
            aspect = DataAspect(),
            limits = ((0, grid.Lx), (0, grid.Ly)),
            title = wxy_title)

ax_wxz = Axis(fig[2, 1:2];
            xlabel = "x (m)",
            ylabel = "z (m)",
            aspect = AxisAspect(2),
            limits = ((0, grid.Lx), (-grid.Lz, 0)),
            title = wxz_title)

ax_uxz = Axis(fig[3, 1:2];
            xlabel = "x (m)",
            ylabel = "z (m)",
            aspect = AxisAspect(2),
            limits = ((0, grid.Lx), (-grid.Lz, 0)),
            title = uxz_title)


wₙ = @lift time_series.w[$n]
uₙ = @lift time_series.u[$n]
Bₙ = @lift view(time_series.B[$n], 1, 1, :)
Uₙ = @lift view(time_series.U[$n], 1, 1, :)
Vₙ = @lift view(time_series.V[$n], 1, 1, :)
wuₙ = @lift view(time_series.wu[$n], 1, 1, :)
wvₙ = @lift view(time_series.wv[$n], 1, 1, :)

k = searchsortedfirst(znodes(grid, Face(); with_halos=true), -8)
wxyₙ = @lift view(time_series.w[$n], :, :, k)
wxzₙ = @lift view(time_series.w[$n], :, 1, :)
uxzₙ = @lift view(time_series.u[$n], :, 1, :)

wlims = (-0.03, 0.03)
ulims = (-0.05, 0.05)

lines!(ax_B, Bₙ)

lines!(ax_U, Uₙ; label = L"\bar{u}")
lines!(ax_U, Vₙ; label = L"\bar{v}")
axislegend(ax_U; position = :rb)

lines!(ax_fluxes, wuₙ; label = L"mean $wu$")
lines!(ax_fluxes, wvₙ; label = L"mean $wv$")
axislegend(ax_fluxes; position = :rb)

hm_wxy = heatmap!(ax_wxy, wxyₙ;
                colorrange = wlims,
                colormap = :balance)

Colorbar(fig[1, 3], hm_wxy; label = "m s⁻¹")

hm_wxz = heatmap!(ax_wxz, wxzₙ;
                colorrange = wlims,
                colormap = :balance)

Colorbar(fig[2, 3], hm_wxz; label = "m s⁻¹")

ax_uxz = heatmap!(ax_uxz, uxzₙ;
                colorrange = ulims,
                colormap = :balance)

Colorbar(fig[3, 3], ax_uxz; label = "m s⁻¹")

fig

frames = 1:length(times)

record(fig, "langmuir_turbulence.mp4", frames, framerate=8) do i
    n[] = i
end

When I run this on Derecho, it fails at generating times_series and I get this:
ERROR: ERROR: ERROR: LoadError: LoadError: LoadError: ERROR: LoadError: BoundsError: attempt to access BoundsError: attempt to access BoundsError: attempt to access BoundsError: attempt to access 0-element Vector{String} at index [1]

The code above creates the desired .jld2 files, but I cannot plot these outputs to verify that what I did worked.

I also tried plotting locally on my computer, but have had no success. #3978 seems to be what I am looking for, but it is not clear to me how they initialized FieldTimeSeries to avoid errors. I tried debugging by myself with no luck. Should I have a separate file that writes the grid and time information? Should I change the way I output the data? It is currently outputting by rank. This complicates things while trying to plot locally.

In summary, I want to know how I should deal with post processing my data.

Thanks in advance for the help!

[glwagner]

Good question and not explained in the docs! One way to make visualizations (and my preferred way) is to download the data to my local machine, and then to make the plots there. This means commenting out or deleting the code after run!(simulation), and then putting the plotting code in a new script and running that on your local machine.

The error you are getting, however, indicates that the file does not exist or does not have any data in it. This is an additional problem. Can you check to make sure that the file is there? It is possible that the simulation has not run. You can also inspect the file directly using JLD2. See the JLD2 docs for how to do that.

Note, you can visualize on Derecho, but it is not very efficient. The issue is that you will have to launch a new job, and use CairoMakie, which can be slow compared to GLMakie. You cannot use GLMakie unless you have a display. (Definitely check out the Makie documentation for more info about that.)

I recently noticed that this is the error we get when files don't exist and I think the situation can be improved.

[anpa5279]

Hello! I took your advice and plotted locally. I realized today that using include("plotting.jl") executes the code differently than just using julia's interactive interface (For example, in my script, FieldTimeSeries(fld_file, "w") would fail but not when I ran FieldTimeSeries(fld_file, "w") in the interactive interface). To resolve this, I still ran include("plotting.jl") but made my entire code a function called plot() so I could run that function in julia's interactive interface. ##3211 helped me understand how to setup FieldTimeSeries locally.

Also, in my code above I changed with_halos = true to with_halos = false because I am not doing any calculations that need the ghost cells.

That being said this is my current code for plotting only:
`using Pkg
using CairoMakie
using Oceananigans
using Oceananigans.Units: minute, minutes, hours

function plot()
Nranks = 4

fld_file="outputs/langmuir_turbulence_fields_0_rank0.jld2"
averages_file="outputs/langmuir_turbulence_averages_0_rank0.jld2"

w_temp = FieldTimeSeries(fld_file, "w")
u_temp = FieldTimeSeries(fld_file, "u")
B_temp = FieldTimeSeries(averages_file, "B")
U_temp = FieldTimeSeries(averages_file, "U")
V_temp = FieldTimeSeries(averages_file, "V")
wu_temp = FieldTimeSeries(averages_file, "wu")
wv_temp = FieldTimeSeries(averages_file, "wv")

Lx = Nranks * u_temp.grid.Lx
Ly = u_temp.grid.Ly
Lz = u_temp.grid.Lz
Nx = Nranks * u_temp.grid.Nx
Ny = u_temp.grid.Ny
Nz = u_temp.grid.Nz
Nt = length(u_temp.times)
grid = RectilinearGrid(size = (Nx, Ny, Nz), extent = (Lx, Ly, Lz))
times = u_temp.times

println(Nx)

w_data = Array{Float64}(undef, (Nx, Ny, Nz + 1, Nt)) #because face value
u_data = Array{Float64}(undef, (Nx, Ny, Nz, Nt))
B_data = Array{Float64}(undef, (1, 1, Nz, Nt))
U_data = Array{Float64}(undef, (1, 1, Nz, Nt))
V_data = Array{Float64}(undef, (1, 1, Nz, Nt))
wu_data = Array{Float64}(undef, (1, 1, Nz + 1, Nt))  
wv_data = Array{Float64}(undef, (1, 1, Nz + 1, Nt))
B_data .= 0
U_data .= 0
V_data .= 0
wu_data .= 0
wv_data .= 0

p = 1
w_data[p:u_temp.grid.Nx, :, :, :] .= w_temp.data
u_data[p:u_temp.grid.Nx, :, :, :] .= u_temp.data
B_data .= B_data .+ B_temp.data
U_data .= U_data .+ U_temp.data
V_data .= V_data .+ V_temp.data
wu_data .= wu_data .+ wu_temp.data
wv_data .= wv_data .+ wv_temp.data

for i in 1:Nranks-1

    p = p + u_temp.grid.Nx

    println("Loading rank $i")

    fld_file="outputs/langmuir_turbulence_fields_$(i)_rank$(i).jld2"
    averages_file="outputs/langmuir_turbulence_averages_$(i)_rank$(i).jld2"

    w_temp = FieldTimeSeries(fld_file, "w")
    u_temp = FieldTimeSeries(fld_file, "u")
    B_temp = FieldTimeSeries(averages_file, "B")
    U_temp = FieldTimeSeries(averages_file, "U")
    V_temp = FieldTimeSeries(averages_file, "V")
    wu_temp = FieldTimeSeries(averages_file, "wu")
    wv_temp = FieldTimeSeries(averages_file, "wv")

    println(size( w_data[p:p + w_temp.grid.Nx-1, :, :, :] ))
    println(size( w_temp.data ))
    w_data[p:p + w_temp.grid.Nx - 1, :, :, :] .= w_temp.data
    u_data[p:p + u_temp.grid.Nx - 1, :, :, :] .= u_temp.data
    B_data .= B_data .+ B_temp.data
    U_data .= U_data .+ U_temp.data
    V_data .= V_data .+ V_temp.data
    wu_data .= wu_data .+ wu_temp.data
    wv_data .= wv_data .+ wv_temp.data
    
end

B_data = B_data ./ Nranks
U_data = U_data ./ Nranks
V_data = V_data ./ Nranks
wu_data = wu_data ./ Nranks
wv_data = wv_data ./ Nranks
#putting everything back into FieldTimeSeries
w = FieldTimeSeries{Center, Center, Face}(grid, times)
u = FieldTimeSeries{Face, Center, Center}(grid, times)
B = FieldTimeSeries{Center, Center, Center}(grid, times)
U = FieldTimeSeries{Center, Center, Center}(grid, times)
V = FieldTimeSeries{Center, Center, Center}(grid, times)
wu = FieldTimeSeries{Center, Center, Face}(grid, times)
wv = FieldTimeSeries{Center, Center, Face}(grid, times)

w .= w_data
u .= u_data
B .= B_data
U .= U_data
V .= V_data
wu .= wu_data
wv .= wv_data

#begin plotting
n = Observable(1)

wxy_title = @lift string("w(x, y, t) at z=-8 m and t = ", prettytime(times[$n]))
wxz_title = @lift string("w(x, z, t) at y=0 m and t = ", prettytime(times[$n]))
uxz_title = @lift string("u(x, z, t) at y=0 m and t = ", prettytime(times[$n]))

fig = Figure(size = (850, 850))

ax_B = Axis(fig[1, 4];
            xlabel = "Buoyancy (m s⁻²)",
            ylabel = "z (m)")

ax_U = Axis(fig[2, 4];
            xlabel = "Velocities (m s⁻¹)",
            ylabel = "z (m)",
            limits = ((-0.07, 0.07), nothing))

ax_fluxes = Axis(fig[3, 4];
                xlabel = "Momentum fluxes (m² s⁻²)",
                ylabel = "z (m)",
                limits = ((-3.5e-5, 3.5e-5), nothing))

ax_wxy = Axis(fig[1, 1:2];
            xlabel = "x (m)",
            ylabel = "y (m)",
            aspect = DataAspect(),
            limits = ((0, Lx), (0, Ly)),
            title = wxy_title)

ax_wxz = Axis(fig[2, 1:2];
            xlabel = "x (m)",
            ylabel = "z (m)",
            aspect = AxisAspect(2),
            limits = ((0, Lx), (-Lz, 0)),
            title = wxz_title)

ax_uxz = Axis(fig[3, 1:2];
            xlabel = "x (m)",
            ylabel = "z (m)",
            aspect = AxisAspect(2),
            limits = ((0, Lx), (-Lz, 0)),
            title = uxz_title)


wₙ = @lift w[$n]
uₙ = @lift u[$n]
Bₙ = @lift view(B[$n], 1, 1, :)
Uₙ = @lift view(U[$n], 1, 1, :)
Vₙ = @lift view(V[$n], 1, 1, :)
wuₙ = @lift view(wu[$n], 1, 1, :)
wvₙ = @lift view(wv[$n], 1, 1, :)

k = searchsortedfirst(znodes(grid, Face(); with_halos=true), -8)
wxyₙ = @lift view(w[$n], :, :, k)
wxzₙ = @lift view(w[$n], :, 1, :)
uxzₙ = @lift view(u[$n], :, 1, :)

wlims = (-0.03, 0.03)
ulims = (-0.05, 0.05)

lines!(ax_B, Bₙ)

lines!(ax_U, Uₙ; label = L"\bar{u}")
lines!(ax_U, Vₙ; label = L"\bar{v}")
axislegend(ax_U; position = :rb)

lines!(ax_fluxes, wuₙ; label = L"mean $wu$")
lines!(ax_fluxes, wvₙ; label = L"mean $wv$")
axislegend(ax_fluxes; position = :rb)

hm_wxy = heatmap!(ax_wxy, wxyₙ;
                colorrange = wlims,
                colormap = :balance)

Colorbar(fig[1, 3], hm_wxy; label = "m s⁻¹")

hm_wxz = heatmap!(ax_wxz, wxzₙ;
                colorrange = wlims,
                colormap = :balance)

Colorbar(fig[2, 3], hm_wxz; label = "m s⁻¹")

ax_uxz = heatmap!(ax_uxz, uxzₙ;
                colorrange = ulims,
                colormap = :balance)

Colorbar(fig[3, 3], ax_uxz; label = "m s⁻¹")

fig

frames = 1:length(times)

record(fig, "langmuir_turbulence.mp4", frames, framerate=8) do i
    n[] = i
end

end`
Here is my result:

langmuir_turbulence.mp4

However, it looks different than my local run:

local_langmuir_turbulence.mp4

Is this because of the periodic boundary conditions I have for x and y? Where did I go wrong?

[glwagner]

The visualizations look similar to me. What's the problem?

It's better to save with halos. I didn't follow your reason for omitting them.

[anpa5279]

Oh, that was a silly question. I forgot that randn() was present in Ξ(z) = randn() * exp(z / 4). It makes sense that they are similar, but not identical.

As for the halos, I was not accounting for ghost cells when extracting my data for post processing and it was causing errors. I am new to julia, so to simplify things, I omitted them. In the future, I will have them in my output files. Hopefully that clears things up!

Thank you @glwagner so much for your help! The discussion posts have been really helpful and your responses have been super informative. I really appreciate it.

[glwagner]

Ah ok, typically if you use FieldTimeSeries for post-processing, it is better to include halos (because they are accounted for automatically, and you may make use of them for certain calculations). If you are using something else, then you likely want to omit them to avoid the offset math.

[anpa5279]

Hi, I have a follow up question. Why do we need to include the halos in the output file?

I am currently running into an error with w_temp = FieldTimeSeries(fld_file, "w") but not with U_temp = FieldTimeSeries(averages_file, "U"). Both these files contain halos, but with w_temp = FieldTimeSeries(fld_file, "w") I get ERROR: type CPU has no field ranks as an error when trying to get the data.

I ran f = jldopen(fld_file) and fa = jldopen(averages_file) to see the differences in f["grid"] and fa["grid"] but did not find anything alarming. Why would I get an error with one file but not the other?

Since this discussion I updated the Stokes drift and added to the fld_file to include the Stokes drift parameters I am interested in:

function save_IC!(file, model)
                 file["IC/friction_velocity"] = u_f
                 file["IC/stokes_velocity"] = stokes_velocity(-grid.z.Δᵃᵃᶜ/2, p.u₁₀)[1]
                 file["IC/wind_speed"] = p.u₁₀
    return nothing
end
simulation.output_writers[:fields] = JLD2OutputWriter(model, fields_to_output,
                                                      schedule = TimeInterval(output_interval),
                                                      filename = "langmuir_turbulence_fields_$rank.jld2",
                                                      overwrite_existing = true, 
                                                      init = save_IC!)

I am unsure what would cause the error and would appreciate some input. I do not think including the halos is the cause of my error, but I am curious if not including the halos causes an error in the output file.