Using RecursiveArrayTools.ArrayPartition to keep track of control parameters

Specific Domains Optimization (Mathematical)
1

I’m making some plans for the QuantumControl package to extend from optimizing arbitrary time-continuous control fields to control fields that are described analytically by a handful of scalar control parameters.

Without going into details about “quantum”, I’ll have a “Hamiltonian” object that encodes the right-hand-side of the ODE for the system dynamics, and that contains multiple control function objects, each of which contains a set of control parameters. I want to extract all of these control parameters into a single AbstractVector that I can feed as x into Optimization.jl or some other optimization toolbox. It looks to me like RecursiveArrayTools.ArrayPartition might be the right tool for the job of keeping track of these parameters.

I imagine this kind of problem is pretty common in other (non-quantum) optimal control setups, so I’m asking for some community feedback on what packages people have used to solve similar problems.

The two things I want:

  • 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 Hamiltonian. Avoiding allocations is a goal in the QuantumControl package, hence the focus here on mutation in lieu of a more functional style that might be better if the “Hamiltonians” weren’t potentially huge objects.
  • 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().

I can illustrate with a MWE. A longer version is at Julia Parametrization MWE · GitHub, which also includes gradient-based optimizations.

To set up the MWE environment:

using Pkg
Pkg.activate(temp=true)
Pkg.add("RecursiveArrayTools")
Pkg.add("ComponentArrays")
Pkg.add("UnPack")
Pkg.add("Optimization")
Pkg.add("OptimizationNLopt")

using ComponentArrays
using RecursiveArrayTools
using UnPack: @unpack
using Optimization
using OptimizationNLopt: NLopt

I’m going to define some “control function” object BoxyField with three control parameters:

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

and set up the “Hamiltonian” to depend on three of these fields:

struct Hamiltonian
    # The actual Hamiltonian contains operators in addition to controls, but we
    # don't need that for this MWE
    controls
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 globals in this MWE 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₃])

TLIST = collect(range(0, 1; length=101));

So now the key I can use ArrayPartition to pull out the control parameters:

# 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

PARAMETERS = get_parameters(HAMILTONIAN)

And then conveniently use that in an optimization (a gradient-free simplex, to keep it easy for the MWE; also: this “optimization problem” is completely contrived, in reality I’ll be simulating dynamics and evaluating a functional based on the final-time quantum state, not matching the shape of the control function to some target):

# 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


prob = OptimizationProblem(func, x0, TLIST; stopval=50,);
res = Optimization.solve(prob, NLopt.LN_NELDERMEAD(), maxiters=1000)

Does all of that look reasonable? Any other alternatives to RecursiveArrayTools I should explore?

2

While everything is type-stable (may need to add controls::T to Hamiltonian) and fixed-length (use Tuples or SVectors instead of Vectors), this should work:

julia> using AccessorsExtra

julia> params = getall(HAMILTONIAN, RecursiveOfType(Number))

julia> params = ...

julia> NEW_HAMILTONIAN = setall(HAMILTONIAN, RecursiveOfType(Number), params)
1 Like