integrator signature refactor -- removes sys dependence

src/integrators/_integrators.jl

@@ -12,6 +12,8 @@ export QuantumStateExponentialIntegrator
 
 export DerivativeIntegrator
 
+export G_bilinear
+
 export jacobian
 export hessian_of_the_lagrangian
 
@@ -23,11 +25,10 @@ export sixth_order_pade
 export eighth_order_pade
 export tenth_order_pade
 
-using ..QuantumSystems
 using ..Isomorphisms
-using ..QuantumObjectUtils
-using ..Losses
-using ..QuantumSystemUtils
+# using ..QuantumSystems
+# using ..QuantumObjectUtils
+# using ..QuantumSystemUtils
 
 using NamedTrajectories
 using TrajectoryIndexingUtils
@@ -36,6 +37,7 @@ using SparseArrays
 using ForwardDiff
 using TestItemRunner
 
+
 abstract type AbstractIntegrator end
 
 abstract type QuantumIntegrator <: AbstractIntegrator end

src/integrators/exponential_integrators.jl

@@ -4,6 +4,14 @@ This file includes expoential integrators for states and unitaries
 
 using ExponentialAction
 
+function exp_eigen(G::AbstractMatrix)
+    Ĥ = Hermitian(Matrix(Isomorphisms.H(G)))
+    λ, V = eigen(Ĥ)
+    expG = iso(sparse(V * Diagonal(exp.(-im * λ)) * V'))
+    droptol!(expG, 1e-12)
+    return expG
+end
+
 # ----------------------------------------------------------------------------- #
 #                         Quantum Exponential Integrators                       #
 # ----------------------------------------------------------------------------- #
@@ -27,15 +35,14 @@ struct UnitaryExponentialIntegrator <: QuantumExponentialIntegrator
     G::Function
 
     function UnitaryExponentialIntegrator(
-        sys::AbstractQuantumSystem,
         unitary_name::Symbol,
         drive_name::Union{Symbol, Tuple{Vararg{Symbol}}},
+        G::Function,
         traj::NamedTrajectory;
-        G::Function=a -> G_bilinear(a, sys.G_drift, sys.G_drives),
         autodiff::Bool=false
     )
-        ketdim = size(sys.H_drift, 1)
-        dim = 2ketdim^2
+        dim = traj.dims[unitary_name]
+        ketdim = Int(sqrt(dim ÷ 2))
 
         unitary_components = traj.components[unitary_name]
 
@@ -73,7 +80,6 @@ struct UnitaryExponentialIntegrator <: QuantumExponentialIntegrator
 end
 
 function (integrator::UnitaryExponentialIntegrator)(
-    sys::AbstractQuantumSystem,
     traj::NamedTrajectory;
     unitary_name::Union{Nothing, Symbol}=nothing,
     drive_name::Union{Nothing, Symbol, Tuple{Vararg{Symbol}}}=nothing,
@@ -83,11 +89,10 @@ function (integrator::UnitaryExponentialIntegrator)(
     @assert !isnothing(unitary_name) "unitary_name must be provided"
     @assert !isnothing(drive_name) "drive_name must be provided"
     return UnitaryExponentialIntegrator(
-        sys,
         unitary_name,
         drive_name,
+        G,
         traj;
-        G=G,
         autodiff=autodiff
     )
 end
@@ -122,14 +127,6 @@ end
     return Ũ⃗ₜ₊₁ - expv(Δtₜ, I(ℰ.ketdim) ⊗ Gₜ, Ũ⃗ₜ)
 end
 
-function hermitian_exp(G::AbstractMatrix)
-    Ĥ = Hermitian(Matrix(Isomorphisms.H(G)))
-    λ, V = eigen(Ĥ)
-    expG = Isomorphisms.iso(sparse(V * Diagonal(exp.(-im * λ)) * V'))
-    droptol!(expG, 1e-12)
-    return expG
-end
-
 @views function jacobian(
     ℰ::UnitaryExponentialIntegrator,
     zₜ::AbstractVector,
@@ -152,7 +149,7 @@ end
 
     Id = I(ℰ.ketdim)
 
-    expĜₜ = Id ⊗ hermitian_exp(Δtₜ * Gₜ)
+    expĜₜ = Id ⊗ exp_eigen(Δtₜ * Gₜ)
 
     ∂Ũ⃗ₜ₊₁ℰ = sparse(I, ℰ.dim, ℰ.dim)
     ∂Ũ⃗ₜℰ = -expĜₜ
@@ -183,15 +180,14 @@ struct QuantumStateExponentialIntegrator <: QuantumExponentialIntegrator
     G::Function
 
     function QuantumStateExponentialIntegrator(
-        sys::AbstractQuantumSystem,
         state_name::Symbol,
         drive_name::Union{Symbol, Tuple{Vararg{Symbol}}},
+        G::Function,
         traj::NamedTrajectory;
-        G::Function=a -> G_bilinear(a, sys.G_drift, sys.G_drives),
         autodiff::Bool=false
     )
-        ketdim = size(sys.H_drift, 1)
-        dim = 2ketdim
+        dim = traj.dims[state_name]
+        ketdim = dim ÷ 2
 
         state_components = traj.components[state_name]
 
@@ -237,7 +233,6 @@ function get_comps(P::QuantumStateExponentialIntegrator, traj::NamedTrajectory)
 end
 
 function (integrator::QuantumStateExponentialIntegrator)(
-    sys::AbstractQuantumSystem,
     traj::NamedTrajectory;
     state_name::Union{Nothing, Symbol}=nothing,
     drive_name::Union{Nothing, Symbol, Tuple{Vararg{Symbol}}}=nothing,
@@ -247,11 +242,10 @@ function (integrator::QuantumStateExponentialIntegrator)(
     @assert !isnothing(state_name) "state_name must be provided"
     @assert !isnothing(drive_name) "drive_name must be provided"
     return QuantumStateExponentialIntegrator(
-        sys,
         state_name,
         drive_name,
+        G,
         traj;
-        G=G,
         autodiff=autodiff
     )
 end
@@ -299,7 +293,7 @@ end
     # compute the generator
     Gₜ = ℰ.G(aₜ)
 
-    expGₜ = hermitian_exp(Δtₜ * Gₜ)
+    expGₜ = exp_eigen(Δtₜ * Gₜ)
 
     ∂ψ̃ₜ₊₁ℰ = sparse(I, ℰ.dim, ℰ.dim)
     ∂ψ̃ₜℰ = -expGₜ
@@ -338,6 +332,7 @@ end
 
     dt = 0.1
 
+
     Z = NamedTrajectory(
         (
             Ũ⃗ = unitary_geodesic(U_goal, T),
@@ -350,7 +345,9 @@ end
         goal=(Ũ⃗ = Ũ⃗_goal,)
     )
 
-    ℰ = UnitaryExponentialIntegrator(system, :Ũ⃗, :a, Z)
+    G = a -> Integrators.G_bilinear(a, system.G_drift, system.G_drives)
+
+    ℰ = UnitaryExponentialIntegrator(:Ũ⃗, :a, G, Z)
 
 
     ∂Ũ⃗ₜℰ, ∂Ũ⃗ₜ₊₁ℰ, ∂aₜℰ, ∂Δtₜℰ = jacobian(ℰ, Z[1].data, Z[2].data, 1)
@@ -397,7 +394,10 @@ end
         goal=(ψ̃ = ψ̃_goal,)
     )
 
-    ℰ = QuantumStateExponentialIntegrator(system, :ψ̃, :a, Z)
+    G = a -> Integrators.G_bilinear(a, system.G_drift, system.G_drives)
+
+    ℰ = QuantumStateExponentialIntegrator(:ψ̃, :a, G, Z)
+
     ∂ψ̃ₜℰ, ∂ψ̃ₜ₊₁ℰ, ∂aₜℰ, ∂Δtₜℰ = jacobian(ℰ, Z[1].data, Z[2].data, 1)
 
     ∂ℰ_forwarddiff = ForwardDiff.jacobian(

src/integrators/pade_integrators.jl

@@ -222,59 +222,50 @@ struct UnitaryPadeIntegrator <: QuantumPadeIntegrator
     G::Function
     ∂G::Function
 
-    """
+    @doc raw"""
         UnitaryPadeIntegrator(
-            sys::AbstractQuantumSystem,
             unitary_name::Symbol,
-            drive_name::Union{Symbol,Tuple{Vararg{Symbol}}};
+            drive_name::Union{Symbol,Tuple{Vararg{Symbol}}},
+            G::Function,
+            ∂G::Function,
+            traj::NamedTrajectory;
             order::Int=4,
-            autodiff::Bool=order != 4
-        ) where R <: Real
-
-    Construct a `UnitaryPadeIntegrator` for the quantum system `sys`.
+            calculate_pade_operators_structure::Bool=true,
+            autodiff::Bool=false
+        )
 
-    # Examples
+    Construct a `UnitaryPadeIntegrator` which computes
 
-    ## a bare integrator
-    ```julia
-        P = UnitaryPadeIntegrator(sys)
+    ```math
+    \text{isovec}\qty(B^{(n)}(a_t) U_{t+1} - F^{(n)}(a_t) U_t)
     ```
 
-    ## for a single drive `a`:
-    ```julia
-        P = UnitaryPadeIntegrator(sys, :Ũ⃗, :a)
-    ```
-
-    ## for two drives `α` and `γ`, order `4`, and autodiffed:
-    ```julia
-        P = UnitaryPadeIntegrator(sys, :Ũ⃗, (:α, :γ); order=4, autodiff=true)
-    ```
+    where `U_t` is the unitary at time `t`, `a_t` is the control at time `t`, and `B^{(n)}(a_t)` and `F^{(n)}(a_t)` are the `n`th order Pade operators of the exponential of the drift operator `G(a_t)`.
 
     # Arguments
-    - `sys::AbstractQuantumSystem`: the quantum system
-    - `unitary_name::Union{Symbol,Nothing}=nothing`: the nameol for the unitary
-    - `drive_name::Union{Symbol,Tuple{Vararg{Symbol}},Nothing}=nothing`: the nameol(s) for the drives
-    - `order::Int=4`: the order of the Pade approximation. Must be in `[4, 6, 8, 10]`. If order is not `4` and `autodiff` is `false`, then the integrator will use the hand-coded fourth order derivatives.
-    - `autodiff::Bool=order != 4`: whether to use automatic differentiation to compute the jacobian and hessian of the lagrangian
+    - `unitary_name::Symbol`: the name of the unitary in the trajectory
+    - `drive_name::Union{Symbol,Tuple{Vararg{Symbol}}}`: the name of the drive(s) in the trajectory
+    - `G::Function`: a function which takes the control vector `a_t` and returns the drive `G(a_t)`, $G(a_t) = \text{iso}(-i H(a_t))$
+    - `∂G::Function`: a function which takes the control vector `a_t` and returns a vector of matrices $\qty(\ldots, \pdv{G}{a^j_t}, \ldots)$
+    - `traj::NamedTrajectory`: the trajectory
+
+    # Keyword Arguments
+    - `order::Int=4`: the order of the Pade approximation. Must be in `[4, 6, 8, 10, 12, 14, 16, 18, 20]`.
 
     """
     function UnitaryPadeIntegrator(
-        sys::AbstractQuantumSystem,
         unitary_name::Symbol,
         drive_name::Union{Symbol,Tuple{Vararg{Symbol}}},
+        G::Function,
+        ∂G::Function,
         traj::NamedTrajectory;
         order::Int=4,
-        G::Function=a -> G_bilinear(a, sys.G_drift, sys.G_drives),
-        ∂G::Function=a -> sys.G_drives,
-        calculate_pade_operators_structure::Bool=true,
         autodiff::Bool=false
     )
         @assert order ∈ keys(PADE_COEFFICIENTS) "order ∉ $(keys(PADE_COEFFICIENTS))"
 
-        ketdim = size(sys.H_drift, 1)
-        dim = 2ketdim^2
-
-        I_2N = sparse(I(2ketdim))
+        dim = traj.dims[unitary_name]
+        ketdim = Int(sqrt(dim ÷ 2))
 
         unitary_components = traj.components[unitary_name]
 
@@ -322,7 +313,6 @@ function get_comps(P::UnitaryPadeIntegrator, traj::NamedTrajectory)
 end
 
 function (integrator::UnitaryPadeIntegrator)(
-    sys::AbstractQuantumSystem,
     traj::NamedTrajectory;
     unitary_name::Union{Symbol, Nothing}=nothing,
     drive_name::Union{Symbol, Tuple{Vararg{Symbol}}, Nothing}=nothing,
@@ -334,13 +324,12 @@ function (integrator::UnitaryPadeIntegrator)(
     @assert !isnothing(unitary_name) "unitary_name must be provided"
     @assert !isnothing(drive_name) "drive_name must be provided"
     return UnitaryPadeIntegrator(
-        sys,
         unitary_name,
         drive_name,
+        G,
+        ∂G,
         traj;
         order=order,
-        G=G,
-        ∂G=∂G,
         autodiff=autodiff
     )
 end
@@ -522,13 +511,12 @@ struct QuantumStatePadeIntegrator <: QuantumPadeIntegrator
     - `autodiff::Bool=false`: whether to use automatic differentiation to compute the jacobian and hessian of the lagrangian
     """
     function QuantumStatePadeIntegrator(
-        sys::AbstractQuantumSystem,
         state_name::Symbol,
         drive_name::Union{Symbol,Tuple{Vararg{Symbol}}},
+        G::Function,
+        ∂G::Function,
         traj::NamedTrajectory;
         order::Int=4,
-        G::Function=a -> G_bilinear(a, sys.G_drift, sys.G_drives),
-        ∂G::Function=a -> sys.G_drives,
         autodiff::Bool=false,
     )
         @assert order ∈ keys(PADE_COEFFICIENTS) "order ∉ $(keys(PADE_COEFFICIENTS))"
@@ -583,7 +571,6 @@ function get_comps(P::QuantumStatePadeIntegrator, traj::NamedTrajectory)
 end
 
 function (integrator::QuantumStatePadeIntegrator)(
-    sys::AbstractQuantumSystem,
     traj::NamedTrajectory;
     state_name::Union{Symbol, Nothing}=nothing,
     drive_name::Union{Symbol, Tuple{Vararg{Symbol}}, Nothing}=nothing,
@@ -595,13 +582,12 @@ function (integrator::QuantumStatePadeIntegrator)(
     @assert !isnothing(state_name) "state_name must be provided"
     @assert !isnothing(drive_name) "drive_name must be provided"
     return QuantumStatePadeIntegrator(
-        sys,
         state_name,
         drive_name,
+        G,
+        ∂G,
         traj;
         order=order,
-        G=G,
-        ∂G=∂G,
         autodiff=autodiff
     )
 end