Fixed Adjoint

src/quantum_systems.jl

@@ -297,67 +297,191 @@ end
 
 # ****************************************************************************** #
 
-
-
 struct ParameterizedQuantumSystem <: AbstractQuantumSystem
     H::Function 
     G::Function
-    Gₐ::Function 
+    Gₐ::Vector{Function} 
     ∂G::Function 
     n_drives::Int 
     levels::Int 
     params::Dict{Symbol, Any}
 
     function ParameterizedQuantumSystem end 
 
-    function ParameterizedQuantumSystem(
-        H_drift::AbstractMatrix{<:Number},
-        H_drives::Vector{<:AbstractMatrix{<:Number}},
-        Hₐ_drives::Vector{<:AbstractMatrix{<:Number}};
-        params::Dict{Symbol, <:Any}=Dict{Symbol, Any}()
-    )
-        levels = size(H_drift, 1)
-        H_drift = sparse(H_drift)
-        G_drift = sparse(Isomorphisms.G(H_drift))
-
-        n_drives = length(H_drives)
-        H_drives = sparse.(H_drives)
-        G_drives = sparse.(Isomorphisms.G.(H_drives))
-
-        Gₐ_drives = sparse.(Isomorphisms.G.(Hₐ_drives))
-
-        if n_drives == 0
-            H = a -> H_drift
-            G = a -> G_drift
-            Gₐ = a -> 0
-            ∂G = a -> 0
-        else
-            H = a -> H_drift + sum(a .* H_drives)
-            G = a -> G_drift + sum(a .* G_drives)
-            Gₐ = a -> sum(a .* Gₐ_drives)
-            ∂G = a -> G_drives
-        end
-
+    function ParameterizedQuantumSystem(base_system::QuantumSystem, sensitivity_systems::Vector{QuantumSystem})
         return new(
-            H,
-            G,
-            Gₐ,
-            ∂G,
-            n_drives,
-            levels,
-            params
+            base_system.H,
+            base_system.G,
+            [sys.G for sys ∈ sensitivity_systems],
+            base_system.∂G,
+            base_system.n_drives,
+            base_system.levels,
+            base_system.params
         )
     end
 
-
-    function ParameterizedQuantumSystem(H_drives::Vector{<:AbstractMatrix{ℂ}},Hₐ_drives::Vector{<:AbstractMatrix{ℂ}}; kwargs...) where ℂ <: Number
-        @assert !isempty(H_drives) "At least one drive is required"
-        @assert length(Hₐ_drives) == length(H_drives) "Size has to be the same"
-        return ParameterizedQuantumSystem(spzeros(ℂ, size(H_drives[1])), H_drives, Hₐ_drives; kwargs...)
+    function ParameterizedQuantumSystem(system::ParameterizedQuantumSystem, indices::AbstractVector{Int64})
+        return new(
+            system.H,
+            system.G,
+            system.Gₐ[indices],
+            system.∂G,
+            system.n_drives,
+            system.levels,
+            system.params
+        )
     end
 
 end
 
