Replace eigendecomposition in rotation_vector with direct Rodrigues extraction (+ fix near-identity NaN)

src/rotation.jl

@@ -27,54 +27,60 @@ function body_frame(r1::SVector{3,T}, r2::SVector{3,T}, r3::SVector{3,T}, L::T)
     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
+function get_all_body_frames!(R_bodyframe, system::Molecules)
+    L   = system.box[1]
     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)
+    @inbounds for m in 1:system.Nmol
+        s = system.start_mol[m]
+        R_bodyframe[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))
+    ax = R[3,2] - R[2,3]
+    ay = R[1,3] - R[3,1]
+    az = R[2,1] - R[1,2]
+    sin_θ = sqrt(ax^2 + ay^2 + az^2) / 2
 
-    # theta 
-    θ     = atan(sin_θ, cos_θ)
-
-    return θ * n
+    if sin_θ > sqrt(eps(T))              # guard sin_θ ≈ 0
+        θ      = atan(sin_θ, cos_θ)
+        inv_2s = θ / (2 * sin_θ)
+        return SVector{3,T}(ax * inv_2s, ay * inv_2s, az * inv_2s)
+
+    elseif cos_θ > 0                     # guard θ ≈ 0, near identity
+        return zero(SVector{3,T})
+
+    else                                  # guard θ ≈ π
+        if R[1,1] >= R[2,2] && R[1,1] >= R[3,3]
+            nx = sqrt(max((R[1,1] + 1) / 2, zero(T)))
+            n  = SVector{3,T}(nx, (R[1,2] + R[2,1]) / (4nx), (R[1,3] + R[3,1]) / (4nx))
+        elseif R[2,2] >= R[3,3]
+            ny = sqrt(max((R[2,2] + 1) / 2, zero(T)))
+            n  = SVector{3,T}((R[1,2] + R[2,1]) / (4ny), ny, (R[2,3] + R[3,2]) / (4ny))
+        else
+            nz = sqrt(max((R[3,3] + 1) / 2, zero(T)))
+            n  = SVector{3,T}((R[1,3] + R[3,1]) / (4nz), (R[2,3] + R[3,2]) / (4nz), nz)
+        end
+        θ = atan(sin_θ, cos_θ)
+        return θ * n / norm(n)
+    end
 end
 
 ##### Definition of the RotationState ie: the accumulating variables #####
 # mutable because we need to update it 
 ##########                                                   ##########
 
 mutable struct RotationState{T}
-    R_ref::Vector{Vector{SMatrix{3,3,T,9}}}   
+    R_ref::Vector{Vector{SMatrix{3,3,T,9}}} 
+    R_bodyframe::Vector{SMatrix{3,3,T,9}} 
     Φ_acc::Vector{Vector{SVector{3,T}}}       
     initialized::Bool
 end
 
 # Start empty, filled during initialise when we first see the system
-RotationState{T}() where {T} = RotationState{T}([], [], false)
+RotationState{T}() where {T} = RotationState{T}([], [], [], false)
 
 ##### The Rotation computation algorithm #####
 ##########                      ##########
@@ -98,23 +104,21 @@ end
 ##########                        ##########
 
 function Arianna.initialise(algorithm::ComputeRotation, simulation::Simulation)
-
     for c in eachindex(simulation.chains)
         system = simulation.chains[c]
         state  = algorithm.states[c]
         T      = typeof(system.temperature)
         N_mol  = system.Nmol
         n_θ    = length(algorithm.theta_T)
 
-        # compute initial body frames
-        R_all = get_all_body_frames(system)
+        state.R_bodyframe = Vector{SMatrix{3,3,T,9}}(undef, N_mol)  # allocate once
+        get_all_body_frames!(state.R_bodyframe, system)              # fill state.R_bodyframe
 
-        # fill state
-        state.R_ref        = [copy(R_all) for _ in 1:n_θ]
-        state.Φ_acc      = [[zero(SVector{3,T}) for _ in 1:N_mol] for _ in 1:n_θ]
-        state.initialized  = true
-
-        resize!(system.Φ, n_θ) # a posteriori as we know n_θ
+        state.R_ref       = [copy(state.R_bodyframe) for _ in 1:n_θ]
+        state.Φ_acc       = [[zero(SVector{3,T}) for _ in 1:N_mol] for _ in 1:n_θ]
+        state.initialized = true
+
+        resize!(system.Φ, n_θ)
         for k in 1:n_θ
             system.Φ[k] = [zero(SVector{3,T}) for _ in 1:N_mol]
         end
@@ -125,20 +129,19 @@ end
 ##########                 ##########
 
 function Arianna.make_step!(simulation::Simulation, algorithm::ComputeRotation)
-
     tcollect(eachindex(simulation.chains) |> Transducers.Map(c -> begin
         system = simulation.chains[c]
         state  = algorithm.states[c]
         N_mol  = system.Nmol
-        R_all  = get_all_body_frames(system)
+        get_all_body_frames!(state.R_bodyframe, system)   # no new vector created updates state.R_bodyframe
         for (k, θ_T) in enumerate(algorithm.theta_T)
             for m in 1:N_mol
-                dR          = state.R_ref[k][m]' * R_all[m]
+                dR        = state.R_ref[k][m]' * state.R_bodyframe[m]
                 Φ_current = rotation_vector(dR)
                 Φ_total   = state.Φ_acc[k][m] + Φ_current
                 if norm(Φ_current) > θ_T
                     state.Φ_acc[k][m] = Φ_total
-                    state.R_ref[k][m]   = R_all[m]
+                    state.R_ref[k][m] = state.R_bodyframe[m]
                 end
                 system.Φ[k][m] = Φ_total
             end
@@ -157,7 +160,7 @@ function write_phi_frame(file::IOStream, t::Int, N_mol::Int, Φs::Vector{<:SVect
     for m in 1:N_mol
         println(file, "$m $(Φs[m][1]) $(Φs[m][2]) $(Φs[m][3])")
     end
-    flush(file)
+    #flush(file) not flushing at everyframe
 end
 
 struct StorePhiTrajectories <: AriannaAlgorithm