ReCo.jl/src/data.jl

using StaticArrays: SVector
using LoopVectorization: @turbo
using JLD2: JLD2

struct Bundle
    t::Vector{Float64}
    c::Matrix{Vector{Float64}}
    φ::Matrix{Float64}
end

function Bundle(N::Int64, n_snapshots::Int64)
    return Bundle(
        zeros(n_snapshots),
        [zeros(2) for i in 1:N, j in 1:n_snapshots],
        zeros(N, n_snapshots),
    )
end

function first_n_snapshots(bundle::Bundle, n::Int64)
    return Bundle(bundle.t[1:n], bundle.c[:, 1:n], bundle.φ[:, 1:n])
end

function save_snapshot!(
    bundle::Bundle, n_snapshot::Int64, t::Float64, particles::Vector{Particle}
)
    bundle.t[n_snapshot] = t

    bundle_c = bundle.c
    bundle_φ = bundle.φ

    @simd for p in particles
        p_id = p.id
        p_c = p.c

        bundle_c_id = bundle_c[p_id, n_snapshot]

        @turbo for j in 1:2
            bundle_c_id[j] = p_c[j]
        end

        bundle_φ[p_id, n_snapshot] = p.φ
    end

    return nothing
end

function set_status(dir::String, n_bundles::Int64, T::Float64)
    open("$dir/status.json", "w") do f
        JSON3.write(f, Dict("n_bundles" => n_bundles, "T" => round(T; digits=3)))
    end

    return nothing
end

function save_bundle(dir::String, bundle::Bundle, n::Int64, T::Float64)
    bundles_dir = "$dir/bundles"
    mkpath(bundles_dir)

    JLD2.save_object("$bundles_dir/bundle_$n.jld2", bundle)

    set_status(dir, n, T)

    return nothing
end