diff --git a/src/eval_grad_discrete_adjoint.jl b/src/eval_grad_discrete_adjoint.jl index f89e045..b2a1c80 100644 --- a/src/eval_grad_discrete_adjoint.jl +++ b/src/eval_grad_discrete_adjoint.jl @@ -66,10 +66,11 @@ function discrete_adjoint( prob::SchrodingerProb{<: AbstractMatrix{Float64}, <: AbstractMatrix{Float64}}, controls, pcof::AbstractVector{Float64}, target::AbstractMatrix{Float64}; - order=2, cost_type=:Infidelity, return_lambda_history=false + order=2, cost_type=:Infidelity, return_lambda_history=false, + kwargs... ) - history = eval_forward(prob, controls, pcof; order=order) + history = eval_forward(prob, controls, pcof; order=order, kwargs...) forcing = compute_guard_forcing(prob, history) @@ -83,8 +84,8 @@ function discrete_adjoint( forcing=forcing[:,end,:] ) - lambda_history = eval_adjoint( - prob, controls, pcof, terminal_condition, order=order, forcing=forcing) + lambda_history = eval_adjoint(prob, controls, pcof, terminal_condition; + order=order, forcing=forcing, kwargs...) grad = zeros(length(pcof)) @@ -114,6 +115,14 @@ function accumulate_gradient!(gradient::AbstractVector{Float64}, order=2 ) + #= + if (order == 2) + accumulate_gradient_order2!(gradient, prob, controls, pcof, history, lambda_history) + return gradient + end + =# + + dt = prob.tf / prob.nsteps N_derivatives = div(order, 2) @@ -123,7 +132,6 @@ function accumulate_gradient!(gradient::AbstractVector{Float64}, RHS = zeros(prob.real_system_size) LHS = zeros(prob.real_system_size) - # If this is really all it takes, that's pretty easy. for n in 0:prob.nsteps-1 # Used for both times @@ -158,6 +166,75 @@ function accumulate_gradient!(gradient::AbstractVector{Float64}, end end +""" +Hard-coded version for order 2 +""" +function accumulate_gradient_order2!(gradient::AbstractVector{Float64}, + prob::SchrodingerProb, controls, pcof::AbstractVector{Float64}, + history::AbstractArray{Float64, 3}, lambda_history::AbstractArray{Float64, 3} + ) + + dt = prob.tf / prob.nsteps + + u = zeros(prob.N_tot_levels) + v = zeros(prob.N_tot_levels) + lambda_u = zeros(prob.N_tot_levels) + lambda_v = zeros(prob.N_tot_levels) + + asym_op_lambda_u = zeros(prob.N_tot_levels) + asym_op_lambda_v = zeros(prob.N_tot_levels) + sym_op_lambda_u = zeros(prob.N_tot_levels) + sym_op_lambda_v = zeros(prob.N_tot_levels) + + for i in 1:prob.N_operators + control = controls[i] + asym_op = prob.asym_operators[i] + sym_op = prob.sym_operators[i] + local_pcof = get_control_vector_slice(pcof, controls, i) + + grad_contrib = zeros(control.N_coeff) + grad_p = zeros(control.N_coeff) + grad_q = zeros(control.N_coeff) + + for n in 0:prob.nsteps-1 + lambda_u .= @view lambda_history[1:prob.N_tot_levels, 1, 1+n+1] + lambda_v .= @view lambda_history[1+prob.N_tot_levels:end, 1, 1+n+1] + + t = n*dt + u .= @view history[1:prob.N_tot_levels, 1, 1+n] + v .= @view history[1+prob.N_tot_levels:end, 1, 1+n] + + eval_grad_p_derivative!(grad_p, control, t, local_pcof, 0) + eval_grad_q_derivative!(grad_q, control, t, local_pcof, 0) + + mul!(asym_op_lambda_u, asym_op, lambda_u) + mul!(asym_op_lambda_v, asym_op, lambda_v) + mul!(sym_op_lambda_u, sym_op, lambda_u) + mul!(sym_op_lambda_v, sym_op, lambda_v) + + grad_contrib .+= grad_q .* -(dot(u, asym_op_lambda_u) + dot(v, asym_op_lambda_v)) + grad_contrib .+= grad_p .* (-dot(u, sym_op_lambda_v) + dot(v, sym_op_lambda_u)) + + t = (n+1)*dt + u .= @view history[1:prob.N_tot_levels, 1, 1+n+1] + v .= @view history[1+prob.N_tot_levels:end, 1, 1+n+1] + + eval_grad_p_derivative!(grad_p, control, t, local_pcof, 0) + eval_grad_q_derivative!(grad_q, control, t, local_pcof, 0) + + grad_contrib .+= grad_q .* -(dot(u, asym_op_lambda_u) + dot(v, asym_op_lambda_v)) + grad_contrib .+= grad_p .* (-dot(u, sym_op_lambda_v) + dot(v, sym_op_lambda_u)) + end + + grad_contrib .*= -0.5*dt + + grad_slice = get_control_vector_slice(gradient, controls, i) + grad_slice .+= grad_contrib + end + + return nothing +end + """ History should be 4D array """ @@ -173,3 +250,4 @@ function compute_guard_forcing(prob, history) return forcing end + diff --git a/src/eval_grad_forced.jl b/src/eval_grad_forced.jl index 3fc9582..e43f45f 100644 --- a/src/eval_grad_forced.jl +++ b/src/eval_grad_forced.jl @@ -1,6 +1,6 @@ function eval_grad_forced(prob::SchrodingerProb{M, VM}, controls, pcof::AbstractVector{Float64}, target::VM; order=2, - cost_type=:Infidelity, return_forcing=false + cost_type=:Infidelity, return_forcing=false, kwargs... ) where {M <: AbstractMatrix{Float64}, VM <: AbstractVecOrMat{Float64}} # Allocate space for gradient @@ -19,7 +19,7 @@ function eval_grad_forced(prob::SchrodingerProb{M, VM}, controls, T = vcat(R[1+prob.N_tot_levels:end,:], -R[1:prob.N_tot_levels,:]) # Get state vector history - history = eval_forward(prob, controls, pcof, order=order) + history = eval_forward(prob, controls, pcof; order=order, kwargs...) final_state = history[:, 1, end, :] @@ -97,8 +97,8 @@ function eval_grad_forced(prob::SchrodingerProb{M, VM}, controls, # Compute the state history of ∂ψ/∂θₖ history_partial_derivative = eval_forward( - diff_prob, controls, pcof, forcing=forcing_ary, - order=order + diff_prob, controls, pcof; forcing=forcing_ary, + order=order, kwargs... ) #=