Skip to content

Instantly share code, notes, and snippets.

@goerz
Created March 20, 2024 21:24
Show Gist options
  • Save goerz/1e226e11c1df29a775fe23d76d0f9ed0 to your computer and use it in GitHub Desktop.
Save goerz/1e226e11c1df29a775fe23d76d0f9ed0 to your computer and use it in GitHub Desktop.
Julia Parametrization MWE
# This MWE is to explore the use of RecursiveArrayTools for optimizing
# parameters for control functions.
#
# We want two things:
#
# * Extract control parameters that might occur in a set of Hamiltonians, each
# Hamiltonian containing multiple controls, and each control containing an
# AbstractVector of parameters (a ComponentArray seems very useful, but it
# could be any AbstractVector). This should be done with a function
# `get_parameters()` that collects all control parameters in a
# single vector. That vector does not need to be contiguous.
# * Given a contiguous Vector `x` that is compatible with what
# `get_parameters()` returns we have to be able to push the value `x` into
# the parameters deep inside the Hamiltonian data structure
#
# The `RecursiveArrayTools.ArrayPartition` seems perfect for this:
#
# * It allows `get_parameters()` to return a non-contiguous `AbstractVector`
# where mutating the vectors mutates the original value deep inside the
# Hamiltian
# * Given a contiguous vector `x`, we can push the value of `x` into the
# Hamiltonian with a simplex `u .= x`, where `u` is `ArrayPartition` returned
# by `get_parameters()`.
using Pkg
Pkg.activate(temp=true)
Pkg.add("RecursiveArrayTools")
Pkg.add("ComponentArrays")
#Pkg.add("Parameters")
Pkg.add("UnPack")
Pkg.add("Optimization")
Pkg.add("OptimizationNLopt")
using ComponentArrays
using RecursiveArrayTools
#using Parameters: @unpack
using UnPack: @unpack
using Optimization
using OptimizationNLopt: NLopt
struct BoxyField
parameters::ComponentVector{Float64,Vector{Float64},Tuple{Axis{(E₀=1, a=2, t₁=3)}}}
# Using a ComponentArray for the `parameters` isn't essential (it could be
# *any* AbstractVector), but ComponentArray does seem very useful for this
# in real life, and we want to make sure here that it composes well with
# `get_parameters`.
end
BoxyField(; E₀=1.0, a=10.0, t₁=0.25) = BoxyField(ComponentArray(; E₀, a, t₁))
BoxyField(p::NamedTuple) = BoxyField(ComponentArray(p))
function (E::BoxyField)(t)
@unpack E₀, a, t₁ = E.parameters
t₂ = 1.0 - t₁
return (E₀ / 2) * (tanh(a * (t - t₁)) - tanh(a * (t - t₂)))
end
struct Hamiltonian
# The actual Hamiltonian contains operators in addition to controls, but we
# don't need that for this MWE
controls
end
# TODO: open issue for ComponentArray: it would be good to have
# Base.convert(::Type{ComponentArray{T, N, A, Ax}}, x::NT) where {T, N, A, Ax, NT<:NamedTuple} = ComponentArray(x)
# get AbstractArray so that mutating the array mutates the parameters deep
# inside the Hamiltonian
function get_parameters(H::Hamiltonian)
parameters = RecursiveArrayTools.ArrayPartition([E.parameters for E in H.controls]...)
return parameters
end
E₁ = BoxyField(E₀=20.00, a=7.5, t₁=0.3)
E₂ = BoxyField(E₀=20.00, a=7.5, t₁=0.3)
E₃ = BoxyField(E₀=20.00, a=7.5, t₁=0.3)
# The following globals would be closures in my actual code. I'm not overly
# concerned about performance-problems due to closures. My actual `func` is
# often expensive, and most of the actual work is done by calling BLAS
# functions.
HAMILTONIAN = Hamiltonian([E₁, E₂, E₃])
PARAMETERS = get_parameters(HAMILTONIAN)
TLIST = collect(range(0, 1; length=101));
# function to minimize
function func(x, tlist)
@assert PARAMETERS !== x # for Optimization.jl; LBFGSB may allow aliasing here
PARAMETERS .= x # XXX
# The above line is is where we really exploit ArrayPartion: it seems like
# the most convention way to push the value in the Vector `x` into the
# parameters inside the controls inside the Hamiltonian
target_pulse(t) = 20 * exp(-20 * (t - 0.5)^2)
mismatch = 0.0
for E in HAMILTONIAN.controls
mismatch += sum(abs2.(E.(tlist) .- target_pulse.(tlist)))
end
println(mismatch)
return mismatch
end
# Guess (3 different possibilities – they all work, at least with Optimization.jl)
x0 = copy(get_parameters(HAMILTONIAN));
@assert x0 !== PARAMETERS;
x0 = Array(x0)
# x0 = PARAMETERS # with this, the optimization will mutate x0
##### Optimization.jl — Simplex
prob = OptimizationProblem(func, x0, TLIST; stopval=50,);
res = Optimization.solve(prob, NLopt.LN_NELDERMEAD(), maxiters=1000)
##### Optimization.jl – Gradient-Based
function grad!(g, x, tlist)
N = length(x)
params_per_control = 3
N_controls = N ÷ params_per_control
for i = 1:N_controls
offset = (i - 1) * params_per_control
E₀, a, t₁ = x[offset+1:offset+params_per_control]
g[offset+1] = sum([d_E0(t, E₀, a, t₁) for t in tlist])
g[offset+2] = sum([d_a(t, E₀, a, t₁) for t in tlist])
g[offset+3] = sum([d_t1(t, E₀, a, t₁) for t in tlist])
end
end
function d_E0(t, E_0, a, t_1) # sympy codegen
t_2 = 1.0 - t_1
return (
E_0 .* (tanh(a .* (t - t_1)) - tanh(a .* (t - t_2))) .* exp(5 * (2 * t - 1) .^ 2) -
40
) .* (tanh(a .* (t - t_1)) - tanh(a .* (t - t_2))) .* exp(-5 * (2 * t - 1) .^ 2) / 2
end
function d_a(t, E_0, a, t_1) # sympy codegen
t_2 = 1.0 - t_1
return (
(E_0 / 2) .*
((t - t_1) ./ cosh(a .* (t - t_1)) .^ 2 - (t - t_2) ./ cosh(a .* (t - t_2)) .^ 2) .*
(
E_0 .* (tanh(a .* (t - t_1)) - tanh(a .* (t - t_2))) .*
exp(5 * (2 * t - 1) .^ 2) - 40
) .* exp(-5 * (2 * t - 1) .^ 2)
)
end
function d_t1(t, E_0, a, t_1) # sympy codegen
t_2 = 1.0 - t_1
return (
(E_0 / 2) .* a .* (
E_0 .* (tanh(a .* (t - t_1)) - tanh(a .* (t + t_1 - 1))) .*
exp(5 * (2 * t - 1) .^ 2) - 40
) .* (tanh(a .* (t - t_1)) .^ 2 + tanh(a .* (t + t_1 - 1)) .^ 2 - 2) .*
exp(-5 * (2 * t - 1) .^ 2)
)
end
ff = OptimizationFunction(func; grad=grad!)
prob = OptimizationProblem(ff, x0, TLIST);
res = Optimization.solve(prob, NLopt.LD_LBFGS(), maxiters=1000)
##### LBFGSB
Pkg.add("LBFGSB")
using LBFGSB
func_lbfgs(x) = func(x, TLIST)
grad_lbfgs!(g, x) = grad!(g, x, TLIST)
@assert x0 isa Vector # LBFGSB does not work with ArrayPartition
fout, xout = lbfgsb(
func_lbfgs,
grad_lbfgs!,
x0,
m=5,
factr=1e7,
pgtol=1e-5,
iprint=-1,
maxfun=15000,
maxiter=15000
)
xout
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment