Separate visualization

visualization/Animation.jl

@@ -0,0 +1,282 @@
+module Animation
+
+export animate
+
+using GLMakie
+using ColorSchemes: ColorSchemes
+using LaTeXStrings: @L_str
+using JLD2: JLD2
+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 animate_bundle!(args, sim_consts::ReCo.SimConsts)
+    bundle_t = args.bundle.t
+    bundle_c = args.bundle.c
+    bundle_φ = args.bundle.φ
+
+    for bundle_snapshot in 1:length(bundle_t)
+        @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
+        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
+
+        if args.n_bundle == 1 && bundle_snapshot == 1
+            scatter!(
+                args.ax,
+                args.particle_xs,
+                args.particle_ys;
+                markersize=2 * sim_consts.particle_radius,
+                markerspace=SceneSpace,
+                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=SceneSpace,
+                    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=SceneSpace,
+                    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=SceneSpace,
+                    color=args.skin_colors,
+                )
+            end
+
+            println("Recording started!")
+        else
+            if args.show_frame_diff && bundle_snapshot > 1
+                @simd for i in 1:(sim_consts.n_particles)
+                    first_ind = 2 * i - 1
+                    second_ind = 2 * i
+                    frame_min_1 = bundle_snapshot - 1
+
+                    args.segment_xs[][first_ind] = bundle_c[i, frame_min_1][1]
+                    args.segment_xs[][second_ind] = bundle_c[i, bundle_snapshot][1]
+
+                    args.segment_ys[][first_ind] = bundle_c[i, frame_min_1][2]
+                    args.segment_ys[][second_ind] = bundle_c[i, bundle_snapshot][2]
+                end
+
+                if bundle_snapshot == 2
+                    linesegments!(
+                        args.ax,
+                        args.segment_xs,
+                        args.segment_ys;
+                        color=args.particle_colors,
+                    )
+                end
+
+                notify(args.segment_xs)
+                notify(args.segment_ys)
+            end
+
+            notify(args.particle_xs)
+            notify(args.particle_ys)
+            notify(args.particle_colors)
+
+            if args.show_center_of_mass
+                notify(args.center_of_mass_point)
+            end
+
+            if args.show_interaction_circle
+                notify(args.interaction_colors)
+            end
+
+            if args.show_skin_circle
+                notify(args.skin_colors)
+            end
+        end
+
+        args.ax.title = "t = $(round(bundle_t[bundle_snapshot], digits=3))"
+
+        recordframe!(args.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::String`: Simulation directory.
+- `framerate::Int64=$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::String;
+    framerate::Int64=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"
+
+    record(fig, animation_path; framerate=framerate) do io
+        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
+
+        bundle_paths = ReCo.sorted_bundle_paths(sim_dir)
+
+        progress = ProgressMeter.Progress(
+            length(bundle_paths); dt=2, enabled=show_progress, desc="Animation: "
+        )
+
+        for (n_bundle, bundle_path) in enumerate(bundle_paths)
+            bundle::ReCo.Bundle = JLD2.load_object(bundle_path)
+
+            args = (;
+                # Input
+                show_center_of_mass,
+                show_interaction_circle,
+                show_skin_circle,
+                show_frame_diff,
+                # Internal
+                io,
+                ax,
+                bundle,
+                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,
+                n_bundle,
+            )
+
+            animate_bundle!(args, sim_consts)
+
+            ProgressMeter.next!(progress)
+        end
+    end
+
+    println("Animation done.")
+
+    return nothing
+end
+
+end # module

visualization/Project.toml

@@ -0,0 +1,8 @@
+[deps]
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+ColorSchemes = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
+GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
+ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
+ReCo = "b25f7548-fcc9-4c91-bc24-841b54f4dd54"

visualization/RewardsPlot.jl

@@ -0,0 +1,54 @@
+module RewardsPlot
+
+export plot_rewards
+
+using CairoMakie
+using JLD2: JLD2
+
+using ReCo: ReCo
+
+include("common_CairoMakie.jl")
+
+const DEFAULT_ENV_HELPER_FILENAME = "env_helper.jld2"
+
+function plot_rewards_from_env_helper(; env_helper::ReCo.RL.EnvHelper, rl_dir::String)
+    rewards = env_helper.shared.hook.rewards
+    n_episodes = length(rewards)
+
+    init_cairomakie!()
+
+    fig = gen_figure(; padding=10)
+
+    ax = Axis(
+        fig[1, 1]; xlabel="Episode", ylabel="Reward", limits=((0, n_episodes), nothing)
+    )
+
+    lines!(ax, 1:n_episodes, rewards)
+
+    set_gaps!(fig)
+
+    save_fig("rewards.pdf", fig; parent_dir=rl_dir)
+
+    return nothing
+end
+
+"""
+    plot_rewards(rl_dir::String, env_helper_filename::String="$DEFAULT_ENV_HELPER_FILENAME")
+
+Plot the rewards of the reinforcement learning process at the directory `rl_dir`.
+
+The output is `rl_dir/rewards.pdf`. `env_helper_filename` can be provided if the name of the `jld2`-file of the environment helper differs from the default `$DEFAULT_ENV_HELPER_FILENAME`.
+
+Return `nothing`.
+"""
+function plot_rewards(
+    rl_dir::String, env_helper_filename::String=DEFAULT_ENV_HELPER_FILENAME
+)
+    env_helper::ReCo.RL.EnvHelper = JLD2.load_object("$rl_dir/$env_helper_filename")
+
+    plot_rewards_from_env_helper(; env_helper, rl_dir)
+
+    return nothing
+end
+
+end # module

visualization/SnapshotPlot.jl

@@ -0,0 +1,115 @@
+module SnapshotPlot
+
+export plot_snapshot
+
+using CairoMakie
+using ColorSchemes: ColorSchemes
+using LaTeXStrings: @L_str
+
+using ..ReCo: ReCo
+
+include("common.jl")
+include("common_CairoMakie.jl")
+
+const DEFAULT_SHOW_CENTER_OF_MASS = false
+const DEFAULT_SHOW_KAPPA = true
+
+function get_wanted_snapshot_number(total_n_snapshots::Int64)
+    print("There are $total_n_snapshots snapshots. Enter the wanted snapshot number: ")
+    answer = readline()
+
+    snapshot = parse(Int64, answer)
+
+    return snapshot
+end
+
+"""
+    plot_snapshot(sim_dir::String; <keyword arguments>)
+
+Plot one snapshot of a simulation.
+
+The function will ask for the number of the snapshot to plot out of the total number of snapshots. The output is `sim_dir/graphics/N.pdf` with `N` as the number of the chosen snapshot.
+
+# Arguments
+- `sim_dir::String`: Simulation directory.
+- `show_center_of_mass::Bool=$DEFAULT_SHOW_CENTER_OF_MASS`: Show the center of mass as a transparent black circle.
+- `show_κ::Bool=$DEFAULT_SHOW_KAPPA`: Show κ as the ratio of the eigenvalues of the gyration tensor.
+"""
+function plot_snapshot(
+    sim_dir::String;
+    show_center_of_mass::Bool=DEFAULT_SHOW_CENTER_OF_MASS,
+    show_κ::Bool=DEFAULT_SHOW_KAPPA,
+)
+    bundles_info = ReCo.BundlesInfo(sim_dir)
+    wanted_snapshot_out_of_total = get_wanted_snapshot_number(
+        bundles_info.total_n_snapshots
+    )
+
+    bundle, bundle_snapshot = ReCo.get_bundle_to_snapshot(
+        bundles_info, wanted_snapshot_out_of_total
+    )
+
+    sim_consts = ReCo.load_sim_consts(sim_dir)
+
+    println("Initializing CairoMakie...")
+
+    init_cairomakie!()
+    fig = gen_figure()
+
+    cs_view = view(bundle.c, :, bundle_snapshot)
+    center_of_mass = ReCo.center_of_mass(cs_view, sim_consts.half_box_len)
+
+    if show_κ
+        eigvals_ratio = ReCo.gyration_tensor_eigvals_ratio(
+            cs_view, sim_consts.half_box_len, center_of_mass
+        )
+        κ = round(eigvals_ratio; digits=2)
+        title = L"\kappa = %$κ"
+
+        ax, color_scheme = gen_axis_and_colorbar(fig, sim_consts; axis_title=title)
+    else
+        ax, color_scheme = gen_axis_and_colorbar(fig, sim_consts)
+    end
+
+    particle_xs = Vector{Float64}(undef, sim_consts.n_particles)
+    particle_ys = Vector{Float64}(undef, sim_consts.n_particles)
+    particle_colors = Vector{ColorSchemes.ColorTypes.RGB}(undef, sim_consts.n_particles)
+
+    @simd for particle_ind in 1:(sim_consts.n_particles)
+        c = bundle.c[particle_ind, bundle_snapshot]
+        φ = bundle.φ[particle_ind, bundle_snapshot]
+
+        color = angle_color(φ, color_scheme)
+
+        particle_xs[particle_ind] = c[1]
+        particle_ys[particle_ind] = c[2]
+        particle_colors[particle_ind] = color
+    end
+
+    scatter!(
+        ax,
+        particle_xs,
+        particle_ys;
+        markersize=2 * sim_consts.particle_radius,
+        markerspace=SceneSpace,
+        color=particle_colors,
+    )
+
+    if show_center_of_mass
+        scatter!(
+            ax,
+            Point(center_of_mass);
+            markersize=6 * sim_consts.particle_radius,
+            markerspace=SceneSpace,
+            color=ColorSchemes.ColorTypes.RGBA(0.0, 0.0, 0.0, 0.6),
+        )
+    end
+
+    set_gaps!(fig)
+
+    save_fig("$wanted_snapshot_out_of_total.pdf", fig; parent_dir="$sim_dir/graphics")
+
+    return nothing
+end
+
+end # module

visualization/common.jl

@@ -0,0 +1,29 @@
+using ColorSchemes
+
+function angle_color(φ::Float64, color_scheme::ColorSchemes.ColorScheme)
+    return get(color_scheme, rem2pi(φ, RoundDown) / (2 * π))
+end
+
+function gen_axis_and_colorbar(
+    fig, sim_consts::ReCo.SimConsts; axis_title::AbstractString=""
+)
+    ax = Axis(
+        fig[1, 1];
+        limits=(
+            -sim_consts.half_box_len,
+            sim_consts.half_box_len,
+            -sim_consts.half_box_len,
+            sim_consts.half_box_len,
+        ),
+        aspect=AxisAspect(1),
+        xlabel=L"x",
+        ylabel=L"y",
+        title=axis_title,
+    )
+
+    color_scheme = ColorSchemes.cyclic_mrybm_35_75_c68_n256_s25
+
+    Colorbar(fig[1, 2]; limits=(0, 2), colormap=color_scheme, label=L"\varphi / \pi")
+
+    return (ax, color_scheme)
+end

visualization/common_CairoMakie.jl

@@ -0,0 +1,31 @@
+function init_cairomakie!()
+    CairoMakie.activate!()
+    set_theme!()
+
+    return nothing
+end
+
+function gen_figure(; padding=5)
+    text_width_in_pt = 405
+
+    return Figure(;
+        resolution=(text_width_in_pt, 0.55 * text_width_in_pt),
+        fontsize=10,
+        figure_padding=padding,
+    )
+end
+
+function set_gaps!(fig::Figure)
+    colgap!(fig.layout, 5)
+    rowgap!(fig.layout, 5)
+
+    return nothing
+end
+
+function save_fig(filename::String, fig::Figure; parent_dir="exports/graphics")
+    mkpath(parent_dir)
+
+    save("$parent_dir/$filename", fig; pt_per_unit=1)
+
+    return nothing
+end