Fixed elliptical_distance

src/RL/EnvHelper.jl

@@ -41,7 +41,7 @@ struct EnvHelperSharedProps{H<:AbstractHook}
 end
 
 function gen_env_helper(::Env, env_helper_params::EnvHelperSharedProps; args)
-    return method_not_implemented()
+    return ReCo.method_not_implemented()
 end
 
 function get_env_agent_hook(env_helper::EnvHelper)

src/RL/Hooks.jl

@@ -1,19 +1,21 @@
+using ..ReCo: Particle
+
 function pre_integration_hook(::EnvHelper)
-    return method_not_implemented()
+    return ReCo.method_not_implemented()
 end
 
 function state_update_helper_hook(
     ::EnvHelper, id1::Int64, id2::Int64, r⃗₁₂::SVector{2,Float64}
 )
-    return method_not_implemented()
+    return ReCo.method_not_implemented()
 end
 
 function state_update_hook(::EnvHelper, particles::Vector{Particle})
-    return method_not_implemented()
+    return ReCo.method_not_implemented()
 end
 
 function update_reward!(::Env, ::EnvHelper, particle::Particle)
-    return method_not_implemented()
+    return ReCo.method_not_implemented()
 end
 
 function update_table_and_actions_hook(

src/RL/LocalCOMEnv.jl

@@ -1,5 +1,7 @@
 export LocalCOMEnv
 
+using ..ReCo: Particle
+
 struct LocalCOMEnv <: Env
     shared::EnvSharedProps
 
@@ -50,6 +52,7 @@ mutable struct LocalCOMEnvHelper <: EnvHelper
 
     add_shape_reward_term::Bool
 
+    center_of_mass::SVector{2,Float64}
     gyration_tensor_eigvec_to_smaller_eigval::SVector{2,Float64}
     gyration_tensor_eigvec_to_bigger_eigval::SVector{2,Float64}
 
@@ -71,6 +74,7 @@ mutable struct LocalCOMEnvHelper <: EnvHelper
             false,
             SVector(0.0, 0.0),
             SVector(0.0, 0.0),
+            SVector(0.0, 0.0),
             half_box_len,
             max_elliptical_distance,
         )
