Add ExponentialUtilities GRAPE support + tests

Author: krugx

Project.toml

@@ -15,19 +15,25 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [weakdeps]
+ExponentialUtilities = "d4d017d3-3776-5f7e-afef-a10c40355c18"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+QuantumGradientGenerators = "a563f35e-61db-434d-8c01-8b9e3ccdfd85"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [extensions]
+QuantumPropagatorsExponentialUtilitiesExt = "ExponentialUtilities"
 QuantumPropagatorsODEExt = "OrdinaryDiffEq"
+QuantumPropagatorsQuantumGradientGeneratorsExt = "QuantumGradientGenerators"
 QuantumPropagatorsRecursiveArrayToolsExt = "RecursiveArrayTools"
 QuantumPropagatorsStaticArraysExt = "StaticArrays"
 
 [compat]
+ExponentialUtilities = "1.11"
 OffsetArrays = "1"
 OrdinaryDiffEq = "6.59"
 ProgressMeter = "1"
+QuantumGradientGenerators = "0.1"
 Random = "1"
 RecursiveArrayTools = "3.12"
 SpecialFunctions = "1, 2"

docs/Project.toml

@@ -5,6 +5,7 @@ DocInventories = "43dc2714-ed3b-44b5-b226-857eda1aa7de"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
 DocumenterInterLinks = "d12716ef-a0f6-4df4-a9f1-a5a34e75c656"
+ExponentialUtilities = "d4d017d3-3776-5f7e-afef-a10c40355c18"
 LiveServer = "16fef848-5104-11e9-1b77-fb7a48bbb589"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"

docs/make.jl

@@ -3,6 +3,7 @@ using Pkg
 using Documenter
 using QuantumPropagators
 using QuantumPropagators: AbstractPropagator, set_t!, set_state!
+import ExponentialUtilities  # ensure ExponentialUtilities extension is loaded
 import OrdinaryDiffEq  # ensure ODE extension is loaded
 using DocumenterCitations
 using DocumenterInterLinks

docs/src/methods.md

