Animate center of mass

src/Animation.jl

@@ -9,9 +9,9 @@ using JSON3: JSON3
 using JLD2: JLD2
 using ProgressMeter: @showprogress
 
-using ..ReCo: Bundle
+using ..ReCo: ReCo, Bundle
 
-function animate_bundle!(args)
+function animate_bundle!(args, sim_consts)
     bundle_t = args.bundle.t
     bundle_c = args.bundle.c
     bundle_φ = args.bundle.φ
@@ -19,27 +19,41 @@ function animate_bundle!(args)
     π2 = 2 * π
 
     for frame in 1:length(bundle_t)
-        @simd for i in 1:(args.n_particles)
+        @simd for i in 1:(sim_consts.n_particles)
             c = bundle_c[i, frame]
-            args.circles[][i] = Circle(Point2(c[1], c[2]), args.particle_radius)
+            args.circles[][i] = Circle(Point2(c[1], c[2]), sim_consts.particle_radius)
 
             color = get(args.color_scheme, rem2pi(bundle_φ[i, frame], RoundDown) / π2)
             args.colors[][i] = RGBAf(color)
 
             if args.debug
                 args.interaction_circles[][i] = Circle(
-                    Point2(c[1], c[2]), args.interaction_r
+                    Point2(c[1], c[2]), sim_consts.interaction_r
                 )
-                args.skin_circles[][i] = Circle(Point2(c[1], c[2]), args.skin_r)
+                args.skin_circles[][i] = Circle(Point2(c[1], c[2]), sim_consts.skin_r)
 
                 args.interaction_colors[][i] = RGBAf(color, 0.08)
                 args.skin_colors[][i] = 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
+            )
+            args.center_of_mass_circle[] = Circle(
+                Point2(center_of_mass[1], center_of_mass[2]),
+                args.center_of_mass_circle_radius,
+            )
+        end
+
         if args.n_bundle == 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))
+            end
+
             if args.debug
                 poly!(args.ax, args.interaction_circles; color=args.interaction_colors)
                 poly!(args.ax, args.skin_circles; color=args.skin_colors)
@@ -48,7 +62,7 @@ function animate_bundle!(args)
             println("Recording started!")
         else
             if args.debug && frame > 1
-                @simd for i in 1:(args.n_particles)
+                @simd for i in 1:(sim_consts.n_particles)
                     first_ind = 2 * i - 1
                     second_ind = 2 * i
                     frame_min_1 = frame - 1
@@ -70,6 +84,10 @@ function animate_bundle!(args)
             notify(args.circles)
             notify(args.colors)
 
+            if args.show_center_of_mass
+                notify(args.center_of_mass_circle)
+            end
+
             if args.debug && frame > 1
                 notify(args.interaction_circles)
                 notify(args.interaction_colors)
@@ -90,7 +108,9 @@ function animate_bundle!(args)
     return nothing
 end
 
-function animate_with_sim_consts(dir::String, sim_consts, framerate::Int64, debug::Bool)
+function animate_with_sim_consts(
+    dir::String, sim_consts, framerate::Int64, show_center_of_mass::Bool, debug::Bool
+)
     set_theme!(theme_black())
 
     fig = Figure(; resolution=(1080, 1080))
@@ -119,10 +139,19 @@ function animate_with_sim_consts(dir::String, sim_consts, framerate::Int64, debu
         circles = Observable(Vector{Circle}(undef, n_particles))
         colors = Observable(Vector{RGBAf}(undef, n_particles))
 
-        segments_x =
+        center_of_mass_circle =
-            segments_y =
+            segments_x =
-                interaction_circles =
+                segments_y =
-                    skin_circles = interaction_colors = skin_colors = nothing
+                    interaction_circles =
+                        skin_circles = interaction_colors = skin_colors = nothing
+
+        center_of_mass_circle_radius = 3 * sim_consts.particle_radius
+
+        if show_center_of_mass
+            center_of_mass_circle = Observable(
+                Circle(Point2(0.0, 0.0), center_of_mass_circle_radius)
+            )
+        end
 
         if debug
             segments_x = Observable(zeros(2 * n_particles))
@@ -135,9 +164,6 @@ function animate_with_sim_consts(dir::String, sim_consts, framerate::Int64, debu
             skin_colors = Observable(Vector{RGBAf}(undef, n_particles))
         end
 
-        particle_radius = sim_consts.particle_radius
-        interaction_r = sim_consts.interaction_r
-
         bundle_paths = readdir("$dir/bundles"; join=true, sort=false)
 
         sort_perm = sortperm([
@@ -148,46 +174,48 @@ function animate_with_sim_consts(dir::String, sim_consts, framerate::Int64, debu
         bundle_paths = bundle_paths[sort_perm]
         sort_perm = nothing
 
-        skin_r = sim_consts.skin_r
-
         @showprogress 1 for (n_bundle, bundle_path) in enumerate(bundle_paths)
             bundle::Bundle = JLD2.load_object(bundle_path)
 
             args = (;
+                # Input
+                show_center_of_mass,
+                debug,
+                # Intern
                 io,
                 ax,
                 bundle,
-                sim_consts,
-                debug,
                 circles,
                 colors,
+                center_of_mass_circle_radius,
+                center_of_mass_circle,
                 segments_x,
                 segments_y,
                 interaction_circles,
                 skin_circles,
                 interaction_colors,
                 skin_colors,
-                particle_radius,
-                skin_r,
-                n_particles,
-                interaction_r,
                 color_scheme,
                 n_bundle,
             )
 
-            animate_bundle!(args)
+            animate_bundle!(args, sim_consts)
         end
     end
 
     return nothing
 end
 
-function animate(dir::String; framerate::Int64=1, debug::Bool=false)
+function animate(
+    dir::String; framerate::Int64=1, show_center_of_mass::Bool=false, debug::Bool=false
+)
     println("Generating animation...")
 
+    set_window_config!(; framerate=framerate)
+
     sim_consts = JSON3.read(read("$dir/sim_consts.json", String))
 
-    animate_with_sim_consts(dir, sim_consts, framerate, debug)
+    animate_with_sim_consts(dir, sim_consts, framerate, show_center_of_mass, debug)
 
     println("Animation done.")