ReCo.jl/visualization/Animation.jl

module Animation

export animate

using ColorSchemes: ColorSchemes
using GLMakie
using JLD2: JLD2
using LaTeXStrings: @L_str
using Observables: Observables
using ProgressMeter: ProgressMeter

using ReCo: ReCo

include("common.jl")

const DEFAULT_FRAMERATE = 10
const DEFAULT_SHOW_CENTER_OF_MASS = false
const DEFAULT_SHOW_INTERACTION_CIRCLE = false
const DEFAULT_SHOW_SKIN_CIRCLE = false
const DEFAULT_SHOW_FRAME_DIFF = false
const DEFAULT_SHOW_PROGRESS = false

function update_values!(
    args, sim_consts::ReCo.SimConsts, bundle::ReCo.Bundle, bundle_snapshot::Integer
)
    @simd for particle_ind in 1:(sim_consts.n_particles)
        c = bundle.c[particle_ind, bundle_snapshot]
        φ = bundle.φ[particle_ind, bundle_snapshot]

        args.particle_xs[][particle_ind] = c[1]
        args.particle_ys[][particle_ind] = c[2]

        color = angle_color(φ, args.color_scheme)
        args.particle_colors[][particle_ind] = color

        if args.show_interaction_circle
            args.interaction_colors[][particle_ind] = ColorSchemes.ColorTypes.RGBA(
                color, 0.08
            )
        end

        if args.show_skin_circle
            args.skin_colors[][particle_ind] = ColorSchemes.ColorTypes.RGBA(color, 0.04)
        end

        if args.show_frame_diff
            first_ind = 2 * particle_ind - 1
            second_ind = 2 * particle_ind

            first_bundle_snapshot = bundle_snapshot - 1
            second_bundle_snapshot = bundle_snapshot

            if bundle_snapshot == 1
                first_bundle_snapshot = 1
            end

            args.segment_xs[][first_ind] = bundle.c[particle_ind, first_bundle_snapshot][1]
            args.segment_xs[][second_ind] = bundle.c[particle_ind, second_bundle_snapshot][1]

            args.segment_ys[][first_ind] = bundle.c[particle_ind, first_bundle_snapshot][2]
            args.segment_ys[][second_ind] = bundle.c[particle_ind, second_bundle_snapshot][2]
        end
    end

    if args.show_center_of_mass
        center_of_mass = ReCo.center_of_mass(
            view(bundle.c, :, bundle_snapshot), sim_consts.half_box_len
        )
        args.center_of_mass_point[] = Point(center_of_mass)
    end
end

function init_animation(args, sim_consts::ReCo.SimConsts, first_bundle::ReCo.Bundle)
    update_values!(args, sim_consts, first_bundle, 1)

    scatter!(
        args.ax,
        args.particle_xs,
        args.particle_ys;
        markersize=2 * sim_consts.particle_radius,
        markerspace=:data,
        color=args.particle_colors,
    )

    if args.show_center_of_mass
        scatter!(
            args.ax,
            args.center_of_mass_point;
            markersize=6 * sim_consts.particle_radius,
            markerspace=:data,
            color=ColorSchemes.ColorTypes.RGBA(1.0, 1.0, 1.0, 0.6),
        )
    end

    if args.show_interaction_circle
        scatter!(
            args.ax,
            args.particle_xs,
            args.particle_ys;
            markersize=2 * sim_consts.interaction_radius,
            markerspace=:data,
            color=args.interaction_colors,
        )
    end

    if args.show_skin_circle
        scatter!(
            args.ax,
            args.particle_xs,
            args.particle_ys;
            markersize=2 * sim_consts.skin_radius,
            markerspace=:data,
            color=args.skin_colors,
        )
    end

    if args.show_frame_diff
        linesegments!(args.ax, args.segment_xs, args.segment_ys; color=args.particle_colors)
    end

    return nothing
end

function animate_bundle!(
    args,
    sim_consts::ReCo.SimConsts,
    io,
    bundle::ReCo.Bundle,
    skip_first_bundle_snapshot::Bool,
)
    for bundle_snapshot in 1:length(bundle.t)
        if !skip_first_bundle_snapshot
            update_values!(args, sim_consts, bundle, bundle_snapshot)

            if args.show_frame_diff
                Observables.notify(args.segment_xs)
                Observables.notify(args.segment_ys)
            end

            Observables.notify(args.particle_xs)
            Observables.notify(args.particle_ys)
            Observables.notify(args.particle_colors)

            if args.show_center_of_mass
                Observables.notify(args.center_of_mass_point)
            end

            if args.show_interaction_circle
                Observables.notify(args.interaction_colors)
            end

            if args.show_skin_circle
                Observables.notify(args.skin_colors)
            end
        end

        args.ax.title = "t = $(round(bundle.t[bundle_snapshot], digits=3))"

        recordframe!(io)
    end

    return nothing
