Various small changes related to real-time evolution, timers and documentation

src/NaturalGradient/NaturalGradient.jl

@@ -173,7 +173,7 @@ function tdvp_relative_error(J::Jacobian, Es::EnergySummary, θdot::Vector)
 end
 
 function NaturalGradient_timeit_wrapper(θ, Oks_and_Eks_; kwargs...)
-    if kwargs[:timer] !== nothing
+    if haskey(kwargs, :timer) && kwargs[:timer] !== nothing
         ng = @timeit kwargs[:timer] "NaturalGradient" NaturalGradient(θ, Oks_and_Eks_; kwargs...)
     else
         ng = NaturalGradient(θ, Oks_and_Eks_; kwargs...)

src/solver/LinearSolveWrapper.jl

@@ -17,6 +17,12 @@ function (solver::LinearSolveWrapper)(M::AbstractMatrix, v::AbstractArray)
     else
         M = Hermitian(M)
     end
+
+    if eltype(M) != eltype(v)
+        if eltype(M) === ComplexF64
+            v = ComplexF64.(v)
+        end
+    end
 
     prob = LinearProblem(M, v)
     sol = solve(prob, solver.alg)

src/solver/eigen_solver.jl

@@ -8,6 +8,11 @@ mutable struct EigenSolver <: AbstractSolver
 end
 
 
+"""
+(solver::AbstractSolver)(M::AbstractMatrix, v::AbstractArray, double::Bool; method=:auto, kwargs...)
+
+returns o = M^-1 v (computed by eigendecomposition)
+"""
 function (solver::EigenSolver)(M::AbstractMatrix, v::AbstractArray)
     #@assert ishermitian(M) "EigenSolver: M is not Hermitian"
     eig = eigen(Hermitian(M))

src/solver/solver.jl

@@ -1,27 +1,41 @@
 abstract type AbstractSolver end
 
+"""
+solve_S(solver::AbstractSolver, J::Jacobian, grad_half::Vector; timer=TimerOutput(), kwargs...)
+
+In the terms of https://arxiv.org/pdf/2503.12557, this returns θdot = (O' * O)^-1 * O' * E_loc   (Eq. 11)
+"""
 function solve_S(solver::AbstractSolver, J::Jacobian, grad_half::Vector; timer=TimerOutput(), kwargs...)
-    @timeit "dense_S" Jd = dense_S(J)
-    @timeit "solve" θdot = -solver(Jd, grad_half; kwargs...)
+    @timeit timer "dense_S" Jd = dense_S(J)
+    @timeit timer "solve" θdot = -solver(Jd, grad_half; kwargs...)
 
     return θdot
 end
 
+"""
+solve_T(solver::AbstractSolver, J::Jacobian, Es::EnergySummary; timer=TimerOutput(), kwargs...)
+
+In the terms of https://arxiv.org/pdf/2503.12557, this returns θdot = O' * (O * O')^-1 * E_loc   (Eq. 13)
+"""
 function solve_T(solver::AbstractSolver, J::Jacobian, Es::EnergySummary; timer=TimerOutput(), kwargs...)
-    @timeit "dense_T" Jd = dense_T(J)
+    @timeit timer "dense_T" Jd = dense_T(J)
     Ekms = centered(Es)
-
-    @timeit "solve" θdot_raw = -solver(Jd, Ekms; kwargs...)
-    θdot = centered(J)' * θdot_raw
+    @timeit timer "solve" θdot_raw = -solver(Jd, Ekms; kwargs...)
+    @timeit timer "mult." θdot = centered(J)' * θdot_raw
 
     return θdot
 end
 
+"""
+(solver::AbstractSolver)(ng::NaturalGradient; method=:auto, compute_error=true, kwargs...)
+
+Computes θdot corresponding to equations (11) or (13) of https://arxiv.org/pdf/2503.12557, depending on the method (solve_T corresponds to eq. (13), while solve_S corresponds to eq. (11)).
+"""
 function (solver::AbstractSolver)(ng::NaturalGradient; method=:auto, compute_error=true, kwargs...)
     if method === :T || (method === :auto && nr_samples(ng.J) < nr_parameters(ng.J))
         ng.θdot = solve_T(solver, ng.J, ng.Es; kwargs...)
     else
-        ng.θdot = solve_S(solver, ng.J, get_gradient(ng) ./ 2; kwargs...)
+        ng.θdot = solve_S(solver, ng.J, get_gradient_timeit_wrapper(ng; kwargs...) ./ 2; kwargs...)
     end
     if compute_error
         tdvp_error!(ng)