diff --git a/Project.toml b/Project.toml index 4daa50b..2a011a6 100644 --- a/Project.toml +++ b/Project.toml @@ -4,6 +4,7 @@ version = "0.1.0" [deps] Arianna = "07692032-97b4-4f8d-80d7-e18df88d31a9" +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" Comonicon = "863f3e99-da2a-4334-8734-de3dacbe5542" ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" @@ -19,6 +20,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] Aqua = "0.8" Arianna = "0.2" +BenchmarkTools = "1.8.0" Comonicon = "1.0" ComponentArrays = "0.15" ConcreteStructs = "0.2" @@ -26,19 +28,19 @@ DataStructures = "0.18" DelimitedFiles = "1.9" Distributions = "0.25" LinearAlgebra = "1.9" +Pkg = "1.9" Printf = "1.9" StaticArrays = "1.9" Statistics = "1.9" TOML = "1" Test = "1.9" julia = "1.9" -Pkg = "1.9" [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab" -Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] test = ["Test", "DelimitedFiles", "Aqua", "Pkg"] diff --git a/src/ParticlesMC.jl b/src/ParticlesMC.jl index f652cb0..6ff9a1b 100644 --- a/src/ParticlesMC.jl +++ b/src/ParticlesMC.jl @@ -20,6 +20,7 @@ include("models.jl") include("molecules.jl") include("atoms.jl") include("moves.jl") +include("msad.jl") """Return the position of particle `i` in `system`. @@ -241,6 +242,7 @@ ParticlesMC implemented in Comonicon. else error("Unsupported policy: $policy for action: $action") end + else error("Unsupported action: $action") end @@ -289,6 +291,33 @@ ParticlesMC implemented in Comonicon. algorithm=eval(Meta.parse(alg)), scheduler=sched, ) + elseif alg == "MSADTracker" + # read theta_T from parameters + parameters = get(output, "parameters", Dict()) + theta_T = Float64(get(parameters, "theta_T", π/4)) + + # build output scheduler from scheduler_params for consistency + output_sched = sched + + # build compute schedule separately + compute_scheduler_params = get(output, "compute_scheduler_params", Dict()) + if "log_base" in keys(compute_scheduler_params) + compute_sched = build_schedule(steps, burn, Float64(compute_scheduler_params["log_base"])) + elseif "linear_interval" in keys(compute_scheduler_params) + compute_sched = build_schedule(steps, burn, compute_scheduler_params["linear_interval"]) + else + compute_sched = build_schedule(steps, burn, 1) # every step by default + end + + # convert to Vector{Int} — build_schedule returns this already + algorithm = ( + algorithm=MSADTracker, + scheduler=compute_sched, # compute schedule + theta_T=theta_T, + output_schedule=output_sched, + path=output_path, + ) + else error("Unsupported output algorithm: $alg") end diff --git a/src/msad.jl b/src/msad.jl new file mode 100644 index 0000000..7d4e94b --- /dev/null +++ b/src/msad.jl @@ -0,0 +1,233 @@ +# msad.jl +# +# On-the-fly Mean Squared Angular Displacement (MSAD) tracker. +# Implements three methods: +# - Integral : accumulates incremental rotation vectors every step +# - Threshold : relays the reference frame when rotation exceeds θ_T +# - Euler : compares directly to initial frame, computed on log schedule + +using LinearAlgebra +using StaticArrays + +##### Basic operations for rotation inspired fromo our python post processing ##### +########## ########## + +function body_frame(r1::SVector{3,T}, r2::SVector{3,T}, r3::SVector{3,T}, L::T) where {T} + e1 = r2 - r1 + v = r3 - r1 + # minimum image convention + e1 = e1 - round.(e1 / L) .* L + v = v - round.( v / L) .* L + # Gram-Schmidt basis right handed + e1 = e1 / norm(e1) + e2 = cross(e1, v) + e2 = e2 / norm(e2) + e3 = cross(e1, e2) + # hcat stacks column vectors → 3×3 matrix + return hcat(e1, e2, e3) # SMatrix{3,3,T} +end + +function get_all_body_frames(system::Molecules) + L = system.box[1] # cubic box + T = typeof(L) + R = Vector{SMatrix{3,3,T,9}}(undef, system.Nmol) # pre-allocate + pos = system.position + @inbounds for m in 1:system.Nmol # iterate over the number of molecules + s = system.start_mol[m] # first bead index for the considered molecule + R[m] = body_frame(pos[s], pos[s+1], pos[s+2], L) + end + return R +end + +function rotation_vector(R::SMatrix{3,3,T}) where {T} + + cos_θ = clamp((tr(R) - 1) / 2, -one(T), one(T)) + + vals, vecs = eigen(R) # eigenvalues and eigenvectors of R + idx = argmin(abs.(real.(vals) .- 1)) # find the idx corresponding to eigenvalue = 1 + n = real.(vecs[:, idx]) # extract the rotation axis vector + n = SVector{3,T}(n[1], n[2], n[3]) + + n_skew = SMatrix{3,3,T}( + 0, n[3], -n[2], + -n[3], 0, n[1], + n[2], -n[1], 0 + ) # skew matrix for n vector + + # Rodrigues formula + sin_θ = clamp(-tr(n_skew * R) / 2, -one(T), one(T)) + + # theta + θ = atan(sin_θ, cos_θ) + + return θ * n +end + +##### Definition of the MSAD State ie: the accumulating variables ##### +# mutable because we need to update it +########## ########## + +mutable struct MSADState{T} + R0::Vector{SMatrix{3,3,T,9}} # (Nmol,3,3) initial body frames at t=0 + R_prev::Vector{SMatrix{3,3,T,9}} # (Nmol,3,3) body frames at previous time + phi_integral::Vector{SVector{3,T}} # (Nmol,3) Accumulated vector for integral method + R_ref_thresh::Vector{SMatrix{3,3,T,9}} # (Nmol,3,3) reference body frames for each molecule in the threshold method + phi_acc::Vector{SVector{3,T}} # (Nmol,3) Accumulated vector for threshold method + initialized::Bool +end + +# Start empty, filled during initialise when we first see the system +MSADState{T}() where {T} = MSADState{T}([], [], [], [], [], false) + +##### The MSAD computation algorithm ##### +########## ########## + +mutable struct MSADTracker{T} <: AriannaAlgorithm + states::Vector{MSADState{T}} + theta_T::T + compute_schedule::Vector{Int} + output_schedule::Vector{Int} + paths_integral::Vector{String} + paths_thresh::Vector{String} + files_integral::Vector{IOStream} + files_thresh::Vector{IOStream} +end + +function MSADTracker(chains; + theta_T::Float64=π/4, + output_schedule::Vector{Int}=Int[], + path::String=".", + scheduler::Vector{Int}=Int[], + kwargs...) + + compute_schedule = scheduler + + if !isempty(compute_schedule) && !isempty(output_schedule) + @assert all(t in compute_schedule for t in output_schedule) """ + output_schedule contains steps not in compute_schedule. + You cannot write output at a step where make_step! was not called. + """ + end + + n = length(chains) + states = [MSADState{Float64}() for _ in 1:n] + dirs = [joinpath(path, "chains", "$c") for c in 1:n] + + paths_integral = [joinpath(d, "phi_integral.dat") for d in dirs] + paths_thresh = [joinpath(d, "phi_thresh.dat") for d in dirs] + + files_integral = Vector{IOStream}(undef, n) + files_thresh = Vector{IOStream}(undef, n) + + return MSADTracker{Float64}(states, theta_T, compute_schedule, + output_schedule, paths_integral, paths_thresh, + files_integral, files_thresh) +end + +##### Function to write phi frame ##### +########## ########## + +function write_phi_frame(file::IOStream, t::Int, N_mol::Int, + phis::Vector{<:SVector}) + println(file, N_mol) + println(file, "t=$t") + for m in 1:N_mol + println(file, "$m $(phis[m][1]) $(phis[m][2]) $(phis[m][3])") + end + flush(file) +end + +##### Initialisation of the simulation ##### +# called once at t = 0, set the system, the mutable struct and the storing paths +########## ########## + +function Arianna.initialise(algorithm::MSADTracker, simulation::Simulation) + + for c in eachindex(simulation.chains) + system = simulation.chains[c] + state = algorithm.states[c] + T = typeof(system.temperature) + + # create output directory if it doesn't exist + mkpath(dirname(algorithm.paths_integral[c])) + + # open output files + algorithm.files_integral[c] = open(algorithm.paths_integral[c], "w") + algorithm.files_thresh[c] = open(algorithm.paths_thresh[c], "w") + + # compute initial body frames + R_all = get_all_body_frames(system) + N_mol = system.Nmol + + # fill state + state.R0 = copy(R_all) + state.R_prev = copy(R_all) + state.phi_integral = [zero(SVector{3,T}) for _ in 1:N_mol] + state.R_ref_thresh = copy(R_all) + state.phi_acc = [zero(SVector{3,T}) for _ in 1:N_mol] + state.initialized = true + + # write t=0, all phi vectors are zero + write_phi_frame(algorithm.files_integral[c], 0, N_mol, + state.phi_integral) + write_phi_frame(algorithm.files_thresh[c], 0, N_mol, + [zero(SVector{3,T}) for _ in 1:N_mol]) + end +end + +##### Finalise the simulation ##### +# close output to not keep unwritten data in memory +########## ########## + +function Arianna.finalise(tracker::MSADTracker, ::Simulation) + for c in eachindex(tracker.files_integral) + close(tracker.files_integral[c]) + close(tracker.files_thresh[c]) + end +end +##### Make a step in simulation ##### +########## ########## + +function Arianna.make_step!(simulation::Simulation, algorithm::MSADTracker) + t = simulation.t + + for c in eachindex(simulation.chains) + system = simulation.chains[c] + state = algorithm.states[c] + N_mol = system.Nmol + + # compute current body frames shared by all three methods + R_all = get_all_body_frames(system) + + # Integral update + for m in 1:N_mol + dR = state.R_prev[m]' * R_all[m] + state.phi_integral[m] = state.phi_integral[m] + rotation_vector(dR) + end + + # Threshold update + phi_current = Vector{SVector{3,eltype(algorithm.theta_T)}}(undef, N_mol) + phi_total = Vector{SVector{3,eltype(algorithm.theta_T)}}(undef, N_mol) + for m in 1:N_mol + dR = state.R_ref_thresh[m]' * R_all[m] + phi_current[m] = rotation_vector(dR) + phi_total[m] = state.phi_acc[m] + phi_current[m] + if norm(phi_current[m]) >= algorithm.theta_T + state.phi_acc[m] = phi_total[m] + state.R_ref_thresh[m] = R_all[m] + end + end + + state.R_prev .= R_all + + # output if time match output schedule + if t in algorithm.output_schedule + + # Integral + write_phi_frame(algorithm.files_integral[c],t,N_mol,state.phi_integral) + + # Threshold + write_phi_frame(algorithm.files_thresh[c],t,N_mol,phi_total) + end + end +end \ No newline at end of file diff --git a/test/bench_rotation.jl b/test/bench_rotation.jl new file mode 100644 index 0000000..deb7aa0 --- /dev/null +++ b/test/bench_rotation.jl @@ -0,0 +1,47 @@ +using StaticArrays, LinearAlgebra, BenchmarkTools + +function rotation_vector_skew(R::SMatrix{3,3,T}) where {T} + cos_θ = clamp((tr(R) - 1) / 2, -one(T), one(T)) + skew = (R - R') / 2 + sin_n = SVector(skew[3,2], skew[1,3], skew[2,1]) + sin_θ = norm(sin_n) + θ = atan(sin_θ, cos_θ) + if sin_θ > 1e-6 + return (θ / sin_θ) * sin_n + elseif cos_θ > 0 + return sin_n + else + RpI = R + one(R) + c1 = SVector(RpI[1,1], RpI[2,1], RpI[3,1]) + c2 = SVector(RpI[1,2], RpI[2,2], RpI[3,2]) + c3 = SVector(RpI[1,3], RpI[2,3], RpI[3,3]) + n = norm(c1) >= norm(c2) ? c1 : c2 + n = norm(n) >= norm(c3) ? n : c3 + return π * (n / norm(n)) + end +end + +function rotation_vector_eig(R::SMatrix{3,3,T}) where {T} + cos_θ = clamp((tr(R) - 1) / 2, -one(T), one(T)) + vals, vecs = eigen(R) + idx = argmin(abs.(real.(vals) .- 1)) + n = real.(vecs[:, idx]) + n = SVector{3,T}(n[1], n[2], n[3]) + n_skew = SMatrix{3,3,T}( + 0, n[3], -n[2], + -n[3], 0, n[1], + n[2], -n[1], 0 + ) + sin_θ = clamp(-tr(n_skew * R) / 2, -one(T), one(T)) + θ = atan(sin_θ, cos_θ) + return θ * n +end + +# 90° rotation around z — normal case +R = SMatrix{3,3,Float64}(0.0, 1.0, 0.0, -1.0, 0.0, 0.0, 0.0, 0.0, 1.0) + +println("=== Skew-symmetric ===") +@btime rotation_vector_skew($R) + +println("=== Eigenvector ===") +@btime rotation_vector_eig($R) diff --git a/test/test_msad.jl b/test/test_msad.jl new file mode 100644 index 0000000..0b12c2b --- /dev/null +++ b/test/test_msad.jl @@ -0,0 +1,157 @@ +using ParticlesMC +using StaticArrays +using LinearAlgebra +using Test + +@testset "body_frame" begin + + # known inputs — worked out by hand + r1 = SVector(0.0, 0.0, 0.0) + r2 = SVector(1.0, 0.0, 0.0) + r3 = SVector(0.0, 1.0, 0.0) + L = 10.0 + + R = ParticlesMC.body_frame(r1, r2, r3, L) + + # expected result from hand calculation: + # e1 = [1, 0, 0] + # e2 = [0, 0, 1] + # e3 = [0, -1, 0] + @test R[:, 1] ≈ SVector( 1.0, 0.0, 0.0) + @test R[:, 2] ≈ SVector( 0.0, 0.0, 1.0) + @test R[:, 3] ≈ SVector( 0.0, -1.0, 0.0) + + # R must be a valid rotation matrix: + # columns orthonormal → R'R = I + # right-handed → det(R) = +1 + @test R' * R ≈ I + @test det(R) ≈ 1.0 + +end + +@testset "body_frame minimum image" begin + + # same molecule but r2 is wrapped across the boundary + # r1 at 0.1, r2 at 9.9 — bond should be -0.2, not 9.8 + r1 = SVector(0.1, 0.0, 0.0) + r2 = SVector(9.9, 0.0, 0.0) + r3 = SVector(0.1, 1.0, 0.0) + L = 10.0 + + R_wrapped = ParticlesMC.body_frame(r1, r2, r3, L) + + # unwrapped version — same molecule, same frame + r2_unwrapped = SVector(-0.1, 0.0, 0.0) + R_unwrapped = ParticlesMC.body_frame(r1, r2_unwrapped, r3, L) + + @test R_wrapped ≈ R_unwrapped + +end + +@testset "rotation_vector" begin + + # identity matrix → zero rotation + R_id = SMatrix{3,3}(1.0I) + @test ParticlesMC.rotation_vector(R_id) ≈ SVector(0.0, 0.0, 0.0) + + # 90° rotation around z-axis + # R = [0 -1 0] + # [1 0 0] + # [0 0 1] + R_90z = SMatrix{3,3}(0.0, 1.0, 0.0, + -1.0, 0.0, 0.0, + 0.0, 0.0, 1.0) + Ω = ParticlesMC.rotation_vector(R_90z) + # should be θ=π/2 around z → Ω = [0, 0, π/2] + @test norm(Ω) ≈ π/2 + @test Ω / norm(Ω) ≈ SVector(0.0, 0.0, 1.0) + +end + +@testset "body_frame PBC - molecule split across each axis" begin + + L = 10.0 + + # --- x axis --- + # r1 near left wall, r2 wrapped to right wall + # true bond is [-0.2, 0, 0], not [9.8, 0, 0] + r1_x = SVector(0.1, 0.0, 0.0) + r2_x = SVector(9.9, 0.0, 0.0) # wrapped + r3_x = SVector(0.1, 1.0, 0.0) + R_wrapped_x = ParticlesMC.body_frame(r1_x, r2_x, r3_x, L) + R_unwrapped_x = ParticlesMC.body_frame(r1_x, SVector(-0.1, 0.0, 0.0), r3_x, L) + @test R_wrapped_x ≈ R_unwrapped_x + + # --- y axis --- + r1_y = SVector(0.0, 0.1, 0.0) + r2_y = SVector(1.0, 0.1, 0.0) + r3_y = SVector(0.0, 9.9, 0.0) # r3 wrapped across y + R_wrapped_y = ParticlesMC.body_frame(r1_y, r2_y, r3_y, L) + R_unwrapped_y = ParticlesMC.body_frame(r1_y, r2_y, SVector(0.0, -0.1, 0.0), L) + @test R_wrapped_y ≈ R_unwrapped_y + + # --- z axis --- + r1_z = SVector(0.0, 0.0, 0.1) + r2_z = SVector(1.0, 0.0, 0.1) + r3_z = SVector(0.0, 1.0, 0.1) + r2_z_wrapped = SVector(1.0, 0.0, 9.9 + 0.1) # wrapped across z... wait + # r1 near top wall, r2 wrapped to bottom + r1_z2 = SVector(0.0, 0.0, 0.1) + r2_z2 = SVector(0.0, 0.0, 9.9) # wrapped + r3_z2 = SVector(1.0, 0.0, 0.1) + R_wrapped_z = ParticlesMC.body_frame(r1_z2, r2_z2, r3_z2, L) + R_unwrapped_z = ParticlesMC.body_frame(r1_z2, SVector(0.0, 0.0, -0.1), r3_z2, L) + @test R_wrapped_z ≈ R_unwrapped_z + +end + +@testset "body_frame PBC - result is always a valid rotation matrix" begin + + L = 10.0 + # test many random molecules, some crossing boundaries + # a valid rotation matrix always satisfies R'R = I and det(R) = +1 + import Random + Random.seed!(42) + for _ in 1:100 + # random positions anywhere in box + r1 = SVector(rand()*L, rand()*L, rand()*L) + r2 = SVector(rand()*L, rand()*L, rand()*L) + r3 = SVector(rand()*L, rand()*L, rand()*L) + R = ParticlesMC.body_frame(r1, r2, r3, L) + @test R' * R ≈ I atol=1e-10 + @test det(R) ≈ 1.0 atol=1e-10 + end + +end + +@testset "rotation_vector PBC - roundtrip" begin + + # if we compute R between two frames, rotation_vector should + # give back an Ω whose norm equals the angle between them + # and applying it back should reconstruct the rotation + + # 180° rotation around x axis + R_180x = SMatrix{3,3}(1.0, 0.0, 0.0, + 0.0, -1.0, 0.0, + 0.0, 0.0, -1.0) + Ω = ParticlesMC.rotation_vector(R_180x) + @test norm(Ω) ≈ π atol=1e-10 + @test Ω / norm(Ω) ≈ SVector(1.0, 0.0, 0.0) atol=1e-10 + + # 180° rotation around y axis + R_180y = SMatrix{3,3}(-1.0, 0.0, 0.0, + 0.0, 1.0, 0.0, + 0.0, 0.0,-1.0) + Ω = ParticlesMC.rotation_vector(R_180y) + @test norm(Ω) ≈ π atol=1e-10 + @test Ω / norm(Ω) ≈ SVector(0.0, 1.0, 0.0) atol=1e-10 + + # small rotation — near-identity, tests the sin≈0 branch + ε = 1e-8 + R_tiny = SMatrix{3,3}(1.0, ε, 0.0, + -ε, 1.0, 0.0, + 0.0, 0.0, 1.0) + Ω = ParticlesMC.rotation_vector(R_tiny) + @test norm(Ω) < 1e-6 # nearly zero rotation + +end \ No newline at end of file diff --git a/test/test_tracker.jl b/test/test_tracker.jl new file mode 100644 index 0000000..ca962c0 --- /dev/null +++ b/test/test_tracker.jl @@ -0,0 +1,21 @@ +using ParticlesMC, StaticArrays, LinearAlgebra + +# check MSADState constructs empty +s = ParticlesMC.MSADState{Float64}() +println("MSADState empty: ", s.initialized == false) + +# check MSADTracker constructor with valid schedules +compute = collect(0:10:100) +output = collect(0:20:100) +chains = [1, 2] +t = ParticlesMC.MSADTracker(chains, pi/4, compute, output, "/tmp/test_msad") +println("MSADTracker created: ", t.theta_T ≈ pi/4) + +# check assertion fires on invalid schedules +try + bad_output = [0, 5, 10] # 5 not in compute schedule + ParticlesMC.MSADTracker(chains, pi/4, compute, bad_output, "/tmp/test_msad") + println("ERROR: should have thrown") +catch e + println("Assertion caught correctly: ", e isa AssertionError) +end \ No newline at end of file