refactor: remove StaticArrays dependency and optimize polarization computations

Project.toml

@@ -12,7 +12,6 @@ FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 SpaceDataModel = "0b37b92c-f0c5-4a52-bd5c-390dec20857c"
-StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [weakdeps]
 DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
@@ -27,5 +26,4 @@ FFTW = "1"
 LinearAlgebra = "1"
 PrecompileTools = "1"
 SpaceDataModel = "0.2.2"
-StaticArrays = "1.9"
 julia = "1.10"

src/PlasmaWaves.jl

@@ -10,7 +10,6 @@ Plasma wave analysis.
 module PlasmaWaves
 using FFTW
 using LinearAlgebra
-using StaticArrays
 using SpaceDataModel: unwrap, times, cadence, SpaceDataModel
 
 using Bumper

src/helicty.jl

@@ -7,11 +7,9 @@ where
 
 ``γ = \\arctan (2 Re(𝐮)^𝐓 Im(𝐮) /(Re(𝐮)^2-Im(𝐮)^2))``
 """
-@inline function phase_factor(u)
-    Re = real(u)
-    Im = imag(u)
-    upper = 2 * Re ⋅ Im
-    lower = Re ⋅ Re - Im ⋅ Im
+@inline function phase_factor(u1, u2, u3)
+    upper = 2 * (real(u1) * imag(u1) + real(u2) * imag(u2) + real(u3) * imag(u3))
+    lower = real(u1)^2 + real(u2)^2 + real(u3)^2 - imag(u1)^2 - imag(u2)^2 - imag(u3)^2
     γ = atan(upper, lower)
     return exp(-im * γ / 2)
 end
@@ -30,47 +28,44 @@ Compute helicity and ellipticity for a single frequency.
 - `ellipticity`: Average ellipticity across the three components
 """
 function wpol_helicity(S, waveangle)
-    # Preallocate arrays for 3 polarization components
     helicity = 0
     ellipticity = 0
-    lam_u = MVector{3, ComplexF64}(undef)
 
     for comp in 1:3
-        # Build state vector λ_u for this polarization component
         alph = sqrt(real(S[comp, comp]))
-        lam_u[comp] = alph
         if comp == 1
-            lam_u[2] = (real(S[1, 2]) / alph) + im * (-imag(S[1, 2]) / alph)
-            lam_u[3] = (real(S[1, 3]) / alph) + im * (-imag(S[1, 3]) / alph)
+            l1 = complex(alph)
+            l2 = (real(S[1, 2]) / alph) + im * (-imag(S[1, 2]) / alph)
+            l3 = (real(S[1, 3]) / alph) + im * (-imag(S[1, 3]) / alph)
         elseif comp == 2
-            lam_u[1] = (real(S[2, 1]) / alph) + im * (-imag(S[2, 1]) / alph)
-            lam_u[3] = (real(S[2, 3]) / alph) + im * (-imag(S[2, 3]) / alph)
+            l1 = (real(S[2, 1]) / alph) + im * (-imag(S[2, 1]) / alph)
+            l2 = complex(alph)
+            l3 = (real(S[2, 3]) / alph) + im * (-imag(S[2, 3]) / alph)
         else
-            lam_u[1] = (real(S[3, 1]) / alph) + im * (-imag(S[3, 1]) / alph)
-            lam_u[2] = (real(S[3, 2]) / alph) + im * (-imag(S[3, 2]) / alph)
+            l1 = (real(S[3, 1]) / alph) + im * (-imag(S[3, 1]) / alph)
+            l2 = (real(S[3, 2]) / alph) + im * (-imag(S[3, 2]) / alph)
+            l3 = complex(alph)
         end
 
-        # Compute the phase rotation (gammay) for this state vector
-        lam_y = phase_factor(lam_u) * lam_u
+        p = phase_factor(l1, l2, l3)
+        y1 = p * l1
+        y2 = p * l2
+        y3 = p * l3
 
