Merge pull request #160 from kestrelquantum/problem-template-options-…

…refactor

Problem template options refactor

Project.toml

@@ -1,7 +1,7 @@
 name = "QuantumCollocation"
 uuid = "0dc23a59-5ffb-49af-b6bd-932a8ae77adf"
 authors = ["Aaron Trowbridge <[email protected]> and contributors"]
-version = "0.2.2"
+version = "0.3.0"
 
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"

docs/src/lib.md

@@ -1,11 +1,21 @@
 #  Library
 
-## QuantumUtils
+## Problem Templates
 ```@autodocs
-Modules = [QuantumCollocation.QuantumUtils]
+Modules = [QuantumCollocation.ProblemTemplates]
 ```
 
-## QuantumSystems 
+## Direct Sums
+```@autodocs
+Modules = [QuantumCollocation.DirectSums]
+```
+
+## Quantum Object Utils
+```@autodocs
+Modules = [QuantumCollocation.QuantumObjectUtils]
+```
+
+## Quantum Systems 
 ```@autodocs
 Modules = [QuantumCollocation.QuantumSystems]
 ```
@@ -14,3 +24,63 @@ Modules = [QuantumCollocation.QuantumSystems]
 ```@autodocs
 Modules = [QuantumCollocation.Integrators]
 ``` 
+
+## Objectives
+```@autodocs
+Modules = [QuantumCollocation.Objectives]
+```
+
+## Losses
+```@autodocs
+Modules = [QuantumCollocation.Losses]
+```
+
+## Embedded Operators
+```@autodocs
+Modules = [QuantumCollocation.EmbeddedOperators]
+```
+
+## Isomorphisms
+```@autodocs
+Modules = [QuantumCollocation.Isomorphisms]
+```
+
+## Options
+```@autodocs
+Modules = [QuantumCollocation.Options]
+```
+
+## Plotting 
+```@autodocs
+Modules = [QuantumCollocation.Plotting]
+```
+
+## Problem Solvers
+```@autodocs
+Modules = [QuantumCollocation.ProblemSolvers]
+```
+
+## Rollouts
+```@autodocs
+Modules = [QuantumCollocation.Rollouts]
+```
+
+## Saving and Loading
+```@autodocs
+Modules = [QuantumCollocation.SaveLoadUtils]
+```
+
+## Structure Utils
+```@autodocs
+Modules = [QuantumCollocation.StructureUtils]
+```
+
+## Trajectory Initialization
+```@autodocs
+Modules = [QuantumCollocation.TrajectoryInitialization]
+```
+
+## Trajectory Interpolations
+```@autodocs
+Modules = [QuantumCollocation.TrajectoryInterpolations]
+```

src/QuantumCollocation.jl

@@ -3,6 +3,9 @@ module QuantumCollocation
 using Reexport
 
 
+include("options.jl")
+@reexport using .Options
+
 include("isomorphisms.jl")
 @reexport using .Isomorphisms
 
@@ -42,9 +45,6 @@ include("dynamics.jl")
 include("evaluators.jl")
 @reexport using .Evaluators
 
-include("options.jl")
-@reexport using .Options
-
 include("problems.jl")
 @reexport using .Problems
 

src/constraints/_constraints.jl

@@ -3,6 +3,7 @@ module Constraints
 export AbstractConstraint
 
 export LinearConstraint
+
 export constrain!
 
 export NonlinearConstraint
@@ -12,6 +13,7 @@ export NonlinearInequalityConstraint
 using ..Losses
 using ..Isomorphisms
 using ..StructureUtils
+using ..Options
 
 using TrajectoryIndexingUtils
 using NamedTrajectories
@@ -126,4 +128,6 @@ struct NonlinearInequalityConstraint <: NonlinearConstraint
     params::Dict{Symbol, Any}
 end
 
+
+
 end

src/constraints/complex_modulus_constraint.jl

@@ -20,13 +20,13 @@ function ComplexModulusContraint(;
     R::Union{Float64, Nothing}=nothing,
     comps::Union{AbstractVector{Int}, Nothing}=nothing,
     times::Union{AbstractVector{Int}, Nothing}=nothing,
-    dim::Union{Int, Nothing}=nothing,
+    zdim::Union{Int, Nothing}=nothing,
     T::Union{Int, Nothing}=nothing,
 )
     @assert !isnothing(R) "must provide a value R, s.t. |z| <= R"
     @assert !isnothing(comps) "must provide components of the complex number"
     @assert !isnothing(times) "must provide times"
