Adaptive step-size Commutator free exponential solver (ACFET)

lib/BloqadeExpr/src/BloqadeExpr.jl

@@ -11,6 +11,7 @@ using Base.Threads: nthreads
 using MLStyle
 using BitBasis
 using LaTeXStrings
+using ForwardDiff
 using Unitful: Quantity, NoUnits, MHz, µm, uconvert
 using InteractiveUtils: subtypes
 using Base.Cartesian: @nexprs
@@ -20,7 +21,16 @@ using BloqadeLattices: BoundedLattice, rydberg_interaction_matrix
 
 include("Lowlevel/Lowlevel.jl")
 
-using .Lowlevel: Hamiltonian, SumOfLinop, ThreadedMatrix, storage_size, to_matrix, precision_type, highest_type, add_I, derivative, RegularLinop, isskewhermitian
+using .Lowlevel: Hamiltonian, 
+    SumOfLinop, 
+    ThreadedMatrix, 
+    storage_size, 
+    to_matrix, 
+    precision_type, 
+    highest_type, 
+    add_I, 
+    RegularLinop, 
+    isskewhermitian
 
 
 export rydberg_h,
@@ -49,10 +59,8 @@ export rydberg_h,
     emulate!,
     precision_type,
     highest_type,
-
     to_matrix,
     add_I,
-    derivative,
     RegularLinop,  # abstype
     SkewHermitian, # abstype
     isskewhermitian

lib/BloqadeExpr/src/Lowlevel/linalg.jl

@@ -201,7 +201,7 @@ end
 #    return SumOfLinop{Hermitian}(ForwardDiff.derivative.(h.fs,t), h.ts)
 #end
 
-function derivative(h::Hamiltonian, t::Real)
+function ForwardDiff.derivative(h::Hamiltonian, t::Real)
     ## remove terms that are zero
     fvals = ForwardDiff.derivative.(h.fs,t)
     mask = collect(fvals .!= 0)

lib/BloqadeExpr/src/Lowlevel/types.jl

@@ -26,7 +26,7 @@
 Base.size(m::ThreadedMatrix, i) = size(m.matrix)[i]
 Base.eltype(m::ThreadedMatrix) = eltype(m.matrix)
 Base.pointer(m::T) where {T <: Diagonal} = pointer(m.diag)
-
+Base.eltype(m::T) where {T <: ThreadedMatrix} = eltype(m.matrix) 
 
 precision_type(m::T) where {T <: Number} = real(typeof(m))
 precision_type(m::T) where {T <: Diagonal} = real(eltype(m))

lib/BloqadeExpr/src/types.jl

@@ -695,6 +695,7 @@ struct DivByTwo{F}
 end
 
 @inline (Func::DivByTwo)(t::Real) = Func.f(t)/2
+#ForwardDiff.derivative(Func::DivByTwo,t) = ForwardDiff.derivative(Func.f,t)/2
 
 
 const RabiTypes = Union{Nothing,SumOfX,SumOfXPhase}

lib/BloqadeExpr/test/linalg_derivative.jl

@@ -3,6 +3,7 @@ using BloqadeExpr
 using LinearAlgebra
 using Random
 using LuxurySparse
+using ForwardDiff
 using BloqadeExpr: Hamiltonian
 
 
@@ -21,7 +22,7 @@ using BloqadeExpr: Hamiltonian
     println(hlist.ts)
 
     for t in 0:0.1:2
-        src = derivative(hlist, t)
+        src = ForwardDiff.derivative(hlist, t)
         tar = htar(t)
 
         @test to_matrix(src) == to_matrix(tar)

lib/BloqadeKrylov/Project.toml

@@ -9,6 +9,7 @@ BloqadeLattices = "bd27d05e-4ce1-5e79-84dd-c5d7d508bbe4"
 BloqadeWaveforms = "bd27d05e-4ce1-5e79-84dd-c5d7d508bbe7"
 Configurations = "5218b696-f38b-4ac9-8b61-a12ec717816d"
 ExponentialUtilities = "d4d017d3-3776-5f7e-afef-a10c40355c18"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GaussQuadrature = "d54b0c1a-921d-58e0-8e36-89d8069c0969"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"

lib/BloqadeKrylov/src/BloqadeKrylov.jl

@@ -6,6 +6,7 @@ using LinearAlgebra
 using Configurations
 using YaoArrayRegister
 using YaoSubspaceArrayReg
+using ForwardDiff
 using BloqadeExpr: Hamiltonian, SumOfLinop
 using ExponentialUtilities
 using ProgressLogging