+
+
+
+
+# struct ParameterizedQuantumSystem <: AbstractQuantumSystem
+#     H::Function 
+#     G::Function
+#     Gₐ::Function 
+#     ∂G::Function 
+#     n_drives::Int 
+#     levels::Int 
+#     params::Dict{Symbol, Any}
+
+#     function ParameterizedQuantumSystem end 
+
+#     function ParameterizedQuantumSystem(
+#         H_drift::AbstractMatrix{<:Number},
+#         H_drives::Vector{<:AbstractMatrix{<:Number}},
+#         Hₐ_drives::Vector{<:AbstractMatrix{<:Number}};
+#         params::Dict{Symbol, <:Any}=Dict{Symbol, Any}()
+#     )
+#         levels = size(H_drift, 1)
+#         H_drift = sparse(H_drift)
+#         G_drift = sparse(Isomorphisms.G(H_drift))
+
+#         n_drives = length(H_drives)
+#         H_drives = sparse.(H_drives)
+#         G_drives = sparse.(Isomorphisms.G.(H_drives))
+
+#         Gₐ_drives = sparse.(Isomorphisms.G.(Hₐ_drives))
+
+#         if n_drives == 0
+#             H = a -> H_drift
+#             G = a -> G_drift
+#             Gₐ = a -> 0
+#             ∂G = a -> 0
+#         else
+#             H = a -> H_drift + sum(a .* H_drives)
+#             G = a -> G_drift + sum(a .* G_drives)
+#             Gₐ = a -> sum(a .* Gₐ_drives)
+#             ∂G = a -> G_drives
+#         end
+
+#         return new(
+#             H,
+#             G,
+#             Gₐ,
+#             ∂G,
+#             n_drives,
+#             levels,
+#             params
+#         )
+#     end
+
+#     function ParameterizedQuantumSystem(
+#         H_drift::AbstractMatrix{<:Number},
+#         H_drives::Vector{<:AbstractMatrix{<:Number}},
+#         Hₐ_drift::AbstractMatrix{<:Number};
+#         params::Dict{Symbol, <:Any}=Dict{Symbol, Any}()
+#     )
+#         levels = size(H_drift, 1)
+#         H_drift = sparse(H_drift)
+#         G_drift = sparse(Isomorphisms.G(H_drift))
+
+#         n_drives = length(H_drives)
+#         H_drives = sparse.(H_drives)
+#         G_drives = sparse.(Isomorphisms.G.(H_drives))
+
+#         Gₐ_drift = sparse(Isomorphisms.G(Hₐ_drift))
+
+#         if n_drives == 0
+#             H = a -> H_drift
+#             G = a -> G_drift
+#             Gₐ = a -> Gₐ_drift
+#             ∂G = a -> 0
+#         else
+#             H = a -> H_drift + sum(a .* H_drives)
+#             G = a -> G_drift + sum(a .* G_drives)
+#             Gₐ = a -> Gₐ_drift
+#             ∂G = a -> G_drives
+#         end
+
+#         return new(
+#             H,
+#             G,
+#             Gₐ,
+#             ∂G,
+#             n_drives,
+#             levels,
+#             params
+#         )
+#     end
+
+#     function ParameterizedQuantumSystem(H_drives::Vector{<:AbstractMatrix{ℂ}},Hₐ_drives::Vector{<:AbstractMatrix{ℂ}}; kwargs...) where ℂ <: Number
+#         @assert !isempty(H_drives) "At least one drive is required"
+#         @assert length(Hₐ_drives) == length(H_drives) "Size has to be the same"
+#         return ParameterizedQuantumSystem(spzeros(ℂ, size(H_drives[1])), H_drives, Hₐ_drives; kwargs...)
+#     end
+
+#     function ParameterizedQuantumSystem(H_drift::AbstractMatrix{ℂ},Hₐ_drives::Vector{<:AbstractMatrix{ℂ}}; kwargs...) where ℂ <: Number
+#         @assert !isempty(H_drives) "At least one drive is required"
+#         @assert length(Hₐ_drives) == length(H_drives) "Size has to be the same"
+#         return ParameterizedQuantumSystem(H_drift, Matrix{ℂ}[], Hₐ_drives; kwargs...)
+#     end
+
+#     function ParameterizedQuantumSystem(H_drives::Vector{<:AbstractMatrix{ℂ}},Hₐ_drift::AbstractMatrix{ℂ}; kwargs...) where ℂ <: Number
+#         @assert !isempty(H_drives) "At least one drive is required"
+#         @assert length(Hₐ_drives) == length(H_drives) "Size has to be the same"
+#         return ParameterizedQuantumSystem(spzeros(ℂ, size(H_drives[1])), H_drives, Hₐ_drift; kwargs...)
+#     end
+
+#     function ParameterizedQuantumSystem(H_drift::AbstractMatrix{ℂ},Hₐ_drift::AbstractMatrix{ℂ}; kwargs...) where ℂ <: Number
+#         @assert !isempty(H_drives) "At least one drive is required"
+#         @assert length(Hₐ_drives) == length(H_drives) "Size has to be the same"
+#         return ParameterizedQuantumSystem(H_drift, Matrix{ℂ}[], Hₐ_drift; kwargs...)
+#     end
+
+#     function ParameterizedQuantumSystem(system::QuantumSystem,Hₐ_drift::AbstractMatrix{ℂ}; kwargs...) where ℂ <: Number
+#         return new(
+#             system.H,
+#             system.G,
+#             a -> sparse(Isomorphisms.G(Hₐ_drift)),
+#             system.∂G,
+#             system.n_drives,
+#             system.levels,
+#             system.params
+#         )
+
+#     end
+
+#     function ParameterizedQuantumSystem(system::QuantumSystem,Hₐ_drives::Vector{<:AbstractMatrix{ℂ}}; kwargs...) where ℂ <: Number
+
+#         Gₐ_drives = sparse.(Isomorphisms.G.(Hₐ_drives))
+#         return new(
+#             system.H,
+#             system.G,
+#             a -> sum(a .* Gₐ_drives),
+#             system.∂G,
+#             system.n_drives,
+#             system.levels,
+#             system.params
+#         )
+
+#     end
+
+# end
+
+
 #***********************************************************************************************#
 
 