-    @assert !isnothing(dim) "must provide a trajectory dimension"
+    @assert !isnothing(zdim) "must provide a trajectory dimension"
     @assert !isnothing(T) "must provide a T"
 
     @assert length(comps) == 2 "component must represent a complex number and have dimension 2"
@@ -37,7 +37,7 @@ function ComplexModulusContraint(;
     params[:R] = R
     params[:comps] = comps
     params[:times] = times
-    params[:dim] = dim
+    params[:zdim] = zdim
     params[:T] = T
 
     gₜ(xₜ, yₜ) = [R^2 - xₜ^2 - yₜ^2]
@@ -47,7 +47,7 @@ function ComplexModulusContraint(;
     @views function g(Z⃗)
         r = zeros(length(times))
         for (i, t) ∈ enumerate(times)
-            zₜ = Z⃗[slice(t, comps, dim)]
+            zₜ = Z⃗[slice(t, comps, zdim)]
             xₜ = zₜ[1]
             yₜ = zₜ[2]
             r[i] = gₜ(xₜ, yₜ)[1]
@@ -58,17 +58,17 @@ function ComplexModulusContraint(;
     ∂g_structure = []
 
     for (i, t) ∈ enumerate(times)
-        push!(∂g_structure, (i, index(t, comps[1], dim)))
-        push!(∂g_structure, (i, index(t, comps[2], dim)))
+        push!(∂g_structure, (i, index(t, comps[1], zdim)))
+        push!(∂g_structure, (i, index(t, comps[2], zdim)))
     end
 
     @views function ∂g(Z⃗; ipopt=true)
         ∂ = spzeros(length(times), length(Z⃗))
         for (i, t) ∈ enumerate(times)
-            zₜ = Z⃗[slice(t, comps, dim)]
+            zₜ = Z⃗[slice(t, comps, zdim)]
             xₜ = zₜ[1]
             yₜ = zₜ[2]
-            ∂[i, slice(t, comps, dim)] = ∂gₜ(xₜ, yₜ)
+            ∂[i, slice(t, comps, zdim)] = ∂gₜ(xₜ, yₜ)
         end
         if ipopt
             return [∂[i, j] for (i, j) in ∂g_structure]
@@ -83,15 +83,15 @@ function ComplexModulusContraint(;
         push!(
             μ∂²g_structure,
             (
-                index(t, comps[1], dim),
-                index(t, comps[1], dim)
+                index(t, comps[1], zdim),
+                index(t, comps[1], zdim)
             )
         )
         push!(
             μ∂²g_structure,
             (
-                index(t, comps[2], dim),
-                index(t, comps[2], dim)
+                index(t, comps[2], zdim),
+                index(t, comps[2], zdim)
             )
         )
     end
@@ -100,7 +100,7 @@ function ComplexModulusContraint(;
         n = length(Z⃗)
         μ∂² = spzeros(n, n)
         for (i, t) ∈ enumerate(times)
-            t_slice = slice(t, comps, dim)
+            t_slice = slice(t, comps, zdim)
             μ∂²[t_slice, t_slice] = μₜ∂²gₜ(μ[i])
         end
         if ipopt
@@ -143,4 +143,4 @@ function ComplexModulusContraint(
         zdim=traj.dim,
         T=traj.T
     )
-end
+end

src/constraints/fidelity_constraint.jl

@@ -161,7 +161,6 @@ function FinalUnitaryFidelityConstraint(
     subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing,
     eval_hessian::Bool=true
 )
-    @assert statesymb ∈ traj.names
     return FinalFidelityConstraint(;
         fidelity_function=iso_vec_unitary_fidelity,
         value=val,
@@ -221,7 +220,7 @@ function FinalUnitaryFreePhaseFidelityConstraint(;
     state_slice::Union{AbstractVector{Int},Nothing}=nothing,
     phase_slice::Union{AbstractVector{Int},Nothing}=nothing,
     goal::Union{AbstractVector{Float64},Nothing}=nothing,
-    phase_operators::Union{AbstractVector{<:AbstractMatrix{<:Complex{Float64}}},Nothing}=nothing,
+    phase_operators::Union{AbstractVector{<:AbstractMatrix{<:Complex}},Nothing}=nothing,
     zdim::Union{Int,Nothing}=nothing,
     subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing,
     eval_hessian::Bool=false
@@ -284,4 +283,25 @@ function FinalUnitaryFreePhaseFidelityConstraint(;
         1,
         params
     )
+end
+
+function FinalUnitaryFreePhaseFidelityConstraint(
+    state_symbol::Symbol,
+    global_symbol::Symbol,
+    phase_operators::AbstractVector{<:AbstractMatrix{<:Complex}},
+    val::Float64,
+    traj::NamedTrajectory;
+    subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing,
+    eval_hessian::Bool=false
+)
+    return FinalUnitaryFreePhaseFidelityConstraint(;
+        value=val,
+        state_slice=slice(traj.T, traj.components[state_symbol], traj.dim),
+        phase_slice=trajectory.global_components[global_symbol],
+        goal=traj.goal[state_symbol],
+        phase_operators=phase_operators,
+        zdim=length(traj),
+        subspace=subspace,
+        eval_hessian=eval_hessian
+    )
 end

src/direct_sums.jl

@@ -83,21 +83,21 @@ to the `reduce` function.
 - `traj1::NamedTrajectory`: The first `NamedTrajectory` object.
 - `traj2::NamedTrajectory`: The second `NamedTrajectory` object.
 - `free_time::Bool=false`: Whether to construct a free time problem.
-- `timestep_symbol::Symbol=:Δt`: The timestep symbol to use for free time problems.
+- `timestep_name::Symbol=:Δt`: The timestep symbol to use for free time problems.
 """
 function direct_sum(
     traj1::NamedTrajectory,
     traj2::NamedTrajectory;
     free_time::Bool=false,
-    timestep_symbol::Symbol=:Δt,
+    timestep_name::Symbol=:Δt,
 )
-    return direct_sum([traj1, traj2]; free_time=free_time, timestep_symbol=timestep_symbol)
+    return direct_sum([traj1, traj2]; free_time=free_time, timestep_name=timestep_name)
 end
 
 function direct_sum(
     trajs::AbstractVector{<:NamedTrajectory};
     free_time::Bool=false,
-    timestep_symbol::Symbol=:Δt,
+    timestep_name::Symbol=:Δt,
 )
     if length(trajs) < 2
         throw(ArgumentError("At least two trajectories must be provided"))
@@ -122,13 +122,13 @@ function direct_sum(
 
     # add timestep to components
     if free_time
-        components = merge_outer(components, NamedTuple{(timestep_symbol,)}([get_timesteps(trajs[1])]))
+        components = merge_outer(components, NamedTuple{(timestep_name,)}([get_timesteps(trajs[1])]))
     end
 
     return NamedTrajectory(
         components,
         controls=merge_outer([traj.control_names for traj in trajs]),
-        timestep=free_time ? timestep_symbol : timestep,
+        timestep=free_time ? timestep_name : timestep,
         bounds=merge_outer([traj.bounds for traj in trajs]),
         initial=merge_outer([traj.initial for traj in trajs]),
         final=merge_outer([traj.final for traj in trajs]),
@@ -670,21 +670,21 @@ end
     pi_false_ops = PiccoloOptions(verbose=false, free_time=false)
     pi_true_ops = PiccoloOptions(verbose=false, free_time=true)
     suffix = "_new"
-    timestep_symbol = :Δt
+    timestep_name = :Δt
 
     prob1 = UnitarySmoothPulseProblem(sys, GATES[:Y], T, Δt, piccolo_options=pi_false_ops, ipopt_options=ops)
     traj1 = direct_sum(prob1.trajectory, add_suffix(prob1.trajectory, suffix), free_time=true)
 
     # Direct sum (shared timestep name)
-    @test get_suffix(traj1, suffix).timestep == timestep_symbol
-    @test get_suffix(traj1, suffix, remove=true).timestep == timestep_symbol
+    @test get_suffix(traj1, suffix).timestep == timestep_name
+    @test get_suffix(traj1, suffix, remove=true).timestep == timestep_name
 
     prob2 = UnitarySmoothPulseProblem(sys, GATES[:Y], T, Δt, ipopt_options=ops, piccolo_options=pi_true_ops)
     traj2 = add_suffix(prob2.trajectory, suffix)
 
     # Trajectory (unique timestep name)
-    @test get_suffix(traj2, suffix).timestep == add_suffix(timestep_symbol, suffix)
-    @test get_suffix(traj2, suffix, remove=true).timestep == timestep_symbol
+    @test get_suffix(traj2, suffix).timestep == add_suffix(timestep_name, suffix)
+    @test get_suffix(traj2, suffix, remove=true).timestep == timestep_name
 end
 
 end # module