Skip to content

algorithms #19

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hmyuuu merged 2 commits into from
Mar 16, 2022
Merged

algorithms #19

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
201a0df
algorithms
zhangmao2018 Mar 16, 2022
c4c758a
revise adam
hmyuuu Mar 16, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
163 changes: 56 additions & 107 deletions src/algorithm/AD.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
#### AD ####
function update!(opt::ControlOpt, alg::AD, obj, dynamics, output)
(; max_episode, ϵ, beta1, beta2) = alg
#### control optimization ####
function update!(opt::ControlOpt, alg::AbstractAD, obj, dynamics, output)
(; max_episode) = alg
ctrl_length = length(dynamics.data.ctrl[1])
ctrl_num = length(dynamics.data.Hc)

Expand All @@ -13,23 +13,21 @@ function update!(opt::ControlOpt, alg::AD, obj, dynamics, output)
set_io!(output, f_noctrl, f_ini)
show(opt, output, obj)

output.f_list = [f_ini]
for ei = 1:(max_episode-1)
# δ = grad(obj, dynamics)
δ = gradient(() -> objective(obj, dynamics)[2], Flux.Params([dynamics.data.ctrl]))
update_ctrl!(alg, obj, dynamics, δ[dynamics.data.ctrl])
bound!(dynamics.data.ctrl, opt.ctrl_bound)
f_out, f_now = objective(obj, dynamics)

set_f!(output, f_out)
set_buffer!(output, [dynamics.data.ctrl])
set_buffer!(output, dynamics.data.ctrl)
set_io!(output, f_out, ei)
show(output, obj)
end
set_io!(output, output.f_list[end])
end

function update_ctrl!(alg::AD{Adam}, obj, dynamics, δ)
function update_ctrl!(alg::AD_Adam, obj, dynamics, δ)
(; ϵ, beta1, beta2) = alg
for ci in 1:length(δ)
mt, vt = 0.0, 0.0
Expand All @@ -40,79 +38,53 @@ function update_ctrl!(alg::AD{Adam}, obj, dynamics, δ)
end
end

function update_ctrl!(alg::AD{GradDescent}, obj, dynamics, δ)
(; ϵ) = alg
dynamics.data.ctrl += ϵ*δ
function update_ctrl!(alg::AD, obj, dynamics, δ)
dynamics.data.ctrl += alg.ϵ*δ
end

#### state optimization ####
function update!(opt::StateOpt, alg::AD, obj, dynamics, output)
function update!(opt::StateOpt, alg::AbstractAD, obj, dynamics, output)
(; max_episode) = alg
f_ini, f_comp = objective(obj, dynamics)
set_f!(output, f_ini)
set_buffer!(output, dynamics.data.ψ0)
set_io!(output, f_ini)
show(opt, output, obj)
for ei in 1:(max_episode-1)
δ = grad(opt, obj, dynamics)
update_state!(alg, obj, dynamics, δ)
δ = gradient(() -> objective(obj, dynamics)[2], Flux.Params([dynamics.data.ψ0]))
update_state!(alg, obj, dynamics, δ[dynamics.data.ψ0])
dynamics.data.ψ0 = dynamics.data.ψ0/norm(dynamics.data.ψ0)
f_out, f_now = objective(obj, dynamics)
set_output!(output, f_out)
set_f!(output, f_out)
set_buffer!(output, dynamics.data.ψ0)
set_io!(output, f_out, ei)
show(output, obj)
end
show(output, output)
set_io!(output, output.f_list[end])
end

function update_state!(alg::AD{Adam}, obj, dynamics, δ)
function update_state!(alg::AD_Adam, obj, dynamics, δ)
(; ϵ, beta1, beta2) = alg
mt, vt = 0.0, 0.0
for ti in 1:length(δ)
dynamics.data.ψ0[ti], mt, vt = Adam(δ[ti], ti, dynamics.data.ψ0[ti], mt, vt, ϵ, beta1, beta2, obj.eps)
end
end

