From b94cf47dbe744eb7d046364b6e99e5ecf50e82f4 Mon Sep 17 00:00:00 2001
From: Kipton Barros <kbarros@gmail.com>
Date: Wed, 27 Nov 2024 11:13:17 -0700
Subject: [PATCH] Lanczos passes smoke test

---
 src/KPM/Lanczos.jl           | 10 +++++-----
 src/KPM/SpinWaveTheoryKPM.jl | 18 +++++++++++-------
 test/test_kpm.jl             |  4 ++--
 3 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/src/KPM/Lanczos.jl b/src/KPM/Lanczos.jl
index 1c8057f09..c57c020da 100644
--- a/src/KPM/Lanczos.jl
+++ b/src/KPM/Lanczos.jl
@@ -64,13 +64,13 @@ end
 function lanczos(mulA!, mulS!, v; niters, lhs=zeros(length(v), 0))
     αs = Float64[]
     βs = Float64[]
+    lhs_adj_Q = Vector{ComplexF64}[]
 
     vp  = zero(v)
     Sv  = zero(v)
     Svp = zero(v)
     w   = zero(v)
     Sw  = zero(v)
-    Q_adj_lhs = zeros(niters, size(lhs, 2))
 
     mulS!(Sv, v)
     c = sqrt(real(dot(v, Sv)))
@@ -83,9 +83,9 @@ function lanczos(mulA!, mulS!, v; niters, lhs=zeros(length(v), 0))
     @. w = w - α * v
     mulS!(Sw, w)
     push!(αs, α)
-    mul!(view(Q_adj_lhs, 1:1, :), v', lhs)
+    push!(lhs_adj_Q, lhs' * v)
 
-    for i in 2:niters
+    for _ in 2:niters
         β = sqrt(real(dot(w, Sw)))
         iszero(β) && break
         @. vp = w / β
@@ -97,10 +97,10 @@ function lanczos(mulA!, mulS!, v; niters, lhs=zeros(length(v), 0))
         @. v = vp
         push!(αs, α)
         push!(βs, β)
-        mul!(view(Q_adj_lhs, i:i, :), v', lhs)
+        push!(lhs_adj_Q, lhs' * v)
     end
 
-    return SymTridiagonal(αs, βs), Q_adj_lhs
+    return SymTridiagonal(αs, βs), reduce(hcat, lhs_adj_Q)
 end
 
 
diff --git a/src/KPM/SpinWaveTheoryKPM.jl b/src/KPM/SpinWaveTheoryKPM.jl
index 6a465a0a9..f551dcdc2 100644
--- a/src/KPM/SpinWaveTheoryKPM.jl
+++ b/src/KPM/SpinWaveTheoryKPM.jl
@@ -185,7 +185,7 @@ end
 function intensities(swt_kpm::SpinWaveTheoryKPM, qpts; energies, kernel::AbstractBroadening, kT=0.0, verbose=false)
     qpts = convert(AbstractQPoints, qpts)
     data = zeros(eltype(swt_kpm.swt.measure), length(energies), length(qpts.qs))
-    return intensities!(data, swt_kpm, qpts; energies, kernel, kT, verbose)
+    return intensities2!(data, swt_kpm, qpts; energies, kernel, kT, verbose)
 end
 
 
@@ -201,19 +201,22 @@ function intensities2!(data, swt_kpm::SpinWaveTheoryKPM, qpts; energies, kernel:
 
     @assert eltype(data) == eltype(measure)
     @assert size(data) == (length(energies), length(qpts.qs))
+    fill!(data, zero(eltype(data)))
 
     Na = nsites(sys)
     Ncells = Na / natoms(cryst)
     Nf = nflavors(swt)
     L = Nf*Na
     Avec_pref = zeros(ComplexF64, Na)
-    u = zeros(ComplexF64, Nobs, 2L)
-    v = zeros(ComplexF64, 2L)
 
     Nobs = size(measure.observables, 1)
     Ncorr = length(measure.corr_pairs)
     corrbuf = zeros(ComplexF64, Ncorr)
 
+    u = zeros(ComplexF64, 2L, Nobs)
+    v = zeros(ComplexF64, 2L)
+    Sv = zeros(ComplexF64, 2L)
+
     for (iq, q) in enumerate(qpts.qs)
         q_reshaped = to_reshaped_rlu(sys, q)
         q_global = cryst.recipvecs * q
@@ -263,18 +266,19 @@ function intensities2!(data, swt_kpm::SpinWaveTheoryKPM, qpts; energies, kernel:
             return w
         end
 
-        data[iω, iq] .= zero(eltype(data))
-
         for ξ in 1:Nobs
             mulA!(v, view(u, :, ξ))
-            tridiag, Q_adj_lhs = lanczos(mulA!, mulS!, v; lhs=u, niters=lanczos_iters)
+            mulS!(Sv, v)
+            c = sqrt(Sv' * v)
+            v ./= c
+            tridiag, lhs_adj_Q = lanczos(mulA!, mulS!, v; lhs=u, niters=lanczos_iters)
 
             (; values, vectors) = eigen(tridiag)
 
             for (iω, ω) in enumerate(energies)
                 f(x) = kernel(x, ω) * thermal_prefactor(x; kT)
 
-                corr_ξ = Q_adj_lhs' * vectors * Diagonal(f.(values)) * (vectors')[:, 1]
+                corr_ξ = c * lhs_adj_Q * vectors * Diagonal(f.(values)) * (vectors'[:, 1])
 
                 # This step assumes that each local observable in the
                 # correlation is Hermitian. In this case, bare correlations
diff --git a/test/test_kpm.jl b/test/test_kpm.jl
index 6ddfc39b7..aad73cc75 100644
--- a/test/test_kpm.jl
+++ b/test/test_kpm.jl
@@ -19,11 +19,11 @@
 
     mulA!(w, v) = mul!(w, A, v)
     mulS!(w, v) = mul!(w, S, v)
-    T, Q_adj_lhs = Sunny.lanczos(mulA!, mulS!, copy(v); lhs, niters=10)
+    T, lhs_adj_Q = Sunny.lanczos(mulA!, mulS!, copy(v); lhs, niters=10)
     ev2 = eigvals(T)
 
     @test ev1 ≈ ev2
-    @test Q' * lhs ≈ Q_adj_lhs
+    @test lhs' * Q ≈ lhs_adj_Q
 
     # Check that extremal eigenvalues match to reasonable accuracy