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

## Summary 

In order to improve performance I propose three fixes : 
- `rotation_vector` : use of Rodrigues formula to leverage knowledge of
orthogonality for our matrices instead of blindly use `eigen`, which
runs an iterative eigensolver and allocates, when the closed-form
Rodrigues extraction needs only a trace and a few subtractions..
- 'R_bodyframe' : declaration of the variable so we are not creating a
new vector at each steps but just redirecting the stored one.
- 'flush' : removed the flush from `make_step!` but keep the one in
`write_phi_frame` which is enough.

To prevent ill defined operation while using the closed-form Rodrigues
extraction which divide by `sin θ` :
- `sin θ > √eps` general case, safe to divide.
- `sin θ ≈ 0, cos θ > 0` θ ≈ 0 (near identity), return the zero vector.
- `sin θ ≈ 0, cos θ ≤ 0` θ ≈ π, recover the axis from the symmetric
part.

EDIT : One can run this code to verify the validity of the closed form,
especially at the boundaries.

```julia
using LinearAlgebra
using StaticArrays
using Printf

# eigen based 
function rotation_vector_eigen(R::SMatrix{3,3,T}) where {T}
    cos_θ      = clamp((tr(R) - 1) / 2, -one(T), one(T))
    vals, vecs = eigen(Matrix(R))
    idx        = argmin(abs.(real.(vals) .- 1))
    n          = real.(vecs[:, idx])
    n          = SVector{3,T}(n[1], n[2], n[3])
    n          = n / norm(n)
    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

# closed form
function rotation_vector_closed(R::SMatrix{3,3,T}) where {T}
    cos_θ = clamp((tr(R) - 1) / 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

    if sin_θ > sqrt(eps(T))                       
        θ      = atan(sin_θ, cos_θ)
        inv_2s = θ / (2 * sin_θ)
        return SVector{3,T}(ax * inv_2s, ay * inv_2s, az * inv_2s)

    elseif cos_θ > 0                             
        return zero(SVector{3,T})

    else                                          # θ ≈ π
        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

# rotation vector to rotation matrix using Rodrigues 
function rot_matrix(axis::SVector{3,T}, θ::T) where {T}
    n = axis / norm(axis)
    K = SMatrix{3,3,T}(0, n[3], -n[2], -n[3], 0, n[1], n[2], -n[1], 0)
    return SMatrix{3,3,T}(1I) + sin(θ)*K + (1 - cos(θ))*(K*K)
end

same(a, b; atol) = isapprox(a, b; atol=atol) || isapprox(a, -b; atol=atol)

const T    = Float64
const AXES = [SVector{3,T}(1,0,0), SVector{3,T}(1,1,0), SVector{3,T}(1,-2,3), SVector{3,T}(-1,0.5,2)] # one can choose whatever axis 

ok = true

# Tcheck aggrement for general case
println("CASE 1  eigen vs closed form, θ ∈ [0.1, 3.0]")
@printf("%-16s %-7s %-12s %-12s %-11s\n", "axis", "θ", "‖closed‖", "‖eigen‖", "|Δ|")
worst1 = 0.0
for axn in AXES, θ in (0.1, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0)
    R  = rot_matrix(axn, T(θ)) # def rotation matrix from axis and θ
    vc = rotation_vector_closed(R) # extract rotation vector from closed form method
    ve = rotation_vector_eigen(R) # from eigen method
    d  = min(maximum(abs.(vc .- ve)), maximum(abs.(vc .+ ve)))
    global worst1 = max(worst1, d)
    global ok &= isapprox(norm(vc), θ; atol=1e-8) && same(vc, ve; atol=1e-6)
    @printf("%-16s %-7.2f %-12.6f %-12.6f %-11.2e\n",
            string(round.(axn; digits=1)), θ, norm(vc), norm(ve), d)
end
@printf("\nworst |Δ| (mod sign): %.2e\n\n", worst1)

# Tcheck aggrement for case near θ = 0 
println("CASE 2  near identity")
@printf("%-16s %-7s %-24s %-24s\n", "axis", "θ", "eigen", "closed")
worst2 = 0.0
for axn in AXES, θ in (1e-3, 1e-6, 1e-9)
    R  = rot_matrix(axn, T(θ))
    vc = rotation_vector_closed(R)
    ve = rotation_vector_eigen(R)
    global worst2 = max(worst2, norm(vc - ve), 0.0)
    global ok &= all(isfinite, vc) && all(isfinite, ve)
    @printf("%-16s %-7.0e %-24s %-24s\n",
            string(round.(axn; digits=1)), θ,
            string(round.(ve; digits=10)), string(round.(vc; digits=10)))
end
@printf("\nworst |Δ| near 0: %.2e\n\n", worst2)

# Tcheck aggrement case near θ = π
println("CASE 3  near θ = π : magnitude and axis")
@printf("%-16s %-26s %-26s\n", "true n̂", "eigen n̂", "closed n̂")
for axn in AXES
    n  = axn / norm(axn)
    R  = rot_matrix(axn, T(π) - 1e-7)
    ve = rotation_vector_eigen(R); vc = rotation_vector_closed(R)
    global ok &= isapprox(norm(vc), π; atol=1e-3) && same(vc/norm(vc), SVector{3,T}(n); atol=1e-2)
    @printf("%-16s %-26s %-26s\n",
            string(round.(n; digits=3)),
            string(round.(ve ./ norm(ve); digits=3)),
            string(round.(vc ./ norm(vc); digits=3)))
end
println("\nboth give |Φ| = θ ≈ π and the correct axis.\n")


println(ok ? "PASSED : closed form matches eigen across all cases." :
             "FAILED")
```

Author: cyfraysse

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