diff --git a/visualization/Animation.jl b/visualization/Animation.jl new file mode 100644 index 0000000..c9ada05 --- /dev/null +++ b/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; ) + +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 diff --git a/visualization/Project.toml b/visualization/Project.toml new file mode 100644 index 0000000..612080c --- /dev/null +++ b/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" diff --git a/visualization/RewardsPlot.jl b/visualization/RewardsPlot.jl new file mode 100644 index 0000000..272f8c0 --- /dev/null +++ b/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 diff --git a/visualization/SnapshotPlot.jl b/visualization/SnapshotPlot.jl new file mode 100644 index 0000000..279c9af --- /dev/null +++ b/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; ) + +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 diff --git a/visualization/common.jl b/visualization/common.jl new file mode 100644 index 0000000..3a27d84 --- /dev/null +++ b/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 diff --git a/visualization/common_CairoMakie.jl b/visualization/common_CairoMakie.jl new file mode 100644 index 0000000..d183db6 --- /dev/null +++ b/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