Julia Parametrization MWE

mwe.jl

# 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