Used modules

analysis/pair_correlation.jl

@@ -2,15 +2,15 @@ if splitdir(pwd())[2] == "analysis"
     cd("..")
 end
 
-if splitdir(pwd())[2] != "ReCo"
-    error("You have to be in the main directeory ReCo!")
+if splitdir(pwd())[2] != "ReCo.jl"
+    error("You have to be in the main directeory ReCo.jl!")
 else
     include("src/analysis/pair_correlation_function.jl")
 end
 
 ##
 
-using JLD2
+using JLD2: JLD2
 
 using CairoMakie
 CairoMakie.activate!()

src/Animation.jl

@@ -1,4 +1,4 @@
-module ReCoAnimation
+module Animation
 
 export animate
 
@@ -9,7 +9,7 @@ using JSON3: JSON3
 using JLD2: JLD2
 using ProgressMeter: @showprogress
 
-using ReCo: Bundle
+using ..ReCo: Bundle
 
 function animate_bundle!(args)
     bundle_t = args.bundle.t
@@ -90,7 +90,7 @@ function animate_bundle!(args)
     return nothing
 end
 
-function animate_after_loading(dir, animation_path, sim_consts, framerate, debug)
+function animate_with_sim_consts(dir::String, sim_consts, framerate::Int64, debug::Bool)
     set_theme!(theme_black())
 
     fig = Figure(; resolution=(1080, 1080))