-        # Helicity: ratio of the norm of the imaginary part to the real part
-        norm_real = norm(real(lam_y))
-        norm_imag = norm(imag(lam_y))
+        norm_real = sqrt(real(y1)^2 + real(y2)^2 + real(y3)^2)
+        norm_imag = sqrt(imag(y1)^2 + imag(y2)^2 + imag(y3)^2)
         helicity += norm_imag / norm_real / 3
 
-        # For ellipticity, use only the first two components
-        u1 = lam_y[1]
-        u2 = lam_y[2]
-
         # TODO: why there is no 2 in front of uppere for `PySPEDAS`
-        uppere = 2 * (imag(u1) * real(u1) + imag(u2) * real(u2))
-        lowere = (-imag(u1)^2 + real(u1)^2 - imag(u2)^2 + real(u2)^2)
+        uppere = 2 * (imag(y1) * real(y1) + imag(y2) * real(y2))
+        lowere = -imag(y1)^2 + real(y1)^2 - imag(y2)^2 + real(y2)^2
         gammarot = atan(uppere, lowere)
         p_rot = exp(-1im * 0.5 * gammarot)
-        lam_urot = SA[p_rot * u1, p_rot * u2]
+        y1_rot = p_rot * y1
+        y2_rot = p_rot * y2
 
-        num = norm(imag(lam_urot))
-        den = norm(real(lam_urot))
+        num = sqrt(imag(y1_rot)^2 + imag(y2_rot)^2)
+        den = sqrt(real(y1_rot)^2 + real(y2_rot)^2)
         ellip_val = (den != 0) ? num / den : NaN
         # Adjust sign using the off-diagonal of ematspec and the wave normal angle
         sign_factor = sign(imag(S[1, 2]) * sin(waveangle))

src/polarization.jl

