diff --git a/Manifest.toml b/Manifest.toml
index 021c7ea7..e2dfc214 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -2,12 +2,12 @@
 
 julia_version = "1.10.5"
 manifest_format = "2.0"
-project_hash = "0e1b56171c4e825157ede62be9481fa3374601e7"
+project_hash = "a8d39c98fbf089d5b72e82a89504c42d025cd3fe"
 
 [[deps.ADTypes]]
-git-tree-sha1 = "99a6f5d0ce1c7c6afdb759daa30226f71c54f6b0"
+git-tree-sha1 = "eea5d80188827b35333801ef97a40c2ed653b081"
 uuid = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
-version = "1.7.1"
+version = "1.9.0"
 
     [deps.ADTypes.extensions]
     ADTypesChainRulesCoreExt = "ChainRulesCore"
@@ -63,24 +63,28 @@ uuid = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 version = "0.4.5"
 
 [[deps.Accessors]]
-deps = ["CompositionsBase", "ConstructionBase", "Dates", "InverseFunctions", "LinearAlgebra", "MacroTools", "Markdown", "Test"]
-git-tree-sha1 = "f61b15be1d76846c0ce31d3fcfac5380ae53db6a"
+deps = ["CompositionsBase", "ConstructionBase", "InverseFunctions", "LinearAlgebra", "MacroTools", "Markdown"]
+git-tree-sha1 = "b392ede862e506d451fc1616e79aa6f4c673dab8"
 uuid = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
-version = "0.1.37"
+version = "0.1.38"
 
     [deps.Accessors.extensions]
     AccessorsAxisKeysExt = "AxisKeys"
+    AccessorsDatesExt = "Dates"
     AccessorsIntervalSetsExt = "IntervalSets"
     AccessorsStaticArraysExt = "StaticArrays"
     AccessorsStructArraysExt = "StructArrays"
+    AccessorsTestExt = "Test"
     AccessorsUnitfulExt = "Unitful"
 
     [deps.Accessors.weakdeps]
     AxisKeys = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5"
+    Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
     IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
     Requires = "ae029012-a4dd-5104-9daa-d747884805df"
     StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
     StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
+    Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
     Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [[deps.Adapt]]
@@ -94,9 +98,9 @@ weakdeps = ["StaticArrays"]
     AdaptStaticArraysExt = "StaticArrays"
 
 [[deps.AdaptivePredicates]]
-git-tree-sha1 = "7d5da5dd472490d048b081ca1bda4a7821b06456"
+git-tree-sha1 = "7e651ea8d262d2d74ce75fdf47c4d63c07dba7a6"
 uuid = "35492f91-a3bd-45ad-95db-fcad7dcfedb7"
-version = "1.1.1"
+version = "1.2.0"
 
 [[deps.AliasTables]]
 deps = ["PtrArrays", "Random"]
@@ -207,21 +211,21 @@ version = "1.1.0"
 
 [[deps.CairoMakie]]
 deps = ["CRC32c", "Cairo", "Cairo_jll", "Colors", "FileIO", "FreeType", "GeometryBasics", "LinearAlgebra", "Makie", "PrecompileTools"]
-git-tree-sha1 = "361dec06290d76b6d70d0c7dc888038eec9df63a"
+git-tree-sha1 = "2b04b60ed9d3e977f93e34952971b608c34b3401"
 uuid = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
-version = "0.12.9"
+version = "0.12.13"
 
 [[deps.Cairo_jll]]
 deps = ["Artifacts", "Bzip2_jll", "CompilerSupportLibraries_jll", "Fontconfig_jll", "FreeType2_jll", "Glib_jll", "JLLWrappers", "LZO_jll", "Libdl", "Pixman_jll", "Xorg_libXext_jll", "Xorg_libXrender_jll", "Zlib_jll", "libpng_jll"]
-git-tree-sha1 = "a2f1c8c668c8e3cb4cca4e57a8efdb09067bb3fd"
+git-tree-sha1 = "009060c9a6168704143100f36ab08f06c2af4642"
 uuid = "83423d85-b0ee-5818-9007-b63ccbeb887a"
-version = "1.18.0+2"
+version = "1.18.2+1"
 
 [[deps.ChainRulesCore]]
 deps = ["Compat", "LinearAlgebra"]
-git-tree-sha1 = "71acdbf594aab5bbb2cec89b208c41b4c411e49f"
+git-tree-sha1 = "3e4b134270b372f2ed4d4d0e936aabaefc1802bc"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "1.24.0"
+version = "1.25.0"
 weakdeps = ["SparseArrays"]
 
     [deps.ChainRulesCore.extensions]
@@ -366,10 +370,10 @@ deps = ["Printf"]
 uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"
 
 [[deps.DelaunayTriangulation]]
-deps = ["AdaptivePredicates", "EnumX", "ExactPredicates", "Random"]
-git-tree-sha1 = "9903123ab7fc5e55053292aff04ff5d7aff92633"
+deps = ["AdaptivePredicates", "EnumX", "ExactPredicates", "PrecompileTools", "Random"]
+git-tree-sha1 = "668bb97ea6df5e654e6288d87d2243591fe68665"
 uuid = "927a84f5-c5f4-47a5-9785-b46e178433df"
-version = "1.3.0"
+version = "1.6.0"
 
 [[deps.DiffResults]]
 deps = ["StaticArraysCore"]
@@ -389,9 +393,9 @@ uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 
 [[deps.Distributions]]
 deps = ["AliasTables", "FillArrays", "LinearAlgebra", "PDMats", "Printf", "QuadGK", "Random", "SpecialFunctions", "Statistics", "StatsAPI", "StatsBase", "StatsFuns"]
-git-tree-sha1 = "e6c693a0e4394f8fda0e51a5bdf5aef26f8235e9"
+git-tree-sha1 = "d7477ecdafb813ddee2ae727afa94e9dcb5f3fb0"
 uuid = "31c24e10-a181-5473-b8eb-7969acd0382f"
-version = "0.25.111"
+version = "0.25.112"
 
     [deps.Distributions.extensions]
     DistributionsChainRulesCoreExt = "ChainRulesCore"