@@ -15,6 +15,7 @@ As discussed in the [Overview](@ref overview_approaches), time propagation can b
    We consider this especially in the piecewise-constant case (`pwc=true` in [`propagate`](@ref)/[`init_prop`](@ref)), which is required for the traditional optimization methods [GRAPE](https://juliaquantumcontrol.github.io/GRAPE.jl/stable/) and [Krotov](https://juliaquantumcontrol.github.io/Krotov.jl/stable/). In these propagations, the time-dependent generator ``\op{H}(t)`` is [evaluated](@ref QuantumPropagators.Controls.evaluate) to a constant operator ``\op{H}`` on each interval of the time grid. The analytical solution to the Schrödinger or Liouville equation is well known, and propagation step simply has to evaluate the application of the time evolution operator ``\op{U} = \exp[-i \op{H} dt]`` to the state ``|Ψ⟩``. The following methods are built in to `QuantumPropagators`:
 
    * [`ExpProp`](@ref method_expprop) – constructs ``\op{U}`` explicitly and then applies it to ``|Ψ⟩``
+   * [`ExponentialUtilities`](@ref method_exponentialutilities) – applies ``\exp(\op{H} dt) |Ψ⟩`` using Krylov methods without explicitly forming ``\op{U}``
    * [`Cheby`](@ref method_cheby) — expansion of ``\op{U} |Ψ⟩`` into Chebychev polynomials, valid if ``\op{H}`` has real eigenvalues
    * [`Newton`](@ref method_newton) – expansion of ``\op{U} |Ψ⟩`` into Newton polynomials, valid if ``\op{H}`` has complex eigenvalues (non-Hermitian Hamiltonian, Liouvillian)
 
@@ -65,6 +66,52 @@ init_prop(state, generator, tlist, method::Val{:ExpProp}; kwargs...)
 * Comparing against another propagator
 
 
+## [`ExponentialUtilities`](@id method_exponentialutilities)
+
+The method requires that the [ExponentialUtilities.jl](https://docs.sciml.ai/ExponentialUtilities/stable/)
+package is loaded
+
+```
+using ExponentialUtilities
+```
+
+and then passed as `method=ExponentialUtilities` (or `method=:expv`) to
+[`propagate`](@ref) or [`init_prop`](@ref):
+
+```@docs
+init_prop(state, generator, tlist, method::Val{:ExponentialUtilities}; kwargs...)
+```
+
+This method evaluates ``\exp(-i \op{H} dt) |Ψ⟩`` via a Krylov expv algorithm
+without explicitly forming the matrix exponential. It is therefore often a
+good fit for larger systems or matrix-free operators where direct matrix
+exponentiation is too costly.
+
+**Advantages**
+
+* Avoids explicit construction of ``\op{U}``
+* Works with matrix-free operators that support `mul!`
+* Good scaling for large sparse systems
+
+**Disadvantages**
+
+* Requires ExponentialUtilities.jl
+* Performance depends on Krylov parameters and operator structure
+
+**When to use**
+
+* Large, sparse, or matrix-free generators
+* Piecewise-constant GRAPE-style workflows that need expv
+
+**GRAPE note**
+
+* When using `method=:expv` with GRAPE, set `gradient_method=:taylor`.
+  The default `:gradgen` path uses `GradVector`/`GradgenOperator`, which is
+  not currently supported by the ExponentialUtilities Krylov backend.
+* For non-Hermitian generators (e.g., Liouvillians), pass
+  `prop_expv_kwargs=(; ishermitian=false)` to avoid Hermitian-only paths.
+
+
 ## [`Cheby`](@id method_cheby)
 
 The method should be loaded with

ext/QuantumPropagatorsExponentialUtilitiesExt.jl

@@ -0,0 +1,201 @@
+module QuantumPropagatorsExponentialUtilitiesExt
+
+using LinearAlgebra
+using TimerOutputs: @timeit_debug, TimerOutput
+using ExponentialUtilities
+
+using QuantumPropagators:
+    QuantumPropagators,
+    PWCPropagator,
+    _pwc_get_max_genop,
+    _pwc_get_genop,
+    _pwc_set_genop!,
+    _pwc_set_t!,
+    _pwc_advance_time!,
+    _pwc_process_parameters
+using QuantumPropagators.Controls: get_controls
+using QuantumPropagators.Interfaces: supports_inplace
+using QuantumPropagators.Generators: ScaledOperator
+import QuantumPropagators: init_prop, prop_step!, set_t!
+
+struct _SizedOp{A}
+    A::A
+end
+
+Base.size(op::_SizedOp) = size(op.A)
+Base.size(op::_SizedOp, dim::Integer) = size(op.A)[dim]
+Base.eltype(op::_SizedOp) = eltype(op.A)
+LinearAlgebra.ishermitian(op::_SizedOp) = LinearAlgebra.ishermitian(op.A)
+LinearAlgebra.mul!(y, op::_SizedOp, x) = mul!(y, op.A, x)
+
+_ensure_size_dim(A) =
+    hasmethod(size, Tuple{typeof(A), Int}) ? A :
+    (hasmethod(size, Tuple{typeof(A)}) ? _SizedOp(A) : A)
+
+
+"""Propagator for Krylov expv propagation via ExponentialUtilities (`method=ExponentialUtilities`).
+
+This is a [`PWCPropagator`](@ref).
+"""
+mutable struct ExpvPropagator{GT,OT,ST} <: PWCPropagator
+    const generator::GT
+    state::ST
+    t::Float64  # time at which current `state` is defined
+    n::Int64 # index of next interval to propagate
+    const tlist::Vector{Float64}
+    parameters::AbstractDict
+    controls
+    genop::OT
+    backward::Bool
+    inplace::Bool
+    convert_state::Type
+    convert_operator::Type
+    expv_kwargs::NamedTuple
+    const timing_data::TimerOutput
+end
+
+
+set_t!(propagator::ExpvPropagator, t) = _pwc_set_t!(propagator, t)
+
+
+_expv_convert_state(state) = typeof(state)
+_expv_convert_operator(::Any) = Any
+_expv_convert_operator(::QuantumPropagators.Generator) = Matrix{ComplexF64}
+_expv_convert_operator(::QuantumPropagators.Operator) = Matrix{ComplexF64}
+_expv_convert_operator(::QuantumPropagators.ScaledOperator) = Matrix{ComplexF64}
+
+
+"""
+```julia
+using ExponentialUtilities
+
+expv_propagator = init_prop(
+    state,
+    generator,
+    tlist;
+    method=ExponentialUtilities,
+    inplace=QuantumPropagators.Interfaces.supports_inplace(state),
+    backward=false,
+    verbose=false,
+    parameters=nothing,
+    expv_kwargs=(;),
+    convert_state=_expv_convert_state(state),
+    convert_operator=_expv_convert_operator(generator),
+    _...
+)
+```
+
+initializes an [`ExpvPropagator`](@ref).
+
+# Method-specific keyword arguments
+
+* `expv_kwargs`: NamedTuple of keyword arguments forwarded to
+  `ExponentialUtilities.expv`.
+* `convert_state`: Type to which to temporarily convert the state before
+  calling `expv`.
+* `convert_operator`: Type to which to convert the operator before calling
+  `expv`.
+"""
+function init_prop(
+    state,
+    generator,
+    tlist,
+    method::Val{:ExponentialUtilities};
+    inplace = supports_inplace(state),
+    backward = false,
+    verbose = false,
+    parameters = nothing,
+    expv_kwargs = (;),
+    convert_state = _expv_convert_state(state),
+    convert_operator = _expv_convert_operator(generator),
+    _...
+)
+    tlist = convert(Vector{Float64}, tlist)
+    controls = get_controls(generator)
+    G = _pwc_get_max_genop(generator, controls, tlist)
+
+    parameters = _pwc_process_parameters(parameters, controls, tlist)
+    timing_data = TimerOutput()
+
+    n = 1
+    t = tlist[1]
+    if backward
+        n = length(tlist) - 1
+        t = float(tlist[n + 1])
+    end
+    GT = typeof(generator)
+    OT = typeof(G)
+    ST = typeof(state)
+    return ExpvPropagator{GT,OT,ST}(
+        generator,
+        inplace ? copy(state) : state,
+        t,
+        n,
+        tlist,
+        parameters,
+        controls,
+        G,
+        backward,
+        inplace,
+        convert_state,
+        convert_operator,
+        expv_kwargs,
+        timing_data,
+    )
+end
+
+
+# Aliases
+init_prop(state, generator, tlist, method::Val{:expv}; kwargs...) =
+    init_prop(state, generator, tlist, Val(:ExponentialUtilities); kwargs...)
+
+
+function prop_step!(propagator::ExpvPropagator)
+    @timeit_debug propagator.timing_data "prop_step!" begin
+        if nameof(typeof(propagator.state)) == :GradVector &&
+           nameof(parentmodule(typeof(propagator.state))) == :QuantumGradientGenerators
+            throw(ArgumentError(
+                "ExponentialUtilities propagation does not support GRAPE `gradient_method=:gradgen`. " *
+                "Use `gradient_method=:taylor` instead."
+            ))
+        end
+        H = propagator.genop
+        n = propagator.n
+        tlist = getfield(propagator, :tlist)
+        (0 < n < length(tlist)) || return nothing
+        dt = tlist[n + 1] - tlist[n]
+        if propagator.backward
+            dt = -dt
+        end
+        dt_expv = complex(dt)
+
+        Ψ = convert(propagator.convert_state, propagator.state)
+        if propagator.inplace
+            if supports_inplace(propagator.genop)
+                _pwc_set_genop!(propagator, n)
+                H = convert(propagator.convert_operator, propagator.genop)
+            else
+                H = convert(propagator.convert_operator, _pwc_get_genop(propagator, n))
+            end
+            H = _ensure_size_dim(H)
+            H = ScaledOperator(-1im, H)
+            @timeit_debug propagator.timing_data "expv" begin
+                Ψ = ExponentialUtilities.expv(dt_expv, H, Ψ; propagator.expv_kwargs...)
+            end
+            copyto!(propagator.state, convert(typeof(propagator.state), Ψ))
+        else
+            H = convert(propagator.convert_operator, _pwc_get_genop(propagator, n))
+            H = _ensure_size_dim(H)
+            H = ScaledOperator(-1im, H)
+            @timeit_debug propagator.timing_data "expv" begin
+                Ψ = ExponentialUtilities.expv(dt_expv, H, Ψ; propagator.expv_kwargs...)
+            end
+            setfield!(propagator, :state, convert(typeof(propagator.state), Ψ))
+        end
+
+        _pwc_advance_time!(propagator)
+        return propagator.state
+    end
+end
+
+end

ext/QuantumPropagatorsQuantumGradientGeneratorsExt.jl

@@ -0,0 +1,11 @@
+module QuantumPropagatorsQuantumGradientGeneratorsExt
+
+using QuantumGradientGenerators: GradgenOperator, GradVector
+
+Base.size(G::GradgenOperator, dim::Integer) = size(G.G, dim)
+Base.similar(::GradVector, ::Type{T}, dims::Tuple{Int,Int}) where {T} =
+    Matrix{T}(undef, dims...)
+Base.similar(::GradVector, ::Type{T}, dims::Tuple{Int}) where {T} =
+    Vector{T}(undef, dims[1])
+
+end

src/generators.jl

@@ -177,6 +177,8 @@ function Base.copyto!(tgt::Operator, src::Operator)
 end
 
 Base.size(O::Operator) = size(O.ops[1])
+Base.size(O::Operator, dim::Integer) = size(O.ops[1], dim)
+Base.eltype(O::Operator) = eltype(O.ops[1])
 
 
 function LinearAlgebra.ishermitian(O::Operator{OT,CT}) where {OT,CT}
@@ -251,6 +253,8 @@ Base.Array(O::ScaledOperator) = Array{ComplexF64}(O)
 
 
 Base.size(O::ScaledOperator) = size(O.operator)
+Base.size(O::ScaledOperator, dim::Integer) = size(O.operator, dim)
+Base.eltype(O::ScaledOperator) = promote_type(eltype(O.operator), typeof(O.coeff))
 
 LinearAlgebra.ishermitian(O::ScaledOperator) = (isreal(O.coeff) && ishermitian(O.operator))
 

src/newton.jl

@@ -34,7 +34,7 @@ mutable struct NewtonWrk{T}
     restarts::Int64
     m_max::Int64
     timing_data::TimerOutput
-    function NewtonWrk(v0::T; m_max::Int64 = 10, _timing_data = TimerOutput()) where {T}
+    function NewtonWrk(v0::T; m_max::Int64=10, _timing_data=TimerOutput()) where {T}
         if m_max <= 2
             error("Newton propagation requires m_max > 2")
         end
@@ -60,8 +60,8 @@ mutable struct NewtonWrk{T}
 end
 
 
-lbound(array::OffsetArray, dim = 1) = first(axes(array)[dim])
-ubound(array::OffsetArray, dim = 1) = last(axes(array)[dim])
+lbound(array::OffsetArray, dim=1) = first(axes(array)[dim])
+ubound(array::OffsetArray, dim=1) = last(axes(array)[dim])
 
 
 function leja_radius(z)
@@ -281,9 +281,9 @@ function newton!(Ψ, H, dt, wrk; kwargs...)
                 wrk.v,
                 H,
                 _dt;
-                extended = true,
-                norm_min = norm_min,
-                _timing_data = wrk.timing_data
+                extended=true,
+                norm_min=norm_min,
+                _timing_data=wrk.timing_data
             )
         end
         if m == 1 && s == 0
@@ -294,7 +294,7 @@ function newton!(Ψ, H, dt, wrk; kwargs...)
             break
         end
         @timeit_debug wrk.timing_data "diagonalize_hessenberg_matrix" begin
-            ritz = diagonalize_hessenberg_matrix(Hess, m, accumulate = true)
+            ritz = diagonalize_hessenberg_matrix(Hess, m, accumulate=true)
         end
 
         # In the first iteration, the radius will be determined

test/Project.toml

@@ -8,6 +8,7 @@ DocInventories = "43dc2714-ed3b-44b5-b226-857eda1aa7de"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
 DocumenterInterLinks = "d12716ef-a0f6-4df4-a9f1-a5a34e75c656"
+ExponentialUtilities = "d4d017d3-3776-5f7e-afef-a10c40355c18"
 GRAPE = "6b52fcaf-80fe-489a-93e9-9f92080510be"
 GRAPELinesearchAnalysis = "290eba36-e2d8-4488-81b6-f66cc44f2186"
 IOCapture = "b5f81e59-6552-4d32-b1f0-c071b021bf89"