ReCo.jl/src/data.jl

using Dates: Dates
using JLD2: JLD2

struct SimConsts
    n_particles::Int64
    v₀::Float64
    δt::Float64
    packing_fraction::Float64
    μₜ::Float64
    Dₜ::Float64
    particle_radius::Float64
    Dᵣ::Float64
    σ::Float64
    ϵ::Float64
    interaction_radius::Float64
    skin_radius::Float64
    n_steps_before_verlet_list_update::Int64
    grid_n::Int64
    half_box_len::Float64
    grid_box_width::Float64
end

struct RunParams
    duration::Float64
    snapshot_at::Float64
    seed::Int64
    n_bundle_snapshots::Int64
    integration_steps::Int64
    n_steps_before_snapshot::Int64
    n_snapshots::Int64
    T::Float64
    T0::Float64
    start_datetime::Dates.DateTime
end

struct SimState
    n_bundles::Int64
    T::Float64
end

struct Bundle
    t::Vector{Float64}
    c::Matrix{SVector{2,Float64}}
    φ::Matrix{Float64}

    n_particles::Int64
    n_snapshots::Int64
end

function Bundle(n_particles::Int64, n_snapshots::Int64)
    return Bundle(
        Vector{Float64}(undef, n_snapshots),
        Matrix{SVector{2,Float64}}(undef, (n_particles, n_snapshots)),
        Matrix{Float64}(undef, (n_particles, n_snapshots)),
        n_particles,
        n_snapshots,
    )
end

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

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

    @simd for p in particles
        bundle.c[p.id, n_snapshot] = p.c

        bundle.φ[p.id, n_snapshot] = p.φ
    end

    return nothing
end

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

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

    JLD2.save_object("$sim_dir/sim_state.jld2", SimState(n_bundle, round(T; digits=3)))

    return nothing
end

function sorted_bundle_paths(sim_dir::String; rev::Bool=false)
    bundles_dir = "$sim_dir/bundles"
    bundle_paths = readdir(bundles_dir; join=true, sort=false)

    n_bundles = length(bundle_paths)

    bundle_nums = Vector{Int64}(undef, n_bundles)

    @simd for bundle_ind in 1:n_bundles
        bundle_path = bundle_paths[bundle_ind]
        bundle_name = splitdir(bundle_path)[2]
        bundle_num_string = splitext(bundle_name)[1]
        bundle_nums[bundle_ind] = parse(Int64, bundle_num_string)
    end

    sort_perm = sortperm(bundle_nums; rev=rev)

    return bundle_paths[sort_perm]
end

function append_bundle_properties!(
    vs::NTuple{N,Vector},
    properties::NTuple{N,Symbol},
    sim_dir::String;
    particle_slice=nothing,
    snapshot_slice,
    first_bundle::Int64=1,
) where {N}
    bundle_paths = ReCo.sorted_bundle_paths(sim_dir)

    for bundle_ind in first_bundle:length(bundle_paths)
        bundle::Bundle = JLD2.load_object(bundle_paths[bundle_ind])

        for (v_ind, v) in enumerate(vs)
            property = properties[v_ind]
            bundle_property = getproperty(bundle, property)

            if !isnothing(particle_slice)
                bundle_property_view = view(bundle_property, particle_slice, snapshot_slice)
            else
                bundle_property_view = view(bundle_property, snapshot_slice)
            end

            append!(v, bundle_property_view)
        end
    end

    return nothing
end

struct BundlesInfo
    n_bundles::Int64
    total_n_snapshots::Int64
    bundle_n_snapshots::Vector{Int64}
    sorted_bundle_paths::Vector{String}
end

function BundlesInfo(sim_dir::String)
    bundle_paths = sorted_bundle_paths(sim_dir)

    bundle_n_snapshots = Vector{Int64}(undef, length(bundle_paths))

    n_bundles = length(bundle_paths)

    total_n_snapshots = 0

    for bundle_ind in 1:n_bundles
        bundle_path = bundle_paths[bundle_ind]
        bundle::Bundle = JLD2.load_object(bundle_path)

        total_n_snapshots += bundle.n_snapshots
        bundle_n_snapshots[bundle_ind] = bundle.n_snapshots
    end

    return BundlesInfo(n_bundles, total_n_snapshots, bundle_n_snapshots, bundle_paths)
end

function load_sim_consts(sim_dir::String)
    sim_consts::ReCo.SimConsts = JLD2.load_object("$sim_dir/sim_consts.jld2")

    return sim_consts
end

function get_bundle_to_snapshot(
    bundles_info::BundlesInfo, snapshot::Int64
)::Tuple{Bundle,Int64}
    @assert 0 < snapshot <= bundles_info.total_n_snapshots

    snapshot_counter = 0

    for bundle_ind in 1:(bundles_info.n_bundles)
        bundle_n_snapshots = bundles_info.bundle_n_snapshots[bundle_ind]
        snapshot_counter += bundle_n_snapshots
        if snapshot_counter >= snapshot
            bundle_snapshot = bundle_n_snapshots - snapshot_counter + snapshot
            bundle_path = bundles_info.sorted_bundle_paths[bundle_ind]

            bundle::ReCo.Bundle = JLD2.load_object(bundle_path)

            return (bundle, bundle_snapshot)
        end
    end
end