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