@@ -125,7 +129,7 @@ function state_update_hook(env_helper::LocalCOMEnvHelper, particles::Vector{Part
             vec_to_local_center_of_mass =
                 env_helper.vec_to_neighbour_sums[id] / n_neighbours
 
-            distance = norm2d(vec_to_local_center_of_mass)
+            distance = ReCo.norm2d(vec_to_local_center_of_mass)
 
             env_helper.distances_to_local_center_of_mass[id] = distance
 
@@ -135,7 +139,7 @@ function state_update_hook(env_helper::LocalCOMEnvHelper, particles::Vector{Part
 
             si, co = sincos(particles[id].φ)
 
-            direction_angle = angle2(SVector(co, si), vec_to_local_center_of_mass)
+            direction_angle = ReCo.angle2(SVector(co, si), vec_to_local_center_of_mass)
 
             direction_angle_state = find_state_interval(
                 direction_angle, env.direction_angle_state_space
@@ -158,7 +162,11 @@ function state_update_hook(env_helper::LocalCOMEnvHelper, particles::Vector{Part
         #println(mean_distance_to_local_center_of_mass / env_helper.max_distance_to_local_center_of_mass) # TODO: Remove
     end
 
-    v1, v2 = Shape.gyration_tensor_eigvecs(particles, env_helper.half_box_len)
+    env_helper.center_of_mass = ReCo.center_of_mass(particles, env_helper.half_box_len)
+
+    v1, v2 = ReCo.gyration_tensor_eigvecs(
+        particles, env_helper.half_box_len, env_helper.center_of_mass
+    )
 
     env_helper.gyration_tensor_eigvec_to_smaller_eigval = v1
     env_helper.gyration_tensor_eigvec_to_bigger_eigval = v2
@@ -191,10 +199,12 @@ function update_reward!(env::LocalCOMEnv, env_helper::LocalCOMEnvHelper, particl
 
         if env_helper.add_shape_reward_term
             elliptical_distance = ReCo.elliptical_distance(
-                particle,
+                particle.c,
+                env_helper.center_of_mass,
                 env_helper.gyration_tensor_eigvec_to_smaller_eigval,
                 env_helper.gyration_tensor_eigvec_to_bigger_eigval,
                 env_helper.shared.goal_gyration_tensor_eigvals_ratio,
+                env_helper.half_box_len,
             )
 
             reward += unnormalized_reward(

src/RL/RL.jl

@@ -12,14 +12,7 @@ using LoopVectorization: @turbo
 using Random: Random
 using ProgressMeter: @showprogress
 
-using ..ReCo:
+using ..ReCo: ReCo
-    ReCo,
-    Particle,
-    angle2,
-    norm2d,
-    Shape,
-    DEFAULT_SKIN_TO_INTERACTION_R_RATIO,
-    method_not_implemented
 
 const INITIAL_STATE_IND = 1
 const INITIAL_REWARD = 0.0
@@ -67,7 +60,7 @@ function run_rl(;
     n_particles::Int64=100,
     seed::Int64=42,
     ϵ_stable::Float64=0.0001,
-    skin_to_interaction_r_ratio::Float64=DEFAULT_SKIN_TO_INTERACTION_R_RATIO,
+    skin_to_interaction_r_ratio::Float64=ReCo.DEFAULT_SKIN_TO_INTERACTION_R_RATIO,
     packing_ratio=0.22,
 ) where {E<:Env}
     @assert 0.0 <= goal_gyration_tensor_eigvals_ratio <= 1.0

src/Shape.jl

@@ -6,7 +6,7 @@ export center_of_mass,
 using StaticArrays: SVector, SMatrix
 using LinearAlgebra: eigvals, eigvecs, Hermitian, dot
 
-using ..ReCo: Particle, restrict_coordinate, restrict_coordinates
+using ..ReCo: Particle
 
 function project_to_unit_circle(x::Float64, half_box_len::Float64)
     φ = (x + half_box_len) * π / half_box_len
@@ -18,7 +18,7 @@ end
 function project_back_from_unit_circle(θ::T, half_box_len::Float64) where {T<:Real}
     x = θ * half_box_len / π - half_box_len
 
-    return restrict_coordinate(x, half_box_len)
+    return ReCo.restrict_coordinate(x, half_box_len)
 end
 
 function center_of_mass_from_proj_sums(
@@ -71,15 +71,15 @@ function center_of_mass(particles::Vector{Particle}, half_box_len::Float64)
     return center_of_mass_from_proj_sums(x_proj_sum, y_proj_sum, half_box_len)
 end
 
-function gyration_tensor(particles::Vector{Particle}, half_box_len::Float64)
+function gyration_tensor(
-    COM = center_of_mass(particles, half_box_len)
+    particles::Vector{Particle}, half_box_len::Float64, COM::SVector{2,Float64}
-
+)
     S11 = 0.0
     S12 = 0.0
     S22 = 0.0
 
     for p in particles
-        shifted_c = restrict_coordinates(p.c - COM, half_box_len)
+        shifted_c = ReCo.restrict_coordinates(p.c - COM, half_box_len)
 
         S11 += shifted_c[1]^2
         S12 += shifted_c[1] * shifted_c[2]
@@ -89,14 +89,22 @@ function gyration_tensor(particles::Vector{Particle}, half_box_len::Float64)
     return Hermitian(SMatrix{2,2}(S11, S12, S12, S22))
 end
 
+function gyration_tensor(particles::Vector{Particle}, half_box_len::Float64)
+    COM = center_of_mass(particles, half_box_len)
+
+    return gyration_tensor(particles, half_box_len, COM)
+end
+
 function gyration_tensor_eigvals_ratio(particles::Vector{Particle}, half_box_len::Float64)
     g_tensor = gyration_tensor(particles, half_box_len)
     ev = eigvals(g_tensor) # Eigenvalues are sorted
     return ev[1] / ev[2]
 end
 
-function gyration_tensor_eigvecs(particles::Vector{Particle}, half_box_len::Float64)
+function gyration_tensor_eigvecs(
-    g_tensor = gyration_tensor(particles, half_box_len)
+    particles::Vector{Particle}, half_box_len::Float64, COM::SVector{2,Float64}
+)
+    g_tensor = gyration_tensor(particles, half_box_len, COM)
     eig_vecs = eigvecs(g_tensor)
 
     v1 = eig_vecs[:, 1]
@@ -106,15 +114,19 @@ function gyration_tensor_eigvecs(particles::Vector{Particle}, half_box_len::Floa
 end
 
 function elliptical_distance(
-    particle::Particle,
+    v::SVector{2,Float64},
+    COM::SVector{2,Float64},
     gyration_tensor_eigvec_to_smaller_eigval::SVector{2,Float64},
     gyration_tensor_eigvec_to_bigger_eigval::SVector{2,Float64},
     goal_gyration_tensor_eigvals_ratio::Float64,
+    half_box_len::Float64,
 )
-    cx′ = dot(particle.c, gyration_tensor_eigvec_to_bigger_eigval)
+    v′ = ReCo.minimum_image(v - COM, half_box_len)
-    cy′ = dot(particle.c, gyration_tensor_eigvec_to_smaller_eigval)
 
-    return cx′^2 + (cy′ / goal_gyration_tensor_eigvals_ratio)^2
+    x = dot(v′, gyration_tensor_eigvec_to_bigger_eigval)
+    y = dot(v′, gyration_tensor_eigvec_to_smaller_eigval)
+
+    return x^2 + (y / goal_gyration_tensor_eigvals_ratio)^2
 end
 
 end # module