@@ -113,6 +113,8 @@ function animate_after_loading(dir, animation_path, sim_consts, framerate, debug
 
     n_particles = sim_consts.n_particles
 
+    animation_path = "$dir/animation.mkv"
+
     record(fig, animation_path; framerate=framerate) do io
         circles = Observable(Vector{Circle}(undef, n_particles))
         colors = Observable(Vector{RGBAf}(undef, n_particles))
@@ -185,11 +187,9 @@ function animate(dir::String; framerate::Int64=1, debug::Bool=false)
 
     sim_consts = JSON3.read(read("$dir/sim_consts.json", String))
 
-    animation_path = "$dir/animation.mkv"
+    animate_with_sim_consts(dir, sim_consts, framerate, debug)
 
-    animate_after_loading(dir, animation_path, sim_consts, framerate, debug)
-
-    println("Animation done and saved to $animation_path.")
+    println("Animation done.")
 
     return nothing
 end

src/Benchmark.jl

@@ -1,3 +1,5 @@
+module Benchmark
+
 using Statistics: mean
 using Dates: now
 using BenchmarkTools: @benchmark
@@ -42,4 +44,6 @@ function run_benchmarks(
     end
 
     return dir
-end
+end
+
+end # module

src/Particle.jl

@@ -1,5 +1,3 @@
-using StaticArrays: SVector
-
 mutable struct Particle
     id::Int64
 

src/PreVectors.jl

@@ -1,3 +1,7 @@
+module PreVectors
+
+export PreVector, push!, reset!, iterate
+
 import Base: push!, iterate
 
 mutable struct PreVector{T}
@@ -27,3 +31,5 @@ function iterate(pv::PreVector{T}, state=1) where {T}
         return (pv.v[state], state + 1)
     end
 end
+
+end # module

src/RL.jl

@@ -10,7 +10,7 @@ using LoopVectorization: @turbo
 using Random: Random
 using ProgressMeter: @showprogress
 
-using ..ReCo
+using ..ReCo: ReCo, Particle
 
 const INITIAL_REWARD = 0.0
 
@@ -143,10 +143,10 @@ end
 
 mutable struct Env <: AbstractEnv
     params::EnvParams
-    particle::ReCo.Particle
+    particle::Particle
     state_ind::Int64
 
-    function Env(params::EnvParams, particle::ReCo.Particle)
+    function Env(params::EnvParams, particle::Particle)
         # initial_state = (nothing, nothing)
         initial_state_ind = params.n_states
 
@@ -154,7 +154,7 @@ mutable struct Env <: AbstractEnv
     end
 end
 
-function reset!(env::Env, particle::ReCo.Particle)
+function reset!(env::Env, particle::Particle)
     env.particle = particle
     env.state_ind = env.params.n_states
 
@@ -274,7 +274,7 @@ function state_hook(
 end
 
 function integration_hook!(
-    particle::ReCo.Particle, rl_params::Params, δt::Float64, si::Float64, co::Float64
+    particle::Particle, rl_params::Params, δt::Float64, si::Float64, co::Float64
 )
     # Apply action
     action = rl_params.actions[particle.id]
@@ -297,7 +297,7 @@ end
 function post_integration_hook(
     rl_params::Params,
     n_particles::Int64,
-    particles::Vector{ReCo.Particle},
+    particles::Vector{Particle},
     half_box_len::Float64,
 )
     # Update reward

src/ReCo.jl

@@ -1,14 +1,37 @@
 module ReCo
 
-export init_sim, run_sim, RL
+export init_sim, run_sim, run_rl, animate
+
+using StaticArrays: SVector
+using OrderedCollections: OrderedDict
+using JLD2: JLD2
+using JSON3: JSON3
+using Distributions: Uniform, Normal
+using ProgressMeter: @showprogress
+using CellListMap: Box, CellList, map_pairwise!, UpdateCellList!
+
+using Random: Random
+using Dates: Dates, now
+
+import Base: wait
+
+include("PreVectors.jl")
+using .PreVectors
 
-include("PreVector.jl")
 include("Particle.jl")
-include("reinforcement_learning.jl")
+
+include("Shape.jl")
+using .Shape
+
+include("RL.jl")
+using .RL
+
 include("data.jl")
 include("setup.jl")
 include("simulation.jl")
-include("shape.jl")
 include("run.jl")
 
+include("Animation.jl")
+using .Animation
+
 end # module

src/Shape.jl

@@ -1,6 +1,12 @@
+module Shape
+
+export center_of_mass, gyration_tensor_eigvals_ratio
+
 using StaticArrays: SVector, SMatrix
 using LinearAlgebra: eigvals, Hermitian
 
+using ..ReCo: Particle, restrict_coordinate, restrict_coordinates
+
 function project_to_unit_circle(x::Float64, half_box_len::Float64)
     φ = (x + half_box_len) * π / half_box_len
     si, co = sincos(φ)
@@ -65,4 +71,6 @@ end
 function gyration_tensor_eigvals_ratio(particles::Vector{Particle}, half_box_len::Float64)
     ev = eigvals(gyration_tensor(particles, half_box_len)) # Eigenvalues are sorted
     return ev[1] / ev[2]
-end
+end
+
+end # module

src/analysis/pair_correlation_function.jl

@@ -1,5 +1,4 @@
 using CairoMakie, LaTeXStrings
-using StaticArrays
 using LoopVectorization: @turbo
 
 using ReCo: minimum_image

src/data.jl

@@ -1,8 +1,3 @@
-using StaticArrays: SVector
-using JLD2: JLD2
-using JSON3: JSON3
-using OrderedCollections: OrderedDict
-
 struct Bundle
     t::Vector{Float64}
     c::Matrix{SVector{2,Float64}}

src/run.jl

@@ -1,8 +1,3 @@
-using Random: Random
-using JSON3: JSON3
-using JLD2: JLD2
-using StaticArrays: SVector
-
 empty_hook(args...) = nothing
 
 function run_sim(

src/setup.jl

@@ -1,6 +1,3 @@
-using Distributions: Uniform
-using Dates: now
-
 const DEFAULT_PACKING_RATIO = 0.5
 const DEFAULT_δt = 1e-5
 const DEFAULT_SKIN_TO_INTERACTION_R_RATIO = 1.5

src/simulation.jl

@@ -1,11 +1,3 @@
-using Dates: Dates, now
-using ProgressMeter: @showprogress
-using Distributions: Normal
-using CellListMap: Box, CellList, map_pairwise!, UpdateCellList!
-using StaticArrays: SVector
-
-import Base: wait
-
 rand_normal01() = rand(Normal(0, 1))
 
 function push_to_verlet_list!(verlet_lists, i, j)

test/runtests.jl

@@ -38,7 +38,7 @@ using StaticArrays: SVector
     end
 end
 
-@testset "shape.jl" begin
+@testset "Shape.jl" begin
     n_particles = 10
     v₀ = 0.0
     sim_consts = ReCo.gen_sim_consts(n_particles, v₀)
@@ -50,9 +50,9 @@ end
     half_box_len = sim_consts.half_box_len
 
     @testset "project_to_unit_circle" begin
-        @test ReCo.project_to_unit_circle(0.0, half_box_len) ≈ SVector(-1.0, 0.0)
-        @test ReCo.project_to_unit_circle(half_box_len, half_box_len) ≈
-            ReCo.project_to_unit_circle(-half_box_len, half_box_len) ≈
+        @test ReCo.Shape.project_to_unit_circle(0.0, half_box_len) ≈ SVector(-1.0, 0.0)
+        @test ReCo.Shape.project_to_unit_circle(half_box_len, half_box_len) ≈
+            ReCo.Shape.project_to_unit_circle(-half_box_len, half_box_len) ≈
             SVector(1.0, 0.0)
     end
 
@@ -61,10 +61,10 @@ end
     )
 
     @testset "center_of_mass" begin
-        @test ReCo.center_of_mass(particles, half_box_len) ≈ SVector(0.0, 0.0)
+        @test ReCo.Shape.center_of_mass(particles, half_box_len) ≈ SVector(0.0, 0.0)
     end
 
     @testset "gyration_tensor" begin
-        @test ReCo.gyration_tensor_eigvals_ratio(particles, half_box_len) == 1.0
+        @test ReCo.Shape.gyration_tensor_eigvals_ratio(particles, half_box_len) == 1.0
     end
 end