Adding script for running lab frame simulation.

leespen1 · leespen1 · commit b659d5820d1b · 2025-11-12T20:23:21.000-05:00
diff --git a/cnot3_labframe/cnot3_labframe.jl b/cnot3_labframe/cnot3_labframe.jl
@@ -0,0 +1,124 @@
+#=
+# Script for testing that setting up cnot3 matches the results of juqbox.
+# (which I assume to be correct)
+=#
+
+using QuantumGateDesign, DelimitedFiles
+using LinearAlgebra: Symmetric
+
+N_osc_levels = 10
+Tmax = 550.0
+nsteps = 1000
+
+subsystem_sizes = (N_osc_levels, 4, 4)
+essential_subsystem_sizes = (1, 2, 2)
+
+fa = 4.10595
+fb = 4.81526
+fs = 7.8447
+xa = 2 * 0.1099
+xb = 2 * 0.1126
+xs = 0.002494^2/xa
+xab = 1.0e-6
+xas = sqrt(xa*xs)
+xbs = sqrt(xb*xs)
+
+transition_freqs = (fs, fb, fa)
+rotation_freqs = transition_freqs # Rotating Frame
+
+kerr_coeffs = Symmetric(
+    [xs   xbs   xas;
+     xbs  xb    xab;
+     xas  xab   xa],
+    :U
+)
+
+
+prob_rot = dispersive_qudits_problem(
+    subsystem_sizes,
+    essential_subsystem_sizes,
+    transition_freqs,
+    rotation_freqs,
+    kerr_coeffs,
+    Tmax,
+    nsteps,
+    rot_frame=true,
+    gmres_abstol=1e-15,
+    gmres_reltol=1e-15,
+)
+
+prob_lab = dispersive_qudits_problem(
+    subsystem_sizes,
+    essential_subsystem_sizes,
+    transition_freqs,
+    rotation_freqs,
+    kerr_coeffs,
+    Tmax,
+    nsteps,
+    rot_frame=false,
+    gmres_abstol=1e-15,
+    gmres_reltol=1e-15,
+)
+
+rot_op_T = compsys_rot_frame_op(rotation_freqs, subsystem_sizes, Tmax)
+
+lab_target = (prob_lab.u0 + im.*prob_lab.v0)*CNOT_gate()
+rot_target = rot_op_T*lab_target
+
+seed=0
+atol=1e-15
+rtol=1e-15
+degree = 14
+D1 = 15
+
+cnot3ret = QuantumGateDesign.setup_cnot3(
+    seed=seed,
+    atol=atol,
+    rtol=rtol,
+    D1=D1,
+    N_osc_levels=N_osc_levels,
+    Tmax=Tmax
+)
+
+Nctrl = 3
+Nfreq = 3
+Cfreq = zeros(Nctrl,Nfreq)
+Cfreq[1:2,2] .= -2.0*pi*xa # carrier freq's for ctrl Hamiltonian 1 & 2
+Cfreq[1:2,3] .= -2.0*pi*xb # carrier freq's for ctrl Hamiltonian 1 & 2
+Cfreq[3,2] = -2.0*pi*xas # carrier freq 2 for ctrl Hamiltonian #3
+Cfreq[3,3] = -2.0*pi*xbs # carrier freq 2 for ctrl Hamiltonian #3
+
+Cfreq_lab_shift = repeat(collect(rotation_freqs), 1, Nfreq) .* 2.0*pi
+Cfreq_lab = Cfreq .+ Cfreq_lab_shift
+
+rot_controls = get_controls(degree, D1, Cfreq, Tmax)
+lab_controls = [DoubleRealPartControl(control)
+                for control in get_controls(degree, D1, Cfreq_lab, Tmax)]
+
+# Alternative implementation, which should be slower
+lab_controls2 = [LabFrameControl(rot_controls[i], 2.0*pi*rotation_freqs[i])
+                 for i in 1:3]
+
+pcof_csv = readdlm("targetError=1e-7_cnot3OptimizationTest_order=6_degree=14_seed=5_nsteps=5409_atol=1.0e-15_rtol=1.0e-15_D1=16_time=6.0_maxiter=10000_nthreads=4_costType=GeneralizedInfidelity_gateDuration=550.0_nCavityLevels=10_pcof.csv", ',')
+pcof = pcof_csv[end,:]
+# Dummy run to compile
+prob_rot.nsteps = 1
+prob_lab.nsteps = 1
+history_rot = eval_forward(prob_rot, rot_controls, pcof, order=6)
+history_lab = eval_forward(prob_lab, lab_controls, pcof, order=6)
+
+# Actual run
+prob_rot.nsteps = 100
+prob_lab.nsteps = 100
+#@time history_rot = eval_forward(prob_rot, rot_controls, pcof, order=6)
+@time history_lab = eval_forward(prob_lab, lab_controls, pcof, order=6)
+
+
+include("stepsize_copy.jl")
+order = 6
+max_walltime = 1 # In hours
+filename_base = "lab_frame"
+nsaves = 10
+initial_nsteps = 50_000 # Nyquist limit should be about here
+collect_data(prob_lab, lab_controls, pcof, order, max_walltime, filename_base,
+             nsaves, initial_nsteps)
diff --git a/cnot3_labframe/stepsize_copy.jl b/cnot3_labframe/stepsize_copy.jl
@@ -0,0 +1,241 @@
+#==========================================
+# Copy of the data collection functions from the cnot3 stepsize test
+==========================================#
+using QuantumGateDesign, DelimitedFiles
+
+
+function collect_data(prob::SchrodingerProb, controls::ControlsType,
+        pcof::AbstractVector{<: Real}, order::Integer, max_walltime::Real,
+        filename_base::AbstractString, N_timestep_saves::Integer, initial_nsteps=2
+    )
+    prob = copy(prob) # Copy problem, just to make sure there are no mutability issues.
+
+    initial_time = time()
+    max_walltime *= 60*60 # convert walltime from hours to seconds
+    
+    filename_csv = filename_base * ".csv"
+    filename_final_states_csv = filename_base * "_finalStates.csv"
+
+    header = hcat(
+        "nsteps", "stepsize", "elapsed_time", "avg_N_gmres_iter",
+        "N_converged_gmres", "avg_gmres_residual", "max_gmres_residual",
+        "R_abs_err_L1", "R_abs_err_L2", "R_rel_err_L1", "R_rel_err_L2",
+        "R_abs_err_Linf",
+    )
+    println(stdout, header)
+
+    writedlm(filename_csv, rpad.(header, 24), ',')
+
+    final_states_vec = Vector{ComplexF64}(undef, length(prob.u0))
+    final_states_vec .= NaN
+    gmres_tracker = GMRESTracker()
+
+    # Run simulation
+    prob.nsteps = initial_nsteps
+    stepsize = prob.tf / prob.nsteps
+
+    # Run simulation once just to get compilation out of the way
+    dummy_history = eval_forward(prob, controls, pcof, order=order)
+
+    is_first_step = true
+    history_2h = nothing
+    history_h = nothing
+    elapsed_time = 0.0
+
+    # Loop until time runs out (with estimator for when we will go overtime on next simulation)
+    while (time()-initial_time) < (max_walltime - 2*elapsed_time)
+        # Run simulation
+        t1 = time()
+        history_h = eval_forward(prob, controls, pcof, order=order,
+                                 gmres_tracker=gmres_tracker, verbose=true)
+        t2 = time()
+        elapsed_time = t2 - t1
+
+        # Collect data into a row
+        if !is_first_step
+            R = RichardsonExtrapolation(history_h[:,1:2:end,:], history_2h, order)
+            csv_row = hcat(
+                prob.nsteps, stepsize, elapsed_time,
+                avg_N_iterations(gmres_tracker), gmres_tracker.N_converged,
+                avg_residual(gmres_tracker), R.abs_err_L1, R.abs_err_L2,
+                R.rel_err_L1, R.rel_err_L2, R.abs_err_Linf,
+            )
+        else
+            csv_row = hcat(
+                prob.nsteps, stepsize, elapsed_time,
+                avg_N_iterations(gmres_tracker), gmres_tracker.N_converged,
+                avg_residual(gmres_tracker), NaN, NaN, NaN, NaN, NaN, 
+            )
+            is_first_step = false
+        end
+
+        # Log data (CSV)
+        open(filename_csv, "a+") do io
+            writedlm(io, rpad.(csv_row, 24), ',')
+        end
+        final_states_vec .= reshape(history_h[:,end,:], :)
+        open(filename_final_states_csv, "a+") do io
+            writedlm(io, rpad.(transpose(final_states_vec), 53), ',')
+        end
+        println(stdout, csv_row) # Print row
+
+        # Prepare for next iteration
+        history_2h = history_h
+        prob.nsteps *= 2
+        stepsize = prob.tf / prob.nsteps
+    end
+
+    return readdlm(filename_csv, ',')
+end
+
+
+function collect_data_grad(prob::SchrodingerProb, controls::ControlsType,
+        pcof::AbstractVector{<: Real}, target::AbstractMatrix{<: Number}, order::Integer, max_walltime::Real,
+        filename_base::AbstractString, N_timestep_saves::Integer
+    )
+    prob = copy(prob) # Copy problem, just to make sure there are no mutability issues.
+
+    initial_time = time()
+    max_walltime *= 60*60 # convert walltime from hours to seconds
+    
+    filename_csv = filename_base * ".csv"
+    filename_final_states_csv = filename_base * "_finalStates.csv"
+
+    header = hcat(
+        "nsteps", "stepsize", "elapsed_time", "forward_time", "adjoint_time",
+        "grad_accum_time", "avg_N_gmres_iter_fwd", "N_converged_gmres_fwd",
+        "avg_gmres_residual_fwd", "avg_N_gmres_iter_adj", "N_converged_gmres_adj",
+        "avg_gmres_residual_adj", "R_abs_err_L1", "R_abs_err_L2", "R_rel_err_L1",
+        "R_rel_err_L2", "R_abs_err_Linf",
+    )
+    println(stdout, header)
+
+    writedlm(filename_csv, rpad.(header, 24), ',')
+
+    final_states_vec = Vector{ComplexF64}(undef, length(prob.u0))
+    final_states_vec .= NaN
+
+    forward_gmres_tracker = GMRESTracker()
+    adjoint_gmres_tracker = GMRESTracker()
+
+    # Run simulation
+    prob.nsteps = 2
+    stepsize = prob.tf / prob.nsteps
+
+    # Run simulation once just to get compilation out of the way
+    dummy_history = eval_forward(prob, controls, pcof, order=order)
+
+    timer = QuantumGateDesign.DiscreteAdjointTimes()
+    is_first_step = true
+    history_2h = nothing
+    history_h = nothing
+    elapsed_time = 0.0
+    N_derivatives = div(order,2)
+    grad = zeros(get_number_of_control_parameters(controls))
+
+    # Loop until time runs out (with estimator for when we will go overtime on next simulation)
+    while (time()-initial_time) < (max_walltime - 2*elapsed_time)
+        # Run simulation
+        history = zeros(prob.real_system_size, 1+N_derivatives, 1+prob.nsteps, prob.N_initial_conditions)
+        lambda_history = zeros(prob.real_system_size, 1+N_derivatives, 1+prob.nsteps, prob.N_initial_conditions)
+        adjoint_forcing = zeros(prob.real_system_size, 1+prob.nsteps, prob.N_initial_conditions)
+
+        QuantumGateDesign.discrete_adjoint!(
+            grad, history, lambda_history, adjoint_forcing, prob, controls,
+            pcof, target, order=order, timer=timer,
+            forward_gmres_tracker=forward_gmres_tracker,
+            adjoint_gmres_tracker=adjoint_gmres_tracker,
+        )
+        history_h = QuantumGateDesign.real_to_complex(history[:,1,:,:])
+
+        # Collect data into a row
+        if !is_first_step
+            R = RichardsonExtrapolation(history_h[:,1:2:end,:], history_2h, order)
+            csv_row = hcat(
+                prob.nsteps, stepsize, QuantumGateDesign.total_time(timer),
+                timer.forward, timer.adjoint, timer.grad_accum,
+                avg_N_iterations(forward_gmres_tracker),
+                forward_gmres_tracker.N_converged,
+                avg_residual(forward_gmres_tracker),
+                avg_N_iterations(adjoint_gmres_tracker),
+                adjoint_gmres_tracker.N_converged,
+                avg_residual(adjoint_gmres_tracker), R.abs_err_L1,
+                R.abs_err_L2, R.rel_err_L1, R.rel_err_L2, R.abs_err_Linf,
+            )
+        else
+            # Rerun to update timing (now that precompilation is out of the way)
+            QuantumGateDesign.discrete_adjoint!(
+                grad, history, lambda_history, adjoint_forcing, prob, controls,
+                pcof, target, order=order, timer=timer,
+                forward_gmres_tracker=forward_gmres_tracker,
+                adjoint_gmres_tracker=adjoint_gmres_tracker,
+            )
+
+            csv_row = hcat(
+                prob.nsteps, stepsize, QuantumGateDesign.total_time(timer),
+                timer.forward, timer.adjoint, timer.grad_accum,
+                avg_N_iterations(forward_gmres_tracker),
+                forward_gmres_tracker.N_converged,
+                avg_residual(forward_gmres_tracker),
+                avg_N_iterations(adjoint_gmres_tracker),
+                adjoint_gmres_tracker.N_converged,
+                avg_residual(adjoint_gmres_tracker), NaN,
+                NaN, NaN, NaN, NaN,
+            )
+            is_first_step = false
+        end
+
+        # Log data (CSV)
+        open(filename_csv, "a+") do io
+            writedlm(io, rpad.(csv_row, 24), ',')
+        end
+        final_states_vec .= reshape(history_h[:,end,:], :)
+        open(filename_final_states_csv, "a+") do io
+            writedlm(io, rpad.(transpose(final_states_vec), 53), ',')
+        end
+        println(stdout, csv_row) # Print row
+
+        # Prepare for next iteration
+        history_2h = history_h
+        prob.nsteps *= 2
+        stepsize = prob.tf / prob.nsteps
+    end
+
+    return readdlm(filename_csv, ',')
+end
+
+
+
+"""
+Turn state vector history into a row matrix with the correct number of timesteps
+saved, for putting into a csv file. 
+
+If the number of timesteps to be saved is greater than the number of timesteps
+in the provided history, use NaN in place of the 'missing' timestep saves
+"""
+function parse_history_for_csv(history::AbstractArray{ComplexF64, 3}, N_timestep_saves::Union{Missing, Integer})
+    if ismissing(N_timestep_saves)
+        return copy(reshape(history, 1, :))
+    end
+
+    parsed_history = Array{ComplexF64, 3}(undef, size(history, 1), 1+N_timestep_saves, size(history, 3))
+    parsed_history .= NaN
+
+    nsteps = size(history, 2) - 1
+
+    if  nsteps >= N_timestep_saves
+        if (nsteps % N_timestep_saves != 0)
+            throw(ArgumentError("Number of timesteps to save ($N_timestep_saves) does not divide the number of timesteps ($nsteps)"))
+        end
+        stride = div(nsteps, N_timestep_saves)
+        parsed_history .= history[:,1:stride:end,:]
+    else
+        if (N_timestep_saves % nsteps != 0)
+            throw(ArgumentError("Number of timesteps ($nsteps) does not divide the number of timesteps to save ($N_timestep_saves)"))
+        end
+        stride = div(N_timestep_saves, nsteps)
+        parsed_history[:,1:stride:end,:] .= history
+    end
+
+    return reshape(parsed_history, 1, :)
+end
