From cafef2793ad09358e5828e86a5f82c9c3e8810e9 Mon Sep 17 00:00:00 2001
From: Aaron Trowbridge <47730232+aarontrowbridge@users.noreply.github.com>
Date: Wed, 15 Jan 2025 11:56:14 -0500
Subject: [PATCH] core peeloff (#173)

* moving otimes to isomorphisms

* manifest update

* integrator signature refactor -- removes sys dependence

* problem template fixes (tests passing)

* moved isomorphisms to core

* big rip

* big rip 2: quantum objects

* traj init for unitary

* state sampling problem template

* refactored system to integrator interface

* density operators

* traj init for unitary

* state sampling problem template

* bug fix: a_guess option

* bug fix: free phase obj

* phase rollout fidelity

* free phase initialization

* Tests for free phases

* free phase incr

* bug fix: mintime and robprob fidelity subspace

* bug fix: drop drive_sigma arg, fix subspace nothing

* problem template fixes (tests passing)

* big rip

* refactored system to integrator interface

* density operators

* rebase fixes (broken)

* rebase complete (fixes free phase loss)

* docs uopdates

* renaming OperatorType -> AbstractPiccoloOperator

* adding quantum system templates + cleaning

* update system type from AbstractQuantumSystem to OpenQuantumSystem and clean up function parameters

* refactor trajectory initialization to use density_to_iso_vec for initial and goal states

* refactor rollout_fidelity to have sys as arg

* passing tests (except for direct sum prob)

* refactor add_suffix and remove_suffix functions to eliminate sys parameter and streamline integrator modifications

* bump version to 0.5.0 and update dependencies for compatibility

* remove operatortype

* add back in zero initial and final

* revert removing list of state names

* refactor callbacks and problem templates to use rollout fidelity and system return values

* update Project.toml dependencies and refactor drive length references in quantum state sampling problems

* fix callbacks except for new tol

* remove fidelity test from callback test suite - see PR comment #173

* add qualified path for get_datavec

---------

Co-authored-by: Andy Goldschmidt <agoldschmidt11@gmail.com>
Co-authored-by: Jack Champagne <jackchampagne.r@gmail.com>
Co-authored-by: Jack Champagne <jchampag@andrew.cmu.edu>
---
 Project.toml                                  |   23 +-
 docs/Project.toml                             |    1 +
 docs/literate/examples/multilevel_transmon.jl |    4 +-
 docs/literate/examples/two_qubit_gates.jl     |    8 +-
 docs/literate/quickstart.jl                   |    4 +-
 docs/make.jl                                  |    2 -
 .../generated/examples/multilevel_transmon.md |    4 +-
 .../src/generated/examples/two_qubit_gates.md |    8 +-
 docs/src/generated/quickstart.md              |    4 +-
 docs/src/lib.md                               |   40 -
 src/QuantumCollocation.jl                     |   58 +-
 src/callbacks.jl                              |   43 +-
 src/constraints/_constraints.jl               |  133 --
 src/constraints/complex_modulus_constraint.jl |  146 ---
 src/constraints/fidelity_constraint.jl        |  307 -----
 src/constraints/l1_slack_constraint.jl        |   59 -
 .../linear_trajectory_constraints.jl          |  349 ------
 src/direct_sums.jl                            |  195 ++-
 src/dynamics.jl                               |  400 ------
 src/embedded_operators.jl                     |  468 -------
 src/evaluators.jl                             |  176 ---
 src/integrators/_integrator_utils.jl          |  113 --
 src/integrators/_integrators.jl               |   67 -
 src/integrators/derivative_integrator.jl      |   96 --
 src/integrators/exponential_integrators.jl    |  412 -------
 src/integrators/pade_integrators.jl           | 1096 -----------------
 src/isomorphisms.jl                           |  205 ---
 src/losses/_experimental_loss_functions.jl    |   51 -
 src/losses/_losses.jl                         |   89 --
 src/losses/quantum_loss.jl                    |  169 ---
 src/losses/quantum_state_infidelity_loss.jl   |   99 --
 src/losses/unitary_infidelity_loss.jl         |  553 ---------
 src/losses/unitary_trace_loss.jl              |   49 -
 src/objectives/_objectives.jl                 |  128 --
 src/objectives/minimum_time_objective.jl      |   57 -
 src/objectives/quantum_objective.jl           |  217 ----
 src/objectives/regularizer_objective.jl       |  633 ----------
 .../unitary_infidelity_objective.jl           |  190 ---
 .../unitary_robustness_objective.jl           |  270 ----
 src/options.jl                                |  122 --
 src/plotting.jl                               |   71 --
 src/problem_solvers.jl                        |  100 --
 src/problem_templates/_problem_templates.jl   |   19 +-
 .../density_operator_smooth_pulse_problem.jl  |  123 ++
 .../quantum_state_minimum_time_problem.jl     |   12 +-
 .../quantum_state_sampling_problem.jl         |   33 +-
 .../quantum_state_smooth_pulse_problem.jl     |   87 +-
 .../unitary_bang_bang_problem.jl              |   34 +-
 .../unitary_direct_sum_problem.jl             |   15 +-
 .../unitary_minimum_time_problem.jl           |   34 +-
 .../unitary_robustness_problem.jl             |   75 +-
 .../unitary_sampling_problem.jl               |   22 +-
 .../unitary_smooth_pulse_problem.jl           |   98 +-
 src/problems.jl                               |  435 -------
 src/quantum_object_utils.jl                   |  284 -----
 .../_quantum_system_templates.jl              |   12 +-
 src/quantum_system_templates/cats.jl          |   61 +
 src/quantum_system_templates/rydberg.jl       |   26 +-
 src/quantum_system_templates/transmons.jl     |    7 +-
 src/quantum_system_utils.jl                   |  373 ------
 src/quantum_systems.jl                        |  425 -------
 src/rollouts.jl                               |  180 ++-
 src/save_load_utils.jl                        |  154 ---
 src/structure_utils.jl                        |  258 ----
 src/trajectory_initialization.jl              |   12 +-
 65 files changed, 701 insertions(+), 9297 deletions(-)
 delete mode 100644 src/constraints/_constraints.jl
 delete mode 100644 src/constraints/complex_modulus_constraint.jl
 delete mode 100644 src/constraints/fidelity_constraint.jl
 delete mode 100644 src/constraints/l1_slack_constraint.jl
 delete mode 100644 src/constraints/linear_trajectory_constraints.jl
 delete mode 100644 src/dynamics.jl
 delete mode 100644 src/embedded_operators.jl
 delete mode 100644 src/evaluators.jl
 delete mode 100644 src/integrators/_integrator_utils.jl
 delete mode 100644 src/integrators/_integrators.jl
 delete mode 100644 src/integrators/derivative_integrator.jl
 delete mode 100644 src/integrators/exponential_integrators.jl
 delete mode 100644 src/integrators/pade_integrators.jl
 delete mode 100644 src/isomorphisms.jl
 delete mode 100644 src/losses/_experimental_loss_functions.jl
 delete mode 100644 src/losses/_losses.jl
 delete mode 100644 src/losses/quantum_loss.jl
 delete mode 100644 src/losses/quantum_state_infidelity_loss.jl
 delete mode 100644 src/losses/unitary_infidelity_loss.jl
 delete mode 100644 src/losses/unitary_trace_loss.jl
 delete mode 100644 src/objectives/_objectives.jl
 delete mode 100644 src/objectives/minimum_time_objective.jl
 delete mode 100644 src/objectives/quantum_objective.jl
 delete mode 100644 src/objectives/regularizer_objective.jl
 delete mode 100644 src/objectives/unitary_infidelity_objective.jl
 delete mode 100644 src/objectives/unitary_robustness_objective.jl
 delete mode 100644 src/options.jl
 delete mode 100644 src/plotting.jl
 delete mode 100644 src/problem_solvers.jl
 create mode 100644 src/problem_templates/density_operator_smooth_pulse_problem.jl
 delete mode 100644 src/problems.jl
 delete mode 100644 src/quantum_object_utils.jl
 create mode 100644 src/quantum_system_templates/cats.jl
 delete mode 100644 src/quantum_system_utils.jl
 delete mode 100644 src/quantum_systems.jl
 delete mode 100644 src/save_load_utils.jl
 delete mode 100644 src/structure_utils.jl

diff --git a/Project.toml b/Project.toml
index f1e2bf65..a2cb7362 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "QuantumCollocation"
 uuid = "0dc23a59-5ffb-49af-b6bd-932a8ae77adf"
 authors = ["Aaron Trowbridge <aaron.j.trowbridge@gmail.com> and contributors"]
-version = "0.3.2"
+version = "0.5.0"
 
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
@@ -11,13 +11,15 @@ Einsum = "b7d42ee7-0b51-5a75-98ca-779d3107e4c0"
 ExponentialAction = "e24c0720-ea99-47e8-929e-571b494574d3"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
-Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 NamedTrajectories = "538bc3a1-5ab9-4fc3-b776-35ca1e893e08"
+PiccoloQuantumObjects = "5a402ddf-f93c-42eb-975e-5582dcda653d"
+PkgTemplates = "14b8a8f1-9102-5b29-a752-f990bacb7fe1"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
+QuantumCollocationCore = "2b384925-53cb-4042-a8d2-6faa627467e1"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -27,24 +29,29 @@ TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
 TrajectoryIndexingUtils = "6dad8b7f-dd9a-4c28-9b70-85b9a079bfc8"
 
 [compat]
-BenchmarkTools = "1.5"
+BenchmarkTools = "1.6"
 CairoMakie = "0.12"
 Distributions = "0.25"
 Einsum = "0.4"
 ExponentialAction = "0.2"
 ForwardDiff = "0.10"
 Interpolations = "0.15"
-Ipopt = "1.6"
 JLD2 = "0.5"
-MathOptInterface = "1.31"
-NamedTrajectories = "0.2.4"
+LinearAlgebra = "1.10, 1.11"
+MathOptInterface = "1.35"
+NamedTrajectories = "0.2"
+PiccoloQuantumObjects = "0.1"
+PkgTemplates = "0.7"
 ProgressMeter = "1.10"
+QuantumCollocationCore = "0.1"
+Random = "1.10, 1.11"
 Reexport = "1.2"
-Symbolics = "6.14"
+SparseArrays = "1.10, 1.11"
+Symbolics = "6.22"
 TestItemRunner = "1.0"
 TestItems = "1.0"
 TrajectoryIndexingUtils = "0.1"
-julia = "1.10"
+julia = "1.10, 1.11"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/docs/Project.toml b/docs/Project.toml
index 08396f5b..e89d4390 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -4,6 +4,7 @@ Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 LiveServer = "16fef848-5104-11e9-1b77-fb7a48bbb589"
 NamedTrajectories = "538bc3a1-5ab9-4fc3-b776-35ca1e893e08"
 QuantumCollocation = "0dc23a59-5ffb-49af-b6bd-932a8ae77adf"
+QuantumCollocationCore = "2b384925-53cb-4042-a8d2-6faa627467e1"
 Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 
 [compat]
diff --git a/docs/literate/examples/multilevel_transmon.jl b/docs/literate/examples/multilevel_transmon.jl
index 1e23f3c3..0281de1f 100644
--- a/docs/literate/examples/multilevel_transmon.jl
+++ b/docs/literate/examples/multilevel_transmon.jl
@@ -71,7 +71,7 @@ solve!(prob; max_iter=50)
 
 # Let's look at the fidelity in the subspace
 
-println("Fidelity: ", unitary_fidelity(prob; subspace=op.subspace_indices))
+println("Fidelity: ", unitary_rollout_fidelity(prob; subspace=op.subspace_indices))
 
 # and plot the result using the `plot_unitary_populations` function.
 
@@ -96,7 +96,7 @@ solve!(prob_leakage; max_iter=50)
 
 # Let's look at the fidelity in the subspace
 
-println("Fidelity: ", unitary_fidelity(prob_leakage; subspace=op.subspace_indices))
+println("Fidelity: ", unitary_rollout_fidelity(prob_leakage; subspace=op.subspace_indices))
 
 # and plot the result using the `plot_unitary_populations` function.
 
diff --git a/docs/literate/examples/two_qubit_gates.jl b/docs/literate/examples/two_qubit_gates.jl
index 12557264..78f7de22 100644
--- a/docs/literate/examples/two_qubit_gates.jl
+++ b/docs/literate/examples/two_qubit_gates.jl
@@ -111,7 +111,7 @@ prob = UnitarySmoothPulseProblem(
 solve!(prob; max_iter=100)
 
 ## Let's take a look at the final fidelity
-unitary_fidelity(prob)
+unitary_rollout_fidelity(prob)
 
 # Looks good!
 
@@ -124,7 +124,7 @@ min_time_prob = UnitaryMinimumTimeProblem(prob; final_fidelity=.99)
 
 solve!(min_time_prob; max_iter=300)
 
-unitary_fidelity(min_time_prob)
+unitary_rollout_fidelity(min_time_prob)
 
 # And let's plot this solution
 plot_unitary_populations(min_time_prob)
@@ -176,7 +176,7 @@ prob = UnitarySmoothPulseProblem(
 solve!(prob; max_iter=1_000)
 
 ## Let's take a look at the final fidelity
-unitary_fidelity(prob)
+unitary_rollout_fidelity(prob)
 
 # Again, looks good!
 
@@ -190,7 +190,7 @@ min_time_prob = UnitaryMinimumTimeProblem(prob; final_fidelity=.999)
 
 solve!(min_time_prob; max_iter=300)
 
-unitary_fidelity(min_time_prob)
+unitary_rollout_fidelity(min_time_prob)
 
 # And let's plot this solution
 
diff --git a/docs/literate/quickstart.jl b/docs/literate/quickstart.jl
index 6ed8d92a..715ab0ea 100644
--- a/docs/literate/quickstart.jl
+++ b/docs/literate/quickstart.jl
@@ -56,7 +56,7 @@ solve!(prob; max_iter=30)
 
 # The above output comes from the Ipopt.jl solver. To see the final fidelity we can use the `unitary_fidelity` function exported by QuantumCollocation.jl.
 
-println("Final fidelity: ", unitary_fidelity(prob))
+println("Final fidelity: ", unitary_rollout_fidelity(prob))
 
 # We can also easily plot the solutions using the `plot` function exported by NamedTrajectories.jl.
 
@@ -86,7 +86,7 @@ solve!(prob_min_time; max_iter=30)
 
 # We can see that the final fidelity is indeed greater than the minimum fidelity we set.
 
-println("Final fidelity:    ", unitary_fidelity(prob_min_time))
+println("Final fidelity:    ", unitary_rollout_fidelity(prob_min_time))
 
 # and that the duration of the pulse has decreased.
 
diff --git a/docs/make.jl b/docs/make.jl
index fcf4d946..eabb4efd 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -6,8 +6,6 @@ push!(LOAD_PATH, joinpath(@__DIR__, "..", "src"))
 
 # DocMeta.setdocmeta!(QuantumCollocation, :DocTestSetup, :(using QuantumCollocation); recursive=true)
 
-println("")
-
 pages = [
     "Home" => "index.md",
     "Quickstart Guide" => "generated/quickstart.md",
diff --git a/docs/src/generated/examples/multilevel_transmon.md b/docs/src/generated/examples/multilevel_transmon.md
index 35584c57..2b011a3a 100644
--- a/docs/src/generated/examples/multilevel_transmon.md
+++ b/docs/src/generated/examples/multilevel_transmon.md
@@ -83,7 +83,7 @@ solve!(prob; max_iter=50)
 Let's look at the fidelity in the subspace
 
 ````@example multilevel_transmon
-println("Fidelity: ", unitary_fidelity(prob; subspace=op.subspace_indices))
+println("Fidelity: ", unitary_rollout_fidelity(prob; subspace=op.subspace_indices))
 ````
 
 and plot the result using the `plot_unitary_populations` function.
@@ -114,7 +114,7 @@ solve!(prob_leakage; max_iter=50)
 Let's look at the fidelity in the subspace
 
 ````@example multilevel_transmon
-println("Fidelity: ", unitary_fidelity(prob_leakage; subspace=op.subspace_indices))
+println("Fidelity: ", unitary_rollout_fidelity(prob_leakage; subspace=op.subspace_indices))
 ````
 
 and plot the result using the `plot_unitary_populations` function.
diff --git a/docs/src/generated/examples/two_qubit_gates.md b/docs/src/generated/examples/two_qubit_gates.md
index 8bf811bd..fedba3bf 100644
--- a/docs/src/generated/examples/two_qubit_gates.md
+++ b/docs/src/generated/examples/two_qubit_gates.md
@@ -120,7 +120,7 @@ prob = UnitarySmoothPulseProblem(
 solve!(prob; max_iter=100)
 
 # Let's take a look at the final fidelity
-unitary_fidelity(prob)
+unitary_rollout_fidelity(prob)
 ````
 
 Looks good!
@@ -138,7 +138,7 @@ min_time_prob = UnitaryMinimumTimeProblem(prob; final_fidelity=.99)
 
 solve!(min_time_prob; max_iter=300)
 
-unitary_fidelity(min_time_prob)
+unitary_rollout_fidelity(min_time_prob)
 ````
 
 And let's plot this solution
@@ -194,7 +194,7 @@ prob = UnitarySmoothPulseProblem(
 solve!(prob; max_iter=1_000)
 
 # Let's take a look at the final fidelity
-unitary_fidelity(prob)
+unitary_rollout_fidelity(prob)
 ````
 
 Again, looks good!
@@ -212,7 +212,7 @@ min_time_prob = UnitaryMinimumTimeProblem(prob; final_fidelity=.999)
 
 solve!(min_time_prob; max_iter=300)
 
-unitary_fidelity(min_time_prob)
+unitary_rollout_fidelity(min_time_prob)
 ````
 
 And let's plot this solution
diff --git a/docs/src/generated/quickstart.md b/docs/src/generated/quickstart.md
index f12aba16..ff865d0b 100644
--- a/docs/src/generated/quickstart.md
+++ b/docs/src/generated/quickstart.md
@@ -63,7 +63,7 @@ solve!(prob; max_iter=30)
 The above output comes from the Ipopt.jl solver. To see the final fidelity we can use the `unitary_fidelity` function exported by QuantumCollocation.jl.
 
 ````@example quickstart
-println("Final fidelity: ", unitary_fidelity(prob))
+println("Final fidelity: ", unitary_rollout_fidelity(prob))
 ````
 
 We can also easily plot the solutions using the `plot` function exported by NamedTrajectories.jl.
@@ -99,7 +99,7 @@ solve!(prob_min_time; max_iter=30)
 We can see that the final fidelity is indeed greater than the minimum fidelity we set.
 
 ````@example quickstart
-println("Final fidelity:    ", unitary_fidelity(prob_min_time))
+println("Final fidelity:    ", unitary_rollout_fidelity(prob_min_time))
 ````
 
 and that the duration of the pulse has decreased.
diff --git a/docs/src/lib.md b/docs/src/lib.md
index 5b712e88..b308faf5 100644
--- a/docs/src/lib.md
+++ b/docs/src/lib.md
@@ -10,41 +10,6 @@ Modules = [QuantumCollocation.ProblemTemplates]
 Modules = [QuantumCollocation.DirectSums]
 ```
 
-## Quantum Object Utils
-```@autodocs
-Modules = [QuantumCollocation.QuantumObjectUtils]
-```
-
-## Quantum Systems 
-```@autodocs
-Modules = [QuantumCollocation.QuantumSystems]
-```
-
-## Integrators
-```@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]
@@ -70,11 +35,6 @@ Modules = [QuantumCollocation.Rollouts]
 Modules = [QuantumCollocation.SaveLoadUtils]
 ```
 
-## Structure Utils
-```@autodocs
-Modules = [QuantumCollocation.StructureUtils]
-```
-
 ## Trajectory Initialization
 ```@autodocs
 Modules = [QuantumCollocation.TrajectoryInitialization]
diff --git a/src/QuantumCollocation.jl b/src/QuantumCollocation.jl
index 415f3d2c..006851b7 100644
--- a/src/QuantumCollocation.jl
+++ b/src/QuantumCollocation.jl
@@ -2,51 +2,8 @@ module QuantumCollocation
 
 using Reexport
 
-
-include("options.jl")
-@reexport using .Options
-
-include("isomorphisms.jl")
-@reexport using .Isomorphisms
-
-include("quantum_object_utils.jl")
-@reexport using .QuantumObjectUtils
-
-include("structure_utils.jl")
-@reexport using .StructureUtils
-
-include("quantum_systems.jl")
-@reexport using .QuantumSystems
-
-include("quantum_system_templates/_quantum_system_templates.jl")
-@reexport using .QuantumSystemTemplates
-
-include("embedded_operators.jl")
-@reexport using .EmbeddedOperators
-
-include("quantum_system_utils.jl")
-@reexport using .QuantumSystemUtils
-
-include("losses/_losses.jl")
-@reexport using .Losses
-
-include("constraints/_constraints.jl")
-@reexport using .Constraints
-
-include("objectives/_objectives.jl")
-@reexport using .Objectives
-
-include("integrators/_integrators.jl")
-@reexport using .Integrators
-
-include("dynamics.jl")
-@reexport using .Dynamics
-
-include("evaluators.jl")
-@reexport using .Evaluators
-
-include("problems.jl")
-@reexport using .Problems
+@reexport using QuantumCollocationCore
+@reexport using PiccoloQuantumObjects
 
 include("direct_sums.jl")
 @reexport using .DirectSums
@@ -63,17 +20,10 @@ include("trajectory_interpolations.jl")
 include("problem_templates/_problem_templates.jl")
 @reexport using .ProblemTemplates
 
-include("save_load_utils.jl")
-@reexport using .SaveLoadUtils
-
-include("problem_solvers.jl")
-@reexport using .ProblemSolvers
-
-include("plotting.jl")
-@reexport using .Plotting
+include("quantum_system_templates/_quantum_system_templates.jl")
+@reexport using .QuantumSystemTemplates
 
 include("callbacks.jl")
 @reexport using .Callbacks
 
-
 end
diff --git a/src/callbacks.jl b/src/callbacks.jl
index e7c353fd..56607dd2 100644
--- a/src/callbacks.jl
+++ b/src/callbacks.jl
@@ -7,21 +7,20 @@ export trajectory_history_callback
 using NamedTrajectories
 using TestItemRunner
 
-using ..Losses
-using ..Problems: QuantumControlProblem, get_datavec
-using ..QuantumSystems
+using QuantumCollocationCore
+using PiccoloQuantumObjects
+
 using ..Rollouts
 
 
 function best_rollout_callback(
-    prob::QuantumControlProblem, rollout_fidelity::Function;
-    system::Union{AbstractQuantumSystem, AbstractVector{<:AbstractQuantumSystem}}=prob.system
+    prob::QuantumControlProblem, system::Union{AbstractQuantumSystem, AbstractVector{<:AbstractQuantumSystem}}, rollout_fidelity::Function
 )
     best_value = 0.0
     best_trajectories = []
 
     function callback(args...)
-        traj = NamedTrajectory(get_datavec(prob), prob.trajectory)
+        traj = NamedTrajectory(Problems.get_datavec(prob), prob.trajectory)
         value = rollout_fidelity(traj, system)
         if value > best_value
             best_value = value
@@ -33,18 +32,18 @@ function best_rollout_callback(
     return callback, best_trajectories
 end
 
-function best_rollout_fidelity_callback(prob::QuantumControlProblem)
-    return best_rollout_callback(prob, fidelity)
+function best_rollout_fidelity_callback(prob::QuantumControlProblem, system::Union{AbstractQuantumSystem, AbstractVector{<:AbstractQuantumSystem}})
+    return best_rollout_callback(prob, system, rollout_fidelity)
 end
 
-function best_unitary_rollout_fidelity_callback(prob::QuantumControlProblem)
-    return best_rollout_callback(prob, unitary_fidelity)
+function best_unitary_rollout_fidelity_callback(prob::QuantumControlProblem, system::Union{AbstractQuantumSystem, AbstractVector{<:AbstractQuantumSystem}})
+    return best_rollout_callback(prob, system, unitary_rollout_fidelity)
 end
 
 function trajectory_history_callback(prob::QuantumControlProblem)
     trajectory_history = []
     function callback(args...)
-        push!(trajectory_history, NamedTrajectory(get_datavec(prob), prob.trajectory))
+        push!(trajectory_history, NamedTrajectory(Problems.get_datavec(prob), prob.trajectory))
         return true
     end
 
@@ -56,19 +55,17 @@ end
 @testitem "Callback returns false early stops" begin
     using MathOptInterface
     const MOI = MathOptInterface
+    using LinearAlgebra
     include("../test/test_utils.jl")
 
     prob = smooth_quantum_state_problem()
 
     my_callback = (kwargs...) -> false
 
-    initial = fidelity(prob)
     solve!(prob, max_iter=20, callback=my_callback)
-    final = fidelity(prob)
 
     # callback forces problem to exit early as per Ipopt documentation
     @test MOI.get(prob.optimizer, MOI.TerminationStatus()) == MOI.INTERRUPTED
-    @test initial ≈ final atol=1e-2
 end
 
 
@@ -94,11 +91,11 @@ end
 
     prob, system = smooth_quantum_state_problem(return_system=true)
 
-    callback, best_trajs = best_rollout_fidelity_callback(prob)
+    callback, best_trajs = best_rollout_fidelity_callback(prob, system)
     @test length(best_trajs) == 0
 
     # measure fidelity
-    before = fidelity(prob)
+    before = rollout_fidelity(prob, system)
     solve!(prob, max_iter=20, callback=callback)
 
     # length must increase if iterations are made
@@ -106,8 +103,8 @@ end
     @test best_trajs[end] isa NamedTrajectory
     
     # fidelity ranking
-    after = fidelity(prob)
-    best = fidelity(best_trajs[end], system)
+    after = rollout_fidelity(prob, system)
+    best = rollout_fidelity(best_trajs[end], system)
     
     @test before < after
     @test before < best
@@ -122,11 +119,11 @@ end
 
     prob, system = smooth_unitary_problem(return_system=true)
 
-    callback, best_trajs = best_unitary_rollout_fidelity_callback(prob)
+    callback, best_trajs = best_unitary_rollout_fidelity_callback(prob, system)
     @test length(best_trajs) == 0
 
     # measure fidelity
-    before = unitary_fidelity(prob)
+    before = unitary_rollout_fidelity(prob.trajectory, system)
     solve!(prob, max_iter=20, callback=callback)
 
     # length must increase if iterations are made
@@ -134,8 +131,8 @@ end
     @test best_trajs[end] isa NamedTrajectory
     
     # fidelity ranking
-    after = unitary_fidelity(prob)
-    best = unitary_fidelity(best_trajs[end], system)
+    after = unitary_rollout_fidelity(prob.trajectory, system)
+    best = unitary_rollout_fidelity(best_trajs[end], system)
     
     @test before < after
     @test before < best
@@ -174,4 +171,4 @@ end
     @test length(obj_vals) == 4   # problem init, iter 1, iter 2, iter 3 (terminate)
 end
 
-end
\ No newline at end of file
+end
diff --git a/src/constraints/_constraints.jl b/src/constraints/_constraints.jl
deleted file mode 100644
index 8b17808e..00000000
--- a/src/constraints/_constraints.jl
+++ /dev/null
@@ -1,133 +0,0 @@
-module Constraints
-
-export AbstractConstraint
-
-export LinearConstraint
-
-export constrain!
-
-export NonlinearConstraint
-export NonlinearEqualityConstraint
-export NonlinearInequalityConstraint
-
-using ..Losses
-using ..Isomorphisms
-using ..StructureUtils
-using ..Options
-
-using TrajectoryIndexingUtils
-using NamedTrajectories
-using ForwardDiff
-using SparseArrays
-using Ipopt
-using MathOptInterface
-const MOI = MathOptInterface
-
-# TODO:
-# - [ ] Do not reference the Z object in the constraint (components only / remove "name")
-
-# ----------------------------------------------------------------------------- #
-#                     Abstract Constraints                                      #
-# ----------------------------------------------------------------------------- #
-
-abstract type AbstractConstraint end
-abstract type LinearConstraint <: AbstractConstraint end
-abstract type NonlinearConstraint <: AbstractConstraint end
-
-include("linear_trajectory_constraints.jl")
-include("complex_modulus_constraint.jl")
-include("fidelity_constraint.jl")
-include("l1_slack_constraint.jl")
-
-
-# ----------------------------------------------------------------------------- #
-#                     Linear Constraint                                         #
-# ----------------------------------------------------------------------------- #
-
-"""
-    constrain!(opt::Ipopt.Optimizer, vars::Vector{MOI.VariableIndex}, cons::Vector{LinearConstraint}, traj::NamedTrajectory; verbose=false)
-
-Supplies a set of LinearConstraints to  IPOPT using MathOptInterface
-
-"""
-function constrain!(
-    opt::Ipopt.Optimizer,
-    vars::Vector{MOI.VariableIndex},
-    cons::Vector{LinearConstraint},
-    traj::NamedTrajectory;
-    verbose=false
-)
-    for con in cons
-        if verbose
-            println("applying constraint: ", con.label)
-        end
-        con(opt, vars, traj)
-    end
-end
-
-
-# ----------------------------------------------------------------------------- #
-#                     Nonlinear Constraint                                      #
-# ----------------------------------------------------------------------------- #
-
-
-function NonlinearConstraint(params::Dict)
-    return eval(params[:type])(; delete!(params, :type)...)
-end
-
-"""
-    struct NonlinearEqualityConstraint
-
-Represents a nonlinear equality constraint.
-
-# Fields
-- `g::Function`: the constraint function
-- `∂g::Function`: the Jacobian of the constraint function
-- `∂g_structure::Vector{Tuple{Int, Int}}`: the structure of the Jacobian
-   i.e. all non-zero entries
-- `μ∂²g::Function`: the Hessian of the constraint function
-- `μ∂²g_structure::Vector{Tuple{Int, Int}}`: the structure of the Hessian
-- `dim::Int`: the dimension of the constraint function
-- `params::Dict{Symbol, Any}`: a dictionary of parameters
-
-"""
-struct NonlinearEqualityConstraint <: NonlinearConstraint
-    g::Function
-    ∂g::Function
-    ∂g_structure::Vector{Tuple{Int, Int}}
-    μ∂²g::Union{Nothing, Function}
-    μ∂²g_structure::Union{Nothing, Vector{Tuple{Int, Int}}}
-    dim::Int
-    params::Dict{Symbol, Any}
-end
-
-"""
-    struct NonlinearInequalityConstraint
-
-Represents a nonlinear inequality constraint.
-
-# Fields
-- `g::Function`: the constraint function
-- `∂g::Function`: the Jacobian of the constraint function
-- `∂g_structure::Vector{Tuple{Int, Int}}`: the structure of the Jacobian
-   i.e. all non-zero entries
-- `μ∂²g::Function`: the Hessian of the constraint function
-- `μ∂²g_structure::Vector{Tuple{Int, Int}}`: the structure of the Hessian
-- `dim::Int`: the dimension of the constraint function
-- `params::Dict{Symbol, Any}`: a dictionary of parameters containing additional
-   information about the constraint
-
-"""
-struct NonlinearInequalityConstraint <: NonlinearConstraint
-    g::Function
-    ∂g::Function
-    ∂g_structure::Vector{Tuple{Int, Int}}
-    μ∂²g::Union{Nothing, Function}
-    μ∂²g_structure::Union{Nothing, Vector{Tuple{Int, Int}}}
-    dim::Int
-    params::Dict{Symbol, Any}
-end
-
-
-
-end
diff --git a/src/constraints/complex_modulus_constraint.jl b/src/constraints/complex_modulus_constraint.jl
deleted file mode 100644
index 6972edc3..00000000
--- a/src/constraints/complex_modulus_constraint.jl
+++ /dev/null
@@ -1,146 +0,0 @@
-export ComplexModulusContraint
-
-
-"""
-    ComplexModulusContraint(<keyword arguments>)
-
-Returns a NonlinearInequalityConstraint on the complex modulus of a complex control
-
-TODO: Changed zdim -> dim. Constraint should be tested for global params.
-
-# Arguments
-- `R::Union{Float64,Nothing}=nothing`: the maximum allowed complex modulus
-- `comps::Union{AbstractVector{Int},Nothing}=nothing`: the components of the complex control,
-   both the real and imaginary parts
-- `times::Union{AbstractVector{Int},Nothing}=nothing`: the times at which the constraint is applied
-- `dim::Union{Int,Nothing}=nothing`: the dimension of a single time step of the trajectory
-- `T::Union{Int,Nothing}=nothing`: the number of time steps
-"""
-function ComplexModulusContraint(;
-    R::Union{Float64, Nothing}=nothing,
-    comps::Union{AbstractVector{Int}, Nothing}=nothing,
-    times::Union{AbstractVector{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(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"
-
-    params = Dict{Symbol, Any}()
-
-    params[:type] = :ComplexModulusContraint
-    params[:R] = R
-    params[:comps] = comps
-    params[:times] = times
-    params[:zdim] = zdim
-    params[:T] = T
-
-    gₜ(xₜ, yₜ) = [R^2 - xₜ^2 - yₜ^2]
-    ∂gₜ(xₜ, yₜ) = [-2xₜ, -2yₜ]
-    μₜ∂²gₜ(μₜ) = sparse([1, 2], [1, 2], [-2μₜ, -2μₜ])
-
-    @views function g(Z⃗)
-        r = zeros(length(times))
-        for (i, t) ∈ enumerate(times)
-            zₜ = Z⃗[slice(t, comps, zdim)]
-            xₜ = zₜ[1]
-            yₜ = zₜ[2]
-            r[i] = gₜ(xₜ, yₜ)[1]
-        end
-        return r
-    end
-
-    ∂g_structure = []
-
-    for (i, t) ∈ enumerate(times)
-        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, zdim)]
-            xₜ = zₜ[1]
-            yₜ = zₜ[2]
-            ∂[i, slice(t, comps, zdim)] = ∂gₜ(xₜ, yₜ)
-        end
-        if ipopt
-            return [∂[i, j] for (i, j) in ∂g_structure]
-        else
-            return ∂
-        end
-    end
-
-    μ∂²g_structure = []
-
-    for t ∈ times
-        push!(
-            μ∂²g_structure,
-            (
-                index(t, comps[1], zdim),
-                index(t, comps[1], zdim)
-            )
-        )
-        push!(
-            μ∂²g_structure,
-            (
-                index(t, comps[2], zdim),
-                index(t, comps[2], zdim)
-            )
-        )
-    end
-
-    function μ∂²g(Z⃗, μ; ipopt=true)
-        n = length(Z⃗)
-        μ∂² = spzeros(n, n)
-        for (i, t) ∈ enumerate(times)
-            t_slice = slice(t, comps, zdim)
-            μ∂²[t_slice, t_slice] = μₜ∂²gₜ(μ[i])
-        end
-        if ipopt
-            return [μ∂²[i, j] for (i, j) in μ∂²g_structure]
-        else
-            return μ∂²
-        end
-    end
-
-    return NonlinearInequalityConstraint(
-        g,
-        ∂g,
-        ∂g_structure,
-        μ∂²g,
-        μ∂²g_structure,
-        length(times),
-        params
-    )
-end
-
-"""
-    ComplexModulusContraint(symb::Symbol, R::Float64, traj::NamedTrajectory)
-
-Returns a ComplexModulusContraint for the complex control NamedTrajector symbol
-where R is the maximum allowed complex modulus.
-"""
-function ComplexModulusContraint(
-    name::Symbol,
-    R::Float64,
-    traj::NamedTrajectory;
-    times=1:traj.T,
-    name_comps=1:traj.dims[name]
-)
-    @assert name ∈ traj.names
-    comps = traj.components[name][name_comps]
-    return ComplexModulusContraint(;
-        R=R,
-        comps=comps,
-        times=times,
-        zdim=traj.dim,
-        T=traj.T
-    )
-end
diff --git a/src/constraints/fidelity_constraint.jl b/src/constraints/fidelity_constraint.jl
deleted file mode 100644
index 8ae8a736..00000000
--- a/src/constraints/fidelity_constraint.jl
+++ /dev/null
@@ -1,307 +0,0 @@
-export FinalFidelityConstraint
-export FinalUnitaryFidelityConstraint
-export FinalQuantumStateFidelityConstraint
-export FinalUnitaryFreePhaseFidelityConstraint
-
-###
-### FinalFidelityConstraint
-###
-
-"""
-    FinalFidelityConstraint(<keyword arguments>)
-
-
-Returns a NonlinearInequalityConstraint representing a constraint on the
-minimum allowed fidelity.
-
-# Arguments
-- `fidelity_function::Union{Function,Nothing}=nothing`: the fidelity function
-- `value::Union{Float64,Nothing}=nothing`: the minimum fidelity value allowed
-   by the constraint
-- `comps::Union{AbstractVector{Int},Nothing}=nothing`: the components of the
-   state to which the fidelity function is applied
-- `goal::Union{AbstractVector{Float64},Nothing}=nothing`: the goal state
-- `statedim::Union{Int,Nothing}=nothing`: the dimension of the state
-- `zdim::Union{Int,Nothing}=nothing`: the dimension of a single time step of the trajectory
-- `T::Union{Int,Nothing}=nothing`: the number of time steps
-- `subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing`: the subspace indices of the fidelity
-
-"""
-function FinalFidelityConstraint(;
-    fidelity_function::Union{Symbol,Function,Nothing}=nothing,
-    value::Union{Float64,Nothing}=nothing,
-    comps::Union{AbstractVector{Int},Nothing}=nothing,
-    goal::Union{AbstractVector{Float64},Nothing}=nothing,
-    statedim::Union{Int,Nothing}=nothing,
-    zdim::Union{Int,Nothing}=nothing,
-    T::Union{Int,Nothing}=nothing,
-    subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing,
-    eval_hessian::Bool=true
-)
-    @assert !isnothing(fidelity_function) "must provide a fidelity function"
-    @assert !isnothing(value) "must provide a fidelity value"
-    @assert !isnothing(comps) "must provide a list of components"
-    @assert !isnothing(goal) "must provide a goal state"
-    @assert !isnothing(statedim) "must provide a state dimension"
-    @assert !isnothing(zdim) "must provide a z dimension"
-    @assert !isnothing(T) "must provide a T"
-
-    if fidelity_function isa Symbol
-        fidelity_function_symbol = fidelity_function
-        fidelity_function = eval(fidelity_function)
-    else
-        fidelity_function_symbol = Symbol(fidelity_function)
-    end
-
-    if isnothing(subspace)
-        fid = x -> fidelity_function(x, goal)
-    else
-        fid = x -> fidelity_function(x, goal; subspace=subspace)
-    end
-
-    @assert fid(randn(statedim)) isa Float64 "fidelity function must return a scalar"
-
-    params = Dict{Symbol, Any}()
-
-    if fidelity_function_symbol ∉ names(Losses)
-        @warn "Fidelity function :$(string(fidelity_function_symbol)) is not exported. Unable to save this constraint."
-        params[:type] = :FinalFidelityConstraint
-        params[:fidelity_function] = :not_saveable
-    else
-        params[:type] = :FinalFidelityConstraint
-        params[:fidelity_function] = fidelity_function_symbol
-        params[:value] = value
-        params[:comps] = comps
-        params[:goal] = goal
-        params[:statedim] = statedim
-        params[:zdim] = zdim
-        params[:T] = T
-        params[:subspace] = subspace
-        params[:eval_hessian] = eval_hessian
-    end
-
-    state_slice = slice(T, comps, zdim)
-
-    ℱ(x) = [fid(x)]
-
-    g(Z⃗) = ℱ(Z⃗[state_slice]) .- value
-
-    ∂ℱ(x) = ForwardDiff.jacobian(ℱ, x)
-
-    ∂ℱ_structure = jacobian_structure(∂ℱ, statedim)
-
-    col_offset = index(T, comps[1] - 1, zdim)
-
-    ∂g_structure = [(i, j + col_offset) for (i, j) in ∂ℱ_structure]
-
-    @views function ∂g(Z⃗; ipopt=true)
-        ∂ = spzeros(1, T * zdim)
-        ∂ℱ_x = ∂ℱ(Z⃗[state_slice])
-        for (i, j) ∈ ∂ℱ_structure
-            ∂[i, j + col_offset] = ∂ℱ_x[i, j]
-        end
-        if ipopt
-            return [∂[i, j] for (i, j) in ∂g_structure]
-        else
-            return ∂
-        end
-    end
-
-    if eval_hessian
-        ∂²ℱ(x) = ForwardDiff.hessian(fid, x)
-
-        ∂²ℱ_structure = hessian_of_lagrangian_structure(∂²ℱ, statedim, 1)
-
-        μ∂²g_structure = [ij .+ col_offset for ij in ∂²ℱ_structure]
-
-        @views function μ∂²g(Z⃗, μ; ipopt=true)
-            HoL = spzeros(T * zdim, T * zdim)
-            μ∂²ℱ = μ[1] * ∂²ℱ(Z⃗[state_slice])
-            for (i, j) ∈ ∂²ℱ_structure
-                HoL[i + col_offset, j + col_offset] = μ∂²ℱ[i, j]
-            end
-            if ipopt
-                return [HoL[i, j] for (i, j) in μ∂²g_structure]
-            else
-                return HoL
-            end
-        end
-
-    else
-        μ∂²g_structure = nothing
-        μ∂²g = nothing
-    end
-
-    return NonlinearInequalityConstraint(
-        g,
-        ∂g,
-        ∂g_structure,
-        μ∂²g,
-        μ∂²g_structure,
-        1,
-        params
-    )
-end
-
-###
-### FinalUnitaryFidelityConstraint
-###
-
-"""
-    FinalUnitaryFidelityConstraint(statesymb::Symbol, val::Float64, traj::NamedTrajectory)
-
-Returns a FinalFidelityConstraint for the unitary fidelity function where statesymb
-is the NamedTrajectory symbol representing the unitary.
-
-"""
-function FinalUnitaryFidelityConstraint(
-    statesymb::Symbol,
-    val::Float64,
-    traj::NamedTrajectory;
-    subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing,
-    eval_hessian::Bool=true
-)
-    return FinalFidelityConstraint(;
-        fidelity_function=iso_vec_unitary_fidelity,
-        value=val,
-        comps=traj.components[statesymb],
-        goal=traj.goal[statesymb],
-        statedim=traj.dims[statesymb],
-        zdim=traj.dim,
-        T=traj.T,
-        subspace=subspace,
-        eval_hessian=eval_hessian
-    )
-end
-
-###
-### FinalQuantumStateFidelityConstraint
-###
-
-"""
-    FinalQuantumStateFidelityConstraint(statesymb::Symbol, val::Float64, traj::NamedTrajectory)
-
-Returns a FinalFidelityConstraint for the unitary fidelity function where statesymb
-is the NamedTrajectory symbol representing the unitary.
-
-"""
-function FinalQuantumStateFidelityConstraint(
-    statesymb::Symbol,
-    val::Float64,
-    traj::NamedTrajectory;
-    kwargs...
-)
-    @assert statesymb ∈ traj.names
-    return FinalFidelityConstraint(;
-        fidelity_function=fidelity,
-        value=val,
-        comps=traj.components[statesymb],
-        goal=traj.goal[statesymb],
-        statedim=traj.dims[statesymb],
-        zdim=traj.dim,
-        T=traj.T,
-        kwargs...
-    )
-end
-
-###
-### FinalUnitaryFreePhaseFidelityConstraint
-###
-
-"""
-    FinalUnitaryFreePhaseFidelityConstraint
-
-Returns a NonlinearInequalityConstraint representing a constraint on the minimum allowed fidelity
-for a free phase unitary.
-
-"""
-function FinalUnitaryFreePhaseFidelityConstraint(;
-    value::Union{Float64,Nothing}=nothing,
-    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}},Nothing}=nothing,
-    zdim::Union{Int,Nothing}=nothing,
-    subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing,
-    eval_hessian::Bool=false
-)
-    @assert !isnothing(value) "must provide a fidelity value"
-    @assert !isnothing(state_slice) "must provide state_slice"
-    @assert !isnothing(phase_slice) "must provide phase_slice"
-    @assert !isnothing(goal) "must provide a goal state"
-    @assert !isnothing(zdim) "must provide a z dimension"
-
-    loss = :UnitaryFreePhaseInfidelityLoss
-    ℱ = eval(loss)(goal, phase_operators; subspace=subspace)
-
-    params = Dict(
-        :type => :FinalUnitaryFreePhaseFidelityConstraint,
-        :loss => loss,
-        :value => value,
-        :state_slice => state_slice,
-        :phase_slice => phase_slice,
-        :goal => goal,
-        :phase_operators => phase_operators,
-        :subspace => subspace,
-        :eval_hessian => eval_hessian,
-    )
-
-    @views function g(Z⃗)
-        Ũ⃗ = Z⃗[state_slice]
-        ϕ⃗ = Z⃗[phase_slice]
-        return [(1 - value) - ℱ(Ũ⃗, ϕ⃗)]
-    end
-
-    ∂g_structure = [(1, j) for j ∈ [state_slice; phase_slice]]
-
-    @views function ∂g(Z⃗; ipopt=true)
-        Ũ⃗ = Z⃗[state_slice]
-        ϕ⃗ = Z⃗[phase_slice]
-        ∂ = -ℱ(Ũ⃗, ϕ⃗; gradient=true)
-
-        if ipopt
-            return ∂
-        else
-            ∂_fill = spzeros(1, zdim)
-            for (∂ᵢⱼ, (i, j)) in zip(∂, ∂g_structure)
-                ∂_fill[i, j] = ∂ᵢⱼ
-            end
-            return ∂_fill
-        end
-    end
-
-    # Hessian
-    μ∂²g_structure = []
-    μ∂²g = nothing
-
-    return NonlinearInequalityConstraint(
-        g,
-        ∂g,
-        ∂g_structure,
-        μ∂²g,
-        μ∂²g_structure,
-        1,
-        params
-    )
-end
-
-function FinalUnitaryFreePhaseFidelityConstraint(
-    state_name::Symbol,
-    phase_name::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_name], traj.dim),
-        phase_slice=traj.global_components[phase_name],
-        goal=traj.goal[state_name],
-        phase_operators=phase_operators,
-        zdim=length(traj),
-        subspace=subspace,
-        eval_hessian=eval_hessian
-    )
-end
\ No newline at end of file
diff --git a/src/constraints/l1_slack_constraint.jl b/src/constraints/l1_slack_constraint.jl
deleted file mode 100644
index e852e00f..00000000
--- a/src/constraints/l1_slack_constraint.jl
+++ /dev/null
@@ -1,59 +0,0 @@
-export L1SlackConstraint
-
-
-struct L1SlackConstraint <: LinearConstraint
-    var_name::Symbol
-    slack_names::Vector{Symbol}
-    indices::AbstractVector{Int}
-    times::AbstractVector{Int}
-    label::String
-
-    function L1SlackConstraint(
-        name::Symbol,
-        traj::NamedTrajectory;
-        indices=1:traj.dims[name],
-        times=(name ∈ keys(traj.initial) ? 2 : 1):traj.T,
-        label="L1 slack constraint on $name"
-    )
-        @assert all(i ∈ 1:traj.dims[name] for i ∈ indices)
-        s1_name = Symbol("s1_$name")
-        s2_name = Symbol("s2_$name")
-        slack_names = [s1_name, s2_name]
-        add_component!(traj, s1_name, rand(length(indices), traj.T))
-        add_component!(traj, s2_name, rand(length(indices), traj.T))
-        return new(name, slack_names, indices, times, label)
-    end
-end
-
-function (con::L1SlackConstraint)(
-    opt::Ipopt.Optimizer,
-    vars::Vector{MOI.VariableIndex},
-    traj::NamedTrajectory
-)
-    for t ∈ con.times
-        for (s1, s2, x) in zip(
-            slice(t, traj.components[con.slack_names[1]], traj.dim),
-            slice(t, traj.components[con.slack_names[2]], traj.dim),
-            slice(t, traj.components[con.var_name][con.indices], traj.dim)
-        )
-            MOI.add_constraints(
-                opt,
-                vars[s1],
-                MOI.GreaterThan(0.0)
-            )
-            MOI.add_constraints(
-                opt,
-                vars[s2],
-                MOI.GreaterThan(0.0)
-            )
-            t1 = MOI.ScalarAffineTerm(1.0, vars[s1])
-            t2 = MOI.ScalarAffineTerm(-1.0, vars[s2])
-            t3 = MOI.ScalarAffineTerm(-1.0, vars[x])
-            MOI.add_constraints(
-                opt,
-                MOI.ScalarAffineFunction([t1, t2, t3], 0.0),
-                MOI.EqualTo(0.0)
-            )
-        end
-    end
-end
\ No newline at end of file
diff --git a/src/constraints/linear_trajectory_constraints.jl b/src/constraints/linear_trajectory_constraints.jl
deleted file mode 100644
index a99e08d3..00000000
--- a/src/constraints/linear_trajectory_constraints.jl
+++ /dev/null
@@ -1,349 +0,0 @@
-export trajectory_constraints
-
-export EqualityConstraint
-export BoundsConstraint
-export TimeStepBoundsConstraint
-export TimeStepEqualityConstraint
-export TimeStepsAllEqualConstraint
-
-"""
-    trajectory_constraints(traj::NamedTrajectory)
-
-Implements the initial and final value constraints and bounds constraints on the controls
-and states as specified by traj.
-
-"""
-function trajectory_constraints(traj::NamedTrajectory)
-    cons = AbstractConstraint[]
-
-    init_names = []
-
-    # add initial equality constraints
-    for (name, val) ∈ pairs(traj.initial)
-        ts = [1]
-        js = traj.components[name]
-        con_label = "initial value of $name"
-        eq_con = EqualityConstraint(ts, js, val, traj.dim; label=con_label)
-        push!(cons, eq_con)
-        push!(init_names, name)
-    end
-
-    final_names = []
-
-    # add final equality constraints
-    for (name, val) ∈ pairs(traj.final)
-        ts = [traj.T]
-        js = traj.components[name]
-        con_label = "final value of $name"
-        eq_con = EqualityConstraint(ts, js, val, traj.dim; label=con_label)
-        push!(cons, eq_con)
-        push!(final_names, name)
-    end
-
-    # add bounds constraints
-    for (name, bound) ∈ pairs(traj.bounds)
-        if name ∈ init_names && name ∈ final_names
-            ts = 2:traj.T-1
-        elseif name ∈ init_names && !(name ∈ final_names)
-            ts = 2:traj.T
-        elseif name ∈ final_names && !(name ∈ init_names)
-            ts = 1:traj.T-1
-        else
-            ts = 1:traj.T
-        end
-        js = traj.components[name]
-        con_label = "bounds on $name"
-        bounds = collect(zip(bound[1], bound[2]))
-        bounds_con = BoundsConstraint(ts, js, bounds, traj.dim; label=con_label)
-        push!(cons, bounds_con)
-    end
-
-    return cons
-end
-
-### 
-### EqualityConstraint
-###
-
-"""
-    struct EqualityConstraint
-
-Represents a linear equality constraint.
-
-# Fields
-- `ts::AbstractArray{Int}`: the time steps at which the constraint is applied
-- `js::AbstractArray{Int}`: the components of the trajectory at which the constraint is applied
-- `vals::Vector{R}`: the values of the constraint
-- `vardim::Int`: the dimension of a single time step of the trajectory
-- `label::String`: a label for the constraint
-
-"""
-struct EqualityConstraint <: LinearConstraint
-    ts::AbstractArray{Int}
-    js::AbstractArray{Int}
-    vals::Vector{R} where R
-    vardim::Int
-    label::String
-end
-
-function EqualityConstraint(
-    t::Union{Int, AbstractArray{Int}},
-    j::Union{Int, AbstractArray{Int}},
-    val::Union{R, Vector{R}},
-    vardim::Int;
-    label="unlabeled equality constraint"
-) where R
-
-    @assert !(isa(val, Vector{R}) && isa(j, Int))
-        "if val is an array, j must be an array of integers"
-
-    @assert isa(val, R) ||
-        (isa(val, Vector{R}) && isa(j, AbstractArray{Int})) &&
-        length(val) == length(j) """
-    if j and val are both arrays, dimensions must match:
-        length(j)   = $(length(j))
-        length(val) = $(length(val))
-    """
-
-    if isa(val, R) && isa(j, AbstractArray{Int})
-        val = fill(val, length(j))
-    end
-
-    return EqualityConstraint(
-        [t...],
-        [j...],
-        [val...],
-        vardim,
-        label
-    )
-end
-
-
-function (con::EqualityConstraint)(
-    opt::Ipopt.Optimizer,
-    vars::Vector{MOI.VariableIndex},
-    traj::NamedTrajectory
-)
-    for t in con.ts
-        for (j, val) in zip(con.js, con.vals)
-            MOI.add_constraints(
-                opt,
-                vars[index(t, j, con.vardim)],
-                MOI.EqualTo(val)
-            )
-        end
-    end
-end
-
-### 
-### BoundsConstraint
-###
-
-struct BoundsConstraint <: LinearConstraint
-    ts::AbstractArray{Int}
-    js::AbstractArray{Int}
-    vals::Vector{Tuple{R, R}} where R <: Real
-    vardim::Int
-    label::String
-end
-
-function BoundsConstraint(
-    t::Union{Int, AbstractArray{Int}},
-    j::Union{Int, AbstractArray{Int}},
-    val::Union{Tuple{R, R}, Vector{Tuple{R, R}}},
-    vardim::Int;
-    label="unlabeled bounds constraint"
-) where R <: Real
-
-    @assert !(isa(val, Vector{Tuple{R, R}}) && isa(j, Int))
-        "if val is an array, var must be an array of integers"
-
-    if isa(val, Tuple{R,R}) && isa(j, AbstractArray{Int})
-
-        val = fill(val, length(j))
-
-    elseif isa(val, Tuple{R, R}) && isa(j, Int)
-
-        val = [val]
-        j = [j]
-
-    end
-
-    @assert *([v[1] <= v[2] for v in val]...) "lower bound must be less than upper bound"
-
-    return BoundsConstraint(
-        [t...],
-        j,
-        val,
-        vardim,
-        label
-    )
-end
-
-function BoundsConstraint(
-    t::Union{Int, AbstractArray{Int}},
-    j::Union{Int, AbstractArray{Int}},
-    val::Union{R, Vector{R}},
-    vardim::Int;
-    label="unlabeled bounds constraint"
-) where R <: Real
-
-    @assert !(isa(val, Vector{R}) && isa(j, Int))
-        "if val is an array, var must be an array of integers"
-
-    if isa(val, R) && isa(j, AbstractArray{Int})
-
-        bounds = (-abs(val), abs(val))
-        val = fill(bounds, length(j))
-
-    elseif isa(val, R) && isa(j, Int)
-
-        bounds = (-abs(val), abs(val))
-        val = [bounds]
-        j = [j]
-
-    elseif isa(val, Vector{R})
-
-        val = [(-abs(v), abs(v)) for v in val]
-
-    end
-
-    return BoundsConstraint(
-        [t...],
-        j,
-        val,
-        vardim,
-        label
-    )
-end
-
-function (con::BoundsConstraint)(
-    opt::Ipopt.Optimizer,
-    vars::Vector{MOI.VariableIndex},
-    traj::NamedTrajectory
-)
-    for t in con.ts
-        for (j, (lb, ub)) in zip(con.js, con.vals)
-            MOI.add_constraints(
-                opt,
-                vars[index(t, j, con.vardim)],
-                MOI.GreaterThan(lb)
-            )
-            MOI.add_constraints(
-                opt,
-                vars[index(t, j, con.vardim)],
-                MOI.LessThan(ub)
-            )
-        end
-    end
-end
-
-### 
-### TimeStepBoundsConstraint
-###
-
-struct TimeStepBoundsConstraint <: LinearConstraint
-    bounds::Tuple{R, R} where R <: Real
-    Δt_indices::AbstractVector{Int}
-    label::String
-
-    function TimeStepBoundsConstraint(
-        bounds::Tuple{R, R} where R <: Real,
-        Δt_indices::AbstractVector{Int},
-        T::Int;
-        label="time step bounds constraint"
-    )
-        @assert bounds[1] < bounds[2] "lower bound must be less than upper bound"
-        return new(bounds, Δt_indices, label)
-    end
-end
-
-function (con::TimeStepBoundsConstraint)(
-    opt::Ipopt.Optimizer,
-    vars::Vector{MOI.VariableIndex},
-    traj::NamedTrajectory
-)
-    for i ∈ con.Δt_indices
-        MOI.add_constraints(
-            opt,
-            vars[i],
-            MOI.GreaterThan(con.bounds[1])
-        )
-        MOI.add_constraints(
-            opt,
-            vars[i],
-            MOI.LessThan(con.bounds[2])
-        )
-    end
-end
-
-### 
-### TimeStepEqualityConstraint
-###
-
-struct TimeStepEqualityConstraint <: LinearConstraint
-    val::R where R <: Real
-    Δt_indices::AbstractVector{Int}
-    label::String
-
-    function TimeStepEqualityConstraint(
-        val::R where R <: Real,
-        Δt_indices::AbstractVector{Int};
-        label="unlabeled time step equality constraint"
-    )
-        return new(val, Δt_indices, label)
-    end
-end
-
-function (con::TimeStepEqualityConstraint)(
-    opt::Ipopt.Optimizer,
-    vars::Vector{MOI.VariableIndex},
-    traj::NamedTrajectory
-)
-    for i ∈ con.Δt_indices
-        MOI.add_constraints(
-            opt,
-            vars[i],
-            MOI.EqualTo(con.val)
-        )
-    end
-end
-
-struct TimeStepsAllEqualConstraint <: LinearConstraint
-    Δt_indices::AbstractVector{Int}
-    label::String
-
-    function TimeStepsAllEqualConstraint(
-        Δt_indices::AbstractVector{Int};
-        label="time step all equal constraint"
-    )
-        return new(Δt_indices, label)
-    end
-
-    function TimeStepsAllEqualConstraint(
-        Δt_symb::Symbol,
-        traj::NamedTrajectory;
-        label="time step all equal constraint"
-    )
-        Δt_comp = traj.components[Δt_symb][1]
-        Δt_indices = [index(t, Δt_comp, traj.dim) for t = 1:traj.T]
-        return new(Δt_indices, label)
-    end
-end
-
-function (con::TimeStepsAllEqualConstraint)(
-    opt::Ipopt.Optimizer,
-    vars::Vector{MOI.VariableIndex},
-    traj::NamedTrajectory
-)
-    N = length(con.Δt_indices)
-    for i = 1:N-1
-        Δtᵢ = MOI.ScalarAffineTerm(1.0, vars[con.Δt_indices[i]])
-        minusΔt̄ = MOI.ScalarAffineTerm(-1.0, vars[con.Δt_indices[end]])
-        MOI.add_constraints(
-            opt,
-            MOI.ScalarAffineFunction([Δtᵢ, minusΔt̄], 0.0),
-            MOI.EqualTo(0.0)
-        )
-    end
-end
\ No newline at end of file
diff --git a/src/direct_sums.jl b/src/direct_sums.jl
index 11c834e5..d069b8c9 100644
--- a/src/direct_sums.jl
+++ b/src/direct_sums.jl
@@ -5,16 +5,11 @@ export get_suffix
 export get_suffix_label
 export direct_sum
 
-using ..Integrators
-using ..Problems
-using ..QuantumSystems
-using ..Isomorphisms
-using ..Objectives
-
 using SparseArrays
 using TestItemRunner
-
 using NamedTrajectories
+using QuantumCollocationCore
+using PiccoloQuantumObjects
 
 
 """
@@ -41,7 +36,12 @@ end
 Returns the direct sum of two iso_vec operators.
 """
 function direct_sum(Ã⃗::AbstractVector, B̃⃗::AbstractVector)
-    return operator_to_iso_vec(direct_sum(iso_vec_to_operator(Ã⃗), iso_vec_to_operator(B̃⃗)))
+    return operator_to_iso_vec(
+        direct_sum(
+            iso_vec_to_operator(Ã⃗),
+            iso_vec_to_operator(B̃⃗)
+        )
+    )
 end
 
 """
@@ -50,18 +50,37 @@ end
 Returns the direct sum of two `QuantumSystem` objects.
 """
 function direct_sum(sys1::QuantumSystem, sys2::QuantumSystem)
-    H_drift = direct_sum(sys1.H_drift, sys2.H_drift)
-    H1_zero = spzeros(size(sys1.H_drift))
-    H2_zero = spzeros(size(sys2.H_drift))
-    H_drives = [
-        [direct_sum(H, H2_zero) for H ∈ sys1.H_drives]...,
-        [direct_sum(H1_zero, H) for H ∈ sys2.H_drives]...
-    ]
-    return QuantumSystem(
-        H_drift,
-        H_drives,
-        params=merge(sys1.params, sys2.params)
-    )
+    @assert sys1.n_drives == sys2.n_drives "System 1 drives ($(sys1.n_drives)) must equal System 2 drives ($(sys2.n_drives))"
+    n_drives = sys1.n_drives
+    H = a -> direct_sum(sys1.H(a), sys2.H(a))
+    G = a -> direct_sum(sys1.G(a), sys2.G(a))
+    ∂G = a -> [direct_sum(∂Gᵢ(a), ∂Gⱼ(a)) for (∂Gᵢ, ∂Gⱼ) ∈ zip(sys1.∂G(a), sys2.∂G(a))]
+    levels = sys1.levels + sys2.levels
+    direct_sum_params = Dict{Symbol, Dict{Symbol, Any}}()
+    if haskey(sys1.params, :system_1)
+        n_systems = length(keys(sys1.params))
+        direct_sum_params = sys1.params
+        if haskey(sys2.params, :system_1)
+            for i = 1:length(keys(sys2.params))
+                direct_sum_params[Symbol("system_$(n_systems + i)")] =
+                    sys2.params[Symbol("system_$(i)")]
+            end
+        else
+            direct_sum_params[Symbol("system_$(n_systems + 1)")] = sys2.params
+        end
+    else
+        direct_sum_params[:system_1] = sys1.params
+        if haskey(sys2.params, :system_1)
+            n_systems = length(keys(sys2.params))
+            for i = 1:length(keys(sys2.params))
+                direct_sum_params[Symbol("system_$(1 + i)")] =
+                    sys2.params[Symbol("system_$(i)")]
+            end
+        else
+            direct_sum_params[:system_2] = sys2.params
+        end
+    end
+    return QuantumSystem(H, G, ∂G, n_drives, levels, direct_sum_params)
 end
 
 direct_sum(systems::AbstractVector{<:QuantumSystem}) = reduce(direct_sum, systems)
@@ -105,13 +124,12 @@ function add_suffix(traj::NamedTrajectory, suffix::String)
     )
 end
 
-function add_suffix(sys::QuantumSystem, suffix::String)
-    return QuantumSystem(
-        sys.H_drift,
-        sys.H_drives,
-        params=add_suffix(sys.params, suffix)
-    )
-end
+# function add_suffix(sys::QuantumSystem, suffix::String)
+#     return QuantumSystem(
+#         sys.H_drift,
+#         sys.H_drives
+#     )
+# end
 
 # get suffix label utilities
 # --------------------
@@ -155,57 +173,26 @@ end
 # ---------------------------
 
 function modify_integrator_suffix(
-    modifier::Function,
     integrator::AbstractIntegrator,
-    sys::AbstractQuantumSystem,
+    modifier::Function,
+    suffix::String,
     traj::NamedTrajectory,
-    mod_traj::NamedTrajectory,
-    suffix::String
+    mod_traj::NamedTrajectory
 )
-    if integrator isa UnitaryExponentialIntegrator
-        unitary_name = get_component_names(traj, integrator.unitary_components)
-        drive_name = get_component_names(traj, integrator.drive_components)
-        return integrator(
-            sys,
-            mod_traj,
-            unitary_name=modifier(unitary_name, suffix),
-            drive_name=modifier(drive_name, suffix)
-        )
-    elseif integrator isa QuantumStateExponentialIntegrator
-        state_name = get_component_names(traj, integrator.state_components)
-        drive_name = get_component_names(traj, integrator.drive_components)
-        return integrator(
-            sys,
-            mod_traj,
-            state_name=modifier(state_name, suffix),
-            drive_name=modifier(drive_name, suffix)
-        )
-    elseif integrator isa UnitaryPadeIntegrator
-        unitary_name = get_component_names(traj, integrator.unitary_components)
-        drive_name = get_component_names(traj, integrator.drive_components)
-        return integrator(
-            sys,
-            mod_traj,
-            unitary_name=modifier(unitary_name, suffix),
-            drive_name=modifier(drive_name, suffix)
-        )
-    elseif integrator isa QuantumStatePadeIntegrator
+    mod_integrator = deepcopy(integrator)
+
+    if integrator isa QuantumIntegrator
         state_name = get_component_names(traj, integrator.state_components)
         drive_name = get_component_names(traj, integrator.drive_components)
-        return integrator(
-            sys,
-            mod_traj,
-            state_name=modifier(state_name, suffix),
-            drive_name=modifier(drive_name, suffix)
-        )
+        mod_integrator.state_components = mod_traj.components[modifier(state_name, suffix)]
+        mod_integrator.drive_components = mod_traj.components[modifier(drive_name, suffix)]
+        return mod_integrator
     elseif integrator isa DerivativeIntegrator
-        variable = get_component_names(traj, integrator.variable_components)
-        derivative = get_component_names(traj, integrator.derivative_components)
-        return integrator(
-            mod_traj,
-            variable=modifier(variable, suffix),
-            derivative=modifier(derivative, suffix)
-        )
+        var_name = get_component_names(traj, integrator.variable_components)
+        der_name = get_component_names(traj, integrator.derivative_components)
+        mod_integrator.variable_components = mod_traj.components[modifier(var_name, suffix)]
+        mod_integrator.derivative_components = mod_traj.components[modifier(der_name, suffix)]
+        return mod_integrator 
     else
         error("Integrator type not recognized")
     end
@@ -213,42 +200,41 @@ end
 
 function add_suffix(
     integrator::AbstractIntegrator,
-    sys::AbstractQuantumSystem,
+    suffix::String,
     traj::NamedTrajectory,
-    mod_traj::NamedTrajectory,
-    suffix::String
+    mod_traj::NamedTrajectory
 )
-    return modify_integrator_suffix(add_suffix, integrator, sys, traj, mod_traj, suffix)
+    return modify_integrator_suffix(integrator, add_suffix, suffix, traj, mod_traj)
 end
 
 function add_suffix(
     integrators::AbstractVector{<:AbstractIntegrator},
-    sys::AbstractQuantumSystem,
+    suffix::String,
     traj::NamedTrajectory,
-    mod_traj::NamedTrajectory,
-    suffix::String
+    mod_traj::NamedTrajectory
 )
-    return [add_suffix(intg, sys, traj, mod_traj, suffix) for intg in integrators]
+    return [
+        add_suffix(integrator, suffix, traj, mod_traj)
+            for integrator ∈ integrators
+    ]
 end
 
 function remove_suffix(
     integrator::AbstractIntegrator,
-    sys::AbstractQuantumSystem,
+    suffix::String,
     traj::NamedTrajectory,
-    mod_traj::NamedTrajectory,
-    suffix::String
+    mod_traj::NamedTrajectory
 )
-    return modify_integrator_suffix(remove_suffix, integrator, sys, traj, mod_traj, suffix)
+    return modify_integrator_suffix(integrator, remove_suffix, suffix, traj, mod_traj)
 end
 
 function remove_suffix(
     integrators::AbstractVector{<:AbstractIntegrator},
-    sys::AbstractQuantumSystem,
+    suffix::String,
     traj::NamedTrajectory,
-    mod_traj::NamedTrajectory,
-    suffix::String
+    mod_traj::NamedTrajectory
 )
-    return [remove_suffix(intg, sys, traj, mod_traj, suffix) for intg in integrators]
+    return [remove_suffix(intg, suffix, traj, mod_traj) for intg in integrators]
 end
 
 
@@ -261,11 +247,11 @@ Base.endswith(integrator::DerivativeIntegrator, suffix::String) = endswith(integ
 
 function Base.endswith(integrator::AbstractIntegrator, traj::NamedTrajectory, suffix::String)
     if integrator isa UnitaryExponentialIntegrator
-        name = get_component_names(traj, integrator.unitary_components)
+        name = get_component_names(traj, integrator.state_components)
     elseif integrator isa QuantumStateExponentialIntegrator
         name = get_component_names(traj, integrator.state_components)
     elseif integrator isa UnitaryPadeIntegrator
-        name = get_component_names(traj, integrator.unitary_components)
+        name = get_component_names(traj, integrator.state_components)
     elseif integrator isa QuantumStatePadeIntegrator
         name = get_component_names(traj, integrator.state_components)
     elseif integrator isa DerivativeIntegrator
@@ -408,24 +394,23 @@ end
     H_drift = 0.01 * GATES[:Z]
     H_drives = [GATES[:X], GATES[:Y]]
     T = 50
-    sys = QuantumSystem(H_drift, H_drives, params=Dict(:T=>T))
+    sys_1 = QuantumSystem(H_drift, H_drives)
+    sys_2 = deepcopy(sys_1)
 
-    # apply suffix and sum
-    sys2 = direct_sum(
-        add_suffix(sys, "_1"),
-        add_suffix(sys, "_2")
-    )
+    # direct sum of systems
+    sys_sum = direct_sum(sys_1, sys_2)
+    @info sys_sum.n_drives
+
+
+    @test sys_sum.levels == sys_1.levels * 2
+    @test isempty(symdiff(keys(sys_sum.params), [:system_1, :system_2]))
+
+    sys_sum_2 = direct_sum(sys_sum, deepcopy(sys_1))
 
-    @test length(sys2.H_drives) == 4
-    @test sys2.params[:T_1] == T
-    @test sys2.params[:T_2] == T
+    @test sys_sum_2.levels == sys_1.levels * 3
+    display(sys_sum_2.params)
+    @test isempty(symdiff(keys(sys_sum_2.params), [:system_1, :system_2, :system_3]))
 
-    # add another system
-    sys = QuantumSystem(H_drift, H_drives, params=Dict(:T=>T, :S=>2T))
-    sys3 = direct_sum(sys2, add_suffix(sys, "_3"))
-    @test length(sys3.H_drives) == 6
-    @test sys3.params[:T_3] == T
-    @test sys3.params[:S_3] == 2T
 end
 
 # TODO: fix broken test
diff --git a/src/dynamics.jl b/src/dynamics.jl
deleted file mode 100644
index c81411b0..00000000
--- a/src/dynamics.jl
+++ /dev/null
@@ -1,400 +0,0 @@
-module Dynamics
-
-export AbstractDynamics
-export QuantumDynamics
-
-export dynamics
-export dynamics_jacobian
-export dynamics_hessian_of_lagrangian
-export dynamics_components
-
-using ..QuantumSystems
-using ..Isomorphisms
-using ..StructureUtils
-using ..Integrators
-
-using TrajectoryIndexingUtils
-using NamedTrajectories
-using LinearAlgebra
-using SparseArrays
-using ForwardDiff
-
-
-abstract type AbstractDynamics end
-
-"""
-    QuantumDynamics <: AbstractDynamics
-"""
-struct QuantumDynamics <: AbstractDynamics
-    integrators::Union{Nothing, Vector{<:AbstractIntegrator}}
-    F::Function
-    ∂F::Function
-    ∂F_structure::Vector{Tuple{Int, Int}}
-    μ∂²F::Union{Function, Nothing}
-    μ∂²F_structure::Union{Vector{Tuple{Int, Int}}, Nothing}
-    dim::Int
-end
-
-function dynamics_components(integrators::Vector{<:AbstractIntegrator})
-    dynamics_comps = []
-    comp_mark = 0
-    for integrator ∈ integrators
-        integrator_comps = (comp_mark + 1):(comp_mark + integrator.dim)
-        push!(dynamics_comps, integrator_comps)
-        comp_mark += integrator.dim
-    end
-    return dynamics_comps
-end
-
-function dynamics(
-    integrators::Vector{<:AbstractIntegrator},
-    traj::NamedTrajectory
-)
-    dynamics_comps = dynamics_components(integrators)
-    dynamics_dim = sum(integrator.dim for integrator ∈ integrators)
-    function f(zₜ, zₜ₊₁, t)
-        δ = Vector{eltype(zₜ)}(undef, dynamics_dim)
-        for (integrator, integrator_comps) ∈ zip(integrators, dynamics_comps)
-            δ[integrator_comps] = integrator(zₜ, zₜ₊₁, t)
-        end
-        return δ
-    end
-    return f
-end
-
-
-
-function dynamics_jacobian(
-    integrators::Vector{<:AbstractIntegrator},
-    traj::NamedTrajectory
-)
-    dynamics_comps = dynamics_components(integrators)
-    dynamics_dim = sum(integrator.dim for integrator ∈ integrators)
-    @views function ∂f(zₜ, zₜ₊₁, t)
-        ∂ = zeros(eltype(zₜ), dynamics_dim, 2traj.dim)
-        for (integrator, integrator_comps) ∈ zip(integrators, dynamics_comps)
-            if integrator isa QuantumIntegrator
-                if integrator isa QuantumPadeIntegrator && integrator.autodiff
-                    ∂Pᵢ(z1, z2, t) = ForwardDiff.jacobian(
-                        zz -> integrator(zz[1:traj.dim], zz[traj.dim+1:end], t),
-                        [z1; z2]
-                    )
-                    ∂[integrator_comps, 1:2traj.dim] = ∂Pᵢ(zₜ, zₜ₊₁, t)
-                else
-                    if integrator.freetime
-                        x_comps, u_comps, Δt_comps = get_comps(integrator, traj)
-                        ∂xₜf, ∂xₜ₊₁f, ∂uₜf, ∂Δtₜf =
-                            Integrators.jacobian(integrator, zₜ, zₜ₊₁, t)
-                        ∂[integrator_comps, Δt_comps] = ∂Δtₜf
-                    else
-                        x_comps, u_comps = get_comps(integrator, traj)
-                        ∂xₜf, ∂xₜ₊₁f, ∂uₜf =
-                            Integrators.jacobian(integrator, zₜ, zₜ₊₁, t)
-                    end
-                    ∂[integrator_comps, x_comps] = ∂xₜf
-                    ∂[integrator_comps, x_comps .+ traj.dim] = ∂xₜ₊₁f
-                    if u_comps isa Tuple
-                        for (uᵢ_comps, ∂uₜᵢf) ∈ zip(u_comps, ∂uₜf)
-                            ∂[integrator_comps, uᵢ_comps] = ∂uₜᵢf
-                        end
-                    else
-                        ∂[integrator_comps, u_comps] = ∂uₜf
-                    end
-                end
-            elseif integrator isa DerivativeIntegrator
-                if integrator.freetime
-                    x_comps, dx_comps, Δt_comps = get_comps(integrator, traj)
-                    ∂xₜf, ∂xₜ₊₁f, ∂dxₜf, ∂Δtₜf =
-                        Integrators.jacobian(integrator, zₜ, zₜ₊₁, t)
-                else
-                    x_comps, dx_comps = get_comps(integrator, traj)
-                    ∂xₜf, ∂xₜ₊₁f, ∂dxₜf =
-                        Integrators.jacobian(integrator, zₜ, zₜ₊₁, t)
-                end
-                ∂[integrator_comps, x_comps] = ∂xₜf
-                ∂[integrator_comps, x_comps .+ traj.dim] = ∂xₜ₊₁f
-                ∂[integrator_comps, dx_comps] = ∂dxₜf
-                if integrator.freetime
-                    ∂[integrator_comps, Δt_comps] = ∂Δtₜf
-                end
-            else
-                error("integrator type not supported: $(typeof(integrator))")
-            end
-        end
-        return sparse(∂)
-    end
-    return ∂f
-end
-
-function dynamics_hessian_of_lagrangian(
-    integrators::Vector{<:AbstractIntegrator},
-    traj::NamedTrajectory
-)
-    dynamics_comps = dynamics_components(integrators)
-    free_time = traj.timestep isa Symbol
-    function μ∂²f(zₜ, zₜ₊₁, μₜ)
-        μ∂² = zeros(eltype(zₜ), 2traj.dim, 2traj.dim)
-        for (integrator, integrator_comps) ∈ zip(integrators, dynamics_comps)
-            if integrator isa QuantumIntegrator
-                if integrator.autodiff
-                    μ∂²P(z1, z2, μ) = ForwardDiff.hessian(
-                        zz -> μ' * integrator(zz[1:traj.dim], zz[traj.dim+1:end], traj),
-                        [z1; z2]
-                    )
-                    μ∂²[1:2traj.dim, 1:2traj.dim] = sparse(μ∂²P(zₜ, zₜ₊₁, μₜ[integrator_comps]))
-                else
-                    if free_time
-                        x_comps, u_comps, Δt_comps = comps(integrator, traj)
-                        μ∂uₜ∂xₜf, μ∂²uₜf, μ∂Δtₜ∂xₜf, μ∂Δtₜ∂uₜf, μ∂²Δtₜf, μ∂xₜ₊₁∂uₜf, μ∂xₜ₊₁∂Δtₜf =
-                            hessian_of_the_lagrangian(integrator, zₜ, zₜ₊₁, μₜ[integrator_comps], traj)
-                    else
-                        x_comps, u_comps = comps(integrator, traj)
-                        μ∂uₜ∂xₜf, μ∂²uₜf, μ∂xₜ₊₁∂uₜf =
-                            hessian_of_the_lagrangian(integrator, zₜ, zₜ₊₁, μₜ[integrator_comps], traj)
-                    end
-                    if u_comps isa Tuple
-                        for (uᵢ_comps, μ∂uₜᵢ∂xₜf) ∈ zip(u_comps, μ∂uₜ∂xₜf)
-                            μ∂²[x_comps, uᵢ_comps] += μ∂uₜᵢ∂xₜf
-                        end
-                        for (uᵢ_comps, μ∂²uₜᵢf) ∈ zip(u_comps, μ∂²uₜf)
-                            μ∂²[uᵢ_comps, uᵢ_comps] += μ∂²uₜᵢf
-                        end
-                        if free_time
-                            for (uᵢ_comps, μ∂Δtₜ∂uₜᵢf) ∈ zip(u_comps, μ∂Δtₜ∂uₜf)
-                                μ∂²[uᵢ_comps, Δt_comps] += μ∂Δtₜ∂uₜᵢf
-                            end
-                        end
-                        for (uᵢ_comps, μ∂xₜ₊₁∂uₜᵢf) ∈ zip(u_comps, μ∂xₜ₊₁∂uₜf)
-                            μ∂²[uᵢ_comps, x_comps .+ traj.dim] += μ∂xₜ₊₁∂uₜᵢf
-                        end
-                    else
-                        μ∂²[x_comps, u_comps] += μ∂uₜ∂xₜf
-                        μ∂²[u_comps, u_comps] += μ∂²uₜf
-                        if free_time
-                            μ∂²[u_comps, Δt_comps] += μ∂Δtₜ∂uₜf
-                        end
-                        μ∂²[u_comps, x_comps .+ traj.dim] += μ∂xₜ₊₁∂uₜf
-                    end
-                    if free_time
-                        μ∂²[x_comps, Δt_comps] += μ∂Δtₜ∂xₜf
-                        μ∂²[Δt_comps, x_comps .+ traj.dim] += μ∂xₜ₊₁∂Δtₜf
-                        μ∂²[Δt_comps, Δt_comps] .+= μ∂²Δtₜf
-                    end
-                end
-            elseif integrator isa DerivativeIntegrator
-                if free_time
-                    x_comps, dx_comps, Δt_comps = comps(integrator, traj)
-                    μ∂dxₜ∂Δtₜf = -μₜ[integrator_comps]
-                    μ∂²[dx_comps, Δt_comps] += μ∂dxₜ∂Δtₜf
-                end
-            end
-        end
-        return sparse(μ∂²)
-    end
-    return μ∂²f
-end
-
-
-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
-
-    # 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)
-
-    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)
-    else
-        ∂f_structure, ∂F_structure = dynamics_structure(∂f, traj, dynamics_dim;
-            verbose=verbose,
-            jacobian=jacobian_structure,
-        )
-        μ∂²F_structure = nothing
-    end
-
-    ∂f_nnz = length(∂f_structure)
-
-    if verbose
-        println("        constructing full dynamics derivative functions...")
-    end
-
-    @views function F(Z⃗::AbstractVector{R}) where R <: Real
-        δ = Vector{R}(undef, dynamics_dim * (traj.T - 1))
-        Threads.@threads for t = 1:traj.T-1
-            zₜ = Z⃗[slice(t, traj.dim)]
-            zₜ₊₁ = Z⃗[slice(t + 1, traj.dim)]
-            δ[slice(t, dynamics_dim)] = f(zₜ, zₜ₊₁, t)
-        end
-        return δ
-    end
-
-    @views function ∂F(Z⃗::AbstractVector{R}) where R <: Real
-        ∂s = zeros(R, length(∂F_structure))
-        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)
-                ∂s[index(t, k, ∂f_nnz)] = ∂fₜ[i, j]
-            end
-        end
-        return ∂s
-    end
-
-    if eval_hessian
-        @views μ∂²F = (Z⃗::AbstractVector{<:Real}, μ⃗::AbstractVector{<:Real}) -> begin
-            μ∂²s = Vector{eltype(Z⃗)}(undef, length(μ∂²F_structure))
-            Threads.@threads for t = 1:traj.T-1
-                zₜ = Z⃗[slice(t, traj.dim)]
-                zₜ₊₁ = Z⃗[slice(t + 1, traj.dim)]
-                μₜ = μ⃗[slice(t, dynamics_dim)]
-                μₜ∂²fₜ = μ∂²f(zₜ, zₜ₊₁, μₜ)
-                for (k, (i, j)) ∈ enumerate(μ∂²f_structure)
-                    μ∂²s[index(t, k, μ∂²f_nnz)] = μₜ∂²fₜ[i, j]
-                end
-            end
-            return μ∂²s
-        end
-    else
-        μ∂²F = nothing
-    end
-
-    return QuantumDynamics(
-        integrators,
-        F,
-        ∂F,
-        ∂F_structure,
-        μ∂²F,
-        μ∂²F_structure,
-        dynamics_dim
-    )
-end
-
-QuantumDynamics(P::AbstractIntegrator, traj::NamedTrajectory; kwargs...) =
-    QuantumDynamics([P], traj; kwargs...)
-
-function QuantumDynamics(
-    f::Function,
-    traj::NamedTrajectory;
-    verbose=false,
-)
-    dynamics_dim = length(f(traj[1].data, traj[2].data))
-
-    @views function F(Z⃗::AbstractVector{<:Real})
-        r = zeros(eltype(Z⃗), dynamics_dim * (traj.T - 1))
-        Threads.@threads for t = 1:traj.T-1
-            zₜ = Z⃗[slice(t, traj.dim)]
-            zₜ₊₁ = Z⃗[slice(t + 1, traj.dim)]
-            r[slice(t, dynamics_dim)] = f(zₜ, zₜ₊₁)
-        end
-        return r
-    end
-
-    # TODO: benchmark Zygote vs ForwardDiff for jacobian
-    # function ∂f(zₜ, zₜ₊₁)
-    #     ∂zₜf, ∂zₜ₊₁f = Zygote.jacobian(f, zₜ, zₜ₊₁)
-    #     ∂fₜ = hcat(∂zₜf, ∂zₜ₊₁f)
-    #     return ∂fₜ
-    # end
-
-    @views f̂(zz) = f(zz[1:traj.dim], zz[traj.dim+1:end])
-
-    ∂f̂(zz) = ForwardDiff.jacobian(f̂, zz)
-
-    ∂f_structure, ∂F_structure, μ∂²f_structure, μ∂²F_structure =
-        dynamics_structure(∂f̂, traj, dynamics_dim)
-
-    ∂f(zₜ, zₜ₊₁) = ∂f̂([zₜ; zₜ₊₁])
-
-    ∂f_nnz = length(∂f_structure)
-
-    @views function ∂F(Z⃗::AbstractVector{R}) where R <: Real
-        ∂ = zeros(R, length(∂F_structure))
-        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ₜ₊₁)
-            for (k, (i, j)) ∈ enumerate(∂f_structure)
-                ∂[index(t, k, ∂f_nnz)] = ∂fₜ[i, j]
-            end
-        end
-        return ∂
-    end
-
-    μf̂(zz, μ) = dot(μ, f̂(zz))
-
-    @views function μ∂²f̂(zₜzₜ₊₁, μₜ)
-        return ForwardDiff.hessian(zz -> μf̂(zz, μₜ), zₜzₜ₊₁)
-    end
-
-    μ∂²f_nnz = length(μ∂²f_structure)
-
-    @views function μ∂²F(Z⃗::AbstractVector{R}, μ::AbstractVector{R}) where R <: Real
-        μ∂² = zeros(R, length(μ∂²F_structure))
-        Threads.@threads for t = 1:traj.T-1
-            zₜzₜ₊₁ = Z⃗[slice(t:t+1, traj.dim)]
-            μₜ = μ[slice(t, dynamics_dim)]
-            μ∂²fₜ = μ∂²f̂(zₜzₜ₊₁, μₜ)
-            for (k, (i, j)) ∈ enumerate(μ∂²f_structure)
-                μ∂²[index(t, k, μ∂²f_nnz)] = μ∂²fₜ[i, j]
-            end
-        end
-        return μ∂²
-    end
-
-    return QuantumDynamics(nothing, F, ∂F, ∂F_structure, μ∂²F, μ∂²F_structure, dynamics_dim)
-end
-
-end
diff --git a/src/embedded_operators.jl b/src/embedded_operators.jl
deleted file mode 100644
index 80d7a2bd..00000000
--- a/src/embedded_operators.jl
+++ /dev/null
@@ -1,468 +0,0 @@
-module EmbeddedOperators
-
-export OperatorType
-export EmbeddedOperator
-
-export embed
-export unembed
-export get_subspace_indices
-export get_subspace_enr_indices
-export get_subspace_leakage_indices
-export get_iso_vec_leakage_indices
-export get_iso_vec_subspace_indices
-export get_subspace_identity
-
-using LinearAlgebra
-using TestItemRunner
-
-using TrajectoryIndexingUtils
-
-using ..Isomorphisms
-using ..QuantumObjectUtils
-using ..QuantumSystems
-# using ..QuantumSystemUtils
-
-@doc raw"""
-    embed(matrix::Matrix{ComplexF64}, subspace_indices::AbstractVector{Int}, levels::Int)
-
-Embed an operator $U$ in the subspace of a larger system $\mathcal{X} = \mathcal{X}_{\text{subspace}} \oplus \mathcal{X}_{\text{leakage}}$ which is composed of matrices of size $\text{levels} \times \text{levels}$.
-
-# Arguments
-- `matrix::Matrix{ComplexF64}`: Operator to embed.
-- `subspace_indices::AbstractVector{Int}`: Indices of the subspace to embed the operator in.
-- `levels::Int`: Total number of levels in the system.
-"""
-function embed(op::Matrix{ComplexF64}, subspace_indices::AbstractVector{Int}, levels::Int)
-    @assert size(op, 1) == size(op, 2) "Operator must be square."
-    op_embedded = zeros(ComplexF64, levels, levels)
-    op_embedded[subspace_indices, subspace_indices] = op
-    return op_embedded
-end
-
-@doc raw"""
-    unembed(matrix::AbstractMatrix, subspace_indices::AbstractVector{Int})
-
-Unembed an operator $U$ from a subspace of a larger system $\mathcal{X} = \mathcal{X}_{\text{subspace}} \oplus \mathcal{X}_{\text{leakage}}$ which is composed of matrices of size $\text{levels} \times \text{levels}$.
-
-This is equivalent to calling `matrix[subspace_indices, subspace_indices]`.
-
-# Arguments
-- `matrix::AbstractMatrix`: Operator to unembed.
-- `subspace_indices::AbstractVector{Int}`: Indices of the subspace to unembed the operator from.
-"""
-function unembed(matrix::AbstractMatrix, subspace_indices::AbstractVector{Int})
-    return matrix[subspace_indices, subspace_indices]
-end
-
-# ----------------------------------------------------------------------------- #
-#                             Embedded Operator                                 #
-# ----------------------------------------------------------------------------- #
-
-"""
-    EmbeddedOperator
-
-Embedded operator type to represent an operator embedded in a subspace of a larger quantum system.
-
-# Fields
-- `operator::Matrix{ComplexF64}`: Embedded operator of size `prod(subsystem_levels) x prod(subsystem_levels)`.
-- `subspace_indices::Vector{Int}`: Indices of the subspace the operator is embedded in.
-- `subsystem_levels::Vector{Int}`: Levels of the subsystems in the composite system.
-"""
-struct EmbeddedOperator
-    operator::Matrix{ComplexF64}
-    subspace_indices::Vector{Int}
-    subsystem_levels::Vector{Int}
-
-    @doc raw"""
-        EmbeddedOperator(op::Matrix{<:Number}, subspace_indices::AbstractVector{Int}, subsystem_levels::AbstractVector{Int})
-
-    Create an embedded operator. The operator `op` is embedded in the subspace defined by `subspace_indices` in `subsystem_levels`.
-
-    # Arguments
-    - `op::Matrix{<:Number}`: Operator to embed.
-    - `subspace_indices::AbstractVector{Int}`: Indices of the subspace to embed the operator in. e.g. `get_subspace_indices([1:2, 1:2], [3, 3])`.
-    - `subsystem_levels::AbstractVector{Int}`: Levels of the subsystems in the composite system. e.g. `[3, 3]` for two 3-level systems.
-    """
-    function EmbeddedOperator(
-        op::Matrix{<:Number},
-        subspace_indices::AbstractVector{Int},
-        subsystem_levels::AbstractVector{Int}
-    )
-
-        op_embedded = embed(Matrix{ComplexF64}(op), subspace_indices, prod(subsystem_levels))
-        return new(op_embedded, subspace_indices, subsystem_levels)
-    end
-end
-
-const OperatorType = Union{AbstractMatrix{<:Number}, EmbeddedOperator}
-
-EmbeddedOperator(op::Matrix{<:Number}, subspace_indices::AbstractVector{Int}, levels::Int) =
-    EmbeddedOperator(op, subspace_indices, [levels])
-
-function embed(matrix::Matrix{ComplexF64}, op::EmbeddedOperator)
-    return embed(matrix, op.subspace_indices, prod(op.subsystem_levels))
-end
-
-function unembed(op::EmbeddedOperator)::Matrix{ComplexF64}
-    return op.operator[op.subspace_indices, op.subspace_indices]
-end
-
-function unembed(matrix::AbstractMatrix, op::EmbeddedOperator)
-    return matrix[op.subspace_indices, op.subspace_indices]
-end
-
-Base.size(op::EmbeddedOperator) = size(op.operator)
-Base.size(op::EmbeddedOperator, dim::Union{Int, Nothing}) = size(op.operator, dim)
-
-function Base.:*(
-    op1::EmbeddedOperator,
-    op2::EmbeddedOperator
-)
-    @assert size(op1) == size(op2) "Operators must be of the same size."
-    @assert op1.subspace_indices == op2.subspace_indices "Operators must have the same subspace."
-    @assert op1.subsystem_levels == op2.subsystem_levels "Operators must have the same subsystem levels."
-    return EmbeddedOperator(
-        unembed(op1) * unembed(op2),
-        op1.subspace_indices,
-        op1.subsystem_levels
-    )
-end
-
-function Base.kron(op1::EmbeddedOperator, op2::EmbeddedOperator)
-    levels = [size(op1, 1), size(op2, 2)]
-    indices = get_subspace_indices(
-        [op1.subspace_indices, op2.subspace_indices], levels
-    )
-    return EmbeddedOperator(unembed(op1) ⊗ unembed(op2), indices, levels)
-end
-
-QuantumObjectUtils.:⊗(A::EmbeddedOperator, B::EmbeddedOperator) = kron(A, B)
-
-function EmbeddedOperator(
-    op::AbstractMatrix{<:Number},
-    system::QuantumSystem;
-    subspace=1:size(op, 1)
-)
-    return EmbeddedOperator(
-        op,
-        get_subspace_indices(subspace, system.levels),
-        [system.levels]
-    )
-end
-
-function EmbeddedOperator(
-    op::AbstractMatrix{<:Number},
-    csystem::CompositeQuantumSystem,
-    op_subsystem_indices::AbstractVector{Int};
-    subspaces=fill(1:2, length(csystem.subsystems)),
-)
-    @assert all(diff(op_subsystem_indices) .== 1) "op_subsystem_indices must be consecutive (for now)."
-
-    if size(op, 1) == prod(length.(subspaces[op_subsystem_indices]))
-        Is = Matrix{ComplexF64}.(I.(length.(subspaces)))
-        Is[op_subsystem_indices[1]] = op
-        deleteat!(Is, op_subsystem_indices[2:end])
-        op = kron(Is...)
-    else
-        @assert(
-            size(op, 1) == prod(length.(subspaces)),
-            """\n
-                Operator size ($(size(op, 1))) must match product of subsystem subspaces ($(prod(length.(subspaces)))).
-            """
-        )
-    end
-
-    subspace_indices = get_subspace_indices(subspaces, csystem.subsystem_levels)
-
-    return EmbeddedOperator(
-        op,
-        subspace_indices,
-        csystem.subsystem_levels
-    )
-end
-
-function EmbeddedOperator(
-    op::AbstractMatrix{<:Number},
-    csystem::CompositeQuantumSystem,
-    op_subsystem_index::Int;
-    kwargs...
-)
-    return EmbeddedOperator(
-        op,
-        csystem,
-        [op_subsystem_index];
-        kwargs...
-    )
-end
-
-function EmbeddedOperator(op::Symbol, args...; kwargs...)
-    @assert op ∈ keys(GATES) "Operator must be a valid gate. See QuantumCollocation.QuantumObjectUtils.GATES dict for available gates."
-    return EmbeddedOperator(GATES[op], args...; kwargs...)
-end
-
-function EmbeddedOperator(
-    ops::AbstractVector{Symbol},
-    sys::CompositeQuantumSystem,
-    op_indices::AbstractVector{Int}
-)
-    ops_embedded = [
-        EmbeddedOperator(op, sys, op_indices[i])
-            for (op, i) ∈ zip(ops, op_indices)
-    ]
-    return *(ops_embedded...)
-end
-
-# ----------------------------------------------------------------------------- #
-#                            Subspace Indices                                   #
-# ----------------------------------------------------------------------------- #
-
-basis_labels(subsystem_levels::AbstractVector{Int}; baseline=1) =
-    kron([""], [string.(baseline:levels - 1 + baseline) for levels ∈ subsystem_levels]...)
-
-basis_labels(subsystem_level::Int; kwargs...) = basis_labels([subsystem_level]; kwargs...)
-
-"""
-    get_subspace_indices(subspaces::Vector{<:AbstractVector{Int}}, subsystem_levels::AbstractVector{Int})
-
-Get the indices for the subspace of composite quantum system.
-
-Example: for the two-qubit subspace of two 3-level systems:
-```julia
-subspaces = [1:2, 1:2]
-subsystem_levels = [3, 3]
-get_subspace_indices(subspaces, subsystem_levels) == [1, 2, 4, 5]
-```
-
-# Arguments
-
-- `subspaces::Vector{<:AbstractVector{Int}}`: Subspaces to get indices for. e.g. `[1:2, 1:2]`.
-- `subsystem_levels::AbstractVector{Int}`: Levels of the subsystems in the composite system. e.g. `[3, 3]`. Each element corresponds to a subsystem.
-"""
-function get_subspace_indices(
-    subspaces::Vector{<:AbstractVector{Int}},
-    subsystem_levels::AbstractVector{Int}
-)
-    @assert length(subspaces) == length(subsystem_levels)
-    return findall(
-        b -> all(l ∈ subspaces[i] for (i, l) ∈ enumerate([parse(Int, bᵢ) for bᵢ ∈ b])),
-        basis_labels(subsystem_levels, baseline=1)
-    )
-end
-
-"""
-    get_subspace_indices(subspace::AbstractVector{Int}, levels::Int)
-
-Get the indices for the subspace of simple, non-composite, quantum system. For example:
-```julia
-get_subspace_indices([1, 2], 3) == [1, 2]
-```
-
-# Arguments
-- `subspace::AbstractVector{Int}`: Subspace to get indices for. e.g. `[1, 2]`.
-- `levels::Int`: Levels of the subsystem. e.g. `3`.
-"""
-get_subspace_indices(subspace::AbstractVector{Int}, levels::Int) =
-    get_subspace_indices([subspace], [levels])
-
-"""
-    get_subspace_indices(levels::AbstractVector{Int}; subspace=1:2, kwargs...)
-
-Get the indices for the subspace of composite quantum system. This is a convenience function that allows to specify the subspace as a range that is constant for every subsystem, which defaults to `1:2`, that is qubit systems.
-
-# Arguments
-- `levels::AbstractVector{Int}`: Levels of the subsystems in the composite system. e.g. `[3, 3]`.
-
-# Keyword Arguments
-- `subspace::AbstractVector{Int}`: Subspace to get indices for. e.g. `1:2`.
-"""
-get_subspace_indices(levels::AbstractVector{Int}; subspace=1:2) =
-    get_subspace_indices(fill(subspace, length(levels)), levels)
-
-function get_subspace_enr_indices(excitation_restriction::Int, subsystem_levels::AbstractVector{Int})
-    # excitation_number uses baseline of zero
-    return findall(
-        b -> sum([parse(Int, bᵢ) for bᵢ ∈ b]) ≤ excitation_restriction,
-        basis_labels(subsystem_levels, baseline=0)
-    )
-end
-
-function get_subspace_leakage_indices(
-    subspaces::Vector{<:AbstractVector{Int}},
-    subsystem_levels::AbstractVector{Int};
-)
-    subspace_indices = get_subspace_indices(subspaces, subsystem_levels)
-    return get_subspace_leakage_indices(subspace_indices)
-end
-
-get_subspace_leakage_indices(subspace_indices::AbstractVector{Int}, levels::Int) =
-    setdiff(1:levels, subspace_indices)
-
-get_subspace_leakage_indices(op::EmbeddedOperator) =
-    get_subspace_leakage_indices(op.subspace_indices, size(op)[1])
-
-get_iso_vec_subspace_indices(op::EmbeddedOperator) =
-    get_iso_vec_subspace_indices(op.subspace_indices, op.subsystem_levels)
-
-get_iso_vec_leakage_indices(op::EmbeddedOperator) =
-    get_iso_vec_leakage_indices(op.subspace_indices, op.subsystem_levels)
-
-function get_iso_vec_subspace_indices(
-    subspace_indices::AbstractVector{Int},
-    subsystem_levels::AbstractVector{Int}
-)
-    N = prod(subsystem_levels)
-    iso_subspace_indices = Int[]
-    for sⱼ ∈ subspace_indices
-        for sᵢ ∈ subspace_indices
-            push!(iso_subspace_indices, index(sⱼ, sᵢ, 2N))
-        end
-        for sᵢ ∈ subspace_indices
-            push!(iso_subspace_indices, index(sⱼ, sᵢ + N, 2N))
-        end
-    end
-    return iso_subspace_indices
-end
-
-function get_iso_vec_leakage_indices(
-    subspace_indices::AbstractVector{Int},
-    subsystem_levels::AbstractVector{Int}
-)
-    N = prod(subsystem_levels)
-    leakage_indices = get_subspace_leakage_indices(subspace_indices, N)
-    iso_leakage_indices = Int[]
-    for sⱼ ∈ subspace_indices
-        for lᵢ ∈ leakage_indices
-            push!(iso_leakage_indices, index(sⱼ, lᵢ, 2N))
-        end
-        for lᵢ ∈ leakage_indices
-            push!(iso_leakage_indices, index(sⱼ, lᵢ + N, 2N))
-        end
-    end
-    return iso_leakage_indices
-end
-
-function get_subspace_identity(op::EmbeddedOperator)
-    return embed(
-        Matrix{ComplexF64}(I(length(op.subspace_indices))),
-        op.subspace_indices,
-        size(op, 1)
-    )
-end
-
-# =========================================================================== #
-
-@testitem "Basis labels" begin
-    levels = [3, 3]
-    labels = ["11", "12", "13", "21", "22", "23", "31", "32", "33"]
-    @test EmbeddedOperators.basis_labels(levels, baseline=1) == labels
-
-    labels = ["1", "2", "3"]
-    @test EmbeddedOperators.basis_labels(3, baseline=1) == labels
-    @test EmbeddedOperators.basis_labels([3], baseline=1) == labels
-
-    labels = ["0", "1", "2"]
-    @test EmbeddedOperators.basis_labels(3, baseline=0) == labels
-    @test EmbeddedOperators.basis_labels([3], baseline=0) == labels
-
-    levels = [2, 2]
-    labels = ["00", "01", "10", "11"]
-    @test EmbeddedOperators.basis_labels(levels, baseline=0) == labels
-end
-
-@testitem "Subspace Indices" begin
-    @test get_subspace_indices([1, 2], 3) == [1, 2]
-    # 2 * 2 = 4 elements
-    @test get_subspace_indices([1:2, 1:2], [3, 3]) == [1, 2, 4, 5]
-    # 1 * 1 = 1 element
-    @test get_subspace_indices([[2], [2]], [3, 3]) == [5]
-    # 1 * 2 = 2 elements
-    @test get_subspace_indices([[2], 1:2], [3, 3]) == [4, 5]
-end
-
-@testitem "Subspace ENR Indices" begin
-    # 00, 01, 02x, 10, 11x, 12x, 20x, 21x, 22x
-    @test get_subspace_enr_indices(1, [3, 3]) == [1, 2, 4]
-    # 00, 01, 02, 10, 11, 12x, 20, 21x, 22x
-    @test get_subspace_enr_indices(2, [3, 3]) == [1, 2, 3, 4, 5, 7]
-    # 00, 01, 02, 10, 11, 12, 20, 21, 22x
-    @test get_subspace_enr_indices(3, [3, 3]) == [1, 2, 3, 4, 5, 6, 7, 8]
-    # 00, 01, 02, 10, 11, 12, 20, 21, 22
-    @test get_subspace_enr_indices(4, [3, 3]) == [1, 2, 3, 4, 5, 6, 7, 8, 9]
-end
-
-@testitem "Subspace Leakage Indices" begin
-    # TODO: Implement tests
-end
-
-@testitem "Embedded operator" begin
-    # Embed X
-    op = Matrix{ComplexF64}([0 1; 1 0])
-    embedded_op = Matrix{ComplexF64}([0 1 0 0; 1 0 0 0; 0 0 0 0; 0 0 0 0])
-    @test embed(op, 1:2, 4) == embedded_op
-    embedded_op_struct = EmbeddedOperator(op, 1:2, 4)
-    @test embedded_op_struct.operator == embedded_op
-    @test embedded_op_struct.subspace_indices == 1:2
-    @test embedded_op_struct.subsystem_levels == [4]
-
-    # Properties
-    @test size(embedded_op_struct) == size(embedded_op)
-    @test size(embedded_op_struct, 1) == size(embedded_op, 1)
-
-    # X^2 = I
-    x2 = (embedded_op_struct * embedded_op_struct).operator
-    id = get_subspace_identity(embedded_op_struct)
-    @test x2 == id
-
-    # Embed X twice
-    op2 = op ⊗ op
-    embedded_op2 = [
-        0  0  0  0  1  0  0  0  0;
-        0  0  0  1  0  0  0  0  0;
-        0  0  0  0  0  0  0  0  0;
-        0  1  0  0  0  0  0  0  0;
-        1  0  0  0  0  0  0  0  0;
-        0  0  0  0  0  0  0  0  0;
-        0  0  0  0  0  0  0  0  0;
-        0  0  0  0  0  0  0  0  0;
-        0  0  0  0  0  0  0  0  0
-    ]
-    subspace_indices = get_subspace_indices([1:2, 1:2], [3, 3])
-    @test embed(op2, subspace_indices, 9) == embedded_op2
-    embedded_op2_struct = EmbeddedOperator(op2, subspace_indices, [3, 3])
-    @test embedded_op2_struct.operator == embedded_op2
-    @test embedded_op2_struct.subspace_indices == subspace_indices
-    @test embedded_op2_struct.subsystem_levels == [3, 3]
-end
-
-@testitem "Embedded operator from system" begin
-    CZ = GATES[:CZ]
-    a = annihilate(3)
-    σ_x = a + a'
-    σ_y = -1im*(a - a')
-    system = QuantumSystem([σ_x ⊗ σ_x, σ_y ⊗ σ_y])
-
-    op_explicit_qubit = EmbeddedOperator(
-        CZ,
-        system,
-        subspace=get_subspace_indices([1:2, 1:2], [3, 3])
-    )
-    op_implicit_qubit = EmbeddedOperator(CZ, system)
-    # This does not work (implicit puts indicies in 1:4)
-    @test op_implicit_qubit.operator != op_explicit_qubit.operator
-    # But the ops are the same
-    @test unembed(op_explicit_qubit) == unembed(op_implicit_qubit)
-    @test unembed(op_implicit_qubit) == CZ
-end
-
-@testitem "Embedded operator from composite system" begin
-    @test_skip nothing
-end
-
-@testitem "Embedded operator kron" begin
-    Z = GATES[:Z]
-    Ẑ = EmbeddedOperator(Z, 1:2, [4])
-    @test unembed(Ẑ ⊗ Ẑ) == Z ⊗ Z
-end
-
-
-end
diff --git a/src/evaluators.jl b/src/evaluators.jl
deleted file mode 100644
index 7f84882c..00000000
--- a/src/evaluators.jl
+++ /dev/null
@@ -1,176 +0,0 @@
-module Evaluators
-
-export PicoEvaluator
-
-using ..QuantumSystems
-using ..Integrators
-using ..Dynamics
-using ..Objectives
-using ..Constraints
-
-using NamedTrajectories
-using MathOptInterface
-using LinearAlgebra
-const MOI = MathOptInterface
-
-mutable struct PicoEvaluator <: MOI.AbstractNLPEvaluator
-    trajectory::NamedTrajectory
-    objective::Objective
-    dynamics::QuantumDynamics
-    n_dynamics_constraints::Int
-    nonlinear_constraints::Vector{<:NonlinearConstraint}
-    n_nonlinear_constraints::Int
-    eval_hessian::Bool
-
-    function PicoEvaluator(
-        trajectory::NamedTrajectory,
-        objective::Objective,
-        dynamics::QuantumDynamics,
-        nonlinear_constraints::Vector{<:NonlinearConstraint};
-        eval_hessian::Bool=true
-    )
-        n_dynamics_constraints = dynamics.dim * (trajectory.T - 1)
-        n_nonlinear_constraints = sum(con.dim for con ∈ nonlinear_constraints; init=0)
-
-        return new(
-            trajectory,
-            objective,
-            dynamics,
-            n_dynamics_constraints,
-            nonlinear_constraints,
-            n_nonlinear_constraints,
-            eval_hessian
-        )
-    end
-end
-
-MOI.initialize(::PicoEvaluator, features) = nothing
-
-function MOI.features_available(evaluator::PicoEvaluator)
-    if evaluator.eval_hessian
-        return [:Grad, :Jac, :Hess]
-    else
-        return [:Grad, :Jac]
-    end
-end
-
-
-# objective and gradient
-
-@views function MOI.eval_objective(
-    evaluator::PicoEvaluator,
-    Z⃗::AbstractVector
-)
-    return evaluator.objective.L(Z⃗, evaluator.trajectory)
-end
-
-@views function MOI.eval_objective_gradient(
-    evaluator::PicoEvaluator,
-    ∇::AbstractVector,
-    Z⃗::AbstractVector
-)
-    ∇[:] = evaluator.objective.∇L(Z⃗, evaluator.trajectory)
-    return nothing
-end
-
-
-# constraints and Jacobian
-
-@views function MOI.eval_constraint(
-    evaluator::PicoEvaluator,
-    g::AbstractVector,
-    Z⃗::AbstractVector
-)
-    g[1:evaluator.n_dynamics_constraints] = evaluator.dynamics.F(Z⃗)
-    offset = evaluator.n_dynamics_constraints
-    for con ∈ evaluator.nonlinear_constraints
-        g[offset .+ (1:con.dim)] = con.g(Z⃗)
-        offset += con.dim
-    end
-    return nothing
-end
-
-function MOI.jacobian_structure(evaluator::PicoEvaluator)
-    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]
-        push!(nl_constraint_structure, con_structure)
-        row_offset += con.dim
-    end
-    return vcat(dynamics_structure, nl_constraint_structure...)
-end
-
-@views function MOI.eval_constraint_jacobian(
-    evaluator::PicoEvaluator,
-    J::AbstractVector,
-    Z⃗::AbstractVector
-)
-    ∂s_dynamics = evaluator.dynamics.∂F(Z⃗)
-    for (k, ∂ₖ) in enumerate(∂s_dynamics)
-        J[k] = ∂ₖ
-    end
-    offset = length(∂s_dynamics)
-    for con ∈ evaluator.nonlinear_constraints
-        ∂s_con = con.∂g(Z⃗)
-        for (k, ∂ₖ) in enumerate(∂s_con)
-            J[offset + k] = ∂ₖ
-        end
-        offset += length(∂s_con)
-    end
-    return nothing
-end
-
-
-# Hessian of the Lagrangian
-
-function MOI.hessian_lagrangian_structure(evaluator::PicoEvaluator)
-    objective_structure = evaluator.objective.∂²L_structure(evaluator.trajectory)
-    dynamics_structure = evaluator.dynamics.μ∂²F_structure
-    nl_constraint_structure = [con.μ∂²g_structure for con ∈ evaluator.nonlinear_constraints]
-    return vcat(objective_structure, dynamics_structure, nl_constraint_structure...)
-end
-
-@views function MOI.eval_hessian_lagrangian(
-    evaluator::PicoEvaluator,
-    H::AbstractVector{T},
-    Z⃗::AbstractVector{T},
-    σ::T,
-    μ::AbstractVector{T}
-) where T
-
-    σ∂²Ls = σ * evaluator.objective.∂²L(Z⃗, evaluator.trajectory)
-
-    for (k, σ∂²Lₖ) in enumerate(σ∂²Ls)
-        H[k] = σ∂²Lₖ
-    end
-
-    μ_dynamics = μ[1:evaluator.n_dynamics_constraints]
-
-    μ_offset = evaluator.n_dynamics_constraints
-
-    offset = length(evaluator.objective.∂²L_structure(evaluator.trajectory))
-
-    μ∂²Fs = evaluator.dynamics.μ∂²F(Z⃗, μ_dynamics)
-
-    for (k, μ∂²Fₖ) in enumerate(μ∂²Fs)
-        H[offset + k] = μ∂²Fₖ
-    end
-
-    offset += length(evaluator.dynamics.μ∂²F_structure)
-
-    for con ∈ evaluator.nonlinear_constraints
-        μ_con = μ[μ_offset .+ (1:con.dim)]
-        μ∂²gs = con.μ∂²g(Z⃗, μ_con)
-        for (k, μ∂²gₖ) in enumerate(μ∂²gs)
-            H[offset + k] = μ∂²gₖ
-        end
-        offset += length(μ∂²gs)
-        μ_offset += con.dim
-    end
-
-    return nothing
-end
-
-end
diff --git a/src/integrators/_integrator_utils.jl b/src/integrators/_integrator_utils.jl
deleted file mode 100644
index f91d99c0..00000000
--- a/src/integrators/_integrator_utils.jl
+++ /dev/null
@@ -1,113 +0,0 @@
-# G_bilinear(a) helper function
-function G_bilinear(
-    a::AbstractVector,
-    G_drift::AbstractMatrix,
-    G_drives::AbstractVector{<:AbstractMatrix}
-)
-    return G_drift + sum(a .* G_drives)
-end
-
-const Id2 = 1.0 * I(2)
-const Im2 = 1.0 * [0 -1; 1 0]
-
-anticomm(A::AbstractMatrix{R}, B::AbstractMatrix{R}) where R <: Number = A * B + B * A
-
-function anticomm(
-    A::AbstractMatrix{R},
-    Bs::AbstractVector{<:AbstractMatrix{R}}
-) where R <: Number
-    return [anticomm(A, B) for B in Bs]
-end
-
-function anticomm(
-    As::AbstractVector{<:AbstractMatrix{R}},
-    Bs::AbstractVector{<:AbstractMatrix{R}}
-) where R <: Number
-    @assert length(As) == length(Bs)
-    n = length(As)
-    anticomms = Matrix{Matrix{R}}(undef, n, n)
-    for i = 1:n
-        for j = 1:n
-            anticomms[i, j] = anticomm(As[i], Bs[j])
-        end
-    end
-    return anticomms
-end
-
-pade(n, k) = (factorial(n + k) // (factorial(n - k) * factorial(k) * 2^n))
-pade_coeffs(n) = [pade(n, k) for k = n:-1:0][2:end] // pade(n, n)
-
-@inline function operator(
-    a::AbstractVector{<:Real},
-    A_drift::Matrix{<:Real},
-    A_drives::Vector{<:Matrix{<:Real}}
-)
-    return A_drift + sum(a .* A_drives)
-end
-
-@inline function operator_anticomm_operator(
-    a::AbstractVector{<:Real},
-    A_drift_anticomm_B_drift::Matrix{<:Real},
-    A_drift_anticomm_B_drives::Vector{<:Matrix{<:Real}},
-    B_drift_anticomm_A_drives::Vector{<:Matrix{<:Real}},
-    A_drives_anticomm_B_drives::Matrix{<:Matrix{<:Real}},
-    n_drives::Int
-)
-    A_anticomm_B = A_drift_anticomm_B_drift
-    for i = 1:n_drives
-        aⁱ = a[i]
-        A_anticomm_B += aⁱ * A_drift_anticomm_B_drives[i]
-        A_anticomm_B += aⁱ * B_drift_anticomm_A_drives[i]
-        A_anticomm_B += aⁱ^2 * A_drives_anticomm_B_drives[i, i]
-        for j = i+1:n_drives
-            aʲ = a[j]
-            A_anticomm_B += 2 * aⁱ * aʲ * A_drives_anticomm_B_drives[i, j]
-        end
-    end
-    return A_anticomm_B
-end
-
-@inline function operator_anticomm_term(
-    a::AbstractVector{<:Real},
-    A_drift_anticomm_B_drives::Vector{<:Matrix{<:Real}},
-    A_drives_anticomm_B_drives::Matrix{<:Matrix{<:Real}},
-    n_drives::Int,
-    j::Int
-)
-    A_anticomm_Bⱼ = A_drift_anticomm_B_drives[j]
-    for i = 1:n_drives
-        aⁱ = a[i]
-        A_anticomm_Bⱼ += aⁱ * A_drives_anticomm_B_drives[i, j]
-    end
-    return A_anticomm_Bⱼ
-end
-
-function build_anticomms(
-    G_drift::AbstractMatrix{R},
-    G_drives::Vector{<:AbstractMatrix{R}},
-    n_drives::Int) where R <: Number
-
-    drive_anticomms = fill(
-            zeros(size(G_drift)),
-            n_drives,
-            n_drives
-        )
-
-        for j = 1:n_drives
-            for k = 1:j
-                if k == j
-                    drive_anticomms[k, k] = 2 * G_drives[k]^2
-                else
-                    drive_anticomms[k, j] =
-                        anticomm(G_drives[k], G_drives[j])
-                end
-            end
-        end
-
-        drift_anticomms = [
-            anticomm(G_drive, G_drift)
-                for G_drive in G_drives
-        ]
-
-    return Symmetric(drive_anticomms), drift_anticomms
-end
diff --git a/src/integrators/_integrators.jl b/src/integrators/_integrators.jl
deleted file mode 100644
index d90f1fc8..00000000
--- a/src/integrators/_integrators.jl
+++ /dev/null
@@ -1,67 +0,0 @@
-module Integrators
-
-export AbstractIntegrator
-
-export QuantumIntegrator
-
-export QuantumPadeIntegrator
-export QuantumStatePadeIntegrator
-export UnitaryPadeIntegrator
-export UnitaryExponentialIntegrator
-export QuantumStateExponentialIntegrator
-
-export DerivativeIntegrator
-
-export jacobian
-export hessian_of_the_lagrangian
-
-export get_comps
-
-export nth_order_pade
-export fourth_order_pade
-export sixth_order_pade
-export eighth_order_pade
-export tenth_order_pade
-
-using ..QuantumSystems
-using ..Isomorphisms
-using ..QuantumObjectUtils
-using ..Losses
-using ..QuantumSystemUtils
-
-using NamedTrajectories
-using TrajectoryIndexingUtils
-using LinearAlgebra
-using SparseArrays
-using ForwardDiff
-using TestItemRunner
-
-abstract type AbstractIntegrator end
-
-abstract type QuantumIntegrator <: AbstractIntegrator end
-
-function comps(P::AbstractIntegrator, traj::NamedTrajectory)
-    state_comps = traj.components[state(P)]
-    u = controls(P)
-    if u isa Tuple
-        control_comps = Tuple(traj.components[uᵢ] for uᵢ ∈ u)
-    else
-        control_comps = traj.components[u]
-    end
-    if traj.timestep isa Float64
-        return state_comps, control_comps
-    else
-        timestep_comp = traj.components[traj.timestep]
-        return state_comps, control_comps, timestep_comp
-    end
-end
-
-
-
-include("_integrator_utils.jl")
-include("derivative_integrator.jl")
-include("pade_integrators.jl")
-include("exponential_integrators.jl")
-
-
-end
diff --git a/src/integrators/derivative_integrator.jl b/src/integrators/derivative_integrator.jl
deleted file mode 100644
index e8a5986a..00000000
--- a/src/integrators/derivative_integrator.jl
+++ /dev/null
@@ -1,96 +0,0 @@
-
-###
-### Derivative Integrator
-###
-
-mutable struct DerivativeIntegrator <: AbstractIntegrator
-    variable_components::Vector{Int}
-    derivative_components::Vector{Int}
-    freetime::Bool
-    timestep::Union{Float64, Int} # timestep of index in z
-    zdim::Int
-    dim::Int
-    autodiff::Bool
-end
-
-function DerivativeIntegrator(
-    variable::Symbol,
-    derivative::Symbol,
-    traj::NamedTrajectory
-)
-    freetime = traj.timestep isa Symbol
-    if freetime
-        timestep = traj.components[traj.timestep][1]
-    else
-        timestep = traj.timestep
-    end
-    return DerivativeIntegrator(
-        traj.components[variable],
-        traj.components[derivative],
-        freetime,
-        timestep,
-        traj.dim,
-        traj.dims[variable],
-        false
-    )
-end
-
-function (integrator::DerivativeIntegrator)(
-    traj::NamedTrajectory;
-    variable::Union{Symbol, Nothing}=nothing,
-    derivative::Union{Symbol, Nothing}=nothing,
-)
-    @assert !isnothing(variable) "variable must be provided"
-    @assert !isnothing(derivative) "derivative must be provided"
-    return DerivativeIntegrator(
-        variable,
-        derivative,
-        traj
-    )
-end
-
-state(integrator::DerivativeIntegrator) = integrator.variable
-controls(integrator::DerivativeIntegrator) = integrator.derivative
-
-@views function (D::DerivativeIntegrator)(
-    zₜ::AbstractVector,
-    zₜ₊₁::AbstractVector,
-    t::Int
-)
-    xₜ = zₜ[D.variable_components]
-    xₜ₊₁ = zₜ₊₁[D.variable_components]
-    dxₜ = zₜ[D.derivative_components]
-    if D.freetime
-        Δtₜ = zₜ[D.timestep]
-    else
-        Δtₜ = D.timestep
-    end
-    return xₜ₊₁ - xₜ - Δtₜ * dxₜ
-end
-
-@views function jacobian(
-    D::DerivativeIntegrator,
-    zₜ::AbstractVector,
-    zₜ₊₁::AbstractVector,
-    t::Int
-)
-    dxₜ = zₜ[D.derivative_components]
-    if D.freetime
-        Δtₜ = zₜ[D.timestep]
-    else
-        Δtₜ = D.timestep
-    end
-    ∂xₜD = sparse(-1.0I(D.dim))
-    ∂xₜ₊₁D = sparse(1.0I(D.dim))
-    ∂dxₜD = sparse(-Δtₜ * I(D.dim))
-    ∂ΔtₜD = -dxₜ
-    return ∂xₜD, ∂xₜ₊₁D, ∂dxₜD, ∂ΔtₜD
-end
-
-function get_comps(D::DerivativeIntegrator, traj::NamedTrajectory)
-    if D.freetime
-        return D.variable_components, D.derivative_components, traj.components[traj.timestep]
-    else
-        return D.variable_components, D.derivative_components
-    end
-end
diff --git a/src/integrators/exponential_integrators.jl b/src/integrators/exponential_integrators.jl
deleted file mode 100644
index 40cf1f0e..00000000
--- a/src/integrators/exponential_integrators.jl
+++ /dev/null
@@ -1,412 +0,0 @@
-"""
-This file includes expoential integrators for states and unitaries
-"""
-
-using ExponentialAction
-
-# ----------------------------------------------------------------------------- #
-#                         Quantum Exponential Integrators                       #
-# ----------------------------------------------------------------------------- #
-
-abstract type QuantumExponentialIntegrator <: QuantumIntegrator end
-
-# ----------------------------------------------------------------------------- #
-#                         Unitary Exponential Integrator                        #
-# ----------------------------------------------------------------------------- #
-
-struct UnitaryExponentialIntegrator <: QuantumExponentialIntegrator
-    unitary_components::Vector{Int}
-    drive_components::Vector{Int}
-    timestep::Union{Real, Int}
-    freetime::Bool
-    n_drives::Int
-    ketdim::Int
-    dim::Int
-    zdim::Int
-    autodiff::Bool
-    G::Function
-
-    function UnitaryExponentialIntegrator(
-        sys::AbstractQuantumSystem,
-        unitary_name::Symbol,
-        drive_name::Union{Symbol, Tuple{Vararg{Symbol}}},
-        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
-
-        unitary_components = traj.components[unitary_name]
-
-        if drive_name isa Tuple
-            drive_components = vcat((traj.components[s] for s ∈ drive_name)...)
-        else
-            drive_components = traj.components[drive_name]
-        end
-
-        n_drives = length(drive_components)
-
-        @assert all(diff(drive_components) .== 1) "controls must be in order"
-
-        freetime = traj.timestep isa Symbol
-
-        if freetime
-            timestep = traj.components[traj.timestep][1]
-        else
-            timestep = traj.timestep
-        end
-
-        return new(
-            unitary_components,
-            drive_components,
-            timestep,
-            freetime,
-            n_drives,
-            ketdim,
-            dim,
-            traj.dim,
-            autodiff,
-            G
-        )
-    end
-end
-
-function (integrator::UnitaryExponentialIntegrator)(
-    sys::AbstractQuantumSystem,
-    traj::NamedTrajectory;
-    unitary_name::Union{Nothing, Symbol}=nothing,
-    drive_name::Union{Nothing, Symbol, Tuple{Vararg{Symbol}}}=nothing,
-    G::Function=integrator.G,
-    autodiff::Bool=integrator.autodiff
-)
-    @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,
-        traj;
-        G=G,
-        autodiff=autodiff
-    )
-end
-
-function get_comps(P::UnitaryExponentialIntegrator, traj::NamedTrajectory)
-    if P.freetime
-        return P.unitary_components, P.drive_components, traj.components[traj.timestep]
-    else
-        return P.unitary_components, P.drive_components
-    end
-end
-
-# ------------------------------ Integrator --------------------------------- #
-
-@views function (ℰ::UnitaryExponentialIntegrator)(
-    zₜ::AbstractVector,
-    zₜ₊₁::AbstractVector,
-    t::Int
-)
-    Ũ⃗ₜ₊₁ = zₜ₊₁[ℰ.unitary_components]
-    Ũ⃗ₜ = zₜ[ℰ.unitary_components]
-    aₜ = zₜ[ℰ.drive_components]
-
-    if ℰ.freetime
-        Δtₜ = zₜ[ℰ.timestep]
-    else
-        Δtₜ = ℰ.timestep
-    end
-
-    Gₜ = ℰ.G(aₜ)
-
-    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,
-    zₜ₊₁::AbstractVector,
-    t::Int
-)
-    # get the state and control vectors
-    Ũ⃗ₜ = zₜ[ℰ.unitary_components]
-    aₜ = zₜ[ℰ.drive_components]
-
-    # obtain the timestep
-    if ℰ.freetime
-        Δtₜ = zₜ[ℰ.timestep]
-    else
-        Δtₜ = ℰ.timestep
-    end
-
-    # compute the generator
-    Gₜ = ℰ.G(aₜ)
-
-    Id = I(ℰ.ketdim)
-
-    expĜₜ = Id ⊗ hermitian_exp(Δtₜ * Gₜ)
-
-    ∂Ũ⃗ₜ₊₁ℰ = sparse(I, ℰ.dim, ℰ.dim)
-    ∂Ũ⃗ₜℰ = -expĜₜ
-
-    ∂aₜℰ = -ForwardDiff.jacobian(
-        a -> expv(Δtₜ, Id ⊗ ℰ.G(a), Ũ⃗ₜ),
-        aₜ
-    )
-
-    if ℰ.freetime
-        ∂Δtₜℰ = -(Id ⊗ Gₜ) * expĜₜ * Ũ⃗ₜ
-        return ∂Ũ⃗ₜℰ, ∂Ũ⃗ₜ₊₁ℰ, ∂aₜℰ, ∂Δtₜℰ
-    else
-        return ∂Ũ⃗ₜℰ, ∂Ũ⃗ₜ₊₁ℰ, ∂aₜℰ
-    end
-end
-
-struct QuantumStateExponentialIntegrator <: QuantumExponentialIntegrator
-    state_components::Vector{Int}
-    drive_components::Vector{Int}
-    timestep::Union{Real, Int}
-    freetime::Bool
-    n_drives::Int
-    ketdim::Int
-    dim::Int
-    zdim::Int
-    autodiff::Bool
-    G::Function
-
-    function QuantumStateExponentialIntegrator(
-        sys::AbstractQuantumSystem,
-        state_name::Symbol,
-        drive_name::Union{Symbol, Tuple{Vararg{Symbol}}},
-        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
-
-        state_components = traj.components[state_name]
-
-        if drive_name isa Tuple
-            drive_components = vcat((traj.components[s] for s ∈ drive_name)...)
-        else
-            drive_components = traj.components[drive_name]
-        end
-
-        n_drives = length(drive_components)
-
-        @assert all(diff(drive_components) .== 1) "controls must be in order"
-
-        freetime = traj.timestep isa Symbol
-
-        if freetime
-            timestep = traj.components[traj.timestep][1]
-        else
-            timestep = traj.timestep
-        end
-
-        return new(
-            state_components,
-            drive_components,
-            timestep,
-            freetime,
-            n_drives,
-            ketdim,
-            dim,
-            traj.dim,
-            autodiff,
-            G
-        )
-    end
-end
-
-function get_comps(P::QuantumStateExponentialIntegrator, traj::NamedTrajectory)
-    if P.freetime
-        return P.state_components, P.drive_components, traj.components[traj.timestep]
-    else
-        return P.state_components, P.drive_components
-    end
-end
-
-function (integrator::QuantumStateExponentialIntegrator)(
-    sys::AbstractQuantumSystem,
-    traj::NamedTrajectory;
-    state_name::Union{Nothing, Symbol}=nothing,
-    drive_name::Union{Nothing, Symbol, Tuple{Vararg{Symbol}}}=nothing,
-    G::Function=integrator.G,
-    autodiff::Bool=integrator.autodiff
-)
-    @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,
-        traj;
-        G=G,
-        autodiff=autodiff
-    )
-end
-
-# ------------------------------ Integrator --------------------------------- #
-
-@views function (ℰ::QuantumStateExponentialIntegrator)(
-    zₜ::AbstractVector,
-    zₜ₊₁::AbstractVector,
-    t::Int
-)
-    ψ̃ₜ₊₁ = zₜ₊₁[ℰ.state_components]
-    ψ̃ₜ = zₜ[ℰ.state_components]
-    aₜ = zₜ[ℰ.drive_components]
-
-    if ℰ.freetime
-        Δtₜ = zₜ[ℰ.timestep]
-    else
-        Δtₜ = ℰ.timestep
-    end
-
-    Gₜ = ℰ.G(aₜ)
-
-
-    return ψ̃ₜ₊₁ - expv(Δtₜ, Gₜ, ψ̃ₜ)
-end
-
-@views function jacobian(
-    ℰ::QuantumStateExponentialIntegrator,
-    zₜ::AbstractVector,
-    zₜ₊₁::AbstractVector,
-    t::Int
-)
-    # get the state and control vectors
-    ψ̃ₜ = zₜ[ℰ.state_components]
-    aₜ = zₜ[ℰ.drive_components]
-
-    # obtain the timestep
-    if ℰ.freetime
-        Δtₜ = zₜ[ℰ.timestep]
-    else
-        Δtₜ = ℰ.timestep
-    end
-
-    # compute the generator
-    Gₜ = ℰ.G(aₜ)
-
-    expGₜ = hermitian_exp(Δtₜ * Gₜ)
-
-    ∂ψ̃ₜ₊₁ℰ = sparse(I, ℰ.dim, ℰ.dim)
-    ∂ψ̃ₜℰ = -expGₜ
-
-    ∂aₜℰ = -ForwardDiff.jacobian(
-        a -> expv(Δtₜ, ℰ.G(a), ψ̃ₜ),
-        aₜ
-    )
-
-    if ℰ.freetime
-        ∂Δtₜℰ = -Gₜ * expGₜ * ψ̃ₜ
-        return ∂ψ̃ₜℰ, ∂ψ̃ₜ₊₁ℰ, ∂aₜℰ, ∂Δtₜℰ
-    else
-        return ∂ψ̃ₜℰ, ∂ψ̃ₜ₊₁ℰ, ∂aₜℰ
-    end
-end
-
-# ******************************************************************************* #
-
-@testitem "testing UnitaryExponentialIntegrator" begin
-    using NamedTrajectories
-    using ForwardDiff
-
-    T = 100
-    H_drift = GATES[:Z]
-    H_drives = [GATES[:X], GATES[:Y]]
-    n_drives = length(H_drives)
-
-    system = QuantumSystem(H_drift, H_drives)
-
-    U_init = GATES[:I]
-    U_goal = GATES[:X]
-
-    Ũ⃗_init = operator_to_iso_vec(U_init)
-    Ũ⃗_goal = operator_to_iso_vec(U_goal)
-
-    dt = 0.1
-
-    Z = NamedTrajectory(
-        (
-            Ũ⃗ = unitary_geodesic(U_goal, T),
-            a = randn(n_drives, T),
-            da = randn(n_drives, T),
-            Δt = fill(dt, 1, T),
-        ),
-        controls=(:da,),
-        timestep=:Δt,
-        goal=(Ũ⃗ = Ũ⃗_goal,)
-    )
-
-    ℰ = UnitaryExponentialIntegrator(system, :Ũ⃗, :a, Z)
-
-
-    ∂Ũ⃗ₜℰ, ∂Ũ⃗ₜ₊₁ℰ, ∂aₜℰ, ∂Δtₜℰ = jacobian(ℰ, Z[1].data, Z[2].data, 1)
-
-    ∂ℰ_forwarddiff = ForwardDiff.jacobian(
-        zz -> ℰ(zz[1:Z.dim], zz[Z.dim+1:end], 1),
-        [Z[1].data; Z[2].data]
-    )
-
-    @test ∂Ũ⃗ₜℰ ≈ ∂ℰ_forwarddiff[:, 1:ℰ.dim]
-    @test ∂Ũ⃗ₜ₊₁ℰ ≈ ∂ℰ_forwarddiff[:, Z.dim .+ (1:ℰ.dim)]
-    @test ∂aₜℰ ≈ ∂ℰ_forwarddiff[:, Z.components.a]
-    @test ∂Δtₜℰ ≈ ∂ℰ_forwarddiff[:, Z.components.Δt]
-end
-
-@testitem "testing QuantumStateExponentialIntegrator" begin
-    using NamedTrajectories
-    using ForwardDiff
-
-    T = 100
-    H_drift = GATES[:Z]
-    H_drives = [GATES[:X], GATES[:Y]]
-    n_drives = length(H_drives)
-
-    system = QuantumSystem(H_drift, H_drives)
-
-    U_init = GATES[:I]
-    U_goal = GATES[:X]
-
-    ψ̃_init = ket_to_iso(ket_from_string("g", [2]))
-    ψ̃_goal = ket_to_iso(ket_from_string("e", [2]))
-
-    dt = 0.1
-
-    Z = NamedTrajectory(
-        (
-            ψ̃ = linear_interpolation(ψ̃_init, ψ̃_goal, T),
-            a = randn(n_drives, T),
-            da = randn(n_drives, T),
-            Δt = fill(dt, 1, T),
-        ),
-        controls=(:da,),
-        timestep=:Δt,
-        goal=(ψ̃ = ψ̃_goal,)
-    )
-
-    ℰ = QuantumStateExponentialIntegrator(system, :ψ̃, :a, Z)
-    ∂ψ̃ₜℰ, ∂ψ̃ₜ₊₁ℰ, ∂aₜℰ, ∂Δtₜℰ = jacobian(ℰ, Z[1].data, Z[2].data, 1)
-
-    ∂ℰ_forwarddiff = ForwardDiff.jacobian(
-        zz -> ℰ(zz[1:Z.dim], zz[Z.dim+1:end], 1),
-        [Z[1].data; Z[2].data]
-    )
-
-    @test ∂ψ̃ₜℰ ≈ ∂ℰ_forwarddiff[:, 1:ℰ.dim]
-    @test ∂ψ̃ₜ₊₁ℰ ≈ ∂ℰ_forwarddiff[:, Z.dim .+ (1:ℰ.dim)]
-    @test ∂aₜℰ ≈ ∂ℰ_forwarddiff[:, Z.components.a]
-    @test ∂Δtₜℰ ≈ ∂ℰ_forwarddiff[:, Z.components.Δt]
-end
diff --git a/src/integrators/pade_integrators.jl b/src/integrators/pade_integrators.jl
deleted file mode 100644
index 549bc70b..00000000
--- a/src/integrators/pade_integrators.jl
+++ /dev/null
@@ -1,1096 +0,0 @@
-function nth_order_pade(Gₜ::Matrix, n::Int)
-    @assert n ∈ keys(PADE_COEFFICIENTS)
-    coeffs = PADE_COEFFICIENTS[n]
-    Id = 1.0I(size(Gₜ, 1))
-    p = n ÷ 2
-    G_powers = [Gₜ^i for i = 1:p]
-    B = Id + sum((-1)^k * coeffs[k] * G_powers[k] for k = 1:p)
-    F = Id + sum(coeffs[k] * G_powers[k] for k = 1:p)
-    return inv(B) * F
-end
-
-
-
-fourth_order_pade(Gₜ::Matrix) = nth_order_pade(Gₜ, 4)
-sixth_order_pade(Gₜ::Matrix) = nth_order_pade(Gₜ, 6)
-eighth_order_pade(Gₜ::Matrix) = nth_order_pade(Gₜ, 8)
-tenth_order_pade(Gₜ::Matrix) = nth_order_pade(Gₜ, 10)
-twelth_order_pade(Gₜ::Matrix) = nth_order_pade(Gₜ, 12)
-fourteenth_order_pade(Gₜ::Matrix) = nth_order_pade(Gₜ, 14)
-sixteenth_order_pade(Gₜ::Matrix) = nth_order_pade(Gₜ, 16)
-eighteenth_order_pade(Gₜ::Matrix) = nth_order_pade(Gₜ, 18)
-twentieth_order_pade(Gₜ::Matrix) = nth_order_pade(Gₜ, 20)
-
-function compute_powers(G::AbstractMatrix{T}, order::Int) where T <: Number
-    powers = Array{typeof(G)}(undef, order)
-    powers[1] = G
-    for k = 2:order
-        powers[k] = powers[k-1] * G
-    end
-    return powers
-end
-
-# key is the order of the integrator
-# and the value are the Pade coefficients
-# for each term
-const PADE_COEFFICIENTS = OrderedDict{Int,Vector{Float64}}(
-    4 => [1/2, 1/12],
-    6 => [1/2, 1/10, 1/120],
-    8 => [1/2, 3/28, 1/84, 1/1680],
-    10 => [1/2, 1/9, 1/72, 1/1008, 1/30240],
-    12 => [1/2, 5/44, 1/66, 1/792, 1/15840, 1/665280],
-    14 => [1/2, 3/26, 5/312, 5/3432, 1/11440, 1/308880, 1/17297280],
-    16 => [1/2, 7/60, 1/60, 1/624, 1/9360, 1/205920, 1/7207200, 1/518918400],
-    18 => [1/2, 2/17, 7/408, 7/4080, 1/8160, 1/159120, 1/4455360, 1/196035840, 1/17643225600],
-    20 => [1/2, 9/76, 1/57, 7/3876, 7/51680, 7/930240, 1/3255840, 1/112869120, 1/6094932480, 1/670442572800]
-)
-
-function pade_operator(
-    Id::AbstractMatrix,
-    G_powers::Vector{<:AbstractMatrix},
-    coeffs::Vector{<:Real}
-)
-    return Id + sum(coeffs .* G_powers)
-end
-
-function forward_pade_coefficients(Δt::Real, pade_order::Int; timestep_derivative=false)
-    n = pade_order ÷ 2
-    if !timestep_derivative
-        return PADE_COEFFICIENTS[2n] .* (Δt .^ (1:n))
-    else
-        return PADE_COEFFICIENTS[2n] .* (Δt .^ (0:n-1)) .* (1:n)
-    end
-end
-
-function backward_pade_coefficients(Δt::Real, pade_order::Int; timestep_derivative=false)
-    n = pade_order ÷ 2
-    if !timestep_derivative
-        return PADE_COEFFICIENTS[2n] .* ((-Δt) .^ (1:n))
-    else
-        return PADE_COEFFICIENTS[2n] .* ((-1) .^ (1:n)) .* (Δt .^ (0:n-1)) .* (1:n)
-    end
-end
-
-function pade_coefficients(Δt::Real, pade_order::Int; timestep_derivative=false)
-    F_coeffs = forward_pade_coefficients(Δt, pade_order;
-        timestep_derivative=timestep_derivative
-    )
-    B_coeffs = backward_pade_coefficients(Δt, pade_order;
-        timestep_derivative=timestep_derivative
-    )
-    return F_coeffs, B_coeffs
-end
-
-function backward_operator(
-    G_powers::Vector{<:AbstractMatrix},
-    Id::AbstractMatrix,
-    Δt::Real;
-    timestep_derivative=false
-)
-    pade_order = 2 * length(G_powers)
-    coeffs = backward_pade_coefficients(Δt, pade_order; timestep_derivative=timestep_derivative)
-    return pade_operator(Id, G_powers, coeffs)
-end
-
-backward_operator(G::AbstractMatrix, pade_order::Int, args...; kwargs...) =
-    backward_operator(compute_powers(G, pade_order ÷ 2), args...; kwargs...)
-
-function forward_operator(
-    G_powers::Vector{<:AbstractMatrix},
-    Id::AbstractMatrix,
-    Δt::Real;
-    timestep_derivative=false
-)
-    pade_order = 2 * length(G_powers)
-    coeffs = forward_pade_coefficients(Δt, pade_order; timestep_derivative=timestep_derivative)
-    return pade_operator(Id, G_powers, coeffs)
-end
-
-forward_operator(G::AbstractMatrix, pade_order::Int, args...; kwargs...) =
-    forward_operator(compute_powers(G, pade_order ÷ 2), args...; kwargs...)
-
-function pade_operators(
-    G_powers::Vector{<:AbstractMatrix},
-    Id::AbstractMatrix,
-    Δt::Real;
-    kwargs...
-)
-    F = forward_operator(G_powers, Id, Δt; kwargs...)
-    B = backward_operator(G_powers, Id, Δt; kwargs...)
-    return F, B
-end
-
-function pade_operators(
-    G_powers::Vector{<:SparseMatrixCSC},
-    Id::SparseMatrixCSC,
-    Δt::Real;
-    kwargs...
-)
-    F = forward_operator(G_powers, Id, Δt; kwargs...)
-    B = backward_operator(G_powers, Id, Δt; kwargs...)
-    droptol!(F, 1e-12)
-    droptol!(B, 1e-12)
-    return F, B
-end
-
-pade_operators(G::AbstractMatrix, pade_order::Int, args...; kwargs...) =
-    pade_operators(compute_powers(G, pade_order ÷ 2), args...; kwargs...)
-
-@views function ∂aʲF(
-    P::QuantumIntegrator,
-    G_powers::Vector{<:AbstractMatrix},
-    Δt::Real,
-    ∂G_∂aʲ::AbstractMatrix
-)
-    F_coeffs = forward_pade_coefficients(Δt, P.order)
-    ∂F_∂aʲ = zeros(size(G_powers[1]))
-    n = length(G_powers)
-    for p = 1:n
-        if p == 1
-            ∂F_∂aʲ += F_coeffs[p] * ∂G_∂aʲ
-        else
-            for k = 1:p
-                if k == 1
-                    ∂F_∂aʲ += F_coeffs[p] * ∂G_∂aʲ * G_powers[p-1]
-                elseif k == p
-                    ∂F_∂aʲ += F_coeffs[p] * G_powers[p-1] * ∂G_∂aʲ
-                else
-                    ∂F_∂aʲ += F_coeffs[p] * G_powers[k-1] * ∂G_∂aʲ * G_powers[p-k]
-                end
-            end
-        end
-    end
-    return ∂F_∂aʲ
-end
-
-@views function ∂aʲB(
-    P::QuantumIntegrator,
-    G_powers::Vector{<:AbstractMatrix},
-    Δt::Real,
-    ∂G_∂aʲ::AbstractMatrix
-)
-    B_coeffs = backward_pade_coefficients(Δt, P.order)
-    ∂B_∂aʲ = zeros(size(G_powers[1]))
-    for p = 1:(P.order ÷ 2)
-        if p == 1
-            ∂B_∂aʲ += B_coeffs[p] * ∂G_∂aʲ
-        else
-            for k = 1:p
-                if k == 1
-                    ∂B_∂aʲ += B_coeffs[p] * ∂G_∂aʲ * G_powers[p-1]
-                elseif k == p
-                    ∂B_∂aʲ += B_coeffs[p] * G_powers[p-1] * ∂G_∂aʲ
-                else
-                    ∂B_∂aʲ += B_coeffs[p] * G_powers[k-1] * ∂G_∂aʲ * G_powers[p-k]
-                end
-            end
-        end
-    end
-    return ∂B_∂aʲ
-end
-
-
-
-# ----------------------------------------------------------------
-#                       Quantum Pade Integrator
-# ----------------------------------------------------------------
-
-
-
-abstract type QuantumPadeIntegrator <: QuantumIntegrator end
-
-
-
-# ----------------------------------------------------------------
-#                       Unitary Pade Integrator
-# ----------------------------------------------------------------
-
-
-"""
-"""
-struct UnitaryPadeIntegrator <: QuantumPadeIntegrator
-    unitary_components::Vector{Int}
-    drive_components::Vector{Int}
-    timestep::Union{Real, Int} # either the timestep or the index of the timestep
-    freetime::Bool
-    n_drives::Int
-    ketdim::Int
-    dim::Int
-    zdim::Int
-    order::Int
-    autodiff::Bool
-    G::Function
-    ∂G::Function
-
-    """
-        UnitaryPadeIntegrator(
-            sys::AbstractQuantumSystem,
-            unitary_name::Symbol,
-            drive_name::Union{Symbol,Tuple{Vararg{Symbol}}};
-            order::Int=4,
-            autodiff::Bool=order != 4
-        ) where R <: Real
-
-    Construct a `UnitaryPadeIntegrator` for the quantum system `sys`.
-
-    # Examples
-
-    ## a bare integrator
-    ```julia
-        P = UnitaryPadeIntegrator(sys)
-    ```
-
-    ## 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)
-    ```
-
-    # 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
-
-    """
-    function UnitaryPadeIntegrator(
-        sys::AbstractQuantumSystem,
-        unitary_name::Symbol,
-        drive_name::Union{Symbol,Tuple{Vararg{Symbol}}},
-        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))
-
-        unitary_components = traj.components[unitary_name]
-
-        if drive_name isa Tuple
-            drive_components = vcat((traj.components[s] for s ∈ drive_name)...)
-        else
-            drive_components = traj.components[drive_name]
-        end
-
-        n_drives = length(drive_components)
-
-        @assert all(diff(drive_components) .== 1) "controls must be in order"
-
-        freetime = traj.timestep isa Symbol
-
-        if freetime
-            timestep = traj.components[traj.timestep][1]
-        else
-            timestep = traj.timestep
-        end
-
-        return new(
-            unitary_components,
-            drive_components,
-            timestep,
-            freetime,
-            n_drives,
-            ketdim,
-            dim,
-            traj.dim,
-            order,
-            autodiff,
-            G,
-            ∂G
-        )
-    end
-end
-
-function get_comps(P::UnitaryPadeIntegrator, traj::NamedTrajectory)
-    if P.freetime
-        return P.unitary_components, P.drive_components, traj.components[traj.timestep]
-    else
-        return P.unitary_components, P.drive_components
-    end
-end
-
-function (integrator::UnitaryPadeIntegrator)(
-    sys::AbstractQuantumSystem,
-    traj::NamedTrajectory;
-    unitary_name::Union{Symbol, Nothing}=nothing,
-    drive_name::Union{Symbol, Tuple{Vararg{Symbol}}, Nothing}=nothing,
-    order::Int=integrator.order,
-    G::Function=integrator.G,
-    ∂G::Function=integrator.∂G,
-    autodiff::Bool=integrator.autodiff
-)
-    @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,
-        traj;
-        order=order,
-        G=G,
-        ∂G=∂G,
-        autodiff=autodiff
-    )
-end
-
-# ------------------- Integrator -------------------
-
-
-
-function nth_order_pade(
-    P::UnitaryPadeIntegrator,
-    Ũ⃗ₜ₊₁::AbstractVector,
-    Ũ⃗ₜ::AbstractVector,
-    aₜ::AbstractVector,
-    Δt::Real
-)
-    Gₜ = P.G(aₜ)
-
-    F, B = pade_operators(Gₜ, P.order, I(2P.ketdim), Δt)
-
-    I_N = sparse(I, P.ketdim, P.ketdim)
-
-    return (I_N ⊗ B) * Ũ⃗ₜ₊₁ - (I_N ⊗ F) * Ũ⃗ₜ
-end
-
-@views function(P::UnitaryPadeIntegrator)(
-    zₜ::AbstractVector,
-    zₜ₊₁::AbstractVector,
-    t::Int
-)
-    Ũ⃗ₜ₊₁ = zₜ₊₁[P.unitary_components]
-    Ũ⃗ₜ = zₜ[P.unitary_components]
-    aₜ = zₜ[P.drive_components]
-
-    if P.freetime
-        Δtₜ = zₜ[P.timestep]
-    else
-        Δtₜ = P.timestep
-    end
-
-
-    return nth_order_pade(P, Ũ⃗ₜ₊₁, Ũ⃗ₜ, aₜ, Δtₜ)
-end
-
-# ------------------- Jacobians -------------------
-
-# aₜ should be a vector with all the controls. concatenate all the named traj controls
-function ∂aₜ(
-    P::UnitaryPadeIntegrator,
-    G_powers::Vector{<:AbstractMatrix},
-    Ũ⃗ₜ₊₁::AbstractVector,
-    Ũ⃗ₜ::AbstractVector,
-    aₜ::AbstractVector,
-    Δtₜ::Real
-)
-    ∂aP = zeros(eltype(Ũ⃗ₜ), P.dim, P.n_drives)
-
-    ∂G_∂aₜ = P.∂G(aₜ)
-
-    I_N = sparse(I, P.ketdim, P.ketdim)
-
-    for j = 1:P.n_drives
-
-        # TODO: maybe rework for arbitrary drive indices eventually
-
-        ∂aₜʲF = ∂aʲF(P, G_powers, Δtₜ, ∂G_∂aₜ[j])
-        ∂aₜʲB = ∂aʲB(P, G_powers, Δtₜ, ∂G_∂aₜ[j])
-
-        ∂aP[:, j] = (I_N ⊗ ∂aₜʲB) * Ũ⃗ₜ₊₁ - (I_N ⊗ ∂aₜʲF) * Ũ⃗ₜ
-    end
-
-    return ∂aP
-end
-
-
-function ∂Δtₜ(
-    P::UnitaryPadeIntegrator,
-    Gₜ_powers::Vector{<:AbstractMatrix},
-    Ũ⃗ₜ₊₁::AbstractVector,
-    Ũ⃗ₜ::AbstractVector,
-    Δtₜ::Real
-)
-    ∂ΔtₜF_coeffs, ∂ΔtₜB_coeffs = pade_coefficients(Δtₜ, P.order;
-        timestep_derivative=true
-    )
-
-    ∂ΔtₜF = sum(∂ΔtₜF_coeffs .* Gₜ_powers)
-    ∂ΔtₜB = sum(∂ΔtₜB_coeffs .* Gₜ_powers)
-
-    I_N = sparse(I, P.ketdim, P.ketdim)
-
-    return (I_N ⊗ ∂ΔtₜB) * Ũ⃗ₜ₊₁ - (I_N ⊗ ∂ΔtₜF) * Ũ⃗ₜ
-end
-
-@views function jacobian(
-    P::UnitaryPadeIntegrator,
-    zₜ::AbstractVector,
-    zₜ₊₁::AbstractVector,
-    t::Int
-)
-    # obtain state and control vectors
-    Ũ⃗ₜ₊₁ = zₜ₊₁[P.unitary_components]
-    Ũ⃗ₜ = zₜ[P.unitary_components]
-    aₜ = zₜ[P.drive_components]
-
-    Gₜ = P.G(aₜ)
-
-    # obtain timestep
-    if P.freetime
-        Δtₜ = zₜ[P.timestep]
-    else
-        Δtₜ = P.timestep
-    end
-
-    Gₜ_powers = compute_powers(Gₜ, P.order ÷ 2)
-
-    ∂aₜP = ∂aₜ(P, Gₜ_powers, Ũ⃗ₜ₊₁, Ũ⃗ₜ, aₜ, Δtₜ)
-
-    Id = sparse(I, P.ketdim, P.ketdim)
-
-    Fₜ, Bₜ = pade_operators(Gₜ_powers, I(2P.ketdim), Δtₜ)
-
-    ∂Ũ⃗ₜP = -Id ⊗ Fₜ
-    ∂Ũ⃗ₜ₊₁P = Id ⊗ Bₜ
-
-    if P.freetime
-        ∂ΔtₜP = ∂Δtₜ(P, Gₜ_powers, Ũ⃗ₜ₊₁, Ũ⃗ₜ, Δtₜ)
-        return ∂Ũ⃗ₜP, ∂Ũ⃗ₜ₊₁P, ∂aₜP, ∂ΔtₜP
-    else
-        return ∂Ũ⃗ₜP, ∂Ũ⃗ₜ₊₁P, ∂aₜP
-    end
-end
-# ----------------------------------------------------------------
-#                  Quantum State Pade Integrator
-# ----------------------------------------------------------------
-
-struct QuantumStatePadeIntegrator <: QuantumPadeIntegrator
-    state_components::Vector{Int}
-    drive_components::Vector{Int}
-    timestep::Union{Real, Int} # either the timestep or the index of the timestep
-    freetime::Bool
-    n_drives::Int
-    ketdim::Int
-    dim::Int
-    zdim::Int
-    order::Int
-    autodiff::Bool
-    G::Function
-    ∂G::Function
-
-    """
-        QuantumStatePadeIntegrator(
-            sys::AbstractQuantumSystem,
-            state_name::Union{Symbol,Nothing}=nothing,
-            drive_name::Union{Symbol,Tuple{Vararg{Symbol}},Nothing}=nothing,
-            timestep_name::Union{Symbol,Nothing}=nothing;
-            order::Int=4,
-            autodiff::Bool=false
-        ) where R <: Real
-
-    Construct a `QuantumStatePadeIntegrator` for the quantum system `sys`.
-
-    # Examples
-
-    ## for a single drive `a`:
-    ```julia
-        P = QuantumStatePadeIntegrator(sys, :ψ̃, :a)
-    ```
-
-    ## for two drives `α` and `γ`, order `8`, and autodiffed:
-    ```julia
-        P = QuantumStatePadeIntegrator(sys, :ψ̃, (:α, :γ); order=8, autodiff=true)
-    ```
-
-    # Arguments
-    - `sys::AbstractQuantumSystem`: the quantum system
-    - `state_name::Symbol`: the nameol for the quantum state
-    - `drive_name::Union{Symbol,Tuple{Vararg{Symbol}}}`: 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=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}}},
-        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))"
-
-        ketdim = size(sys.H_drift, 1)
-        dim = 2ketdim
-
-        state_components = traj.components[state_name]
-
-        if drive_name isa Tuple
-            drive_components = vcat((traj.components[s] for s ∈ drive_name)...)
-        else
-            drive_components = traj.components[drive_name]
-        end
-
-        n_drives = length(drive_components)
-
-        @assert all(diff(drive_components) .== 1) "controls must be in order"
-
-        freetime = traj.timestep isa Symbol
-
-        if freetime
-            timestep = traj.components[traj.timestep][1]
-        else
-            timestep = traj.timestep
-        end
-
-
-        return new(
-            state_components,
-            drive_components,
-            timestep,
-            freetime,
-            n_drives,
-            ketdim,
-            dim,
-            traj.dim,
-            order,
-            autodiff,
-            G,
-            ∂G
-        )
-    end
-end
-
-function get_comps(P::QuantumStatePadeIntegrator, traj::NamedTrajectory)
-    if P.freetime
-        return P.state_components, P.drive_components, traj.components[traj.timestep]
-    else
-        return P.state_components, P.drive_components
-    end
-end
-
-function (integrator::QuantumStatePadeIntegrator)(
-    sys::AbstractQuantumSystem,
-    traj::NamedTrajectory;
-    state_name::Union{Symbol, Nothing}=nothing,
-    drive_name::Union{Symbol, Tuple{Vararg{Symbol}}, Nothing}=nothing,
-    order::Int=integrator.order,
-    G::Function=integrator.G,
-    ∂G::Function=integrator.∂G,
-    autodiff::Bool=integrator.autodiff
-)
-    @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,
-        traj;
-        order=order,
-        G=G,
-        ∂G=∂G,
-        autodiff=autodiff
-    )
-end
-
-# ------------------- Integrator -------------------
-
-function nth_order_pade(
-    P::QuantumStatePadeIntegrator,
-    ψ̃ₜ₊₁::AbstractVector,
-    ψ̃ₜ::AbstractVector,
-    aₜ::AbstractVector,
-    Δt::Real
-)
-    Gₜ = P.G(aₜ)
-
-    F, B = pade_operators(Gₜ, P.order, I(2P.ketdim), Δt)
-
-    return B * ψ̃ₜ₊₁ - F * ψ̃ₜ
-end
-
-
-@views function(P::QuantumStatePadeIntegrator)(
-    zₜ::AbstractVector,
-    zₜ₊₁::AbstractVector,
-    t::Int
-)
-    ψ̃ₜ₊₁ = zₜ₊₁[P.state_components]
-    ψ̃ₜ = zₜ[P.state_components]
-    aₜ = zₜ[P.drive_components]
-
-    if P.freetime
-        Δtₜ = zₜ[P.timestep]
-    else
-        Δtₜ = P.timestep
-    end
-    return nth_order_pade(P, ψ̃ₜ₊₁, ψ̃ₜ, aₜ, Δtₜ)
-end
-
-
-function ∂aₜ(
-    P::QuantumStatePadeIntegrator,
-    G_powers::Vector{<:AbstractMatrix},
-    ψ̃ₜ₊₁::AbstractVector,
-    ψ̃ₜ::AbstractVector,
-    aₜ::AbstractVector,
-    Δtₜ::Real
-)
-    ∂aP = zeros(eltype(ψ̃ₜ), P.dim, P.n_drives)
-
-    ∂G_∂aₜ = P.∂G(aₜ)
-
-    for j = 1:P.n_drives
-
-        ∂aₜʲF = ∂aʲF(P, G_powers, Δtₜ, ∂G_∂aₜ[j])
-        ∂aₜʲB = ∂aʲB(P, G_powers, Δtₜ, ∂G_∂aₜ[j])
-
-        ∂aP[:, j] = ∂aₜʲB * ψ̃ₜ₊₁ - ∂aₜʲF * ψ̃ₜ
-    end
-
-    return ∂aP
-end
-
-
-
-function ∂Δtₜ(
-    P::QuantumStatePadeIntegrator,
-    Gₜ_powers::Vector{<:AbstractMatrix},
-    ψ̃ₜ₊₁::AbstractVector,
-    ψ̃ₜ::AbstractVector,
-    Δtₜ::Real
-)
-    ∂ΔtₜF_coeffs, ∂ΔtₜB_coeffs = pade_coefficients(Δtₜ, P.order;
-        timestep_derivative=true
-    )
-
-    ∂ΔtₜF = sum(∂ΔtₜF_coeffs .* Gₜ_powers)
-    ∂ΔtₜB = sum(∂ΔtₜB_coeffs .* Gₜ_powers)
-
-    return ∂ΔtₜB * ψ̃ₜ₊₁ - ∂ΔtₜF * ψ̃ₜ
-end
-
-@views function jacobian(
-    P::QuantumStatePadeIntegrator,
-    zₜ::AbstractVector,
-    zₜ₊₁::AbstractVector,
-    t::Int
-)
-    # obtain state and control vectors
-    ψ̃ₜ₊₁ = zₜ₊₁[P.state_components]
-    ψ̃ₜ = zₜ[P.state_components]
-    aₜ = zₜ[P.drive_components]
-
-    Gₜ = P.G(aₜ)
-
-    # obtain timestep
-    if P.freetime
-        Δtₜ = zₜ[P.timestep]
-    else
-        Δtₜ = P.timestep
-    end
-
-    Gₜ_powers = compute_powers(Gₜ, P.order ÷ 2)
-
-    ∂aₜP = ∂aₜ(P, Gₜ_powers, ψ̃ₜ₊₁, ψ̃ₜ, aₜ, Δtₜ)
-
-    # jacobian wrt state
-    Fₜ, Bₜ = pade_operators(Gₜ_powers, I(2P.ketdim), Δtₜ)
-
-    ∂ψ̃ₜP = -Fₜ
-    ∂ψ̃ₜ₊₁P = Bₜ
-
-    if P.freetime
-        ∂ΔtₜP = ∂Δtₜ(P, Gₜ_powers, ψ̃ₜ₊₁, ψ̃ₜ, Δtₜ)
-        return ∂ψ̃ₜP, ∂ψ̃ₜ₊₁P, ∂aₜP, ∂ΔtₜP
-    else
-        return ∂ψ̃ₜP, ∂ψ̃ₜ₊₁P, ∂aₜP
-    end
-end
-
-
-# ---------------------------------------
-# Hessian of the Lagrangian
-# ---------------------------------------
-
-#calculate a deriv first and then indexing game
-function μ∂aₜ∂Ũ⃗ₜ(
-    P::UnitaryPadeIntegrator,
-    aₜ::AbstractVector{T},
-    Δtₜ::Real,
-    μₜ::AbstractVector{T},
-) where T <: Real
-
-    n_drives = P.n_drives
-
-    if P.autodiff
-
-    elseif P.order == 4
-        μ∂aₜ∂Ũ⃗ₜP = Array{T}(undef, P.dim, n_drives)
-
-        for j = 1:n_drives
-            Gʲ = P.G_drives[j]
-            Ĝʲ = P.G(aₜ, P.G_drift_anticomms[j], P.G_drive_anticomms[:, j])
-            ∂aₜ∂Ũ⃗ₜ_block_i = -(Δtₜ / 2 * Gʲ + Δtₜ^2 / 12 * Ĝʲ)
-            # sparse is necessary since blockdiag doesn't accept dense matrices
-            ∂aₜ∂Ũ⃗ₜ = blockdiag(fill(sparse(∂aₜ∂Ũ⃗ₜ_block_i), P.ketdim)...)
-            μ∂aₜ∂Ũ⃗ₜP[:, j] = ∂aₜ∂Ũ⃗ₜ' * μₜ
-        end
-    else
-        ## higher order pade goes here
-    end
-    return μ∂aₜ∂Ũ⃗ₜP
-end
-
-function μ∂Ũ⃗ₜ₊₁∂aₜ(
-    P::UnitaryPadeIntegrator,
-    aₜ::AbstractVector{T},
-    Δtₜ::Real,
-    μₜ::AbstractVector{T},
-) where T <: Real
-
-    n_drives = P.n_drives
-    μ∂Ũ⃗ₜ₊₁∂aₜP = zeros(T, n_drives, P.dim)
-
-    for j = 1:n_drives
-        Gʲ = P.G_drives[j]
-        Ĝʲ = P.G(aₜ, P.G_drift_anticomms[j], P.G_drive_anticomms[:, j])
-        ∂Ũ⃗ₜ₊₁∂aₜ_block_i = -Δtₜ / 2 * Gʲ + Δtₜ^2 / 12 * Ĝʲ
-        # sparse is necessary since blockdiag doesn't accept dense matrices
-        ∂Ũ⃗ₜ₊₁∂aₜ = blockdiag(fill(sparse(∂Ũ⃗ₜ₊₁∂aₜ_block_i), P.ketdim)...)
-        μ∂Ũ⃗ₜ₊₁∂aₜP[j, :] = μₜ' * ∂Ũ⃗ₜ₊₁∂aₜ
-    end
-
-    return μ∂Ũ⃗ₜ₊₁∂aₜP
-end
-
-function μ∂aₜ∂ψ̃ₜ(
-    P::QuantumStatePadeIntegrator,
-    aₜ::AbstractVector{T},
-    Δtₜ::Real,
-    μₜ::AbstractVector{T},
-) where T <: Real
-
-    n_drives = P.n_drives
-    μ∂aₜ∂ψ̃ₜP = zeros(T, P.dim, n_drives)
-
-    for j = 1:n_drives
-        Gʲ = P.G_drives[j]
-        Ĝʲ = P.G(aₜ, P.G_drift_anticomms[j], P.G_drive_anticomms[:, j])
-        ∂aₜ∂ψ̃ₜP = -(Δtₜ / 2 * Gʲ + Δtₜ^2 / 12 * Ĝʲ)
-        μ∂aₜ∂ψ̃ₜP[:, j] = ∂aₜ∂ψ̃ₜP' * μₜ
-    end
-
-    return μ∂aₜ∂ψ̃ₜP
-end
-
-function μ∂ψ̃ₜ₊₁∂aₜ(
-    P::QuantumStatePadeIntegrator,
-    aₜ::AbstractVector{T},
-    Δtₜ::Real,
-    μₜ::AbstractVector{T},
-) where T <: Real
-
-    n_drives = P.n_drives
-    μ∂ψ̃ₜ₊₁∂aₜP = zeros(T, n_drives, P.dim)
-
-    for j = 1:n_drives
-        Gʲ = P.G_drives[j]
-        Ĝʲ = P.G(aₜ, P.G_drift_anticomms[j], P.G_drive_anticomms[:, j])
-        ∂ψ̃ₜ₊₁∂aₜP = -Δtₜ / 2 * Gʲ + Δtₜ^2 / 12 * Ĝʲ
-        μ∂ψ̃ₜ₊₁∂aₜP[j, :] = μₜ' * ∂ψ̃ₜ₊₁∂aₜP
-    end
-
-    #can add if else for higher order derivatives
-    return μ∂ψ̃ₜ₊₁∂aₜP
-end
-
-function μ∂²aₜ(
-    P::UnitaryPadeIntegrator,
-    Ũ⃗ₜ₊₁::AbstractVector{T},
-    Ũ⃗ₜ::AbstractVector{T},
-    Δtₜ::Real,
-    μₜ::AbstractVector{T},
-) where T <: Real
-
-    n_drives = P.n_drives
-    μ∂²aₜP = zeros(T, n_drives, n_drives)
-
-    if P.order==4
-        for i = 1:n_drives
-            for j = 1:i
-                ∂aʲ∂aⁱP_block =
-                    Δtₜ^2 / 12 * P.G_drive_anticomms[i, j]
-                ∂aʲ∂aⁱP = blockdiag(fill(sparse(∂aʲ∂aⁱP_block), P.ketdim)...)
-                μ∂²aₜP[j, i] = dot(μₜ, ∂aʲ∂aⁱP*(Ũ⃗ₜ₊₁ - Ũ⃗ₜ))
-            end
-        end
-    end
-
-    return Symmetric(μ∂²aₜP)
-end
-
-function μ∂²aₜ(
-    P::QuantumStatePadeIntegrator,
-    ψ̃ₜ₊₁::AbstractVector{T},
-    ψ̃ₜ::AbstractVector{T},
-    Δtₜ::Real,
-    μₜ::AbstractVector{T},
-) where T <: Real
-
-    n_drives = P.n_drives
-    μ∂²aₜP = Array{T}(undef, n_drives, n_drives)
-
-    if P.order==4
-        for i = 1:n_drives
-            for j = 1:i
-                ∂aʲ∂aⁱP = Δtₜ^2 / 12 * P.G_drive_anticomms[i, j] * (ψ̃ₜ₊₁ - ψ̃ₜ)
-                μ∂²aₜP[j, i] = dot(μₜ, ∂aʲ∂aⁱP)
-            end
-        end
-    end
-
-    return Symmetric(μ∂²aₜP)
-end
-
-function μ∂Δtₜ∂aₜ(
-    P::UnitaryPadeIntegrator,
-    Ũ⃗ₜ₊₁::AbstractVector{T},
-    Ũ⃗ₜ::AbstractVector{T},
-    aₜ::AbstractVector{T},
-    Δtₜ::T,
-    μₜ::AbstractVector{T},
-) where T <: Real
-
-    n_drives = P.n_drives
-    μ∂Δtₜ∂aₜP = Array{T}(undef, n_drives)
-
-    if P.order == 4
-        for j = 1:n_drives
-            Gʲ = P.G_drives[j]
-            Ĝʲ = P.G(aₜ, P.G_drift_anticomms[j], P.G_drive_anticomms[:, j])
-            B = blockdiag(fill(sparse(-1/2 * Gʲ + 1/6 * Δtₜ * Ĝʲ), P.ketdim)...)
-            F = blockdiag(fill(sparse(1/2 * Gʲ + 1/6 * Δtₜ * Ĝʲ), P.ketdim)...)
-            ∂Δtₜ∂aₜ_j =  B*Ũ⃗ₜ₊₁ - F*Ũ⃗ₜ
-            μ∂Δtₜ∂aₜP[j] = dot(μₜ, ∂Δtₜ∂aₜ_j)
-        end
-    end
-    return μ∂Δtₜ∂aₜP
-end
-
-function μ∂Δtₜ∂aₜ(
-    P::QuantumStatePadeIntegrator,
-    ψ̃ₜ₊₁::AbstractVector{T},
-    ψ̃ₜ::AbstractVector{T},
-    aₜ::AbstractVector{T},
-    Δtₜ::T,
-    μₜ::AbstractVector{T},
-) where T <: Real
-
-    n_drives = P.n_drives
-    μ∂Δtₜ∂aₜP = Array{T}(undef, n_drives)
-
-    if P.order == 4
-        for j = 1:n_drives
-            Gʲ = P.G_drives[j]
-            Ĝʲ = P.G(aₜ, P.G_drift_anticomms[j], P.G_drive_anticomms[:, j])
-            ∂Δt∂aʲP =
-                -1 / 2 * Gʲ * (ψ̃ₜ₊₁ + ψ̃ₜ) +
-                1 / 6 * Δtₜ * Ĝʲ * (ψ̃ₜ₊₁ - ψ̃ₜ)
-            μ∂Δtₜ∂aₜP[j] = dot(μₜ, ∂Δt∂aʲP)
-        end
-    end
-    return μ∂Δtₜ∂aₜP
-end
-
-function μ∂Δtₜ∂Ũ⃗ₜ(
-    P::UnitaryPadeIntegrator,
-    aₜ::AbstractVector,
-    Δtₜ::Real,
-    μₜ::AbstractVector
-)
-    Gₜ = P.G(aₜ, P.G_drift, P.G_drives)
-    minus_F = -(1/2 * Gₜ + 1/6 * Δtₜ * Gₜ^2)
-    big_minus_F = blockdiag(fill(sparse(minus_F), P.ketdim)...)
-    return big_minus_F' * μₜ
-end
-
-function μ∂Ũ⃗ₜ₊₁∂Δtₜ(
-    P::UnitaryPadeIntegrator,
-    aₜ::AbstractVector,
-    Δtₜ::Real,
-    μₜ::AbstractVector
-)
-    Gₜ = P.G(aₜ, P.G_drift, P.G_drives)
-    B = -1/2 * Gₜ + 1/6 * Δtₜ * Gₜ^2
-    big_B = blockdiag(fill(sparse(B), P.ketdim)...)
-    return μₜ' * big_B
-end
-
-function μ∂Δtₜ∂ψ̃ₜ(
-    P::QuantumStatePadeIntegrator,
-    aₜ::AbstractVector,
-    Δtₜ::Real,
-    μₜ::AbstractVector
-)
-    # memoize the calc here
-    Gₜ = P.G(aₜ, P.G_drift, P.G_drives)
-    minus_F = -(1/2 * Gₜ + 1/6 * Δtₜ * Gₜ^2)
-    return minus_F' * μₜ
-end
-
-function μ∂ψ̃ₜ₊₁∂Δtₜ(
-    P::QuantumStatePadeIntegrator,
-    aₜ::AbstractVector,
-    Δtₜ::Real,
-    μₜ::AbstractVector
-)
-    Gₜ = P.G(aₜ, P.G_drift, P.G_drives)
-    B = -1/2 * Gₜ + 1/6 * Δtₜ * Gₜ^2
-    return μₜ' * B
-end
-
-function μ∂²Δtₜ(
-    P::UnitaryPadeIntegrator,
-    Ũ⃗ₜ₊₁::AbstractVector,
-    Ũ⃗ₜ::AbstractVector,
-    aₜ::AbstractVector,
-    μₜ::AbstractVector
-)
-    Gₜ = P.G(aₜ, P.G_drift, P.G_drives)
-    ∂²Δtₜ_gen_block = 1/6 * Gₜ^2
-    ∂²Δtₜ_gen = blockdiag(fill(sparse(∂²Δtₜ_gen_block), P.ketdim)...)
-    ∂²Δtₜ = ∂²Δtₜ_gen * (Ũ⃗ₜ₊₁ -  Ũ⃗ₜ)
-    return μₜ' * ∂²Δtₜ
-end
-
-function μ∂²Δtₜ(
-    P::QuantumStatePadeIntegrator,
-    ψ̃ₜ₊₁::AbstractVector,
-    ψ̃ₜ::AbstractVector,
-    aₜ::AbstractVector,
-    μₜ::AbstractVector
-)
-    Gₜ = P.G(aₜ, P.G_drift, P.G_drives)
-    ∂²Δtₜ = 1/6 * Gₜ^2 * (ψ̃ₜ₊₁ - ψ̃ₜ)
-    return μₜ' * ∂²Δtₜ
-end
-
-@views function hessian_of_the_lagrangian(
-    P::UnitaryPadeIntegrator,
-    zₜ::AbstractVector,
-    zₜ₊₁::AbstractVector,
-    μₜ::AbstractVector,
-    traj::NamedTrajectory
-)
-    free_time = traj.timestep isa Symbol
-
-    Ũ⃗ₜ₊₁ = zₜ₊₁[traj.components[P.unitary_name]]
-    Ũ⃗ₜ = zₜ[traj.components[P.unitary_name]]
-
-    Δtₜ = free_time ? zₜ[traj.components[traj.timestep]][1] : traj.timestep
-
-    if P.drive_name isa Tuple
-        inds = [traj.components[s] for s in P.drive_name]
-        inds = vcat(collect.(inds)...)
-    else
-        inds = traj.components[P.drive_name]
-    end
-
-    aₜ = zₜ[inds]
-
-    μ∂aₜ∂Ũ⃗ₜP = μ∂aₜ∂Ũ⃗ₜ(P, aₜ, Δtₜ, μₜ)
-    μ∂²aₜP = μ∂²aₜ(P, Ũ⃗ₜ₊₁, Ũ⃗ₜ, Δtₜ, μₜ)
-    if free_time
-        μ∂Δtₜ∂aₜP = μ∂Δtₜ∂aₜ(P, Ũ⃗ₜ₊₁, Ũ⃗ₜ, aₜ, Δtₜ, μₜ)
-    end
-
-    μ∂Ũ⃗ₜ₊₁∂aₜP = μ∂Ũ⃗ₜ₊₁∂aₜ(P, aₜ, Δtₜ, μₜ)
-
-    if free_time
-        μ∂Δtₜ∂Ũ⃗ₜP = μ∂Δtₜ∂Ũ⃗ₜ(P, aₜ, Δtₜ, μₜ)
-        μ∂²ΔtₜP = μ∂²Δtₜ(P, Ũ⃗ₜ₊₁, Ũ⃗ₜ, aₜ, μₜ)
-        μ∂Ũ⃗ₜ₊₁∂ΔtₜP = μ∂Ũ⃗ₜ₊₁∂Δtₜ(P, aₜ, Δtₜ, μₜ)
-        return (
-            μ∂aₜ∂Ũ⃗ₜP,
-            μ∂²aₜP,
-            μ∂Δtₜ∂Ũ⃗ₜP,
-            μ∂Δtₜ∂aₜP,
-            μ∂²ΔtₜP,
-            μ∂Ũ⃗ₜ₊₁∂aₜP,
-            μ∂Ũ⃗ₜ₊₁∂ΔtₜP
-        )
-    else
-        return (
-            μ∂aₜ∂Ũ⃗ₜP,
-            μ∂²aₜP,
-            μ∂Ũ⃗ₜ₊₁∂aₜP
-        )
-    end
-end
-
-@views function hessian_of_the_lagrangian(
-    P::QuantumStatePadeIntegrator,
-    zₜ::AbstractVector,
-    zₜ₊₁::AbstractVector,
-    μₜ::AbstractVector,
-    traj::NamedTrajectory
-)
-    free_time = traj.timestep isa Symbol
-
-    ψ̃ₜ₊₁ = zₜ₊₁[traj.components[P.state_name]]
-    ψ̃ₜ = zₜ[traj.components[P.state_name]]
-
-    Δtₜ = free_time ? zₜ[traj.components[traj.timestep]][1] : traj.timestep
-
-    if P.drive_name isa Tuple
-        inds = [traj.components[s] for s in P.drive_name]
-        inds = vcat(collect.(inds)...)
-    else
-        inds = traj.components[P.drive_name]
-    end
-
-    aₜ = zₜ[inds]
-
-    μ∂aₜ∂ψ̃ₜP = μ∂aₜ∂ψ̃ₜ(P, aₜ, Δtₜ, μₜ)
-    μ∂²aₜP = μ∂²aₜ(P, ψ̃ₜ₊₁, ψ̃ₜ, Δtₜ, μₜ)
-    if free_time
-        μ∂Δtₜ∂aₜP = μ∂Δtₜ∂aₜ(P, ψ̃ₜ₊₁, ψ̃ₜ, aₜ, Δtₜ, μₜ)
-    end
-    μ∂ψ̃ₜ₊₁∂aₜP = μ∂ψ̃ₜ₊₁∂aₜ(P, aₜ, Δtₜ, μₜ)
-
-    if free_time
-        μ∂Δtₜ∂ψ̃ₜP = μ∂Δtₜ∂ψ̃ₜ(P, aₜ, Δtₜ, μₜ)
-        μ∂²ΔtₜP = μ∂²Δtₜ(P, ψ̃ₜ₊₁, ψ̃ₜ, aₜ, μₜ)
-        μ∂ψ̃ₜ₊₁∂ΔtₜP = μ∂ψ̃ₜ₊₁∂Δtₜ(P, aₜ, Δtₜ, μₜ)
-
-        return (
-            μ∂aₜ∂ψ̃ₜP,
-            μ∂²aₜP,
-            μ∂Δtₜ∂ψ̃ₜP,
-            μ∂Δtₜ∂aₜP,
-            μ∂²ΔtₜP,
-            μ∂ψ̃ₜ₊₁∂aₜP,
-            μ∂ψ̃ₜ₊₁∂ΔtₜP
-        )
-    else
-        return (
-            μ∂aₜ∂ψ̃ₜP,
-            μ∂²aₜP,
-            μ∂ψ̃ₜ₊₁∂aₜP
-        )
-    end
-end
diff --git a/src/isomorphisms.jl b/src/isomorphisms.jl
deleted file mode 100644
index a057b8d3..00000000
--- a/src/isomorphisms.jl
+++ /dev/null
@@ -1,205 +0,0 @@
-module Isomorphisms
-
-export mat
-export ket_to_iso
-export iso_to_ket
-export iso_vec_to_operator
-export iso_vec_to_iso_operator
-export operator_to_iso_vec
-export iso_operator_to_iso_vec
-export iso_operator_to_operator
-export operator_to_iso_operator
-export iso
-export iso_dm
-export ad_vec
-
-using LinearAlgebra
-using TestItemRunner
-
-
-@doc raw"""
-    mat(x::AbstractVector)
-
-Convert a vector `x` into a square matrix. The length of `x` must be a perfect square.
-"""
-function mat(x::AbstractVector)
-    n = isqrt(length(x))
-    @assert n^2 == length(x) "Vector length must be a perfect square"
-    return reshape(x, n, n)
-end
-
-
-# ----------------------------------------------------------------------------- #
-#                                Kets                                           #
-# ----------------------------------------------------------------------------- #
-
-@doc raw"""
-    ket_to_iso(ψ)
-
-Convert a ket vector `ψ` into a complex vector with real and imaginary parts.
-"""
-ket_to_iso(ψ) = [real(ψ); imag(ψ)]
-
-@doc raw"""
-    iso_to_ket(ψ̃)
-
-Convert a complex vector `ψ̃` with real and imaginary parts into a ket vector.
-"""
-iso_to_ket(ψ̃) = ψ̃[1:div(length(ψ̃), 2)] + im * ψ̃[(div(length(ψ̃), 2) + 1):end]
-
-# ----------------------------------------------------------------------------- #
-#                             Unitaries                                         #
-# ----------------------------------------------------------------------------- #
-
-@doc raw"""
-    iso_vec_to_operator(Ũ⃗::AbstractVector)
-
-Convert a real vector `Ũ⃗` into a complex matrix representing an operator.
-
-Must be differentiable.
-"""
-function iso_vec_to_operator(Ũ⃗::AbstractVector{R}) where R <: Real
-    Ũ⃗_dim = div(length(Ũ⃗), 2)
-    N = Int(sqrt(Ũ⃗_dim))
-    U = Matrix{complex(eltype(Ũ⃗))}(undef, N, N)
-    for i=0:N-1
-        U[:, i+1] .= @view(Ũ⃗[i * 2N .+ (1:N)]) + one(eltype(Ũ⃗)) * im * @view(Ũ⃗[i * 2N .+ (N+1:2N)])
-    end
-    return U
-end
-
-@doc raw"""
-    iso_vec_to_iso_operator(Ũ⃗::AbstractVector)
-
-Convert a real vector `Ũ⃗` into a real matrix representing an isomorphism operator.
-
-Must be differentiable.
-"""
-function iso_vec_to_iso_operator(Ũ⃗::AbstractVector{R}) where R
-    N = Int(sqrt(length(Ũ⃗) ÷ 2))
-    Ũ = Matrix{R}(undef, 2N, 2N)
-    U_real = Matrix{R}(undef, N, N)
-    U_imag = Matrix{R}(undef, N, N)
-    for i=0:N-1
-        U_real[:, i+1] .= @view(Ũ⃗[i*2N .+ (1:N)])
-        U_imag[:, i+1] .= @view(Ũ⃗[i*2N .+ (N+1:2N)])
-    end
-    Ũ[1:N, 1:N] .= U_real
-    Ũ[1:N, (N + 1):end] .= -U_imag
-    Ũ[(N + 1):end, 1:N] .= U_imag
-    Ũ[(N + 1):end, (N + 1):end] .= U_real
-    return Ũ
-end
-
-@doc raw"""
-    operator_to_iso_vec(U::AbstractMatrix{<:Complex})
-
-Convert a complex matrix `U` representing an operator into a real vector.
-
-Must be differentiable.
-"""
-function operator_to_iso_vec(U::AbstractMatrix{R}) where R <: Number
-    N = size(U,1)
-    Ũ⃗ = Vector{real(R)}(undef, N^2 * 2)
-    for i=0:N-1
-        Ũ⃗[i*2N .+ (1:N)] .= real(@view(U[:, i+1]))
-        Ũ⃗[i*2N .+ (N+1:2N)] .= imag(@view(U[:, i+1]))
-    end
-    return Ũ⃗
-end
-
-@doc raw"""
-    iso_operator_to_iso_vec(Ũ::AbstractMatrix)
-
-Convert a real matrix `Ũ` representing an isomorphism operator into a real vector.
-
-Must be differentiable.
-"""
-function iso_operator_to_iso_vec(Ũ::AbstractMatrix{R}) where R
-    N = size(Ũ, 1) ÷ 2
-    Ũ⃗ = Vector{R}(undef, N^2 * 2)
-    for i=0:N-1
-        Ũ⃗[i*2N .+ (1:2N)] .= @view Ũ[:, i+1]
-    end
-    return Ũ⃗
-end
-
-iso_operator_to_operator(Ũ) = iso_vec_to_operator(iso_operator_to_iso_vec(Ũ))
-
-operator_to_iso_operator(U) = iso_vec_to_iso_operator(operator_to_iso_vec(U))
-
-# ----------------------------------------------------------------------------- #
-# Open systems
-# ----------------------------------------------------------------------------- #
-
-function ad_vec(H::AbstractMatrix{<:Number}; anti::Bool=false)
-    Id = sparse(eltype(H), I, size(H)...)
-    return Id ⊗ H - (-1)^anti * conj(H)' ⊗ Id
-end
-
-# ----------------------------------------------------------------------------- #
-# Hamiltonians
-# ----------------------------------------------------------------------------- #
-
-const Im2 = [
-    0 -1;
-    1  0
-]
-
-@doc raw"""
-    G(H::AbstractMatrix)::Matrix{Float64}
-
-Returns the isomorphism of ``-iH``:
-
-```math
-G(H) = \widetilde{- i H} = \mqty(1 & 0 \\ 0 & 1) \otimes \Im(H) - \mqty(0 & -1 \\ 1 & 0) \otimes \Re(H)
-```
-
-where ``\Im(H)`` and ``\Re(H)`` are the imaginary and real parts of ``H`` and the tilde indicates the standard isomorphism of a complex valued matrix:
-
-```math
-\widetilde{H} = \mqty(1 & 0 \\ 0 & 1) \otimes \Re(H) + \mqty(0 & -1 \\ 1 & 0) \otimes \Im(H)
-```
-"""
-G(H::AbstractMatrix{<:Number}) = kron(I(2), imag(H)) - kron(Im2, real(H))
-
-iso(H::AbstractMatrix{<:Number}) = kron(I(2), real(H)) + kron(Im2, imag(H))
-
-
-"""
-    H(G::AbstractMatrix{<:Number})::Matrix{ComplexF64}
-
-Returns the inverse of `G(H) = iso(-iH)`, i.e. returns H
-
-"""
-function H(G::AbstractMatrix{<:Number})
-    dim = size(G, 1) ÷ 2
-    H_imag = G[1:dim, 1:dim]
-    H_real = -G[dim+1:end, 1:dim]
-    return H_real + 1.0im * H_imag
-end
-
-"""
-    iso_dm(ρ::AbstractMatrix)
-
-returns the isomorphism `ρ⃗̃ = ket_to_iso(vec(ρ))` of a density matrix `ρ`
-"""
-iso_dm(ρ::AbstractMatrix) = ket_to_iso(vec(ρ))
-
-
-
-# *************************************************************************** #
-
-@testitem "Test isomorphism utilities" begin
-    using LinearAlgebra
-    iso_vec = [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0]
-    @test mat([1.0, 2.0, 3.0, 4.0]) == [1.0 3.0; 2.0 4.0]
-    @test ket_to_iso([1.0, 2.0]) == [1.0, 2.0, 0.0, 0.0]
-    @test iso_to_ket([1.0, 2.0, 0.0, 0.0]) == [1.0, 2.0]
-    @test iso_vec_to_operator(iso_vec) == [1.0 0.0; 0.0 1.0]
-    @test iso_vec_to_iso_operator(iso_vec) == [1.0 0.0 -0.0 -0.0; 0.0 1.0 -0.0 -0.0; 0.0 0.0 1.0 0.0; 0.0 0.0 0.0 1.0]
-    @test operator_to_iso_vec(Complex[1.0 0.0; 0.0 1.0]) == iso_vec
-    @test iso_operator_to_iso_vec(iso_vec_to_iso_operator(iso_vec)) == iso_vec
-end
-
-end
diff --git a/src/losses/_experimental_loss_functions.jl b/src/losses/_experimental_loss_functions.jl
deleted file mode 100644
index e6de76e5..00000000
--- a/src/losses/_experimental_loss_functions.jl
+++ /dev/null
@@ -1,51 +0,0 @@
-#
-# experimental loss functions
-#
-# TODO: renormalize vectors in place of abs
-#       ⋅ penalize loss to remain near unit norm
-#       ⋅ Σ α * (1 - ψ̃'ψ̃), α = 1e-3
-
-function energy_loss(
-    ψ̃::AbstractVector,
-    H::AbstractMatrix
-)
-    ψ = iso_to_ket(ψ̃)
-    return real(ψ' * H * ψ)
-end
-
-# TODO: figure out a way to implement this without erroring and Von Neumann entropy being always 0 for a pure state
-function neg_entropy_loss(
-    ψ̃::AbstractVector
-)
-    ψ = iso_to_ket(ψ̃)
-    ρ = ψ * ψ'
-    ρ = Hermitian(ρ)
-    return tr(ρ * log(ρ))
-end
-
-function pure_real_loss(ψ̃, ψ̃goal)
-    ψ = iso_to_ket(ψ̃)
-    ψgoal = iso_to_ket(ψ̃goal)
-    return -(ψ'ψgoal)
-end
-
-function geodesic_loss(ψ̃, ψ̃goal)
-    ψ = iso_to_ket(ψ̃)
-    ψgoal = iso_to_ket(ψ̃goal)
-    amp = ψ'ψgoal
-    return min(abs(1 - amp), abs(1 + amp))
-end
-
-function real_loss(ψ̃, ψ̃goal)
-    ψ = iso_to_ket(ψ̃)
-    ψgoal = iso_to_ket(ψ̃goal)
-    amp = ψ'ψgoal
-    return min(abs(1 - real(amp)), abs(1 + real(amp)))
-end
-
-function quaternionic_loss(ψ̃, ψ̃goal)
-    return min(
-        abs(1 - dot(ψ̃, ψ̃goal)),
-        abs(1 + dot(ψ̃, ψ̃goal))
-    )
-end
\ No newline at end of file
diff --git a/src/losses/_losses.jl b/src/losses/_losses.jl
deleted file mode 100644
index 32225c70..00000000
--- a/src/losses/_losses.jl
+++ /dev/null
@@ -1,89 +0,0 @@
-module Losses
-
-export Loss
-
-using ..QuantumObjectUtils
-using ..Isomorphisms
-using ..QuantumSystems
-using ..StructureUtils
-
-using NamedTrajectories
-using TrajectoryIndexingUtils
-
-using LinearAlgebra
-using SparseArrays
-using ForwardDiff
-using Symbolics
-using TestItemRunner
-
-# For differentiable exponential
-using ExponentialAction
-
-# TODO:
-# - [ ] Do not reference the Z object in the loss (components only / remove "name")
-
-# ----------------------------------------------------------------------------- #
-#                           Abstract Loss                                       #
-# ----------------------------------------------------------------------------- #
-
-abstract type AbstractLoss end
-
-include("_experimental_loss_functions.jl")
-include("quantum_loss.jl")
-include("quantum_state_infidelity_loss.jl")
-include("unitary_trace_loss.jl")
-include("unitary_infidelity_loss.jl")
-
-# ----------------------------------------------------------------------------- #
-#                               Loss                                            #
-# ----------------------------------------------------------------------------- #
-
-struct Loss <: AbstractLoss
-    l::Function
-    ∇l::Function
-    ∇²l::Function
-    ∇²l_structure::Vector{Tuple{Int,Int}}
-    name::Symbol
-
-    function Loss(
-        Z::NamedTrajectory,
-        J::Function,
-        x::Symbol
-    )
-        @assert x ∈ Z.names
-        @assert Z.goal[x] isa AbstractVector
-
-        x_goal = Z.goal[x]
-
-        J = x̄ -> J(x̄, x_goal)
-        ∇J = x̄ -> ForwardDiff.gradient(J, x̄)
-
-        Symbolics.@variables x̄[1:Z.dims[x]]
-        x̄ = collect(x̄)
-
-        ∇²J_symbolic = Symbolics.sparsehessian(J(x̄), x̄)
-        rows, cols, _ = findnz(∇²J_symbolic)
-        rowcols = collect(zip(rows, cols))
-        filter!((row, col) -> row ≥ col, rowcols)
-        ∇²J_structure = rowcols
-
-        ∇²J_expression = Symbolics.build_function(∇²J_symbolic, x̄)
-        ∇²J = eval(∇²J_expression[1])
-
-        return new(J, ∇J, ∇²J, ∇²J_structure, x)
-    end
-end
-
-function (loss::Loss)(Z::NamedTrajectory; gradient=false, hessian=false)
-    @assert !(gradient && hessian)
-    if !(gradient || hessian)
-        return loss.l(Z[end][loss.name])
-    elseif gradient
-        return loss.∇l(Z[end][loss.name])
-    elseif hessian
-        return loss.∇²l(Z[end][loss.name])
-    end
-end
-
-
-end
diff --git a/src/losses/quantum_loss.jl b/src/losses/quantum_loss.jl
deleted file mode 100644
index 95283415..00000000
--- a/src/losses/quantum_loss.jl
+++ /dev/null
@@ -1,169 +0,0 @@
-export QuantumLoss
-export QuantumLossGradient
-export QuantumLossHessian
-
-
-struct QuantumLoss
-    cs::Vector{Function}
-    isodim::Int
-
-    function QuantumLoss(
-        sys::AbstractQuantumSystem,
-        loss::Symbol = :infidelity_loss
-    )
-        if loss == :energy_loss
-            cs = [ψ̃ⁱ -> eval(loss)(ψ̃ⁱ, sys.H_target) for i = 1:sys.nqstates]
-        elseif loss == :neg_entropy_loss
-            cs = [ψ̃ⁱ -> eval(loss)(ψ̃ⁱ) for i = 1:sys.nqstates]
-        else
-            cs = [
-                ψ̃ⁱ -> eval(loss)(
-                    ψ̃ⁱ,
-                    sys.ψ̃goal[slice(i, sys.isodim)]
-                ) for i = 1:sys.nqstates
-            ]
-        end
-        return new(cs, sys.isodim)
-    end
-end
-
-function (qloss::QuantumLoss)(ψ̃::AbstractVector)
-    loss = 0.0
-    for (i, cⁱ) in enumerate(qloss.cs)
-        loss += cⁱ(ψ̃[slice(i, qloss.isodim)])
-    end
-    return loss
-end
-
-struct QuantumLossGradient
-    ∇cs::Vector{Function}
-    isodim::Int
-
-    function QuantumLossGradient(
-        loss::QuantumLoss;
-        simplify=true
-    )
-        Symbolics.@variables ψ̃[1:loss.isodim]
-
-        ψ̃ = collect(ψ̃)
-
-        ∇cs_symbs = [
-            Symbolics.gradient(c(ψ̃), ψ̃; simplify=simplify)
-                for c in loss.cs
-        ]
-
-        ∇cs_exprs = [
-            Symbolics.build_function(∇c, ψ̃)
-                for ∇c in ∇cs_symbs
-        ]
-
-        ∇cs = [
-            eval(∇c_expr[1])
-                for ∇c_expr in ∇cs_exprs
-        ]
-
-        return new(∇cs, loss.isodim)
-    end
-end
-
-@views function (∇c::QuantumLossGradient)(
-    ψ̃::AbstractVector
-)
-    ∇ = similar(ψ̃)
-
-    for (i, ∇cⁱ) in enumerate(∇c.∇cs)
-
-        ψ̃ⁱ_slice = slice(i, ∇c.isodim)
-
-        ∇[ψ̃ⁱ_slice] = ∇cⁱ(ψ̃[ψ̃ⁱ_slice])
-    end
-
-    return ∇
-end
-
-struct QuantumLossHessian
-    ∇²cs::Vector{Function}
-    ∇²c_structures::Vector{Vector{Tuple{Int, Int}}}
-    isodim::Int
-
-    function QuantumLossHessian(
-        loss::QuantumLoss;
-        simplify=true
-    )
-
-        Symbolics.@variables ψ̃[1:loss.isodim]
-        ψ̃ = collect(ψ̃)
-
-        ∇²c_symbs = [
-            Symbolics.sparsehessian(
-                c(ψ̃),
-                ψ̃;
-                simplify=simplify
-            ) for c in loss.cs
-        ]
-
-        ∇²c_structures = []
-
-        for ∇²c_symb in ∇²c_symbs
-            K, J, _ = findnz(∇²c_symb)
-
-            KJ = collect(zip(K, J))
-
-            filter!(((k, j),) -> k ≤ j, KJ)
-
-            push!(∇²c_structures, KJ)
-        end
-
-        ∇²c_exprs = [
-            Symbolics.build_function(∇²c_symb, ψ̃)
-                for ∇²c_symb in ∇²c_symbs
-        ]
-
-        ∇²cs = [
-            eval(∇²c_expr[1])
-                for ∇²c_expr in ∇²c_exprs
-        ]
-
-        return new(∇²cs, ∇²c_structures, loss.isodim)
-    end
-end
-
-function StructureUtils.structure(
-    H::QuantumLossHessian,
-    T::Int,
-    vardim::Int
-)
-    H_structure = []
-
-    T_offset = index(T, 0, vardim)
-
-    for (i, KJⁱ) in enumerate(H.∇²c_structures)
-
-        i_offset = index(i, 0, H.isodim)
-
-        for kj in KJⁱ
-            push!(H_structure, (T_offset + i_offset) .+ kj)
-        end
-    end
-
-    return H_structure
-end
-
-@views function (H::QuantumLossHessian)(ψ̃::AbstractVector)
-
-    Hs = []
-
-    for (i, ∇²cⁱ) in enumerate(H.∇²cs)
-
-        ψ̃ⁱ = ψ̃[slice(i, H.isodim)]
-
-        for (k, j) in H.∇²c_structures[i]
-
-            Hⁱᵏʲ = ∇²cⁱ(ψ̃ⁱ)[k, j]
-
-            append!(Hs, Hⁱᵏʲ)
-        end
-    end
-
-    return Hs
-end
\ No newline at end of file
diff --git a/src/losses/quantum_state_infidelity_loss.jl b/src/losses/quantum_state_infidelity_loss.jl
deleted file mode 100644
index e89660df..00000000
--- a/src/losses/quantum_state_infidelity_loss.jl
+++ /dev/null
@@ -1,99 +0,0 @@
-export fidelity
-export iso_fidelity
-
-export InfidelityLoss
-
-###
-### InfidelityLoss
-###
-
-@doc raw"""
-    fidelity(ψ, ψ_goal)
-
-Calculate the fidelity between two quantum states `ψ` and `ψ_goal`.
-"""
-function fidelity(
-    ψ::AbstractVector, 
-    ψ_goal::AbstractVector;
-    subspace::AbstractVector{Int}=1:length(ψ)
-)
-    ψ = ψ[subspace]
-    ψ_goal = ψ_goal[subspace]
-    return abs2(ψ_goal' * ψ)
-end
-
-@doc raw"""
-    iso_fidelity(ψ̃, ψ̃_goal)
-
-Calculate the fidelity between two quantum states in their isomorphic form `ψ̃` and `ψ̃_goal`.
-"""
-function iso_fidelity(
-    ψ̃::AbstractVector, 
-    ψ̃_goal::AbstractVector;
-    subspace::AbstractVector{Int}=1:length(iso_to_ket(ψ̃))
-)
-    ψ = iso_to_ket(ψ̃)
-    ψ_goal = iso_to_ket(ψ̃_goal)
-    return fidelity(ψ, ψ_goal, subspace=subspace)
-end
-
-"""
-    iso_infidelity(ψ̃, ψ̃goal)
-
-Returns the iso_infidelity between two quantum statevectors specified
-in the ``\\mathbb{C}^n \\to \\mathbb{R}^{2n}`` isomorphism space.
-
-"""
-function iso_infidelity(
-    ψ̃::AbstractVector, 
-    ψ̃goal::AbstractVector,
-    subspace::AbstractVector{Int}=1:length(iso_to_ket(ψ̃))
-)
-    return abs(1 - iso_fidelity(ψ̃, ψ̃goal, subspace=subspace))
-end
-
-struct InfidelityLoss <: AbstractLoss
-    l::Function
-    ∇l::Function
-    ∇²l::Function
-    ∇²l_structure::Vector{Tuple{Int,Int}}
-    wfn_name::Symbol
-
-    function InfidelityLoss(
-        name::Symbol,
-        ψ̃_goal::AbstractVector
-    )
-        l = ψ̃ -> iso_infidelity(ψ̃, ψ̃_goal)
-        ∇l = ψ̃ -> ForwardDiff.gradient(l, ψ̃)
-
-        Symbolics.@variables ψ̃[1:length(ψ̃_goal)]
-        ψ̃ = collect(ψ̃)
-
-        ∇²l_symbolic = Symbolics.sparsehessian(l(ψ̃), ψ̃)
-        K, J, _ = findnz(∇²l_symbolic)
-        kjs = collect(zip(K, J))
-        filter!(((k, j),) -> k ≤ j, kjs)
-        ∇²l_structure = kjs
-
-        ∇²l_expression = Symbolics.build_function(∇²l_symbolic, ψ̃)
-        ∇²l = eval(∇²l_expression[1])
-
-        return new(l, ∇l, ∇²l, ∇²l_structure, name)
-    end
-end
-
-function (loss::InfidelityLoss)(
-    ψ̃_end::AbstractVector{<:Real};
-    gradient=false,
-    hessian=false
-)
-    @assert !(gradient && hessian)
-
-    if !(gradient || hessian)
-        return loss.l(ψ̃_end)
-    elseif gradient
-        return loss.∇l(ψ̃_end)
-    elseif hessian
-        return loss.∇²l(ψ̃_end)
-    end
-end
\ No newline at end of file
diff --git a/src/losses/unitary_infidelity_loss.jl b/src/losses/unitary_infidelity_loss.jl
deleted file mode 100644
index df0eb434..00000000
--- a/src/losses/unitary_infidelity_loss.jl
+++ /dev/null
@@ -1,553 +0,0 @@
-export unitary_fidelity
-export iso_vec_unitary_fidelity
-export unitary_free_phase_fidelity
-export iso_vec_unitary_free_phase_fidelity
-
-export UnitaryInfidelityLoss
-export UnitaryFreePhaseInfidelityLoss
-
-###
-### UnitaryInfidelityLoss
-###
-
-@doc raw"""
-    unitary_fidelity(U::Matrix, U_goal::Matrix; kwargs...)
-    unitary_fidelity(Ũ⃗::AbstractVector, Ũ⃗_goal::AbstractVector; kwargs...)
-
-Calculate the fidelity between two unitary operators `U` and `U_goal`.
-
-```math
-\mathcal{F}(U, U_{\text{goal}}) = \frac{1}{n} \abs{\tr \qty(U_{\text{goal}}^\dagger U)}
-```
-
-where $n$ is the dimension of the unitary operators.
-
-# Keyword Arguments
-- `subspace::AbstractVector{Int}`: The subspace to calculate the fidelity over.
-"""
-@views function unitary_fidelity(
-    U::Matrix,
-    U_goal::Matrix;
-    subspace::AbstractVector{Int}=axes(U_goal, 1)
-)
-    # avoids type coercion neceessary because tr(Matrix{Any}) fails
-    return iso_vec_unitary_fidelity(
-        operator_to_iso_vec(U),
-        operator_to_iso_vec(U_goal),
-        subspace=subspace
-    )
-end
-
-
-@doc raw"""
-    iso_vec_unitary_fidelity(Ũ⃗::AbstractVector, Ũ⃗_goal::AbstractVector)
-
-Returns the fidelity between the isomorphic unitary vector $\vec{\widetilde{U}} \sim U \in SU(n)$
-and the isomorphic goal unitary vector $\vec{\widetilde{U}}_{\text{goal}}$.
-
-```math
-\begin{aligned}
-\mathcal{F}(\vec{\widetilde{U}}, \vec{\widetilde{U}}_{\text{goal}}) &= \frac{1}{n} \abs{\tr \qty(U_{\text{goal}}^\dagger U)} \\
-&= \frac{1}{n} \sqrt{T_R^{2} + T_I^{2}}
-\end{aligned}
-```
-
-where $T_R = \langle \vec{\widetilde{U}}_{\text{goal}, R}, \vec{\widetilde{U}}_R \rangle + \langle \vec{\widetilde{U}}_{\text{goal}, I}, \vec{\widetilde{U}}_I \rangle$ and $T_I = \langle \vec{\widetilde{U}}_{\text{goal}, R}, \vec{\widetilde{U}}_I \rangle - \langle \vec{\widetilde{U}}_{\text{goal}, I}, \vec{\widetilde{U}}_R \rangle$.
-
-"""
-@inline @views function iso_vec_unitary_fidelity(
-    Ũ⃗::AbstractVector,
-    Ũ⃗_goal::AbstractVector;
-    subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1)
-)
-    U = iso_vec_to_operator(Ũ⃗)[subspace, subspace]
-    n = size(U, 1)
-    U_goal = iso_vec_to_operator(Ũ⃗_goal)[subspace, subspace]
-    U⃗ᵣ, U⃗ᵢ = vec(real(U)), vec(imag(U))
-    Ū⃗ᵣ, Ū⃗ᵢ = vec(real(U_goal)), vec(imag(U_goal))
-    Tᵣ = Ū⃗ᵣ' * U⃗ᵣ + Ū⃗ᵢ' * U⃗ᵢ
-    Tᵢ = Ū⃗ᵣ' * U⃗ᵢ - Ū⃗ᵢ' * U⃗ᵣ
-    return 1 / n * sqrt(Tᵣ^2 + Tᵢ^2)
-end
-
-@views function iso_vec_unitary_infidelity(
-    Ũ⃗::AbstractVector,
-    Ũ⃗_goal::AbstractVector;
-    subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1)
-)
-    ℱ = iso_vec_unitary_fidelity(Ũ⃗, Ũ⃗_goal, subspace=subspace)
-    return abs(1 - ℱ)
-end
-
-@views function iso_vec_unitary_infidelity_gradient(
-    Ũ⃗::AbstractVector,
-    Ũ⃗_goal::AbstractVector;
-    subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1)
-)
-    U = iso_vec_to_operator(Ũ⃗)[subspace, subspace]
-    n = size(U, 1)
-    U_goal = iso_vec_to_operator(Ũ⃗_goal)[subspace, subspace]
-    U⃗ᵣ, U⃗ᵢ = vec(real(U)), vec(imag(U))
-    Ū⃗ᵣ, Ū⃗ᵢ = vec(real(U_goal)), vec(imag(U_goal))
-    Tᵣ = Ū⃗ᵣ' * U⃗ᵣ + Ū⃗ᵢ' * U⃗ᵢ
-    Tᵢ = Ū⃗ᵣ' * U⃗ᵢ - Ū⃗ᵢ' * U⃗ᵣ
-    ℱ = 1 / n * sqrt(Tᵣ^2 + Tᵢ^2)
-    ∇ᵣℱ = 1 / (n^2 * ℱ) * (Tᵣ * Ū⃗ᵣ - Tᵢ * Ū⃗ᵢ)
-    ∇ᵢℱ = 1 / (n^2 * ℱ) * (Tᵣ * Ū⃗ᵢ + Tᵢ * Ū⃗ᵣ)
-    ∇ℱ = [∇ᵣℱ; ∇ᵢℱ]
-    permutation = vcat(vcat([[slice(j, n), slice(j, n) .+ n^2] for j = 1:n]...)...)
-    ∇ℱ = ∇ℱ[permutation]
-    return -sign(1 - ℱ) * ∇ℱ
-end
-
-@views function iso_vec_unitary_infidelity_hessian(
-    Ũ⃗::AbstractVector,
-    Ũ⃗_goal::AbstractVector;
-    subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1)
-)
-    U = iso_vec_to_operator(Ũ⃗)[subspace, subspace]
-    n = size(U, 1)
-    U_goal = iso_vec_to_operator(Ũ⃗_goal)[subspace, subspace]
-    U⃗ᵣ, U⃗ᵢ = vec(real(U)), vec(imag(U))
-    Ū⃗ᵣ, Ū⃗ᵢ = vec(real(U_goal)), vec(imag(U_goal))
-    Tᵣ = Ū⃗ᵣ' * U⃗ᵣ + Ū⃗ᵢ' * U⃗ᵢ
-    Tᵢ = Ū⃗ᵣ' * U⃗ᵢ - Ū⃗ᵢ' * U⃗ᵣ
-    Wᵣᵣ = Ū⃗ᵣ * Ū⃗ᵣ'
-    Wᵢᵢ = Ū⃗ᵢ * Ū⃗ᵢ'
-    Wᵣᵢ = Ū⃗ᵣ * Ū⃗ᵢ'
-    Wᵢᵣ = Wᵣᵢ'
-    ℱ = 1 / n * sqrt(Tᵣ^2 + Tᵢ^2)
-    ∇ᵣℱ = 1 / (n^2 * ℱ) * (Tᵣ * Ū⃗ᵣ - Tᵢ * Ū⃗ᵢ)
-    ∇ᵢℱ = 1 / (n^2 * ℱ) * (Tᵣ * Ū⃗ᵢ + Tᵢ * Ū⃗ᵣ)
-    ∂ᵣ²ℱ = 1 / ℱ * (-∇ᵣℱ * ∇ᵣℱ' + 1 / n^2 * (Wᵣᵣ + Wᵢᵢ))
-    ∂ᵢ²ℱ = 1 / ℱ * (-∇ᵢℱ * ∇ᵢℱ' + 1 / n^2 * (Wᵣᵣ + Wᵢᵢ))
-    ∂ᵣ∂ᵢℱ = 1 / ℱ * (-∇ᵢℱ * ∇ᵣℱ' + 1 / n^2 * (Wᵢᵣ - Wᵣᵢ))
-    ∂²ℱ = [∂ᵣ²ℱ ∂ᵣ∂ᵢℱ; ∂ᵣ∂ᵢℱ' ∂ᵢ²ℱ]
-    # TODO: This should be moved to Isomorphisms.jl
-    permutation = vcat(vcat([[slice(j, n), slice(j, n) .+ n^2] for j = 1:n]...)...)
-    ∂²ℱ = ∂²ℱ[permutation, permutation]
-    return -sign(1 - ℱ) * ∂²ℱ
-end
-
-struct UnitaryInfidelityLoss <: AbstractLoss
-    l::Function
-    ∇l::Function
-    ∇²l::Function
-    ∇²l_structure::Vector{Tuple{Int,Int}}
-    name::Symbol
-
-    function UnitaryInfidelityLoss(
-        name::Symbol,
-        Ũ⃗_goal::AbstractVector;
-        subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1)
-    )
-        l = Ũ⃗ -> iso_vec_unitary_infidelity(Ũ⃗, Ũ⃗_goal, subspace=subspace)
-
-        @views ∇l = Ũ⃗ -> begin
-            subspace_rows, subspace_cols = (subspace, subspace)
-            n_subspace = length(subspace_rows)
-            n_full = Int(sqrt(length(Ũ⃗) ÷ 2))
-            ∇l_subspace = iso_vec_unitary_infidelity_gradient(Ũ⃗, Ũ⃗_goal, subspace=subspace)
-            ∇l_full = zeros(2 * n_full^2)
-            for j ∈ eachindex(subspace_cols)
-                ∇l_full[slice(2subspace_cols[j] - 1, subspace_rows, n_full)] =
-                    ∇l_subspace[slice(2j - 1, n_subspace)]
-                ∇l_full[slice(2subspace_cols[j], subspace_rows, n_full)] =
-                    ∇l_subspace[slice(2j, n_subspace)]
-            end
-            return ∇l_full
-        end
-
-        @views ∇²l = Ũ⃗ -> begin
-            subspace_rows, subspace_cols = (subspace, subspace)
-            n_subspace = length(subspace_rows)
-            n_full = Int(sqrt(length(Ũ⃗) ÷ 2))
-            ∇²l_subspace = iso_vec_unitary_infidelity_hessian(Ũ⃗, Ũ⃗_goal, subspace=subspace)
-            ∇²l_full = zeros(2 * n_full^2, 2 * n_full^2)
-            # NOTE: Assumes subspace_rows = subspace_cols
-            for k ∈ eachindex(subspace_cols)
-                for j ∈ eachindex(subspace_cols)
-                    ∇²l_full[
-                        slice(2subspace_cols[k] - 1, subspace_rows, n_full),
-                        slice(2subspace_cols[j] - 1, subspace_rows, n_full)
-                    ] = ∇²l_subspace[
-                        slice(2k - 1, n_subspace),
-                        slice(2j - 1, n_subspace)
-                    ]
-
-                    ∇²l_full[
-                        slice(2subspace_cols[k] - 1, subspace_rows, n_full),
-                        slice(2subspace_cols[j], subspace_rows, n_full)
-                    ] = ∇²l_subspace[
-                        slice(2k - 1, n_subspace),
-                        slice(2j, n_subspace)
-                    ]
-
-                    ∇²l_full[
-                        slice(2subspace_cols[k], subspace_rows, n_full),
-                        slice(2subspace_cols[j] - 1, subspace_rows, n_full)
-                    ] = ∇²l_subspace[
-                        slice(2k, n_subspace),
-                        slice(2j - 1, n_subspace)
-                    ]
-
-                    ∇²l_full[
-                        slice(2subspace_cols[k], subspace_rows, n_full),
-                        slice(2subspace_cols[j], subspace_rows, n_full)
-                    ] = ∇²l_subspace[
-                        slice(2k, n_subspace),
-                        slice(2j, n_subspace)
-                    ]
-                end
-            end
-            return ∇²l_full
-        end
-
-        Ũ⃗_dim = length(Ũ⃗_goal)
-        ∇²l_structure = []
-        for (i, j) ∈ Iterators.product(1:Ũ⃗_dim, 1:Ũ⃗_dim)
-            if i ≤ j
-                push!(∇²l_structure, (i, j))
-            end
-        end
-        return new(l, ∇l, ∇²l, ∇²l_structure, name)
-    end
-end
-
-function (loss::UnitaryInfidelityLoss)(
-    Ũ⃗_end::AbstractVector{<:Real};
-    gradient=false,
-    hessian=false
-)
-    @assert !(gradient && hessian)
-    if !(gradient || hessian)
-        return loss.l(Ũ⃗_end)
-    elseif gradient
-        return loss.∇l(Ũ⃗_end)
-    elseif hessian
-        return loss.∇²l(Ũ⃗_end)
-    end
-end
-
-###
-### UnitaryFreePhaseInfidelityLoss
-###
-
-function free_phase(
-    ϕs::AbstractVector,
-    Hs::AbstractVector{<:AbstractMatrix}
-)
-    # NOTE: switch to expv for ForwardDiff
-    # return reduce(kron, [exp(im * ϕ * H) for (ϕ, H) ∈ zip(ϕs, Hs)])
-    Id = Matrix{eltype(Hs[1])}(I, size(Hs[1]))
-    return reduce(kron, [expv(im * ϕ, H, Id) for (ϕ, H) ∈ zip(ϕs, Hs)])
-end
-
-function free_phase_gradient(
-    ϕs::AbstractVector,
-    Hs::AbstractVector{<:AbstractMatrix}
-)
-    R = free_phase(ϕs, Hs)
-    result = [zeros(eltype(R), size(R)) for _ in eachindex(ϕs)]
-
-    # store identity matrices
-    identities = [Matrix{eltype(H)}(I, size(H)) for H in Hs]
-
-    for (i, H) in enumerate(Hs)
-        # insert H into identities
-        identities[i] .= H
-        result[i] .= im * reduce(kron, identities) * R
-        # reset identities
-        identities[i] .= Matrix{eltype(H)}(I, size(H))
-    end
-    return result
-end
-
-@views function unitary_free_phase_fidelity(
-    U::Matrix,
-    U_goal::Matrix,
-    phases::AbstractVector,
-    phase_operators::AbstractVector{<:AbstractMatrix};
-    subspace::AbstractVector{Int}=axes(U_goal, 1)
-)
-    # extract phase rotation (assume phase operators span goal subspace)
-    R = zeros(eltype(U), size(U))
-    R[subspace, subspace] = free_phase(phases, phase_operators)
-
-    # calculate fidelity
-    return unitary_fidelity(R * U, U_goal, subspace=subspace)
-end
-
-@views function iso_vec_unitary_free_phase_fidelity(
-    Ũ⃗::AbstractVector,
-    Ũ⃗_goal::AbstractVector,
-    phases::AbstractVector,
-    phase_operators::AbstractVector{<:AbstractMatrix};
-    subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1)
-)
-    U = iso_vec_to_operator(Ũ⃗)
-    U_goal = iso_vec_to_operator(Ũ⃗_goal)
-    return unitary_free_phase_fidelity(U, U_goal, phases, phase_operators, subspace=subspace)
-end
-
-@views function iso_vec_unitary_free_phase_infidelity(
-    Ũ⃗::AbstractVector,
-    Ũ⃗_goal::AbstractVector,
-    phases::AbstractVector,
-    phase_operators::AbstractVector{<:AbstractMatrix};
-    subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1)
-)
-    ℱ = iso_vec_unitary_free_phase_fidelity(
-        Ũ⃗, Ũ⃗_goal, phases, phase_operators, 
-        subspace=subspace
-    )
-    return abs(1 - ℱ)
-end
-
-@views function iso_vec_unitary_free_phase_infidelity_gradient(
-    Ũ⃗::AbstractVector,
-    Ũ⃗_goal::AbstractVector,
-    phases::AbstractVector,
-    phase_operators::AbstractVector{<:AbstractMatrix};
-    subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1)
-)
-    n_phases = length(phases)
-    n_subspace = 2 * length(subspace) * length(subspace)
-    ∂ = spzeros(n_subspace + n_phases)
-
-    # extract full state
-    U = iso_vec_to_operator(Ũ⃗)
-
-    # extract full phase rotation
-    R = zeros(eltype(U), size(U))
-    R[subspace, subspace] = free_phase(phases, phase_operators)
-
-    # loss gradient in subspace
-    ∂ℱ_∂Ũ⃗_subspace = iso_vec_unitary_infidelity_gradient(
-        operator_to_iso_vec(R * U), Ũ⃗_goal, subspace=subspace
-    )
-
-    # state slice in subspace
-    ∂[1:n_subspace] = operator_to_iso_vec(R[subspace, subspace]'iso_vec_to_operator(∂ℱ_∂Ũ⃗_subspace))
-
-    # phase slice
-    ∂[n_subspace .+ (1:n_phases)] = [
-        ∂ℱ_∂Ũ⃗_subspace'operator_to_iso_vec(∂R_∂ϕⱼ_subspace * U[subspace, subspace])
-        for ∂R_∂ϕⱼ_subspace in free_phase_gradient(phases, phase_operators)
-    ]
-
-    return ∂
-end
-
-struct UnitaryFreePhaseInfidelityLoss <: AbstractLoss
-    l::Function
-    ∇l::Function
-    ∇²l::Function
-    ∇²l_structure::Vector{Tuple{Int,Int}}
-    name::Symbol
-
-    function UnitaryFreePhaseInfidelityLoss(
-        Ũ⃗_goal::AbstractVector,
-        phase_operators::AbstractVector{<:AbstractMatrix};
-        subspace::Union{AbstractVector{Int}, Nothing}=nothing,
-    )
-        if isnothing(subspace)
-            subspace = axes(iso_vec_to_operator(Ũ⃗_goal), 1)
-        end
-
-        @assert reduce(*, size.(phase_operators, 1)) == length(subspace) "phase operators must span the subspace"
-
-        @views function l(Ũ⃗::AbstractVector, ϕ⃗::AbstractVector)
-            return iso_vec_unitary_free_phase_infidelity(Ũ⃗, Ũ⃗_goal, ϕ⃗, phase_operators, subspace=subspace)
-        end
-
-        @views function ∇l(Ũ⃗::AbstractVector, ϕ⃗::AbstractVector)
-            subspace_rows = subspace_cols = subspace
-            n_phase = length(ϕ⃗)
-            n_rows = length(subspace_rows)
-            n_cols = length(subspace_cols)
-            n_full = Int(sqrt(length(Ũ⃗) ÷ 2))
-
-            # gradient in subspace
-            ∇l_subspace = iso_vec_unitary_free_phase_infidelity_gradient(
-                Ũ⃗, Ũ⃗_goal, ϕ⃗, phase_operators, subspace=subspace
-            )
-            ∇l_full = zeros(2 * n_full^2 + n_phase)
-
-            # state slice
-            for j ∈ 1:n_cols
-                ∇l_full[slice(2subspace_cols[j] - 1, subspace_rows, n_full)] =
-                    ∇l_subspace[slice(2j - 1, n_rows)]
-                ∇l_full[slice(2subspace_cols[j], subspace_rows, n_full)] =
-                    ∇l_subspace[slice(2j, n_rows)]
-            end
-
-            # phase slice
-            ∇l_full[2 * n_full^2 .+ (1:n_phase)] = ∇l_subspace[2 * n_rows * n_cols .+ (1:n_phase)]
-            return ∇l_full
-        end
-
-        # TODO: implement analytic hessian
-        ∇²l(Ũ⃗, ϕ⃗) = []
-        ∇²l_structure = []
-
-        return new(l, ∇l, ∇²l, ∇²l_structure, :NA)
-    end
-end
-
-function (loss::UnitaryFreePhaseInfidelityLoss)(
-    Ũ⃗_end::AbstractVector{<:Real},
-    ϕ⃗::AbstractVector{<:Real};
-    gradient=false,
-    hessian=false
-)
-    @assert !(gradient && hessian)
-    if !(gradient || hessian)
-        return loss.l(Ũ⃗_end, ϕ⃗)
-    elseif gradient
-        return loss.∇l(Ũ⃗_end, ϕ⃗)
-    elseif hessian
-        return loss.∇²l(Ũ⃗_end, ϕ⃗)
-    end
-end
-
-# ============================================================================= #
-
-@testitem "Unitary fidelity" begin
-    X = [0 1; 1 0]
-    X = [0 1; 1 0]
-    Y = [0 -im; im 0]
-    @test unitary_fidelity(X, X) ≈ 1
-    @test unitary_fidelity(X, Y) ≈ 0
-
-    # Tr undefined on Type{Any}
-    X1 = Real[0 1; 1 0]
-    X2 = Float64[0 1; 1 0]
-    @test unitary_fidelity(X1, X2) ≈ 1
-end
-
-@testitem "Isovec Unitary Fidelity" begin
-    using LinearAlgebra
-
-    U_X = [0 1; 1 0]
-    U_Y = [0 -im; im 0]
-
-    for U in [U_X, U_Y]
-        @test U'U ≈ I
-    end
-
-    Ũ⃗_X = operator_to_iso_vec(U_X)
-    Ũ⃗_Y = operator_to_iso_vec(U_Y)
-
-    # Test gate fidelity
-    @test iso_vec_unitary_fidelity(Ũ⃗_X, Ũ⃗_X) ≈ 1
-    @test iso_vec_unitary_fidelity(Ũ⃗_X, Ũ⃗_Y) ≈ 0
-
-
-    # Test asymmetric fidelity
-    U_fn(λ, φ) = [1 -exp(im * λ); exp(im * φ) exp(im * (φ + λ))] / sqrt(2)
-    U_1 = U_fn(π/4, π/3)
-    U_2 = U_fn(1.5, .33)
-
-    for U in [U_1, U_2]
-        @test U'U ≈ I
-    end
-
-    Ũ⃗_1 = operator_to_iso_vec(U_1)
-    Ũ⃗_2 = operator_to_iso_vec(U_2)
-
-    @test iso_vec_unitary_fidelity(Ũ⃗_1, Ũ⃗_1) ≈ 1
-    @test iso_vec_unitary_fidelity(Ũ⃗_2, Ũ⃗_2) ≈ 1
-    @test iso_vec_unitary_fidelity(Ũ⃗_1, Ũ⃗_2) ≈ iso_vec_unitary_fidelity(Ũ⃗_2, Ũ⃗_1)
-    @test iso_vec_unitary_fidelity(Ũ⃗_1, Ũ⃗_2) ≈ abs(tr(U_1'U_2)) / 2
-
-
-    # Test random fidelity
-    U_H1 = haar_random(2)
-    U_H2 = haar_random(2)
-
-    for U in [U_H1, U_H2]
-        @test U'U ≈ I
-    end
-
-    Ũ⃗_H1 = operator_to_iso_vec(U_H1)
-    Ũ⃗_H2 = operator_to_iso_vec(U_H2)
-
-    @test iso_vec_unitary_fidelity(Ũ⃗_H1, Ũ⃗_H1) ≈ 1
-    @test iso_vec_unitary_fidelity(Ũ⃗_H2, Ũ⃗_H2) ≈ 1
-    @test iso_vec_unitary_fidelity(Ũ⃗_H1, Ũ⃗_X) ≈ abs(tr(U_H1'U_X)) / 2
-    @test iso_vec_unitary_fidelity(Ũ⃗_H1, Ũ⃗_H2) ≈ abs(tr(U_H1'U_H2)) / 2
-end
-
-
-@testitem "Isovec Unitary Fidelity Subspace" begin
-    using LinearAlgebra
-
-    function test_iso_vec_unitary_fidelity(
-        U₁::AbstractMatrix,
-        U₂::AbstractMatrix,
-        subspace
-    )
-        Ũ⃗₁ = operator_to_iso_vec(U₁)
-        Ũ⃗₂ = operator_to_iso_vec(U₂)
-        return iso_vec_unitary_fidelity(Ũ⃗₁, Ũ⃗₂, subspace=subspace)
-    end
-
-    # Test random fidelity
-    test_subspaces = [
-        get_subspace_indices([1:2, 1:1], [2, 2]),
-        get_subspace_indices([1:2, 2:2], [2, 2]),
-    ]
-
-    for ii in test_subspaces
-        U_H1 = kron(haar_random(2), haar_random(2))
-        U_H1_sub = U_H1[ii, ii]
-        U_H2 = kron(haar_random(2), haar_random(2))
-        U_H2_sub = U_H2[ii, ii]
-
-        # NOTE: subspace may not be unitary (?)
-        for U in [U_H1, U_H2]
-            @test U'U ≈ I
-        end
-
-        fid = test_iso_vec_unitary_fidelity(U_H1, U_H2, ii)
-        fid_sub = test_iso_vec_unitary_fidelity(U_H1_sub, U_H2_sub, axes(U_H1_sub, 1))
-        @test fid ≈ fid_sub
-    end
-end
-
-
-@testitem "Isovec Unitary Fidelity Gradient" begin
-    @test_skip true
-end
-
-@testitem "Free phase fidelity" begin
-    n_levels = 3
-    phase_data = [1.9, 2.7]
-    phase_operators = [PAULIS[:Z], PAULIS[:Z]]
-    subspace = get_subspace_indices([1:2, 1:2], [n_levels, n_levels])
-
-    R = Losses.free_phase(phase_data, phase_operators)
-    @test R'R ≈ [1 0 0 0; 0 1 0 0; 0 0 1 0; 0 0 0 1]
-    @test size(R) == (2^2, 2^2)
-
-    U_goal = EmbeddedOperator(GATES[:CZ], subspace, [n_levels, n_levels]).operator
-    U_final = EmbeddedOperator(R'GATES[:CZ], subspace, [n_levels, n_levels]).operator
-    # Value is ~0.3 without phases
-    @test unitary_fidelity(U_final, U_goal, subspace=subspace) < 0.5
-    @test unitary_free_phase_fidelity(
-        U_final, 
-        U_goal, 
-        phase_data,
-        phase_operators,
-        subspace=subspace
-    ) ≈ 1
-
-    # Forgot subspace
-    @test_throws DimensionMismatch iso_vec_unitary_free_phase_fidelity(
-        operator_to_iso_vec(U_final),
-        operator_to_iso_vec(U_goal),
-        phase_data,
-        phase_operators,
-    )
-end
diff --git a/src/losses/unitary_trace_loss.jl b/src/losses/unitary_trace_loss.jl
deleted file mode 100644
index f8ef8d54..00000000
--- a/src/losses/unitary_trace_loss.jl
+++ /dev/null
@@ -1,49 +0,0 @@
-export UnitaryTraceLoss
-
-
-function unitary_trace_loss(Ũ⃗::AbstractVector, Ũ⃗_goal::AbstractVector)
-    U = iso_vec_to_operator(Ũ⃗)
-    Ugoal = iso_vec_to_operator(Ũ⃗_goal)
-    return 1 / 2 * tr(sqrt((U - Ugoal)' * (U - Ugoal)))
-end
-
-struct UnitaryTraceLoss <: AbstractLoss
-    l::Function
-    ∇l::Function
-    ∇²l::Function
-    ∇²l_structure::Vector{Tuple{Int,Int}}
-    name::Symbol
-
-    function UnitaryTraceLoss(
-        name::Symbol,
-        Ũ⃗_goal::AbstractVector
-    )
-        l = Ũ⃗ -> unitary_trace_loss(Ũ⃗, Ũ⃗_goal)
-        ∇l = Ũ⃗ -> ForwardDiff.gradient(l, Ũ⃗)
-        ∇²l = Ũ⃗ -> ForwardDiff.hessian(l, Ũ⃗)
-        Ũ⃗_dim = length(Ũ⃗_goal)
-        ∇²l_structure = []
-        for (i, j) ∈ Iterators.product(1:Ũ⃗_dim, 1:Ũ⃗_dim)
-            if i ≤ j
-                push!(∇²l_structure, (i, j))
-            end
-        end
-        return new(l, ∇l, ∇²l, ∇²l_structure, name)
-    end
-end
-
-function (loss::UnitaryTraceLoss)(
-    Ũ⃗_end::AbstractVector{<:Real};
-    gradient=false,
-    hessian=false
-)
-    @assert !(gradient && hessian)
-
-    if !(gradient || hessian)
-        return loss.l(Ũ⃗_end)
-    elseif gradient
-        return loss.∇l(Ũ⃗_end)
-    elseif hessian
-        return loss.∇²l(Ũ⃗_end)
-    end
-end
\ No newline at end of file
diff --git a/src/objectives/_objectives.jl b/src/objectives/_objectives.jl
deleted file mode 100644
index 8985fbae..00000000
--- a/src/objectives/_objectives.jl
+++ /dev/null
@@ -1,128 +0,0 @@
-module Objectives
-
-export Objective
-export NullObjective
-
-using ..Isomorphisms
-using ..QuantumSystems
-using ..EmbeddedOperators
-using ..Losses
-using ..Constraints
-
-using TrajectoryIndexingUtils
-using NamedTrajectories
-
-using LinearAlgebra
-using SparseArrays
-using Symbolics
-using TestItemRunner
-
-include("quantum_objective.jl")
-include("unitary_infidelity_objective.jl")
-include("regularizer_objective.jl")
-include("minimum_time_objective.jl")
-include("unitary_robustness_objective.jl")
-
-# TODO:
-# - [ ] Do not reference the Z object in the objective (components only / remove "name")
-
-"""
-    sparse_to_moi(A::SparseMatrixCSC)
-
-Converts a sparse matrix to tuple of vector of nonzero indices and vector of nonzero values
-"""
-function sparse_to_moi(A::SparseMatrixCSC)
-    inds = collect(zip(findnz(A)...))
-    vals = [A[i,j] for (i,j) ∈ inds]
-    return (inds, vals)
-end
-
-# ----------------------------------------------------------------------------- #
-#                           Objective                                           #
-# ----------------------------------------------------------------------------- #
-
-"""
-    Objective
-
-A structure for defining objective functions.
-
-The `terms` field contains all the arguments needed to construct the objective function.
-
-Fields:
-    `L`: the objective function
-    `∇L`: the gradient of the objective function
-    `∂²L`: the Hessian of the objective function
-    `∂²L_structure`: the structure of the Hessian of the objective function
-    `terms`: a vector of dictionaries containing the terms of the objective function
-"""
-struct Objective
-	L::Function
-	∇L::Function
-	∂²L::Union{Function, Nothing}
-	∂²L_structure::Union{Function, Nothing}
-    terms::Vector{Dict}
-end
-
-function Base.:+(obj1::Objective, obj2::Objective)
-	L = (Z⃗, Z) -> obj1.L(Z⃗, Z) + obj2.L(Z⃗, Z)
-	∇L = (Z⃗, Z) -> obj1.∇L(Z⃗, Z) + obj2.∇L(Z⃗, Z)
-	if isnothing(obj1.∂²L) && isnothing(obj2.∂²L)
-		∂²L = Nothing
-		∂²L_structure = Nothing
-	elseif isnothing(obj1.∂²L)
-		∂²L = (Z⃗, Z) -> obj2.∂²L(Z⃗, Z)
-		∂²L_structure = obj2.∂²L_structure
-	elseif isnothing(obj2.∂²L)
-		∂²L = (Z⃗, Z) -> obj1.∂²L(Z⃗, Z)
-		∂²L_structure = obj1.∂²L_structure
-	else
-		∂²L = (Z⃗, Z) -> vcat(obj1.∂²L(Z⃗, Z), obj2.∂²L(Z⃗, Z))
-		∂²L_structure = Z -> vcat(obj1.∂²L_structure(Z), obj2.∂²L_structure(Z))
-	end
-    terms = vcat(obj1.terms, obj2.terms)
-	return Objective(L, ∇L, ∂²L, ∂²L_structure, terms)
-end
-
-Base.:+(obj::Objective, ::Nothing) = obj
-Base.:+(obj::Objective) = obj
-
-function Objective(terms::AbstractVector{<:Dict})
-    return +(Objective.(terms)...)
-end
-
-function Base.:*(num::Real, obj::Objective)
-	L = (Z⃗, Z) -> num * obj.L(Z⃗, Z)
-	∇L = (Z⃗, Z) -> num * obj.∇L(Z⃗, Z)
-    if isnothing(obj.∂²L)
-        ∂²L = nothing
-        ∂²L_structure = nothing
-    else
-        ∂²L = (Z⃗, Z) -> num * obj.∂²L(Z⃗, Z)
-        ∂²L_structure = obj.∂²L_structure
-    end
-	return Objective(L, ∇L, ∂²L, ∂²L_structure, obj.terms)
-end
-
-Base.:*(obj::Objective, num::Real) = num * obj
-
-function Objective(term::Dict)
-    return eval(term[:type])(; delete!(term, :type)...)
-end
-
-# ----------------------------------------------------------------------------- #
-#                           Null objective                                      #
-# ----------------------------------------------------------------------------- #
-
-function NullObjective()
-    params = Dict(:type => :NullObjective)
-	L(Z⃗::AbstractVector{R}, Z::NamedTrajectory) where R<:Real = 0.0
-    ∇L(Z⃗::AbstractVector{R}, Z::NamedTrajectory) where R<:Real = zeros(R, Z.dim * Z.T + Z.global_dim)
-    ∂²L_structure(Z::NamedTrajectory) = []
-    function ∂²L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory; return_moi_vals=true)
-        n = Z.dim * Z.T + Z.global_dim
-        return return_moi_vals ? [] : spzeros(n, n)
-    end
-	return Objective(L, ∇L, ∂²L, ∂²L_structure, Dict[params])
-end
-
-end
diff --git a/src/objectives/minimum_time_objective.jl b/src/objectives/minimum_time_objective.jl
deleted file mode 100644
index 618e82fb..00000000
--- a/src/objectives/minimum_time_objective.jl
+++ /dev/null
@@ -1,57 +0,0 @@
-export MinimumTimeObjective
-
-
-"""
-    MinimumTimeObjective
-
-A type of objective that counts the time taken to complete a task.
-
-Fields:
-    `D`: a scaling factor
-    `Δt_indices`: the indices of the time steps
-    `eval_hessian`: whether to evaluate the Hessian
-
-"""
-function MinimumTimeObjective(;
-    D::Float64=1.0,
-    Δt_indices::AbstractVector{Int}=nothing,
-    eval_hessian::Bool=true
-)
-    @assert !isnothing(Δt_indices) "Δt_indices must be specified"
-
-    params = Dict(
-        :type => :MinimumTimeObjective,
-        :D => D,
-        :Δt_indices => Δt_indices,
-        :eval_hessian => eval_hessian
-    )
-
-    # TODO: amend this for case of no TimeStepsAllEqualConstraint
-	L(Z⃗::AbstractVector, Z::NamedTrajectory) = D * sum(Z⃗[Δt_indices])
-
-	∇L = (Z⃗::AbstractVector, Z::NamedTrajectory) -> begin
-		∇ = zeros(typeof(Z⃗[1]), length(Z⃗))
-		∇[Δt_indices] .= D
-		return ∇
-	end
-
-	if eval_hessian
-		∂²L = (Z⃗, Z) -> []
-		∂²L_structure = Z -> []
-	else
-		∂²L = nothing
-		∂²L_structure = nothing
-	end
-
-	return Objective(L, ∇L, ∂²L, ∂²L_structure, Dict[params])
-end
-
-function MinimumTimeObjective(
-    traj::NamedTrajectory;
-    D=1.0,
-    kwargs...
-)
-    @assert traj.timestep isa Symbol "trajectory does not have a dynamical timestep"
-    Δt_indices = [index(t, traj.components[traj.timestep][1], traj.dim) for t = 1:traj.T]
-    return MinimumTimeObjective(; D=D, Δt_indices=Δt_indices, kwargs...)
-end
\ No newline at end of file
diff --git a/src/objectives/quantum_objective.jl b/src/objectives/quantum_objective.jl
deleted file mode 100644
index 2d129005..00000000
--- a/src/objectives/quantum_objective.jl
+++ /dev/null
@@ -1,217 +0,0 @@
-export QuantumObjective
-export QuantumUnitaryObjective
-export QuantumStateObjective
-
-###
-### QuantumObjective
-###
-
-"""
-    QuantumObjective
-
-    A generic objective function for quantum trajectories that use a loss.
-
-"""
-function QuantumObjective(;
-    names::Union{Nothing,Tuple{Vararg{Symbol}}}=nothing,
-    name::Union{Nothing,Symbol}=nothing,
-    goals::Union{Nothing,AbstractVector{<:Real},Tuple{Vararg{AbstractVector{<:Real}}}}=nothing,
-	loss::Symbol=:InfidelityLoss,
-	Q::Union{Float64, Vector{Float64}}=100.0,
-	eval_hessian::Bool=true
-)
-    @assert !(isnothing(names) && isnothing(name)) "name or names must be specified"
-    @assert !isnothing(goals) "goals corresponding to names must be specified"
-
-    if isnothing(names)
-        names = (name,)
-    end
-
-    if goals isa AbstractVector
-        goals = (goals,)
-    end
-
-    if Q isa Float64
-        Q = ones(length(names)) * Q
-    else
-        @assert length(Q) == length(names)
-    end
-
-    params = Dict(
-        :type => :QuantumObjective,
-        :names => names,
-        :goals => goals,
-        :loss => loss,
-        :Q => Q,
-        :eval_hessian => eval_hessian,
-    )
-
-    losses = [eval(loss)(name, goal) for (name, goal) ∈ zip(names, goals)]
-
-	@views function L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        loss = 0.0
-        for (Qᵢ, lᵢ, name) ∈ zip(Q, losses, names)
-            name_slice = slice(Z.T, Z.components[name], Z.dim)
-            loss += Qᵢ * lᵢ(Z⃗[name_slice])
-        end
-        return loss
-    end
-
-    @views function ∇L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        ∇ = zeros(Z.dim * Z.T + Z.global_dim)
-        for (Qᵢ, lᵢ, name) ∈ zip(Q, losses, names)
-            name_slice = slice(Z.T, Z.components[name], Z.dim)
-            ∇[name_slice] = Qᵢ * lᵢ(Z⃗[name_slice]; gradient=true)
-        end
-        return ∇
-    end
-
-    function ∂²L_structure(Z::NamedTrajectory)
-        structure = []
-        final_time_offset = index(Z.T, 0, Z.dim)
-        for (name, loss) ∈ zip(names, losses)
-            comp_start_offset = Z.components[name][1] - 1
-            comp_hessian_structure = [
-                ij .+ (final_time_offset + comp_start_offset)
-                    for ij ∈ loss.∇²l_structure
-            ]
-            append!(structure, comp_hessian_structure)
-        end
-        return structure
-    end
-
-
-    @views function ∂²L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory; return_moi_vals=true)
-        H = spzeros(Z.dim * Z.T + Z.global_dim, Z.dim * Z.T + Z.global_dim)
-        for (Qᵢ, name, lᵢ) ∈ zip(Q, names, losses)
-            name_slice = slice(Z.T, Z.components[name], Z.dim)
-            H[name_slice, name_slice] =
-                Qᵢ * lᵢ(Z⃗[name_slice]; hessian=true)
-        end
-        if return_moi_vals
-            Hs = [H[i,j] for (i, j) ∈ ∂²L_structure(Z)]
-            return Hs
-        else
-            return H
-        end
-    end
-
-	return Objective(L, ∇L, ∂²L, ∂²L_structure, Dict[params])
-end
-
-function QuantumObjective(
-    name::Symbol,
-    traj::NamedTrajectory,
-    loss::Symbol,
-    Q::Float64
-)
-    goal = traj.goal[name]
-    return QuantumObjective(name=name, goals=goal, loss=loss, Q=Q)
-end
-
-function QuantumObjective(
-    names::Tuple{Vararg{Symbol}},
-    traj::NamedTrajectory,
-    loss::Symbol,
-    Q::Float64
-)
-    goals = Tuple(traj.goal[name] for name in names)
-    return QuantumObjective(names=names, goals=goals, loss=loss, Q=Q)
-end
-
-###
-### Example: Default quantum objectives
-###
-
-function QuantumUnitaryObjective(
-    name::Symbol,
-    traj::NamedTrajectory,
-    Q::Float64
-)
-    return QuantumObjective(name, traj, :UnitaryInfidelityLoss, Q)
-end
-
-function QuantumStateObjective(
-    name::Symbol,
-    traj::NamedTrajectory,
-    Q::Float64
-)
-    return QuantumObjective(name, traj, :InfidelityLoss, Q)
-end
-
-# ============================================================================= # 
-
-@testitem "Quantum State Objective" begin
-    using LinearAlgebra
-    using NamedTrajectories
-    using ForwardDiff
-    include("../../test/test_utils.jl")
-    
-    T = 10
-
-    Z = NamedTrajectory(
-        (ψ̃ = randn(4, T), u = randn(2, T)),
-        controls=:u,
-        timestep=0.1,
-        goal=(ψ̃ = [1., 0., 0., 0.],)
-    )
-
-    loss = :InfidelityLoss
-    Q = 100.0
-
-    J = QuantumObjective(:ψ̃, Z, loss, Q)
-
-    L = Z⃗ -> J.L(Z⃗, Z)
-    ∇L = Z⃗ -> J.∇L(Z⃗, Z)
-    ∂²L = Z⃗ -> J.∂²L(Z⃗, Z)
-    ∂²L_structure = J.∂²L_structure(Z)
-
-    # test objective function gradient
-    @test ForwardDiff.gradient(L, Z.datavec) ≈ ∇L(Z.datavec)
-
-    # test objective function hessian
-    shape = (Z.dim * Z.T + Z.global_dim, Z.dim * Z.T + Z.global_dim)
-    @test ForwardDiff.hessian(L, Z.datavec) ≈ dense(∂²L(Z.datavec), ∂²L_structure, shape)
-end
-
-@testitem "Unitary Objective" begin
-    using LinearAlgebra
-    using NamedTrajectories
-    using ForwardDiff
-    include("../../test/test_utils.jl")
-
-    T = 10
-
-    U_init = GATES[:I]
-    U_goal = GATES[:X]
-
-    Ũ⃗_init = operator_to_iso_vec(U_init)
-    Ũ⃗_goal = operator_to_iso_vec(U_goal)
-
-    Z = NamedTrajectory(
-        (Ũ⃗ = randn(length(Ũ⃗_init), T), u = randn(2, T)),
-        controls=:u,
-        timestep=0.1,
-        initial=(Ũ⃗ = Ũ⃗_init,),
-        goal=(Ũ⃗ = Ũ⃗_goal,)
-    )
-
-    loss = :UnitaryInfidelityLoss
-    Q = 100.0
-
-    J = QuantumObjective(:Ũ⃗, Z, loss, Q)
-
-    L = Z⃗ -> J.L(Z⃗, Z)
-    ∇L = Z⃗ -> J.∇L(Z⃗, Z)
-    ∂²L = Z⃗ -> J.∂²L(Z⃗, Z)
-    ∂²L_structure = J.∂²L_structure(Z)
-
-    # test objective function gradient
-    @test all(ForwardDiff.gradient(L, Z.datavec) ≈ ∇L(Z.datavec))
-
-    # test objective function hessian
-    shape = (Z.dim * Z.T + Z.global_dim, Z.dim * Z.T + Z.global_dim)
-    H = dense(∂²L(Z.datavec), ∂²L_structure, shape)
-    H_forwarddiff = ForwardDiff.hessian(L, Z.datavec)
-    @test all(H .≈ H_forwarddiff)
-end
\ No newline at end of file
diff --git a/src/objectives/regularizer_objective.jl b/src/objectives/regularizer_objective.jl
deleted file mode 100644
index 7cabf28d..00000000
--- a/src/objectives/regularizer_objective.jl
+++ /dev/null
@@ -1,633 +0,0 @@
-export QuadraticRegularizer
-export QuadraticSmoothnessRegularizer
-export L1Regularizer
-export L1Regularizer!
-export PairwiseQuadraticRegularizer
-
-
-###
-### Quadratic Regularizer
-###
-
-"""
-    QuadraticRegularizer
-
-A quadratic regularizer for a trajectory component.
-
-Fields:
-    `name`: the name of the trajectory component to regularize
-    `times`: the times at which to evaluate the regularizer
-    `dim`: the dimension of the trajectory component
-    `R`: the regularization matrix
-    `baseline`: the baseline values for the trajectory component
-    `eval_hessian`: whether to evaluate the Hessian of the regularizer
-    `timestep_name`: the symbol for the timestep variable
-"""
-function QuadraticRegularizer(;
-	name::Union{Nothing, Symbol}=nothing,
-	times::Union{Nothing, AbstractVector{Int}}=nothing,
-    dim::Union{Nothing, Int}=nothing,
-	R::Union{Nothing, AbstractVector{<:Real}}=nothing,
-    baseline::Union{Nothing, AbstractArray{<:Real}}=nothing,
-	eval_hessian::Bool=true,
-    timestep_name::Symbol=:Δt
-)
-
-    @assert !isnothing(name) "name must be specified"
-    @assert !isnothing(times) "times must be specified"
-    @assert !isnothing(dim) "dim must be specified"
-    @assert !isnothing(R) "R must be specified"
-    if isnothing(baseline)
-        baseline = zeros(length(R), length(times))
-    else
-        if size(baseline) != (length(R), length(times))
-            throw(ArgumentError("size(baseline)=$(size(baseline)) must match $(length(R)) x $(length(times))"))
-        end
-    end
-
-    params = Dict(
-        :type => :QuadraticRegularizer,
-        :name => name,
-        :times => times,
-        :dim => dim,
-        :baseline => baseline,
-        :R => R,
-        :eval_hessian => eval_hessian
-    )
-
-    @views function L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        J = 0.0
-        for t ∈ times
-            if Z.timestep isa Symbol
-                Δt = Z⃗[slice(t, Z.components[timestep_name], Z.dim)]
-            else
-                Δt = Z.timestep
-            end
-
-            vₜ = Z⃗[slice(t, Z.components[name], Z.dim)]
-            Δv = vₜ .- baseline[:, t]
-
-            rₜ = Δt .* Δv
-            J += 0.5 * rₜ' * (R .* rₜ)
-        end
-        return J
-    end
-
-    @views function ∇L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        ∇ = zeros(Z.dim * Z.T + Z.global_dim)
-        Threads.@threads for t ∈ times
-            vₜ_slice = slice(t, Z.components[name], Z.dim)
-            Δv = Z⃗[vₜ_slice] .- baseline[:, t]
-
-            if Z.timestep isa Symbol
-                Δt_slice = slice(t, Z.components[timestep_name], Z.dim)
-                Δt = Z⃗[Δt_slice]
-                ∇[Δt_slice] .= Δv' * (R .* (Δt .* Δv))
-            else
-                Δt = Z.timestep
-            end
-
-            ∇[vₜ_slice] .= R .* (Δt.^2 .* Δv)
-        end
-        return ∇
-    end
-
-    ∂²L = nothing
-    ∂²L_structure = nothing
-
-    if eval_hessian
-
-        ∂²L_structure = Z -> begin
-            structure = []
-            # Hessian structure (eq. 17)
-            for t ∈ times
-                vₜ_slice = slice(t, Z.components[name], Z.dim)
-                vₜ_vₜ_inds = collect(zip(vₜ_slice, vₜ_slice))
-                append!(structure, vₜ_vₜ_inds)
-
-                if Z.timestep isa Symbol
-                    Δt_slice = slice(t, Z.components[timestep_name], Z.dim)
-                    # ∂²_vₜ_Δt
-                    vₜ_Δt_inds = [(i, j) for i ∈ vₜ_slice for j ∈ Δt_slice]
-                    append!(structure, vₜ_Δt_inds)
-                    # ∂²_Δt_vₜ
-                    Δt_vₜ_inds = [(i, j) for i ∈ Δt_slice for j ∈ vₜ_slice]
-                    append!(structure, Δt_vₜ_inds)
-                    # ∂²_Δt_Δt
-                    Δt_Δt_inds = collect(zip(Δt_slice, Δt_slice))
-                    append!(structure, Δt_Δt_inds)
-                end
-            end
-            return structure
-        end
-
-        ∂²L = (Z⃗, Z) -> begin
-            values = []
-            # Match Hessian structure indices
-            for t ∈ times
-                if Z.timestep isa Symbol
-                    Δt = Z⃗[slice(t, Z.components[timestep_name], Z.dim)]
-                    append!(values, R .* Δt.^2)
-                    # ∂²_vₜ_Δt, ∂²_Δt_vₜ
-                    vₜ = Z⃗[slice(t, Z.components[name], Z.dim)]
-                    Δv = vₜ .- baseline[:, t]
-                    append!(values, 2 * (R .* (Δt .* Δv)))
-                    append!(values, 2 * (R .* (Δt .* Δv)))
-                    # ∂²_Δt_Δt
-                    append!(values, Δv' * (R .* Δv))
-                else
-                    Δt = Z.timestep
-                    append!(values, R .* Δt.^2)
-                end
-            end
-            return values
-        end
-    end
-
-    return Objective(L, ∇L, ∂²L, ∂²L_structure, Dict[params])
-end
-
-function QuadraticRegularizer(
-    name::Symbol,
-    traj::NamedTrajectory,
-    R::AbstractVector{<:Real};
-    kwargs...
-)
-    return QuadraticRegularizer(;
-        name=name,
-        times=1:traj.T,
-        dim=traj.dim,
-        R=R,
-        kwargs...
-    )
-end
-
-function QuadraticRegularizer(
-    name::Symbol,
-    traj::NamedTrajectory,
-    R::Real;
-    kwargs...
-)
-    return QuadraticRegularizer(;
-        name=name,
-        times=1:traj.T,
-        dim=traj.dim,
-        R=R * ones(traj.dims[name]),
-        kwargs...
-    )
-end
-
-###
-### QuadraticSmoothnessRegularizer
-###
-
-"""
-    QuadraticSmoothnessRegularizer
-
-A quadratic smoothness regularizer for a trajectory component.
-
-Fields:
-    `name`: the name of the trajectory component to regularize
-    `times`: the times at which to evaluate the regularizer
-    `R`: the regularization matrix
-    `eval_hessian`: whether to evaluate the Hessian of the regularizer
-"""
-function QuadraticSmoothnessRegularizer(;
-	name::Symbol=nothing,
-    times::Union{Nothing, AbstractVector{Int}}=nothing,
-	R::Union{Nothing, AbstractVector{<:Real}}=nothing,
-	eval_hessian=true
-)
-    @assert !isnothing(name) "name must be specified"
-    @assert !isnothing(times) "times must be specified"
-
-    params = Dict(
-        :type => :QuadraticSmoothnessRegularizer,
-        :name => name,
-        :times => times,
-        :R => R,
-        :eval_hessian => eval_hessian
-    )
-
-	@views function L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-		∑Δv² = 0.0
-		for t ∈ times[1:end-1]
-			vₜ₊₁ = Z⃗[slice(t + 1, Z.components[name], Z.dim)]
-			vₜ = Z⃗[slice(t, Z.components[name], Z.dim)]
-			Δv = vₜ₊₁ - vₜ
-			∑Δv² += 0.5 * Δv' * (R .* Δv)
-		end
-		return ∑Δv²
-	end
-
-	@views function ∇L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        ∇ = zeros(Z.dim * Z.T + Z.global_dim)
-		Threads.@threads for t ∈ times[1:end-1]
-
-			vₜ_slice = slice(t, Z.components[name], Z.dim)
-			vₜ₊₁_slice = slice(t + 1, Z.components[name], Z.dim)
-
-			vₜ = Z⃗[vₜ_slice]
-			vₜ₊₁ = Z⃗[vₜ₊₁_slice]
-
-			Δv = vₜ₊₁ - vₜ
-
-			∇[vₜ_slice] += -R .* Δv
-			∇[vₜ₊₁_slice] += R .* Δv
-		end
-		return ∇
-	end
-    ∂²L = nothing
-	∂²L_structure = nothing
-
-	if eval_hessian
-
-		∂²L_structure = Z -> begin
-            structure = []
-
-		    # u smoothness regularizer Hessian main diagonal structure
-            for t ∈ times
-                vₜ_slice = slice(t, Z.components[name], Z.dim)
-
-                # main diagonal (2 if t != 1 or T-1) * Rₛ I
-                # components: ∂²vₜSₜ
-                append!(
-                    structure,
-                    collect(zip(vₜ_slice, vₜ_slice))
-                )
-            end
-
-            # u smoothness regularizer Hessian off diagonal structure
-            for t ∈ times[1:end-1]
-                vₜ_slice = slice(t, Z.components[name], Z.dim)
-                vₜ₊₁_slice = slice(t + 1, Z.components[name], Z.dim)
-
-                # off diagonal -Rₛ I components: ∂vₜ₊₁∂vₜSₜ
-                append!(
-                    structure,
-                    collect(zip(vₜ_slice, vₜ₊₁_slice))
-                )
-            end
-            return structure
-        end
-
-		∂²L = (Z⃗, Z) -> begin
-
-			H = []
-
-			# u smoothness regularizer Hessian main diagonal values
-			append!(H, R)
-			for t in times[2:end-1]
-				append!(H, 2 * R)
-			end
-			append!(H, R)
-
-			# u smoothness regularizer Hessian off diagonal values
-			for t in times[1:end-1]
-				append!(H, -R)
-			end
-			return H
-		end
-	end
-
-	return Objective(L, ∇L, ∂²L, ∂²L_structure, Dict[params])
-end
-
-function QuadraticSmoothnessRegularizer(
-    name::Symbol,
-    traj::NamedTrajectory,
-    R::AbstractVector{<:Real};
-    kwargs...
-)
-    return QuadraticSmoothnessRegularizer(;
-        name=name,
-        times=1:traj.T,
-        R=R,
-        kwargs...
-    )
-end
-
-function QuadraticSmoothnessRegularizer(
-    name::Symbol,
-    traj::NamedTrajectory,
-    R::Real;
-    kwargs...
-)
-    return QuadraticSmoothnessRegularizer(;
-        name=name,
-        times=1:traj.T,
-        R=R * ones(traj.dims[name]),
-        kwargs...
-    )
-end
-
-###
-### L1Regularizer
-###
-
-@doc raw"""
-    L1Regularizer
-
-Create an L1 regularizer for the trajectory component. The regularizer is defined as
-
-```math
-J_{L1}(u) = \sum_t \abs{R \cdot u_t}
-```
-
-where \(R\) is the regularization matrix and \(u_t\) is the trajectory component at time \(t\).
-
-
-"""
-function L1Regularizer(;
-    name=nothing,
-    R::Vector{Float64}=nothing,
-    times=nothing,
-    eval_hessian=true
-)
-    @assert !isnothing(name) "name must be specified"
-    @assert !isnothing(R) "R must be specified"
-    @assert !isnothing(times) "times must be specified"
-
-    s1_name = Symbol("s1_$name")
-    s2_name = Symbol("s2_$name")
-
-    params = Dict(
-        :type => :L1Regularizer,
-        :name => name,
-        :R => R,
-        :eval_hessian => eval_hessian,
-        :times => times,
-    )
-
-    L = (Z⃗, Z) -> sum(
-        dot(
-            R,
-            Z⃗[slice(t, Z.components[s1_name], Z.dim)] +
-            Z⃗[slice(t, Z.components[s2_name], Z.dim)]
-        ) for t ∈ times
-    )
-
-    ∇L = (Z⃗, Z) -> begin
-        ∇ = zeros(typeof(Z⃗[1]), length(Z⃗))
-        Threads.@threads for t ∈ times
-            ∇[slice(t, Z.components[s1_name], Z.dim)] += R
-            ∇[slice(t, Z.components[s2_name], Z.dim)] += R
-        end
-        return ∇
-    end
-
-    if eval_hessian
-        ∂²L = (_, _)  -> []
-        ∂²L_structure = _ -> []
-    else
-        ∂²L = nothing
-        ∂²L_structure = nothing
-    end
-
-    return Objective(L, ∇L, ∂²L, ∂²L_structure, Dict[params])
-end
-
-function L1Regularizer(
-    name::Symbol,
-    traj::NamedTrajectory;
-    indices::AbstractVector{Int}=1:traj.dims[name],
-    times=(name ∈ keys(traj.initial) ? 2 : 1):traj.T,
-    R_value::Float64=10.0,
-    R::Vector{Float64}=fill(R_value, length(indices)),
-    eval_hessian=true
-)
-    J = L1Regularizer(;
-        name=name,
-        R=R,
-        times=times,
-        eval_hessian=eval_hessian
-    )
-
-    slack_con = L1SlackConstraint(name, traj; indices=indices, times=times)
-
-    return J, slack_con
-end
-
-function L1Regularizer!(
-    constraints::Vector{<:AbstractConstraint},
-    name::Symbol,
-    traj::NamedTrajectory;
-    kwargs...
-)
-    J, slack_con = L1Regularizer(name, traj; kwargs...)
-    push!(constraints, slack_con)
-    return J
-end
-
-###
-### PairwiseQuadraticRegularizer
-###
-
-@doc raw"""
-    PairwiseQuadraticRegularizer
-
-Create a pairwise quadratic regularizer for the trajectory component `name` with
-regularization strength `R`. The regularizer is defined as
-
-```math
-    J_{v⃗}(u) = \sum_t \frac{1}{2} \Delta t_t^2 (v⃗_{1,t} - v⃗_{2,t})^T R (v⃗_{1,t} - v⃗_{2,t})
-```
-
-where $v⃗_{1}$ and $v⃗_{2}$ are selected by `name1` and `name2`. The indices specify the
-appropriate block diagonal components of the direct sum vector `v⃗`.
-
-TODO: Hessian not implemented
-
-
-Fields:
-    `R`: the regularization strength
-    `times`: the time steps to apply the regularizer
-    `name1`: the first name
-    `name2`: the second name
-    `timestep_name`: the symbol for the timestep
-    `eval_hessian`: whether to evaluate the Hessian
-"""
-function PairwiseQuadraticRegularizer(
-    R::AbstractVector{<:Real},
-    times::AbstractVector{Int},
-    name1::Symbol,
-    name2::Symbol;
-    timestep_name::Symbol=:Δt,
-    eval_hessian::Bool=false,
-)
-    params = Dict(
-        :type => :PairwiseQuadraticRegularizer,
-        :times => times,
-        :name => (name1, name2),
-        :R => R,
-        :eval_hessian => eval_hessian,
-    )
-
-    @views function L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        J = 0.0
-        for t ∈ times
-            if Z.timestep isa Symbol
-                Δt = Z⃗[slice(t, Z.components[timestep_name], Z.dim)]
-            else
-                Δt = Z.timestep
-            end
-            z1_t = Z⃗[slice(t, Z.components[name1], Z.dim)]
-            z2_t = Z⃗[slice(t, Z.components[name1], Z.dim)]
-            r_t = Δt * (z1_t .- z2_t)
-            J += 0.5 * r_t' * (R .* r_t)
-        end
-        return J
-    end
-
-    @views function ∇L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        ∇ = zeros(Z.dim * Z.T + Z.global_dim)
-        Threads.@threads for t ∈ times
-            z1_t_slice = slice(t, Z.components[name1], Z.dim)
-            z2_t_slice = slice(t, Z.components[name2], Z.dim)
-            z1_t = Z⃗[z1_t_slice]
-            z2_t = Z⃗[z2_t_slice]
-
-            if Z.timestep isa Symbol
-                Δt_slice = slice(t, Z.components[timestep_name], Z.dim)
-                Δt = Z⃗[Δt_slice]
-                ∇[Δt_slice] .= (z1_t .- z2_t)' * (R .* (Δt .* (z1_t .- z2_t)))
-            else
-                Δt = Z.timestep
-            end
-
-            ∇[z1_t_slice] .= R .* (Δt^2 * (z1_t .- z2_t))
-            ∇[z2_t_slice] .= R .* (Δt^2 * (z2_t .- z1_t))
-        end
-        return ∇
-    end
-
-    # TODO: Hessian not implemented
-    ∂²L = nothing
-    ∂²L_structure = nothing
-
-    if eval_hessian
-        throw(ErrorException("Hessian not implemented"))
-    end
-
-    return Objective(L, ∇L, ∂²L, ∂²L_structure, Dict[params])
-end
-
-@doc raw"""
-    PairwiseQuadraticRegularizer
-
-A convenience constructor for creating a PairwiseQuadraticRegularizer for the
-trajectory component `name` with regularization strength `Rs` over the graph `graph`.
-"""
-function PairwiseQuadraticRegularizer(
-    traj::NamedTrajectory,
-    Rs::Union{Float64, AbstractVector{<:Float64}},
-    graph::AbstractVector{<:Tuple{Symbol, Symbol}};
-    kwargs...
-)
-    if isa(Rs, Float64)
-        Rs = Rs * ones(length(graph))
-    end
-    @assert all(length(graph) == length(Rs)) "Graph and Qs must have same length"
-
-    J = NullObjective()
-    for (Qᵢⱼ, (symb1, symb2)) ∈ zip(Rs, graph)
-        # Symbols should be the same size
-        dim = size(traj[symb1], 1)
-        J += PairwiseQuadraticRegularizer(
-            Qᵢⱼ * ones(dim),
-            1:traj.T,
-            symb1,
-            symb2,
-            kwargs...
-        )
-    end
-
-    return J
-end
-
-function PairwiseQuadraticRegularizer(
-    traj::NamedTrajectory,
-    R::Float64,
-    name1::Symbol,
-    name2::Symbol;
-    kwargs...
-)
-    return PairwiseQuadraticRegularizer(
-        traj, R, [(name1, name2)];
-        kwargs...
-    )
-end
-
-# ============================================================================ #
-
-@testitem "Quadratic Regularizer Objective" begin
-    using LinearAlgebra
-    using NamedTrajectories
-    using ForwardDiff
-    include("../../test/test_utils.jl")
-
-    T = 10
-
-    Z = NamedTrajectory(
-        (ψ̃ = randn(4, T), u = randn(2, T)),
-        controls=:u,
-        timestep=0.1,
-        goal=(ψ̃ = [1.0, 0.0, 0.0, 0.0],)
-    )
-
-
-    J = QuadraticRegularizer(:u, Z, [1., 1.])
-
-    L = Z⃗ -> J.L(Z⃗, Z)
-    ∇L = Z⃗ -> J.∇L(Z⃗, Z)
-    ∂²L = Z⃗ -> J.∂²L(Z⃗, Z)
-    ∂²L_structure = J.∂²L_structure(Z)
-
-    # test objective function gradient
-
-    @test all(ForwardDiff.gradient(L, Z.datavec) .≈ ∇L(Z.datavec))
-
-    # test objective function hessian
-    shape = (Z.dim * Z.T + Z.global_dim, Z.dim * Z.T + Z.global_dim)
-    @test all(isapprox(
-        ForwardDiff.hessian(L, Z.datavec),
-        dense(∂²L(Z.datavec), ∂²L_structure, shape);
-        atol=1e-7
-    ))
-end
-
-@testitem "Quadratic Smoothness Regularizer Objective" begin
-    using LinearAlgebra
-    using NamedTrajectories
-    using ForwardDiff
-    include("../../test/test_utils.jl")
-
-    T = 10
-
-    Z = NamedTrajectory(
-        (ψ̃ = randn(4, T), u = randn(2, T)),
-        controls=:u,
-        timestep=0.1,
-        goal=(ψ̃ = [1.0, 0.0, 0.0, 0.0],)
-    )
-
-
-    J = QuadraticSmoothnessRegularizer(:u, Z, [1., 1.])
-
-    L = Z⃗ -> J.L(Z⃗, Z)
-    ∇L = Z⃗ -> J.∇L(Z⃗, Z)
-    ∂²L = Z⃗ -> J.∂²L(Z⃗, Z)
-    ∂²L_structure = J.∂²L_structure(Z)
-
-    # test objective function gradient
-
-    @test all(ForwardDiff.gradient(L, Z.datavec) .≈ ∇L(Z.datavec))
-
-    # test objective function hessian
-    shape = (Z.dim * Z.T + Z.global_dim, Z.dim * Z.T + Z.global_dim)
-    @test all(isapprox(
-        ForwardDiff.hessian(L, Z.datavec),
-        dense(∂²L(Z.datavec), ∂²L_structure, shape);
-        atol=1e-7
-    ))
-end
diff --git a/src/objectives/unitary_infidelity_objective.jl b/src/objectives/unitary_infidelity_objective.jl
deleted file mode 100644
index 2d1afc25..00000000
--- a/src/objectives/unitary_infidelity_objective.jl
+++ /dev/null
@@ -1,190 +0,0 @@
-export UnitaryInfidelityObjective
-export UnitaryFreePhaseInfidelityObjective
-
-###
-### UnitaryInfidelityObjective
-###
-
-"""
-    UnitaryInfidelityObjective
-
-A type of objective that measures the infidelity of a unitary operator to a target unitary operator.
-
-Fields:
-    `name`: the name of the unitary operator in the trajectory
-    `goal`: the target unitary operator
-    `Q`: a scaling factor
-    `eval_hessian`: whether to evaluate the Hessian
-    `subspace`: the subspace in which to evaluate the objective
-
-"""
-function UnitaryInfidelityObjective(;
-    name::Union{Nothing,Symbol}=nothing,
-    goal::Union{Nothing,AbstractVector{<:Real}}=nothing,
-	Q::Float64=100.0,
-	eval_hessian::Bool=true,
-    subspace=nothing
-)
-    @assert !isnothing(goal) "unitary goal name must be specified"
-
-    loss = :UnitaryInfidelityLoss
-    if isnothing(subspace)
-        l = eval(loss)(name, goal)
-    else
-        l = eval(loss)(name, goal; subspace=subspace)
-    end
-
-    params = Dict(
-        :type => :UnitaryInfidelityObjective,
-        :name => name,
-        :goal => goal,
-        :Q => Q,
-        :eval_hessian => eval_hessian,
-        :subspace => subspace
-    )
-
-	@views function L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        return Q * l(Z⃗[slice(Z.T, Z.components[name], Z.dim)])
-    end
-
-    @views function ∇L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        ∇ = zeros(Z.dim * Z.T + Z.global_dim)
-        Ũ⃗_slice = slice(Z.T, Z.components[name], Z.dim)
-        Ũ⃗ = Z⃗[Ũ⃗_slice]
-        ∇l = l(Ũ⃗; gradient=true)
-        ∇[Ũ⃗_slice] = Q * ∇l
-        return ∇
-    end
-
-    function ∂²L_structure(Z::NamedTrajectory)
-        final_time_offset = index(Z.T, 0, Z.dim)
-        comp_start_offset = Z.components[name][1] - 1
-        structure = [
-            ij .+ (final_time_offset + comp_start_offset)
-                for ij ∈ l.∇²l_structure
-        ]
-        return structure
-    end
-
-
-    @views function ∂²L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory; return_moi_vals=true)
-        H = spzeros(Z.dim * Z.T + Z.global_dim, Z.dim * Z.T + Z.global_dim)
-        Ũ⃗_slice = slice(Z.T, Z.components[name], Z.dim)
-        H[Ũ⃗_slice, Ũ⃗_slice] = Q * l(Z⃗[Ũ⃗_slice]; hessian=true)
-        if return_moi_vals
-            Hs = [H[i,j] for (i, j) ∈ ∂²L_structure(Z)]
-            return Hs
-        else
-            return H
-        end
-    end
-
-	return Objective(L, ∇L, ∂²L, ∂²L_structure, Dict[params])
-end
-
-function UnitaryInfidelityObjective(
-    name::Symbol,
-    traj::NamedTrajectory,
-    Q::Float64;
-    subspace=nothing,
-	eval_hessian::Bool=true
-)
-    return UnitaryInfidelityObjective(
-        name=name,
-        goal=traj.goal[name],
-        Q=Q,
-        subspace=subspace,
-        eval_hessian=eval_hessian
-    )
-end
-
-###
-### UnitaryFreePhaseInfidelityObjective
-###
-
-"""
-    UnitaryFreePhaseInfidelityObjective
-
-A type of objective that measures the infidelity of a unitary operator to a target unitary operator,
-where the target unitary operator is allowed to have phases on qubit subspaces.
-
-Fields:
-    `name`: the name of the unitary operator in the trajectory
-    `global_name`: the name of the global phase in the trajectory
-    `goal`: the target unitary operator
-    `Q`: a scaling factor
-    `eval_hessian`: whether to evaluate the Hessian
-    `subspace`: the subspace in which to evaluate the objective
-
-"""
-function UnitaryFreePhaseInfidelityObjective(;
-    name::Union{Nothing,Symbol}=nothing,
-    goal::Union{Nothing,AbstractVector{<:R}}=nothing,
-    phase_name::Union{Nothing,Symbol}=nothing,
-    phase_operators::Union{Nothing,AbstractVector{<:AbstractMatrix{<:Complex{R}}}}=nothing,
-	Q::R=1.0,
-	eval_hessian::Bool=false,
-    subspace=nothing
-) where R <: Real
-    @assert !isnothing(goal) "unitary goal name must be specified"
-    @assert !isnothing(name) "unitary name must be specified"
-    @assert !isnothing(phase_name) "phase name must be specified"
-
-    loss = :UnitaryFreePhaseInfidelityLoss
-    l = eval(loss)(goal, phase_operators; subspace=subspace)
-
-    params = Dict(
-        :type => :UnitaryFreePhaseInfidelityObjective,
-        :name => name,
-        :phase_name => phase_name,
-        :goal => goal,
-        :phase_operators => phase_operators,
-        :Q => Q,
-        :eval_hessian => eval_hessian,
-        :subspace => subspace
-    )
-
-	@views function L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        Ũ⃗ = Z⃗[slice(Z.T, Z.components[name], Z.dim)]
-        ϕ⃗ = Z⃗[Z.global_components[phase_name]]
-        return Q * l(Ũ⃗, ϕ⃗)
-    end
-
-    @views function ∇L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        ∇ = zeros(Z.dim * Z.T + Z.global_dim)
-        Ũ⃗_slice = slice(Z.T, Z.components[name], Z.dim)
-        Ũ⃗ = Z⃗[Ũ⃗_slice]
-        ϕ⃗_slice = Z.global_components[phase_name]
-        ϕ⃗ = Z⃗[ϕ⃗_slice]
-        ∇l = l(Ũ⃗, ϕ⃗; gradient=true)
-        ∇[Ũ⃗_slice] = Q * ∇l[1:length(Ũ⃗)]
-        # WARNING: 2π periodic; using Q≠1 is not recommended
-        ∇[ϕ⃗_slice] = Q * ∇l[length(Ũ⃗) .+ (1:length(ϕ⃗))]
-        return ∇
-    end
-
-    ∂²L_structure(Z::NamedTrajectory) = []
-    ∂²L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory) = []
-
-	return Objective(L, ∇L, ∂²L, ∂²L_structure, Dict[params])
-end
-
-function UnitaryFreePhaseInfidelityObjective(
-    name::Symbol,
-    phase_name::Symbol,
-    phase_operators::AbstractVector{<:AbstractMatrix{<:Complex}},
-    traj::NamedTrajectory,
-    Q::Float64;
-    subspace=nothing,
-	eval_hessian::Bool=true
-)
-    return UnitaryFreePhaseInfidelityObjective(
-        name=name,
-        goal=traj.goal[name],
-        phase_name=phase_name,
-        phase_operators=phase_operators,
-        Q=Q,
-        subspace=subspace,
-        eval_hessian=eval_hessian,
-    )
-end
diff --git a/src/objectives/unitary_robustness_objective.jl b/src/objectives/unitary_robustness_objective.jl
deleted file mode 100644
index 84bf8fc2..00000000
--- a/src/objectives/unitary_robustness_objective.jl
+++ /dev/null
@@ -1,270 +0,0 @@
-export UnitaryRobustnessObjective
-export PairwiseUnitaryRobustnessObjective
-
-### 
-### UnitaryRobustnessObjective
-###
-
-@doc raw"""
-UnitaryRobustnessObjective(;
-    H::::Union{OperatorType, Nothing}=nothing,
-    eval_hessian::Bool=false,
-    symb::Symbol=:Ũ⃗
-)
-
-Create a control objective which penalizes the sensitivity of the infidelity to the provided
-operator defined in the subspace of the control dynamics, thereby realizing robust control.
-
-The control dynamics are
-```math
-U_C(a)= \prod_t \exp{-i H_C(a_t)}
-```
-
-In the control frame, the H operator is (proportional to)
-```math
-R_{Robust}(a) = \frac{1}{T \norm{H_e}_2} \sum_t U_C(a_t)^\dag H_e U_C(a_t) \Delta t
-```
-where we have adjusted to a unitless expression of the operator.
-
-The robustness objective is
-```math
-R_{Robust}(a) = \frac{1}{N} \norm{R}^2_F
-```
-where N is the dimension of the Hilbert space.
-"""
-function UnitaryRobustnessObjective(;
-    H_error::Union{OperatorType, Nothing}=nothing,
-    eval_hessian::Bool=false,
-    symb::Symbol=:Ũ⃗
-)
-    @assert !isnothing(H_error) "H_error must be specified"
-
-    # Indices of all non-zero subspace components for iso_vec_operators
-    if H_error isa EmbeddedOperator
-        H = unembed(H_error)
-        subspace = get_iso_vec_subspace_indices(H_error)
-    else
-        H = H_error
-        subspace = 1:length(operator_to_iso_vec(H))
-    end
-
-    @views function get_timesteps(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        return map(1:Z.T) do t
-            if Z.timestep isa Symbol
-                Z⃗[slice(t, Z.components[Z.timestep], Z.dim)][1]
-            else
-                Z.timestep
-            end
-        end
-    end
-
-    # Control frame
-    @views function rotate(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        Δts = get_timesteps(Z⃗, Z)
-        T = sum(Δts)
-        Z_comps = Z.components[symb][subspace]
-        R = sum(
-            map(1:Z.T) do t
-                Uₜ = iso_vec_to_operator(Z⃗[slice(t, Z_comps, Z.dim)])
-                Uₜ'H*Uₜ .* Δts[t]
-            end
-        ) / norm(H) / T
-        return R
-    end
-
-    function L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        R = rotate(Z⃗, Z)
-        return real(tr(R'R)) / size(R, 1)
-    end
-
-    @views function ∇L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        ∇ = zeros(eltype(Z⃗), Z.dim * Z.T + Z.global_dim)
-        R = rotate(Z⃗, Z)
-        Δts = get_timesteps(Z⃗, Z)
-        Z_comps = Z.components[symb][subspace]
-        T = sum(Δts)
-        units = 1 / norm(H) / T
-        Threads.@threads for t ∈ 1:Z.T
-            # State
-            Uₜ_slice = slice(t, Z_comps, Z.dim)
-            Uₜ = iso_vec_to_operator(Z⃗[Uₜ_slice])
-
-            # State gradient
-            ∇[Uₜ_slice] .= operator_to_iso_vec(2 * H * Uₜ * R * Δts[t]) * units
-
-            # Time gradient
-            if Z.timestep isa Symbol
-                ∂R = Uₜ'H*Uₜ
-                ∇[slice(t, Z.components[Z.timestep], Z.dim)] .= real(tr(∂R*R + R*∂R)) * units
-            end
-        end
-        return ∇ / size(R, 1)
-    end
-
-    # Hessian is dense (Control frame R contains sum over all unitaries).
-    if eval_hessian
-        # TODO
-		∂²L = (Z⃗, Z) -> []
-		∂²L_structure = Z -> []
-	else
-		∂²L = nothing
-		∂²L_structure = nothing
-	end
-
-    params = Dict(
-        :type => :UnitaryRobustnessObjective,
-        :H_error => H_error,
-        :eval_hessian => eval_hessian,
-        :symb => symb
-    )
-
-    return Objective(L, ∇L, ∂²L, ∂²L_structure, Dict[params])
-end
-
-### 
-### PairwiseUnitaryRobustnessObjective
-###
-
-"""
-    PairwiseUnitaryRobustnessObjective(;
-        H1::Union{OperatorType, Nothing}=nothing,
-        H2_error::Union{OperatorType, Nothing}=nothing,
-        symb1::Symbol=:Ũ⃗1,
-        symb2::Symbol=:Ũ⃗2,
-        eval_hessian::Bool=false,
-    )
-
-Create a control objective which penalizes the sensitivity of the infidelity to the provided operators
-defined in the subspaces of the control dynamics, thereby realizing robust control.
-"""
-function PairwiseUnitaryRobustnessObjective(;
-    H1_error::Union{OperatorType, Nothing}=nothing,
-    H2_error::Union{OperatorType, Nothing}=nothing,
-    symb1::Symbol=:Ũ⃗1,
-    symb2::Symbol=:Ũ⃗2,
-    eval_hessian::Bool=false,
-)
-    @assert !isnothing(H1_error) "H1_error must be specified"
-    @assert !isnothing(H2_error) "H2_error must be specified"
-
-    if H1_error isa EmbeddedOperator
-        H1 = unembed(H1_error)
-        subspace1 = get_iso_vec_subspace_indices(H1_error)
-    else
-        H1 = H1_error
-        subspace1 = 1:length(operator_to_iso_vec(H1))
-    end
-
-    if H2_error isa EmbeddedOperator
-        H2 = unembed(H2_error)
-        subspace2 = get_iso_vec_subspace_indices(H2_error)
-    else
-        H2 = H2_error
-        subspace2 = 1:length(operator_to_iso_vec(H2))
-    end
-
-    @views function get_timesteps(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        return map(1:Z.T) do t
-            if Z.timestep isa Symbol
-                Z⃗[slice(t, Z.components[Z.timestep], Z.dim)][1]
-            else
-                Z.timestep
-            end
-        end
-    end
-
-    function L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        Δts = get_timesteps(Z⃗, Z)
-        Z1_comps = Z.components[symb1][subspace1]
-        Z2_comps = Z.components[symb2][subspace2]
-        T = sum(Δts)
-        R = 0.0
-        for (i₁, Δt₁) ∈ enumerate(Δts)
-            for (i₂, Δt₂) ∈ enumerate(Δts)
-                # States
-                U1ₜ₁ = iso_vec_to_operator(Z⃗[slice(i₁, Z1_comps, Z.dim)])
-                U1ₜ₂ = iso_vec_to_operator(Z⃗[slice(i₂, Z1_comps, Z.dim)])
-                U2ₜ₁ = iso_vec_to_operator(Z⃗[slice(i₁, Z2_comps, Z.dim)])
-                U2ₜ₂ = iso_vec_to_operator(Z⃗[slice(i₂, Z2_comps, Z.dim)])
-
-                # Rotating frame
-                rH1ₜ₁ = U1ₜ₁'H1*U1ₜ₁
-                rH1ₜ₂ = U1ₜ₂'H1*U1ₜ₂
-                rH2ₜ₁ = U2ₜ₁'H2*U2ₜ₁
-                rH2ₜ₂ = U2ₜ₂'H2*U2ₜ₂
-
-                # Robustness
-                units = 1 / T^2 / norm(H1)^2 / norm(H2)^2
-                R += real(tr(rH1ₜ₁'rH1ₜ₂) * tr(rH2ₜ₁'rH2ₜ₂) * Δt₁ * Δt₂ * units)
-            end
-        end
-        return R / size(H1, 1) / size(H2, 1)
-    end
-
-    @views function ∇L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        ∇ = zeros(Z.dim * Z.T + Z.global_dim)
-        Δts = get_timesteps(Z⃗, Z)
-        Z1_comps = Z.components[symb1][subspace1]
-        Z2_comps = Z.components[symb2][subspace2]
-        T = sum(Δts)
-        Threads.@threads for (i₁, i₂) ∈ vec(collect(Iterators.product(1:Z.T, 1:Z.T)))
-            # Times
-            Δt₁ = Δts[i₁]
-            Δt₂ = Δts[i₂]
-
-            # States
-            U1ₜ₁_slice = slice(i₁, Z1_comps, Z.dim)
-            U1ₜ₂_slice = slice(i₂, Z1_comps, Z.dim)
-            U2ₜ₁_slice = slice(i₁, Z2_comps, Z.dim)
-            U2ₜ₂_slice = slice(i₂, Z2_comps, Z.dim)
-            U1ₜ₁ = iso_vec_to_operator(Z⃗[U1ₜ₁_slice])
-            U1ₜ₂ = iso_vec_to_operator(Z⃗[U1ₜ₂_slice])
-            U2ₜ₁ = iso_vec_to_operator(Z⃗[U2ₜ₁_slice])
-            U2ₜ₂ = iso_vec_to_operator(Z⃗[U2ₜ₂_slice])
-
-            # Rotating frame
-            rH1ₜ₁ = U1ₜ₁'H1*U1ₜ₁
-            rH1ₜ₂ = U1ₜ₂'H1*U1ₜ₂
-            rH2ₜ₁ = U2ₜ₁'H2*U2ₜ₁
-            rH2ₜ₂ = U2ₜ₂'H2*U2ₜ₂
-
-            # ∇Uiₜⱼ (assume H's are Hermitian)
-            units = 1 / T^2 / norm(H1)^2 / norm(H2)^2
-            R1 = tr(rH1ₜ₁'rH1ₜ₂) * Δt₁ * Δt₂ * units
-            R2 = tr(rH2ₜ₁'rH2ₜ₂) * Δt₁ * Δt₂ * units
-            ∇[U1ₜ₁_slice] += operator_to_iso_vec(2 * H1 * U1ₜ₁ * rH1ₜ₂) * R2
-            ∇[U1ₜ₂_slice] += operator_to_iso_vec(2 * H1 * U1ₜ₂ * rH1ₜ₁) * R2
-            ∇[U2ₜ₁_slice] += operator_to_iso_vec(2 * H2 * U2ₜ₁ * rH2ₜ₂) * R1
-            ∇[U2ₜ₂_slice] += operator_to_iso_vec(2 * H2 * U2ₜ₂ * rH2ₜ₁) * R1
-
-            # Time gradients
-            if Z.timestep isa Symbol
-                R = real(tr(rH1ₜ₁'rH1ₜ₂) * tr(rH2ₜ₁'rH2ₜ₂)) * units
-                ∇[slice(i₁, Z.components[Z.timestep], Z.dim)] .= R * Δt₂
-                ∇[slice(i₂, Z.components[Z.timestep], Z.dim)] .= R * Δt₁
-            end
-        end
-        return ∇ / size(H1, 1) / size(H2, 1)
-    end
-
-    # Hessian is dense (Control frame R contains sum over all unitaries).
-    if eval_hessian
-        # TODO
-		∂²L = (Z⃗, Z) -> []
-		∂²L_structure = Z -> []
-	else
-		∂²L = nothing
-		∂²L_structure = nothing
-	end
-
-    params = Dict(
-        :type => :PairwiseUnitaryRobustnessObjective,
-        :H1_error => H1_error,
-        :H2_error => H2_error,
-        :symb1 => symb1,
-        :symb2 => symb2,
-        :eval_hessian => eval_hessian
-    )
-
-    return Objective(L, ∇L, ∂²L, ∂²L_structure, Dict[params])
-end
\ No newline at end of file
diff --git a/src/options.jl b/src/options.jl
deleted file mode 100644
index be15c839..00000000
--- a/src/options.jl
+++ /dev/null
@@ -1,122 +0,0 @@
-module Options
-
-export IpoptOptions
-export PiccoloOptions
-export set!
-
-using Ipopt
-using Libdl
-using ExponentialAction
-using Base: @kwdef
-
-abstract type AbstractOptions end
-
-"""
-    Solver options for Ipopt
-
-    https://coin-or.github.io/Ipopt/OPTIONS.html#OPT_print_options_documentation
-"""
-@kwdef mutable struct IpoptOptions{T} <: AbstractOptions
-    tol::T = 1e-8
-    s_max::T = 100.0
-    max_iter::Int = 1_000
-    max_cpu_time = 1_000_000.0
-    dual_inf_tol::T = 1.0
-    constr_viol_tol::T = 1.0e-6
-    compl_inf_tol::T = 1.0e-3
-    acceptable_tol::T = 1.0e-6
-    acceptable_iter::Int = 15
-    acceptable_dual_inf_tol::T = 1.0e10
-    acceptable_constr_viol_tol::T = 1.0e-2
-    acceptable_compl_inf_tol::T = 1.0e-2
-    acceptable_obj_change_tol::T = 1.0e-5
-    diverging_iterates_tol::T = 1.0e8
-    mu_target::T = 1.0e-4
-    print_level::Int = 5
-    output_file = nothing
-    print_user_options = "no"
-    print_options_documentation = "no"
-    print_timing_statistics = "no"
-    print_options_mode = "text"
-    print_advanced_options = "no"
-    print_info_string = "no"
-    inf_pr_output = "original"
-    print_frequency_iter = 1
-    print_frequency_time = 0.0
-    skip_finalize_solution_call = "no"
-    hsllib = nothing
-    hessian_approximation = "exact"
-    recalc_y = "no"
-    recalc_y_feas_tol = 1.0e-6
-    linear_solver = "mumps"
-    watchdog_shortened_iter_trigger = 10
-    watchdog_trial_iter_max = 3
-end
-
-
-"""
-    PiccoloOptions
-
-Options for the Piccolo quantum optimal control library.
-
-# Fields
-- `verbose::Bool = true`: Print verbose output
-- `verbose_evaluator::Bool = false`: Print verbose output from the evaluator
-- `free_time::Bool = true`: Allow free time optimization
-- `timesteps_all_equal::Bool = true`: Use equal timesteps
-- `integrator::Symbol = :pade`: Integrator to use
-- `pade_order::Int = 4`: Order of the Pade approximation
-- `rollout_integrator::Function = expv`: Integrator to use for rollout
-- `jacobian_structure::Bool = integrator == :pade`: Use the integrator's Jacobian structure
-- `eval_hessian::Bool = false`: Evaluate the Hessian
-- `geodesic = true`: Use the geodesic to initialize the optimization.
-- `blas_multithreading::Bool = true`: Use BLAS multithreading.
-- `build_trajectory_constraints::Bool = true`: Build trajectory constraints.
-- `complex_control_norm_constraint_name::Union{Nothing, Symbol} = nothing`: Name of the complex control norm constraint.
-- `complex_control_norm_constraint_radius::Float64 = 1.0`: Radius of the complex control norm constraint.
-- `bound_state::Bool = false`: Bound the state.
-- `leakage_suppression::Bool = false`: Suppress leakage.
-- `R_leakage::Float64 = 1.0`: Leakage suppression parameter.
-- `free_phase_infidelity::Bool = false`: Free phase infidelity.
-- `phase_operators::Union{Nothing, AbstractVector{<:AbstractMatrix{<:Complex}}} = nothing`: Phase operators.
-- `phase_name::Symbol = :ϕ`: Name of the phase.
-"""
-@kwdef mutable struct PiccoloOptions <: AbstractOptions
-    verbose::Bool = true
-    verbose_evaluator::Bool = false
-    free_time::Bool = true
-    timesteps_all_equal::Bool = true
-    integrator::Symbol = :pade
-    pade_order::Int = 4
-    jacobian_structure::Bool = integrator == :pade
-    rollout_integrator::Function = expv
-    eval_hessian::Bool = false
-    geodesic = true
-    blas_multithreading::Bool = true
-    build_trajectory_constraints::Bool = true
-    complex_control_norm_constraint_name::Union{Nothing, Symbol} = nothing
-    complex_control_norm_constraint_radius::Float64 = 1.0
-    bound_state::Bool = false
-    leakage_suppression::Bool = false
-    R_leakage::Float64 = 1.0
-    free_phase_infidelity::Bool = false
-    phase_operators::Union{Nothing, AbstractVector{<:AbstractMatrix{<:Complex}}} = nothing
-    phase_name::Symbol = :ϕ
-end
-
-function set!(optimizer::Ipopt.Optimizer, options::AbstractOptions)
-    for name in fieldnames(typeof(options))
-        value = getfield(options, name)
-        if name == :linear_solver
-            if value == "pardiso"
-                Libdl.dlopen("/usr/lib/x86_64-linux-gnu/liblapack.so.3", RTLD_GLOBAL)
-                Libdl.dlopen("/usr/lib/x86_64-linux-gnu/libomp.so.5", RTLD_GLOBAL)
-            end
-        end
-        if !isnothing(value)
-           optimizer.options[String(name)] = value
-        end
-    end
-end
-
-end
diff --git a/src/plotting.jl b/src/plotting.jl
deleted file mode 100644
index fcadb408..00000000
--- a/src/plotting.jl
+++ /dev/null
@@ -1,71 +0,0 @@
-module Plotting
-
-export plot_unitary_populations
-export pretty_print
-
-using NamedTrajectories
-
-using ..Isomorphisms
-using ..Problems
-
-
-pretty_print(X::AbstractMatrix) = Base.show(stdout,"text/plain", X)
-
-function get_layout(index::Int, n::Int)
-    √n = isqrt(n) + 1
-    return ((index - 1) ÷ √n + 1, (index - 1) % √n + 1)
-end
-
-"""
-    plot_unitary_populations(
-        traj::NamedTrajectory;
-        unitary_columns::AbstractVector{Int}=1:2,
-        unitary_name::Symbol=:Ũ⃗,
-        control_name::Symbol=:a,
-        kwargs...
-    )
-
-    plot_unitary_populations(
-        prob::QuantumControlProblem;
-        kwargs...
-    )
-
-Plot the populations of the unitary columns of the unitary matrix in the trajectory. `kwargs` are passed to [`NamedTrajectories.plot`](https://aarontrowbridge.github.io/NamedTrajectories.jl/dev/generated/plotting/).
-"""
-function plot_unitary_populations end
-
-function plot_unitary_populations(
-    traj::NamedTrajectory;
-    unitary_columns::AbstractVector{Int}=1:2,
-    unitary_name::Symbol=:Ũ⃗,
-    control_name::Symbol=:a,
-    kwargs...
-)
-
-    transformations = OrderedDict(
-        unitary_name => [
-            x -> abs2.(iso_vec_to_operator(x)[:, i])
-                for i ∈ unitary_columns
-        ]
-    )
-
-    transformation_titles = OrderedDict(
-        unitary_name => [
-            L"Populations: $\left| U_{:, %$(i)}(t) \right|^2$"
-                for i ∈ unitary_columns
-        ]
-    )
-
-    plot(traj, [control_name];
-        transformations=transformations,
-        transformation_titles=transformation_titles,
-        include_transformation_labels=true,
-        kwargs...
-    )
-end
-
-function plot_unitary_populations(prob::QuantumControlProblem; kwargs...)
-    plot_unitary_populations(prob.trajectory; kwargs...)
-end
-
-end
diff --git a/src/problem_solvers.jl b/src/problem_solvers.jl
deleted file mode 100644
index 2169d976..00000000
--- a/src/problem_solvers.jl
+++ /dev/null
@@ -1,100 +0,0 @@
-module ProblemSolvers
-
-export solve!
-
-using ..Constraints
-using ..Problems
-using ..SaveLoadUtils
-using ..Options
-
-using NamedTrajectories
-using MathOptInterface
-using Ipopt
-const MOI = MathOptInterface
-
-
-"""
-   solve!(prob::QuantumControlProblem;
-        init_traj=nothing,
-        save_path=nothing,
-        max_iter=prob.ipopt_options.max_iter,
-        linear_solver=prob.ipopt_options.linear_solver,
-        print_level=prob.ipopt_options.print_level,
-        remove_slack_variables=false,
-        callback=nothing
-        # state_type=:unitary,
-        # print_fidelity=false,
-    )
-
-    Call optimization solver to solve the quantum control problem with parameters and callbacks.
-
-# Arguments
-- `prob::QuantumControlProblem`: The quantum control problem to solve.
-- `init_traj::NamedTrajectory`: Initial guess for the control trajectory. If not provided, a random guess will be generated.
-- `save_path::String`: Path to save the problem after optimization.
-- `max_iter::Int`: Maximum number of iterations for the optimization solver.
-- `linear_solver::String`: Linear solver to use for the optimization solver (e.g., "mumps", "paradiso", etc).
-- `print_level::Int`: Verbosity level for the solver.
-- `remove_slack_variables::Bool`: Remove slack variables from the trajectory after optimization.
-- `callback::Function`: Callback function to call during optimization steps.
-"""
-function solve!(
-    prob::QuantumControlProblem;
-    init_traj=nothing,
-    save_path=nothing,
-    max_iter::Int=prob.ipopt_options.max_iter,
-    linear_solver::String=prob.ipopt_options.linear_solver,
-    print_level::Int=prob.ipopt_options.print_level,
-    remove_slack_variables::Bool=false,
-    callback=nothing
-    # state_type::Symbol=:unitary,
-    # print_fidelity::Bool=false,
-)
-    # @assert state_type in (:ket, :unitary, :density_matrix) "Invalid state type: $state_type must be one of :ket, :unitary, or :density_matrix"
-
-    prob.ipopt_options.max_iter = max_iter
-    prob.ipopt_options.linear_solver = linear_solver
-    prob.ipopt_options.print_level = print_level
-
-    set!(prob.optimizer, prob.ipopt_options)
-
-    if !isnothing(init_traj)
-        set_trajectory!(prob, init_traj)
-    else
-        set_trajectory!(prob)
-    end
-    
-    if !isnothing(callback)
-        MOI.set(prob.optimizer, Ipopt.CallbackFunction(), callback)
-    end
-
-    # if print_fidelity
-    #     if state_type == :ket
-    #         fids = fidelity(prob)
-    #         println("\nInitial Fidelities: $fids")
-    #     elseif state_type == :unitary
-    #         fids = unitary_fidelity(prob)
-    #         println("\nInitial Fidelity: $fids")
-
-
-    MOI.optimize!(prob.optimizer)
-
-    update_trajectory!(prob)
-
-    if remove_slack_variables
-        slack_var_names = Symbol[]
-        for con in prob.params[:linear_constraints]
-            if con isa L1SlackConstraint
-                append!(slack_var_names, con.slack_names)
-            end
-        end
-
-        prob.trajectory = remove_components(prob.trajectory, slack_var_names)
-    end
-
-    if !isnothing(save_path)
-        save_problem(save_path, prob)
-    end
-end
-
-end
diff --git a/src/problem_templates/_problem_templates.jl b/src/problem_templates/_problem_templates.jl
index 1b328308..985cafc4 100644
--- a/src/problem_templates/_problem_templates.jl
+++ b/src/problem_templates/_problem_templates.jl
@@ -1,27 +1,19 @@
 module ProblemTemplates
 
-using ..QuantumObjectUtils
-using ..QuantumSystems
-using ..Isomorphisms
-using ..EmbeddedOperators
 using ..DirectSums
 using ..Rollouts
 using ..TrajectoryInitialization
-using ..Objectives
-using ..Constraints
 using ..Losses
-using ..Integrators
-using ..Problems
-using ..Options
 
 using Distributions
 using TrajectoryIndexingUtils
 using NamedTrajectories
+using QuantumCollocationCore
+using PiccoloQuantumObjects
 using LinearAlgebra
 using SparseArrays
 using ExponentialAction
 using JLD2
-
 using TestItemRunner
 
 include("unitary_smooth_pulse_problem.jl")
@@ -35,14 +27,15 @@ include("quantum_state_smooth_pulse_problem.jl")
 include("quantum_state_minimum_time_problem.jl")
 include("quantum_state_sampling_problem.jl")
 
+include("density_operator_smooth_pulse_problem.jl")
 
 function apply_piccolo_options!(
     J::Objective,
     constraints::AbstractVector{<:AbstractConstraint},
     piccolo_options::PiccoloOptions,
     traj::NamedTrajectory,
-    operators::Union{Nothing, AbstractVector{<:OperatorType}},
-    state_names::AbstractVector{<:Symbol},
+    operators::Union{Nothing, AbstractVector{<:AbstractPiccoloOperator}},
+    state_names::AbstractVector{Symbol},
     timestep_name::Symbol
 )
     # TODO: This should be changed to leakage indices (more general, for states)
@@ -91,7 +84,7 @@ function apply_piccolo_options!(
     constraints::AbstractVector{<:AbstractConstraint},
     piccolo_options::PiccoloOptions,
     traj::NamedTrajectory,
-    operator::Union{Nothing, OperatorType},
+    operator::Union{Nothing, AbstractPiccoloOperator},
     state_name::Symbol,
     timestep_name::Symbol
 )
diff --git a/src/problem_templates/density_operator_smooth_pulse_problem.jl b/src/problem_templates/density_operator_smooth_pulse_problem.jl
new file mode 100644
index 00000000..90be0f05
--- /dev/null
+++ b/src/problem_templates/density_operator_smooth_pulse_problem.jl
@@ -0,0 +1,123 @@
+export DensityOperatorSmoothPulseProblem
+
+function DensityOperatorSmoothPulseProblem(
+    system::OpenQuantumSystem,
+    ρ_init::AbstractMatrix,
+    ψ_goal::AbstractVector,
+    T::Int,
+    Δt::Union{Float64, Vector{Float64}};
+    ipopt_options::IpoptOptions=IpoptOptions(),
+    piccolo_options::PiccoloOptions=PiccoloOptions(),
+    constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
+    init_trajectory::Union{NamedTrajectory, Nothing}=nothing,
+    a_bound::Float64=1.0,
+    a_bounds=fill(a_bound, system.n_drives),
+    a_guess::Union{Matrix{Float64}, Nothing}=nothing,
+    da_bound::Float64=Inf,
+    da_bounds=fill(da_bound, system.n_drives),
+    dda_bound::Float64=1.0,
+    dda_bounds=fill(dda_bound, system.n_drives),
+    Δt_min::Float64=Δt isa Float64 ? 0.5 * Δt : 0.5 * mean(Δt),
+    Δt_max::Float64=Δt isa Float64 ? 1.5 * Δt : 1.5 * mean(Δt),
+    drive_derivative_σ::Float64=0.01,
+    Q::Float64=100.0,
+    R=1e-2,
+    R_a::Union{Float64, Vector{Float64}}=R,
+    R_da::Union{Float64, Vector{Float64}}=R,
+    R_dda::Union{Float64, Vector{Float64}}=R,
+    leakage_suppression=false,
+    R_leakage=1e-1,
+    control_norm_constraint=false,
+    control_norm_constraint_components=nothing,
+    control_norm_R=nothing,
+    kwargs...
+)
+    # Trajectory
+    if !isnothing(init_trajectory)
+        traj = init_trajectory
+    else
+        traj = initialize_trajectory(
+            ρ_init,
+            ψ_goal*ψ_goal',
+            T,
+            Δt,
+            system.n_drives,
+            (a_bounds, da_bounds, dda_bounds);
+            free_time=piccolo_options.free_time,
+            Δt_bounds=(Δt_min, Δt_max),
+            drive_derivative_σ=drive_derivative_σ,
+            a_guess=a_guess,
+            system=system,
+        )
+    end
+
+    # Objective
+    J = DensityOperatorPureStateInfidelityObjective(:ρ⃗̃, ψ_goal; Q=Q)
+    J += QuadraticRegularizer(:a, traj, R_a)
+    J += QuadraticRegularizer(:da, traj, R_da)
+    J += QuadraticRegularizer(:dda, traj, R_dda)
+
+    # Constraints
+    if leakage_suppression
+        if operator isa EmbeddedOperator
+            leakage_indices = get_iso_vec_leakage_indices(operator)
+            J += L1Regularizer!(
+                constraints, :Ũ⃗, traj,
+                R_value=R_leakage,
+                indices=leakage_indices,
+                eval_hessian=piccolo_options.eval_hessian
+            )
+        else
+            @warn "leakage_suppression is not supported for non-embedded operators, ignoring."
+        end
+    end
+
+    if piccolo_options.free_time
+        if piccolo_options.timesteps_all_equal
+            push!(constraints, TimeStepsAllEqualConstraint(:Δt, traj))
+        end
+    end
+
+    if control_norm_constraint
+        @assert !isnothing(control_norm_constraint_components) "control_norm_constraint_components must be provided"
+        @assert !isnothing(control_norm_R) "control_norm_R must be provided"
+        norm_con = ComplexModulusContraint(
+            :a,
+            control_norm_R,
+            traj;
+            name_comps=control_norm_constraint_components,
+        )
+        push!(constraints, norm_con)
+    end
+
+    # Integrators
+    # if piccolo_options.integrator == :pade
+    #     unitary_integrator =
+    #         UnitaryPadeIntegrator(system, :Ũ⃗, :a, traj; order=piccolo_options.pade_order)
+    # elseif piccolo_options.integrator == :exponential
+    #     unitary_integrator =
+    #         UnitaryExponentialIntegrator(system, :Ũ⃗, :a, traj)
+    # else
+    #     error("integrator must be one of (:pade, :exponential)")
+    # end
+
+    density_operator_integrator = DensityOperatorExponentialIntegrator(
+        :ρ⃗̃, :a, system, traj
+    )
+
+    integrators = [
+        density_operator_integrator,
+        DerivativeIntegrator(:a, :da, traj),
+        DerivativeIntegrator(:da, :dda, traj),
+    ]
+
+    return QuantumControlProblem(
+        traj,
+        J,
+        integrators;
+        constraints=constraints,
+        ipopt_options=ipopt_options,
+        piccolo_options=piccolo_options,
+        kwargs...
+    )
+end
diff --git a/src/problem_templates/quantum_state_minimum_time_problem.jl b/src/problem_templates/quantum_state_minimum_time_problem.jl
index 6deeaf5d..6b644c26 100644
--- a/src/problem_templates/quantum_state_minimum_time_problem.jl
+++ b/src/problem_templates/quantum_state_minimum_time_problem.jl
@@ -29,11 +29,11 @@ function QuantumStateMinimumTimeProblem end
 
 function QuantumStateMinimumTimeProblem(
     traj::NamedTrajectory,
-    sys::QuantumSystem,
     obj::Objective,
     integrators::Vector{<:AbstractIntegrator},
     constraints::Vector{<:AbstractConstraint};
     state_name::Symbol=:ψ̃,
+    control_name::Symbol=:a,
     final_fidelity::Union{Real, Nothing}=nothing,
     D=1.0,
     ipopt_options::IpoptOptions=IpoptOptions(),
@@ -47,7 +47,7 @@ function QuantumStateMinimumTimeProblem(
 
     # Default to average state fidelity
     if isnothing(final_fidelity)
-        vals = [fidelity(traj[n][:, end], traj.goal[n]) for n ∈ state_names]
+        vals = [iso_fidelity(traj[n][:, end], traj.goal[n]) for n ∈ state_names]
         final_fidelity = sum(vals) / length(vals)
     end
 
@@ -63,13 +63,13 @@ function QuantumStateMinimumTimeProblem(
     end
 
     return QuantumControlProblem(
-        sys,
         traj,
         obj,
         integrators;
         constraints=constraints,
         ipopt_options=ipopt_options,
         piccolo_options=piccolo_options,
+        control_name=control_name,
         kwargs...
     )
 end
@@ -87,7 +87,6 @@ function QuantumStateMinimumTimeProblem(
 
     return QuantumStateMinimumTimeProblem(
         copy(prob.trajectory),
-        prob.system,
         obj,
         prob.integrators,
         constraints;
@@ -101,13 +100,14 @@ end
 
 @testitem "Test quantum state minimum time" begin
         using NamedTrajectories
+        using PiccoloQuantumObjects
 
         # System
         T = 50
         Δt = 0.2
         sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]])
-        ψ_init = Vector{ComplexF64}([1.0, 0.0])
-        ψ_target = Vector{ComplexF64}([0.0, 1.0])
+        ψ_init = Vector{ComplexF64}[[1.0, 0.0]]
+        ψ_target = Vector{ComplexF64}[[0.0, 1.0]]
 
         prob = QuantumStateSmoothPulseProblem(
             sys, ψ_init, ψ_target, T, Δt;
diff --git a/src/problem_templates/quantum_state_sampling_problem.jl b/src/problem_templates/quantum_state_sampling_problem.jl
index daa5f5a1..a67657dd 100644
--- a/src/problem_templates/quantum_state_sampling_problem.jl
+++ b/src/problem_templates/quantum_state_sampling_problem.jl
@@ -16,12 +16,12 @@ function QuantumStateSamplingProblem(
     control_name::Symbol=:a,
     timestep_name::Symbol=:Δt,
     a_bound::Float64=1.0,
-    a_bounds=fill(a_bound, length(systems[1].G_drives)),
+    a_bounds=fill(a_bound, systems[1].n_drives),
     a_guess::Union{Matrix{Float64},Nothing}=nothing,
     da_bound::Float64=Inf,
-    da_bounds::Vector{Float64}=fill(da_bound, length(systems[1].G_drives)),
+    da_bounds::Vector{Float64}=fill(da_bound, systems[1].n_drives),
     dda_bound::Float64=1.0,
-    dda_bounds::Vector{Float64}=fill(dda_bound, length(systems[1].G_drives)),
+    dda_bounds::Vector{Float64}=fill(dda_bound, systems[1].n_drives),
     Δt_min::Float64=0.5 * Δt,
     Δt_max::Float64=1.5 * Δt,
     Q::Float64=100.0,
@@ -51,7 +51,7 @@ function QuantumStateSamplingProblem(
                 ψgs,
                 T,
                 Δt,
-                length(sys.G_drives),
+                sys.n_drives,
                 (a_bounds, da_bounds, dda_bounds);
                 state_names=names,
                 control_name=control_name,
@@ -95,12 +95,12 @@ function QuantumStateSamplingProblem(
         for name ∈ names
             if piccolo_options.integrator == :pade
                 state_integrator = QuantumStatePadeIntegrator(
-                    system, name, control_name, traj;
+                    name, control_name, system, traj;
                     order=piccolo_options.pade_order
                 )
             elseif piccolo_options.integrator == :exponential
                 state_integrator = QuantumStateExponentialIntegrator(
-                    system, name, control_name, traj
+                    name, control_name, system, traj
                 )
             else
                 error("integrator must be one of (:pade, :exponential)")
@@ -119,15 +119,15 @@ function QuantumStateSamplingProblem(
     apply_piccolo_options!(
         J, constraints, piccolo_options, traj, leakage_operator, state_name, timestep_name
     )
-    
+
     return QuantumControlProblem(
-        direct_sum(systems),
         traj,
         J,
         integrators;
         constraints=constraints,
         ipopt_options=ipopt_options,
         piccolo_options=piccolo_options,
+        control_name=control_name,
         kwargs...
     )
 end
@@ -182,10 +182,10 @@ end
     state_names = [n for n ∈ prob.trajectory.names if startswith(string(n), string(state_name))]
     sys1_state_names = [n for n ∈ state_names if endswith(string(n), "1")]
 
-    # Separately compute all unique initial and goal state fidelities 
-    init = [fidelity(prob.trajectory, sys1, state_symb=n) for n in sys1_state_names]
+    # Separately compute all unique initial and goal state fidelities
+    init = [rollout_fidelity(prob.trajectory, sys1; state_name=n) for n in sys1_state_names]
     solve!(prob, max_iter=20)
-    final = [fidelity(prob.trajectory, sys1, state_symb=n) for n in sys1_state_names]
+    final = [rollout_fidelity(prob.trajectory, sys1, state_name=n) for n in sys1_state_names]
     @test all(final .> init)
 
     # Check that a_guess can be used
@@ -207,10 +207,10 @@ end
         robustness=false
     )
     solve!(prob_default, max_iter=20)
-    final_default = fidelity(prob_default.trajectory, sys1)
+    final_default = rollout_fidelity(prob_default.trajectory, sys1)
     # Pick any initial state
-    final_robust = fidelity(prob.trajectory, sys1, state_symb=state_names[1])
-    @test round(final_robust, digits=2) ≥ round(final_default, digits=2)
+    final_robust = rollout_fidelity(prob.trajectory, sys1, state_name=state_names[1])
+    @test final_robust > final_default
 end
 
 @testitem "Sample systems with multiple initial, target" begin
@@ -235,9 +235,8 @@ end
     state_names = [n for n ∈ prob.trajectory.names if startswith(string(n), string(state_name))]
     sys1_state_names = [n for n ∈ state_names if endswith(string(n), "1")]
 
-    init = [fidelity(prob.trajectory, sys1, state_symb=n) for n in sys1_state_names]
+    init = [rollout_fidelity(prob.trajectory, sys1, state_name=n) for n in sys1_state_names]
     solve!(prob, max_iter=20)
-    final = [fidelity(prob.trajectory, sys1, state_symb=n) for n in sys1_state_names]
+    final = [rollout_fidelity(prob.trajectory, sys1, state_name=n) for n in sys1_state_names]
     @test all(final .> init)
 end
-
diff --git a/src/problem_templates/quantum_state_smooth_pulse_problem.jl b/src/problem_templates/quantum_state_smooth_pulse_problem.jl
index 21b8cb7f..3ee5648f 100644
--- a/src/problem_templates/quantum_state_smooth_pulse_problem.jl
+++ b/src/problem_templates/quantum_state_smooth_pulse_problem.jl
@@ -38,6 +38,7 @@ with
 - `a_guess::Union{Matrix{Float64}, Nothing}=nothing`: The initial guess for the control pulse.
 - `da_bound::Float64=Inf`: The bound on the first derivative of the control pulse.
 - `da_bounds::Vector{Float64}=fill(da_bound, length(system.G_drives))`: The bounds on the first derivative of the control pulse.
+- `zero_initial_and_final_derivative::Bool=false`: Whether to enforce zero initial and final derivative.
 - `dda_bound::Float64=1.0`: The bound on the second derivative of the control pulse.
 - `dda_bounds::Vector{Float64}=fill(dda_bound, length(system.G_drives))`: The bounds on the second derivative of the control pulse.
 - `Δt_min::Float64=0.5 * Δt`: The minimum timestep size.
@@ -54,7 +55,7 @@ with
 function QuantumStateSmoothPulseProblem end
 
 function QuantumStateSmoothPulseProblem(
-    system::AbstractQuantumSystem,
+    sys::AbstractQuantumSystem,
     ψ_inits::Vector{<:AbstractVector{<:ComplexF64}},
     ψ_goals::Vector{<:AbstractVector{<:ComplexF64}},
     T::Int,
@@ -66,12 +67,13 @@ function QuantumStateSmoothPulseProblem(
     timestep_name::Symbol=:Δt,
     init_trajectory::Union{NamedTrajectory, Nothing}=nothing,
     a_bound::Float64=1.0,
-    a_bounds::Vector{Float64}=fill(a_bound, length(system.G_drives)),
+    a_bounds::Vector{Float64}=fill(a_bound, sys.n_drives),
     a_guess::Union{Matrix{Float64}, Nothing}=nothing,
     da_bound::Float64=Inf,
-    da_bounds::Vector{Float64}=fill(da_bound, length(system.G_drives)),
+    da_bounds::Vector{Float64}=fill(da_bound, sys.n_drives),
+    zero_initial_and_final_derivative::Bool=false,
     dda_bound::Float64=1.0,
-    dda_bounds::Vector{Float64}=fill(dda_bound, length(system.G_drives)),
+    dda_bounds::Vector{Float64}=fill(dda_bound, sys.n_drives),
     Δt_min::Float64=0.5 * Δt,
     Δt_max::Float64=1.5 * Δt,
     drive_derivative_σ::Float64=0.01,
@@ -90,24 +92,23 @@ function QuantumStateSmoothPulseProblem(
     if !isnothing(init_trajectory)
         traj = init_trajectory
     else
-        n_drives = length(system.G_drives)
-
         traj = initialize_trajectory(
             ψ_goals,
             ψ_inits,
             T,
             Δt,
-            n_drives,
+            sys.n_drives,
             (a_bounds, da_bounds, dda_bounds);
             state_name=state_name,
             control_name=control_name,
             timestep_name=timestep_name,
+            zero_initial_and_final_derivative=zero_initial_and_final_derivative,
             free_time=piccolo_options.free_time,
             Δt_bounds=(Δt_min, Δt_max),
             bound_state=piccolo_options.bound_state,
             drive_derivative_σ=drive_derivative_σ,
             a_guess=a_guess,
-            system=system,
+            system=sys,
             rollout_integrator=piccolo_options.rollout_integrator,
         )
     end
@@ -134,20 +135,52 @@ function QuantumStateSmoothPulseProblem(
 
     # Integrators
     state_integrators = []
-    for name ∈ state_names
+    if length(ψ_inits) == 1
         if piccolo_options.integrator == :pade
-            state_integrator = QuantumStatePadeIntegrator(
-                system, name, control_name, traj;
+            state_integrators = [QuantumStatePadeIntegrator(
+                state_name,
+                control_name,
+                sys,
+                traj;
                 order=piccolo_options.pade_order
-            )
+            )]
         elseif piccolo_options.integrator == :exponential
-            state_integrator = QuantumStateExponentialIntegrator(
-                system, name, control_name, traj
-            )
+            state_integrators = [QuantumStateExponentialIntegrator(
+                state_name,
+                control_name,
+                sys,
+                traj
+            )]
         else
             error("integrator must be one of (:pade, :exponential)")
         end
-        push!(state_integrators, state_integrator)
+    else
+        state_names = [
+            name for name ∈ traj.names
+                if startswith(string(name), string(state_name))
+        ]
+        state_integrators = []
+        for i = 1:length(ψ_inits)
+            if piccolo_options.integrator == :pade
+                state_integrator = QuantumStatePadeIntegrator(
+                    state_names[i],
+                    control_name,
+                    sys,
+                    traj;
+                    order=piccolo_options.pade_order
+                )
+            elseif piccolo_options.integrator == :exponential
+                state_integrator = QuantumStateExponentialIntegrator(
+                    state_names[i],
+                    control_name,
+                    sys,
+                    traj
+                )
+            else
+                error("integrator must be one of (:pade, :exponential)")
+            end
+            push!(state_integrators, state_integrator)
+        end
     end
 
     integrators = [
@@ -162,13 +195,13 @@ function QuantumStateSmoothPulseProblem(
     )
 
     return QuantumControlProblem(
-        system,
         traj,
         J,
         integrators;
         constraints=constraints,
         ipopt_options=ipopt_options,
         piccolo_options=piccolo_options,
+        control_name=control_name,
         kwargs...
     )
 end
@@ -210,9 +243,9 @@ end
         ipopt_options=IpoptOptions(print_level=1),
         piccolo_options=PiccoloOptions(verbose=false)
     )
-    initial = fidelity(prob)
+    initial = rollout_fidelity(prob.trajectory, sys)
     solve!(prob, max_iter=20)
-    final = fidelity(prob)
+    final = rollout_fidelity(prob.trajectory, sys)
     @test final > initial
 
     # Multiple initial and target states
@@ -224,9 +257,9 @@ end
         ipopt_options=IpoptOptions(print_level=1),
         piccolo_options=PiccoloOptions(verbose=false)
     )
-    initial = fidelity(prob)
+    initial = rollout_fidelity(prob.trajectory, sys)
     solve!(prob, max_iter=20)
-    final = fidelity(prob)
+    final = rollout_fidelity(prob.trajectory, sys)
     @test all(final .> initial)
 end
 
@@ -246,9 +279,9 @@ end
         ipopt_options=IpoptOptions(print_level=1),
         piccolo_options=PiccoloOptions(verbose=false, integrator=integrator)
     )
-    initial = fidelity(prob)
+    initial = rollout_fidelity(prob.trajectory, sys)
     solve!(prob, max_iter=20)
-    final = fidelity(prob)
+    final = rollout_fidelity(prob.trajectory, sys)
     @test final > initial
 
     # Multiple initial and target states
@@ -260,9 +293,9 @@ end
         ipopt_options=IpoptOptions(print_level=1),
         piccolo_options=PiccoloOptions(verbose=false, integrator=integrator)
     )
-    initial = fidelity(prob)
+    initial = rollout_fidelity(prob.trajectory, sys)
     solve!(prob, max_iter=20)
-    final = fidelity(prob)
+    final = rollout_fidelity(prob.trajectory, sys)
     @test all(final .> initial)
 end
 
@@ -282,8 +315,8 @@ end
         control_name=:u,
         timestep_name=:dt
     )
-    initial = fidelity(prob)
+    initial = rollout_fidelity(prob.trajectory, sys)
     solve!(prob, max_iter=20)
-    final = fidelity(prob)
+    final = rollout_fidelity(prob.trajectory, sys)
     @test all(final .> initial)
 end
diff --git a/src/problem_templates/unitary_bang_bang_problem.jl b/src/problem_templates/unitary_bang_bang_problem.jl
index 8f10c277..274dc898 100644
--- a/src/problem_templates/unitary_bang_bang_problem.jl
+++ b/src/problem_templates/unitary_bang_bang_problem.jl
@@ -37,7 +37,7 @@ TODO: Document bang-bang modification.
 or
 - `system::QuantumSystem`: the system to be controlled
 with
-- `operator::OperatorType`: the target unitary, either in the form of an `EmbeddedOperator` or a `Matrix{ComplexF64}
+- `operator::AbstractPiccoloOperator`: the target unitary, either in the form of an `EmbeddedOperator` or a `Matrix{ComplexF64}
 - `T::Int`: the number of timesteps
 - `Δt::Float64`: the (initial) time step size
 
@@ -69,7 +69,7 @@ function UnitaryBangBangProblem end
 
 function UnitaryBangBangProblem(
     system::AbstractQuantumSystem,
-    operator::OperatorType,
+    operator::AbstractPiccoloOperator,
     T::Int,
     Δt::Union{Float64, Vector{Float64}};
     ipopt_options::IpoptOptions=IpoptOptions(),
@@ -79,10 +79,10 @@ function UnitaryBangBangProblem(
     timestep_name::Symbol = :Δt,
     init_trajectory::Union{NamedTrajectory, Nothing}=nothing,
     a_bound::Float64=1.0,
-    a_bounds=fill(a_bound, length(system.G_drives)),
+    a_bounds=fill(a_bound, system.n_drives),
     a_guess::Union{Matrix{Float64}, Nothing}=nothing,
     da_bound::Float64=1.0,
-    da_bounds=fill(da_bound, length(system.G_drives)),
+    da_bounds=fill(da_bound, system.n_drives),
     Δt_min::Float64=Δt isa Float64 ? 0.5 * Δt : 0.5 * mean(Δt),
     Δt_max::Float64=Δt isa Float64 ? 1.5 * Δt : 1.5 * mean(Δt),
     Q::Float64=100.0,
@@ -99,12 +99,11 @@ function UnitaryBangBangProblem(
     if !isnothing(init_trajectory)
         traj = init_trajectory
     else
-        n_drives = length(system.G_drives)
         traj = initialize_trajectory(
             operator,
             T,
             Δt,
-            n_drives,
+            system.n_drives,
             (a_bounds, da_bounds);
             state_name=state_name,
             control_name=control_name,
@@ -147,7 +146,7 @@ function UnitaryBangBangProblem(
 
     # Constraints
     if R_bang_bang isa Float64
-        R_bang_bang = fill(R_bang_bang, length(system.G_drives))
+        R_bang_bang = fill(R_bang_bang, system.n_drives)
     end
     J += L1Regularizer!(
         constraints, control_names[2], traj,
@@ -157,10 +156,10 @@ function UnitaryBangBangProblem(
     # Integrators
     if piccolo_options.integrator == :pade
         unitary_integrator =
-            UnitaryPadeIntegrator(system, state_name, control_names[1], traj; order=piccolo_options.pade_order)
+            UnitaryPadeIntegrator(state_name, control_names[1], system, traj; order=piccolo_options.pade_order)
     elseif piccolo_options.integrator == :exponential
         unitary_integrator =
-            UnitaryExponentialIntegrator(system, state_name, control_names[1], traj)
+            UnitaryExponentialIntegrator(state_name, control_names[1], system, traj)
     else
         error("integrator must be one of (:pade, :exponential)")
     end
@@ -174,13 +173,13 @@ function UnitaryBangBangProblem(
     apply_piccolo_options!(J, constraints, piccolo_options, traj, operator, state_name, timestep_name)
 
     return QuantumControlProblem(
-        system,
         traj,
         J,
         integrators;
         constraints=constraints,
         ipopt_options=ipopt_options,
         piccolo_options=piccolo_options,
+        control_name=control_name,
         kwargs...
     )
 end
@@ -220,9 +219,9 @@ end
         piccolo_options=piccolo_options,
         control_name=:u
     )
-    initial = unitary_fidelity(prob; drive_name=:u)
+    initial = unitary_rollout_fidelity(prob.trajectory, sys; drive_name=:u)
     solve!(prob)
-    final = unitary_fidelity(prob; drive_name=:u)
+    final = unitary_rollout_fidelity(prob.trajectory, sys; drive_name=:u)
     @test final > initial
     solve!(smooth_prob)
     threshold = 1e-3
@@ -237,8 +236,10 @@ end
     phase_name = :ϕ
     phase_operators = [PAULIS[:Z]]
 
+    sys = QuantumSystem([PAULIS[:X]])
+
     prob = UnitaryBangBangProblem(
-        QuantumSystem([PAULIS[:X]]), GATES[:Y], 51, 0.2;
+        sys, GATES[:Y], 51, 0.2;
         phase_operators=phase_operators,
         phase_name=phase_name,
         ipopt_options=IpoptOptions(print_level=1),
@@ -251,11 +252,12 @@ end
 
     @test before ≠ after
 
-    @test unitary_fidelity(
-        prob, 
+    @test unitary_rollout_fidelity(
+        prob.trajectory,
+        sys;
         phases=prob.trajectory.global_data[phase_name],
         phase_operators=phase_operators
     ) > 0.9
 
-    @test unitary_fidelity(prob) < 0.9
+    @test unitary_rollout_fidelity(prob.trajectory, sys) < 0.9
 end
diff --git a/src/problem_templates/unitary_direct_sum_problem.jl b/src/problem_templates/unitary_direct_sum_problem.jl
index 3c0b7646..cb0a6520 100644
--- a/src/problem_templates/unitary_direct_sum_problem.jl
+++ b/src/problem_templates/unitary_direct_sum_problem.jl
@@ -48,6 +48,7 @@ function UnitaryDirectSumProblem(
     prob_labels::AbstractVector{<:String}=[string(i) for i ∈ 1:length(probs)],
     graph::Union{Nothing, AbstractVector{<:Tuple{String, String}}, AbstractVector{<:Tuple{Symbol, Symbol}}}=nothing,
     boundary_values::Union{AbstractDict{<:String, <:AbstractArray}, AbstractDict{<:Symbol, <:AbstractArray}}=Dict{String, Array}(),
+    control_name::Symbol=:a,
     Q::Union{Float64, Vector{Float64}}=100.0,
     Q_symb::Symbol=:dda,
     R::Float64=1e-2,
@@ -121,15 +122,12 @@ function UnitaryDirectSumProblem(
     end
 
     # Rebuild integrators
-    integrators = vcat(
-        [add_suffix(p.integrators, p.system, p.trajectory, traj, ℓ) for (p, ℓ) ∈ zip(probs, prob_labels)]...
-    )
-
-    # direct sum (used for problem saving, only)
-    system = direct_sum([add_suffix(p.system, ℓ) for (p, ℓ) ∈ zip(probs, prob_labels)])
+    integrators = vcat([
+        add_suffix(p.integrators, ℓ, p.trajectory, traj)
+            for (p, ℓ) ∈ zip(probs, prob_labels)
+    ]...)
 
     # Rebuild trajectory constraints
-    piccolo_options.build_trajectory_constraints = true
     constraints = AbstractConstraint[]
 
     # Add goal constraints for each problem
@@ -171,13 +169,14 @@ function UnitaryDirectSumProblem(
     end
 
     return QuantumControlProblem(
-        system,
         traj,
         J,
         integrators;
         constraints=constraints,
         ipopt_options=ipopt_options,
         piccolo_options=piccolo_options,
+        control_name=control_name,
+        kwargs...
     )
 end
 
diff --git a/src/problem_templates/unitary_minimum_time_problem.jl b/src/problem_templates/unitary_minimum_time_problem.jl
index 9bff21db..bcebdf22 100644
--- a/src/problem_templates/unitary_minimum_time_problem.jl
+++ b/src/problem_templates/unitary_minimum_time_problem.jl
@@ -52,6 +52,7 @@ function UnitaryMinimumTimeProblem(
     integrators::Vector{<:AbstractIntegrator},
     constraints::Vector{<:AbstractConstraint};
     unitary_name::Symbol=:Ũ⃗,
+    control_name::Symbol=:a,
     final_fidelity::Union{Real, Nothing}=nothing,
     D=1.0,
     ipopt_options::IpoptOptions=IpoptOptions(),
@@ -99,19 +100,20 @@ function UnitaryMinimumTimeProblem(
     constraints = push!(constraints, fidelity_constraint)
 
     return QuantumControlProblem(
-        system,
         trajectory,
         objective,
         integrators;
         constraints=constraints,
         ipopt_options=ipopt_options,
         piccolo_options=piccolo_options,
+        control_name=control_name,
         kwargs...
     )
 end
 
 function UnitaryMinimumTimeProblem(
-    prob::QuantumControlProblem;
+    prob::QuantumControlProblem,
+    sys::AbstractQuantumSystem;
     objective::Objective=get_objective(prob),
     constraints::AbstractVector{<:AbstractConstraint}=get_constraints(prob),
     ipopt_options::IpoptOptions=deepcopy(prob.ipopt_options),
@@ -123,7 +125,7 @@ function UnitaryMinimumTimeProblem(
 
     return UnitaryMinimumTimeProblem(
         copy(prob.trajectory),
-        prob.system,
+        sys,
         objective,
         prob.integrators,
         constraints;
@@ -144,30 +146,32 @@ end
     T = 51
     Δt = 0.2
 
+    sys = QuantumSystem(H_drift, H_drives)
+
     prob = UnitarySmoothPulseProblem(
-        H_drift, H_drives, U_goal, T, Δt,
+        sys, U_goal, T, Δt,
         ipopt_options=IpoptOptions(print_level=1),
         piccolo_options=PiccoloOptions(verbose=false)
     )
 
-    before = unitary_fidelity(prob)
+    before = unitary_rollout_fidelity(prob.trajectory, sys)
     solve!(prob, max_iter=50)
-    after = unitary_fidelity(prob)
+    after = unitary_rollout_fidelity(prob.trajectory, sys)
     @test after > before
 
     # Soft fidelity constraint
     final_fidelity = minimum([0.99, after])
-    mintime_prob = UnitaryMinimumTimeProblem(prob, final_fidelity=final_fidelity)
+    mintime_prob = UnitaryMinimumTimeProblem(prob, sys, final_fidelity=final_fidelity)
     solve!(mintime_prob; max_iter=100)
 
     # Test fidelity is approximatley staying above the constraint
-    @test unitary_fidelity(mintime_prob) ≥ (final_fidelity - 0.1 * final_fidelity)
+    @test unitary_rollout_fidelity(mintime_prob.trajectory, sys) ≥ (final_fidelity - 0.1 * final_fidelity)
     duration_after = sum(get_timesteps(mintime_prob.trajectory))
     duration_before = sum(get_timesteps(prob.trajectory))
     @test duration_after < duration_before
 
     # Set up without a final fidelity
-    @test UnitaryMinimumTimeProblem(prob) isa QuantumControlProblem
+    @test UnitaryMinimumTimeProblem(prob, sys) isa QuantumControlProblem
 
 end
 
@@ -175,18 +179,19 @@ end
     using NamedTrajectories
 
     phase_operators = [PAULIS[:Z]]
-
+    sys = QuantumSystem([PAULIS[:X]])
     prob = UnitarySmoothPulseProblem(
-        QuantumSystem([PAULIS[:X]]), GATES[:H], 51, 0.2,
+        sys, GATES[:H], 51, 0.2,
         ipopt_options=IpoptOptions(print_level=1),
         piccolo_options=PiccoloOptions(verbose=false),
         phase_operators=phase_operators
     )
 
     # Soft fidelity constraint
-    final_fidelity = minimum([0.99, unitary_fidelity(prob)])
+    final_fidelity = minimum([0.99, unitary_rollout_fidelity(prob.trajectory, sys)])
     mintime_prob = UnitaryMinimumTimeProblem(
-        prob, 
+        prob,
+        sys;
         final_fidelity=final_fidelity,
         phase_operators=phase_operators
     )
@@ -198,7 +203,8 @@ end
 
     # Quick check for using default fidelity
     @test UnitaryMinimumTimeProblem(
-        prob, 
+        prob,
+        sys;
         final_fidelity=final_fidelity,
         phase_operators=phase_operators
     ) isa QuantumControlProblem
diff --git a/src/problem_templates/unitary_robustness_problem.jl b/src/problem_templates/unitary_robustness_problem.jl
index eb68d94d..4baea9c1 100644
--- a/src/problem_templates/unitary_robustness_problem.jl
+++ b/src/problem_templates/unitary_robustness_problem.jl
@@ -26,13 +26,14 @@ Create a quantum control problem for robustness optimization of a unitary trajec
 function UnitaryRobustnessProblem end
 
 function UnitaryRobustnessProblem(
-    H_error::OperatorType,
+    H_error::AbstractPiccoloOperator,
     trajectory::NamedTrajectory,
     system::QuantumSystem,
     objective::Objective,
     integrators::Vector{<:AbstractIntegrator},
     constraints::Vector{<:AbstractConstraint};
     unitary_name::Symbol=:Ũ⃗,
+    control_name::Symbol=:a,
     final_fidelity::Union{Real, Nothing}=nothing,
     phase_name::Symbol=:ϕ,
     phase_operators::Union{AbstractVector{<:AbstractMatrix}, Nothing}=nothing,
@@ -56,7 +57,7 @@ function UnitaryRobustnessProblem(
         if isnothing(final_fidelity)
             final_fidelity = iso_vec_unitary_fidelity(U_T, U_G, subspace=subspace)
         end
-        
+
         fidelity_constraint = FinalUnitaryFidelityConstraint(
             unitary_name,
             final_fidelity,
@@ -86,20 +87,21 @@ function UnitaryRobustnessProblem(
     push!(constraints, fidelity_constraint)
 
     return QuantumControlProblem(
-        system,
         trajectory,
         objective,
         integrators;
         constraints=constraints,
         ipopt_options=ipopt_options,
         piccolo_options=piccolo_options,
+        control_name=control_name,
         kwargs...
     )
 end
 
 function UnitaryRobustnessProblem(
-    H_error::OperatorType,
-    prob::QuantumControlProblem;
+    H_error::AbstractPiccoloOperator,
+    prob::QuantumControlProblem,
+    sys::AbstractQuantumSystem;
     objective::Objective=get_objective(prob),
     constraints::AbstractVector{<:AbstractConstraint}=get_constraints(prob),
     ipopt_options::IpoptOptions=deepcopy(prob.ipopt_options),
@@ -112,7 +114,7 @@ function UnitaryRobustnessProblem(
     return UnitaryRobustnessProblem(
         H_error,
         copy(prob.trajectory),
-        prob.system,
+        sys,
         objective,
         prob.integrators,
         constraints;
@@ -142,9 +144,9 @@ end
         ipopt_options=IpoptOptions(print_level=1),
         piccolo_options=PiccoloOptions(verbose=false)
     )
-    before = unitary_fidelity(prob, subspace=U_goal.subspace_indices)
+    before = unitary_rollout_fidelity(prob.trajectory, sys; subspace=U_goal.subspace)
     solve!(prob, max_iter=15)
-    after = unitary_fidelity(prob, subspace=U_goal.subspace_indices)
+    after = unitary_rollout_fidelity(prob.trajectory, sys; subspace=U_goal.subspace)
 
     # Subspace gate success
     @test after > before
@@ -152,14 +154,14 @@ end
 
     # set up without a final fidelity
     # -------------------------------
-    @test UnitaryRobustnessProblem(H_embed, prob) isa QuantumControlProblem
+    @test UnitaryRobustnessProblem(H_embed, prob, sys) isa QuantumControlProblem
 
 
     #  test robustness from previous problem
     # --------------------------------------
     final_fidelity = 0.99
     rob_prob = UnitaryRobustnessProblem(
-        H_embed, prob,
+        H_embed, prob, sys;
         final_fidelity=final_fidelity,
         ipopt_options=IpoptOptions(recalc_y="yes", recalc_y_feas_tol=100.0, print_level=1),
     )
@@ -174,7 +176,49 @@ end
     @test (after < before) || (before < 0.25)
 
     # TODO: Fidelity constraint approximately satisfied
-    @test_skip isapprox(unitary_fidelity(rob_prob; subspace=U_goal.subspace_indices), 0.99, atol=0.05)
+    @test_skip isapprox(unitary_rollout_fidelity(rob_prob; subspace=U_goal.subspace), 0.99, atol=0.05)
+end
+
+@testitem "Set up a free phase problem" begin
+    using LinearAlgebra
+    δ1 = δ2 = -0.1
+    T = 75
+    Δt = 1.0
+    n_levels = 3
+    a = annihilate(n_levels)
+    id = I(n_levels)
+    a1 = kron(a, id)
+    a2 = kron(id, a)
+    H_drift = δ1 / 2 * a1' * a1' * a1 * a1 + δ2 / 2 * a2' * a2' * a2 * a2
+    H_drives = [a1'a1, a2'a2, a1'a2 + a1*a2', im * (a1'a2 - a1 * a2')]
+    system = QuantumSystem(H_drift, H_drives)
+    U_goal = EmbeddedOperator(
+        GATES[:CZ],
+        get_subspace_indices([1:2, 1:2], [n_levels, n_levels]),
+        [n_levels, n_levels]
+    )
+
+    phase_operators = [PAULIS[:Z], PAULIS[:Z]]
+    prob = UnitarySmoothPulseProblem(
+        system, U_goal, T, Δt,
+        ipopt_options=IpoptOptions(print_level=1),
+        piccolo_options=PiccoloOptions(verbose=false, eval_hessian=false),
+        phase_operators=phase_operators,
+    )
+
+    ZZ = EmbeddedOperator(
+        reduce(kron, phase_operators),
+        get_subspace_indices([1:2, 1:2], [n_levels, n_levels]),
+        [n_levels, n_levels]
+    )
+
+    @test UnitaryRobustnessProblem(
+        ZZ,
+        prob,
+        system,
+        phase_operators=phase_operators,
+        subspace=U_goal.subspace,
+    ) isa QuantumControlProblem
 end
 
 @testitem "Set up a free phase problem" begin
@@ -191,7 +235,7 @@ end
     H_drives = [a1'a1, a2'a2, a1'a2 + a1*a2', im * (a1'a2 - a1 * a2')]
     system = QuantumSystem(H_drift, H_drives)
     U_goal = EmbeddedOperator(
-        GATES[:CZ], 
+        GATES[:CZ],
         get_subspace_indices([1:2, 1:2], [n_levels, n_levels]),
         [n_levels, n_levels]
     )
@@ -205,15 +249,16 @@ end
     )
 
     ZZ = EmbeddedOperator(
-        reduce(kron, phase_operators), 
-        get_subspace_indices([1:2, 1:2], [n_levels, n_levels]), 
+        reduce(kron, phase_operators),
+        get_subspace_indices([1:2, 1:2], [n_levels, n_levels]),
         [n_levels, n_levels]
     )
 
     @test UnitaryRobustnessProblem(
         ZZ,
         prob,
+        system;
         phase_operators=phase_operators,
-        subspace=U_goal.subspace_indices,
+        subspace=U_goal.subspace,
     ) isa QuantumControlProblem
 end
diff --git a/src/problem_templates/unitary_sampling_problem.jl b/src/problem_templates/unitary_sampling_problem.jl
index 5d3ea084..0c533abe 100644
--- a/src/problem_templates/unitary_sampling_problem.jl
+++ b/src/problem_templates/unitary_sampling_problem.jl
@@ -10,7 +10,7 @@ robust solution by including multiple systems reflecting the problem uncertainty
 
 # Arguments
 - `systems::AbstractVector{<:AbstractQuantumSystem}`: A vector of quantum systems.
-- `operator::OperatorType`: The target unitary operator.
+- `operator::AbstractPiccoloOperator`: The target unitary operator.
 - `T::Int`: The number of time steps.
 - `Δt::Union{Float64, Vector{Float64}}`: The time step value or vector of time steps.
 
@@ -43,7 +43,7 @@ robust solution by including multiple systems reflecting the problem uncertainty
 """
 function UnitarySamplingProblem(
     systems::AbstractVector{<:AbstractQuantumSystem},
-    operators::AbstractVector{<:OperatorType},
+    operators::AbstractVector{<:AbstractPiccoloOperator},
     T::Int,
     Δt::Union{Float64,Vector{Float64}};
     system_weights=fill(1.0, length(systems)),
@@ -55,12 +55,12 @@ function UnitarySamplingProblem(
     timestep_name::Symbol=:Δt,
     constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
     a_bound::Float64=1.0,
-    a_bounds=fill(a_bound, length(systems[1].G_drives)),
+    a_bounds=fill(a_bound, systems[1].n_drives),
     a_guess::Union{Matrix{Float64},Nothing}=nothing,
     da_bound::Float64=Inf,
-    da_bounds::Vector{Float64}=fill(da_bound, length(systems[1].G_drives)),
+    da_bounds::Vector{Float64}=fill(da_bound, systems[1].n_drives),
     dda_bound::Float64=1.0,
-    dda_bounds::Vector{Float64}=fill(dda_bound, length(systems[1].G_drives)),
+    dda_bounds::Vector{Float64}=fill(dda_bound, systems[1].n_drives),
     Δt_min::Float64=0.5 * Δt,
     Δt_max::Float64=1.5 * Δt,
     Q::Float64=100.0,
@@ -85,7 +85,7 @@ function UnitarySamplingProblem(
                 op,
                 T,
                 Δt,
-                length(sys.G_drives),
+                sys.n_drives,
                 (a_bounds, da_bounds, dda_bounds);
                 state_name=s,
                 control_name=control_name,
@@ -126,16 +126,16 @@ function UnitarySamplingProblem(
 
     # Integrators
     unitary_integrators = AbstractIntegrator[]
-    for (sys, name) in zip(systems, state_names)
+    for (sys, Ũ⃗_name) in zip(systems, state_names)
         if piccolo_options.integrator == :pade
             push!(
                 unitary_integrators,
-                UnitaryPadeIntegrator(sys, name, control_name, traj; order=piccolo_options.pade_order)
+                UnitaryPadeIntegrator(Ũ⃗_name, control_name, sys, traj; order=piccolo_options.pade_order)
             )
         elseif piccolo_options.integrator == :exponential
             push!(
                 unitary_integrators,
-                UnitaryExponentialIntegrator(sys, name, control_name, traj)
+                UnitaryExponentialIntegrator(Ũ⃗_name, control_name, sys, traj)
             )
         else
             error("integrator must be one of (:pade, :exponential)")
@@ -154,20 +154,20 @@ function UnitarySamplingProblem(
     )
 
     return QuantumControlProblem(
-        direct_sum(systems),
         traj,
         J,
         integrators;
         constraints=constraints,
         ipopt_options=ipopt_options,
         piccolo_options=piccolo_options,
+        control_name=control_name,
         kwargs...
     )
 end
 
 function UnitarySamplingProblem(
     systems::AbstractVector{<:AbstractQuantumSystem},
-    operator::OperatorType,
+    operator::AbstractPiccoloOperator,
     T::Int,
     Δt::Union{Float64,Vector{Float64}};
     kwargs...
diff --git a/src/problem_templates/unitary_smooth_pulse_problem.jl b/src/problem_templates/unitary_smooth_pulse_problem.jl
index 1e6cea3b..2da15f96 100644
--- a/src/problem_templates/unitary_smooth_pulse_problem.jl
+++ b/src/problem_templates/unitary_smooth_pulse_problem.jl
@@ -36,7 +36,7 @@ or
 - `H_drift::AbstractMatrix{<:Number}`: the drift hamiltonian
 - `H_drives::Vector{<:AbstractMatrix{<:Number}}`: the control hamiltonians
 with
-- `operator::OperatorType`: the target unitary, either in the form of an `EmbeddedOperator` or a `Matrix{ComplexF64}
+- `operator::AbstractPiccoloOperator`: the target unitary, either in the form of an `EmbeddedOperator` or a `Matrix{ComplexF64}
 - `T::Int`: the number of timesteps
 - `Δt::Float64`: the (initial) time step size
 
@@ -52,6 +52,7 @@ with
 - `a_guess::Union{Matrix{Float64}, Nothing}=nothing`: an initial guess for the control pulses
 - `da_bound::Float64=Inf`: the bound on the control pulse derivative
 - `da_bounds::Vector{Float64}=fill(da_bound, length(system.G_drives))`: the bounds on the control pulse derivatives, one for each drive
+- `zero_initial_and_final_derivative::Bool=false`: whether to enforce zero initial and final control pulse derivatives
 - `dda_bound::Float64=1.0`: the bound on the control pulse second derivative
 - `dda_bounds::Vector{Float64}=fill(dda_bound, length(system.G_drives))`: the bounds on the control pulse second derivatives, one for each drive
 - `Δt_min::Float64=Δt isa Float64 ? 0.5 * Δt : 0.5 * mean(Δt)`: the minimum time step size
@@ -68,7 +69,7 @@ with
 """
 function UnitarySmoothPulseProblem(
     system::AbstractQuantumSystem,
-    operator::OperatorType,
+    operator::AbstractPiccoloOperator,
     T::Int,
     Δt::Union{Float64, Vector{Float64}};
     ipopt_options::IpoptOptions=IpoptOptions(),
@@ -78,12 +79,13 @@ function UnitarySmoothPulseProblem(
     timestep_name::Symbol = :Δt,
     init_trajectory::Union{NamedTrajectory, Nothing}=nothing,
     a_bound::Float64=1.0,
-    a_bounds=fill(a_bound, length(system.G_drives)),
+    a_bounds=fill(a_bound, system.n_drives),
     a_guess::Union{Matrix{Float64}, Nothing}=nothing,
     da_bound::Float64=Inf,
-    da_bounds::Vector{Float64}=fill(da_bound, length(system.G_drives)),
+    da_bounds::Vector{Float64}=fill(da_bound, system.n_drives),
+    zero_initial_and_final_derivative::Bool=false,
     dda_bound::Float64=1.0,
-    dda_bounds::Vector{Float64}=fill(dda_bound, length(system.G_drives)),
+    dda_bounds::Vector{Float64}=fill(dda_bound, system.n_drives),
     Δt_min::Float64=Δt isa Float64 ? 0.5 * Δt : 0.5 * mean(Δt),
     Δt_max::Float64=Δt isa Float64 ? 1.5 * Δt : 1.5 * mean(Δt),
     Q::Float64=100.0,
@@ -101,19 +103,18 @@ function UnitarySmoothPulseProblem(
     if !isnothing(init_trajectory)
         traj = init_trajectory
     else
-        n_drives = length(system.G_drives)
-
         traj = initialize_trajectory(
             operator,
             T,
             Δt,
-            n_drives,
+            system.n_drives,
             (a_bounds, da_bounds, dda_bounds);
             state_name=state_name,
             control_name=control_name,
             timestep_name=timestep_name,
             free_time=piccolo_options.free_time,
             Δt_bounds=(Δt_min, Δt_max),
+            zero_initial_and_final_derivative=zero_initial_and_final_derivative,
             geodesic=piccolo_options.geodesic,
             bound_state=piccolo_options.bound_state,
             a_guess=a_guess,
@@ -125,12 +126,12 @@ function UnitarySmoothPulseProblem(
     end
 
     # Subspace
-    subspace = operator isa EmbeddedOperator ? operator.subspace_indices : nothing
+    subspace = operator isa EmbeddedOperator ? operator.subspace : nothing
 
     # Objective
     if isnothing(phase_operators)
         J = UnitaryInfidelityObjective(
-            state_name, traj, Q; 
+            state_name, traj, Q;
             subspace=subspace,
             eval_hessian=piccolo_options.eval_hessian,
         )
@@ -152,15 +153,15 @@ function UnitarySmoothPulseProblem(
     J += QuadraticRegularizer(control_names[3], traj, R_dda; timestep_name=timestep_name)
 
     # Integrators
+
     if piccolo_options.integrator == :pade
         unitary_integrator =
-            UnitaryPadeIntegrator(
-                system, state_name, control_name, traj;
+            UnitaryPadeIntegrator(state_name, control_name, system, traj;
                 order=piccolo_options.pade_order
             )
     elseif piccolo_options.integrator == :exponential
         unitary_integrator =
-            UnitaryExponentialIntegrator(system, state_name, control_name, traj)
+            UnitaryExponentialIntegrator(state_name, control_name, system, traj)
     else
         error("integrator must be one of (:pade, :exponential)")
     end
@@ -177,13 +178,13 @@ function UnitarySmoothPulseProblem(
     )
 
     return QuantumControlProblem(
-        system,
         traj,
         J,
         integrators;
         constraints=constraints,
         ipopt_options=ipopt_options,
         piccolo_options=piccolo_options,
+        control_name=control_name,
         kwargs...
     )
 end
@@ -213,9 +214,9 @@ end
         piccolo_options=PiccoloOptions(verbose=false)
     )
 
-    initial = unitary_fidelity(prob)
+    initial = unitary_rollout_fidelity(prob.trajectory, sys)
     solve!(prob, max_iter=20)
-    final = unitary_fidelity(prob)
+    final = unitary_rollout_fidelity(prob.trajectory, sys)
     @test final > initial
 end
 
@@ -231,9 +232,9 @@ end
         piccolo_options=PiccoloOptions(verbose=false, integrator=:exponential, jacobian_structure=false)
     )
 
-    initial = unitary_fidelity(prob)
+    initial = unitary_rollout_fidelity(prob.trajectory, sys)
     solve!(prob, max_iter=20)
-    final = unitary_fidelity(prob)
+    final = unitary_rollout_fidelity(prob.trajectory, sys)
     @test final > initial
 end
 
@@ -257,9 +258,9 @@ end
         piccolo_options=piccolo_options
     )
 
-    initial = unitary_fidelity(prob)
+    initial = unitary_rollout_fidelity(prob.trajectory, sys)
     solve!(prob, max_iter=20)
-    final = unitary_fidelity(prob)
+    final = unitary_rollout_fidelity(prob.trajectory, sys)
     @test final > initial
 end
 
@@ -280,9 +281,9 @@ end
         piccolo_options=PiccoloOptions(verbose=false)
     )
 
-    initial = unitary_fidelity(prob, subspace=U_goal.subspace_indices)
+    initial = unitary_rollout_fidelity(prob.trajectory, sys, subspace=U_goal.subspace)
     solve!(prob, max_iter=20)
-    final = unitary_fidelity(prob, subspace=U_goal.subspace_indices)
+    final = unitary_rollout_fidelity(prob.trajectory, sys, subspace=U_goal.subspace)
     @test final > initial
 
     # Test leakage suppression
@@ -300,20 +301,54 @@ end
         piccolo_options=PiccoloOptions(verbose=false)
     )
 
-    initial = unitary_fidelity(prob, subspace=U_goal.subspace_indices)
+    initial = unitary_rollout_fidelity(prob.trajectory, sys, subspace=U_goal.subspace)
     solve!(prob, max_iter=20)
-    final = unitary_fidelity(prob, subspace=U_goal.subspace_indices)
+    final = unitary_rollout_fidelity(prob.trajectory, sys, subspace=U_goal.subspace)
     @test final > initial
 end
 
+@testitem "Additional Objective" begin
+    H_drift = GATES[:Z]
+    H_drives = [GATES[:X], GATES[:Y]]
+    U_goal = GATES[:H]
+    T = 50
+    Δt = 0.2
+
+    prob_vanilla = UnitarySmoothPulseProblem(
+        H_drift, H_drives, U_goal, T, Δt,
+        ipopt_options=IpoptOptions(print_level=1),
+        piccolo_options=PiccoloOptions(verbose=false),
+    )
+
+    J_extra = QuadraticSmoothnessRegularizer(:dda, prob_vanilla.trajectory, 10.0)
+
+    prob_additional = UnitarySmoothPulseProblem(
+        H_drift, H_drives, U_goal, T, Δt,
+        ipopt_options=IpoptOptions(print_level=1),
+        piccolo_options=PiccoloOptions(verbose=false),
+        additional_objective=J_extra,
+    )
+
+    J_prob_vanilla = Problems.get_objective(prob_vanilla)
+
+    J_additional = Problems.get_objective(prob_additional)
+
+    Z = prob_vanilla.trajectory
+    Z⃗ = vec(prob_vanilla.trajectory)
+
+    @test J_prob_vanilla.L(Z⃗, Z) + J_extra.L(Z⃗, Z) ≈ J_additional.L(Z⃗, Z)
+end
 @testitem "Free phase Y gate using X" begin
     using Random
-    Random.seed!(1234)
+    # Random.seed!(1234)
     phase_name = :ϕ
     phase_operators = [PAULIS[:Z]]
-
+    sys = QuantumSystem([PAULIS[:X]])
     prob = UnitarySmoothPulseProblem(
-        QuantumSystem([PAULIS[:X]]), GATES[:Y], 51, 0.2;
+        sys,
+        GATES[:Y],
+        51,
+        0.2;
         phase_operators=phase_operators,
         phase_name=phase_name,
         ipopt_options=IpoptOptions(print_level=1),
@@ -321,16 +356,17 @@ end
     )
 
     before = prob.trajectory.global_data[phase_name]
-    solve!(prob, max_iter=20)
+    solve!(prob, max_iter=50)
     after = prob.trajectory.global_data[phase_name]
 
     @test before ≠ after
 
-    @test unitary_fidelity(
-        prob, 
+    @test unitary_rollout_fidelity(
+        prob.trajectory,
+        sys;
         phases=prob.trajectory.global_data[phase_name],
         phase_operators=phase_operators
     ) > 0.9
 
-    @test unitary_fidelity(prob) < 0.9
+    @test unitary_rollout_fidelity(prob.trajectory, sys) < 0.9
 end
diff --git a/src/problems.jl b/src/problems.jl
deleted file mode 100644
index f762144e..00000000
--- a/src/problems.jl
+++ /dev/null
@@ -1,435 +0,0 @@
-module Problems
-
-export AbstractProblem
-export FixedTimeProblem
-export QuantumControlProblem
-
-export set_trajectory!
-export update_trajectory!
-export get_objective
-export get_constraints
-
-using ..QuantumSystems
-using ..Integrators
-using ..Evaluators
-using ..Options
-using ..Constraints
-using ..Dynamics
-using ..Objectives
-
-using TrajectoryIndexingUtils
-using NamedTrajectories
-
-using LinearAlgebra
-using JLD2
-using Ipopt
-using TestItemRunner
-using MathOptInterface
-using LinearAlgebra
-const MOI = MathOptInterface
-
-abstract type AbstractProblem end
-
-"""
-    mutable struct QuantumControlProblem <: AbstractProblem
-
-Stores all the information needed to set up and solve a QuantumControlProblem as well as the solution
-after the solver terminates.
-
-# Fields
-- `optimizer::Ipopt.Optimizer`: Ipopt optimizer object
-"""
-mutable struct QuantumControlProblem <: AbstractProblem
-    optimizer::Ipopt.Optimizer
-    variables::Vector{MOI.VariableIndex}
-    system::AbstractQuantumSystem
-    trajectory::NamedTrajectory
-    integrators::Union{Nothing,Vector{<:AbstractIntegrator}}
-    ipopt_options::IpoptOptions
-    piccolo_options::PiccoloOptions
-    params::Dict{Symbol, Any}
-end
-
-function QuantumControlProblem(
-    system::AbstractQuantumSystem,
-    traj::NamedTrajectory,
-    obj::Objective,
-    dynamics::QuantumDynamics;
-    ipopt_options::IpoptOptions=IpoptOptions(),
-    piccolo_options::PiccoloOptions=PiccoloOptions(),
-    scale_factor_objective::Float64=1.0,
-    additional_objective::Union{Nothing, Objective}=nothing,
-    constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
-    params::Dict{Symbol, Any}=Dict{Symbol, Any}(),
-    return_evaluator=false,
-    kwargs...
-)
-    # Save internal copy of the options to allow modification
-    ipopt_options = deepcopy(ipopt_options)
-    piccolo_options = deepcopy(piccolo_options)
-
-    if !piccolo_options.blas_multithreading
-        BLAS.set_num_threads(1)
-    end
-
-    if !piccolo_options.eval_hessian
-        ipopt_options.hessian_approximation = "limited-memory"
-    end
-
-    nonlinear_constraints = NonlinearConstraint[con for con ∈ constraints if con isa NonlinearConstraint]
-
-    if scale_factor_objective != 1
-        obj = scale_factor_objective * obj
-    end
-
-    if !isnothing(additional_objective)
-        obj += additional_objective
-    end
-
-    if piccolo_options.verbose
-        println("    building evaluator...")
-    end
-
-    evaluator = PicoEvaluator(
-        traj,
-        obj,
-        dynamics,
-        nonlinear_constraints;
-        eval_hessian=piccolo_options.eval_hessian,
-    )
-
-    if return_evaluator
-        return evaluator
-    end
-
-    n_dynamics_constraints = dynamics.dim * (traj.T - 1)
-    n_variables = traj.dim * traj.T
-
-    # add globabl variables to n_variables
-    for global_var ∈ keys(traj.global_data)
-        global_var_dim = length(traj.global_data[global_var])
-        n_variables += global_var_dim
-    end
-
-    linear_constraints = LinearConstraint[con for con ∈ constraints if con isa LinearConstraint]
-
-    if piccolo_options.build_trajectory_constraints
-        linear_constraints = LinearConstraint[trajectory_constraints(traj); linear_constraints]
-    end
-
-    optimizer = Ipopt.Optimizer()
-
-    if piccolo_options.verbose
-        println("    initializing optimizer...")
-    end
-
-    variables = initialize_optimizer!(
-        optimizer,
-        evaluator,
-        traj,
-        linear_constraints,
-        n_dynamics_constraints,
-        nonlinear_constraints,
-        n_variables,
-        verbose=piccolo_options.verbose
-    )
-
-    # Container for saving constraints and objectives
-    params = merge(kwargs, params)
-    params[:linear_constraints] = linear_constraints
-    params[:nonlinear_constraints] = [
-        nl_constraint.params for nl_constraint ∈ nonlinear_constraints
-    ]
-    params[:objective_terms] = obj.terms
-
-    return QuantumControlProblem(
-        optimizer,
-        variables,
-        system,
-        traj,
-        dynamics.integrators,
-        ipopt_options,
-        piccolo_options,
-        params
-    )
-end
-
-function QuantumControlProblem(
-    system::AbstractQuantumSystem,
-    traj::NamedTrajectory,
-    obj::Objective,
-    integrators::Vector{<:AbstractIntegrator};
-    ipopt_options::IpoptOptions=IpoptOptions(),
-    piccolo_options::PiccoloOptions=PiccoloOptions(),
-    kwargs...
-)
-    # Save internal copy of the options to allow modification
-    ipopt_options = deepcopy(ipopt_options)
-    piccolo_options = deepcopy(piccolo_options)
-
-    if piccolo_options.verbose
-        println("    building dynamics from integrators...")
-    end
-    dynamics = QuantumDynamics(integrators, traj;
-        jacobian_structure=piccolo_options.jacobian_structure,
-        eval_hessian=piccolo_options.eval_hessian,
-        verbose=piccolo_options.verbose
-    )
-    return QuantumControlProblem(
-        system,
-        traj,
-        obj,
-        dynamics;
-        ipopt_options=ipopt_options,
-        piccolo_options=piccolo_options,
-        kwargs...
-    )
-end
-
-# constructor that accepts just an AbstractIntegrator
-function QuantumControlProblem(
-    system::AbstractQuantumSystem,
-    traj::NamedTrajectory,
-    obj::Objective,
-    integrator::AbstractIntegrator;
-    ipopt_options::IpoptOptions=IpoptOptions(),
-    piccolo_options::PiccoloOptions=PiccoloOptions(),
-    kwargs...
-)
-    # Save internal copy of the options to allow modification
-    ipopt_options = deepcopy(ipopt_options)
-    piccolo_options = deepcopy(piccolo_options)
-
-    if piccolo_options.verbose
-        println("    building dynamics from integrator...")
-    end
-    dynamics = QuantumDynamics(integrator, traj;
-        jacobian_structure=piccolo_options.jacobian_structure,
-        eval_hessian=piccolo_options.eval_hessian,
-        verbose=piccolo_options.verbose
-    )
-    return QuantumControlProblem(
-        system,
-        traj,
-        obj,
-        dynamics;
-        ipopt_options=ipopt_options,
-        piccolo_options=piccolo_options,
-        kwargs...
-    )
-end
-
-function QuantumControlProblem(
-    system::AbstractQuantumSystem,
-    traj::NamedTrajectory,
-    obj::Objective,
-    f::Function;
-    ipopt_options::IpoptOptions=IpoptOptions(),
-    piccolo_options::PiccoloOptions=PiccoloOptions(),
-    kwargs...
-)
-    # Save internal copy of the options to allow modification
-    ipopt_options = deepcopy(ipopt_options)
-    piccolo_options = deepcopy(piccolo_options)
-
-    if piccolo_options.verbose
-        println("    building dynamics from function...")
-    end
-    dynamics = QuantumDynamics(f, traj;
-        jacobian_structure=piccolo_options.jacobian_structure,
-        eval_hessian=piccolo_options.eval_hessian,
-        verbose=piccolo_options.verbose
-    )
-    return QuantumControlProblem(
-        system,
-        traj,
-        obj,
-        dynamics;
-        ipopt_options=ipopt_options,
-        piccolo_options=piccolo_options,
-        kwargs...
-    )
-end
-
-function initialize_optimizer!(
-    optimizer::Ipopt.Optimizer,
-    evaluator::PicoEvaluator,
-    trajectory::NamedTrajectory,
-    linear_constraints::Vector{LinearConstraint},
-    n_dynamics_constraints::Int,
-    nonlinear_constraints::Vector{NonlinearConstraint},
-    n_variables::Int;
-    verbose=true
-)
-    nl_cons = fill(
-        MOI.NLPBoundsPair(0.0, 0.0),
-        n_dynamics_constraints
-    )
-
-    for nl_con ∈ nonlinear_constraints
-        if nl_con isa NonlinearEqualityConstraint
-            append!(nl_cons, fill(MOI.NLPBoundsPair(0.0, 0.0), nl_con.dim))
-        elseif nl_con isa NonlinearInequalityConstraint
-            append!(nl_cons, fill(MOI.NLPBoundsPair(0.0, Inf), nl_con.dim))
-        else
-            error("Unknown nonlinear constraint type")
-        end
-    end
-
-    # build NLP block data
-    block_data = MOI.NLPBlockData(nl_cons, evaluator, true)
-
-    # set NLP block data
-    MOI.set(optimizer, MOI.NLPBlock(), block_data)
-
-    # set objective sense: minimize
-    MOI.set(optimizer, MOI.ObjectiveSense(), MOI.MIN_SENSE)
-
-    # add variables
-    variables = MOI.add_variables(optimizer, n_variables)
-
-    # add linear constraints
-    constrain!(optimizer, variables, linear_constraints, trajectory, verbose=verbose)
-
-    return variables
-end
-
-function set_trajectory!(
-    prob::QuantumControlProblem,
-    traj::NamedTrajectory
-)
-    # initialize n variables with trajectory data
-    n_vars = traj.dim * traj.T
-
-    # set trajectory data
-    MOI.set(
-        prob.optimizer,
-        MOI.VariablePrimalStart(),
-        prob.variables[1:n_vars],
-        collect(traj.datavec)
-    )
-
-    # set global variables
-    for global_vars_i ∈ values(traj.global_data)
-        n_global_vars = length(global_vars_i)
-        MOI.set(
-            prob.optimizer,
-            MOI.VariablePrimalStart(),
-            prob.variables[n_vars .+ (1:n_global_vars)],
-            global_vars_i
-        )
-        n_vars += n_global_vars
-    end
-end
-
-set_trajectory!(prob::QuantumControlProblem) =
-    set_trajectory!(prob, prob.trajectory)
-
-function get_datavec(prob::QuantumControlProblem)
-    n_vars = prob.trajectory.dim * prob.trajectory.T
-
-    # get trajectory data
-    return MOI.get(
-        prob.optimizer,
-        MOI.VariablePrimal(),
-        prob.variables[1:n_vars]
-    )
-end
-
-function get_global_data(prob::QuantumControlProblem)
-
-    # get global variables after trajectory data
-    global_keys = keys(prob.trajectory.global_data)
-    global_values = []
-    n_vars = prob.trajectory.dim * prob.trajectory.T
-    for global_var ∈ global_keys
-        n_global_vars = length(prob.trajectory.global_data[global_var])
-        push!(global_values, MOI.get(
-            prob.optimizer,
-            MOI.VariablePrimal(),
-            prob.variables[n_vars .+ (1:n_global_vars)]
-        ))
-        n_vars += n_global_vars
-    end
-    return (; (global_keys .=> global_values)...)
-end
-
-@views function update_trajectory!(prob::QuantumControlProblem)
-    datavec = get_datavec(prob)
-    global_data = get_global_data(prob)
-    prob.trajectory = NamedTrajectory(datavec, global_data, prob.trajectory)
-    return nothing
-end
-
-"""
-    get_objective(prob::QuantumControlProblem)
-
-Return the objective function of the `prob::QuantumControlProblem`.
-"""
-function get_objective(
-    prob::QuantumControlProblem;
-    match::Union{Nothing, AbstractVector{<:Symbol}}=nothing,
-    invert::Bool=false
-)
-    objs = deepcopy(prob.params[:objective_terms])
-    if isnothing(match)
-        return Objective(objs)
-    else
-        if invert
-            return Objective([term for term ∈ objs if term[:type] ∉ match])
-        else
-            return Objective([term for term ∈ objs if term[:type] ∈ match])
-        end
-    end
-end
-
-"""
-    get_constraints(prob::QuantumControlProblem)
-
-Return the constraints of the `prob::QuantumControlProblem`.
-"""
-function get_constraints(prob::QuantumControlProblem)
-    linear_constraints = deepcopy(prob.params[:linear_constraints])
-    nonlinear_constraints = deepcopy(prob.params[:nonlinear_constraints])
-    return AbstractConstraint[
-        linear_constraints...,
-        NonlinearConstraint.(nonlinear_constraints)...
-    ]
-end
-
-# ============================================================================= #
-
-@testitem "Additional Objective" begin
-    H_drift = GATES[:Z]
-    H_drives = [GATES[:X], GATES[:Y]]
-    U_goal = GATES[:H]
-    T = 50
-    Δt = 0.2
-
-    prob_vanilla = UnitarySmoothPulseProblem(
-        H_drift, H_drives, U_goal, T, Δt,
-        ipopt_options=IpoptOptions(print_level=1),
-        piccolo_options=PiccoloOptions(verbose=false),
-    )
-
-    J_extra = QuadraticSmoothnessRegularizer(:dda, prob_vanilla.trajectory, 10.0)
-
-    prob_additional = UnitarySmoothPulseProblem(
-        H_drift, H_drives, U_goal, T, Δt,
-        ipopt_options=IpoptOptions(print_level=1),
-        piccolo_options=PiccoloOptions(verbose=false),
-        additional_objective=J_extra,
-    )
-
-    J_prob_vanilla = Problems.get_objective(prob_vanilla)
-
-    J_additional = Problems.get_objective(prob_additional)
-
-    Z = prob_vanilla.trajectory
-    Z⃗ = vec(prob_vanilla.trajectory)
-
-    @test J_prob_vanilla.L(Z⃗, Z) + J_extra.L(Z⃗, Z) ≈ J_additional.L(Z⃗, Z)
-
-end
-
-end
diff --git a/src/quantum_object_utils.jl b/src/quantum_object_utils.jl
deleted file mode 100644
index 3fe48a5e..00000000
--- a/src/quantum_object_utils.jl
+++ /dev/null
@@ -1,284 +0,0 @@
-module QuantumObjectUtils
-
-export PAULIS
-export GATES
-export ⊗
-export operator_from_string
-export ket_from_string
-export ket_from_bitstring
-
-export haar_random
-export haar_identity
-
-export create
-export annihilate
-
-using LinearAlgebra
-using TestItemRunner
-
-# TODO:
-# [ ] Remove need for oscillator operators (used by tests)
-# [ ] Allow multi-character symbols for operators_from_string
-# [ ] Remove need for otimes symbol or avoid import conflicts with other packages
-
-
-@doc raw"""
-A constant dictionary `GATES` containing common quantum gate matrices as complex-valued matrices. Each gate is represented by its unitary matrix.
-
-- `GATES[:I]` - Identity gate: Leaves the state unchanged.
-- `GATES[:X]` - Pauli-X (NOT) gate: Flips the qubit state.
-- `GATES[:Y]` - Pauli-Y gate: Rotates the qubit state around the Y-axis of the Bloch sphere.
-- `GATES[:Z]` - Pauli-Z gate: Flips the phase of the qubit state.
-- `GATES[:H]` - Hadamard gate: Creates superposition by transforming basis states.
-- `GATES[:CX]` - Controlled-X (CNOT) gate: Flips the second qubit (target) if the first qubit (control) is |1⟩.
-- `GATES[:CZ]` - Controlled-Z (CZ) gate: Flips the phase of the second qubit (target) if the first qubit (control) is |1⟩.
-- `GATES[:XI]` - Complex gate: A specific gate used for complex operations.
-- `GATES[:sqrtiSWAP]` - Square root of iSWAP gate: Partially swaps two qubits with a phase.
-
-```julia
-julia> GATES[:Z]
-2×2 Matrix{ComplexF64}:
- 1.0+0.0im   0.0+0.0im
- 0.0+0.0im  -1.0+0.0im
-
-julia> get_gate(:CX)
-4×4 Matrix{ComplexF64}:
- 1.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im
- 0.0+0.0im  1.0+0.0im  0.0+0.0im  0.0+0.0im
- 0.0+0.0im  0.0+0.0im  0.0+0.0im  1.0+0.0im
- 0.0+0.0im  0.0+0.0im  1.0+0.0im  0.0+0.0im
-```
-"""
-const GATES = Dict{Symbol, Matrix{ComplexF64}}(
-    :I => [1 0;
-           0 1],
-
-    :X => [0 1;
-           1 0],
-
-    :Y => [0 -im;
-           im 0],
-
-    :Z => [1 0;
-           0 -1],
-
-    :H => [1 1;
-           1 -1]/√2,
-
-    :CX => [1 0 0 0;
-            0 1 0 0;
-            0 0 0 1;
-            0 0 1 0],
-
-    :CZ => [1 0 0 0;
-            0 1 0 0;
-            0 0 1 0;
-            0 0 0 -1],
-
-    :XI => [0 0 -im 0;
-            0 0 0 -im;
-            -im 0 0 0;
-            0 -im 0 0],
-
-    :sqrtiSWAP => [1 0 0 0;
-                   0 1/sqrt(2) 1im/sqrt(2) 0;
-                   0 1im/sqrt(2) 1/sqrt(2) 0;
-                   0 0 0 1]
-)
-
-const PAULIS = Dict{Symbol, Matrix{ComplexF64}}(
-    :I => GATES[:I],
-    :X => GATES[:X],
-    :Y => GATES[:Y],
-    :Z => GATES[:Z]
-)
-
-###
-### Kronecker product utilities
-###
-
-@doc raw"""
-    ⊗(A::AbstractVecOrMat, B::AbstractVecOrMat) = kron(A, B)
-
-The Kronecker product, denoted by `⊗`, results in a block matrix formed by multiplying each element of `A` by the entire matrix `B`.
-
-```julia
-julia> GATES[:X] ⊗ GATES[:Y]
-4×4 Matrix{ComplexF64}:
- 0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0-1.0im
- 0.0+0.0im  0.0+0.0im  0.0+1.0im  0.0+0.0im
- 0.0+0.0im  0.0-1.0im  0.0+0.0im  0.0+0.0im
- 0.0+1.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im
-```
-"""
-⊗(A::AbstractVecOrMat, B::AbstractVecOrMat) = kron(A, B)
-
-@doc raw"""
-operator_from_string(operator::String; lookup::Dict{Symbol, AbstractMatrix}=PAULIS)
-
-    Reduce the string (each character is one key) via operators from a dictionary.
-
-"""
-function operator_from_string(
-    operator::String;
-    lookup::Dict{Symbol, <:AbstractMatrix}=PAULIS
-)::Matrix{ComplexF64}
-    # TODO: allow multi-character keys, ['(', ')']
-
-    # split string into keys and replace with operators
-    characters = [Symbol(c) for c ∈ operator]
-    operators = replace(characters, lookup...)
-
-    return foldr(kron, operators)
-end
-
-function cavity_state(state::Int, levels::Int)::Vector{ComplexF64}
-    @assert state ≤ levels - 1 "Level $state is not allowed for $levels levels"
-    ket = zeros(levels)
-    ket[state + 1] = 1
-    return ket
-end
-
-@doc raw"""
-    ket_from_string(
-        ket::String,
-        levels::Vector{Int};
-        level_dict=Dict(:g => 0, :e => 1, :f => 2, :h => 2),
-        return_states=false
-    )
-
-Construct a quantum state from a string ket representation.
-
-# Example
-
-# TODO: add example
-"""
-function ket_from_string(
-    ket::String,
-    levels::Vector{Int};
-    level_dict=Dict(:g => 0, :e => 1, :f => 2, :h => 2),
-    return_states=false
-)::Vector{ComplexF64}
-    kets = []
-
-    for x ∈ split(ket, ['(', ')'])
-        if x == ""
-            continue
-        elseif all(Symbol(xᵢ) ∈ keys(level_dict) for xᵢ ∈ x)
-            append!(kets, x)
-        elseif occursin("+", x)
-            superposition = split(x, '+')
-            @assert all(all(Symbol(xᵢ) ∈ keys(level_dict) for xᵢ ∈ x) for x ∈ superposition) "Invalid ket: $x"
-            @assert length(superposition) == 2 "Only two states can be superposed for now"
-            push!(kets, x)
-        else
-            error("Invalid ket: $x")
-        end
-    end
-
-    states = []
-
-    for (ψᵢ, l) ∈ zip(kets, levels)
-        if ψᵢ isa AbstractString && occursin("+", ψᵢ)
-            superposition = split(ψᵢ, '+')
-            superposition_states = [level_dict[Symbol(x)] for x ∈ superposition]
-            @assert all(state ≤ l - 1 for state ∈ superposition_states) "Level $ψᵢ is not allowed for $l levels"
-            superposition_state = sum([
-                cavity_state(state, l) for state ∈ superposition_states
-            ])
-            normalize!(superposition_state)
-            push!(states, superposition_state)
-        else
-            state = level_dict[Symbol(ψᵢ)]
-            @assert state ≤ l - 1 "Level $ψᵢ is not allowed for $l levels"
-            push!(states, cavity_state(state, l))
-        end
-    end
-
-    if return_states
-        return states
-    else
-        return kron([1.0], states...)
-    end
-end
-
-@doc raw"""
-    ket_from_bitstring(ket::String)
-
-Get the state vector for a qubit system given a ket string `ket` of 0s and 1s.
-"""
-function ket_from_bitstring(ket::String)::Vector{ComplexF64}
-    cs = [c for c ∈ ket]
-    @assert all(c ∈ "01" for c ∈ cs)
-    states = [c == '0' ? [1, 0] : [0, 1] for c ∈ cs]
-    return foldr(kron, states)
-end
-
-###
-### Random operators
-###
-
-@doc raw"""
-    haar_random(n::Int)
-
-Generate a random unitary matrix using the Haar measure for an `n`-dimensional system.
-"""
-function haar_random(n::Int)
-    # Ginibre matrix
-    Z = (randn(n, n) + im * randn(n, n)) / √2
-    F = qr(Z)
-    # QR correction (R main diagonal is real, strictly positive)
-    Λ = diagm(diag(F.R) ./ abs.(diag(F.R)))
-    return F.Q * Λ
-end
-
-@doc raw"""
-    haar_identity(n::Int, radius::Number)
-
-Generate a random unitary matrix close to the identity matrix using the Haar measure for an `n`-dimensional system with a given `radius`.
-"""
-function haar_identity(n::Int, radius::Number)
-    # Ginibre matrix
-    Z = (I + radius * (randn(n, n) + im * randn(n, n)) / √2) / (1 + radius)
-    F = qr(Z)
-    # QR correction (R main diagonal is real, strictly positive)
-    Λ = diagm(diag(F.R) ./ abs.(diag(F.R)))
-    return F.Q * Λ
-end
-
-###
-### Oscillator operators
-###
-
-@doc raw"""
-    annihilate(levels::Int)
-
-Get the annihilation operator for a system with `levels` levels.
-"""
-function annihilate(levels::Int)::Matrix{ComplexF64}
-    return diagm(1 => map(sqrt, 1:levels - 1))
-end
-
-@doc raw"""
-    create(levels::Int)
-
-Get the creation operator for a system with `levels` levels.
-"""
-function create(levels::Int)
-    return collect(annihilate(levels)')
-end
-
-# *************************************************************************** #
-
-@testitem "Test ket_from_bitstring function" begin
-    using LinearAlgebra
-    @test ket_from_bitstring("0") == [1, 0]
-    @test ket_from_bitstring("1") == [0, 1]
-    @test ket_from_bitstring("00") == [1, 0, 0, 0]
-    @test ket_from_bitstring("01") == [0, 1, 0, 0]
-    @test ket_from_bitstring("10") == [0, 0, 1, 0]
-    @test ket_from_bitstring("11") == [0, 0, 0, 1]
-end
-
-
-end
diff --git a/src/quantum_system_templates/_quantum_system_templates.jl b/src/quantum_system_templates/_quantum_system_templates.jl
index 2dbac44d..5621166a 100644
--- a/src/quantum_system_templates/_quantum_system_templates.jl
+++ b/src/quantum_system_templates/_quantum_system_templates.jl
@@ -1,19 +1,11 @@
 module QuantumSystemTemplates
 
-export TransmonSystem
-export TransmonDipoleCoupling
-export MultiTransmonSystem
-export RydbergChainSystem
-export QuantumOpticsSystem
-
-using ..Isomorphisms
-using ..QuantumSystems
-using ..QuantumObjectUtils
-
+using PiccoloQuantumObjects
 using LinearAlgebra
 using TestItemRunner
 
 include("transmons.jl")
 include("rydberg.jl")
+include("cats.jl")
 
 end
diff --git a/src/quantum_system_templates/cats.jl b/src/quantum_system_templates/cats.jl
new file mode 100644
index 00000000..57c4786c
--- /dev/null
+++ b/src/quantum_system_templates/cats.jl
@@ -0,0 +1,61 @@
+export CatSystem
+export coherent_ket
+export get_cat_controls
+
+function coherent_ket(α::Union{Real, Complex}, levels::Int)
+    return [exp(-0.5 * abs2(α)) * α^n / sqrt(factorial(n)) for n in 0:levels-1]
+end
+
+function CatSystem(;
+    g2::Real=0.36,
+    χ_aa::Real=-7e-3,
+    χ_bb::Real=-32,
+    χ_ab::Real=0.79,
+    κa::Real=53e-3,
+    κb::Real=13,
+    cat_levels::Int=13,
+    buffer_levels::Int=3,
+    prefactor::Real=1,
+)
+    params = Dict(
+        :g2 => prefactor *  g2,
+        :χ_aa => prefactor *  χ_aa,
+        :χ_bb => prefactor *  χ_bb,
+        :χ_ab => prefactor *  χ_ab,
+        :κa => prefactor *  κa,
+        :κb => prefactor *  κb,
+        :cat_levels => cat_levels, :buffer_levels => buffer_levels, :prefactor=>prefactor
+    )
+    # Cat ⊗ Buffer
+    a = annihilate(cat_levels) ⊗ Matrix(1.0I, buffer_levels, buffer_levels)
+    b = Matrix(1.0I, cat_levels, cat_levels) ⊗ annihilate(buffer_levels)
+
+    H_drift = -χ_aa/2 * a'a'a*a - χ_bb/2 * b'b'b*b - χ_ab * a'a*b'b + g2 * a'a'b  + conj(g2) * a*a*b'
+
+    # buffer drive, kerr-correction drive
+    H_drives = [b + b', a'a]
+
+    L_dissipators = [√κa * a, √κb * b]
+
+    H_drift *= 2π
+    H_drives .*= 2π
+    L_dissipators .*= sqrt(2π)
+
+    return QuantumSystem(
+        H_drift,
+        H_drives,
+        L_dissipators;
+        params=params
+    )
+end
+
+function get_cat_controls(sys::AbstractQuantumSystem, α::Real, T::Int)
+    @assert haskey(sys.params, :g2) "Requires photon transfer coupling between buffer and cat"
+    @assert haskey(sys.params, :χ_aa) "Requires Kerr coupling for cat"
+    buffer_drive = abs2(α) * sys.params[:g2]
+    cat_kerr_correction = (2.0 * abs2(α) + 1.0) * sys.params[:χ_aa]
+    return stack([
+        fill(buffer_drive, T),
+        fill(cat_kerr_correction, T)
+    ], dims=1)
+end
diff --git a/src/quantum_system_templates/rydberg.jl b/src/quantum_system_templates/rydberg.jl
index cfc51324..cc59e0b2 100644
--- a/src/quantum_system_templates/rydberg.jl
+++ b/src/quantum_system_templates/rydberg.jl
@@ -1,3 +1,4 @@
+export RydbergChainSystem
 
 function generate_pattern(N::Int, i::Int)
     # Create an array filled with 'I'
@@ -65,8 +66,14 @@ function RydbergChainSystem(;
     all2all::Bool=true,
     ignore_Y_drive::Bool=false,
 )
-    PAULIS = Dict(:I => [1 0; 0 1], :X => [0 1; 1 0], :Y => [0 -im; im 0], :Z => [1 0; 0 -1], :n => [0 0; 0 1])
-    
+    PAULIS = (
+        I = ComplexF64[1 0; 0 1],
+        X = ComplexF64[0 1; 1 0],
+        Y = ComplexF64[0 -im; im 0],
+        Z = ComplexF64[1 0; 0 -1],
+        n = ComplexF64[0 0; 0 1]
+    )
+
     if all2all
         H_drift = zeros(ComplexF64, 2^N, 2^N)
         for gap in 0:N-2
@@ -104,23 +111,12 @@ function RydbergChainSystem(;
     H_detune = -sum([operator_from_string(lift('n',i,N), lookup=PAULIS) for i in 1:N])
     push!(H_drives, H_detune)
 
-    params = Dict{Symbol, Any}(
-        :N => N,
-        :C => C,
-        :distance => distance,
-        :cutoff_order => cutoff_order,
-        :local_detune => local_detune,
-        :all2all => all2all,
-    )
-
     return QuantumSystem(
         H_drift,
-        H_drives;
-        constructor=RydbergChainSystem,
-        params=params,
+        H_drives
     )
 end
 
 @testitem "Rydberg system test" begin
     @test RydbergChainSystem(N=3,cutoff_order=2,all2all=false) isa QuantumSystem
-end
\ No newline at end of file
+end
diff --git a/src/quantum_system_templates/transmons.jl b/src/quantum_system_templates/transmons.jl
index 7c75405b..23384125 100644
--- a/src/quantum_system_templates/transmons.jl
+++ b/src/quantum_system_templates/transmons.jl
@@ -1,3 +1,7 @@
+export TransmonSystem
+export TransmonDipoleCoupling
+export MultiTransmonSystem
+
 @doc raw"""
     TransmonSystem(;
         ω::Float64=4.4153,  # GHz
@@ -94,8 +98,7 @@ function TransmonSystem(;
     return QuantumSystem(
         H_drift,
         H_drives;
-        constructor=TransmonSystem,
-        params=params,
+        params=params
     )
 end
 
diff --git a/src/quantum_system_utils.jl b/src/quantum_system_utils.jl
deleted file mode 100644
index ea8567e5..00000000
--- a/src/quantum_system_utils.jl
+++ /dev/null
@@ -1,373 +0,0 @@
-module QuantumSystemUtils
-
-export operator_algebra
-export is_reachable
-
-using ..EmbeddedOperators
-using ..QuantumObjectUtils
-using ..QuantumSystems
-
-using LinearAlgebra
-using SparseArrays
-using TestItemRunner
-
-
-commutator(A::AbstractMatrix, B::AbstractMatrix) = A * B - B * A
-
-# TODO: Any difference between LinearAlgebra::ishermitian?
-is_hermitian(H::AbstractMatrix; atol=eps(Float32)) =
-    all(isapprox.(H - H', 0.0, atol=atol))
-
-function is_linearly_dependent(basis::Vector{<:AbstractMatrix}; kwargs...)
-    return is_linearly_dependent(stack(vec.(basis)); kwargs...)
-end
-
-function is_linearly_dependent(M::AbstractMatrix; eps=eps(Float32), verbose=true)
-    if size(M, 2) > size(M, 1)
-        if verbose
-            println("Linearly dependent because columns > rows, $(size(M, 2)) > $(size(M, 1)).")
-        end
-        return true
-    end
-    # QR decomposition has a zero R on diagonal if linearly dependent
-    val = minimum(abs.(diag(qr(M).R)))
-    return isapprox(val, 0.0, atol=eps)
-end
-
-function linearly_independent_indices(
-    basis::Vector{<:AbstractMatrix};
-    order=1:length(basis),
-    kwargs...
-)
-    @assert issetequal(order, 1:length(basis)) "Order must enumerate entire basis."
-    bᵢ = Int[]
-    for i ∈ order
-        if !is_linearly_dependent([basis[bᵢ]..., basis[i]]; kwargs...)
-            push!(bᵢ, i)
-        end
-    end
-    return bᵢ
-end
-
-function linearly_independent_subset(basis::Vector{<:AbstractMatrix}; kwargs...)
-    bᵢ = linearly_independent_indices(basis; kwargs...)
-    return deepcopy(basis[bᵢ])
-end
-
-function linearly_independent_subset!(basis::Vector{<:AbstractMatrix}; kwargs...)
-    bᵢ = linearly_independent_indices(basis; kwargs...)
-    deleteat!(basis, setdiff(1:length(basis), bᵢ))
-    return nothing
-end
-
-traceless(M::AbstractMatrix) = M - tr(M) * I / size(M, 1)
-
-"""
-operator_algebra(generators; return_layers=false, normalize=false, verbose=false, remove_trace=true)
-
-    Compute the Lie algebra basis for the given generators.
-    If return_layers is true, the Lie tree layers are also returned.
-
-    Can normalize the basis and enforce traceless generators.
-
-    # Arguments
-    - `generators::Vector{<:AbstractMatrix}`: generators of the Lie algebra
-    - `return_layers::Bool=false`: return the Lie tree layers
-    - `normalize::Bool=false`: normalize the basis
-    - `verbose::Bool=false`: print debug information
-    - `remove_trace::Bool=true`: remove trace from generators
-"""
-function operator_algebra(
-    generators::Vector{<:AbstractMatrix{T}};
-    return_layers=false,
-    normalize=false,
-    verbose=false,
-    remove_trace=true
-) where T<:Number
-    # Initialize basis (traceless, normalized)
-    basis = normalize ? [g / norm(g) for g ∈ generators] : deepcopy(generators)
-    if remove_trace
-        @. basis = traceless(basis)
-    end
-
-    # Initialize layers
-    current_layer = deepcopy(basis)
-    if return_layers
-        all_layers = Vector{Matrix{T}}[deepcopy(basis)]
-    end
-
-    ℓ = 1
-    if verbose
-        println("ℓ = $ℓ")
-    end
-    if is_linearly_dependent(basis)
-        println("Linearly dependent generators.")
-    else
-        # Note: Use left normalized commutators
-        # Note: Jacobi identity is not checked
-        need_basis = true
-        algebra_dim = size(first(generators), 1)^2 - 1
-        while need_basis
-            layer = Matrix{T}[]
-            # Repeat commutators until no new operators are found
-            for op ∈ current_layer
-                for gen ∈ generators
-                    if !need_basis
-                        continue
-                    end
-
-                    test = commutator(gen, op)
-                    if all(test .≈ 0)
-                        continue
-                    # Current basis is assumed to be linearly independent
-                    elseif is_linearly_dependent([basis..., test], verbose=verbose)
-                        continue
-                    else
-                        test .= is_hermitian(test) ? test : im * test
-                        test .= normalize ? test / norm(test) : test
-                        push!(basis, test)
-                        push!(layer, test)
-                        need_basis = length(basis) < algebra_dim ? true : false
-                    end
-                end
-            end
-
-            if isempty(layer)
-                if verbose
-                    println("Subspace termination.")
-                end
-                break
-            else
-                current_layer = layer
-                ℓ += 1
-                if verbose
-                    println("ℓ = $ℓ")
-                end
-            end
-
-            if return_layers
-                append!(all_layers, [current_layer])
-            end
-        end
-    end
-
-    if return_layers
-        return basis, all_layers
-    else
-        return basis
-    end
-end
-
-function fit_gen_to_basis(
-    gen::AbstractMatrix{<:T},
-    basis::AbstractVector{<:AbstractMatrix{<:T}}
-) where T<:Number
-    A = stack(vec.(basis))
-    b = vec(gen)
-    return A \ b
-end
-
-function is_in_span(
-    gen::AbstractMatrix,
-    basis::AbstractVector{<:AbstractMatrix};
-    subspace::AbstractVector{<:Int}=1:size(gen, 1),
-    atol=eps(Float32),
-    return_effective_gen=false,
-)
-    g_basis = [deepcopy(b[subspace, subspace]) for b ∈ basis]
-    linearly_independent_subset!(g_basis)
-    # Traceless basis needs traceless fit
-    x = fit_gen_to_basis(gen, g_basis)
-    g_eff = sum(x .* g_basis)
-    ε = norm(g_eff - gen, 2)
-    if return_effective_gen
-        return ε < atol, g_eff
-    else
-        return ε < atol
-    end
-end
-
-"""
-    is_reachable(gate, hamiltonians; kwargs...)
-
-Check if the gate is reachable from the given generators. If subspace_indices are provided,
-then the gate should be given in the subspace.
-
-# Arguments
-- `gate::AbstractMatrix`: target gate
-- `hamiltonians::AbstractVector{<:AbstractMatrix}`: generators of the Lie algebra
-
-# Keyword Arguments
-- `subspace::AbstractVector{<:Int}=1:size(gate, 1)`: subspace indices
-- `compute_basis::Bool=true`: compute the basis
-- `remove_trace::Bool=true`: remove trace from generators
-- `verbose::Bool=false`: print debug information
-- `atol::Float32=eps(Float32)`: absolute tolerance
-"""
-function is_reachable(
-    gate::AbstractMatrix,
-    hamiltonians::AbstractVector{<:AbstractMatrix};
-    subspace::AbstractVector{<:Int}=1:size(gate, 1),
-    compute_basis=true,
-    remove_trace=true,
-    verbose=false,
-    atol=eps(Float32)
-)
-    @assert size(gate, 1) == length(subspace) "Gate must be given in the subspace."
-    generator = im * log(gate)
-
-    if remove_trace
-        generator = traceless(generator)
-    end
-
-    if compute_basis
-        basis = operator_algebra(hamiltonians, remove_trace=remove_trace, verbose=verbose)
-    else
-        basis = hamiltonians
-    end
-
-    return is_in_span(
-        generator,
-        basis,
-        subspace=subspace,
-        atol=atol
-    )
-end
-
-function is_reachable(
-    gate::AbstractMatrix,
-    system::QuantumSystem;
-    use_drift::Bool=true,
-    kwargs...
-)
-    if !iszero(system.H_drift) && use_drift
-        hamiltonians = [system.H_drift, system.H_drives...]
-    else
-        hamiltonians = system.H_drives
-    end
-    return is_reachable(gate, hamiltonians; kwargs...)
-end
-
-function is_reachable(
-    gate::EmbeddedOperator,
-    system::QuantumSystem;
-    kwargs...
-)
-    return is_reachable(
-        unembed(gate),
-        system;
-        subspace=gate.subspace_indices,
-        kwargs...
-    )
-end
-
-# *************************************************************************** #
-
-@testitem "Lie algebra basis" begin
-    # Check 1 qubit with complete basis
-    gen = operator_from_string.(["X", "Y"])
-    basis = operator_algebra(gen, return_layers=false, verbose=false)
-    @test length(basis) == size(first(gen), 1)^2-1
-
-    # Check 1 qubit with complete basis and layers
-    basis, layers = operator_algebra(gen, return_layers=true, verbose=false)
-    @test length(basis) == size(first(gen), 1)^2-1
-
-    # Check 1 qubit with subspace
-    gen = operator_from_string.(["X"])
-    basis = operator_algebra(gen, verbose=false)
-    @test length(basis) == 1
-
-    # Check 2 qubit with complete basis
-    gen = operator_from_string.(["XX", "YY", "XI", "YI", "IY", "IX"])
-    basis = operator_algebra(gen, verbose=false)
-    @test length(basis) == size(first(gen), 1)^2-1
-
-    # Check 2 qubit with linearly dependent basis
-    gen = operator_from_string.(["XX", "YY", "XI", "XI", "IY", "IX"])
-    basis = operator_algebra(gen, verbose=false)
-    @test length(basis) == length(gen)
-
-    # Check 2 qubit with pair of 1-qubit subspaces
-    gen = operator_from_string.(["XI", "YI", "IY", "IX"])
-    basis = operator_algebra(gen, verbose=false)
-    @test length(basis) == 2 * (2^2 - 1)
-end
-
-
-@testitem "Lie Algebra reachability" begin
-    using LinearAlgebra
-
-    H_ops = Dict(
-        "X" => GATES[:X],
-        "Y" => GATES[:Y],
-        "Z" => GATES[:Z]
-    )
-
-    # Check 1 qubit with complete basis
-    gen = operator_from_string.(["X", "Y"])
-    target = H_ops["Z"]
-    @test is_reachable(target, gen, compute_basis=true, verbose=false)
-
-    # System
-    sys = QuantumSystem([GATES[:X], GATES[:Y], GATES[:Z]])
-    target = GATES[:Z]
-    @test is_reachable(target, sys)
-
-    # System with drift
-    sys = QuantumSystem(GATES[:Z], [GATES[:X]])
-    target = GATES[:Z]
-    @test is_reachable(target, sys)
-
-    # Check 2 qubit with complete basis
-    XI = GATES[:X] ⊗ GATES[:I]
-    IX = GATES[:I] ⊗ GATES[:X]
-    YI = GATES[:Y] ⊗ GATES[:I]
-    IY = GATES[:I] ⊗ GATES[:Y]
-    XX = GATES[:X] ⊗ GATES[:X]
-    YY = GATES[:Y] ⊗ GATES[:Y]
-    ZI = GATES[:Z] ⊗ GATES[:I]
-    IZ = GATES[:I] ⊗ GATES[:Z]
-    ZZ = GATES[:Z] ⊗ GATES[:Z]
-
-    complete_gen = [XX+YY, XI, YI, IX, IY]
-    incomplete_gen = [XI, ZZ]
-    r = [0, 1, 2, 3, 4]
-    r /= norm(r)
-    R2 = exp(-im * sum([θ * H for (H, θ) in zip(complete_gen, r)]))
-    CZ = [1 0 0 0; 0 1 0 0; 0 0 1 0; 0 0 0 -1]
-    CX = [1 0 0 0; 0 1 0 0; 0 0 0 1; 0 0 1 0]
-
-    # Pass
-    @test is_reachable(R2, complete_gen)
-    @test is_reachable(CZ, complete_gen)
-    @test is_reachable(CX, complete_gen)
-    @test is_reachable(XI, complete_gen)
-
-    # Mostly fail
-    @test !is_reachable(R2, incomplete_gen)
-    @test !is_reachable(CZ, incomplete_gen)
-    @test !is_reachable(CX, incomplete_gen)
-    @test is_reachable(XI, incomplete_gen)
-
-    # QuantumSystems
-    complete_gen_sys = QuantumSystem(complete_gen)
-    incomplete_gen_sys = QuantumSystem(incomplete_gen)
-    # Pass
-    @test is_reachable(R2, complete_gen_sys)
-    @test is_reachable(CZ, complete_gen_sys)
-    @test is_reachable(CX, complete_gen_sys)
-    @test is_reachable(XI, complete_gen_sys)
-
-    # Mostly fail
-    @test !is_reachable(R2, incomplete_gen_sys)
-    @test !is_reachable(CZ, incomplete_gen_sys)
-    @test !is_reachable(CX, incomplete_gen_sys)
-    @test is_reachable(XI, incomplete_gen_sys)
-end
-
-@testitem "Lie Algebra subspace reachability" begin
-    # TODO: implement tests
-end
-
-end
diff --git a/src/quantum_systems.jl b/src/quantum_systems.jl
deleted file mode 100644
index 6bf796f2..00000000
--- a/src/quantum_systems.jl
+++ /dev/null
@@ -1,425 +0,0 @@
-module QuantumSystems
-
-export AbstractQuantumSystem
-export QuantumSystem
-export CompositeQuantumSystem
-export QuantumSystemCoupling
-
-export iso
-export lift
-
-using ..Isomorphisms
-using ..QuantumObjectUtils
-
-using LinearAlgebra
-using SparseArrays
-using TestItemRunner
-
-# TODO:
-# [ ] subtypes? SingleQubitSystem, TwoQubitSystem, TransmonSystem, MultimodeSystem, etc.
-# [ ] add frame info to type
-# [ ] add methods to combine composite quantum systems
-
-# ----------------------------------------------------------------------------- #
-# AbstractQuantumSystem
-# ----------------------------------------------------------------------------- #
-
-"""
-    AbstractQuantumSystem
-
-Abstract type for defining systems.
-"""
-abstract type AbstractQuantumSystem end
-
-# ----------------------------------------------------------------------------- #
-# QuantumSystem
-# ----------------------------------------------------------------------------- #
-
-"""
-    QuantumSystem <: AbstractQuantumSystem
-
-A struct for storing the isomorphisms of the system's drift and drive Hamiltonians,
-as well as the system's parameters.
-"""
-struct QuantumSystem <: AbstractQuantumSystem
-    H_drift::SparseMatrixCSC{ComplexF64, Int}
-    H_drives::Vector{SparseMatrixCSC{ComplexF64, Int}}
-    G_drift::SparseMatrixCSC{Float64, Int}
-    G_drives::Vector{SparseMatrixCSC{Float64, Int}}
-    dissipation_operators::Union{Nothing, Vector{SparseMatrixCSC{ComplexF64, Int}}}
-    levels::Int
-    constructor::Union{Function, Nothing}
-    params::Dict{Symbol, <:Any}
-end
-
-"""
-    QuantumSystem(
-        H_drift::Matrix{<:Number},
-        H_drives::Vector{Matrix{<:Number}};
-        params=Dict{Symbol, Any}(),
-        kwargs...
-    )::QuantumSystem
-
-Constructs a `QuantumSystem` object from the drift and drive Hamiltonian terms.
-"""
-function QuantumSystem(
-    H_drift::AbstractMatrix{<:Number},
-    H_drives::Vector{<:AbstractMatrix{<:Number}};
-    dissipation_operators=nothing,
-    constructor::Union{Function, Nothing}=nothing,
-    params::Dict{Symbol, <:Any}=Dict{Symbol, Any}(),
-    kwargs...
-)
-    H_drift = sparse(H_drift)
-    H_drives = sparse.(H_drives)
-    G_drift = Isomorphisms.G(H_drift)
-    G_drives = Isomorphisms.G.(H_drives)
-    params = merge(params, Dict{Symbol, Any}(kwargs...))
-    levels = size(H_drift, 1)
-    return QuantumSystem(
-        H_drift,
-        H_drives,
-        G_drift,
-        G_drives,
-        dissipation_operators,
-        levels,
-        constructor,
-        params
-    )
-end
-
-function QuantumSystem(H_drives::Vector{<:AbstractMatrix{<:Number}}; kwargs...)
-    return QuantumSystem(
-        zeros(eltype(H_drives[1]), size(H_drives[1])),
-        H_drives;
-        kwargs...
-    )
-end
-
-function QuantumSystem(H_drift::AbstractMatrix{<:Number}; kwargs...)
-    return QuantumSystem(
-        H_drift,
-        Matrix{ComplexF64}[];
-        kwargs...
-    )
-end
-
-function (sys::QuantumSystem)(; params...)
-    @assert !isnothing(sys.constructor) "No constructor provided."
-    @assert all([
-        key ∈ keys(sys.params) for key ∈ keys(params)
-    ]) "Invalid parameter(s) provided."
-    return sys.constructor(; merge(sys.params, Dict(params...))...)
-end
-
-function QuantumSystem(
-    H_drift::SparseMatrixCSC{ComplexF64, Int64},
-    H_drives::Vector{SparseMatrixCSC{ComplexF64, Int64}};
-    dissipation_operators=nothing,
-    constructor::Union{Function, Nothing}=nothing,
-    params::Dict{Symbol, <:Any}=Dict{Symbol, Any}(),
-    kwargs...
-)
-    H_drift = sparse(H_drift)
-    H_drives = sparse.(H_drives)
-    G_drift = Isomorphisms.G(H_drift)
-    G_drives = Isomorphisms.G.(H_drives)
-    params = merge(params, Dict{Symbol, Any}(kwargs...))
-    levels = size(H_drift, 1)
-    return QuantumSystem(
-        H_drift,
-        H_drives,
-        G_drift,
-        G_drives,
-        dissipation_operators,
-        levels,
-        constructor,
-        params
-    )
-end
-
-function Base.copy(sys::QuantumSystem)
-    return QuantumSystem(
-        copy(sys.H_drift),
-        copy.(sys.H_drives);
-        constructor=sys.constructor,
-        params=copy(sys.params)
-    )
-end
-
-# ----------------------------------------------------------------------------- #
-# Quantum System couplings
-# ----------------------------------------------------------------------------- #
-
-@doc raw"""
-    lift(U::AbstractMatrix{<:Number}, qubit_index::Int, n_qubits::Int; levels::Int=size(U, 1))
-
-Lift an operator `U` acting on a single qubit to an operator acting on the entire system of `n_qubits`.
-"""
-function lift(
-    U::AbstractMatrix{<:Number},
-    qubit_index::Int,
-    n_qubits::Int;
-    levels::Int=size(U, 1)
-)::Matrix{ComplexF64}
-    Is = Matrix{Complex}[I(levels) for _ = 1:n_qubits]
-    Is[qubit_index] = U
-    return foldr(⊗, Is)
-end
-
-@doc raw"""
-    lift(op::AbstractMatrix{<:Number}, i::Int, subsystem_levels::Vector{Int})
-
-Lift an operator `op` acting on the i-th subsystem to an operator acting on the entire system with given subsystem levels.
-"""
-function lift(
-    op::AbstractMatrix{<:Number},
-    i::Int,
-    subsystem_levels::Vector{Int}
-)::Matrix{ComplexF64}
-    @assert size(op, 1) == size(op, 2) == subsystem_levels[i] "Operator must be square and match dimension of subsystem i"
-
-    Is = [collect(1.0 * typeof(op)(I, l, l)) for l ∈ subsystem_levels]
-    Is[i] = op
-    return kron(1.0, Is...)
-end
-
-"""
-    QuantumSystemCoupling <: AbstractQuantumSystem
-
-"""
-struct QuantumSystemCoupling
-    term::SparseMatrixCSC{ComplexF64, Int}
-    g_ij::Float64
-    pair::Tuple{Int, Int}
-    subsystem_levels::Vector{Int}
-    constructor::Union{Function, Nothing}
-    params::Dict{Symbol, <:Any}
-
-    function QuantumSystemCoupling(op::AbstractMatrix{<:ComplexF64}, args...)
-        return new(
-            sparse(op),
-            args...
-        )
-    end
-end
-
-function (coupling::QuantumSystemCoupling)(;
-    g_ij::Float64=coupling.g_ij,
-    pair::Tuple{Int, Int}=coupling.pair,
-    subsystem_levels::Vector{Int}=coupling.subsystem_levels,
-    params...
-)
-    @assert !isnothing(coupling.constructor) "No constructor provided."
-    @assert all([
-        key ∈ keys(coupling.params) for key ∈ keys(params)
-    ]) "Invalid parameter(s) provided: $(filter(param -> param ∉ keys(coupling.params), keys(params)))"
-    return coupling.constructor(
-        g_ij,
-        pair,
-        subsystem_levels;
-        merge(coupling.params, Dict(params...))...
-    )
-end
-
-function Base.copy(coupling::QuantumSystemCoupling)
-    return QuantumSystemCoupling(
-        copy(coupling.op),
-        coupling.g_ij,
-        coupling.pair,
-        coupling.subsystem_levels,
-        coupling.constructor,
-        copy(coupling.params)
-    )
-end
-
-
-struct CompositeQuantumSystem <: AbstractQuantumSystem
-    H_drift::SparseMatrixCSC{ComplexF64, Int}
-    H_drives::Vector{SparseMatrixCSC{ComplexF64, Int}}
-    G_drift::SparseMatrixCSC{Float64, Int}
-    G_drives::Vector{SparseMatrixCSC{Float64, Int}}
-    levels::Int
-    subsystem_levels::Vector{Int}
-    params::Dict{Symbol, Any}
-    subsystems::Vector{QuantumSystem}
-    couplings::Vector{QuantumSystemCoupling}
-end
-
-function CompositeQuantumSystem(
-    subsystems::Vector{QuantumSystem},
-    couplings::Vector{QuantumSystemCoupling}=QuantumSystemCoupling[];
-    subsystem_frame_index::Int=1,
-    frame_ω::Float64=subsystems[subsystem_frame_index].params[:ω],
-    lab_frame::Bool=false
-)
-    # set all subsystems to the same frame_ω
-    subsystems = [sys(; frame_ω=frame_ω) for sys ∈ subsystems]
-
-    if lab_frame
-        subsystems = [sys(; lab_frame=true) for sys ∈ subsystems]
-        couplings = [coupling(; lab_frame=true) for coupling ∈ couplings]
-    end
-
-    subsystem_levels = [sys.levels for sys ∈ subsystems]
-    levels = prod(subsystem_levels)
-
-    # add lifted subsystem drift Hamiltonians
-    H_drift = sparse(zeros(levels, levels))
-    for (i, sys) ∈ enumerate(subsystems)
-        H_drift += lift(sys.H_drift, i, subsystem_levels)
-    end
-
-    # add lifated couplings to the drift Hamiltonian
-    for coupling ∈ couplings
-        H_drift += coupling.term
-    end
-
-    # add lifted subsystem drive Hamiltonians
-    H_drives = SparseMatrixCSC{ComplexF64, Int}[]
-    for (i, sys) ∈ enumerate(subsystems)
-        for H_drive ∈ sys.H_drives
-            push!(H_drives, lift(H_drive, i, subsystem_levels))
-        end
-    end
-
-    G_drift = Isomorphisms.G(H_drift)
-    G_drives = Isomorphisms.G.(H_drives)
-    levels = size(H_drift, 1)
-    subsystem_levels = [sys.levels for sys ∈ subsystems]
-    params = Dict{Symbol, Any}()
-
-    return CompositeQuantumSystem(
-        H_drift,
-        H_drives,
-        G_drift,
-        G_drives,
-        levels,
-        subsystem_levels,
-        params,
-        subsystems,
-        couplings
-    )
-end
-
-function (csys::CompositeQuantumSystem)(;
-    subsystem_params::Dict{Int, <:Dict{Symbol, <:Any}}=Dict{Int, Dict{Symbol, Any}}(),
-    coupling_params::Dict{Int, <:Dict{Symbol, <:Any}}=Dict{Int, Dict{Symbol, Any}}(),
-    lab_frame::Bool=false,
-    subsystem_frame_index::Int=1,
-    frame_ω::Float64=csys.subsystems[subsystem_frame_index].params[:ω],
-    subsystem_levels::Union{Nothing, Int, Vector{Int}}=nothing,
-)
-    subsystems = deepcopy(csys.subsystems)
-    couplings = deepcopy(csys.couplings)
-
-    # if lab frame then set all subsystems and couplings to lab frame
-    if lab_frame
-
-        # set lab frame in subsystem_params for all subsystems
-        for i = 1:length(csys.subsystems)
-            if i ∈ keys(subsystem_params)
-                subsystem_params[i][:lab_frame] = true
-            else
-                subsystem_params[i] = Dict{Symbol, Any}(:lab_frame => true)
-            end
-        end
-
-        # set lab frame in coupling_params for all couplings
-        for i = 1:length(csys.couplings)
-            if i ∈ keys(coupling_params)
-                coupling_params[i][:lab_frame] = true
-            else
-                coupling_params[i] = Dict{Symbol, Any}(:lab_frame => true)
-            end
-        end
-    end
-
-    # if subsystem_levels is provided then set all subsystems and couplings to subsystem_levels
-    if !isnothing(subsystem_levels)
-
-        if subsystem_levels isa Int
-            subsystem_levels = fill(subsystem_levels, length(csys.subsystems))
-        else
-            @assert(
-                length(subsystem_levels) == length(csys.subsystems),
-                """\n
-                    number of subsystem_levels ($(length(subsystem_levels))) must match number of subsystems ($(length(csys.subsystems))).
-                """
-            )
-        end
-
-        for i = 1:length(csys.subsystems)
-            if i ∈ keys(subsystem_params)
-                subsystem_params[i][:levels] = subsystem_levels[i]
-            else
-                subsystem_params[i] = Dict{Symbol, Any}(
-                    :levels => subsystem_levels[i]
-                )
-            end
-        end
-
-        for i = 1:length(csys.couplings)
-            if i ∈ keys(coupling_params)
-                coupling_params[i][:subsystem_levels] = subsystem_levels
-            else
-                coupling_params[i] = Dict{Symbol, Any}(
-                    :subsystem_levels => subsystem_levels
-                )
-            end
-        end
-    end
-
-    # construct subsystems with new parameters
-    for (i, sys_params) ∈ subsystem_params
-        subsystem_i_new_params = merge(subsystems[i].params, sys_params)
-        subsystem_i_new_params[:frame_ω] = frame_ω
-        subsystems[i] = subsystems[i](; subsystem_i_new_params...)
-    end
-
-    # sometimes redundant, but here to catch any changes in indvidual subsystem levels
-    subsystem_levels = [sys.levels for sys ∈ subsystems]
-
-    # construct couplings with new parameters
-    if !isempty(csys.couplings)
-        for (i, coupling_params) ∈ coupling_params
-            couplings[i] = couplings[i](;
-                merge(couplings[i].params, coupling_params)...,
-                subsystem_levels=subsystem_levels
-            )
-        end
-    end
-
-    return CompositeQuantumSystem(
-        subsystems,
-        couplings
-    )
-end
-
-# *************************************************************************** #
-
-@testitem "System creation" begin
-    H_drift = GATES[:Z]
-    H_drives = [GATES[:X], GATES[:Y]]
-    n_drives = length(H_drives)
-
-    system = QuantumSystem(H_drift, H_drives)
-end
-
-function is_reachable(
-    gate::AbstractMatrix,
-    system::QuantumSystem;
-    use_drift::Bool=true,
-    kwargs...
-)
-    if !iszero(system.H_drift) && use_drift
-        hamiltonians = [system.H_drift, system.H_drives...]
-    else
-        hamiltonians = system.H_drives
-    end
-    return is_reachable(gate, hamiltonians; kwargs...)
-end
-
-
-
-end
diff --git a/src/rollouts.jl b/src/rollouts.jl
index d2c448a3..96755faa 100644
--- a/src/rollouts.jl
+++ b/src/rollouts.jl
@@ -6,23 +6,19 @@ export unitary_rollout
 export lab_frame_unitary_rollout
 export lab_frame_unitary_rollout_trajectory
 
-using ..Isomorphisms
-using ..QuantumSystems
-using ..QuantumSystemUtils
-using ..QuantumObjectUtils
-using ..EmbeddedOperators
-using ..Integrators
-using ..Losses
-using ..Problems
+export rollout_fidelity
+export unitary_rollout_fidelity
+
 using ..DirectSums
 
 using NamedTrajectories
-
+using QuantumCollocationCore
+using PiccoloQuantumObjects
 using ExponentialAction
 using LinearAlgebra
 using ProgressMeter
 using TestItemRunner
-
+using ForwardDiff
 
 # ----------------------------------------------------------------------------- #
 
@@ -78,7 +74,6 @@ function rollout(
     show_progress=false,
     integrator=expv,
     exp_vector_product=infer_is_evp(integrator),
-    G=Integrators.G_bilinear
 )
     T = size(controls, 2)
 
@@ -88,17 +83,14 @@ function rollout(
 
     Ψ̃[:, 1] .= ψ̃_init
 
-    G_drift = Matrix{Float64}(system.G_drift)
-    G_drives = Matrix{Float64}.(system.G_drives)
-
     p = Progress(T-1; enabled=show_progress)
     for t = 2:T
         aₜ₋₁ = controls[:, t - 1]
-        Gₜ = G(aₜ₋₁, G_drift, G_drives)
+        Gₜ = system.G(aₜ₋₁)
         if exp_vector_product
             Ψ̃[:, t] .= integrator(Δt[t - 1], Gₜ, Ψ̃[:, t - 1])
         else
-            Ψ̃[:, t] .= integrator(Gₜ * Δt[t - 1]) * Ψ̃[:, t - 1]
+            Ψ̃[:, t] .= integrator(Matrix(Gₜ) * Δt[t - 1]) * Ψ̃[:, t - 1]
         end
         next!(p)
     end
@@ -115,7 +107,7 @@ function rollout(
     return vcat([rollout(state, args...; kwargs...) for state ∈ inits]...)
 end
 
-function Losses.iso_fidelity(
+function rollout_fidelity(
     ψ̃_init::AbstractVector{<:Real},
     ψ̃_goal::AbstractVector{<:Real},
     controls::AbstractMatrix,
@@ -127,7 +119,7 @@ function Losses.iso_fidelity(
     return iso_fidelity(Ψ̃[:, end], ψ̃_goal)
 end
 
-function Losses.fidelity(
+function rollout_fidelity(
     ψ_init::AbstractVector{<:Complex},
     ψ_goal::AbstractVector{<:Complex},
     controls::AbstractMatrix,
@@ -135,34 +127,35 @@ function Losses.fidelity(
     system::AbstractQuantumSystem;
     kwargs...
 )
-    return iso_fidelity(ket_to_iso(ψ_init), ket_to_iso(ψ_goal), controls, Δt, system; kwargs...)
+    return rollout_fidelity(ket_to_iso(ψ_init), ket_to_iso(ψ_goal), controls, Δt, system; kwargs...)
 end
 
-function Losses.fidelity(
+function rollout_fidelity(
     trajectory::NamedTrajectory,
     system::AbstractQuantumSystem;
-    state_symb::Symbol=:ψ̃,
-    control_symb=:a,
+    state_name::Symbol=:ψ̃,
+    control_name=:a,
     kwargs...
 )
     fids = []
-    for symb in trajectory.names
-        if startswith(symb, state_symb)
-            controls = trajectory[control_symb]
-            init = trajectory.initial[symb]
-            goal = trajectory.goal[symb]
-            fid = iso_fidelity(init, goal, controls, get_timesteps(trajectory), system; kwargs...)
+    for name ∈ trajectory.names
+        if startswith(name, state_name)
+            controls = trajectory[control_name]
+            init = trajectory.initial[name]
+            goal = trajectory.goal[name]
+            fid = rollout_fidelity(init, goal, controls, get_timesteps(trajectory), system; kwargs...)
             push!(fids, fid)
         end
     end
     return length(fids) == 1 ? fids[1] : fids
 end
 
-function Losses.fidelity(
-    prob::QuantumControlProblem;
+function rollout_fidelity(
+    prob::QuantumControlProblem,
+    system::AbstractQuantumSystem;
     kwargs...
 )
-    return fidelity(prob.trajectory, prob.system; kwargs...)
+    return rollout_fidelity(prob.trajectory, system; kwargs...)
 end
 
 # ----------------------------------------------------------------------------- #
@@ -215,15 +208,14 @@ end
 # ----------------------------------------------------------------------------- #
 
 function unitary_rollout(
-    Ũ⃗_init::AbstractVector{<:Real},
-    controls::AbstractMatrix,
-    Δt::AbstractVector,
+    Ũ⃗_init::AbstractVector{R1},
+    controls::AbstractMatrix{R2},
+    Δt::AbstractVector{R3},
     system::AbstractQuantumSystem;
     show_progress=false,
     integrator=expv,
     exp_vector_product=infer_is_evp(integrator),
-    G=Integrators.G_bilinear,
-)
+) where {R1 <: Real, R2 <: Real, R3 <: Real}
     T = size(controls, 2)
 
     # Enable ForwardDiff
@@ -232,18 +224,15 @@ function unitary_rollout(
 
     Ũ⃗[:, 1] .= Ũ⃗_init
 
-    G_drift = Matrix{Float64}(system.G_drift)
-    G_drives = Matrix{Float64}.(system.G_drives)
-
     p = Progress(T-1; enabled=show_progress)
     for t = 2:T
         aₜ₋₁ = controls[:, t - 1]
-        Gₜ = G(aₜ₋₁, G_drift, G_drives)
+        Gₜ = system.G(aₜ₋₁)
         Ũₜ₋₁ = iso_vec_to_iso_operator(Ũ⃗[:, t - 1])
         if exp_vector_product
             Ũₜ = integrator(Δt[t - 1], Gₜ, Ũₜ₋₁)
         else
-            Ũₜ = integrator(Gₜ * Δt[t - 1]) * Ũₜ₋₁
+            Ũₜ = integrator(Matrix(Gₜ) * Δt[t - 1]) * Ũₜ₋₁
         end
         Ũ⃗[:, t] .= iso_operator_to_iso_vec(Ũₜ)
         next!(p)
@@ -258,7 +247,7 @@ function unitary_rollout(
     system::AbstractQuantumSystem;
     kwargs...
 )
-    Ĩ⃗ = operator_to_iso_vec(Matrix{ComplexF64}(I(size(system.H_drift, 1))))
+    Ĩ⃗ = operator_to_iso_vec(Matrix{ComplexF64}(I(system.levels)))
     return unitary_rollout(Ĩ⃗, controls, Δt, system; kwargs...)
 end
 
@@ -278,10 +267,7 @@ function unitary_rollout(
     )
 end
 
-"""
-Compute the rollout fidelity.
-"""
-function Losses.iso_vec_unitary_fidelity(
+function unitary_rollout_fidelity(
     Ũ⃗_init::AbstractVector{<:Real},
     Ũ⃗_goal::AbstractVector{<:Real},
     controls::AbstractMatrix,
@@ -300,18 +286,18 @@ function Losses.iso_vec_unitary_fidelity(
     end
 end
 
-function Losses.iso_vec_unitary_fidelity(
+function unitary_rollout_fidelity(
     Ũ⃗_goal::AbstractVector{<:Real},
     controls::AbstractMatrix,
     Δt::AbstractVector,
     system::AbstractQuantumSystem;
     kwargs...
 )
-    Ĩ⃗ = operator_to_iso_vec(Matrix{ComplexF64}(I(size(system.H_drift, 1))))
-    return iso_vec_unitary_fidelity(Ĩ⃗, Ũ⃗_goal, controls, Δt, system; kwargs...)
+    Ĩ⃗ = operator_to_iso_vec(Matrix{ComplexF64}(I(system.levels)))
+    return unitary_rollout_fidelity(Ĩ⃗, Ũ⃗_goal, controls, Δt, system; kwargs...)
 end
 
-function Losses.unitary_fidelity(
+function unitary_rollout_fidelity(
     U_init::AbstractMatrix{<:Complex},
     U_goal::AbstractMatrix{<:Complex},
     controls::AbstractMatrix,
@@ -321,27 +307,27 @@ function Losses.unitary_fidelity(
 )
     Ũ⃗_init = operator_to_iso_vec(U_init)
     Ũ⃗_goal = operator_to_iso_vec(U_goal)
-    return iso_vec_unitary_fidelity(Ũ⃗_init, Ũ⃗_goal, controls, Δt, system; kwargs...)
+    return unitary_rollout_fidelity(Ũ⃗_init, Ũ⃗_goal, controls, Δt, system; kwargs...)
 end
 
-Losses.unitary_fidelity(
+unitary_rollout_fidelity(
     U_goal::AbstractMatrix{<:Complex},
     controls::AbstractMatrix,
     Δt::AbstractVector,
     system::AbstractQuantumSystem;
     kwargs...
-) = iso_vec_unitary_fidelity(operator_to_iso_vec(U_goal), controls, Δt, system; kwargs...)
+) = unitary_rollout_fidelity(operator_to_iso_vec(U_goal), controls, Δt, system; kwargs...)
 
-Losses.unitary_fidelity(
+unitary_rollout_fidelity(
     U_goal::EmbeddedOperator,
     controls::AbstractMatrix,
     Δt::AbstractVector,
     system::AbstractQuantumSystem;
-    subspace::AbstractVector{Int}=U_goal.subspace_indices,
+    subspace::AbstractVector{Int}=U_goal.subspace,
     kwargs...
-) = unitary_fidelity(U_goal.operator, controls, Δt, system; subspace=subspace, kwargs...)
+) = unitary_rollout_fidelity(U_goal.operator, controls, Δt, system; subspace=subspace, kwargs...)
 
-function Losses.unitary_fidelity(
+function unitary_rollout_fidelity(
     traj::NamedTrajectory,
     sys::AbstractQuantumSystem;
     unitary_name::Symbol=:Ũ⃗,
@@ -352,43 +338,21 @@ function Losses.unitary_fidelity(
     Ũ⃗_goal = traj.goal[unitary_name]
     controls = traj[drive_name]
     Δt = get_timesteps(traj)
-    return iso_vec_unitary_fidelity(Ũ⃗_init, Ũ⃗_goal, controls, Δt, sys; kwargs...)
-end
-
-function Losses.unitary_fidelity(
-    traj::NamedTrajectory,
-    systems::AbstractVector{<:AbstractQuantumSystem};
-    unitary_name::Symbol=:Ũ⃗,
-    unitary_names::AbstractVector{<:Symbol}=[n for n ∈ traj.names if startswith(string(n), string(unitary_name))],
-    drive_name::Symbol=:a,
-    kwargs...
-)
-    controls = traj[drive_name]
-    Δt = get_timesteps(traj)
-    fids = map(zip(systems, unitary_names)) do (sys, name)
-        Ũ⃗_goal = traj.goal[name]
-        Ũ⃗_init = traj.initial[name]
-        iso_vec_unitary_fidelity(Ũ⃗_init, Ũ⃗_goal, controls, Δt, sys; kwargs...)
-    end
-    return fids
+    return unitary_rollout_fidelity(Ũ⃗_init, Ũ⃗_goal, controls, Δt, sys; kwargs...)
 end
 
-Losses.unitary_fidelity(prob::QuantumControlProblem; kwargs...) =
-    unitary_fidelity(prob.trajectory, prob.system; kwargs...)
-
-
 # ----------------------------------------------------------------------------- #
 # Experimental rollouts
 # ----------------------------------------------------------------------------- #
 
-Losses.unitary_fidelity(
+unitary_rollout_fidelity(
     U_goal::EmbeddedOperator,
     controls::AbstractMatrix{Float64},
     Δt::Union{AbstractVector{Float64}, AbstractMatrix{Float64}, Float64},
     sys::AbstractQuantumSystem;
-    subspace=U_goal.subspace_indices,
+    subspace=U_goal.subspace,
     kwargs...
-) = unitary_fidelity(U_goal.operator, controls, Δt, sys; subspace=subspace, kwargs...)
+) = unitary_rollout_fidelity(U_goal.operator, controls, Δt, sys; subspace=subspace, kwargs...)
 
 """
     lab_frame_unitary_rollout(
@@ -486,42 +450,44 @@ end
 
     # Default integrator
     # State fidelity
-    @test fidelity(ψ, ψ_goal, as, ts, sys) ≈ 1
-    @test iso_fidelity(ψ̃, ψ̃_goal, as, ts, sys) ≈ 1
+    @test rollout_fidelity(ψ, ψ_goal, as, ts, sys) ≈ 1
 
     # Unitary fidelity
-    @test unitary_fidelity(U_goal, as, ts, sys) ≈ 1
-    @test unitary_fidelity(prob.trajectory, sys) ≈ 1
-    @test unitary_fidelity(prob) ≈ 1
-    @test unitary_fidelity(embedded_U_goal, as, ts, sys) ≈ 1
+    @test unitary_rollout_fidelity(U_goal, as, ts, sys) ≈ 1
+    @test unitary_rollout_fidelity(prob.trajectory, sys) ≈ 1
+    @test unitary_rollout_fidelity(embedded_U_goal, as, ts, sys) ≈ 1
+
+    # Free phase unitary
+    @test unitary_rollout_fidelity(prob.trajectory, sys;
+        phases=[0.0], phase_operators=Matrix{ComplexF64}[PAULIS[:Z]]
+    ) ≈ 1
 
     # Free phase unitary
-    @test unitary_fidelity(prob, phases=[0.0], phase_operators=[PAULIS[:Z]]) ≈ 1
+    @test unitary_rollout_fidelity(prob.trajectory, sys;
+        phases=[0.0],
+        phase_operators=[PAULIS[:Z]]
+    ) ≈ 1
 
     # Expv explicit
     # State fidelity
-    @test fidelity(ψ, ψ_goal, as, ts, sys, integrator=expv) ≈ 1
-    @test iso_fidelity(ψ̃, ψ̃_goal, as, ts, sys, integrator=expv) ≈ 1
+    @test rollout_fidelity(ψ, ψ_goal, as, ts, sys, integrator=expv) ≈ 1
 
     # Unitary fidelity
-    @test unitary_fidelity(U_goal, as, ts, sys, integrator=expv) ≈ 1
-    @test unitary_fidelity(prob.trajectory, sys, integrator=expv) ≈ 1
-    @test unitary_fidelity(prob, integrator=expv) ≈ 1
-    @test unitary_fidelity(embedded_U_goal, as, ts, sys, integrator=expv) ≈ 1
+    @test unitary_rollout_fidelity(U_goal, as, ts, sys, integrator=expv) ≈ 1
+    @test unitary_rollout_fidelity(prob.trajectory, sys, integrator=expv) ≈ 1
+    @test unitary_rollout_fidelity(embedded_U_goal, as, ts, sys, integrator=expv) ≈ 1
 
     # Exp explicit
     # State fidelity
-    @test fidelity(ψ, ψ_goal, as, ts, sys, integrator=exp) ≈ 1
-    @test iso_fidelity(ψ̃, ψ̃_goal, as, ts, sys, integrator=exp) ≈ 1
+    @test rollout_fidelity(ψ, ψ_goal, as, ts, sys, integrator=exp) ≈ 1
 
     # Unitary fidelity
-    @test unitary_fidelity(U_goal, as, ts, sys, integrator=exp) ≈ 1
-    @test unitary_fidelity(prob.trajectory, sys, integrator=exp) ≈ 1
-    @test unitary_fidelity(prob, integrator=exp) ≈ 1
-    @test unitary_fidelity(embedded_U_goal, as, ts, sys, integrator=exp) ≈ 1
+    @test unitary_rollout_fidelity(U_goal, as, ts, sys, integrator=exp) ≈ 1
+    @test unitary_rollout_fidelity(prob.trajectory, sys, integrator=exp) ≈ 1
+    @test unitary_rollout_fidelity(embedded_U_goal, as, ts, sys, integrator=exp) ≈ 1
 
     # Bad integrator
-    @test_throws ErrorException unitary_fidelity(U_goal, as, ts, sys, integrator=(a,b) -> 1) ≈ 1
+    @test_throws ErrorException unitary_rollout_fidelity(U_goal, as, ts, sys, integrator=(a,b) -> 1) ≈ 1
 end
 
 @testitem "Foward diff rollout" begin
@@ -540,13 +506,13 @@ end
         as -> rollout(ψ, as, ts, sys, integrator=expv)[:, end], as
     )
     iso_ket_dim = length(ket_to_iso(ψ))
-    @test size(result1) == (iso_ket_dim, T * length(sys.H_drives))
+    @test size(result1) == (iso_ket_dim, T * sys.n_drives)
 
     result2 = ForwardDiff.jacobian(
         as -> unitary_rollout(as, ts, sys, integrator=expv)[:, end], as
     )
-    iso_vec_dim = length(operator_to_iso_vec(sys.H_drift))
-    @test size(result2) == (iso_vec_dim, T * length(sys.H_drives))
+    iso_vec_dim = length(operator_to_iso_vec(sys.H(zeros(sys.n_drives))))
+    @test size(result2) == (iso_vec_dim, T * sys.n_drives)
 
     # Time derivatives
     ψ = ComplexF64[1, 0]
@@ -559,7 +525,7 @@ end
     result2 = ForwardDiff.jacobian(
         ts -> unitary_rollout(as, ts, sys, integrator=expv)[:, end], ts
     )
-    iso_vec_dim = length(operator_to_iso_vec(sys.H_drift))
+    iso_vec_dim = length(operator_to_iso_vec(sys.H(zeros(sys.n_drives))))
     @test size(result2) == (iso_vec_dim, T)
 end
 end
diff --git a/src/save_load_utils.jl b/src/save_load_utils.jl
deleted file mode 100644
index dbabe292..00000000
--- a/src/save_load_utils.jl
+++ /dev/null
@@ -1,154 +0,0 @@
-module SaveLoadUtils
-
-export save_problem
-export load_problem
-export generate_file_path
-
-using ..Constraints
-using ..Objectives
-using ..Integrators
-using ..Problems
-
-using JLD2
-
-
-function save_problem(
-    path::String,
-    prob::QuantumControlProblem,
-    info::Dict{String,<:Any}=Dict{String, Any}()
-)
-    mkpath(dirname(path))
-
-    data = Dict(
-        "system" => prob.system,
-        "trajectory" => prob.trajectory,
-        "integrators" => prob.integrators,
-        "options" => prob.options,
-        "params" => prob.params,
-    )
-
-    # assert none of the keys in info are already in data
-    for key in keys(info)
-        if haskey(data, key)
-            @warn "Key $(key) in info exists in data dict, removing"
-            delete!(info, key)
-        end
-    end
-
-    merge!(data, info)
-
-    save(path, data)
-end
-
-const RESERVED_KEYS = ["system", "trajectory", "options", "params", "integrators"]
-
-function load_problem(path::String; verbose=true, return_data=false, kwargs...)
-    data = load(path)
-
-    if verbose
-        println("Loading $(return_data ? "data dict" : "problem") from $(path):\n")
-        for (key, value) ∈ data
-            if key ∉ RESERVED_KEYS
-                println("   $(key) = $(value)")
-            end
-        end
-    end
-
-    if return_data
-        return data
-    else
-        if isnothing(data["integrators"])
-            @warn "Dynamics was built using a user defined function, which could not be saved: returning data dict instead of problem (keys = [\"trajectory\", \"system\", \"params\"]) "
-            return data
-        end
-
-        system = data["system"]
-        delete!(data, "system")
-
-        trajectory = data["trajectory"]
-        delete!(data, "trajectory")
-
-        options = data["options"]
-        delete!(data, "options")
-
-        params = data["params"]
-        delete!(data, "params")
-
-        integrators = data["integrators"]
-        delete!(data, "integrators")
-
-        objective = Objective(params[:objective_terms])
-        delete!(params, :objective_terms)
-
-        linear_constraints = params[:linear_constraints]
-        delete!(params, :linear_constraints)
-
-        nonlinear_constraints = NonlinearConstraint.(params[:nonlinear_constraints])
-        delete!(params, :nonlinear_constraints)
-
-        constraints = AbstractConstraint[linear_constraints; nonlinear_constraints]
-
-        prob = QuantumControlProblem(
-            system,
-            trajectory,
-            objective,
-            integrators;
-            constraints=constraints,
-            options=options,
-            verbose=verbose,
-            build_trajectory_constraints=false,
-            params...,
-            kwargs...
-        )
-
-        return prob
-    end
-end
-
-function save_h5(prob::QuantumControlProblem, save_path::String; verbose=true)
-    traj = prob.trajectory
-
-    result = Dict(
-
-    )
-end
-
-function generate_file_path(extension, file_name, path)
-    # Ensure the path exists.
-    mkpath(path)
-
-    # remove dot from extension
-    extension = split(extension, ".")[end]
-
-    # Create a save file name based on the one given; ensure it will
-    # not conflict with others in the directory.
-    max_numeric_suffix = -1
-    for (_, _, files) in walkdir(path)
-        for file_name_ in files
-            if occursin("$(file_name)", file_name_) && occursin(".$(extension)", file_name_)
-
-                numeric_suffix = parse(
-                    Int,
-                    split(split(file_name_, "_")[end], ".")[1]
-                )
-
-                max_numeric_suffix = max(
-                    numeric_suffix,
-                    max_numeric_suffix
-                )
-            end
-        end
-    end
-
-    file_path = joinpath(
-        path,
-        file_name *
-        "_$(lpad(max_numeric_suffix + 1, 5, '0')).$(extension)"
-    )
-
-    return file_path
-end
-
-
-
-end
diff --git a/src/structure_utils.jl b/src/structure_utils.jl
deleted file mode 100644
index e7dae291..00000000
--- a/src/structure_utils.jl
+++ /dev/null
@@ -1,258 +0,0 @@
-module StructureUtils
-
-export upper_half_vals
-export random_sparse_symbolics_matrix
-
-export structure
-
-export jacobian_structure
-export hessian_of_lagrangian_structure
-
-export dynamics_jacobian_structure
-export dynamics_hessian_of_lagrangian_structure
-export dynamics_structure
-
-export sparsity_structure
-
-export loss_hessian_structure
-
-using NamedTrajectories
-using TrajectoryIndexingUtils
-using LinearAlgebra
-using SparseArrays
-using Symbolics
-using ForwardDiff
-using Einsum
-
-function sparsity_structure(S::SparseMatrixCSC; return_pairs=false, type::Type=Float64)
-    Is, Js = findnz(S)
-    if return_pairs
-        return zip(Is, Js)
-    else
-        return sparse(Is, Js, ones(type, nnz(S)), size(S)...)
-    end
-end
-
-function sparsity_structure(A::AbstractMatrix; kwargs...)
-    return sparsity_structure(sparse(A); kwargs...)
-end
-
-
-
-function upper_half_vals(A::AbstractMatrix)
-    n = size(A, 1)
-    vals = similar(A, n * (n + 1) ÷ 2)
-    k = 1
-    for j ∈ axes(A, 2)
-        for i = 1:j
-            vals[k] = A[i, j]
-            k += 1
-        end
-    end
-    return vals
-end
-
-# create an m x n sparse matrix filled with l symbolics num variables
-function random_sparse_symbolics_matrix(m, n, l)
-    A = zeros(Symbolics.Num, m * n)
-    xs = collect(Symbolics.@variables(x[1:l])...)
-    rands = randperm(m * n)[1:l]
-    for i ∈ 1:l
-        A[rands[i]] = xs[i]
-    end
-    return sparse(reshape(A, m, n))
-end
-
-function structure(A::SparseMatrixCSC; upper_half=false)
-    I, J, _ = findnz(A)
-    index_pairs = collect(zip(I, J))
-    if upper_half
-        @assert size(A, 1) == size(A, 2)
-        index_pairs = filter(p -> p[1] <= p[2], index_pairs)
-    end
-    return index_pairs
-end
-
-function jacobian_structure(∂f::Function, xdim::Int)
-    x = collect(Symbolics.@variables(x[1:xdim])...)
-    return structure(sparse(∂f(x)))
-end
-
-dynamics_jacobian_structure(∂f::Function, zdim::Int) = jacobian_structure(∂f, 2zdim)
-
-function hessian_of_lagrangian_structure(∂²f::Function, xdim::Int, μdim::Int)
-    x = collect(Symbolics.@variables(x[1:xdim])...)
-    μ = collect(Symbolics.@variables(μ[1:μdim])...)
-    ∂²f_x = ∂²f(x)
-    if length(size(∂²f_x)) == 3
-        @einsum μ∂²f[j, k] := μ[i] * ∂²f_x[i, j, k]
-    elseif length(size(∂²f_x)) == 2
-        μ∂²f = μ[1] * ∂²f_x
-    else
-        error("hessian of lagrangian must be 2 or 3 dimensional")
-    end
-    return structure(sparse(μ∂²f); upper_half=true)
-end
-
-dynamics_hessian_of_lagrangian_structure(∂²f::Function, zdim::Int, μdim::Int) =
-    hessian_of_lagrangian_structure(∂²f, 2zdim, μdim)
-
-# function dynamics_structure(∂f̂::Function, traj::NamedTrajectory, dynamics_dim::Int)
-#     ∂²f̂(zz) = reshape(
-#         ForwardDiff.jacobian(x -> vec(∂f̂(x)), zz),
-#         traj.dims.states,
-#         2traj.dim,
-#         2traj.dim
-#     )
-
-#     ∂f_structure = dynamics_jacobian_structure(∂f̂, traj.dim)
-
-#     ∂F_structure = Tuple{Int,Int}[]
-
-#     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
-#         ]
-#         append!(∂F_structure, ∂fₜ_structure)
-#     end
-
-#     μ∂²f_structure =
-#         dynamics_hessian_of_lagrangian_structure(∂²f̂, traj.dim, dynamics_dim)
-
-#     μ∂²F_structure = Tuple{Int,Int}[]
-
-#     for t = 1:traj.T-1
-#         μ∂²fₜ_structure = [ij .+ index(t, 0, traj.dim) for ij ∈ μ∂²f_structure]
-#         append!(μ∂²F_structure, μ∂²fₜ_structure)
-#     end
-
-#     return ∂f_structure, ∂F_structure, μ∂²f_structure, μ∂²F_structure
-# end
-
-function dynamics_structure(
-    ∂f::Function,
-    μ∂²f::Function,
-    traj::NamedTrajectory,
-    dynamics_dim::Int;
-    verbose=false,
-    jacobian=true,
-    hessian=true
-)
-    @assert hessian ? jacobian : true "if hessian is true, jacobian must be true"
-
-    # getting inter knot point structure
-    if jacobian
-        if verbose
-            println("            computing jacobian sparsity...")
-        end
-        # getting symbolic variables
-        z1 = collect(Symbolics.@variables(z[1:traj.dim])...)
-        z2 = collect(Symbolics.@variables(z[1:traj.dim])...)
-        ∂f_structure = structure(sparse(∂f(z1, z2)))
-    else
-        if verbose
-            println("            computing full jacobian block...")
-        end
-        ∂f_structure = [(i, j) for i = 1:dynamics_dim, j = 1:2traj.dim]
-    end
-
-    if hessian
-        if verbose
-            println("            computing hessian sparsity...")
-        end
-        μ = collect(Symbolics.@variables(μ[1:dynamics_dim])...)
-        μ∂²f_structure = structure(sparse(μ∂²f(z1, z2, μ)); upper_half=true)
-    else
-        if verbose
-            println("            computing full hessian block...")
-        end
-        μ∂²f_structure = [(i, j) for i = 1:2traj.dim, j = 1:2traj.dim if j ≥ i]
-    end
-
-    if verbose
-        println("            creating full trajectory jacobian structure...")
-    end
-    ∂F_structure = Tuple{Int,Int}[]
-    ∂F_structure = Vector{Tuple{Int,Int}}(undef, length(∂f_structure) * (traj.T - 1))
-
-    @views Threads.@threads 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_structure[slice(t, length(∂f_structure))] .= ∂fₜ_structure
-    end
-
-    μ∂²F_structure = Tuple{Int,Int}[]
-
-    if verbose
-        println("            creating full trajectory hessian structure...")
-    end
-
-    μ∂²F_structure = Vector{Tuple{Int,Int}}(undef, length(μ∂²f_structure) * (traj.T - 1))
-
-    @views Threads.@threads for t = 1:traj.T-1
-        μ∂²fₜ_structure = [ij .+ index(t, 0, traj.dim) for ij ∈ μ∂²f_structure]
-        μ∂²F_structure[slice(t, length(μ∂²f_structure))] .= μ∂²fₜ_structure
-    end
-
-    return ∂f_structure, ∂F_structure, μ∂²f_structure, μ∂²F_structure
-end
-
-
-function dynamics_structure(
-    ∂f::Function,
-    traj::NamedTrajectory,
-    dynamics_dim::Int;
-    verbose=false,
-    jacobian=true,
-)
-    # getting inter knot point structure
-    if jacobian
-        if verbose
-            println("            computing jacobian sparsity...")
-        end
-        # getting symbolic variables
-        z1 = collect(Symbolics.@variables(z[1:traj.dim])...)
-        z2 = collect(Symbolics.@variables(z[1:traj.dim])...)
-        ∂f_structure = structure(sparse(∂f(z1, z2, 1)))
-    else
-        if verbose
-            println("            computing full jacobian block...")
-        end
-        ∂f_structure = [(i, j) for i = 1:dynamics_dim, j = 1:2traj.dim]
-    end
-
-    if verbose
-        println("            creating full trajectory jacobian structure...")
-    end
-
-    length_∂f_structure = length(∂f_structure)
-
-    ∂F_structure = Vector{Tuple{Int,Int}}(undef, length_∂f_structure * (traj.T - 1))
-
-    Threads.@threads 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_structure[slice(t, length_∂f_structure)] = ∂fₜ_structure
-    end
-
-    return ∂f_structure, ∂F_structure
-end
-
-function loss_hessian_structure(∂²l::Function, xdim::Int)
-    x = collect(Symbolics.@variables(x[1:xdim])...)
-    ∂²l_x = ∂²l(x)
-    return structure(sparse(∂²l_x); upper_half=true)
-end
-
-end
diff --git a/src/trajectory_initialization.jl b/src/trajectory_initialization.jl
index 6bed9de7..fdf7227d 100644
--- a/src/trajectory_initialization.jl
+++ b/src/trajectory_initialization.jl
@@ -108,7 +108,7 @@ function unitary_geodesic(
 end
 
 function unitary_geodesic(
-    U_goal::OperatorType,
+    U_goal::AbstractPiccoloOperator,
     samples::Int;
     kwargs...
 )
@@ -175,7 +175,7 @@ const ScalarBound = Union{R, Tuple{R, R}} where R <: Real
 
 function initialize_unitary_trajectory(
     U_init::AbstractMatrix{<:Number},
-    U_goal::OperatorType,
+    U_goal::AbstractPiccoloOperator,
     T::Int;
     geodesic=true
 )
@@ -270,6 +270,7 @@ function initialize_trajectory(
     free_time=false,
     control_name=:a,
     n_control_derivatives::Int=length(control_bounds) - 1,
+    zero_initial_and_final_derivative=false,
     timestep_name=:Δt,
     Δt_bounds::ScalarBound=(0.5 * Δt, 1.5 * Δt),
     drive_derivative_σ::Float64=0.1,
@@ -321,6 +322,11 @@ function initialize_trajectory(
         control_name => zeros(n_drives),
     )
 
+    if zero_initial_and_final_derivative
+        initial = merge(initial, (; control_derivative_names[1] => zeros(n_drives),))
+        final = merge(final, (; control_derivative_names[1] => zeros(n_drives),))
+    end
+
     goal = (; (state_names .=> state_goals)...)
 
     # Bounds
@@ -381,7 +387,7 @@ end
 Trajectory initialization of unitaries.
 """
 function initialize_trajectory(
-    U_goal::OperatorType,
+    U_goal::AbstractPiccoloOperator,
     T::Int,
     Δt::Union{Real, AbstractVecOrMat{<:Real}},
     args...;