function update_state!(alg::AD{GradDescent}, obj, dynamics, δ)
(;ϵ) = alg
dynamics.data.ψ0 += ϵ*δ
end

function grad(opt, obj, dynamics::Lindblad{noiseless,free,ket})
(; H0, dH, ψ0, tspan) = dynamics.data
δ = gradient(x->objective(obj, H0, dH, x, tspan), ψ0).|>real |>sum
end

function grad(opt, obj, dynamics::Lindblad{noisy,free,ket})
(;H0, dH, ψ0, tspan, decay_opt, γ) = dynamics.data
δ = gradient(x->objective(obj, H0, dH, x, tspan, decay_opt, γ), ψ0).|>real |>sum
end

function grad(opt, obj, dynamics::Lindblad{noiseless,timedepend,ket})
(; H0, dH, ψ0, tspan) = dynamics.data
δ = gradient(x->objective(obj, H0, dH, x, tspan), ψ0).|>real |>sum
end

function grad(opt, obj, dynamics::Lindblad{noisy,timedepend,ket})
(;H0, dH, ψ0, tspan, decay_opt, γ) = dynamics.data
δ = gradient(x->objective(obj, H0, dH, x, tspan, decay_opt, γ), ψ0).|>real |>sum
end

function grad(opt, obj, dynamics::Kraus{noiseless,free,ket})
(;K, dK, ψ0) = dynamics.data
δ = gradient(x->objective(obj, K, dK, x), ψ0).|>real |>sum
function update_state!(alg::AD, obj, dynamics, δ)
dynamics.data.ψ0 += alg.ϵ*δ
end

#### measurement optimization (linear conbination) ####
function update!(opt::Mopt_LinearComb, alg::AD, obj, dynamics, output, POVM_basis, M_num)
(; max_episode, rng) = alg
#### find the optimal linear combination of a given set of POVM ####
function update!(opt::Mopt_LinearComb, alg::AbstractAD, obj, dynamics, output)
(; max_episode) = alg
(; POVM_basis, M_num) = opt
basis_num = length(POVM_basis)
# initialize
B = [rand(rng, basis_num) for i in 1:M_num]
B = bound_LC_coeff!(B)

M = [sum([B[i][j]*POVM_basis[j] for j in 1:basis_num]) for i in 1:M_num]
f_opt, f_comp = objective(obj::QFIM{SLD}, dynamics)
bound_LC_coeff!(opt.B)
M = [sum([opt.B[i][j]*POVM_basis[j] for j in 1:basis_num]) for i in 1:M_num]
obj_QFIM = QFIM_Obj(obj)
f_opt, f_comp = objective(obj_QFIM, dynamics)
obj_POVM = set_M(obj, POVM_basis)
f_povm, f_comp = objective(obj_POVM, dynamics)
obj_copy = set_M(obj, M)
Expand All @@ -122,60 +94,48 @@ function update!(opt::Mopt_LinearComb, alg::AD, obj, dynamics, output, POVM_basi
set_io!(output, f_opt, f_povm, f_ini)
show(opt, output, obj)
for ei in 1:(max_episode-1)
δ = grad(opt, obj, dynamics, B, POVM_basis, M_num)
update_M!(alg, obj, dynamics, δ, B)
B = bound_LC_coeff!(B)
M = [sum([B[i][j]*POVM_basis[j] for j in 1:basis_num]) for i in 1:M_num]
δ = gradient(() -> objective(opt, obj, dynamics)[2], Flux.Params([opt.B]))
update_M!(opt, alg, obj, δ[opt.B])
bound_LC_coeff!(opt.B)
M = [sum([opt.B[i][j]*POVM_basis[j] for j in 1:basis_num]) for i in 1:M_num]
obj_copy = set_M(obj, M)
f_out, f_now = objective(obj_copy, dynamics)
set_output!(output, f_out)
set_f!(output, f_out)
set_buffer!(output, M)
set_io!(output, f_out, ei)
show(output, obj)
end
show(output, output)
set_io!(output, output.f_list[end])
end

function update_M!(alg::AD{Adam}, obj, dynamics, δ, B::Vector{Vector{Float64}})
function update_M!(opt::Mopt_LinearComb, alg::AD_Adam, obj, δ)
(; ϵ, beta1, beta2) = alg
for ci in 1:length(δ)
mt, vt = 0.0, 0.0
for ti in 1:length(δ[1])
B[ci][ti], mt, vt = Adam(δ[ci][ti], ti, B[ci][ti], mt, vt, ϵ, beta1, beta2, obj.eps)
opt.B[ci][ti], mt, vt = Adam(δ[ci][ti], ti, opt.B[ci][ti], mt, vt, ϵ, beta1, beta2, obj.eps)
end
end
end

function update_M!(alg::AD{GradDescent}, obj, dynamics, δ, B::Vector{Vector{Float64}})
(;ϵ) = alg
B += ϵ*δ
function update_M!(opt::Mopt_LinearComb, alg::AD, obj, δ)
opt.B += alg.ϵ*δ
end

function grad(opt, obj, dynamics, B, POVM_basis, M_num::Number)
(;H0, dH, ρ0, tspan, decay_opt, γ) = dynamics.data
basis_num = length(POVM_basis)
δ = gradient(x->objective(obj, [sum([x[i][j]*POVM_basis[j] for j in 1:basis_num]) for i in 1:M_num], H0, dH, ρ0, tspan, decay_opt, γ), B).|>real |>sum
end

function grad(opt, obj, dynamics, B, POVM_basis, M_num::Number)
(;K, dK, ρ0) = dynamics.data
basis_num = length(POVM_basis)
δ = gradient(x->objective(obj, [sum([x[i][j]*POVM_basis[j] for j in 1:basis_num]) for i in 1:M_num], K, dK, ρ0), B).|>real |>sum
end

#### measurement optimization (rotation) ####
function update!(opt::Mopt_Rotation, alg::AD, obj, dynamics, output, POVM_basis)
(; max_episode, rng) = alg
#### find the optimal rotated measurement of a given set of POVM ####
function update!(opt::Mopt_Rotation, alg::AbstractAD, obj, dynamics, output)
(; max_episode) = alg
(; POVM_basis) = opt
dim = size(dynamics.data.ρ0)[1]
suN = suN_generator(dim)
Lambda = [Matrix{ComplexF64}(I,dim,dim)]
append!(Lambda, [suN[i] for i in 1:length(suN)])

M_num = length(POVM_basis)
suN = suN_generator(dim)
append!(opt.Lambda, [Matrix{ComplexF64}(I,dim,dim)])
append!(opt.Lambda, [suN[i] for i in 1:length(suN)])

s = rand(rng, dim*dim)
U = rotation_matrix(s, Lambda)
U = rotation_matrix(opt.s, opt.Lambda)
M = [U*POVM_basis[i]*U' for i in 1:M_num]
f_opt, f_comp = objective(obj::QFIM{SLD}, dynamics)
obj_QFIM = QFIM_Obj(obj)
f_opt, f_comp = objective(obj_QFIM, dynamics)
obj_POVM = set_M(obj, POVM_basis)
f_povm, f_comp = objective(obj_POVM, dynamics)
obj_copy = set_M(obj, M)
Expand All @@ -185,40 +145,29 @@ function update!(opt::Mopt_Rotation, alg::AD, obj, dynamics, output, POVM_basis)
set_io!(output, f_opt, f_povm, f_ini)
show(opt, output, obj)
for ei in 1:(max_episode-1)
δ = grad(opt, obj, dynamics, s, POVM_basis, Lambda)
update_M!(alg, obj, dynamics, δ, s)
s = bound_rot_coeff!(s)
U = rotation_matrix(s, Lambda)
δ = gradient(() -> objective(opt, obj, dynamics)[2], Flux.Params([opt.s]))
update_M!(opt, alg, obj, δ[opt.s])
bound_rot_coeff!(opt.s)
U = rotation_matrix(opt.s, opt.Lambda)
M = [U*POVM_basis[i]*U' for i in 1:M_num]
obj_copy = set_M(obj, M)
f_out, f_now = objective(obj_copy, dynamics)

set_output!(output, f_out)
set_f!(output, f_out)
set_buffer!(output, M)
set_io!(output, f_out, ei)
show(output, obj)
end
show(output, output)
set_io!(output, output.f_list[end])
end

function update_M!(alg::AD{Adam}, obj, dynamics, δ, s::Vector{Float64})
function update_M!(opt::Mopt_Rotation, alg::AD_Adam, obj, δ)
(; ϵ, beta1, beta2) = alg
mt, vt = 0.0, 0.0
for ti in 1:length(δ)
s[ti], mt, vt = Adam(δ[ti], ti, s[ti], mt, vt, ϵ, beta1, beta2, obj.eps)
opt.s[ti], mt, vt = Adam(δ[ti], ti, opt.s[ti], mt, vt, ϵ, beta1, beta2, obj.eps)
end
end

function update_M!(alg::AD{GradDescent}, obj, dynamics, δ, s::Vector{Float64})
(;ϵ) = alg
s += ϵ*δ
end

function grad(opt, obj, dynamics, s, POVM_basis, Lambda::AbstractArray)
(;H0, dH, ρ0, tspan, decay_opt, γ) = dynamics.data
δ = gradient(x->objective(obj, x, POVM_basis, Lambda, H0, dH, ρ0, tspan, decay_opt, γ), s).|>real |>sum
end

function grad(opt, obj, dynamics, s, POVM_basis, Lambda::AbstractArray)
(;K, dK, ρ0) = dynamics.data
δ = gradient(x->objective(obj, x, POVM_basis, Lambda, K, dK, ρ0), s).|>real |>sum
function update_M!(opt::Mopt_Rotation, alg::AD, obj, δ)
opt.s += alg.ϵ*δ
end
20 changes: 13 additions & 7 deletions src/algorithm/DDPG.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,6 @@ using StableRNGs
using Flux.Losses
using IntervalSets

RLBase.action_space(env) = env.action_space
RLBase.state_space(env) = env.state_space
RLBase.reward(env) = env.reward
RLBase.is_terminated(env) = env.done
RLBase.state(env) = env.state

function Base.rsplit( v, l::Int)
u = reshape(v,l,length(v)÷l)
[u[:,i] for i=1:size(u,2)]
Expand Down Expand Up @@ -164,6 +158,12 @@ function update!(opt::ControlOpt, alg::DDPG, obj, dynamics, output)
end
end

RLBase.action_space(env::ControlEnv) = env.action_space
RLBase.state_space(env::ControlEnv) = env.state_space
RLBase.reward(env::ControlEnv) = env.reward
RLBase.is_terminated(env::ControlEnv) = env.done
RLBase.state(env::ControlEnv) = env.state

function RLBase.reset!(env::ControlEnv)
state = env.dynamics.data.ρ0
env.dstate = [state |> zero for _ = 1:(env.para_num)]
Expand Down Expand Up @@ -224,6 +224,12 @@ mutable struct StateEnv{T<:Complex, M<:Real, R<:AbstractRNG}
episode::Int
end

RLBase.action_space(env::StateEnv) = env.action_space
RLBase.state_space(env::StateEnv) = env.state_space
RLBase.reward(env::StateEnv) = env.reward
RLBase.is_terminated(env::StateEnv) = env.done
RLBase.state(env::StateEnv) = env.state

function update!(Sopt::StateOpt, alg::DDPG, obj, dynamics, output)
(; max_episode, layer_num, layer_dim, rng) = alg
episode = 1
Expand Down Expand Up @@ -271,7 +277,7 @@ function update!(Sopt::StateOpt, alg::DDPG, obj, dynamics, output)

stop_condition = StopAfterStep(max_episode, is_show_progress=false)
hook = TotalRewardPerEpisode(is_display_on_exit=false)
run(agent, env, stop_condition, hook)
RLBase.run(agent, env, stop_condition, hook)

show(output, output)
if save_type(output) == :no_save
Expand Down
Loading