Update operator interface

docs/make.jl

@@ -21,29 +21,34 @@ links = InterLinks(
     "QuantumControl" => "https://juliaquantumcontrol.github.io/QuantumControl.jl/$DEV_OR_STABLE",
 )
 
+fallbacks = ExternalFallbacks(
+    "QuantumPropagators.Interfaces.supports_inplace" => "@extref QuantumPropagators :jl:function:`QuantumPropagators.Interfaces.supports_inplace`",
+    automatic = false,
+)
+
 println("Starting makedocs")
 
 makedocs(;
-    plugins=[links],
-    authors=AUTHORS,
-    sitename="QuantumGradientGenerators.jl",
-    doctest=false,
-    format=Documenter.HTML(;
-        prettyurls=true,
-        canonical="https://juliaquantumcontrol.github.io/QuantumGradientGenerators.jl",
-        assets=[
+    plugins = [links, fallbacks],
+    authors = AUTHORS,
+    sitename = "QuantumGradientGenerators.jl",
+    doctest = false,
+    format = Documenter.HTML(;
+        prettyurls = true,
+        canonical = "https://juliaquantumcontrol.github.io/QuantumGradientGenerators.jl",
+        assets = [
             asset(
                 "https://juliaquantumcontrol.github.io/QuantumControl.jl/dev/assets/topbar/topbar.css"
             ),
             asset(
                 "https://juliaquantumcontrol.github.io/QuantumControl.jl/dev/assets/topbar/topbar.js"
             ),
         ],
-        footer="[$NAME.jl]($GITHUB) v$VERSION docs powered by [Documenter.jl](https://github.com/JuliaDocs/Documenter.jl)."
+        footer = "[$NAME.jl]($GITHUB) v$VERSION docs powered by [Documenter.jl](https://github.com/JuliaDocs/Documenter.jl)."
     ),
-    pages=["Home" => "index.md", "API" => "api.md",]
+    pages = ["Home" => "index.md", "API" => "api.md",]
 )
 
 println("Finished makedocs")
 
-deploydocs(; repo="github.com/JuliaQuantumControl/QuantumGradientGenerators.jl",)
+deploydocs(; repo = "github.com/JuliaQuantumControl/QuantumGradientGenerators.jl",)

src/evaluate.jl

@@ -6,7 +6,12 @@ function evaluate(O::GradgenOperator, args...; kwargs...)
 end
 
 
-function evaluate!(G::GradgenOperator, gradgen::GradGenerator, args...; vals_dict=IdDict())
+function evaluate!(
+    G::GradgenOperator,
+    gradgen::GradGenerator,
+    args...;
+    vals_dict = IdDict()
+)
     evaluate!(G.G, gradgen.G, args...; vals_dict)
     for (i, control) in enumerate(gradgen.controls)
         μ = gradgen.control_derivs[i]
@@ -18,7 +23,7 @@ function evaluate!(G::GradgenOperator, gradgen::GradGenerator, args...; vals_dic
 end
 
 
-function evaluate(gradgen::GradGenerator, args...; vals_dict=IdDict())
+function evaluate(gradgen::GradGenerator, args...; vals_dict = IdDict())
     G = evaluate(gradgen.G, args...; vals_dict)
     control_deriv_ops = [evaluate(μ, args...; vals_dict) for μ ∈ gradgen.control_derivs]
     num_controls = length(control_deriv_ops)

src/gradgen_operator.jl

@@ -30,7 +30,7 @@ function get_controls(O1::GradgenOperator)
 end
 
 
-function random_state(H::GradgenOperator; rng=GLOBAL_RNG, _...)
+function random_state(H::GradgenOperator; rng = GLOBAL_RNG, _...)
     state = random_state(H.G; rng)
     num_controls = length(H.control_deriv_ops)
     grad_states = [random_state(H.G; rng) for i ∈ eachindex(H.control_deriv_ops)]

src/linalg.jl

@@ -75,8 +75,16 @@ function Base.length(Ψ::GradVector)
 end
 
 
-function Base.size(G::GradgenOperator)
-    return Base.size(G.G)
+function Base.size(O::GradgenOperator{num_controls,GT,CGT}) where {num_controls,GT,CGT}
+    return (num_controls + 1) .* size(O.G)
+end
+
+
+function Base.size(
+    O::GradgenOperator{num_controls,GT,CGT},
+    dim::Integer
+) where {num_controls,GT,CGT}
+    return (num_controls + 1) * size(O.G, dim)
 end
 
 
@@ -88,6 +96,9 @@ function Base.similar(G::GradgenOperator{num_controls,GT,CGT}) where {num_contro
     return GradgenOperator{num_controls,GT,CGT}(similar(G.G), similar(G.control_deriv_ops))
 end
 
+function Base.eltype(O::GradgenOperator{num_controls,GT,CGT}) where {num_controls,GT,CGT}
+    return promote_type(eltype(GT), eltype(CGT))
+end
 
 function Base.copyto!(dest::GradgenOperator, src::GradgenOperator)
     copyto!(dest.G, src.G)
@@ -161,23 +172,23 @@ function *(
 end
 
 
-@inline function convert_gradgen_to_dense(G)
+@inline function convert_gradgen_to_dense(G::GradGenerator)
     N = size(G.G)[1]
     L = length(G.control_derivs)
     G_full = zeros(eltype(G.G), N * (L + 1), N * (L + 1))
     convert_gradgen_to_dense!(G_full, G)
 end
 
 
-@inline function convert_gradgen_to_dense!(G_full, G)
+@inline function convert_gradgen_to_dense!(G_full, G::GradGenerator)
     N = size(G.G)[1]
     L = length(G.control_derivs)
-    @inbounds for i = 1:L+1
-        G_full[(i-1)*N+1:i*N, (i-1)*N+1:i*N] .= G.G
+    @inbounds for i = 1:(L+1)
+        G_full[((i-1)*N+1):(i*N), ((i-1)*N+1):(i*N)] .= G.G
     end
     # Set the control-derivatives in the last (block-)column
     @inbounds for i = 1:L
-        G_full[(i-1)*N+1:i*N, L*N+1:(L+1)*N] .= G.control_derivs[i]
+        G_full[((i-1)*N+1):(i*N), (L*N+1):((L+1)*N)] .= G.control_derivs[i]
     end
     return G_full
 end
@@ -195,9 +206,9 @@ end
     N = length(Ψ.state)
     L = length(Ψ.grad_states)
     @inbounds for i = 1:L
-        Ψ_full[(i-1)*N+1:i*N] .= Ψ.grad_states[i]
+        Ψ_full[((i-1)*N+1):(i*N)] .= Ψ.grad_states[i]
     end
-    @inbounds Ψ_full[L*N+1:(L+1)*N] .= Ψ.state
+    @inbounds Ψ_full[(L*N+1):((L+1)*N)] .= Ψ.state
     return Ψ_full
 end
 
@@ -206,9 +217,9 @@ end
     N = length(Ψ.state)
     L = length(Ψ.grad_states)
     @inbounds for i = 1:L
-        Ψ.grad_states[i] .= Ψ_full[(i-1)*N+1:i*N]
+        Ψ.grad_states[i] .= Ψ_full[((i-1)*N+1):(i*N)]
     end
-    @inbounds Ψ.state .= Ψ_full[L*N+1:(L+1)*N]
+    @inbounds Ψ.state .= Ψ_full[(L*N+1):((L+1)*N)]
     return Ψ
 end
 
@@ -225,8 +236,8 @@ function Base.convert(
     L = num_controls
     N = length(vec) ÷ (L + 1)  # dimension of state
     @assert length(vec) == (L + 1) * N
-    grad_states = [convert(T, vec[(i-1)*N+1:i*N]) for i = 1:L]
-    state = convert(T, vec[L*N+1:(L+1)*N])
+    grad_states = [convert(T, vec[((i-1)*N+1):(i*N)]) for i = 1:L]
+    state = convert(T, vec[(L*N+1):((L+1)*N)])
     return GradVector{num_controls,T}(state, grad_states)
 end
 
@@ -236,8 +247,7 @@ function Base.Array{T}(G::GradgenOperator) where {T}
     𝟘 = zeros(T, N, M)
     μ = G.control_deriv_ops
     block_rows = [
-        hcat([𝟘 for j = 1:i-1]..., Array{T}(G.G), [𝟘 for j = i+1:L]..., Array{T}(μ[i]))
-        for i = 1:L
+        hcat([𝟘 for j = 1:(i-1)]..., Array{T}(G.G), [𝟘 for j = (i+1):L]..., Array{T}(μ[i])) for i = 1:L
     ]
     last_block_row = hcat([𝟘 for j = 1:L]..., Array{T}(G.G))
     return Base.Array{T}(vcat(block_rows..., last_block_row))

test/clean.jl

@@ -4,10 +4,10 @@
 Clean up build/doc/testing artifacts. Restore to clean checkout state
 (distclean)
 """
-function clean(; distclean=false, _exit=true)
+function clean(; distclean = false, _exit = true)
 
     _glob(folder, ending) =
-        [name for name in readdir(folder; join=true) if (name |> endswith(ending))]
+        [name for name in readdir(folder; join = true) if (name |> endswith(ending))]
     _exists(name) = isfile(name) || isdir(name)
     _push!(lst, name) = _exists(name) && push!(lst, name)
 
@@ -33,12 +33,12 @@ function clean(; distclean=false, _exit=true)
 
     for name in CLEAN
         @info "rm $name"
-        rm(name, force=true, recursive=true)
+        rm(name, force = true, recursive = true)
     end
     if distclean
         for name in DISTCLEAN
             @info "rm $name"
-            rm(name, force=true, recursive=true)
+            rm(name, force = true, recursive = true)
         end
         if _exit
             @info "Exiting"
@@ -48,4 +48,4 @@ function clean(; distclean=false, _exit=true)
 
 end
 
-distclean() = clean(distclean=true)
+distclean() = clean(distclean = true)

test/test_gradgen.jl

@@ -81,17 +81,20 @@ using QuantumPropagators.Controls: evaluate
         Ψ̃.state
     ]
 
+    @test size(G̃) == size(G̃_full)
+    @test eltype(G̃) == eltype(G̃_full)
+
     # proper initialization? grad_states should be zero
     @test norm(Ψ̃_full) == norm(Ψ̃.state) == norm(Ψ)
 
     Ψ̃_out_full = exp(-𝕚 * G̃_full * dt) * Ψ̃_full
     # propagation correct?
-    @test norm(Ψ̃_out_full[2N+1:3N] - Û_Ψ) < 1e-10
+    @test norm(Ψ̃_out_full[(2N+1):3N] - Û_Ψ) < 1e-10
 
     # do we get the same results as from newton?
-    @test norm(Ψ̃_out_full[2N+1:3N] - Ψ̃_out.state) < 1e-10
+    @test norm(Ψ̃_out_full[(2N+1):3N] - Ψ̃_out.state) < 1e-10
     @test norm(Ψ̃_out_full[1:N] - Ψ̃_out.grad_states[1]) < 1e-10
-    @test norm(Ψ̃_out_full[N+1:2N] - Ψ̃_out.grad_states[2]) < 1e-10
+    @test norm(Ψ̃_out_full[(N+1):2N] - Ψ̃_out.grad_states[2]) < 1e-10
 
     ###########################################################################
     # Test custom expprop
@@ -108,12 +111,12 @@ using QuantumPropagators.Controls: evaluate
 
     Ψ̃_out_full = exp(-𝕚 * G̃_full * dt) * Ψ̃_full
     # propagation correct?
-    @test norm(Ψ̃_out_full[2N+1:3N] - Û_Ψ) < 1e-10
+    @test norm(Ψ̃_out_full[(2N+1):3N] - Û_Ψ) < 1e-10
 
     # do we get the same results as from newton?
-    @test norm(Ψ̃_out_full[2N+1:3N] - Ψ̃_out.state) < 1e-10
+    @test norm(Ψ̃_out_full[(2N+1):3N] - Ψ̃_out.state) < 1e-10
     @test norm(Ψ̃_out_full[1:N] - Ψ̃_out.grad_states[1]) < 1e-10
-    @test norm(Ψ̃_out_full[N+1:2N] - Ψ̃_out.grad_states[2]) < 1e-10
+    @test norm(Ψ̃_out_full[(N+1):2N] - Ψ̃_out.grad_states[2]) < 1e-10
 
     ###########################################################################
     # Test standard expprop
@@ -122,10 +125,10 @@ using QuantumPropagators.Controls: evaluate
         Ψ̃,
         G̃,
         [0, dt];
-        method=:expprop,
-        inplace=true,
-        convert_state=Vector{ComplexF64},
-        convert_operator=Matrix{ComplexF64}
+        method = :expprop,
+        inplace = true,
+        convert_state = Vector{ComplexF64},
+        convert_operator = Matrix{ComplexF64}
     )
     Ψ̃_out_exp = prop_step!(propagator)
     @test norm(Ψ̃_out_exp - Ψ̃_out) < 1e-11
@@ -159,14 +162,14 @@ using QuantumPropagators.Controls: evaluate
     Ψ̃_out_full2 = exp(-𝕚 * G̃_full2 * dt) * Ψ̃_full2
 
     # propagation correct?
-    @test norm(Ψ̃_out_full1[N+1:2N] - Û_Ψ) < 1e-12
-    @test norm(Ψ̃_out_full2[N+1:2N] - Û_Ψ) < 1e-12
+    @test norm(Ψ̃_out_full1[(N+1):2N] - Û_Ψ) < 1e-12
+    @test norm(Ψ̃_out_full2[(N+1):2N] - Û_Ψ) < 1e-12
 
     # do we get the same results as with the combined grad-gen?
     @test norm(Ψ̃_out_full1[1:N] - Ψ̃_out_full[1:N]) < 1e-12
-    @test norm(Ψ̃_out_full2[1:N] - Ψ̃_out_full[N+1:2N]) < 1e-12
-    @test norm(Ψ̃_out_full1[N+1:2N] - Ψ̃_out_full[2N+1:3N]) < 1e-12
-    @test norm(Ψ̃_out_full2[N+1:2N] - Ψ̃_out_full[2N+1:3N]) < 1e-12
+    @test norm(Ψ̃_out_full2[1:N] - Ψ̃_out_full[(N+1):2N]) < 1e-12
+    @test norm(Ψ̃_out_full1[(N+1):2N] - Ψ̃_out_full[(2N+1):3N]) < 1e-12
+    @test norm(Ψ̃_out_full2[(N+1):2N] - Ψ̃_out_full[(2N+1):3N]) < 1e-12
 
     ###########################################################################
     # Compare against Zygote

test/test_interface.jl

@@ -36,42 +36,42 @@ end
 @testset "GradGenerator Interface" begin
 
     N = 10
-    Ĥ₀ = random_matrix(N, hermitian=true)
-    Ĥ₁ = random_matrix(N, hermitian=true)
-    Ĥ₂ = random_matrix(N, hermitian=true)
+    Ĥ₀ = random_matrix(N, hermitian = true)
+    Ĥ₁ = random_matrix(N, hermitian = true)
+    Ĥ₂ = random_matrix(N, hermitian = true)
     ϵ₁(t) = 1.0
     ϵ₂(t) = 1.0
     Ĥ_of_t = hamiltonian(Ĥ₀, (Ĥ₁, ϵ₁), (Ĥ₂, ϵ₂))
 
-    tlist = collect(range(0, 10; length=101))
+    tlist = collect(range(0, 10; length = 101))
 
     G̃_of_t = GradGenerator(Ĥ_of_t)
 
     Ψ = random_state_vector(N)
     Ψ̃ = GradVector(Ψ, length(get_controls(G̃_of_t)))
 
-    @test check_generator(G̃_of_t; state=Ψ̃, tlist, for_gradient_optimization=false)
+    @test check_generator(G̃_of_t; state = Ψ̃, tlist, for_gradient_optimization = false)
 
 end
 
 
 @testset "GradGenerator Interface (Static)" begin
 
     N = 10
-    Ĥ₀ = SMatrix{N,N,ComplexF64}(random_matrix(N, hermitian=true))
-    Ĥ₁ = SMatrix{N,N,ComplexF64}(random_matrix(N, hermitian=true))
-    Ĥ₂ = SMatrix{N,N,ComplexF64}(random_matrix(N, hermitian=true))
+    Ĥ₀ = SMatrix{N,N,ComplexF64}(random_matrix(N, hermitian = true))
+    Ĥ₁ = SMatrix{N,N,ComplexF64}(random_matrix(N, hermitian = true))
+    Ĥ₂ = SMatrix{N,N,ComplexF64}(random_matrix(N, hermitian = true))
     ϵ₁(t) = 1.0
     ϵ₂(t) = 1.0
     Ĥ_of_t = hamiltonian(Ĥ₀, (Ĥ₁, ϵ₁), (Ĥ₂, ϵ₂))
 
-    tlist = collect(range(0, 10; length=101))
+    tlist = collect(range(0, 10; length = 101))
 
     G̃_of_t = GradGenerator(Ĥ_of_t)
 
     Ψ = SVector{N,ComplexF64}(random_state_vector(N))
     Ψ̃ = GradVector(Ψ, length(get_controls(G̃_of_t)))
 
-    @test check_generator(G̃_of_t; state=Ψ̃, tlist, for_gradient_optimization=false)
+    @test check_generator(G̃_of_t; state = Ψ̃, tlist, for_gradient_optimization = false)
 
 end