refactor: simplified integrator jacobians

Project.toml

@@ -1,7 +1,7 @@
 name = "QuantumCollocationCore"
 uuid = "2b384925-53cb-4042-a8d2-6faa627467e1"
 authors = ["Aaron Trowbridge <[email protected]> and contributors"]
-version = "0.1.3"
+version = "0.2.0"
 
 [deps]
 Einsum = "b7d42ee7-0b51-5a75-98ca-779d3107e4c0"
@@ -17,7 +17,6 @@ NamedTrajectories = "538bc3a1-5ab9-4fc3-b776-35ca1e893e08"
 PiccoloQuantumObjects = "5a402ddf-f93c-42eb-975e-5582dcda653d"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
-Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
 TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
 TrajectoryIndexingUtils = "6dad8b7f-dd9a-4c28-9b70-85b9a079bfc8"
@@ -36,7 +35,6 @@ NamedTrajectories = "0.2"
 PiccoloQuantumObjects = "0.1"
 Reexport = "1.2"
 SparseArrays = "1.10, 1.11"
-Symbolics = "6.25"
 TestItemRunner = "1.0"
 TestItems = "1.0"
 TrajectoryIndexingUtils = "0.1"

src/QuantumCollocationCore.jl

@@ -5,9 +5,6 @@ using Reexport
 include("options.jl")
 @reexport using .Options
 
-include("structure_utils.jl")
-@reexport using .StructureUtils
-
 include("losses/_losses.jl")
 @reexport using .Losses
 

src/constraints/_constraints.jl

@@ -11,7 +11,6 @@ export NonlinearEqualityConstraint
 export NonlinearInequalityConstraint
 
 using ..Losses
-using ..StructureUtils
 using ..Options
 
 using TrajectoryIndexingUtils
@@ -59,7 +58,7 @@ function constrain!(
 )
     for con in cons
         if verbose
-            println("applying constraint: ", con.label)
+            println("        applying constraint: ", con.label)
         end
         con(opt, vars, traj)
     end

src/constraints/fidelity_constraint.jl