end

"""
    animate(sim_dir::String; <keyword arguments>)

Animate a simulation.

The output is `sim_dir/animation.mkv`.

Return `nothing`.

# Arguments
- `sim_dir::AbstractString`: Simulation directory.
- `framerate::Integer=$DEFAULT_FRAMERATE`: Framerate
- `show_center_of_mass::Bool=$DEFAULT_SHOW_CENTER_OF_MASS`: Show the center of mass as transparent white circle.
- `show_interaction_circle::Bool=$DEFAULT_SHOW_INTERACTION_CIRCLE`: Show the interaction radius with a circle around every particle.
- `show_skin_circle::Bool=$DEFAULT_SHOW_SKIN_CIRCLE`: Show the skin radius with a circle around every particle.
- `show_frame_diff::Bool=$DEFAULT_SHOW_FRAME_DIFF`: Show the translation of particles between two frames as lines connecting the old and new position. This is helpful to recognize unwanted jumps.
- `show_progress::Bool=$DEFAULT_SHOW_PROGRESS`: Show animation progress bar.
"""
function animate(
    sim_dir::AbstractString;
    framerate::Integer=DEFAULT_FRAMERATE,
    show_center_of_mass::Bool=DEFAULT_SHOW_CENTER_OF_MASS,
    show_interaction_circle::Bool=DEFAULT_SHOW_INTERACTION_CIRCLE,
    show_skin_circle::Bool=DEFAULT_SHOW_SKIN_CIRCLE,
    show_frame_diff::Bool=DEFAULT_SHOW_FRAME_DIFF,
    show_progress::Bool=DEFAULT_SHOW_PROGRESS,
)
    println("Initializing GLMakie...")

    GLMakie.activate!()
    set_theme!(theme_black())

    fig = Figure(; resolution=(1080, 1080))

    sim_consts = ReCo.load_sim_consts(sim_dir)

    ax, color_scheme = gen_axis_and_colorbar(fig, sim_consts)

    n_particles = sim_consts.n_particles

    animation_path = "$sim_dir/animation.mkv"

    particle_xs = Observable(Vector{Float64}(undef, n_particles))
    particle_ys = Observable(Vector{Float64}(undef, n_particles))

    particle_colors = Observable(
        Vector{ColorSchemes.ColorTypes.RGB{Float64}}(undef, n_particles)
    )

    center_of_mass_point =
        segment_xs =
            segment_ys =
                interaction_circle_xs =
                    interaction_circle_ys =
                        skin_circle_xs =
                            skin_circle_ys = interaction_colors = skin_colors = nothing

    if show_center_of_mass
        center_of_mass_point = Observable(Point2(0.0, 0.0))
    end

    if show_interaction_circle
        interaction_circle_xs = Observable(Vector{Float64}(undef, n_particles))
        interaction_circle_ys = Observable(Vector{Float64}(undef, n_particles))

        interaction_colors = Observable(
            Vector{ColorSchemes.ColorTypes.RGBA{Float64}}(undef, n_particles)
        )
    end

    if show_skin_circle
        skin_circle_xs = Observable(Vector{Float64}(undef, n_particles))
        skin_circle_ys = Observable(Vector{Float64}(undef, n_particles))

        skin_colors = Observable(
            Vector{ColorSchemes.ColorTypes.RGBA{Float64}}(undef, n_particles)
        )
    end

    if show_frame_diff
        segment_xs = Observable(zeros(Float64, 2 * n_particles))
        segment_ys = Observable(zeros(Float64, 2 * n_particles))
    end

    args = (;
        # Input
        show_center_of_mass,
        show_interaction_circle,
        show_skin_circle,
        show_frame_diff,
        # Internal
        ax,
        particle_xs,
        particle_ys,
        particle_colors,
        center_of_mass_point,
        segment_xs,
        segment_ys,
        interaction_circle_xs,
        interaction_circle_ys,
        skin_circle_xs,
        skin_circle_ys,
        interaction_colors,
        skin_colors,
        color_scheme,
    )

    bundle_paths = ReCo.sorted_bundle_paths(sim_dir)

    bundle::ReCo.Bundle = JLD2.load_object(bundle_paths[1])
    init_animation(args, sim_consts, bundle)

    progress = ProgressMeter.Progress(
        length(bundle_paths); dt=2, enabled=show_progress, desc="Animation: "
    )

    record(fig, animation_path; framerate=framerate) do io
        println("Started recording!")

        for (n_bundle, bundle_path) in enumerate(bundle_paths)
            if n_bundle != 1
                bundle = JLD2.load_object(bundle_path)
            end

            skip_first_bundle_snapshot = n_bundle == 1

            animate_bundle!(args, sim_consts, io, bundle, skip_first_bundle_snapshot)

            ProgressMeter.next!(progress)
        end
    end

    println("Animation done.")

    return nothing
end

end # module