Snapshot graphic with CairoMakie

analysis/radial_distribution_function/radial_distribution_function.jl

@@ -48,7 +48,7 @@ end
 function radial_distribution(;
     sim_dirs::Vector{String}, n_radii::Int64, n_last_snapshots::Int64, n_particles::Int64
 )
-    sim_consts::ReCo.SimConsts = JLD2.load_object("$(sim_dirs[1])/sim_consts.jld2")
+    sim_consts = ReCo.load_sim_consts(sim_dirs[1])
 
     particle_radius = sim_consts.particle_radius
 

graphics/common_CairoMakie.jl

@@ -5,7 +5,7 @@ function init_cairomakie!()
     return nothing
 end
 
-function gen_figure(; padding=4)
+function gen_figure(; padding=5)
     text_width_in_pt = 405
 
     return Figure(;
@@ -22,8 +22,10 @@ function set_gaps!(fig::Figure)
     return nothing
 end
 
-function save_fig(filename::String, fig::Figure, parent_dir="exports/graphics")
+function save_fig(filename::String, fig::Figure; parent_dir="exports/graphics")
     mkpath(parent_dir)
 
-    return save("$parent_dir/$filename", fig; pt_per_unit=1)
+    save("$parent_dir/$filename", fig; pt_per_unit=1)
+
+    return nothing
 end

src/Animation.jl

@@ -1,47 +1,60 @@
 module Animation
 
-export animate
+export animate_with_gl, animate_with_cairo
 
+using CairoMakie
 using GLMakie
-using ColorSchemes: cyclic_mrybm_35_75_c68_n256_s25
+using ColorSchemes: ColorSchemes
 using LaTeXStrings: @L_str
 using JLD2: JLD2
 using ProgressMeter: ProgressMeter
 
 using ..ReCo: ReCo
 
+include("../graphics/common_CairoMakie.jl")
+
+const DEFAULT_CENTER_OF_MASS_CIRCLE_COLOR = RGBAf(1, 1, 1, 0.5)
+
+function angle_color(φ::Float64, color_scheme::ColorSchemes.ColorScheme)
+    return get(color_scheme, rem2pi(φ, RoundDown) / (2 * π))
+end
+
 function animate_bundle!(args, sim_consts::ReCo.SimConsts)
     bundle_t = args.bundle.t
     bundle_c = args.bundle.c
     bundle_φ = args.bundle.φ
 
-    π2 = 2 * π
+    for snapshot in 1:length(bundle_t)
+        @simd for particle_ind in 1:(sim_consts.n_particles)
+            c = bundle_c[particle_ind, snapshot]
+            φ = bundle_φ[particle_ind, snapshot]
 
-    for frame in 1:length(bundle_t)
+            args.circles[][particle_ind] = Circle(
-        @simd for i in 1:(sim_consts.n_particles)
+                Point2(c[1], c[2]), sim_consts.particle_radius
-            c = bundle_c[i, frame]
+            )
-            args.circles[][i] = Circle(Point2(c[1], c[2]), sim_consts.particle_radius)
 
-            color = get(args.color_scheme, rem2pi(bundle_φ[i, frame], RoundDown) / π2)
+            color = angle_color(φ, args.color_scheme)
-            args.colors[][i] = RGBAf(color)
+            args.colors[][particle_ind] = RGBAf(color)
 
             if args.show_interaction_circle
-                args.interaction_circles[][i] = Circle(
+                args.interaction_circles[][particle_ind] = Circle(
                     Point2(c[1], c[2]), sim_consts.interaction_r
                 )
-                args.interaction_colors[][i] = RGBAf(color, 0.08)
+                args.interaction_colors[][particle_ind] = RGBAf(color, 0.08)
             end
 
             if args.show_skin_circle
-                args.skin_circles[][i] = Circle(Point2(c[1], c[2]), sim_consts.skin_r)
+                args.skin_circles[][particle_ind] = Circle(
+                    Point2(c[1], c[2]), sim_consts.skin_r
+                )
 
-                args.skin_colors[][i] = RGBAf(color, 0.04)
+                args.skin_colors[][particle_ind] = RGBAf(color, 0.04)
             end
         end
 
         if args.show_center_of_mass
             center_of_mass = ReCo.center_of_mass(
-                view(bundle_c, :, frame), sim_consts.half_box_len
+                view(bundle_c, :, snapshot), sim_consts.half_box_len
             )
             args.center_of_mass_circle[] = Circle(
                 Point2(center_of_mass[1], center_of_mass[2]),
@@ -49,11 +62,15 @@ function animate_bundle!(args, sim_consts::ReCo.SimConsts)
             )
         end
 
-        if args.n_bundle == 1 # First and only frame of first bundle
+        if args.n_bundle == 1 && snapshot == 1
             poly!(args.ax, args.circles; color=args.colors)
 
             if args.show_center_of_mass
-                poly!(args.ax, args.center_of_mass_circle; color=RGBAf(1, 1, 1, 0.5))
+                poly!(
+                    args.ax,
+                    args.center_of_mass_circle;
+                    color=DEFAULT_CENTER_OF_MASS_CIRCLE_COLOR,
+                )
             end
 
             if args.show_interaction_circle
@@ -66,20 +83,20 @@ function animate_bundle!(args, sim_consts::ReCo.SimConsts)
 
             println("Recording started!")
         else
-            if args.show_frame_diff && frame > 1
+            if args.show_frame_diff && snapshot > 1
                 @simd for i in 1:(sim_consts.n_particles)
                     first_ind = 2 * i - 1
                     second_ind = 2 * i
-                    frame_min_1 = frame - 1
+                    frame_min_1 = snapshot - 1
 
                     args.segments_x[][first_ind] = bundle_c[i, frame_min_1][1]
-                    args.segments_x[][second_ind] = bundle_c[i, frame][1]
+                    args.segments_x[][second_ind] = bundle_c[i, snapshot][1]
 
                     args.segments_y[][first_ind] = bundle_c[i, frame_min_1][2]
-                    args.segments_y[][second_ind] = bundle_c[i, frame][2]
+                    args.segments_y[][second_ind] = bundle_c[i, snapshot][2]
                 end
 
-                if frame == 2
+                if snapshot == 2
                     linesegments!(
                         args.ax, args.segments_x, args.segments_y; color=args.colors
                     )
@@ -107,7 +124,7 @@ function animate_bundle!(args, sim_consts::ReCo.SimConsts)
             end
         end
 
-        args.ax.title = "t = $(round(bundle_t[frame], digits=3))"
+        args.ax.title = "t = $(round(bundle_t[snapshot], digits=3))"
 
         recordframe!(args.io)
     end
@@ -115,20 +132,7 @@ function animate_bundle!(args, sim_consts::ReCo.SimConsts)
     return nothing
 end
 
-function animate_with_sim_consts(
+function gen_axis_and_colorbar(fig, sim_consts::ReCo.SimConsts)
-    dir::String,
-    sim_consts::ReCo.SimConsts,
-    framerate::Int64,
-    show_center_of_mass::Bool,
-    show_interaction_circle::Bool,
-    show_skin_circle::Bool,
-    show_frame_diff::Bool,
-    show_progress::Bool,
-)
-    GLMakie.activate!()
-    set_theme!(theme_black())
-
-    fig = Figure(; resolution=(1080, 1080))
     ax = Axis(
         fig[1, 1];
         limits=(
@@ -142,13 +146,36 @@ function animate_with_sim_consts(
         ylabel=L"y",
     )
 
-    color_scheme = cyclic_mrybm_35_75_c68_n256_s25
+    color_scheme = ColorSchemes.cyclic_mrybm_35_75_c68_n256_s25
 
-    Colorbar(fig[1, 2]; limits=(0, 2), colormap=color_scheme, label=L"\frac{\varphi}{\pi}")
+    Colorbar(fig[1, 2]; limits=(0, 2), colormap=color_scheme, label=L"\varphi / \pi")
+
+    return (ax, color_scheme)
+end
+
+function animate_with_gl(
+    sim_dir::String;
+    framerate::Int64=1,
+    show_center_of_mass::Bool=false,
+    show_interaction_circle::Bool=false,
+    show_skin_circle::Bool=false,
+    show_frame_diff::Bool=false,
+    show_progress::Bool=true,
+)
+    println("Generating animation...")
+
+    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 = "$dir/animation.mkv"
+    animation_path = "$sim_dir/animation.mkv"
 
     record(fig, animation_path; framerate=framerate) do io
         circles = Observable(Vector{Circle}(undef, n_particles))
@@ -183,7 +210,7 @@ function animate_with_sim_consts(
             segments_y = Observable(zeros(2 * n_particles))
         end
 
-        bundle_paths = ReCo.sorted_bundle_paths(dir)
+        bundle_paths = ReCo.sorted_bundle_paths(sim_dir)
 
         progress = ProgressMeter.Progress(
             length(bundle_paths); dt=2, enabled=show_progress, desc="Animation: "
@@ -222,36 +249,71 @@ function animate_with_sim_consts(
         end
     end
 
-    return nothing
-end
-
-function animate(
-    dir::String;
-    framerate::Int64=1,
-    show_center_of_mass::Bool=false,
-    show_interaction_circle::Bool=false,
-    show_skin_circle::Bool=false,
-    show_frame_diff::Bool=false,
-    show_progress::Bool=true,
-)
-    println("Generating animation...")
-
-    sim_consts::ReCo.SimConsts = JLD2.load_object("$dir/sim_consts.jld2")
-
-    animate_with_sim_consts(
-        dir,
-        sim_consts,
-        framerate,
-        show_center_of_mass,
-        show_interaction_circle,
-        show_skin_circle,
-        show_frame_diff,
-        show_progress,
-    )
-
     println("Animation done.")
 
     return nothing
 end
 
+function get_wanted_snapshot_number(n_snapshots::Int64)
+    print("There are $n_snapshots snapshots. Enter the wanted snapshot number: ")
+    answer = readline()
+
+    snapshot = parse(Int64, answer)
+
+    @assert 0 < snapshot <= n_snapshots
+
+    return snapshot
+end
+
+function animate_with_cairo(sim_dir::String; show_center_of_mass::Bool=false)
+    bundles_info = ReCo.BundlesInfo(sim_dir)
+    snapshot = get_wanted_snapshot_number(bundles_info.total_n_snapshots)
+
+    bundle_path = ""
+    snapshot_counter = 1
+
+    for bundle_ind in bundles_info.n_bundles
+        snapshot_counter += bundles_info.bundle_n_snapshots[bundle_ind]
+        if snapshot_counter >= snapshot
+            bundle_path = bundles_info.sorted_bundle_paths[bundle_ind]
+            break
+        end
+    end
+
+    bundle::ReCo.Bundle = JLD2.load_object(bundle_path)
+    sim_consts = ReCo.load_sim_consts(sim_dir)
+
+    println("Initializing")
+
+    init_cairomakie!()
+    fig = gen_figure()
+
+    ax, color_scheme = gen_axis_and_colorbar(fig, sim_consts)
+
+    @simd for particle_ind in 1:(sim_consts.n_particles)
+        c = bundle.c[particle_ind, snapshot]
+        φ = bundle.φ[particle_ind, snapshot]
+
+        color = angle_color(φ, color_scheme)
+
+        poly!(ax, Circle(Point2(c[1], c[2]), sim_consts.particle_radius); color=color)
+    end
+
+    if show_center_of_mass
+        center_of_mass = ReCo.center_of_mass(
+            view(bundle.c, :, snapshot), sim_consts.half_box_len
+        )
+        circle = Circle(
+            Point2(center_of_mass[1], center_of_mass[2]), 3 * sim_consts.particle_radius
+        )
+        poly!(ax, circle; color=DEFAULT_CENTER_OF_MASS_CIRCLE_COLOR)
+    end
+
+    set_gaps!(fig)
+
+    save_fig("$snapshot.pdf", fig; parent_dir="$sim_dir/graphics")
+
+    return nothing
+end
+
 end # module

src/ReCo.jl

@@ -1,6 +1,6 @@
 module ReCo
 
-export init_sim, run_sim, run_rl, animate, LocalCOMEnv
+export init_sim, run_sim, run_rl, animate_with_gl, animate_with_cairo, LocalCOMEnv
 
 using StaticArrays: SVector
 using JLD2: JLD2

src/data.jl

@@ -114,7 +114,7 @@ function append_bundle_properties!(
     bundle_paths = ReCo.sorted_bundle_paths(sim_dir)
 
     for bundle_ind in first_bundle:length(bundle_paths)
-        bundle::ReCo.Bundle = JLD2.load_object(bundle_paths[bundle_ind])
+        bundle::Bundle = JLD2.load_object(bundle_paths[bundle_ind])
 
         for (v_ind, v) in enumerat(vs)
             property = properties[v_ind]
@@ -132,3 +132,36 @@ function append_bundle_properties!(
 
     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(dir::String)
+    sim_consts::ReCo.SimConsts = JLD2.load_object("$dir/sim_consts.jld2")
+
+    return sim_consts
+end

src/run.jl

@@ -23,7 +23,7 @@ function run_sim(
 
     Random.seed!(seed)
 
-    sim_consts::SimConsts = JLD2.load_object("$dir/sim_consts.jld2")
+    sim_consts = load_sim_consts(dir)
 
     integration_steps = floor(Int64, duration / sim_consts.δt) + 1