@@ -14,7 +15,7 @@ using GaussQuadrature
 export emulate!, emulate_step!
 export KrylovEvolution, Magnus4Evolution
 
-export CFETEvolution
+export CFETEvolution, ACFETEvolution
 export CFET2_1, CFET4_2, CFET6_5, CFET8_11
 
 # utils.jl introduce a new type of Hamiltonian called ValHamiltonian:
@@ -38,6 +39,9 @@ include("cfet.jl")
 include("tables/cfet_tbl.jl")
 
 
+## adaptive CFET 
+include("adapt_cfet.jl")
+
 
 if VERSION < v"1.7"
     include("patch.jl")

lib/BloqadeKrylov/src/adapt_cfet.jl

@@ -68,14 +68,18 @@ mutable struct ACFETEvolution{Reg<:AbstractRegister,T<:Real,H<:Hamiltonian,ALGTB
     hamiltonian::H
     options::KrylovOptions
     alg_table::ALGTBL
+    step_tol::T
 
-    function ACFETEvolution{Reg,T,H,ALGTBL}(reg, start_clock, end_clock, step_size,hamiltonian, options, algo) where {Reg,T,H,ALGTBL}
+    function ACFETEvolution{Reg,T,H,ALGTBL}(reg, start_clock, end_clock, step_size, hamiltonian, options, algo,step_tol) where {Reg,T,H,ALGTBL}
         start_clock ≥ 0 || throw(ArgumentError("start clock must not be negative"))
         end_clock ≥ 0 || throw(ArgumentError("end clock must not be negative"))
 
         end_clock ≥ start_clock || throw(ArgumentError("end clock must be larger than start clock"))
         step_size ≤ end_clock-start_clock || throw(ArgumentError("initial step size cannot be larger than the whole evolution time"))
-        return new{Reg,T,H,ALGTBL}(reg, start_clock, end_clock, step_size, hamiltonian, options, algo)
+        step_tol ≥ 0 || throw(ArgumentError("step tolerance must be positive"))
+
+        #min_step_size > 0 || throw(ArgumentError("minimum step size must be positive and >0."))
+        return new{Reg,T,H,ALGTBL}(reg, start_clock, end_clock, step_size, hamiltonian, options, algo,step_tol)
     end
 end
 
@@ -92,7 +96,7 @@ Create a `ACFETEvolution` object.
 - `hamiltonian`: low-level hamiltonian object of type [`Hamiltonian`](@ref).
 - `options`: options of the evolution in type [`KrylovOptions`](@ref).
 """
-function ACFETEvolution(reg, start_clock, end_clock, step_size,hamiltonian, options, algo)
+function ACFETEvolution(reg, start_clock, end_clock, step_size, hamiltonian, options, algo,step_tol)
     return ACFETEvolution{typeof(reg),typeof(start_clock),typeof(hamiltonian),typeof(algo)}(
         reg,
         start_clock,
@@ -101,25 +105,27 @@ function ACFETEvolution(reg, start_clock, end_clock, step_size,hamiltonian, opti
         hamiltonian,
         options,
         algo,
+        step_tol
     )
 end
 
 function Adapt.adapt_structure(to, x::ACFETEvolution)
-    return ACFETEvolution(adapt(to, x.reg), x.start_clock, x.end_clock, x.step_size, adapt(to, x.hamiltonian), x.options, x.alg_table)
+    return ACFETEvolution(adapt(to, x.reg), x.start_clock, x.end_clock, x.step_size, adapt(to, x.hamiltonian), x.options, x.alg_table,x.step_tol)
 end
 
-function ACFETEvolution(reg::AbstractRegister, start_clock, end_clock , h, algo::CFETTables = CFET2_1(); step_size=1e-7, kw...)
-    all(≥(0), clocks) || throw(ArgumentError("clocks must not be negative"))
+function ACFETEvolution(reg::AbstractRegister, start_clock, end_clock , h, algo::CFETTables = CFET2_1(); step_size=1e-7, step_tol=1e-12, kw...)
+    #all(≥(0), clocks) || throw(ArgumentError("clocks must not be negative"))
     options = from_kwargs(KrylovOptions; kw...)
     P = real(eltype(statevec(reg)))
     T = isreal(h) ? P : Complex{P}
 
-    return ACFETEvolution(reg, start_clock, end_clock, step_size, Hamiltonian(T, h, space(reg)), options, algo)
+    return ACFETEvolution(reg, start_clock, end_clock, step_size, Hamiltonian(T, h, space(reg)), options, algo, step_tol)
 end
 
 ## given Hamiltonian, current time `t` and duration `dt`, plus a CFETTable, 
 # output generator Ω at certain ETstep (exponential time propogator) 
 ## t + xs[i]dt
+#=
 function __construct_Ω(h::Hamiltonian, t::Real, dt::Real, Tbl::CFETTables, ETStep::Int)
 
     gs = Tbl.Gs[ETStep]
@@ -133,16 +139,17 @@ function __construct_Ω(h::Hamiltonian, t::Real, dt::Real, Tbl::CFETTables, ETSt
 
     return SumOfLinop{LinearAlgebra.Hermitian}(fs, h.ts)
 end
+=#
 
 function __construct_dΩ(h::Hamiltonian, t::Real, dt::Real, Tbl::CFETTables, ETStep::Int)
 
     gs = Tbl.Gs[ETStep]
     xs = Tbl.xs 
 
-    fs = gs[1]*xs[1]*derivative.(h.fs,t + xs[1]*dt)
+    fs = gs[1]*xs[1]*collect(ForwardDiff.derivative.(h.fs,t + xs[1]*dt))
 
     for i in 2:length(gs)      
-        fs += gs[i]*xs[i]*derivative.(h.fs,t + xs[i]*dt)
+        fs += gs[i]*xs[i]*collect(ForwardDiff.derivative.(h.fs,t + xs[i]*dt))
     end
 
     return SumOfLinop{LinearAlgebra.Hermitian}(fs, h.ts)
@@ -153,7 +160,8 @@ end
 end
 
 ## here, since generic algo for defect is unclear, we specialize it for each support types:
-function emulate_step!(prob::ACFETEvolution{<:Any,<:Any,<:Any,ALGTBL<:CFET2_1}, step::Int, clock::Real, tol::Real)
+function emulate_step!(prob::ACFETEvolution{<:Any,<:Any,<:Any,<:CFET2_1}, step::Int, clock::Real, tol::Real)
+    p = 2
     state = statevec(prob.reg)
     Ham = prob.hamiltonian
 
@@ -185,28 +193,179 @@ function emulate_step!(prob::ACFETEvolution{<:Any,<:Any,<:Any,ALGTBL<:CFET2_1},
     dest .+= (duration^2/2)*(Bv - dBv)
 
     ## -A(t+τ)v
-    mul!(tmp, Ham(clock+duration), state)
+    t_curr = clock+duration
+    mul!(tmp, -im*Ham(clock+duration), state)
     dest .-= tmp
 
-    ϵ = norm(dest)
+    ϵ = norm(dest)*duration/(p+1)
+
+    sp = scale_factor(0.25, 4.0 , 0.9, ϵ, tol, p)
+    ## calculate and update next step size:
+    @debug println("eps: $ϵ scale factor: $sp")
+    prob.step_size = duration*sp
+
+    ## if the next step_size will exceed the end_clock then set to the remainder
+    if t_curr + prob.step_size > prob.end_clock
+        prob.step_size = prob.end_clock - t_curr
+    end
+
+
+    # normalization
+    if mod(step, prob.options.normalize_step) == 0
+        normalize!(prob.reg)
+    end
+
+    if prob.options.normalize_finally && t_curr == prob.end_clock
+        normalize!(prob.reg)
+    end
+
+    return prob
+end
+
+
+## here, we just implemented in dump way, since we only need a few orders.
+## we can implement a generic way to do this, using recursive function.
+
+## tmp is the workspace for intermediate results
+function _comm_rk1(X,Y,state)
+    # this caluclate [X,Y]state
+    u = similar(state)
+    v = similar(state)
+    tmp = similar(v)
+
+    mul!(v,X,state); # v = Xψ 
+    mul!(u,Y,state); # u = Yψ
+
+    mul!(tmp,X,u); copyto!(u,tmp) # Xu = XYψ -> u
+    mul!(tmp,Y,v); copyto!(v,tmp) # Yv = YXψ -> v
+
+    tmp = nothing
+
+    return u .- v # u - v = [X,Y]ψ
+end
+
+function _comm_rk2(X,Y,state)
+    # this calculate [X,[X,Y]]v
+    tmp = similar(state)
+
+    w = _comm_rk1(X,Y,state) # w = [X,Y]state
+    mul!(tmp,X,w); copyto!(w,tmp) # w = X[X,Y]state
+
+    mul!(tmp,X,state) # tmp = Xstate 
+    r = _comm_rk1(X,Y,tmp) # r = [X,Y]Xstate
+
+    return w.-r 
+
+end
+
+function _comm_rk3(X,Y,state)
+    # this calculate [X,[X,[X,Y]]]v
+    tmp = similar(state)
+
+    w = _comm_rk2(X,Y,state) 
+    mul!(tmp,X,w); copyto!(w,tmp)
+
+    mul!(tmp,X,state)  
+    r = _comm_rk2(X,Y,tmp) 
+
+    return w.-r 
+
+end
+
+function _comm_rk4(X,Y,state)
+    # this calculate [X,[X,[X,Y]]]v
+    tmp = similar(state)
+
+    w = _comm_rk3(X,Y,state) 
+    mul!(tmp,X,w); copyto!(w,tmp)
+
+    mul!(tmp,X,state)  
+    r = _comm_rk3(X,Y,tmp) 
+
+    return w.-r 
+
+end
+
+
+function _gamma_p4(X,Y,t,state)
+    # Y = X'
+    tmp = similar(state)
+    dest = similar(state)
+
+    mul!(dest,X,state) #Xv
+    mul!(tmp,Y,state)
 
+    dest .+= t*tmp
+    tmp = nothing
+
+    dest .+= t^2/2*_comm_rk1(X,Y,state)
+    dest .+= t^3/6*_comm_rk2(X,Y,state)
+    dest .+= t^4/24*_comm_rk3(X,Y,state)
+    return dest
+end
+
+
+
+## here, since generic algo for defect is unclear, we specialize it for each support types:
+function emulate_step!(prob::ACFETEvolution{<:Any,<:Any,<:Any,<:CFET4_2}, step::Int, clock::Real, tol::Real)
+    p = 4
+    state = statevec(prob.reg)
+    Ham = prob.hamiltonian
+
+    duration = prob.step_size # get duration 
+
+    #construct Ωi:
+    Ω1 = -im*__construct_Ω(Ham, clock, duration, prob.alg_table, 1)
+    dΩ1 = -im*__construct_dΩ(Ham, clock, duration, prob.alg_table, 1)
+    Ω2 = -im*__construct_Ω(Ham, clock, duration, prob.alg_table, 2)
+    dΩ2 = -im*__construct_dΩ(Ham, clock, duration, prob.alg_table, 2)
+
+    # perform evolution:
+    ## each exponential-time prop. 
+    prob.options.expmv_backend(duration, Ω1, state; prob.options.tol)
+
+    d = _gamma_p4(Ω1, dΩ1, duration, state)
+    prob.options.expmv_backend(duration, Ω2, d; prob.options.tol)
+
+    prob.options.expmv_backend(duration, Ω2, state; prob.options.tol) ## this update state to the final results 
+
+
+
+    tmp = similar(state)
+    mul!(tmp, -im*Ham(clock+duration), state)
+    d .-= tmp
+    tmp = nothing
+    d .+= _gamma_p4(Ω2, dΩ2, duration, state)
+
+    t_curr = clock+duration
+    ϵ = norm(d)*duration/(p+1)
+
+    sp = scale_factor(0.25, 4.0 , 0.9, ϵ, tol, p)
     ## calculate and update next step size:
-    prob.step_size = duration*scale_factor(0.25, 4.0 , 0.9, ϵ, tol, p)
+    @debug println("eps: $ϵ scale factor: $sp")
+    prob.step_size = duration*sp
+
+    ## if the next step_size will exceed the end_clock then set to the remainder
+    if t_curr + prob.step_size > prob.end_clock
+        prob.step_size = prob.end_clock - t_curr
+    end
 
 
-    # do we need this normalization? 
+    # normalization
     if mod(step, prob.options.normalize_step) == 0
         normalize!(prob.reg)
     end
 
-    if prob.options.normalize_finally && step == length(prob.durations)
+    if prob.options.normalize_finally && t_curr == prob.end_clock
         normalize!(prob.reg)
     end
 
     return prob
 end
 
-function emulate_step!(prob::ACFETEvolution, step::Int, clock::Real, duration::Real)
-    error("unsupported CFET algo for adaptive step-size")
+
+
+function emulate_step!(prob::ACFETEvolution{<:Any,<:Any,<:Any,<:Any}, step::Int, clock::Real, duration::Real)
+    error("unsupported CFET algo for adaptive step-size. please choose from [CFET2_1, CFET4_2]")
 end