@@ -20,8 +20,14 @@ p^2  &= 1-\\frac{(tr 𝐒)^2-(tr 𝐒^2)}{(tr 𝐒)^2-n^{-1}(tr 𝐒)^2} \\\\
 """
 function polarization(S)
     n = size(S, 1)
-    trS2 = tr(S * S)
-    trS = tr(S)
+    trS = zero(eltype(S))
+    trS2 = zero(eltype(S))
+    @inbounds for i in axes(S, 1)
+        trS += S[i, i]
+        for j in axes(S, 2)
+            trS2 += S[i, j] * S[j, i]
+        end
+    end
     return real((n * trS2 - trS^2) / ((n - 1) * trS^2))
 end
 
@@ -110,7 +116,6 @@ function _wavpol(X::AbstractMatrix{T}, fs = 1; nfft = 256, noverlap = div(nfft,
 
     plan = plan_rfft(zeros(T, nfft, n), 1)
 
-    SfType = SMatrix{n, n, Complex{T}}
     Threads.@threads for j in 1:nsteps
         @no_escape begin
             Xw = @alloc(T, nfft, n)
@@ -127,7 +132,7 @@ function _wavpol(X::AbstractMatrix{T}, fs = 1; nfft = 256, noverlap = div(nfft,
             smooth_spectral_matrix!(Sm, S, smooth_f)
             # Compute the following polarization parameters from the spectral matrix ``S``:
             for f in 1:Nfreq
-                Sf = SfType(view(Sm, :, :, f))
+                Sf = @view Sm[:, :, f]
                 power[j, f] = real(tr(Sf))
                 degpol[j, f] = polarization(Sf)
                 waveangle[j, f] = wave_normal_angle(Sf)

src/svd.jl

@@ -1,29 +1,69 @@
-svd_S(S) = begin
-    A = @SMatrix [
-        real(S[1, 1])   real(S[1, 2])   real(S[1, 3]);
-        real(S[1, 2])   real(S[2, 2])   real(S[2, 3]);
-        real(S[1, 3])   real(S[2, 3])   real(S[3, 3]);
-        0               -imag(S[1, 2]) -imag(S[1, 3]);
-        imag(S[1, 2])   0               -imag(S[2, 3]);
-        imag(S[1, 3])   imag(S[2, 3])   0;
-    ]
-    return svd(A; full = false)
+function _smallest_eigenvector(a11, a22, a33, a12, a13, a23, λ)
+    b11, b22, b33 = a11 - λ, a22 - λ, a33 - λ
+    # Three columns of adj(A − λI); each lies in the null space. Pick the largest for stability.
+    c1 = (b22 * b33 - a23^2, a23 * a13 - a12 * b33, a12 * a23 - b22 * a13)
+    c2 = (a12 * b33 - a13 * a23, a13^2 - b11 * b33, b11 * a23 - a12 * a13)
+    c3 = (a12 * a23 - a13 * b22, a13 * a12 - b11 * a23, b11 * b22 - a12^2)
+    v = argmax(c -> c[1]^2 + c[2]^2 + c[3]^2, (c1, c2, c3))
+    invn = inv(sqrt(v[1]^2 + v[2]^2 + v[3]^2))
+    return v[1] * invn, v[2] * invn, v[3] * invn
+end
+
+# Upper triangle of B^T B, where B is the 6×3 real matrix [Re(S); skew(Im(S))] for 3×3 Hermitian S.
+# This equals the Gram matrix whose eigenvalues give SVD singular values squared.
+@inbounds function _gram_upper(S)
+    r11, r22, r33 = real(S[1, 1]), real(S[2, 2]), real(S[3, 3])
+    r12, i12 = reim(S[1, 2])
+    r13, i13 = reim(S[1, 3])
+    r23, i23 = reim(S[2, 3])
+
+    a11 = r11^2 + r12^2 + r13^2 + i12^2 + i13^2
+    a22 = r12^2 + r22^2 + r23^2 + i12^2 + i23^2
+    a33 = r13^2 + r23^2 + r33^2 + i13^2 + i23^2
+    a12 = r11 * r12 + r12 * r22 + r13 * r23 + i13 * i23
+    a13 = r11 * r13 + r12 * r23 + r13 * r33 - i12 * i23
+    a23 = r12 * r13 + r22 * r23 + r23 * r33 + i12 * i13
+    return a11, a22, a33, a12, a13, a23
 end
 
 function svd_polarization(S::AbstractMatrix)
-    _, Svals, V = svd_S(S)
-    # singular values in Svals (vector length = 3)
-    # take v = column 3 of V (associated with smallest singular value?) or as in your algorithm
-    v = V[:, 3]
-    v = v .* sign(v[3])
-    # compute θ, φ
-    theta = atan(sqrt(v[1]^2 + v[2]^2) / v[3])
-    phi = atan(v[2], v[1])
-    # compute planarity and ellipticity
-    # ellipticity: ratio of the two axes of polarization ellipse * sign of polarization
-    w1, w2, w3 = sort(Svals)
-    planarity = 1.0 - sqrt(w1 / w3)
-    ellipticity = (w2 / w3) * sign(imag(S[1, 2]))
+    a11, a22, a33, a12, a13, a23 = _gram_upper(S)
+
+    q = (a11 + a22 + a33) / 3
+    p1 = a12^2 + a13^2 + a23^2
+    if p1 == 0
+        λ1 = max(a11, a22, a33)
+        λ3 = min(a11, a22, a33)
+        λ2 = a11 + a22 + a33 - λ1 - λ3
+        imin = argmin((a11, a22, a33))
+        v1 = imin == 1 ? one(a11) : zero(a11)
+        v2 = imin == 2 ? one(a22) : zero(a22)
+        v3 = imin == 3 ? one(a33) : zero(a33)
+    else
+        p2 = (a11 - q)^2 + (a22 - q)^2 + (a33 - q)^2 + 2p1
+        p = sqrt(p2 / 6)
+        b11 = (a11 - q) / p
+        b22 = (a22 - q) / p
+        b33 = (a33 - q) / p
+        b12 = a12 / p
+        b13 = a13 / p
+        b23 = a23 / p
+        r = (b11 * b22 * b33 + 2 * b12 * b13 * b23 - b11 * b23^2 - b22 * b13^2 - b33 * b12^2) / 2
+        ϕ = acos(clamp(r, -1, 1)) / 3
+        λ1 = q + 2p * cos(ϕ)
+        λ3 = q + 2p * cos(ϕ + 2π / 3)
+        λ2 = 3q - λ1 - λ3
+        v1, v2, v3 = _smallest_eigenvector(a11, a22, a33, a12, a13, a23, λ3)
+    end
+
+    s = ifelse(v3 < 0, -one(v3), one(v3))
+    v1 *= s
+    v2 *= s
+    v3 *= s
+    theta = atan(sqrt(v1^2 + v2^2), v3)
+    phi = atan(v2, v1)
+    planarity = 1.0 - (max(λ3, 0) / λ1)^(1 // 4)
+    ellipticity = sqrt(max(λ2, 0) / λ1) * sign(imag(S[1, 2]))
     return (; theta, phi, planarity, ellipticity)
 end
 
@@ -50,15 +90,12 @@ function _wavpol_svd(X::AbstractMatrix{T}, fs = 1; nfft = 256, noverlap = div(nf
 
     plan = plan_rfft(zeros(T, nfft, n), 1)
 
-    # tforeach(1:nsteps) do j
     Threads.@threads for j in 1:nsteps
-        # tforeach(1:nsteps) do j
         @no_escape begin
             Xw = @alloc(T, nfft, n)
             Xf = @alloc(Complex{T}, Nfreq, n)
             S = @alloc(Complex{T}, n, n, Nfreq)
             Sm = @alloc(Complex{T}, n, n, Nfreq)
-
             start_idx = 1 + (j - 1) * noverlap
             end_idx = start_idx + nfft - 1
             Xw .= view(X, start_idx:end_idx, :) .* smooth_t

test/runtests.jl

@@ -13,31 +13,26 @@ end
     @test_call PlasmaWaves.workload()
 end
 
-@testset "svd of spectral matrix" begin
-    a = 3.5
-    b = 1.2
+@testset "svd_polarization" begin
+    # RHCP wave in xy-plane propagating along z: Xf = [a, ib, 0]
+    # → S[1,2] = -iab, wave normal = ẑ, planarity = 1, ellipticity = -b/a
+    a, b = 3.5, 1.2
     Xf = zeros(ComplexF64, 1, 3)
     Xf[1, 1] = a
     Xf[1, 2] = 1im * b
-    S = spectral_matrix(Xf)
-    Sf = @view S[:, :, 1]
+    S3d = spectral_matrix(Xf)
+    Sf = @view S3d[:, :, 1]
+    @test Sf ≈ [a^2 -im*a*b 0; im*a*b b^2 0; 0 0 0]
+    res = PlasmaWaves.svd_polarization(Sf)
+    @test res.theta ≈ 0
+    @test res.planarity ≈ 1.0
+    @test res.ellipticity ≈ -b / a
 
-    @test Sf ≈ [a^2 -im * a * b 0; im * a * b b^2 0; 0 0 0]
-    A = [
-        a^2     0     0;
-        0       b^2   0;
-        0       0     0;
-        0       a * b 0;
-        -a * b  0     0;
-        0       0     0
-    ]
-    U, S, V = PlasmaWaves.svd_S(Sf)
-    @test U * Diagonal(S) * V' ≈ A
-    @test S ≈ [a * sqrt(a^2 + b^2), b * sqrt(a^2 + b^2), 0]
-    s = Base.sign(U[1, 1])
-    @test s * U[1, 1] ≈ a / sqrt(a^2 + b^2) + 0im
-    @test s * U[2, 2] ≈ b / sqrt(a^2 + b^2) + 0im
-    @test s * U[4, 2] ≈ a / sqrt(a^2 + b^2) + 0im
+    # Wave normal along x: Gram matrix has a11=λ₃=0, a12=a13=0
+    S_x = ComplexF64[0 0 0; 0 4.0 2.0; 0 2.0 8.0]
+    res_x = PlasmaWaves.svd_polarization(S_x)
+    @test !any(isnan, (res_x.theta, res_x.phi, res_x.planarity, res_x.ellipticity))
+    @test res_x.theta ≈ π / 2
 end
 
 @testset "Spectral matrix from time sequence" begin