src/rollouts.jl

@@ -639,31 +639,31 @@ end
 const ⊗ = kron
 
 function adjoint_unitary_rollout(trajectory::NamedTrajectory{Float64},system::ParameterizedQuantumSystem;
-    unitary_name = :Ũ⃗,drive_name=:a)
+    unitary_name::Symbol = :Ũ⃗, drive_name::Symbol=:a, indices::AbstractVector{Int64}=1:length(system.Gₐ))
 
     Ĩ⃗ = operator_to_iso_vec(Matrix{ComplexF64}(I(system.levels)))
 
     Ũ⃗ = zeros(trajectory[unitary_name].size)
-    Ũ⃗ₐ = zeros(trajectory[unitary_name].size)
-
-    Ũ⃗[:,1] = Ĩ⃗
-    for t = 2:trajectory.T
-        aₜ₋₁ = trajectory[drive_name][:, t - 1]
-
-
-        Ũₜ₋₁ = (Ũ⃗[:, t - 1])
-        Ũₐₜ₋₁ = (Ũ⃗ₐ[:, t - 1])
-
-        Ĝ = a_ ->  [I(system.levels)⊗(system.G(a_)) I(system.levels)⊗(system.Gₐ(a_)) ;  I(system.levels)⊗(system.G(a_)*0) I(system.levels)⊗(system.G(a_)) ]
-
-
-        out = expv(get_timesteps(trajectory)[t-1], Ĝ(aₜ₋₁), vcat((Ũₐₜ₋₁),(Ũₜ₋₁)))
-
-        Ũ⃗[:, t] = (out[end÷2+1:end])
-        Ũ⃗ₐ[:, t] = (out[1:end÷2])
+    Ũ⃗ₐ = [zeros(trajectory[unitary_name].size) for i ∈ indices]
+
+    for (idx,i) ∈ enumerate(indices)
+        Ũ⃗[:,1] = Ĩ⃗
+        for t = 2:trajectory.T
+            aₜ₋₁ = trajectory[drive_name][:, t - 1]
+
+
+            Ũₜ₋₁ = (Ũ⃗[:, t - 1])
+            Ũₐₜ₋₁ = (Ũ⃗ₐ[idx][:, t - 1])
+
+            Ĝ = a_ ->  [I(system.levels)⊗(system.G(a_)) I(system.levels)⊗(system.Gₐ[i](a_)) ;  I(system.levels)⊗(system.G(a_)*0) I(system.levels)⊗(system.G(a_)) ]
+
+            out = expv(get_timesteps(trajectory)[t-1], Ĝ(aₜ₋₁), vcat((Ũₐₜ₋₁),(Ũₜ₋₁)))
+
+            Ũ⃗[:, t] = (out[end÷2+1:end])
+            Ũ⃗ₐ[idx][:, t] = (out[1:end÷2])
+        end
     end
     return Ũ⃗,Ũ⃗ₐ
-
 end
 
 # ----------------------------------------------------------------------------- #