Post-processing of tracer output

jl

using Plots
using ADIOS2
using Random
#using BenchmarkTools

sim_name = "pe459_v7"
#sim_path = f"/global/cfs/cdirs/m499/rkube/xgc1-runs/tracers-em/{sim_name}"
sim_path = "/global/cscratch1/sd/rkube/xgc_runs/tracers-em/$(sim_name)"

# Number of particles that we track.
num_ptl = 8000
# Time steps
tstep_rg = 2750:4330

# Sub-cycles per time-step
ncycle_e = 200 
ncycle_i = 1 

adios = adios_init_serial()

"""
    Extract GIDs of particle we are tracking from the first time-step.
"""
function get_gids_tracked(num_ptl)
    # Open .bp file, 
    io = declare_io(adios, "IO_first_$(rand(Int))")
    engine = open(io, "$(sim_path)/xgc.particle.ele.02750.bp", mode_read)
    gid_var = inquire_variable(io, "gid")
    gid_tracked = Array{Int64}(undef, shape(gid_var)...);
    get(engine, gid_var, gid_tracked)
    perform_gets(engine)
    close(engine)
    # These are the GIDs of the particles we are tracking
    # Pick num_ptl random particles and return their GID
    tracked_gids = shuffle(gid_tracked)[1:num_ptl]
end


# The GIDs are located at different places in subsequent time-steps.
# This expression returns their index in the gid variable at the new time-step.

function read_collected!(output_arr, tstep, tracked_gids)
    # output_arr: bp-file processed data is stored in this array
    # tstep: Int, time step
    # tracked_gids: Index array with indices corresponding to the GIDs that we track
    #   
    #io = declare_io(adios, "IO$(rand(Int, 10000000))")
    io = declare_io(adios, "IO$(tstep)")
    engine = open(io, "$(sim_path)/xgc.particle.ele.0$(tstep).bp", mode_read)
    phase_var = inquire_variable(io, "phase")
    gid_var = inquire_variable(io, "gid")
    nsteps = steps(phase_var)
    # Phase over all sub-cycles
    #ephase = zeros(Float32, 2, length(tracked_gids), nsteps);
    # Read-out for GIDs
    gid = Array{Int64}(undef, shape(gid_var)...);
    # Read-out for phase
    tmp = Array{Float32}(undef, shape(phase_var)...);

    # Fetch phase data from all steps
    for icycle = 1:nsteps
        begin_step(engine)
        # At the first step we need to read the gid.
        if icycle == 1
            get(engine, gid_var, gid)
        end 
        # Read the phase into a temporary array.
        get(engine, phase_var, tmp) 
        end_step(engine)
        perform_gets(engine)
        # Access all variables only after perform_gets has been called
        # First, obtain the indices of the tracked_gids in the GIDs of the current step.
        tracked_gid_idx = [findfirst(gid .== g) for g in tracked_gids]   
        tracked_gid_idx[tracked_gid_idx .== nothing] .= 1
        if icycle == 1
            @show sum(tracked_gid_idx .== 1)
        end 

        # Store the tracked particles only in gid. Use tracked_gid_idx 
        # to permute the scattered, tracked particles in-order in ephase
        output_arr[:, :, icycle] = Float32.(tmp[1:2, tracked_gid_idx])
    end    
    close(engine)
end

# Track 8000 particles
gids_tracked = get_gids_tracked(8000);

# Store only R and Z
phase_all = zeros(Float32, 2, num_ptl, ncycle_e, length(tstep_rg));

for (idx, t) in enumerate(tstep_rg)
    @show t
    read_collected!(@view(phase_all[:, :, :, idx]), t, gids_tracked)
end


"""
    Store phase_all in a bp file
"""
function write_phase(phase)
    io = declare_io(adios, "io_ph2")
    engine = open(io, "$(sim_name)_phase_all.bp", mode_write)
    phase_var = define_variable(io, "phase_all", phase)
    put!(engine, phase_var, phase)
    perform_puts!(engine)
    close(engine)
end
write_phase(phase_all)