@@ -495,9 +499,9 @@ version = "1.8.0"
 
 [[deps.FFTW_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
-git-tree-sha1 = "c6033cc3892d0ef5bb9cd29b7f2f0331ea5184ea"
+git-tree-sha1 = "4d81ed14783ec49ce9f2e168208a12ce1815aa25"
 uuid = "f5851436-0d7a-5f13-b9de-f02708fd171a"
-version = "3.3.10+0"
+version = "3.3.10+1"
 
 [[deps.FileIO]]
 deps = ["Pkg", "Requires", "UUIDs"]
@@ -616,9 +620,9 @@ version = "0.4.2"
 
 [[deps.GeoInterface]]
 deps = ["Extents", "GeoFormatTypes"]
-git-tree-sha1 = "5921fc0704e40c024571eca551800c699f86ceb4"
+git-tree-sha1 = "2f6fce56cdb8373637a6614e14a5768a88450de2"
 uuid = "cf35fbd7-0cd7-5166-be24-54bfbe79505f"
-version = "1.3.6"
+version = "1.3.7"
 
 [[deps.GeometryBasics]]
 deps = ["EarCut_jll", "Extents", "GeoInterface", "IterTools", "LinearAlgebra", "StaticArrays", "StructArrays", "Tables"]
@@ -634,9 +638,9 @@ version = "0.21.0+0"
 
 [[deps.Glib_jll]]
 deps = ["Artifacts", "Gettext_jll", "JLLWrappers", "Libdl", "Libffi_jll", "Libiconv_jll", "Libmount_jll", "PCRE2_jll", "Zlib_jll"]
-git-tree-sha1 = "7c82e6a6cd34e9d935e9aa4051b66c6ff3af59ba"
+git-tree-sha1 = "674ff0db93fffcd11a3573986e550d66cd4fd71f"
 uuid = "7746bdde-850d-59dc-9ae8-88ece973131d"
-version = "2.80.2+0"
+version = "2.80.5+0"
 
 [[deps.Graphics]]
 deps = ["Colors", "LinearAlgebra", "NaNMath"]
@@ -669,9 +673,9 @@ version = "8.3.1+0"
 
 [[deps.Hwloc_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl"]
-git-tree-sha1 = "5e19e1e4fa3e71b774ce746274364aef0234634e"
+git-tree-sha1 = "dd3b49277ec2bb2c6b94eb1604d4d0616016f7a6"
 uuid = "e33a78d0-f292-5ffc-b300-72abe9b543c8"
-version = "2.11.1+0"
+version = "2.11.2+0"
 
 [[deps.HypergeometricFunctions]]
 deps = ["LinearAlgebra", "OpenLibm_jll", "SpecialFunctions"]
@@ -763,9 +767,9 @@ weakdeps = ["Unitful"]
 
 [[deps.IntervalArithmetic]]
 deps = ["CRlibm_jll", "MacroTools", "RoundingEmulator"]
-git-tree-sha1 = "fe30dec78e68f27fc416901629c6e24e9d5f057b"
+git-tree-sha1 = "8e125d40cae3a9f4276cdfeb4fcdb1828888a4b3"
 uuid = "d1acc4aa-44c8-5952-acd4-ba5d80a2a253"
-version = "0.22.16"
+version = "0.22.17"
 weakdeps = ["DiffRules", "ForwardDiff", "IntervalSets", "LinearAlgebra", "RecipesBase"]
 
     [deps.IntervalArithmetic.extensions]
@@ -787,9 +791,9 @@ weakdeps = ["Random", "RecipesBase", "Statistics"]
     IntervalSetsStatisticsExt = "Statistics"
 
 [[deps.InverseFunctions]]
-git-tree-sha1 = "2787db24f4e03daf859c6509ff87764e4182f7d1"
+git-tree-sha1 = "a779299d77cd080bf77b97535acecd73e1c5e5cb"
 uuid = "3587e190-3f89-42d0-90ee-14403ec27112"
-version = "0.1.16"
+version = "0.1.17"
 weakdeps = ["Dates", "Test"]
 
     [deps.InverseFunctions.extensions]
@@ -831,9 +835,9 @@ version = "1.0.0"
 
 [[deps.JLD2]]
 deps = ["FileIO", "MacroTools", "Mmap", "OrderedCollections", "PrecompileTools", "Requires", "TranscodingStreams"]
-git-tree-sha1 = "1dbe274ad1c199490b19a153b60a9e46c5386f8d"
+git-tree-sha1 = "aeab5c68eb2cf326619bf71235d8f4561c62fe22"
 uuid = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
-version = "0.5.1"
+version = "0.5.5"
 
 [[deps.JLLWrappers]]
 deps = ["Artifacts", "Preferences"]
@@ -855,9 +859,9 @@ version = "0.1.5"
 
 [[deps.JpegTurbo_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl"]
-git-tree-sha1 = "c84a835e1a09b289ffcd2271bf2a337bbdda6637"
+git-tree-sha1 = "25ee0be4d43d0269027024d75a24c24d6c6e590c"
 uuid = "aacddb02-875f-59d6-b918-886e6ef4fbf8"
-version = "3.0.3+0"
+version = "3.0.4+0"
 
 [[deps.KernelDensity]]
 deps = ["Distributions", "DocStringExtensions", "FFTW", "Interpolations", "StatsBase"]
@@ -873,15 +877,15 @@ version = "3.100.2+0"
 
 [[deps.LLVMOpenMP_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl"]
-git-tree-sha1 = "e16271d212accd09d52ee0ae98956b8a05c4b626"
+git-tree-sha1 = "78211fb6cbc872f77cad3fc0b6cf647d923f4929"
 uuid = "1d63c593-3942-5779-bab2-d838dc0a180e"
-version = "17.0.6+0"
+version = "18.1.7+0"
 
 [[deps.LZO_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl"]
-git-tree-sha1 = "70c5da094887fd2cae843b8db33920bac4b6f07d"
+git-tree-sha1 = "854a9c268c43b77b0a27f22d7fab8d33cdb3a731"
 uuid = "dd4b983a-f0e5-5f8d-a1b7-129d4a5fb1ac"
-version = "2.10.2+0"
+version = "2.10.2+1"
 
 [[deps.LaTeXStrings]]
 git-tree-sha1 = "50901ebc375ed41dbf8058da26f9de442febbbec"
@@ -1034,16 +1038,16 @@ uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 version = "0.5.13"
 
 [[deps.Makie]]
-deps = ["Animations", "Base64", "CRC32c", "ColorBrewer", "ColorSchemes", "ColorTypes", "Colors", "Contour", "Dates", "DelaunayTriangulation", "Distributions", "DocStringExtensions", "Downloads", "FFMPEG_jll", "FileIO", "FilePaths", "FixedPointNumbers", "Format", "FreeType", "FreeTypeAbstraction", "GeometryBasics", "GridLayoutBase", "ImageIO", "InteractiveUtils", "IntervalSets", "Isoband", "KernelDensity", "LaTeXStrings", "LinearAlgebra", "MacroTools", "MakieCore", "Markdown", "MathTeXEngine", "Observables", "OffsetArrays", "Packing", "PlotUtils", "PolygonOps", "PrecompileTools", "Printf", "REPL", "Random", "RelocatableFolders", "Scratch", "ShaderAbstractions", "Showoff", "SignedDistanceFields", "SparseArrays", "Statistics", "StatsBase", "StatsFuns", "StructArrays", "TriplotBase", "UnicodeFun", "Unitful"]
-git-tree-sha1 = "204f06860af9008fa08b3a4842f48116e1209a2c"
+deps = ["Animations", "Base64", "CRC32c", "ColorBrewer", "ColorSchemes", "ColorTypes", "Colors", "Contour", "Dates", "DelaunayTriangulation", "Distributions", "DocStringExtensions", "Downloads", "FFMPEG_jll", "FileIO", "FilePaths", "FixedPointNumbers", "Format", "FreeType", "FreeTypeAbstraction", "GeometryBasics", "GridLayoutBase", "ImageBase", "ImageIO", "InteractiveUtils", "Interpolations", "IntervalSets", "Isoband", "KernelDensity", "LaTeXStrings", "LinearAlgebra", "MacroTools", "MakieCore", "Markdown", "MathTeXEngine", "Observables", "OffsetArrays", "Packing", "PlotUtils", "PolygonOps", "PrecompileTools", "Printf", "REPL", "Random", "RelocatableFolders", "Scratch", "ShaderAbstractions", "Showoff", "SignedDistanceFields", "SparseArrays", "Statistics", "StatsBase", "StatsFuns", "StructArrays", "TriplotBase", "UnicodeFun", "Unitful"]
+git-tree-sha1 = "50ebda951efaa11b6db0413c1128726b8eab3bf0"
 uuid = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
-version = "0.21.9"
+version = "0.21.13"
 
 [[deps.MakieCore]]
 deps = ["ColorTypes", "GeometryBasics", "IntervalSets", "Observables"]
-git-tree-sha1 = "b0e2e3473af351011e598f9219afb521121edd2b"
+git-tree-sha1 = "4604f03e5b057e8e62a95a44929cafc9585b0fe9"
 uuid = "20f20a25-4f0e-4fdf-b5d1-57303727442b"
-version = "0.8.6"
+version = "0.8.9"
 
 [[deps.MappedArrays]]
 git-tree-sha1 = "2dab0221fe2b0f2cb6754eaa743cc266339f527e"
@@ -1104,9 +1108,9 @@ version = "0.5.6"
 
 [[deps.MutableArithmetics]]
 deps = ["LinearAlgebra", "SparseArrays", "Test"]
-git-tree-sha1 = "d0a6b1096b584a2b88efb70a92f8cb8c881eb38a"
+git-tree-sha1 = "3eba928678787843e504c153a9b8e80d7d73ab17"
 uuid = "d8a4904e-b15c-11e9-3269-09a3773c0cb0"
-version = "1.4.6"
+version = "1.5.0"
 
 [[deps.NaNMath]]
 deps = ["OpenLibm_jll"]
@@ -1180,9 +1184,9 @@ version = "0.8.1+2"
 
 [[deps.OpenSSL_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl"]
-git-tree-sha1 = "1b35263570443fdd9e76c76b7062116e2f374ab8"
+git-tree-sha1 = "7493f61f55a6cce7325f197443aa80d32554ba10"
 uuid = "458c3c95-2e84-50aa-8efc-19380b2a3a95"
-version = "3.0.15+0"
+version = "3.0.15+1"
 
 [[deps.OpenSpecFun_jll]]
 deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
@@ -1315,9 +1319,9 @@ version = "1.0.0"
 
 [[deps.QuadGK]]
 deps = ["DataStructures", "LinearAlgebra"]
-git-tree-sha1 = "1d587203cf851a51bf1ea31ad7ff89eff8d625ea"
+git-tree-sha1 = "cda3b045cf9ef07a08ad46731f5a3165e56cf3da"
 uuid = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
-version = "2.11.0"
+version = "2.11.1"
 
     [deps.QuadGK.extensions]
     QuadGKEnzymeExt = "Enzyme"
@@ -1399,15 +1403,15 @@ version = "1.3.0"
 
 [[deps.Rmath]]
 deps = ["Random", "Rmath_jll"]
-git-tree-sha1 = "f65dcb5fa46aee0cf9ed6274ccbd597adc49aa7b"
+git-tree-sha1 = "852bd0f55565a9e973fcfee83a84413270224dc4"
 uuid = "79098fc4-a85e-5d69-aa6a-4863f24498fa"
-version = "0.7.1"
+version = "0.8.0"
 
 [[deps.Rmath_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl"]
-git-tree-sha1 = "e60724fd3beea548353984dc61c943ecddb0e29a"
+git-tree-sha1 = "58cdd8fb2201a6267e1db87ff148dd6c1dbd8ad8"
 uuid = "f50d1b31-88e8-58de-be2c-1cc44531875f"
-version = "0.4.3+0"
+version = "0.5.1+0"
 
 [[deps.RoundingEmulator]]
 git-tree-sha1 = "40b9edad2e5287e05bd413a38f61a8ff55b9557b"
@@ -1426,9 +1430,9 @@ version = "0.7.0"
 
 [[deps.SIMD]]
 deps = ["PrecompileTools"]
-git-tree-sha1 = "2803cab51702db743f3fda07dd1745aadfbf43bd"
+git-tree-sha1 = "98ca7c29edd6fc79cd74c61accb7010a4e7aee33"
 uuid = "fdea26ae-647d-5447-a871-4b548cad5224"
-version = "3.5.0"
+version = "3.6.0"
 
 [[deps.SPRAL_jll]]
 deps = ["Artifacts", "CompilerSupportLibraries_jll", "Hwloc_jll", "JLLWrappers", "Libdl", "METIS_jll", "libblastrampoline_jll"]
@@ -1438,9 +1442,9 @@ version = "2024.5.8+0"
 
 [[deps.SciMLBase]]
 deps = ["ADTypes", "Accessors", "ArrayInterface", "CommonSolve", "ConstructionBase", "Distributed", "DocStringExtensions", "EnumX", "Expronicon", "FunctionWrappersWrappers", "IteratorInterfaceExtensions", "LinearAlgebra", "Logging", "Markdown", "PrecompileTools", "Preferences", "Printf", "RecipesBase", "RecursiveArrayTools", "Reexport", "RuntimeGeneratedFunctions", "SciMLOperators", "SciMLStructures", "StaticArraysCore", "Statistics", "SymbolicIndexingInterface", "Tables"]
-git-tree-sha1 = "303a73db99326a8be43e695fbab9e076b02118ca"
+git-tree-sha1 = "d01eebc2dbd30c83a857ae8fcc7f12ea6bd5b10c"
 uuid = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
-version = "2.52.2"
+version = "2.56.0"
 
     [deps.SciMLBase.extensions]
     SciMLBaseChainRulesCoreExt = "ChainRulesCore"
@@ -1598,9 +1602,9 @@ version = "0.34.3"
 
 [[deps.StatsFuns]]
 deps = ["HypergeometricFunctions", "IrrationalConstants", "LogExpFunctions", "Reexport", "Rmath", "SpecialFunctions"]
-git-tree-sha1 = "cef0472124fab0695b58ca35a77c6fb942fdab8a"
+git-tree-sha1 = "b423576adc27097764a90e163157bcfc9acf0f46"
 uuid = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
-version = "1.3.1"
+version = "1.3.2"
 weakdeps = ["ChainRulesCore", "InverseFunctions"]
 
     [deps.StatsFuns.extensions]
@@ -1631,9 +1635,9 @@ version = "7.2.1+1"
 
 [[deps.SymbolicIndexingInterface]]
 deps = ["Accessors", "ArrayInterface", "RuntimeGeneratedFunctions", "StaticArraysCore"]
-git-tree-sha1 = "988e04b34a4c3b824fb656f542473df99a4f610d"
+git-tree-sha1 = "0225f7c62f5f78db35aae6abb2e5cabe38ce578f"
 uuid = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
-version = "0.3.30"
+version = "0.3.31"
 
 [[deps.SymbolicLimits]]
 deps = ["SymbolicUtils"]
@@ -1643,9 +1647,9 @@ version = "0.2.2"
 
 [[deps.SymbolicUtils]]
 deps = ["AbstractTrees", "ArrayInterface", "Bijections", "ChainRulesCore", "Combinatorics", "ConstructionBase", "DataStructures", "DocStringExtensions", "DynamicPolynomials", "IfElse", "LinearAlgebra", "MultivariatePolynomials", "NaNMath", "Setfield", "SparseArrays", "SpecialFunctions", "StaticArrays", "SymbolicIndexingInterface", "TermInterface", "TimerOutputs", "Unityper"]
-git-tree-sha1 = "9d983078d9e99421fcca44c373e4304b8421fdde"
+git-tree-sha1 = "3927e02dc7648a45ec6aa592bcd8374094a44740"
 uuid = "d1185830-fcd6-423d-90d6-eec64667417b"
-version = "3.6.0"
+version = "3.7.1"
 
     [deps.SymbolicUtils.extensions]
     SymbolicUtilsLabelledArraysExt = "LabelledArrays"
@@ -1657,9 +1661,9 @@ version = "3.6.0"
 
 [[deps.Symbolics]]
 deps = ["ADTypes", "ArrayInterface", "Bijections", "CommonWorldInvalidations", "ConstructionBase", "DataStructures", "DiffRules", "Distributions", "DocStringExtensions", "DomainSets", "DynamicPolynomials", "IfElse", "LaTeXStrings", "LambertW", "Latexify", "Libdl", "LinearAlgebra", "LogExpFunctions", "MacroTools", "Markdown", "NaNMath", "PrecompileTools", "Primes", "RecipesBase", "Reexport", "RuntimeGeneratedFunctions", "SciMLBase", "Setfield", "SparseArrays", "SpecialFunctions", "StaticArraysCore", "SymbolicIndexingInterface", "SymbolicLimits", "SymbolicUtils", "TermInterface"]
-git-tree-sha1 = "2226d810512c678d2ec9c2a9b2e227c2ebc43573"
+git-tree-sha1 = "6a7c7cd9bd8c051877a5a29fb597b18362dbc4e4"
 uuid = "0c5d862f-8b57-4792-8d23-62f2024744c7"
-version = "6.11.0"
+version = "6.13.1"
 
     [deps.Symbolics.extensions]
     SymbolicsForwardDiffExt = "ForwardDiff"
@@ -1743,9 +1747,9 @@ uuid = "6dad8b7f-dd9a-4c28-9b70-85b9a079bfc8"
 version = "0.1.0"
 
 [[deps.TranscodingStreams]]
-git-tree-sha1 = "e84b3a11b9bece70d14cce63406bbc79ed3464d2"
+git-tree-sha1 = "0c45878dcfdcfa8480052b6ab162cdd138781742"
 uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa"
-version = "0.11.2"
+version = "0.11.3"
 
 [[deps.TriplotBase]]
 git-tree-sha1 = "4d4ed7f294cda19382ff7de4c137d24d16adc89b"
@@ -1907,21 +1911,21 @@ version = "2.0.3+0"
 
 [[deps.libpng_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Zlib_jll"]
-git-tree-sha1 = "d7015d2e18a5fd9a4f47de711837e980519781a4"
+git-tree-sha1 = "b70c870239dc3d7bc094eb2d6be9b73d27bef280"
 uuid = "b53b4c65-9356-5827-b1ea-8c7a1a84506f"
-version = "1.6.43+1"
+version = "1.6.44+0"
 
 [[deps.libsixel_jll]]
 deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "Libdl", "Pkg", "libpng_jll"]
-git-tree-sha1 = "d4f63314c8aa1e48cd22aa0c17ed76cd1ae48c3c"
+git-tree-sha1 = "7dfa0fd9c783d3d0cc43ea1af53d69ba45c447df"
 uuid = "075b6546-f08a-558a-be8f-8157d0f608a5"
-version = "1.10.3+0"
+version = "1.10.3+1"
 
 [[deps.libsodium_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
-git-tree-sha1 = "848ab3d00fe39d6fbc2a8641048f8f272af1c51e"
+git-tree-sha1 = "f76d682d87eefadd3f165d8d9fda436464213142"
 uuid = "a9144af2-ca23-56d9-984f-0d03f7b5ccf8"
-version = "1.0.20+0"
+version = "1.0.20+1"
 
 [[deps.libvorbis_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Ogg_jll", "Pkg"]
diff --git a/Project.toml b/Project.toml
index 10bdbede..5feef8dc 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.2.2"
+version = "0.3.0"
 
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
diff --git a/docs/src/lib.md b/docs/src/lib.md
index b517605a..5b712e88 100644
--- a/docs/src/lib.md
+++ b/docs/src/lib.md
@@ -1,11 +1,21 @@
 #  Library
 
-## QuantumUtils
+## Problem Templates
 ```@autodocs
-Modules = [QuantumCollocation.QuantumUtils]
+Modules = [QuantumCollocation.ProblemTemplates]
 ```
 
-## QuantumSystems 
+## Direct Sums
+```@autodocs
+Modules = [QuantumCollocation.DirectSums]
+```
+
+## Quantum Object Utils
+```@autodocs
+Modules = [QuantumCollocation.QuantumObjectUtils]
+```
+
+## Quantum Systems 
 ```@autodocs
 Modules = [QuantumCollocation.QuantumSystems]
 ```
@@ -14,3 +24,63 @@ Modules = [QuantumCollocation.QuantumSystems]
 ```@autodocs
 Modules = [QuantumCollocation.Integrators]
 ``` 
+
+## Objectives
+```@autodocs
+Modules = [QuantumCollocation.Objectives]
+```
+
+## Losses
+```@autodocs
+Modules = [QuantumCollocation.Losses]
+```
+
+## Embedded Operators
+```@autodocs
+Modules = [QuantumCollocation.EmbeddedOperators]
+```
+
+## Isomorphisms
+```@autodocs
+Modules = [QuantumCollocation.Isomorphisms]
+```
+
+## Options
+```@autodocs
+Modules = [QuantumCollocation.Options]
+```
+
+## Plotting 
+```@autodocs
+Modules = [QuantumCollocation.Plotting]
+```
+
+## Problem Solvers
+```@autodocs
+Modules = [QuantumCollocation.ProblemSolvers]
+```
+
+## Rollouts
+```@autodocs
+Modules = [QuantumCollocation.Rollouts]
+```
+
+## Saving and Loading
+```@autodocs
+Modules = [QuantumCollocation.SaveLoadUtils]
+```
+
+## Structure Utils
+```@autodocs
+Modules = [QuantumCollocation.StructureUtils]
+```
+
+## Trajectory Initialization
+```@autodocs
+Modules = [QuantumCollocation.TrajectoryInitialization]
+```
+
+## Trajectory Interpolations
+```@autodocs
+Modules = [QuantumCollocation.TrajectoryInterpolations]
+```
\ No newline at end of file
diff --git a/src/QuantumCollocation.jl b/src/QuantumCollocation.jl
index d5ce2604..efaf4055 100644
--- a/src/QuantumCollocation.jl
+++ b/src/QuantumCollocation.jl
@@ -3,6 +3,9 @@ module QuantumCollocation
 using Reexport
 
 
+include("options.jl")
+@reexport using .Options
+
 include("isomorphisms.jl")
 @reexport using .Isomorphisms
 
@@ -42,9 +45,6 @@ include("dynamics.jl")
 include("evaluators.jl")
 @reexport using .Evaluators
 
-include("options.jl")
-@reexport using .Options
-
 include("problems.jl")
 @reexport using .Problems
 
diff --git a/src/constraints/_constraints.jl b/src/constraints/_constraints.jl
index a059ad88..8b17808e 100644
--- a/src/constraints/_constraints.jl
+++ b/src/constraints/_constraints.jl
@@ -3,6 +3,7 @@ module Constraints
 export AbstractConstraint
 
 export LinearConstraint
+
 export constrain!
 
 export NonlinearConstraint
@@ -12,6 +13,7 @@ export NonlinearInequalityConstraint
 using ..Losses
 using ..Isomorphisms
 using ..StructureUtils
+using ..Options
 
 using TrajectoryIndexingUtils
 using NamedTrajectories
@@ -126,4 +128,6 @@ struct NonlinearInequalityConstraint <: NonlinearConstraint
     params::Dict{Symbol, Any}
 end
 
+
+
 end
diff --git a/src/constraints/complex_modulus_constraint.jl b/src/constraints/complex_modulus_constraint.jl
index 9e0de5fc..6972edc3 100644
--- a/src/constraints/complex_modulus_constraint.jl
+++ b/src/constraints/complex_modulus_constraint.jl
@@ -20,13 +20,13 @@ function ComplexModulusContraint(;
     R::Union{Float64, Nothing}=nothing,
     comps::Union{AbstractVector{Int}, Nothing}=nothing,
     times::Union{AbstractVector{Int}, Nothing}=nothing,
-    dim::Union{Int, Nothing}=nothing,
+    zdim::Union{Int, Nothing}=nothing,
     T::Union{Int, Nothing}=nothing,
 )
     @assert !isnothing(R) "must provide a value R, s.t. |z| <= R"
     @assert !isnothing(comps) "must provide components of the complex number"
     @assert !isnothing(times) "must provide times"
-    @assert !isnothing(dim) "must provide a trajectory dimension"
+    @assert !isnothing(zdim) "must provide a trajectory dimension"
     @assert !isnothing(T) "must provide a T"
 
     @assert length(comps) == 2 "component must represent a complex number and have dimension 2"
@@ -37,7 +37,7 @@ function ComplexModulusContraint(;
     params[:R] = R
     params[:comps] = comps
     params[:times] = times
-    params[:dim] = dim
+    params[:zdim] = zdim
     params[:T] = T
 
     gₜ(xₜ, yₜ) = [R^2 - xₜ^2 - yₜ^2]
@@ -47,7 +47,7 @@ function ComplexModulusContraint(;
     @views function g(Z⃗)
         r = zeros(length(times))
         for (i, t) ∈ enumerate(times)
-            zₜ = Z⃗[slice(t, comps, dim)]
+            zₜ = Z⃗[slice(t, comps, zdim)]
             xₜ = zₜ[1]
             yₜ = zₜ[2]
             r[i] = gₜ(xₜ, yₜ)[1]
@@ -58,17 +58,17 @@ function ComplexModulusContraint(;
     ∂g_structure = []
 
     for (i, t) ∈ enumerate(times)
-        push!(∂g_structure, (i, index(t, comps[1], dim)))
-        push!(∂g_structure, (i, index(t, comps[2], dim)))
+        push!(∂g_structure, (i, index(t, comps[1], zdim)))
+        push!(∂g_structure, (i, index(t, comps[2], zdim)))
     end
 
     @views function ∂g(Z⃗; ipopt=true)
         ∂ = spzeros(length(times), length(Z⃗))
         for (i, t) ∈ enumerate(times)
-            zₜ = Z⃗[slice(t, comps, dim)]
+            zₜ = Z⃗[slice(t, comps, zdim)]
             xₜ = zₜ[1]
             yₜ = zₜ[2]
-            ∂[i, slice(t, comps, dim)] = ∂gₜ(xₜ, yₜ)
+            ∂[i, slice(t, comps, zdim)] = ∂gₜ(xₜ, yₜ)
         end
         if ipopt
             return [∂[i, j] for (i, j) in ∂g_structure]
@@ -83,15 +83,15 @@ function ComplexModulusContraint(;
         push!(
             μ∂²g_structure,
             (
-                index(t, comps[1], dim),
-                index(t, comps[1], dim)
+                index(t, comps[1], zdim),
+                index(t, comps[1], zdim)
             )
         )
         push!(
             μ∂²g_structure,
             (
-                index(t, comps[2], dim),
-                index(t, comps[2], dim)
+                index(t, comps[2], zdim),
+                index(t, comps[2], zdim)
             )
         )
     end
@@ -100,7 +100,7 @@ function ComplexModulusContraint(;
         n = length(Z⃗)
         μ∂² = spzeros(n, n)
         for (i, t) ∈ enumerate(times)
-            t_slice = slice(t, comps, dim)
+            t_slice = slice(t, comps, zdim)
             μ∂²[t_slice, t_slice] = μₜ∂²gₜ(μ[i])
         end
         if ipopt
@@ -143,4 +143,4 @@ function ComplexModulusContraint(
         zdim=traj.dim,
         T=traj.T
     )
-end
\ No newline at end of file
+end
diff --git a/src/constraints/fidelity_constraint.jl b/src/constraints/fidelity_constraint.jl
index 095b3627..2a495147 100644
--- a/src/constraints/fidelity_constraint.jl
+++ b/src/constraints/fidelity_constraint.jl
@@ -161,7 +161,6 @@ function FinalUnitaryFidelityConstraint(
     subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing,
     eval_hessian::Bool=true
 )
-    @assert statesymb ∈ traj.names
     return FinalFidelityConstraint(;
         fidelity_function=iso_vec_unitary_fidelity,
         value=val,
@@ -221,7 +220,7 @@ function FinalUnitaryFreePhaseFidelityConstraint(;
     state_slice::Union{AbstractVector{Int},Nothing}=nothing,
     phase_slice::Union{AbstractVector{Int},Nothing}=nothing,
     goal::Union{AbstractVector{Float64},Nothing}=nothing,
-    phase_operators::Union{AbstractVector{<:AbstractMatrix{<:Complex{Float64}}},Nothing}=nothing,
+    phase_operators::Union{AbstractVector{<:AbstractMatrix{<:Complex}},Nothing}=nothing,
     zdim::Union{Int,Nothing}=nothing,
     subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing,
     eval_hessian::Bool=false
@@ -284,4 +283,25 @@ function FinalUnitaryFreePhaseFidelityConstraint(;
         1,
         params
     )
+end
+
+function FinalUnitaryFreePhaseFidelityConstraint(
+    state_symbol::Symbol,
+    global_symbol::Symbol,
+    phase_operators::AbstractVector{<:AbstractMatrix{<:Complex}},
+    val::Float64,
+    traj::NamedTrajectory;
+    subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing,
+    eval_hessian::Bool=false
+)
+    return FinalUnitaryFreePhaseFidelityConstraint(;
+        value=val,
+        state_slice=slice(traj.T, traj.components[state_symbol], traj.dim),
+        phase_slice=trajectory.global_components[global_symbol],
+        goal=traj.goal[state_symbol],
+        phase_operators=phase_operators,
+        zdim=length(traj),
+        subspace=subspace,
+        eval_hessian=eval_hessian
+    )
 end
\ No newline at end of file
diff --git a/src/direct_sums.jl b/src/direct_sums.jl
index 8496fc94..47f0c441 100644
--- a/src/direct_sums.jl
+++ b/src/direct_sums.jl
@@ -83,21 +83,21 @@ to the `reduce` function.
 - `traj1::NamedTrajectory`: The first `NamedTrajectory` object.
 - `traj2::NamedTrajectory`: The second `NamedTrajectory` object.
 - `free_time::Bool=false`: Whether to construct a free time problem.
-- `timestep_symbol::Symbol=:Δt`: The timestep symbol to use for free time problems.
+- `timestep_name::Symbol=:Δt`: The timestep symbol to use for free time problems.
 """
 function direct_sum(
     traj1::NamedTrajectory,
     traj2::NamedTrajectory;
     free_time::Bool=false,
-    timestep_symbol::Symbol=:Δt,
+    timestep_name::Symbol=:Δt,
 )
-    return direct_sum([traj1, traj2]; free_time=free_time, timestep_symbol=timestep_symbol)
+    return direct_sum([traj1, traj2]; free_time=free_time, timestep_name=timestep_name)
 end
 
 function direct_sum(
     trajs::AbstractVector{<:NamedTrajectory};
     free_time::Bool=false,
-    timestep_symbol::Symbol=:Δt,
+    timestep_name::Symbol=:Δt,
 )
     if length(trajs) < 2
         throw(ArgumentError("At least two trajectories must be provided"))
@@ -122,13 +122,13 @@ function direct_sum(
 
     # add timestep to components
     if free_time
-        components = merge_outer(components, NamedTuple{(timestep_symbol,)}([get_timesteps(trajs[1])]))
+        components = merge_outer(components, NamedTuple{(timestep_name,)}([get_timesteps(trajs[1])]))
     end
 
     return NamedTrajectory(
         components,
         controls=merge_outer([traj.control_names for traj in trajs]),
-        timestep=free_time ? timestep_symbol : timestep,
+        timestep=free_time ? timestep_name : timestep,
         bounds=merge_outer([traj.bounds for traj in trajs]),
         initial=merge_outer([traj.initial for traj in trajs]),
         final=merge_outer([traj.final for traj in trajs]),
@@ -670,21 +670,21 @@ end
     pi_false_ops = PiccoloOptions(verbose=false, free_time=false)
     pi_true_ops = PiccoloOptions(verbose=false, free_time=true)
     suffix = "_new"
-    timestep_symbol = :Δt
+    timestep_name = :Δt
 
     prob1 = UnitarySmoothPulseProblem(sys, GATES[:Y], T, Δt, piccolo_options=pi_false_ops, ipopt_options=ops)
     traj1 = direct_sum(prob1.trajectory, add_suffix(prob1.trajectory, suffix), free_time=true)
 
     # Direct sum (shared timestep name)
-    @test get_suffix(traj1, suffix).timestep == timestep_symbol
-    @test get_suffix(traj1, suffix, remove=true).timestep == timestep_symbol
+    @test get_suffix(traj1, suffix).timestep == timestep_name
+    @test get_suffix(traj1, suffix, remove=true).timestep == timestep_name
 
     prob2 = UnitarySmoothPulseProblem(sys, GATES[:Y], T, Δt, ipopt_options=ops, piccolo_options=pi_true_ops)
     traj2 = add_suffix(prob2.trajectory, suffix)
 
     # Trajectory (unique timestep name)
-    @test get_suffix(traj2, suffix).timestep == add_suffix(timestep_symbol, suffix)
-    @test get_suffix(traj2, suffix, remove=true).timestep == timestep_symbol
+    @test get_suffix(traj2, suffix).timestep == add_suffix(timestep_name, suffix)
+    @test get_suffix(traj2, suffix, remove=true).timestep == timestep_name
 end
 
 end # module
diff --git a/src/objectives/regularizer_objective.jl b/src/objectives/regularizer_objective.jl
index c9339fa4..7cabf28d 100644
--- a/src/objectives/regularizer_objective.jl
+++ b/src/objectives/regularizer_objective.jl
@@ -13,7 +13,7 @@ export PairwiseQuadraticRegularizer
     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
@@ -21,7 +21,7 @@ Fields:
     `R`: the regularization matrix
     `baseline`: the baseline values for the trajectory component
     `eval_hessian`: whether to evaluate the Hessian of the regularizer
-    `timestep_symbol`: the symbol for the timestep variable
+    `timestep_name`: the symbol for the timestep variable
 """
 function QuadraticRegularizer(;
 	name::Union{Nothing, Symbol}=nothing,
@@ -30,7 +30,7 @@ function QuadraticRegularizer(;
 	R::Union{Nothing, AbstractVector{<:Real}}=nothing,
     baseline::Union{Nothing, AbstractArray{<:Real}}=nothing,
 	eval_hessian::Bool=true,
-    timestep_symbol::Symbol=:Δt
+    timestep_name::Symbol=:Δt
 )
 
     @assert !isnothing(name) "name must be specified"
@@ -59,7 +59,7 @@ function QuadraticRegularizer(;
         J = 0.0
         for t ∈ times
             if Z.timestep isa Symbol
-                Δt = Z⃗[slice(t, Z.components[timestep_symbol], Z.dim)]
+                Δt = Z⃗[slice(t, Z.components[timestep_name], Z.dim)]
             else
                 Δt = Z.timestep
             end
@@ -80,7 +80,7 @@ function QuadraticRegularizer(;
             Δv = Z⃗[vₜ_slice] .- baseline[:, t]
 
             if Z.timestep isa Symbol
-                Δt_slice = slice(t, Z.components[timestep_symbol], Z.dim)
+                Δt_slice = slice(t, Z.components[timestep_name], Z.dim)
                 Δt = Z⃗[Δt_slice]
                 ∇[Δt_slice] .= Δv' * (R .* (Δt .* Δv))
             else
@@ -106,7 +106,7 @@ function QuadraticRegularizer(;
                 append!(structure, vₜ_vₜ_inds)
 
                 if Z.timestep isa Symbol
-                    Δt_slice = slice(t, Z.components[timestep_symbol], Z.dim)
+                    Δ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)
@@ -126,7 +126,7 @@ function QuadraticRegularizer(;
             # Match Hessian structure indices
             for t ∈ times
                 if Z.timestep isa Symbol
-                    Δt = Z⃗[slice(t, Z.components[timestep_symbol], Z.dim)]
+                    Δ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)]
@@ -433,7 +433,7 @@ regularization strength `R`. The regularizer is defined as
     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 
+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
@@ -444,7 +444,7 @@ Fields:
     `times`: the time steps to apply the regularizer
     `name1`: the first name
     `name2`: the second name
-    `timestep_symbol`: the symbol for the timestep
+    `timestep_name`: the symbol for the timestep
     `eval_hessian`: whether to evaluate the Hessian
 """
 function PairwiseQuadraticRegularizer(
@@ -452,7 +452,7 @@ function PairwiseQuadraticRegularizer(
     times::AbstractVector{Int},
     name1::Symbol,
     name2::Symbol;
-    timestep_symbol::Symbol=:Δt,
+    timestep_name::Symbol=:Δt,
     eval_hessian::Bool=false,
 )
     params = Dict(
@@ -467,7 +467,7 @@ function PairwiseQuadraticRegularizer(
         J = 0.0
         for t ∈ times
             if Z.timestep isa Symbol
-                Δt = Z⃗[slice(t, Z.components[timestep_symbol], Z.dim)]
+                Δt = Z⃗[slice(t, Z.components[timestep_name], Z.dim)]
             else
                 Δt = Z.timestep
             end
@@ -488,7 +488,7 @@ function PairwiseQuadraticRegularizer(
             z2_t = Z⃗[z2_t_slice]
 
             if Z.timestep isa Symbol
-                Δt_slice = slice(t, Z.components[timestep_symbol], Z.dim)
+                Δ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
@@ -515,7 +515,7 @@ end
 @doc raw"""
     PairwiseQuadraticRegularizer
 
-A convenience constructor for creating a PairwiseQuadraticRegularizer for the 
+A convenience constructor for creating a PairwiseQuadraticRegularizer for the
 trajectory component `name` with regularization strength `Rs` over the graph `graph`.
 """
 function PairwiseQuadraticRegularizer(
@@ -567,7 +567,7 @@ end
     include("../../test/test_utils.jl")
 
     T = 10
-    
+
     Z = NamedTrajectory(
         (ψ̃ = randn(4, T), u = randn(2, T)),
         controls=:u,
diff --git a/src/options.jl b/src/options.jl
index 76e9addb..39fe0f59 100644
--- a/src/options.jl
+++ b/src/options.jl
@@ -70,9 +70,16 @@ end
     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)
diff --git a/src/problem_templates/_problem_templates.jl b/src/problem_templates/_problem_templates.jl
index e82e9d2e..abb41cf9 100644
--- a/src/problem_templates/_problem_templates.jl
+++ b/src/problem_templates/_problem_templates.jl
@@ -33,4 +33,59 @@ include("unitary_bang_bang_problem.jl")
 include("quantum_state_smooth_pulse_problem.jl")
 include("quantum_state_minimum_time_problem.jl")
 
+
+function apply_piccolo_options!(
+    J::Objective,
+    constraints::AbstractVector{<:AbstractConstraint},
+    piccolo_options::PiccoloOptions,
+    traj::NamedTrajectory,
+    operator::Union{Nothing, OperatorType},
+    state_name::Symbol,
+    timestep_name::Symbol
+)
+    if !isnothing(operator) && piccolo_options.leakage_suppression
+        state_names = [
+            name for name ∈ traj.names
+                if startswith(string(name), string(state_name))
+        ]
+
+        if operator isa EmbeddedOperator
+            leakage_indices = get_iso_vec_leakage_indices(operator)
+            for state_name in state_names
+                J += L1Regularizer!(
+                    constraints,
+                    state_name,
+                    traj;
+                    R_value=piccolo_options.R_leakage,
+                    indices=leakage_indices,
+                    eval_hessian=piccolo_options.eval_hessian
+                )
+            end
+        else
+            @warn "leakage_suppression is only supported for embedded operators, ignoring."
+        end
+    end
+
+    if piccolo_options.free_time
+        if piccolo_options.timesteps_all_equal
+            push!(
+                constraints,
+                TimeStepsAllEqualConstraint(timestep_name, traj)
+            )
+        end
+    end
+
+    if !isnothing(piccolo_options.complex_control_norm_constraint_name)
+        norm_con = ComplexModulusContraint(
+            piccolo_options.complex_control_norm_constraint_name,
+            piccolo_options.complex_control_norm_constraint_radius,
+            traj;
+        )
+        push!(constraints, norm_con)
+    end
+
+    return
+end
+
+
 end
diff --git a/src/problem_templates/quantum_state_minimum_time_problem.jl b/src/problem_templates/quantum_state_minimum_time_problem.jl
index d3b0fcf1..6deeaf5d 100644
--- a/src/problem_templates/quantum_state_minimum_time_problem.jl
+++ b/src/problem_templates/quantum_state_minimum_time_problem.jl
@@ -2,9 +2,28 @@ export QuantumStateMinimumTimeProblem
 
 
 """
-    QuantumStateMinimumTimeProblem
+    QuantumStateMinimumTimeProblem(traj, sys, obj, integrators, constraints; kwargs...)
+    QuantumStateMinimumTimeProblem(prob; kwargs...)
+
+Construct a `QuantumControlProblem` for the minimum time problem of reaching a target state.
+
+# Arguments
+- `traj::NamedTrajectory`: The initial trajectory.
+- `sys::QuantumSystem`: The quantum system.
+- `obj::Objective`: The objective function.
+- `integrators::Vector{<:AbstractIntegrator}`: The integrators.
+- `constraints::Vector{<:AbstractConstraint}`: The constraints.
+or
+- `prob::QuantumControlProblem`: The quantum control problem.
+
+# Keyword Arguments
+- `state_name::Symbol=:ψ̃`: The symbol for the state variables.
+- `final_fidelity::Union{Real, Nothing}=nothing`: The final fidelity.
+- `D=1.0`: The cost weight on the time.
+- `ipopt_options::IpoptOptions=IpoptOptions()`: The Ipopt options.
+- `piccolo_options::PiccoloOptions=PiccoloOptions()`: The Piccolo options.
+- `kwargs...`: Additional keyword arguments, passed to [`QuantumControlProblem`](@ref).
 
-TODO: Add documentation
 """
 function QuantumStateMinimumTimeProblem end
 
@@ -14,14 +33,14 @@ function QuantumStateMinimumTimeProblem(
     obj::Objective,
     integrators::Vector{<:AbstractIntegrator},
     constraints::Vector{<:AbstractConstraint};
-    state_symbol::Symbol=:ψ̃,
+    state_name::Symbol=:ψ̃,
     final_fidelity::Union{Real, Nothing}=nothing,
     D=1.0,
     ipopt_options::IpoptOptions=IpoptOptions(),
     piccolo_options::PiccoloOptions=PiccoloOptions(),
     kwargs...
 )
-    state_names = [name for name in traj.names if startswith(name, state_symbol)]
+    state_names = [name for name in traj.names if startswith(name, state_name)]
     @assert length(state_names) ≥ 1 "No matching states found in trajectory"
 
     obj += MinimumTimeObjective(traj; D=D, eval_hessian=piccolo_options.eval_hessian)
@@ -87,8 +106,8 @@ end
         T = 50
         Δt = 0.2
         sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]])
-        ψ_init = [1.0, 0.0]
-        ψ_target = [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_smooth_pulse_problem.jl b/src/problem_templates/quantum_state_smooth_pulse_problem.jl
index b3fa7613..027097b8 100644
--- a/src/problem_templates/quantum_state_smooth_pulse_problem.jl
+++ b/src/problem_templates/quantum_state_smooth_pulse_problem.jl
@@ -2,36 +2,55 @@ export QuantumStateSmoothPulseProblem
 
 
 """
-    QuantumStateSmoothPulseProblem(
-        system::AbstractQuantumSystem,
-        ψ_init::Union{AbstractVector{<:Number}, Vector{<:AbstractVector{<:Number}}},
-        ψ_goal::Union{AbstractVector{<:Number}, Vector{<:AbstractVector{<:Number}}},
-        T::Int,
-        Δt::Float64;
-        kwargs...
-    )
+    QuantumStateSmoothPulseProblem(system, ψ_inits, ψ_goals, T, Δt; kwargs...)
 
-    QuantumStateSmoothPulseProblem(
-        H_drift::AbstractMatrix{<:Number},
-        H_drives::Vector{<:AbstractMatrix{<:Number}},
-        args...;
-        kwargs...
-    )
+Create a quantum state smooth pulse problem. The goal is to find a control pulse
+`a(t)` that drives all of the initial states `ψ_inits` to the corresponding
+target states `ψ_goals` using `T` timesteps of size `Δt`. This problem also controls the first and second derivatives of the control pulse, `da(t)` and `dda(t)`, to ensure smoothness.
 
-Create a quantum control problem for smooth pulse optimization of a quantum state trajectory.
+# Arguments
+- `system::AbstractQuantumSystem`: The quantum system.
+- `ψ_inits::Vector{<:AbstractVector{<:ComplexF64}}`: The initial states.
+- `ψ_goals::Vector{<:AbstractVector{<:ComplexF64}}`: The target states.
+- `T::Int`: The number of timesteps.
+- `Δt::Float64`: The timestep size.
 
-TODO: Document args
+# Keyword Arguments
+- `ipopt_options::IpoptOptions=IpoptOptions()`: The IPOPT options.
+- `piccolo_options::PiccoloOptions=PiccoloOptions()`: The Piccolo options.
+- `state_name::Symbol=:ψ̃`: The name of the state variable.
+- `control_name::Symbol=:a`: The name of the control variable.
+- `timestep_name::Symbol=:Δt`: The name of the timestep variable.
+- `init_trajectory::Union{NamedTrajectory, Nothing}=nothing`: The initial trajectory.
+- `a_bound::Float64=1.0`: The bound on the control pulse.
+- `a_bounds::Vector{Float64}=fill(a_bound, length(system.G_drives))`: The bounds on the control pulse.
+- `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.
+- `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.
+- `Δt_max::Float64=1.5 * Δt`: The maximum timestep size.
+- `drive_derivative_σ::Float64=0.01`: The standard deviation of the drive derivative random initialization.
+- `Q::Float64=100.0`: The weight on the state objective.
+- `R=1e-2`: The weight on the control pulse and its derivatives.
+- `R_a::Union{Float64, Vector{Float64}}=R`: The weight on the control pulse.
+- `R_da::Union{Float64, Vector{Float64}}=R`: The weight on the first derivative of the control pulse.
+- `R_dda::Union{Float64, Vector{Float64}}=R`: The weight on the second derivative of the control pulse.
+- `leakage_operator::Union{Nothing, EmbeddedOperator}=nothing`: The leakage operator, if leakage suppression is desired.
+- `constraints::Vector{<:AbstractConstraint}=AbstractConstraint[]`: The constraints.
 """
-function QuantumStateSmoothPulseProblem end
-
 function QuantumStateSmoothPulseProblem(
     system::AbstractQuantumSystem,
-    ψ_init::Union{AbstractVector{<:Number}, Vector{<:AbstractVector{<:Number}}},
-    ψ_goal::Union{AbstractVector{<:Number}, Vector{<:AbstractVector{<:Number}}},
+    ψ_inits::Vector{<:AbstractVector{<:ComplexF64}},
+    ψ_goals::Vector{<:AbstractVector{<:ComplexF64}},
     T::Int,
     Δt::Float64;
     ipopt_options::IpoptOptions=IpoptOptions(),
     piccolo_options::PiccoloOptions=PiccoloOptions(),
+    state_name::Symbol=:ψ̃,
+    control_name::Symbol=:a,
+    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)),
@@ -48,39 +67,28 @@ function QuantumStateSmoothPulseProblem(
     R_a::Union{Float64, Vector{Float64}}=R,
     R_da::Union{Float64, Vector{Float64}}=R,
     R_dda::Union{Float64, Vector{Float64}}=R,
-    R_L1::Float64=20.0,
+    leakage_operator::Union{Nothing, EmbeddedOperator}=nothing,
     constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
-    L1_regularized_names=Symbol[],
-    L1_regularized_indices::NamedTuple=NamedTuple(),
     kwargs...
 )
-    @assert all(name ∈ L1_regularized_names for name in keys(L1_regularized_indices) if !isempty(L1_regularized_indices[name]))
-
-    if ψ_init isa AbstractVector{<:Number} && ψ_goal isa AbstractVector{<:Number}
-        ψ_inits = [ψ_init]
-        ψ_goals = [ψ_goal]
-    else
-        @assert length(ψ_init) == length(ψ_goal)
-        ψ_inits = ψ_init
-        ψ_goals = ψ_goal
-    end
-
-    ψ̃_inits = ket_to_iso.(Vector{ComplexF64}.(ψ_inits))
-    ψ̃_goals = ket_to_iso.(Vector{ComplexF64}.(ψ_goals))
-
-    n_drives = length(system.G_drives)
+    @assert length(ψ_inits) == length(ψ_goals)
 
     # Trajectory
     if !isnothing(init_trajectory)
         traj = init_trajectory
     else
-        traj = initialize_quantum_state_trajectory(
-            ψ̃_goals,
-            ψ̃_inits,
+        n_drives = length(system.G_drives)
+
+        traj = initialize_trajectory(
+            ψ_goals,
+            ψ_inits,
             T,
             Δt,
             n_drives,
-            (a = a_bounds, da = da_bounds, dda = dda_bounds);
+            (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),
             drive_derivative_σ=drive_derivative_σ,
@@ -91,50 +99,94 @@ function QuantumStateSmoothPulseProblem(
     end
 
     # Objective
-    J = QuadraticRegularizer(:a, traj, R_a)
-    J += QuadraticRegularizer(:da, traj, R_da)
-    J += QuadraticRegularizer(:dda, traj, R_dda)
+    control_names = [
+        name for name ∈ traj.names
+            if endswith(string(name), string(control_name))
+    ]
 
-    for i = 1:length(ψ_inits)
-        J += QuantumStateObjective(Symbol("ψ̃$i"), traj, Q)
-    end
+    J = QuadraticRegularizer(control_names[1], traj, R_a; timestep_name=timestep_name)
+    J += QuadraticRegularizer(control_names[2], traj, R_da; timestep_name=timestep_name)
+    J += QuadraticRegularizer(control_names[3], traj, R_dda; timestep_name=timestep_name)
 
-    # Constraints
-    for name in L1_regularized_names
-        if name in keys(L1_regularized_indices)
-            J += L1Regularizer!(
-                constraints, name, traj,
-                R_value=R_L1,
-                indices=L1_regularized_indices[name],
-                eval_hessian=piccolo_options.eval_hessian
-            )
-        else
-            J += L1Regularizer!(
-                constraints, name, traj;
-                R_value=R_L1,
-                eval_hessian=piccolo_options.eval_hessian
-            )
+    if length(ψ_inits) == 1
+        J += QuantumStateObjective(state_name, traj, Q)
+    else
+        state_names = [
+            name for name ∈ traj.names
+                if startswith(string(name), string(state_name))
+        ]
+        @assert length(state_names) == length(ψ_inits) "Number of states must match number of initial states"
+        for i = 1:length(ψ_inits)
+            J += QuantumStateObjective(state_names[i], traj, Q)
         end
     end
 
-    if piccolo_options.free_time
-        if piccolo_options.timesteps_all_equal
-            push!(constraints, TimeStepsAllEqualConstraint(:Δt, traj))
+    # Integrators
+    if length(ψ_inits) == 1
+        if piccolo_options.integrator == :pade
+            state_integrators = [QuantumStatePadeIntegrator(
+                system,
+                state_name,
+                control_name,
+                traj;
+                order=piccolo_options.pade_order
+            )]
+        elseif piccolo_options.integrator == :exponential
+            state_integrators = [QuantumStateExponentialIntegrator(
+                system,
+                state_name,
+                control_name,
+                traj
+            )]
+        else
+            error("integrator must be one of (:pade, :exponential)")
+        end
+    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(
+                    system,
+                    state_names[i],
+                    control_name,
+                    traj;
+                    order=piccolo_options.pade_order
+                )
+            elseif piccolo_options.integrator == :exponential
+                state_integrator = QuantumStateExponentialIntegrator(
+                    system,
+                    state_names[i],
+                    control_name,
+                    traj
+                )
+            else
+                error("integrator must be one of (:pade, :exponential)")
+            end
+            push!(state_integrators, state_integrator)
         end
     end
 
-    # Integrators
-    ψ̃_integrators = [
-        QuantumStatePadeIntegrator(system, Symbol("ψ̃$i"), :a, traj)
-            for i = 1:length(ψ_inits)
-    ]
-
     integrators = [
-        ψ̃_integrators...,
-        DerivativeIntegrator(:a, :da, traj),
-        DerivativeIntegrator(:da, :dda, traj)
+        state_integrators...,
+        DerivativeIntegrator(control_name, control_names[2], traj),
+        DerivativeIntegrator(control_names[2], control_names[3], traj)
     ]
 
+    # Optional Piccolo constraints and objectives
+    apply_piccolo_options!(
+        J,
+        constraints,
+        piccolo_options,
+        traj,
+        leakage_operator,
+        state_name,
+        timestep_name
+    )
+
     return QuantumControlProblem(
         system,
         traj,
@@ -147,6 +199,16 @@ function QuantumStateSmoothPulseProblem(
     )
 end
 
+function QuantumStateSmoothPulseProblem(
+    system::AbstractQuantumSystem,
+    ψ_init::AbstractVector{<:ComplexF64},
+    ψ_goal::AbstractVector{<:ComplexF64},
+    args...;
+    kwargs...
+)
+    return QuantumStateSmoothPulseProblem(system, [ψ_init], [ψ_goal], args...; kwargs...)
+end
+
 function QuantumStateSmoothPulseProblem(
     H_drift::AbstractMatrix{<:Number},
     H_drives::Vector{<:AbstractMatrix{<:Number}},
@@ -164,8 +226,8 @@ end
     T = 50
     Δt = 0.2
     sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]])
-    ψ_init = [1.0, 0.0]
-    ψ_target = [0.0, 1.0]
+    ψ_init = Vector{ComplexF64}([1.0, 0.0])
+    ψ_target = Vector{ComplexF64}([0.0, 1.0])
 
     # Single initial and target states
     # --------------------------------
@@ -181,8 +243,8 @@ end
 
     # Multiple initial and target states
     # ----------------------------------
-    ψ_inits = [[1.0, 0.0], [0.0, 1.0]]
-    ψ_targets = [[0.0, 1.0], [1.0, 0.0]]
+    ψ_inits = Vector{ComplexF64}.([[1.0, 0.0], [0.0, 1.0]])
+    ψ_targets = Vector{ComplexF64}.([[0.0, 1.0], [1.0, 0.0]])
     prob = QuantumStateSmoothPulseProblem(
         sys, ψ_inits, ψ_targets, T, Δt;
         ipopt_options=IpoptOptions(print_level=1),
@@ -199,8 +261,8 @@ end
     T = 50
     Δt = 0.2
     sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]])
-    ψ_init = [1.0, 0.0]
-    ψ_target = [0.0, 1.0]
+    ψ_init = Vector{ComplexF64}([1.0, 0.0])
+    ψ_target = Vector{ComplexF64}([0.0, 1.0])
     integrator=:exponential
 
     # Single initial and target states
@@ -217,8 +279,8 @@ end
 
     # Multiple initial and target states
     # ----------------------------------
-    ψ_inits = [[1.0, 0.0], [0.0, 1.0]]
-    ψ_targets = [[0.0, 1.0], [1.0, 0.0]]
+    ψ_inits = Vector{ComplexF64}.([[1.0, 0.0], [0.0, 1.0]])
+    ψ_targets = Vector{ComplexF64}.([[0.0, 1.0], [1.0, 0.0]])
     prob = QuantumStateSmoothPulseProblem(
         sys, ψ_inits, ψ_targets, T, Δt;
         ipopt_options=IpoptOptions(print_level=1),
@@ -229,3 +291,25 @@ end
     final = fidelity(prob)
     @test all(final .> initial)
 end
+
+@testitem "Test quantum state with multiple initial states and final states" begin
+    # System
+    T = 50
+    Δt = 0.2
+    sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]])
+    ψ_inits = Vector{ComplexF64}.([[1.0, 0.0], [0.0, 1.0]])
+    ψ_targets = Vector{ComplexF64}.([[0.0, 1.0], [1.0, 0.0]])
+
+    prob = QuantumStateSmoothPulseProblem(
+        sys, ψ_inits, ψ_targets, T, Δt;
+        ipopt_options=IpoptOptions(print_level=1),
+        piccolo_options=PiccoloOptions(verbose=false),
+        state_name=:psi,
+        control_name=:u,
+        timestep_name=:dt
+    )
+    initial = fidelity(prob)
+    solve!(prob, max_iter=20)
+    final = fidelity(prob)
+    @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 f3de2650..ac05a252 100644
--- a/src/problem_templates/unitary_bang_bang_problem.jl
+++ b/src/problem_templates/unitary_bang_bang_problem.jl
@@ -43,28 +43,26 @@ with
 # Keyword Arguments
 - `ipopt_options::IpoptOptions=IpoptOptions()`: the options for the Ipopt solver
 - `piccolo_options::PiccoloOptions=PiccoloOptions()`: the options for the Piccolo solver
-- `constraints::Vector{<:AbstractConstraint}=AbstractConstraint[]`: the constraints to enforce
+- `state_name::Symbol = :Ũ⃗`: the name of the state variable
+- `control_name::Symbol = :a`: the name of the control variable
+- `timestep_name::Symbol = :Δt`: the name of the timestep variable
 - `init_trajectory::Union{NamedTrajectory, Nothing}=nothing`: an initial trajectory to use
 - `a_bound::Float64=1.0`: the bound on the control pulse
-- `a_bounds::Vector{Float64}=fill(a_bound, length(system.G_drives))`: the bounds on the control pulses, one for each drive
+- `a_bounds=fill(a_bound, length(system.G_drives))`: the bounds on the control pulses, one for each drive
 - `a_guess::Union{Matrix{Float64}, Nothing}=nothing`: an initial guess for the control pulses
 - `da_bound::Float64=1.0`: 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
-- `Δt_min::Float64=0.5 * Δt`: the minimum time step size
-- `Δt_max::Float64=1.5 * Δt`: the maximum time step size
+- `da_bounds=fill(da_bound, length(system.G_drives))`: the bounds on the control pulse derivatives, one for each drive
+- `Δt_min::Float64=Δt isa Float64 ? 0.5 * Δt : 0.5 * mean(Δt)`: the minimum time step size
+- `Δt_max::Float64=Δt isa Float64 ? 1.5 * Δt : 1.5 * mean(Δt)`: the maximum time step size
 - `drive_derivative_σ::Float64=0.01`: the standard deviation of the initial guess for the control pulse derivatives
 - `Q::Float64=100.0`: the weight on the infidelity objective
 - `R=1e-2`: the weight on the regularization terms
+- `quadratic_control_regularization=false`: whether or not to use quadratic regularization for the control pulses
 - `R_a::Union{Float64, Vector{Float64}}=R`: the weight on the regularization term for the control pulses
 - `R_da::Union{Float64, Vector{Float64}}=R`: the weight on the regularization term for the control pulse derivatives
-- `R_bang_bang::Union{Float64, Vector{Float64}}=1.0`: the weight on the bang-bang regularization term
-- `leakage_suppression::Bool=false`: whether or not to suppress leakage to higher energy states
-- `R_leakage=1e-1`: the weight on the leakage suppression term
-- `bound_unitary=integrator == :exponential`: whether or not to bound the unitary
-- `control_norm_constraint=false`: whether or not to enforce a constraint on the control pulse norm
-- `control_norm_constraint_components=nothing`: the components of the control pulse to use for the norm constraint
-- `control_norm_R=nothing`: the weight on the control pulse norm constraint
-
+- `R_bang_bang::Union{Float64, Vector{Float64}}=1e-1`: the weight on the bang-bang regularization term
+- `global_data::Union{NamedTuple, Nothing}=nothing`: global data to be used in the problem
+- `constraints::Vector{<:AbstractConstraint}=AbstractConstraint[]`: the constraints to enforce
 
 TODO: control modulus norm, advanced feature, needs documentation
 
@@ -76,7 +74,9 @@ function UnitaryBangBangProblem(
     Δt::Union{Float64, Vector{Float64}};
     ipopt_options::IpoptOptions=IpoptOptions(),
     piccolo_options::PiccoloOptions=PiccoloOptions(),
-    constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
+    state_name::Symbol = :Ũ⃗,
+    control_name::Symbol = :a,
+    timestep_name::Symbol = :Δt,
     init_trajectory::Union{NamedTrajectory, Nothing}=nothing,
     a_bound::Float64=1.0,
     a_bounds=fill(a_bound, length(system.G_drives)),
@@ -88,16 +88,12 @@ function UnitaryBangBangProblem(
     drive_derivative_σ::Float64=0.01,
     Q::Float64=100.0,
     R=1e-2,
+    quadratic_control_regularization=false,
     R_a::Union{Float64, Vector{Float64}}=R,
     R_da::Union{Float64, Vector{Float64}}=R,
     R_bang_bang::Union{Float64, Vector{Float64}}=1e-1,
-    leakage_suppression=false,
-    R_leakage=1e-1,
-    control_norm_constraint=false,
-    control_norm_constraint_components=nothing,
-    control_norm_R=nothing,
-    bound_unitary=piccolo_options.integrator == :exponential,
     global_data::Union{NamedTuple, Nothing}=nothing,
+    constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
     kwargs...
 )
     # Trajectory
@@ -105,17 +101,19 @@ function UnitaryBangBangProblem(
         traj = init_trajectory
     else
         n_drives = length(system.G_drives)
-        traj = initialize_unitary_trajectory(
+        traj = initialize_trajectory(
             operator,
             T,
             Δt,
             n_drives,
-            (a = a_bounds, da = da_bounds);
-            n_derivatives=1,
+            (a_bounds, da_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),
             geodesic=piccolo_options.geodesic,
-            bound_unitary=bound_unitary,
+            bound_state=piccolo_options.bound_state,
             drive_derivative_σ=drive_derivative_σ,
             a_guess=a_guess,
             system=system,
@@ -126,81 +124,63 @@ function UnitaryBangBangProblem(
 
     # Objective
     if isnothing(global_data)
-        J = UnitaryInfidelityObjective(:Ũ⃗, traj, Q;
+        J = UnitaryInfidelityObjective(state_name, traj, Q;
             subspace=operator isa EmbeddedOperator ? operator.subspace_indices : nothing,
         )
     else
         # TODO: remove hardcoded args
         J = UnitaryFreePhaseInfidelityObjective(
-            name=:Ũ⃗,
-            phase_name=:ϕ,
-            goal=operator_to_iso_vec(operator isa EmbeddedOperator ? operator.operator : operator),
-            phase_operators=[GATES[:Z] for _ in eachindex(traj.global_components[:ϕ])],
+            name=state_name,
+            phase_name=piccolo_options.phase_name,
+            goal=operator_to_iso_vec(
+                operator isa EmbeddedOperator ? operator.operator : operator
+            ),
+            phase_operators=[GATES[:Z] for _ in eachindex(traj.global_components[:piccolo_options.phase_name])],
             Q=Q,
             eval_hessian=piccolo_options.eval_hessian,
             subspace=operator isa EmbeddedOperator ? operator.subspace_indices : nothing
         )
     end
-    J += QuadraticRegularizer(:a, traj, R_a)
-    J += QuadraticRegularizer(:da, traj, R_da)
+
+    control_names = [
+        name for name ∈ traj.names
+            if endswith(string(name), string(control_name))
+    ]
+
+    # TODO: do we need these regularizers?
+    if quadratic_control_regularization
+        J += QuadraticRegularizer(control_names[1], traj, R_a; timestep_name=timestep_name)
+        J += QuadraticRegularizer(control_names[2], traj, R_da; timestep_name=timestep_name)
+    end
 
     # Constraints
     if R_bang_bang isa Float64
         R_bang_bang = fill(R_bang_bang, length(system.G_drives))
     end
     J += L1Regularizer!(
-        constraints, :da, traj,
+        constraints, control_names[2], traj,
         R=R_bang_bang, eval_hessian=piccolo_options.eval_hessian
     )
 
-    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)
+            UnitaryPadeIntegrator(system, state_name, control_names[1], traj; order=piccolo_options.pade_order)
     elseif piccolo_options.integrator == :exponential
         unitary_integrator =
-            UnitaryExponentialIntegrator(system, :Ũ⃗, :a, traj)
+            UnitaryExponentialIntegrator(system, state_name, control_names[1], traj)
     else
         error("integrator must be one of (:pade, :exponential)")
     end
 
     integrators = [
         unitary_integrator,
-        DerivativeIntegrator(:a, :da, traj),
+        DerivativeIntegrator(control_names[1], control_names[2], traj),
     ]
 
+    # Optional Piccolo constraints and objectives
+    apply_piccolo_options!(J, constraints, piccolo_options, traj, operator, state_name, timestep_name)
+
     return QuantumControlProblem(
         system,
         traj,
@@ -245,25 +225,29 @@ end
     Δt = 0.2
 
     ipopt_options = IpoptOptions(print_level=1, max_iter=25)
-    piccolo_options = PiccoloOptions(verbose=false)
+    piccolo_options = PiccoloOptions(verbose=false, pade_order=12)
 
     prob = UnitaryBangBangProblem(
         sys, U_goal, T, Δt,
         R_bang_bang=10.,
-        ipopt_options=ipopt_options, piccolo_options=piccolo_options
+        ipopt_options=ipopt_options,
+        piccolo_options=piccolo_options,
+        control_name=:u
     )
 
     smooth_prob = UnitarySmoothPulseProblem(
-        sys, U_goal, T, Δt,
-        ipopt_options=ipopt_options, piccolo_options=piccolo_options
+        sys, U_goal, T, Δt;
+        ipopt_options=ipopt_options,
+        piccolo_options=piccolo_options,
+        control_name=:u
     )
-    initial = unitary_fidelity(prob)
+    initial = unitary_fidelity(prob; drive_name=:u)
     solve!(prob)
-    final = unitary_fidelity(prob)
+    final = unitary_fidelity(prob; drive_name=:u)
     @test final > initial
     solve!(smooth_prob)
     threshold = 1e-3
-    a_sparse = sum(prob.trajectory.da .> 5e-2)
-    a_dense = sum(smooth_prob.trajectory.da .> 5e-2)
+    a_sparse = sum(prob.trajectory.du .> 5e-2)
+    a_dense = sum(smooth_prob.trajectory.du .> 5e-2)
     @test a_sparse < a_dense
 end
diff --git a/src/problem_templates/unitary_direct_sum_problem.jl b/src/problem_templates/unitary_direct_sum_problem.jl
index bb9e8665..2c4b6813 100644
--- a/src/problem_templates/unitary_direct_sum_problem.jl
+++ b/src/problem_templates/unitary_direct_sum_problem.jl
@@ -4,15 +4,15 @@ export UnitaryDirectSumProblem
 @doc """
     UnitaryDirectSumProblem(probs, final_fidelity; kwargs...)
 
-Construct a `QuantumControlProblem` as a direct sum of unitary gate problems. The 
-purpose is to find solutions that are as close as possible with respect to one of 
+Construct a `QuantumControlProblem` as a direct sum of unitary gate problems. The
+purpose is to find solutions that are as close as possible with respect to one of
 their components. In particular, this is useful for finding interpolatable control solutions.
 
 A graph of edges (specified by problem labels) will enforce a `PairwiseQuadraticRegularizer` between
 the component trajectories of the problem in `probs` corresponding to the names of the edge in `edges`
 with corresponding edge weight `Q`.
 
-Boundary values can be included to enforce a `QuadraticRegularizer` on edges where one of the nodes is 
+Boundary values can be included to enforce a `QuadraticRegularizer` on edges where one of the nodes is
 not optimized. The boundary values are specified as a dictionary with keys corresponding to the edge
 labels and values corresponding to the boundary values.
 
@@ -107,11 +107,11 @@ function UnitaryDirectSumProblem(
         for ℓ in prob_labels
             a_symb, da_symb, dda_symb = add_suffix(:a, ℓ), add_suffix(:da, ℓ), add_suffix(:dda, ℓ)
             n_drives = length(traj.components[a_symb])
-            a, da, dda = TrajectoryInitialization.initialize_controls(
+            a, da, dda = TrajectoryInitialization.initialize_control_trajectory(
                 n_drives,
                 n_derivatives,
-                traj.T, 
-                traj.bounds[a_symb], 
+                traj.T,
+                traj.bounds[a_symb],
                 drive_derivative_σ
             )
             update!(traj, a_symb, (1 - drive_reset_ratio) * traj[a_symb] + drive_reset_ratio * a)
@@ -164,7 +164,7 @@ function UnitaryDirectSumProblem(
             Q_fid = isa(Q, Number) ? Q : Q[1]
             J += UnitaryInfidelityObjective(
                 add_suffix(:Ũ⃗, ℓ), traj, Q_fid,
-                subspace=subspace, 
+                subspace=subspace,
                 eval_hessian=piccolo_options.eval_hessian
             )
         end
@@ -212,11 +212,11 @@ end
 
     # Test label graph
     direct_sum_prob2 = UnitaryDirectSumProblem(
-        [prob1, prob2], 
+        [prob1, prob2],
         0.99,
         prob_labels=["a", "b"],
         graph=[("a", "b")],
-        verbose=false, 
+        verbose=false,
         ipopt_options=ops)
     state_names_ab = vcat(
         add_suffix(prob1.trajectory.state_names, "a")...,
@@ -231,20 +231,20 @@ end
 
     # Test bad graph
     @test_throws ArgumentError UnitaryDirectSumProblem(
-        [prob1, prob2], 
-        0.99, 
+        [prob1, prob2],
+        0.99,
         prob_labels=["a", "b"],
         graph=[("x", "b")],
-        verbose=false, 
+        verbose=false,
         ipopt_options=ops
     )
 
     # Test symbol graph
     direct_sum_prob3 = UnitaryDirectSumProblem(
-        [prob1, prob2], 
-        0.99, 
+        [prob1, prob2],
+        0.99,
         graph=[(:a1, :a2)],
-        verbose=false, 
+        verbose=false,
         ipopt_options=ops
     )
     @test issetequal(direct_sum_prob3.trajectory.state_names, state_names)
@@ -259,13 +259,13 @@ end
 
     # Test boundary values
     direct_sum_prob5 = UnitaryDirectSumProblem(
-        [prob1, prob2], 
+        [prob1, prob2],
         0.99,
         graph=[("x", "1"), ("1", "2")],
         R_b=1e3,
         Q_symb=:dda,
         boundary_values=Dict("x"=>copy(prob1.trajectory[:dda])),
-        verbose=false, 
+        verbose=false,
         ipopt_options=ops
     )
     # # TODO: Check for objectives?
diff --git a/src/problem_templates/unitary_minimum_time_problem.jl b/src/problem_templates/unitary_minimum_time_problem.jl
index 7409c706..42b3b372 100644
--- a/src/problem_templates/unitary_minimum_time_problem.jl
+++ b/src/problem_templates/unitary_minimum_time_problem.jl
@@ -36,7 +36,7 @@ J(\vec{\tilde{U}}, a, \dot{a}, \ddot{a}) + D \sum_t \Delta t_t \\
 - `constraints::Vector{<:AbstractConstraint}`: The constraints.
 
 # Keyword Arguments
-- `unitary_symbol::Symbol=:Ũ⃗`: The symbol for the unitary control.
+- `unitary_name::Symbol=:Ũ⃗`: The symbol for the unitary control.
 - `final_fidelity::Float64=0.99`: The final fidelity.
 - `D=1.0`: The weight for the minimum-time objective.
 - `ipopt_options::IpoptOptions=IpoptOptions()`: The options for the Ipopt solver.
@@ -51,8 +51,8 @@ function UnitaryMinimumTimeProblem(
     objective::Objective,
     integrators::Vector{<:AbstractIntegrator},
     constraints::Vector{<:AbstractConstraint};
-    unitary_symbol::Symbol=:Ũ⃗,
-    global_symbol::Union{Nothing, Symbol}=nothing,
+    unitary_name::Symbol=:Ũ⃗,
+    global_name::Union{Nothing, Symbol}=nothing,
     final_fidelity::Union{Real, Nothing}=nothing,
     D=1.0,
     ipopt_options::IpoptOptions=IpoptOptions(),
@@ -60,39 +60,40 @@ function UnitaryMinimumTimeProblem(
     subspace=nothing,
     kwargs...
 )
-    @assert unitary_symbol ∈ trajectory.names
+    @assert unitary_name ∈ trajectory.names
 
     if isnothing(final_fidelity)
         final_fidelity = iso_vec_unitary_fidelity(
-            trajectory[unitary_symbol][:, end], trajectory.goal[unitary_symbol]
+            trajectory[unitary_name][:, end], trajectory.goal[unitary_name]
         )
     end
 
     objective += MinimumTimeObjective(trajectory; D=D, eval_hessian=piccolo_options.eval_hessian)
 
-    if isnothing(global_symbol)
+    if isnothing(global_name)
         fidelity_constraint = FinalUnitaryFidelityConstraint(
-            unitary_symbol,
+            unitary_name,
             final_fidelity,
             trajectory;
             subspace=subspace,
             eval_hessian=piccolo_options.eval_hessian
         )
     else
-        # TODO: remove hardcoded args
+        phase_operators= [
+            GATES[:Z] for _ in eachindex(trajectory.global_components[global_name])
+        ]
         fidelity_constraint = FinalUnitaryFreePhaseFidelityConstraint(
-            value=final_fidelity,
-            state_slice=slice(trajectory.T, trajectory.components[unitary_symbol], trajectory.dim),
-            phase_slice=trajectory.global_components[global_symbol],
-            goal=trajectory.goal[unitary_symbol],
-            phase_operators=[GATES[:Z] for _ in eachindex(trajectory.global_components[global_symbol])],
-            zdim=trajectory.dim * trajectory.T + trajectory.global_dim,
+            unitary_name,
+            global_name,
+            phase_operators,
+            final_fidelity,
+            trajectory;
             subspace=subspace,
             eval_hessian=piccolo_options.eval_hessian
         )
     end
 
-    constraints = push!(constraints, fidelity_constraint) 
+    constraints = push!(constraints, fidelity_constraint)
 
     return QuantumControlProblem(
         system,
@@ -116,7 +117,7 @@ function UnitaryMinimumTimeProblem(
     kwargs...
 )
     piccolo_options.build_trajectory_constraints = build_trajectory_constraints
-    
+
     return UnitaryMinimumTimeProblem(
         copy(prob.trajectory),
         prob.system,
@@ -147,14 +148,14 @@ end
     )
 
     before = unitary_fidelity(prob)
-    solve!(prob, max_iter=20)
+    solve!(prob, max_iter=50)
     after = unitary_fidelity(prob)
     @test after > before
 
     # Soft fidelity constraint
     final_fidelity = minimum([0.99, after])
     mintime_prob = UnitaryMinimumTimeProblem(prob, final_fidelity=final_fidelity)
-    solve!(mintime_prob; max_iter=40)
+    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)
diff --git a/src/problem_templates/unitary_robustness_problem.jl b/src/problem_templates/unitary_robustness_problem.jl
index 991f5d84..de807dd8 100644
--- a/src/problem_templates/unitary_robustness_problem.jl
+++ b/src/problem_templates/unitary_robustness_problem.jl
@@ -3,18 +3,25 @@ export UnitaryRobustnessProblem
 
 @doc raw"""
     UnitaryRobustnessProblem(
-        H_error, trajectory, system, objective, integrators, constraints;
-        unitary_symbol=:Ũ⃗,
-        final_fidelity=nothing,
-        subspace=nothing,
-        ipopt_options=IpoptOptions(),
-        piccolo_options=PiccoloOptions(),
+        H_error,
+        trajectory,
+        system,
+        objective,
+        integrators,
+        constraints;
         kwargs...
     )
 
     UnitaryRobustnessProblem(Hₑ, prob::QuantumControlProblem; kwargs...)
 
 Create a quantum control problem for robustness optimization of a unitary trajectory.
+
+# Keyword Arguments
+- `unitary_symbol::Symbol=:Ũ⃗`: The symbol for the unitary trajectory in `trajectory`.
+- `final_fidelity::Union{Real, Nothing}=nothing`: The target fidelity for the final unitary.
+- `ipopt_options::IpoptOptions=IpoptOptions()`: Options for the Ipopt solver.
+- `piccolo_options::PiccoloOptions=PiccoloOptions()`: Options for the Piccolo solver.
+- `kwargs...`: Additional keyword arguments passed to `QuantumControlProblem`.
 """
 function UnitaryRobustnessProblem end
 
@@ -96,12 +103,12 @@ end
     H_drift = zeros(3, 3)
     H_drives = [create(3) + annihilate(3), im * (create(3) - annihilate(3))]
     sys = QuantumSystem(H_drift, H_drives)
-    
+
     U_goal = EmbeddedOperator(:X, sys)
     H_embed = EmbeddedOperator(:Z, sys)
     T = 51
     Δt = 0.2
-    
+
     #  test initial problem
     # ---------------------
     prob = UnitarySmoothPulseProblem(
@@ -126,7 +133,7 @@ end
         final_fidelity=final_fidelity,
         ipopt_options=IpoptOptions(recalc_y="yes", recalc_y_feas_tol=100.0, print_level=1),
     )
-    solve!(rob_prob, max_iter=50)    
+    solve!(rob_prob, max_iter=50)
 
     loss(Z⃗) = UnitaryRobustnessObjective(H_error=H_embed).L(Z⃗, prob.trajectory)
 
diff --git a/src/problem_templates/unitary_sampling_problem.jl b/src/problem_templates/unitary_sampling_problem.jl
index 539a2b61..6d32fd20 100644
--- a/src/problem_templates/unitary_sampling_problem.jl
+++ b/src/problem_templates/unitary_sampling_problem.jl
@@ -2,7 +2,7 @@ export UnitarySamplingProblem
 
 
 @doc raw"""
-    UnitarySamplingProblem
+    UnitarySamplingProblem(systemns, operator, T, Δt; kwargs...)
 
 A `UnitarySamplingProblem` is a quantum control problem where the goal is to find a control pulse that generates a target unitary operator for a set of quantum systems.
 The controls are shared among all systems, and the optimization seeks to find the control pulse that achieves the operator for each system. The idea is to enforce a
@@ -12,34 +12,33 @@ robust solution by including multiple systems reflecting the problem uncertainty
 - `systems::AbstractVector{<:AbstractQuantumSystem}`: A vector of quantum systems.
 - `operator::OperatorType`: The target unitary operator.
 - `T::Int`: The number of time steps.
-- `Δt::Union{Float64, Vector{Float64}}`: The time step size.
-- `system_labels::Vector{String}`: The labels for each system.
-- `system_weights::Vector{Float64}`: The weights for each system.
-- `init_trajectory::Union{NamedTrajectory, Nothing}`: The initial trajectory.
-- `ipopt_options::IpoptOptions`: The IPOPT options.
-- `piccolo_options::PiccoloOptions`: The Piccolo options.
-- `constraints::Vector{<:AbstractConstraint}`: The constraints.
-- `a_bound::Float64`: The bound for the control amplitudes.
-- `a_bounds::Vector{Float64}`: The bounds for the control amplitudes.
-- `a_guess::Union{Matrix{Float64}, Nothing}`: The initial guess for the control amplitudes.
-- `da_bound::Float64`: The bound for the control first derivatives.
-- `da_bounds::Vector{Float64}`: The bounds for the control first derivatives.
-- `dda_bound::Float64`: The bound for the control second derivatives.
-- `dda_bounds::Vector{Float64}`: The bounds for the control second derivatives.
-- `Δt_min::Float64`: The minimum time step size.
-- `Δt_max::Float64`: The maximum time step size.
-- `drive_derivative_σ::Float64`: The standard deviation for the drive derivative noise.
-- `Q::Float64`: The fidelity weight.
-- `R::Float64`: The regularization weight.
-- `R_a::Union{Float64, Vector{Float64}}`: The regularization weight for the control amplitudes.
-- `R_da::Union{Float64, Vector{Float64}}`: The regularization weight for the control first derivatives.
-- `R_dda::Union{Float64, Vector{Float64}}`: The regularization weight for the control second derivatives.
-- `leakage_suppression::Bool`: Whether to suppress leakage.
-- `R_leakage::Float64`: The regularization weight for the leakage.
-- `bound_unitary::Bool`: Whether to bound the unitary.
-- `control_norm_constraint::Bool`: Whether to enforce a control norm constraint.
-- `control_norm_constraint_components`: The components for the control norm constraint.
-- `control_norm_R`: The regularization weight for the control norm constraint.
+- `Δt::Union{Float64, Vector{Float64}}`: The time step value or vector of time steps.
+
+# Keyword Arguments
+- `system_labels::Vector{String} = string.(1:length(systems))`: The labels for each system.
+- `system_weights::Vector{Float64} = fill(1.0, length(systems))`: The weights for each system.
+- `init_trajectory::Union{NamedTrajectory, Nothing} = nothing`: The initial trajectory.
+- `ipopt_options::IpoptOptions = IpoptOptions()`: The IPOPT options.
+- `piccolo_options::PiccoloOptions = PiccoloOptions()`: The Piccolo options.
+- `state_name::Symbol = :Ũ⃗`: The name of the state variable.
+- `control_name::Symbol = :a`: The name of the control variable.
+- `timestep_name::Symbol = :Δt`: The name of the timestep variable.
+- `constraints::Vector{<:AbstractConstraint} = AbstractConstraint[]`: The constraints.
+- `a_bound::Float64 = 1.0`: The bound for the control amplitudes.
+- `a_bounds::Vector{Float64} = fill(a_bound, length(systems[1].G_drives))`: The bounds for the control amplitudes.
+- `a_guess::Union{Matrix{Float64}, Nothing} = nothing`: The initial guess for the control amplitudes.
+- `da_bound::Float64 = Inf`: The bound for the control first derivatives.
+- `da_bounds::Vector{Float64} = fill(da_bound, length(systems[1].G_drives))`: The bounds for the control first derivatives.
+- `dda_bound::Float64 = 1.0`: The bound for the control second derivatives.
+- `dda_bounds::Vector{Float64} = fill(dda_bound, length(systems[1].G_drives))`: The bounds for the control second derivatives.
+- `Δt_min::Float64 = 0.5 * Δt`: The minimum time step size.
+- `Δt_max::Float64 = 1.5 * Δt`: The maximum time step size.
+- `drive_derivative_σ::Float64 = 0.01`: The standard deviation for the drive derivative noise.
+- `Q::Float64 = 100.0`: The fidelity weight.
+- `R::Float64 = 1e-2`: The regularization weight.
+- `R_a::Union{Float64, Vector{Float64}} = R`: The regularization weight for the control amplitudes.
+- `R_da::Union{Float64, Vector{Float64}} = R`: The regularization weight for the control first derivatives.
+- `R_dda::Union{Float64, Vector{Float64}} = R`: The regularization weight for the control second derivatives.
 - `kwargs...`: Additional keyword arguments.
 
 """
@@ -47,16 +46,19 @@ function UnitarySamplingProblem(
     systems::AbstractVector{<:AbstractQuantumSystem},
     operator::OperatorType,
     T::Int,
-    Δt::Union{Float64, Vector{Float64}};
+    Δt::Union{Float64,Vector{Float64}};
     system_labels=string.(1:length(systems)),
     system_weights=fill(1.0, length(systems)),
-    init_trajectory::Union{NamedTrajectory, Nothing}=nothing,
+    init_trajectory::Union{NamedTrajectory,Nothing}=nothing,
     ipopt_options::IpoptOptions=IpoptOptions(),
     piccolo_options::PiccoloOptions=PiccoloOptions(),
+    state_name::Symbol=:Ũ⃗,
+    control_name::Symbol=:a,
+    timestep_name::Symbol=:Δt,
     constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
     a_bound::Float64=1.0,
     a_bounds=fill(a_bound, length(systems[1].G_drives)),
-    a_guess::Union{Matrix{Float64}, Nothing}=nothing,
+    a_guess::Union{Matrix{Float64},Nothing}=nothing,
     da_bound::Float64=Inf,
     da_bounds::Vector{Float64}=fill(da_bound, length(systems[1].G_drives)),
     dda_bound::Float64=1.0,
@@ -66,63 +68,66 @@ function UnitarySamplingProblem(
     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,
-    bound_unitary=piccolo_options.integrator == :exponential,
-    control_norm_constraint=false,
-    control_norm_constraint_components=nothing,
-    control_norm_R=nothing,
+    R_a::Union{Float64,Vector{Float64}}=R,
+    R_da::Union{Float64,Vector{Float64}}=R,
+    R_dda::Union{Float64,Vector{Float64}}=R,
     kwargs...
 )
     # Create keys for multiple systems
-    Ũ⃗_keys = [add_suffix(:Ũ⃗, ℓ) for ℓ ∈ system_labels]
+    Ũ⃗_names = [add_suffix(state_name, ℓ) for ℓ ∈ system_labels]
 
     # Trajectory
     if !isnothing(init_trajectory)
         traj = init_trajectory
     else
         n_drives = length(systems[1].G_drives)
-        traj = initialize_unitary_trajectory(
+
+        traj = initialize_trajectory(
             operator,
             T,
             Δt,
             n_drives,
-            (a = a_bounds, da = da_bounds, dda = dda_bounds);
+            (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),
             geodesic=piccolo_options.geodesic,
-            bound_unitary=bound_unitary,
+            bound_state=piccolo_options.bound_state,
             drive_derivative_σ=drive_derivative_σ,
             a_guess=a_guess,
             system=systems,
             rollout_integrator=piccolo_options.rollout_integrator,
-            Ũ⃗_keys=Ũ⃗_keys
+            state_names=Ũ⃗_names
         )
     end
 
+    control_names = [
+        name for name ∈ traj.names
+        if endswith(string(name), string(control_name))
+    ]
+
     # Objective
     J = NullObjective()
-    for (wᵢ, Ũ⃗_key) in zip(system_weights, Ũ⃗_keys)
+    for (wᵢ, Ũ⃗_name) in zip(system_weights, Ũ⃗_names)
         J += wᵢ * UnitaryInfidelityObjective(
-            Ũ⃗_key, traj, Q;
+            Ũ⃗_name, traj, Q;
             subspace=operator isa EmbeddedOperator ? operator.subspace_indices : nothing
         )
     end
-    J += QuadraticRegularizer(:a, traj, R_a)
-    J += QuadraticRegularizer(:da, traj, R_da)
-    J += QuadraticRegularizer(:dda, traj, R_dda)
+    J += QuadraticRegularizer(control_names[1], traj, R_a; timestep_name=timestep_name)
+    J += QuadraticRegularizer(control_names[2], traj, R_da; timestep_name=timestep_name)
+    J += QuadraticRegularizer(control_names[3], traj, R_dda; timestep_name=timestep_name)
 
     # Constraints
-    if leakage_suppression
+    if piccolo_options.leakage_suppression
         if operator isa EmbeddedOperator
             leakage_indices = get_iso_vec_leakage_indices(operator)
-            for Ũ⃗_key in Ũ⃗_keys
+            for Ũ⃗_name in Ũ⃗_names
                 J += L1Regularizer!(
-                    constraints, Ũ⃗_key, traj,
-                    R_value=R_leakage,
+                    constraints, Ũ⃗_name, traj,
+                    R_value=piccolo_options.R_leakage,
                     indices=leakage_indices,
                     eval_hessian=piccolo_options.eval_hessian
                 )
@@ -138,30 +143,27 @@ function UnitarySamplingProblem(
         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"
+    if !isnothing(piccolo_options.complex_control_norm_constraint_name)
         norm_con = ComplexModulusContraint(
-            :a,
-            control_norm_R,
+            piccolo_options.complex_control_norm_constraint_name,
+            piccolo_options.complex_control_norm_constraint_radius,
             traj;
-            name_comps=control_norm_constraint_components,
         )
         push!(constraints, norm_con)
     end
 
     # Integrators
     unitary_integrators = AbstractIntegrator[]
-    for (sys, Ũ⃗_key) in zip(systems, Ũ⃗_keys)
+    for (sys, Ũ⃗_name) in zip(systems, Ũ⃗_names)
         if piccolo_options.integrator == :pade
             push!(
                 unitary_integrators,
-                UnitaryPadeIntegrator(sys, Ũ⃗_key, :a, traj; order=piccolo_options.pade_order)
+                UnitaryPadeIntegrator(sys, Ũ⃗_name, control_name, traj; order=piccolo_options.pade_order)
             )
         elseif piccolo_options.integrator == :exponential
             push!(
                 unitary_integrators,
-                UnitaryExponentialIntegrator(sys, Ũ⃗_key, :a, traj)
+                UnitaryExponentialIntegrator(sys, Ũ⃗_name, control_name, traj)
             )
         else
             error("integrator must be one of (:pade, :exponential)")
@@ -170,8 +172,8 @@ function UnitarySamplingProblem(
 
     integrators = [
         unitary_integrators...,
-        DerivativeIntegrator(:a, :da, traj),
-        DerivativeIntegrator(:da, :dda, traj),
+        DerivativeIntegrator(control_name, control_names[2], traj),
+        DerivativeIntegrator(control_names[2], control_names[3], traj),
     ]
 
     return QuantumControlProblem(
@@ -192,7 +194,7 @@ function UnitarySamplingProblem(
     num_samples::Int,
     operator::OperatorType,
     T::Int,
-    Δt::Union{Float64, Vector{Float64}};
+    Δt::Union{Float64,Vector{Float64}};
     kwargs...
 )
     samples = rand(distribution, num_samples)
@@ -220,7 +222,7 @@ end
     operator = GATES[:H]
     systems(ζ) = QuantumSystem(ζ * GATES[:Z], [GATES[:X], GATES[:Y]])
 
-    ip_ops = IpoptOptions(print_level=1, recalc_y = "yes", recalc_y_feas_tol = 1e1)
+    ip_ops = IpoptOptions(print_level=1, recalc_y="yes", recalc_y_feas_tol=1e1)
     pi_ops = PiccoloOptions(verbose=false)
 
     prob = UnitarySamplingProblem(
diff --git a/src/problem_templates/unitary_smooth_pulse_problem.jl b/src/problem_templates/unitary_smooth_pulse_problem.jl
index 3420f70a..f0d30f04 100644
--- a/src/problem_templates/unitary_smooth_pulse_problem.jl
+++ b/src/problem_templates/unitary_smooth_pulse_problem.jl
@@ -2,7 +2,7 @@ export UnitarySmoothPulseProblem
 
 
 @doc raw"""
-    UnitarySmoothPulseProblem(system::QuantumSystem, operator, T, Δt; kwargs...)
+    UnitarySmoothPulseProblem(system, operator, T, Δt; kwargs...)
 
 Construct a `QuantumControlProblem` for a free-time unitary gate problem with smooth control pulses enforced by constraining the second derivative of the pulse trajectory, i.e.,
 
@@ -42,32 +42,27 @@ with
 # Keyword Arguments
 - `ipopt_options::IpoptOptions=IpoptOptions()`: the options for the Ipopt solver
 - `piccolo_options::PiccoloOptions=PiccoloOptions()`: the options for the Piccolo solver
-- `constraints::Vector{<:AbstractConstraint}=AbstractConstraint[]`: the constraints to enforce
+- `state_name::Symbol = :Ũ⃗`: the name of the state
+- `control_name::Symbol = :a`: the name of the control
+- `timestep_name::Symbol = :Δt`: the name of the timestep
 - `init_trajectory::Union{NamedTrajectory, Nothing}=nothing`: an initial trajectory to use
 - `a_bound::Float64=1.0`: the bound on the control pulse
 - `a_bounds::Vector{Float64}=fill(a_bound, length(system.G_drives))`: the bounds on the control pulses, one for each drive
 - `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
-- `dda_bound::Float64=1.0`: the bound on the control pulse derivative
-- `dda_bounds::Vector{Float64}=fill(dda_bound, length(system.G_drives))`: the bounds on the control pulse derivatives, one for each drive
-- `Δt_min::Float64=0.5 * Δt`: the minimum time step size
-- `Δt_max::Float64=1.5 * Δt`: the maximum time step size
+- `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
+- `Δt_max::Float64=Δt isa Float64 ? 1.5 * Δt : 1.5 * mean(Δt)`: the maximum time step size
 - `drive_derivative_σ::Float64=0.01`: the standard deviation of the initial guess for the control pulse derivatives
 - `Q::Float64=100.0`: the weight on the infidelity objective
 - `R=1e-2`: the weight on the regularization terms
 - `R_a::Union{Float64, Vector{Float64}}=R`: the weight on the regularization term for the control pulses
 - `R_da::Union{Float64, Vector{Float64}}=R`: the weight on the regularization term for the control pulse derivatives
 - `R_dda::Union{Float64, Vector{Float64}}=R`: the weight on the regularization term for the control pulse second derivatives
-- `leakage_suppression::Bool=false`: whether or not to suppress leakage to higher energy states
-- `R_leakage=1e-1`: the weight on the leakage suppression term
-- `bound_unitary=integrator == :exponential`: whether or not to bound the unitary
-- `control_norm_constraint=false`: whether or not to enforce a constraint on the control pulse norm
-- `control_norm_constraint_components=nothing`: the components of the control pulse to use for the norm constraint
-- `control_norm_R=nothing`: the weight on the control pulse norm constraint
-
-
-TODO: control modulus norm, advanced feature, needs documentation
+- `global_data::Union{NamedTuple, Nothing}=nothing`: global data for the problem
+- `constraints::Vector{<:AbstractConstraint}=AbstractConstraint[]`: the constraints to enforce
 
 """
 function UnitarySmoothPulseProblem(
@@ -77,7 +72,9 @@ function UnitarySmoothPulseProblem(
     Δt::Union{Float64, Vector{Float64}};
     ipopt_options::IpoptOptions=IpoptOptions(),
     piccolo_options::PiccoloOptions=PiccoloOptions(),
-    constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
+    state_name::Symbol = :Ũ⃗,
+    control_name::Symbol = :a,
+    timestep_name::Symbol = :Δt,
     init_trajectory::Union{NamedTrajectory, Nothing}=nothing,
     a_bound::Float64=1.0,
     a_bounds=fill(a_bound, length(system.G_drives)),
@@ -94,30 +91,30 @@ function UnitarySmoothPulseProblem(
     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,
-    bound_unitary=piccolo_options.integrator == :exponential,
     global_data::Union{NamedTuple, Nothing}=nothing,
+    constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
     kwargs...
 )
+
     # Trajectory
     if !isnothing(init_trajectory)
         traj = init_trajectory
     else
         n_drives = length(system.G_drives)
-        traj = initialize_unitary_trajectory(
+
+        traj = initialize_trajectory(
             operator,
             T,
             Δt,
             n_drives,
-            (a = a_bounds, da = da_bounds, dda = dda_bounds);
+            (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),
             geodesic=piccolo_options.geodesic,
-            bound_unitary=bound_unitary,
+            bound_state=piccolo_options.bound_state,
             drive_derivative_σ=drive_derivative_σ,
             a_guess=a_guess,
             system=system,
@@ -128,75 +125,57 @@ function UnitarySmoothPulseProblem(
 
     # Objective
     if isnothing(global_data)
-        J = UnitaryInfidelityObjective(:Ũ⃗, traj, Q;
+        J = UnitaryInfidelityObjective(state_name, traj, Q;
             subspace=operator isa EmbeddedOperator ? operator.subspace_indices : nothing,
         )
     else
         # TODO: remove hardcoded args
         J = UnitaryFreePhaseInfidelityObjective(
-            name=:Ũ⃗,
+            name=state_name,
             phase_name=:ϕ,
-            goal=operator_to_iso_vec(operator isa EmbeddedOperator ? operator.operator : operator),
-            phase_operators=[GATES[:Z] for _ in eachindex(traj.global_components[:ϕ])],
+            goal=operator_to_iso_vec(
+                operator isa EmbeddedOperator ?
+                operator.operator :
+                operator
+            ),
+            phase_operators=fill(GATES[:Z], length(traj.global_components[:ϕ])),
             Q=Q,
             eval_hessian=piccolo_options.eval_hessian,
             subspace=operator isa EmbeddedOperator ? operator.subspace_indices : nothing
         )
     end
-    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
+    control_names = [
+        name for name ∈ traj.names
+            if endswith(string(name), string(control_name))
+    ]
 
-    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
+    J += QuadraticRegularizer(control_names[1], traj, R_a; timestep_name=timestep_name)
+    J += QuadraticRegularizer(control_names[2], traj, R_da; timestep_name=timestep_name)
+    J += QuadraticRegularizer(control_names[3], traj, R_dda; timestep_name=timestep_name)
 
     # Integrators
     if piccolo_options.integrator == :pade
         unitary_integrator =
-            UnitaryPadeIntegrator(system, :Ũ⃗, :a, traj; order=piccolo_options.pade_order)
+            UnitaryPadeIntegrator(system, state_name, control_name, traj;
+                order=piccolo_options.pade_order
+            )
     elseif piccolo_options.integrator == :exponential
         unitary_integrator =
-            UnitaryExponentialIntegrator(system, :Ũ⃗, :a, traj)
+            UnitaryExponentialIntegrator(system, state_name, control_name, traj)
     else
         error("integrator must be one of (:pade, :exponential)")
     end
 
     integrators = [
         unitary_integrator,
-        DerivativeIntegrator(:a, :da, traj),
-        DerivativeIntegrator(:da, :dda, traj),
+        DerivativeIntegrator(control_name, control_names[2], traj),
+        DerivativeIntegrator(control_names[2], control_names[3], traj),
     ]
 
+    # Optional Piccolo constraints and objectives
+    apply_piccolo_options!(J, constraints, piccolo_options, traj, operator, state_name, timestep_name)
+
     return QuantumControlProblem(
         system,
         traj,
@@ -271,6 +250,34 @@ end
     @test final > initial
 end
 
+@testitem "Hadamard gate with exponential integrator, bounded states, and control norm constraint" begin
+    sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]])
+    U_goal = GATES[:H]
+    T = 51
+    Δt = 0.2
+
+    piccolo_options = PiccoloOptions(
+        verbose=false,
+        integrator=:exponential,
+        # jacobian_structure=false,
+        bound_state=true,
+        complex_control_norm_constraint_name=:a
+    )
+
+    prob = UnitarySmoothPulseProblem(
+        sys, U_goal, T, Δt,
+        ipopt_options=IpoptOptions(print_level=1),
+        piccolo_options=piccolo_options
+    )
+
+    initial = unitary_fidelity(prob)
+    solve!(prob, max_iter=20)
+    final = unitary_fidelity(prob)
+    @test final > initial
+end
+
+
+
 @testitem "EmbeddedOperator Hadamard gate" begin
     a = annihilate(3)
     sys = QuantumSystem([(a + a')/2, (a - a')/(2im)])
diff --git a/src/quantum_systems.jl b/src/quantum_systems.jl
index fe52c2a3..105ba50d 100644
--- a/src/quantum_systems.jl
+++ b/src/quantum_systems.jl
@@ -25,15 +25,7 @@ using TestItemRunner
 # ----------------------------------------------------------------------------- #
 
 """
-```julia
-AbstractQuantumSystem
-```
-   |
-    -> EmbeddedOperator
-   |                  |
-   |                   -> QuantumObjective
-   |                  |
-    -> Matrix (goal) -
+    AbstractQuantumSystem
 
 Abstract type for defining systems.
 """
diff --git a/src/rollouts.jl b/src/rollouts.jl
index 693094be..91eaec7a 100644
--- a/src/rollouts.jl
+++ b/src/rollouts.jl
@@ -27,14 +27,14 @@ using TestItemRunner
 # ----------------------------------------------------------------------------- #
 
 """
-    infer_is_jvp(integrator::Function)
+    infer_is_evp(integrator::Function)
 
-Infer whether the integrator is a jacobian-vector product (JVP) function.
+Infer whether the integrator is a exponential-vector product (EVP) function.
 
 If `true`, the integrator is expected to have a signature like the exponential action,
 `expv`. Otherwise, it is expected to have a signature like `exp`.
 """
-function infer_is_jvp(integrator::Function)
+function infer_is_evp(integrator::Function)
     # name + args
     ns = fieldcount.([m.sig for m ∈ methods(integrator)])
     is_exp = 2 ∈ ns
@@ -65,7 +65,7 @@ end
 Rollout a quantum state `ψ̃_init` under the control `controls` for a time `Δt`
 using the system `system`.
 
-If `jacobian_vector_product` is `true`, the integrator is expected to have a signature like
+If `exp_vector_product` is `true`, the integrator is expected to have a signature like
 the exponential action, `expv`. Otherwise, it is expected to have a signature like `exp`.
 
 Types should allow for autodifferentiable controls and times.
@@ -77,7 +77,7 @@ function rollout(
     system::AbstractQuantumSystem;
     show_progress=false,
     integrator=expv,
-    jacobian_vector_product=infer_is_jvp(integrator),
+    exp_vector_product=infer_is_evp(integrator),
     G=Integrators.G_bilinear
 )
     T = size(controls, 2)
@@ -93,7 +93,7 @@ function rollout(
     for t = 2:T
         aₜ₋₁ = controls[:, t - 1]
         Gₜ = G(aₜ₋₁, G_drift, G_drives)
-        if jacobian_vector_product
+        if exp_vector_product
             Ψ̃[:, t] .= integrator(Δt[t - 1], Gₜ, Ψ̃[:, t - 1])
         else
             Ψ̃[:, t] .= integrator(Gₜ * Δt[t - 1]) * Ψ̃[:, t - 1]
@@ -174,7 +174,7 @@ function open_rollout(
     system::AbstractQuantumSystem;
     show_progress=false,
     integrator=expv,
-    jacobian_vector_product=infer_is_jvp(integrator),
+    exp_vector_product=infer_is_evp(integrator),
     H=a -> Integrators.G_bilinear(a, system.H_drift, system.H_drives),
 )
     T = size(controls, 2)
@@ -195,7 +195,7 @@ function open_rollout(
     for t = 2:T
         aₜ₋₁ = controls[:, t - 1]
         adGₜ = Isomorphisms.G(ad_vec(H(aₜ₋₁)))
-        if jacobian_vector_product
+        if exp_vector_product
             ρ⃗̃[:, t] = integrator(Δt[t - 1], adGₜ + iso(L), ρ⃗̃[:, t - 1])
         else
             ρ⃗̃[:, t] = integrator(Δt[t - 1], adGₜ + iso(L)) * ρ⃗̃[:, t - 1]
@@ -217,7 +217,7 @@ function unitary_rollout(
     system::AbstractQuantumSystem;
     show_progress=false,
     integrator=expv,
-    jacobian_vector_product=infer_is_jvp(integrator),
+    exp_vector_product=infer_is_evp(integrator),
     G=Integrators.G_bilinear,
 )
     T = size(controls, 2)
@@ -234,7 +234,7 @@ function unitary_rollout(
         aₜ₋₁ = controls[:, t - 1]
         Gₜ = G(aₜ₋₁, G_drift, G_drives)
         Ũₜ₋₁ = iso_vec_to_iso_operator(Ũ⃗[:, t - 1])
-        if jacobian_vector_product
+        if exp_vector_product
             Ũₜ = integrator(Δt[t - 1], Gₜ, Ũₜ₋₁)
         else
             Ũₜ = integrator(Gₜ * Δt[t - 1]) * Ũₜ₋₁
diff --git a/src/trajectory_initialization.jl b/src/trajectory_initialization.jl
index 32f3c567..f6e1bc6b 100644
--- a/src/trajectory_initialization.jl
+++ b/src/trajectory_initialization.jl
@@ -3,8 +3,7 @@ module TrajectoryInitialization
 export unitary_geodesic
 export linear_interpolation
 export unitary_linear_interpolation
-export initialize_unitary_trajectory
-export initialize_quantum_state_trajectory
+export initialize_trajectory
 export convert_fixed_time
 export convert_free_time
 
@@ -118,31 +117,40 @@ function unitary_geodesic(
     if U_goal isa EmbeddedOperator
         U_goal = U_goal.operator
     end
-    return unitary_geodesic(Matrix{ComplexF64}(I(size(U_goal, 1))), U_goal, samples; kwargs...)
+    return unitary_geodesic(
+        Matrix{ComplexF64}(I(size(U_goal, 1))),
+        U_goal,
+        samples;
+        kwargs...
+    )
 end
 
+"""
+    unitary_geodesic(U_init, U_goal, times; kwargs...)
+
+Compute the geodesic connecting U_init and U_goal at the specified times. Allows for the possibility of unequal times and ranges outside [0,1].
+
+# Arguments
+- `U_init::AbstractMatrix{<:Number}`: The initial unitary operator.
+- `U_goal::AbstractMatrix{<:Number}`: The goal unitary operator.
+- `times::AbstractVector{<:Number}`: The times at which to evaluate the geodesic.
+
+# Keyword Arguments
+- `return_unitary_isos::Bool=true`: If true returns a matrix where each column is a unitary isovec, i.e. vec(vcat(real(U), imag(U))). If false, returns a vector of unitary matrices.
+- `return_generator::Bool=false`: If true, returns the effective Hamiltonian generating the geodesic.
+"""
 function unitary_geodesic(
     U_init::AbstractMatrix{<:Number},
     U_goal::AbstractMatrix{<:Number},
-    timesteps::AbstractVector{<:Number};
+    times::AbstractVector{<:Number};
     return_unitary_isos=true,
     return_generator=false
 )
-    """
-    Compute the effective generator of the geodesic connecting U₀ and U₁.
-        U_goal = exp(-im * H * T) U_init
-        log(U_goal * U_init') = -im * H * T
-
-    Allow for the possibiltiy of unequal timesteps and ranges outside [0,1].
-
-    Returns the geodesic.
-    Optionally returns the effective Hamiltonian generating the geodesic.
-    """
-    t₀ = timesteps[1]
-    T = timesteps[end] - t₀
+    t₀ = times[1]
+    T = times[end] - t₀
     H = im * log(U_goal * U_init') / T
     # -im prefactor is not included in H
-    U_geo = [exp(-im * H * (t - t₀)) * U_init for t ∈ timesteps]
+    U_geo = [exp(-im * H * (t - t₀)) * U_init for t ∈ times]
     if !return_unitary_isos
         if return_generator
             return U_geo, H
@@ -167,7 +175,7 @@ linear_interpolation(x::AbstractVector, y::AbstractVector, n::Int) =
 const VectorBound = Union{AbstractVector{R}, Tuple{AbstractVector{R}, AbstractVector{R}}} where R <: Real
 const ScalarBound = Union{R, Tuple{R, R}} where R <: Real
 
-function initialize_unitaries(
+function initialize_unitary_trajectory(
     U_init::AbstractMatrix{<:Number},
     U_goal::OperatorType,
     T::Int;
@@ -185,7 +193,7 @@ end
 #                           Initial controls                                    #
 # ----------------------------------------------------------------------------- #
 
-function initialize_controls(
+function initialize_control_trajectory(
     n_drives::Int,
     n_derivatives::Int,
     T::Int,
@@ -216,22 +224,28 @@ function initialize_controls(
     return controls
 end
 
-function initialize_controls(a::AbstractMatrix, Δt::AbstractVecOrMat, n_derivatives::Int)
+function initialize_control_trajectory(
+    a::AbstractMatrix,
+    Δt::AbstractVecOrMat,
+    n_derivatives::Int
+)
     controls = Matrix{Float64}[a]
+
     for n in 1:n_derivatives
         # next derivative
         push!(controls,  derivative(controls[end], Δt))
 
         # to avoid constraint violation error at initial iteration for da, dda, ...
         if n > 1
-            controls[end-1][:, end] = controls[end-1][:, end-1] + Δt[end-1] * controls[end][:, end-1]
+            controls[end-1][:, end] =
+                controls[end-1][:, end-1] + Δt[end-1] * controls[end][:, end-1]
         end
     end
     return controls
 end
 
-initialize_controls(a::AbstractMatrix, Δt::Real, n_derivatives::Int) =
-    initialize_controls(a, fill(Δt, size(a, 2)), n_derivatives)
+initialize_control_trajectory(a::AbstractMatrix, Δt::Real, n_derivatives::Int) =
+    initialize_control_trajectory(a, fill(Δt, size(a, 2)), n_derivatives)
 
 # ----------------------------------------------------------------------------- #
 #                           Trajectory initialization                           #
@@ -245,28 +259,45 @@ Initialize a trajectory for a unitary control problem. The trajectory is initial
 data that should be consistently the same type (in this case, Float64).
 
 """
-function initialize_unitary_trajectory(
-    U_goal::OperatorType,
+function initialize_trajectory(
+    state_data::Vector{Matrix{Float64}},
+    state_inits::Vector{<:AbstractVector{Float64}},
+    state_goals::Vector{<:AbstractVector{Float64}},
+    state_names::AbstractVector{Symbol},
     T::Int,
     Δt::Union{Float64, AbstractVecOrMat{<:Float64}},
     n_drives::Int,
-    all_a_bounds::NamedTuple{anames, <:Tuple{Vararg{VectorBound}}} where anames;
-    U_init::AbstractMatrix{<:Number}=Matrix{ComplexF64}(I(size(U_goal, 1))),
-    n_derivatives::Int=2,
-    geodesic=true,
-    bound_unitary=false,
+    control_bounds::Tuple{Vararg{VectorBound}};
+    bound_state=false,
     free_time=false,
+    control_name=:a,
+    n_control_derivatives::Int=length(control_bounds) - 1,
+    timestep_name=:Δt,
     Δt_bounds::ScalarBound=(0.5 * Δt, 1.5 * Δt),
     drive_derivative_σ::Float64=0.1,
     a_guess::Union{AbstractMatrix{<:Float64}, Nothing}=nothing,
-    system::Union{AbstractQuantumSystem, AbstractVector{<:AbstractQuantumSystem}, Nothing}=nothing,
     global_data::Union{NamedTuple, Nothing}=nothing,
-    rollout_integrator::Function=expv,
-    Ũ⃗_keys::AbstractVector{<:Symbol}=[:Ũ⃗],
-    a_keys::AbstractVector{<:Symbol}=[Symbol("d"^i * "a") for i in 0:n_derivatives]
 )
-    @assert length(a_keys) == n_derivatives + 1 "a_keys must have the same length as n_derivatives + 1"
+    @assert !isnothing(state_inits) "state_init must be provided"
+    @assert !isnothing(state_goals) "state_goal must be provided"
+    @assert length(state_data) == length(state_names) == length(state_inits) == length(state_goals) "state_data, state_names, state_inits, and state_goals must have the same length"
+
+    @assert length(control_bounds) == n_control_derivatives + 1 "control_bounds must have $n_control_derivatives + 1 elements"
 
+    # assert that state names are unique
+    @assert length(state_names) == length(Set(state_names)) "state_names must be unique"
+
+    # Control data
+    control_derivative_names = [
+        Symbol("d"^i * string(control_name))
+            for i = 1:n_control_derivatives
+    ]
+
+    control_names = (control_name, control_derivative_names...)
+
+    control_bounds = NamedTuple{control_names}(control_bounds)
+
+    # Timestep data
     if free_time
         if Δt isa Real
             Δt = fill(Δt, 1, T)
@@ -275,94 +306,64 @@ function initialize_unitary_trajectory(
         else
             @assert size(Δt) == (1, T) "Δt must be a Real, AbstractVector, or 1x$(T) AbstractMatrix"
         end
-        timestep = :Δt
+        timestep = timestep_name
     else
         @assert Δt isa Real "Δt must be a Real if free_time is false"
         timestep = Δt
     end
 
-    Ũ⃗_init = operator_to_iso_vec(U_init)
-    if U_goal isa EmbeddedOperator
-        Ũ⃗_goal = operator_to_iso_vec(U_goal.operator)
-    else
-        Ũ⃗_goal = operator_to_iso_vec(U_goal)
-    end
 
     # Constraints
-    Ũ⃗_inits = repeat([Ũ⃗_init], length(Ũ⃗_keys))
     initial = (;
-        (Ũ⃗_keys .=> Ũ⃗_inits)...,
-        a = zeros(n_drives),
+        (state_names .=> state_inits)...,
+        control_name => zeros(n_drives),
     )
 
-    final = (
-        a = zeros(n_drives),
+    final = (;
+        control_name => zeros(n_drives),
     )
 
-    Ũ⃗_goals = repeat([Ũ⃗_goal], length(Ũ⃗_keys))
-    goal = (; (Ũ⃗_keys .=> Ũ⃗_goals)...)
+    goal = (; (state_names .=> state_goals)...)
 
     # Bounds
-    bounds = all_a_bounds
+    bounds = control_bounds
 
-    if bound_unitary
-        Ũ⃗_dim = length(Ũ⃗_init)
-        Ũ⃗_bounds = repeat([(-ones(Ũ⃗_dim), ones(Ũ⃗_dim))], length(Ũ⃗_keys))
-        bounds = merge(bounds, (; (Ũ⃗_keys .=> Ũ⃗_bounds)...))
+    # Put unit box bounds on the state if bound_state is true
+    if bound_state
+        state_dim = length(state_inits[1])
+        state_bounds = repeat([(-ones(state_dim), ones(state_dim))], length(state_names))
+        bounds = merge(bounds, (; (state_names .=> state_bounds)...))
     end
 
-    # Initial state and control values
+    # Trajectory
     if isnothing(a_guess)
-        Ũ⃗ = initialize_unitaries(U_init, U_goal, T, geodesic=geodesic)
-        Ũ⃗_values = repeat([Ũ⃗], length(Ũ⃗_keys))
-        a_values = initialize_controls(
+        # Randomly sample controls
+        a_values = initialize_control_trajectory(
             n_drives,
-            n_derivatives,
+            n_control_derivatives,
             T,
-            bounds[a_keys[1]],
+            bounds[control_name],
             drive_derivative_σ
         )
     else
-        @assert size(a_guess, 1) == n_drives "a_guess must have the same number of drives as n_drives"
-        @assert size(a_guess, 2) == T "a_guess must have the same number of timesteps as T"
-        @assert !isnothing(system) "system must be provided if a_guess is provided"
-
-        if Δt isa AbstractMatrix
-            ts = vec(Δt)
-        elseif Δt isa Float64
-            ts = fill(Δt, T)
-        else
-            ts = Δt
-        end
-
-        if system isa AbstractVector
-            @assert length(system) == length(Ũ⃗_keys) "systems must have the same length as Ũ⃗_keys"
-            Ũ⃗_values = map(system) do sys
-                unitary_rollout(Ũ⃗_init, a_guess, ts, sys; integrator=rollout_integrator)
-            end
-        else
-            Ũ⃗ = unitary_rollout(Ũ⃗_init, a_guess, ts, system; integrator=rollout_integrator)
-            Ũ⃗_values = repeat([Ũ⃗], length(Ũ⃗_keys))
-        end
-        Ũ⃗_values = Matrix{Float64}.(Ũ⃗_values)
-        a_values = initialize_controls(a_guess, ts, n_derivatives)
+        # Use provided controls and take derivatives
+        a_values = initialize_control_trajectory(a_guess, Δt, n_control_derivatives)
     end
 
-    # Trajectory
-    keys = [Ũ⃗_keys..., a_keys...]
-    values = [Ũ⃗_values..., a_values...]
+    names = [state_names..., control_names...]
+    values = [state_data..., a_values...]
 
     if free_time
-        push!(keys, :Δt)
+        push!(names, timestep_name)
         push!(values, Δt)
-        controls = (a_keys[end], :Δt)
-        bounds = merge(bounds, (Δt = Δt_bounds,))
+        controls = (control_names[end], timestep_name)
+        bounds = merge(bounds, (; timestep_name => Δt_bounds,))
     else
-        controls = (a_keys[end],)
+        controls = (control_names[end],)
     end
 
     return NamedTrajectory(
-        (; (keys .=> values)...);
+        (; (names .=> values)...);
         controls=controls,
         timestep=timestep,
         bounds=bounds,
@@ -373,98 +374,180 @@ function initialize_unitary_trajectory(
     )
 end
 
-function initialize_quantum_state_trajectory(
-    ψ̃_goals::AbstractVector{<:AbstractVector{<:Real}},
-    ψ̃_inits::AbstractVector{<:AbstractVector{<:Real}},
+function initialize_trajectory(
+    U_goal::OperatorType,
     T::Int,
-    Δt::Union{Real, AbstractVector{<:Real}},
-    n_drives::Int,
-    all_a_bounds::NamedTuple{anames, <:Tuple{Vararg{VectorBound}}} where anames;
-    n_derivatives::Int=2,
-    free_time=false,
-    Δt_bounds::ScalarBound=(0.5 * Δt, 1.5 * Δt),
-    drive_derivative_σ::Float64=0.1,
+    Δt::Union{Real, AbstractVecOrMat{<:Real}},
+    args...;
+    state_name::Symbol=:Ũ⃗,
+    state_names::AbstractVector{<:Symbol}=[state_name],
+    U_init::AbstractMatrix{<:Number}=Matrix{ComplexF64}(I(size(U_goal, 1))),
     a_guess::Union{AbstractMatrix{<:Float64}, Nothing}=nothing,
     system::Union{AbstractQuantumSystem, AbstractVector{<:AbstractQuantumSystem}, Nothing}=nothing,
-    global_data::Union{NamedTuple, Nothing}=nothing,
-    rollout_integrator::Function=exp,
-    ψ̃_keys::AbstractVector{<:Symbol}=[Symbol("ψ̃$i") for i = 1:length(ψ̃_goals)],
-    a_keys::AbstractVector{<:Symbol}=[Symbol("d"^i * "a") for i = 0:n_derivatives]
+    rollout_integrator::Function=expv,
+    geodesic=true,
+    kwargs...
 )
-    @assert length(ψ̃_inits) == length(ψ̃_goals) "ψ̃_inits and ψ̃_goals must have the same length"
-    @assert length(ψ̃_keys) == length(ψ̃_goals) "ψ̃_keys and ψ̃_goals must have the same length"
+    Ũ⃗_init = operator_to_iso_vec(U_init)
 
-    if free_time
-        if Δt isa Real
-            Δt = fill(Δt, 1, T)
-        elseif Δt isa AbstractVector
-            Δt = reshape(Δt, 1, :)
+    if U_goal isa EmbeddedOperator
+        Ũ⃗_goal = operator_to_iso_vec(U_goal.operator)
+    else
+        Ũ⃗_goal = operator_to_iso_vec(U_goal)
+    end
+
+    if isnothing(a_guess)
+        # No guess provided, initialize a geodesic and randomly sample controls
+
+        Ũ⃗_traj = initialize_unitary_trajectory(U_init, U_goal, T; geodesic=geodesic)
+
+        state_data = repeat([Ũ⃗_traj], length(state_names))
+    else
+        if system isa AbstractQuantumSystem
+            @assert size(a_guess, 1) == length(system.H_drives) "a_guess must have the same number of drives as n_drives"
+        elseif system isa AbstractVector
+            @assert size(a_guess, 1) == length(system[1].H_drives) "a_guess must have the same number of drives as n_drives"
+        end
+
+        @assert size(a_guess, 2) == T "a_guess must have the same number of timesteps as T"
+        @assert !isnothing(system) "system must be provided if a_guess is provided"
+
+        if Δt isa AbstractMatrix
+            timesteps = vec(Δt)
+        elseif Δt isa Float64
+            timesteps = fill(Δt, T)
         else
-            @assert size(Δt) == (1, T) "Δt must be a Real, AbstractVector, or 1x$(T) AbstractMatrix"
+            timesteps = Δt
+        end
+
+        if system isa AbstractVector
+            @assert length(system) == length(state_names) "systems must have the same length as state_names"
+            state_data = map(system) do sys
+                unitary_rollout(Ũ⃗_init, a_guess, timesteps, sys; integrator=rollout_integrator)
+            end
+        else
+            state = unitary_rollout(Ũ⃗_init, a_guess, timesteps, system; integrator=rollout_integrator)
+            state_data = [state]
         end
+        state_data = Matrix{Float64}.(state_data)
     end
 
-    # Constraints
-    state_initial = (; (ψ̃_keys .=> ψ̃_inits)...)
-    control_initial = (a = zeros(n_drives),)
-    initial = merge(state_initial, control_initial)
+    if system isa AbstractVector && length(state_names) != length(system)
+        state_names = [string(state_name) * "_system_$i" for i = 1:length(system)]
+        @warn "length of state_names and number of systems ($(length(system))) are not equal, created state names for each system: " state_names
+    end
 
-    final = (a = zeros(n_drives),)
+    state_inits = repeat([Ũ⃗_init], length(state_names))
+    state_goals = repeat([Ũ⃗_goal], length(state_names))
+
+    return initialize_trajectory(
+        state_data,
+        state_inits,
+        state_goals,
+        state_names,
+        T,
+        Δt,
+        args...;
+        a_guess=a_guess,
+        kwargs...
+    )
+end
 
-    goal = (; (ψ̃_keys .=> ψ̃_goals)...)
+function initialize_trajectory(
+    ψ_goals::AbstractVector{<:AbstractVector{ComplexF64}},
+    ψ_inits::AbstractVector{<:AbstractVector{ComplexF64}},
+    T::Int,
+    Δt::Union{Real, AbstractVector{<:Real}},
+    args...;
+    state_name=:ψ̃,
+    state_names::AbstractVector{<:Symbol}=length(ψ_goals) == 1 ?
+        [state_name] :
+        [Symbol(string(state_name) * "$i") for i = 1:length(ψ_goals)],
+    a_guess::Union{AbstractMatrix{<:Float64}, Nothing}=nothing,
+    system::Union{AbstractQuantumSystem, AbstractVector{<:AbstractQuantumSystem}, Nothing}=nothing,
+    rollout_integrator::Function=expv,
+    kwargs...
+)
+    @assert length(ψ_inits) == length(ψ_goals) "ψ_inits and ψ_goals must have the same length"
+    @assert length(state_names) == length(ψ_goals) "state_names and ψ_goals must have the same length"
 
-    # Bounds
-    bounds = all_a_bounds
+    ψ̃_goals = ket_to_iso.(ψ_goals)
+    ψ̃_inits = ket_to_iso.(ψ_inits)
 
-    # Initial state and control values
     if isnothing(a_guess)
-        ψ̃_values = NamedTuple([
-            k => linear_interpolation(ψ̃_init, ψ̃_goal, T)
-                for (k, ψ̃_init, ψ̃_goal) in zip(ψ̃_keys, ψ̃_inits, ψ̃_goals)
-        ])
-        a_values = initialize_controls(
-            n_drives,
-            n_derivatives,
-            T,
-            bounds[a_keys[1]],
-            drive_derivative_σ
-        )
+        state_data = []
+        for (ψ̃_init, ψ̃_goal) ∈ zip(ψ̃_inits, ψ̃_goals)
+            ψ̃_traj = linear_interpolation(ψ̃_init, ψ̃_goal, T)
+            push!(state_data, ψ̃_traj)
+        end
+        if system isa AbstractVector
+            state_data = repeat(state_data, length(system))
+        end
     else
-        ψ̃_values = NamedTuple([
-            k => rollout(ψ̃_init, a_guess, Δt, system, integrator=rollout_integrator)
-                for (k, ψ̃_init) in zip(ψ̃_keys, ψ̃_inits)
-        ])
-        a_values = initialize_controls(a_guess, Δt, n_derivatives)
+        @assert size(a_guess, 1) == n_drives "a_guess must have n_drives = $(n_drives) drives"
+        @assert size(a_guess, 2) == T "a_guess must have T = $(T) timesteps"
+        @assert !isnothing(system) "system must be provided if a_guess is provided"
+
+        if Δt isa AbstractMatrix
+            timesteps = vec(Δt)
+        elseif Δt isa Float64
+            timesteps = fill(Δt, T)
+        else
+            timesteps = Δt
+        end
+
+        if system isa AbstractVector
+            state_data = []
+            for sys ∈ system
+                for ψ̃_init ∈ ψ̃_inits
+                    ψ̃_traj = rollout(ψ̃_init, a_guess, timesteps, sys;
+                        integrator=rollout_integrator
+                    )
+                    push!(state_data, ψ̃_traj)
+                end
+            end
+        else
+            state_data = []
+            for ψ̃_init ∈ ψ̃_inits
+                ψ̃_traj = rollout(ψ̃_init, a_guess, timesteps, system;
+                    integrator=rollout_integrator
+                )
+                push!(state_data, ψ̃_traj)
+            end
+        end
     end
 
-    # Trajectory
-    keys = [ψ̃_keys..., a_keys...]
-    values = [ψ̃_values..., a_values...]
+    state_data = Matrix{Float64}.(state_data)
 
-    if free_time
-        push!(keys, :Δt)
-        push!(values, Δt)
-        controls = (a_keys[end], :Δt)
-        timestep = :Δt
-        bounds = merge(bounds, (Δt = Δt_bounds,))
-    else
-        controls = (a_keys[end],)
-        @assert Δt isa Real "Δt must be a Real if free_time is false"
-        timestep = Δt
+    if system isa AbstractVector
+        if lenth(state_names) != length(system) * length(ψ_goals)
+            state_names = vcat([
+                Symbol.(string.(state_names) .* "_system_$i")
+                    for i = 1:length(system)
+            ]...)
+            @warn "length of state_names and number of systems ($(length(system))) * number of states ($(length(ψ_goals))) are not equal, created state names for each system: " state_names
+        end
+        ψ̃_inits = repeat(ψ̃_inits, length(system))
+        ψ̃_goals = repeat(ψ̃_goals, length(system))
     end
 
-    return NamedTrajectory(
-        (; (keys .=> values)...);
-        controls=controls,
-        timestep=timestep,
-        bounds=bounds,
-        initial=initial,
-        final=final,
-        goal=goal,
-        global_data=isnothing(global_data) ? (;) : global_data
+    state_inits = ψ̃_inits
+    state_goals = ψ̃_goals
+
+    return initialize_trajectory(
+        state_data,
+        state_inits,
+        state_goals,
+        state_names,
+        T,
+        Δt,
+        args...;
+        a_guess=a_guess,
+        kwargs...
     )
 end
 
+
 # ============================================================================= #
 
 remove_component(
@@ -536,7 +619,7 @@ end
     drive_bounds = [1.0, 2.0]
     drive_derivative_σ = 0.01
 
-    a, da, dda = TrajectoryInitialization.initialize_controls(n_drives, n_derivates, T, drive_bounds, drive_derivative_σ)
+    a, da, dda = TrajectoryInitialization.initialize_control_trajectory(n_drives, n_derivates, T, drive_bounds, drive_derivative_σ)
 
     @test size(a) == (n_drives, T)
     @test size(da) == (n_drives, T)
@@ -625,10 +708,10 @@ end
     T = 10
     Δt = 0.1
     n_drives = 2
-    all_a_bounds = (a = [1.0, 1.0],)
+    a_bounds = ([1.0, 1.0],)
 
-    traj = initialize_unitary_trajectory(
-        U_goal, T, Δt, n_drives, all_a_bounds
+    traj = initialize_trajectory(
+        U_goal, T, Δt, n_drives, a_bounds
     )
 
     @test traj isa NamedTrajectory
@@ -637,16 +720,16 @@ end
 @testitem "quantum state trajectory initialization" begin
     using NamedTrajectories
 
-    ψ̃_init = ket_to_iso([0.0, 1.0])
-    ψ̃_goal = ket_to_iso([1.0, 0.0])
+    ψ_init = Vector{ComplexF64}([0.0, 1.0])
+    ψ_goal = Vector{ComplexF64}([1.0, 0.0])
 
     T = 10
     Δt = 0.1
     n_drives = 2
-    all_a_bounds = (a = [1.0, 1.0],)
+    all_a_bounds = ([1.0, 1.0],)
 
-    traj = initialize_quantum_state_trajectory(
-        [ψ̃_goal], [ψ̃_init], T, Δt, n_drives, all_a_bounds
+    traj = initialize_trajectory(
+        [ψ_goal], [ψ_init], T, Δt, n_drives, all_a_bounds
     )
 
     @test traj isa NamedTrajectory
diff --git a/src/trajectory_interpolations.jl b/src/trajectory_interpolations.jl
index 365a9b43..317590d2 100644
--- a/src/trajectory_interpolations.jl
+++ b/src/trajectory_interpolations.jl
@@ -16,9 +16,11 @@ struct DataInterpolation
     values_components::Vector{Int}
 
     function DataInterpolation(
-        times::AbstractVector{Float64}, values::AbstractMatrix{Float64}; 
-        timestep_components::AbstractVector{Int}=Int[], kind::Symbol=:linear
-    )   
+        times::AbstractVector{Float64},
+        values::AbstractMatrix{Float64};
+        timestep_components::AbstractVector{Int}=Int[],
+        kind::Symbol=:linear
+    )
         comps = setdiff(1:size(values, 1), timestep_components)
         if kind == :linear
             interpolants = [linear_interpolation(times, values[c, :]) for c in comps]
@@ -38,10 +40,10 @@ struct DataInterpolation
     end
 
     function DataInterpolation(
-        traj::NamedTrajectory; timestep_symbol::Symbol=:Δt, kwargs...
+        traj::NamedTrajectory; timestep_name::Symbol=:Δt, kwargs...
     )
-        if timestep_symbol ∈ keys(traj.components)
-            timestep_components = traj.components[timestep_symbol]
+        if timestep_name ∈ keys(traj.components)
+            timestep_components = traj.components[timestep_name]
         else
             timestep_components = Int[]
         end
@@ -90,7 +92,7 @@ end
     interp = DataInterpolation(free_traj)
     new_free_data = interp(get_times(traj))
 
-    # Replace the final timestep with the original value (can't be known a priori) 
+    # Replace the final timestep with the original value (can't be known a priori)
     new_free_data[free_traj.components.Δt, end] = free_traj.data[free_traj.components.Δt, end]
     @test new_free_data ≈ free_traj.data
 
@@ -116,4 +118,4 @@ end
 end
 
 
-end
\ No newline at end of file
+end