on-the-fly computation of MSAD for trimers

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,26 +20,27 @@ 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"
 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"]

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

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

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)