@@ -88,7 +88,7 @@ function FinalFidelityConstraint(;
 
     ∂ℱ(x) = ForwardDiff.jacobian(ℱ, x)
 
-    ∂ℱ_structure = jacobian_structure(∂ℱ, statedim)
+    ∂ℱ_structure = Iterators.product(1, 1:statedim)
 
     col_offset = index(T, comps[1] - 1, zdim)
 

src/dynamics.jl

@@ -8,7 +8,6 @@ export dynamics_jacobian
 export dynamics_hessian_of_lagrangian
 export dynamics_components
 
-using ..StructureUtils
 using ..Integrators
 
 using TrajectoryIndexingUtils
@@ -27,7 +26,7 @@ struct QuantumDynamics <: AbstractDynamics
     integrators::Union{Nothing, Vector{<:AbstractIntegrator}}
     F::Function
     ∂F::Function
-    ∂F_structure::Vector{Tuple{Int, Int}}
+    ∂F_structure::Function
     μ∂²F::Union{Function, Nothing}
     μ∂²F_structure::Union{Vector{Tuple{Int, Int}}, Nothing}
     dim::Int
@@ -45,8 +44,7 @@ function dynamics_components(integrators::Vector{<:AbstractIntegrator})
 end
 
 function dynamics(
-    integrators::Vector{<:AbstractIntegrator},
-    traj::NamedTrajectory
+    integrators::Vector{<:AbstractIntegrator}
 )
     dynamics_comps = dynamics_components(integrators)
     dynamics_dim = sum(integrator.dim for integrator ∈ integrators)
@@ -60,21 +58,19 @@ function dynamics(
     return f
 end
 
-
-
 function dynamics_jacobian(
-    integrators::Vector{<:AbstractIntegrator},
-    traj::NamedTrajectory
+    integrators::Vector{<:AbstractIntegrator}
 )
     dynamics_comps = dynamics_components(integrators)
     dynamics_dim = sum(integrator.dim for integrator ∈ integrators)
+    zdim = integrators[1].zdim
     @views function ∂f(zₜ, zₜ₊₁, t)
-        ∂ = zeros(eltype(zₜ), dynamics_dim, 2traj.dim)
+        ∂ = spzeros(eltype(zₜ), dynamics_dim, 2zdim)
         for (integrator, integrator_comps) ∈ zip(integrators, dynamics_comps)
-            ∂[integrator_comps, 1:traj.dim] = 
+            ∂[integrator_comps, 1:2zdim] =
                 Integrators.jacobian(integrator, zₜ, zₜ₊₁, t)
         end
-        return sparse(∂)
+        return ∂
     end
     return ∂f
 end
@@ -152,73 +148,54 @@ function QuantumDynamics(
     integrators::Vector{<:AbstractIntegrator},
     traj::NamedTrajectory;
     eval_hessian=true,
-    jacobian_structure=true,
     verbose=false
 )
     if verbose
         println("        constructing knot point dynamics functions...")
     end
 
-    free_time = traj.timestep isa Symbol
+    f = dynamics(integrators)
 
-    # if free_time
-    #     @assert all([
-    #         !isnothing(state(integrator)) &&
-    #         !isnothing(controls(integrator))
-    #             for integrator ∈ integrators
-    #     ])
-    # else
-    #     @assert all([
-    #         !isnothing(state(integrator)) &&
-    #         !isnothing(controls(integrator))
-    #             for integrator ∈ integrators
-    #     ])
-    # end
-
-    # for integrator ∈ integrators
-    #     if integrator isa QuantumIntegrator && controls(integrator) isa Tuple
-    #         drive_comps = [traj.components[s] for s ∈ integrator.drive_symb]
-    #         number_of_drives = sum(length.(drive_comps))
-    #         @assert number_of_drives == integrator.n_drives "number of drives ($(number_of_drives)) does not match number of drive terms in Hamiltonian ($(integrator.n_drives))"
-    #     end
-    # end
-
-    f = dynamics(integrators, traj)
-
-    ∂f = dynamics_jacobian(integrators, traj)
+    ∂f = dynamics_jacobian(integrators)
 
     if eval_hessian
         μ∂²f = dynamics_hessian_of_lagrangian(integrators, traj)
     else
         μ∂²f = nothing
     end
 
-    if verbose
-        println("        determining dynamics derivative structure...")
-    end
-
     dynamics_dim = sum(integrator.dim for integrator ∈ integrators)
 
     if eval_hessian
-        ∂f_structure, ∂F_structure, μ∂²f_structure, μ∂²F_structure =
-            dynamics_structure(∂f, μ∂²f, traj, dynamics_dim;
-                verbose=verbose,
-                jacobian=jacobian_structure,
-                hessian=!any(
-                    integrator.autodiff for integrator ∈ integrators if integrator isa QuantumIntegrator
-                )
-            )
-        μ∂²f_nnz = length(μ∂²f_structure)
+        @error "hessians not implemented"
+        # ∂f_structure, ∂F_structure, μ∂²f_structure, μ∂²F_structure =
+        #     dynamics_structure(∂f, μ∂²f, traj, dynamics_dim;
+        #         verbose=verbose,
+        #         jacobian=jacobian_structure,
+        #         hessian=!any(
+        #             integrator.autodiff for integrator ∈ integrators if integrator isa QuantumIntegrator
+        #         )
+        #     )
+        # μ∂²f_nnz = length(μ∂²f_structure)
     else
-        ∂f_structure, ∂F_structure = dynamics_structure(∂f, traj, dynamics_dim;
-            verbose=verbose,
-            jacobian=jacobian_structure,
-        )
+        function ∂F_structure()
+            ∂f_structure = [(i, j) for i = 1:dynamics_dim, j = 1:2traj.dim]
+
+            ∂F_ = Vector{Tuple{Int,Int}}(undef, length(∂f_structure) * (traj.T - 1))
+            for t = 1:traj.T-1
+                ∂fₜ_structure = [
+                    (
+                        i + index(t, 0, dynamics_dim),
+                        j + index(t, 0, traj.dim)
+                    ) for (i, j) ∈ ∂f_structure
+                ]
+                ∂F_[slice(t, length(∂f_structure))] = ∂fₜ_structure
+            end
+            return ∂F_
+        end
         μ∂²F_structure = nothing
     end
 
-    ∂f_nnz = length(∂f_structure)
-
     if verbose
         println("        constructing full dynamics derivative functions...")
     end
@@ -234,12 +211,13 @@ function QuantumDynamics(
     end
 
     @views function ∂F(Z⃗::AbstractVector{R}) where R <: Real
-        ∂s = zeros(R, length(∂F_structure))
+        ∂f_nnz = dynamics_dim * 2traj.dim
+        ∂s = zeros(R, ∂f_nnz * (traj.T - 1))
         Threads.@threads for t = 1:traj.T-1
             zₜ = Z⃗[slice(t, traj.dim)]
             zₜ₊₁ = Z⃗[slice(t + 1, traj.dim)]
             ∂fₜ = ∂f(zₜ, zₜ₊₁, t)
-            for (k, (i, j)) ∈ enumerate(∂f_structure)
+            for (k, (i, j)) ∈ enumerate(Iterators.product(1:dynamics_dim, 1:2traj.dim))
                 ∂s[index(t, k, ∂f_nnz)] = ∂fₜ[i, j]
             end
         end

src/evaluators.jl

@@ -90,11 +90,11 @@ end
 end
 
 function MOI.jacobian_structure(evaluator::PicoEvaluator)
-    dynamics_structure = evaluator.dynamics.∂F_structure
+    dynamics_structure = evaluator.dynamics.∂F_structure()
     row_offset = evaluator.n_dynamics_constraints
     nl_constraint_structure = []
     for con ∈ evaluator.nonlinear_constraints
-        con_structure = [(i + row_offset, j) for (i, j) in con.∂g_structure]
+        con_structure = [(i + row_offset, j) for (i, j) ∈ con.∂g_structure]
         push!(nl_constraint_structure, con_structure)
         row_offset += con.dim
     end

src/integrators/exponential_integrators.jl

@@ -138,8 +138,7 @@ end
     # obtain the timestep
     if ℰ.freetime
         Δtₜ = zₜ[ℰ.timestep]
-
-   else
+    else
         Δtₜ = ℰ.timestep
     end
 
@@ -498,9 +497,9 @@ end
     ∂ℰ = jacobian(ℰ, Z[1].data, Z[2].data, 1)
 
     ∂Ũ⃗ₜℰ = ∂ℰ[:, ℰ.state_components]
-    ∂Ũ⃗ₜ₊₁ℰ = ∂ℰ[:, Z.dim .+ ℰ.state_components] 
+    ∂Ũ⃗ₜ₊₁ℰ = ∂ℰ[:, Z.dim .+ ℰ.state_components]
     ∂aₜℰ = ∂ℰ[:, ℰ.drive_components]
-    ∂Δtₜℰ = ∂ℰ[:, Z.components.Δt]  
+    ∂Δtₜℰ = ∂ℰ[:, Z.components.Δt]
 
     ∂ℰ_finitediff= FiniteDiff.finite_difference_jacobian(
         zz -> ℰ(zz[1:Z.dim], zz[Z.dim+1:end], 1),