# Optimal control module

This section contains instructions on setting up optimal control calculations in Spinach. The package contains an implementation of the gradient ascent pulse engineering (GRAPE) method, as well as its LBFGS-GRAPE and Newton-GRAPE enhancements. Ensemble optimal control is supported with a variety of secondary constrains, penalties, and options. Multiple examples are available in examples/optimal_control directory of the standard distribution.

## Stage 1: specify the spin system

The usual spin system specification should be done. For example, if we are interested in performing magnetisation transfer from one spin to another in a peptide chain in a liquid, the NMR calculation for that chain needs to be set up:

    % Magnetic field
sys.magnet=9.4;

% Spin system
sys.isotopes={'15N','1H','13C','13C','13C','15N'};

% Chemical shifts, ppm
inter.zeeman.scalar={0.0, 0.0, 55.0, 30.0, 175.0, 0.0};

% Scalar couplings, Hz (literature values)
inter.coupling.scalar=cell(6);
inter.coupling.scalar{1,3}=-11;
inter.coupling.scalar{2,3}=140;
inter.coupling.scalar{3,4}=35;
inter.coupling.scalar{3,5}=55;
inter.coupling.scalar{3,6}=7;
inter.coupling.scalar{5,6}=-15;

% Basis set
bas.formalism='sphten-liouv';
bas.approximation='IK-0';
bas.level=4;

% Spinach housekeeping
spin_system=create(sys,inter);
spin_system=basis(spin_system,bas);


## Stage 2: get initial and target states

You need to request your initial state and your desired target states from Spinach. The states must then be normalised. For example, if we are looking to move all magnetisation from Lz on spin 2 to Lz on spin 5, the following needs to be written:

    % Set up and normalise the initial state
rho_init=state(spin_system,{'Lz'},{2});
rho_init=rho_init/norm(full(rho_init),2);

    % Set up and normalise the target state
rho_targ=state(spin_system,{'Lz'},{5});
rho_targ=rho_targ/norm(full(rho_targ),2);


The first line of each block calls the usual state.m function, the the second line divides the resulting state vector by its 2-norm. If you have multiple source states that must be transferred to multiple destination states, put the corresponding state vectors into cell arrays.

## Stage 3: get drift and control operators

In an optimal control problem, the Hamiltonian is split into the part that the instrument cannot change (called drift Hamiltonian) and the operators whose coefficients the instrument can control (called control operators):

$\mathbf{H}=\mathbf{H}_0+\sum\limits_{k=1}^K c_k(t)\mathbf{H}_k$

where $\mathbf{H}_0$ is the drift Hamiltonian, and $\mathbf{H}_k$ are the control operators with time-dependent control coefficients $c_k(t)$. In a typical NMR experiment, the drift Hamiltonian contains chemical shifts and J-couplings, and control operators are Lx and Ly spin operators. All of these should be requested from Spinach as usual: the Hamiltonian is obtained from hamiltonian.m (if relaxaton is present, you should also call relaxation.m), and spin operators are obtained by calling operator.m function. In our peptide system, we would like Lx and Ly operators on carbon, nitrogen and protons, therefore:

    % Drift Hamiltonian
H0=hamiltonian(assume(spin_system,'nmr'));

% Control operators
LpH=operator(spin_system,'L+','1H');
LpC=operator(spin_system,'L+','13C');
LpF=operator(spin_system,'L+','19F');
LxH=(LpH+LpH')/2; LyH=(LpH-LpH')/2i;
LxC=(LpC+LpC')/2; LyC=(LpC-LpC')/2i;
LxF=(LpF+LpF')/2; LyF=(LpF-LpF')/2i;


If you have an ensemble with a different drift Liouvillian in each system, put those Hamiltonians into a cell array structured as {{La},{Lb},...}}. For powder averaged calculations, the drifts.m function would do that for you, for example (static powder):

    % Drift Liouvillians for the ensemble
control.drifts=drifts(spin_system,@powder,parameters,'labframe');


If your drift Liouvillian is time-dependent, the inner cell arrays should contain as many Liouvillians as there are points in your time grid: {{La1,La2,...},{Lb1,Lb2,...},...}}.

## Stage 4: put together the control structure

Spinach requires the user to specify all optimal control related options in a single structure called control. The following fields are available:

Optimal control parameter structure
Field name Mandatory/Optional Content
control.fidelity O A character string specifying the type of fidelity functional. Specify 'real' for the real part of the overlap between the final state and the target state, 'imag' for the imaginary part, and 'square' for the absolute square of the overlap. The default is 'real'.
control.operators M A cell array of matrices of control operators. If phase-amplitude optimisation or plotting is intended, these must be in the order {L1x,L1y,L2x,L2y,...}.
control.drifts M A cell array containing drift Hamiltonian matrices. If a single system is treated, there should be only one matrix; if there is an ensemble (with, for example, different J-couplings in each member), the cell array should contain multiple matrices.
control.rho_init M A cell array of initial states. For example, if state vector A should be moved into state vector X, and state vector B into state vector Y, your specification for the initial states should be {A,B}.
control.rho_targ M A cell array of target states. For example, if state vector A should be moved into state vector X, and state vector B into state vector Y, your specification for the target states should be {X,Y}.
control.prefix O A function handle containing a prefix pulse sequence that is applied to the initial condition before the optimal control stage begins. The syntax of the prefix function is
    rho_init=prefix(spin_system,drift,rho_init);


where drift will be changed automatically for each ensemble member, and the first element of the drift array is used in the case of time-dependent drifts.

control.suffix O A function handle containing a suffix pulse sequence that is applied to the target state before the scalar product with the final system state is taken. The syntax of the suffix function is
    rho_targ=suffix(spin_system,drift,rho_targ);


where drift will be changed automatically for each ensemble member, and the last element of the drift array is used in the case of time-dependent drifts.

control.pulse_dur M Control pulse duration in seconds.
control.pulse_nsteps M Number of time slices in the control pulse, an even integer.
control.pulse_dt O Instead of the previous two fields, which imply a uniform time grid, an explicit array of pulse slice durations may be supplied in this variable.
control.freeze M Freeze mask for the waveform. A logical array of the same shape as the initial guess with 1 for the points to be excluded from the optimisation and 0 for the points to be included.
control.dead_time O Dead time in seconds. The target state is to be reached in this time after the pulse has ended. Only the drift hamiltonian is active during the dead time.
control.keyholes O An array of linear function handles to be applied to the state vector after each time slice. Typically, these are coherence and correlation order projectors. For example, to filter out two-spin correlations after step 20 in a 50-step waveform, use:
    control.keyholes=cell(1,50);
control.keyholes{20}=@(rho)correlation(spin_system,rho,2,'all');

control.pwr_levels M A vector of control power levels across the ensemble, in rad/s. These correspond to things like non-uniform B1 field: the amplifier outputs some nominal power, but it is felt differently in different parts of the ensemble.
control.offsets O A vector of generalised offsets across the ensemble, in Hz. These correspond to things like resonance offsets that may be different for different signals in the spectrum.
control.offset_operator O Offset operator, usually something like Lz on the spins of interest. For each member of the ensemble, this operator will be multiplied by the corresponding element of control.offsets array, converted to rad/s and added to the drift Hamiltonian of that ensemble member.
control.phase_cycle O An array of phase combinations that the resulting pulse must obey. The first element of each line is the phase to be applied to the initial condition, the next element is the phase to be applied to the first pair of controls, the next element is the phase to be applied to the next pair of controls, etc. until the last element which is the phase to be applied to the target state. Multiple lines in such an array constitute a phase cycle.
control.basis O It is occasionally advantageous to run waveform optimisations not point-by-point, but with respect to coefficients in front of of some smooth functions, for example sinusoids. Those functions should be specified as rows of this array, which should have dimensions [nfunctions x control.pulse_npoints]. The functions must be orthonormal. The initial guess then becomes the set of time-dependent coefficients in front of the functions in this basis.
control.penalties O A cell array of character strings specifying the penalty functionals to be added to the fidelity functional. These are needed to make sure the waveform does not fly out into infinity. See penalty.m for details. The allowed values are 'NS' (norm square), 'SNS' (spillout norm square) and 'DNS' (derivative norm square). Default is 'SNS'.
control.p_weights O A vector of weight multipliers for the penalty functinals specified in control.penalties - default is 'SNS' penalty with weight 100.
control.amplitudes O For phase-only optimisations using grape_phase.m, the amplitude array (a stack of row vectors, on per {Lx,Ly} pair) must be specified here.
control.method O Optimisation method: 'lbfgs' for limited-memory BFGS, 'newton' for RFO regularised Newton-Raphson. In all cases, bracketing line search is used. Default is 'lbfgs'.
control.max_iter O Maximum number of optimisation iterations, default is 100.
control.tol_x O Termination tolerance on the step norm. Default is 1e-3.
control.tol_g O Termination tolerance on the gradient norm. Default is 1e-6.
control.u_bound O Upper bound for the 'SNS' penalty, as a fraction of the power transmitted by the amplifier. Default value is +1.
control.l_bound O Lower bound for the 'SNS' penalty, as a fraction of the power transmitted by the amplifier. Default value is -1.
control.plotting O A cell array of character strings specifying the plotting options for the diagnostic output. Possible values are: 'xy_controls' (coefficient of each control as a function of time), 'phi_controls' (phase of each X,Y control pair as a function of time), 'amp_controls' (amplitude of each X,Y control pair as a function of time), 'correlation_order' (correlation order dynamics computed by trajan.m), 'coherence_order' (coherence order dynamics computed by trajan.m, 'local_each_spin' (local probability density at each spin computed by trajan.m), 'local_each_spin' (probability density involving each spin computed by trajan.m), 'level polulations' (populations of individual energy levels computed by trajan.m), 'robustness' (fidelity profiles (1D ensemble), surfaces (2D ensemble), and histograms (nD ensemble)).
control.checkpoint O Checkpoint file name. The file is saved at each optimiser iteration. Optimisation will be restarted from this file if it exists.

## Stage 5: ensemble control considerations

Spinach supports five types of ensembles:

1. Source-target pairs. If multiple sources and multiple targets are specified, the optimisation will simultaneously drive each source to its corresponding target.
2. Drift Liouvillians. If multiple drifts are specified, the optimisation will aim to meet the target(s) for each drift independently.
3. Control powers. If multiple powers are specified, the optimisation will aim to meet the target(s) for each power level independently.
4. Offsets. If an offset operator and a list of coefficients are specified, the optimisation will aim to meet the target(s) for the offset operator added to the drift with each of those coefficients.
5. Phase cycle. If multiple lines are specified in the phase cycle array, the optimisation will aim to meet the target(s) for each of those phase lines simultaneously.

By default, these ensembles are kroneckered together. So, if you have two state-target pairs, ten different drift Hamiltonians, five power levels, five offsets, and four phase combinations, your ensemble would have 2x10x5x5x4=2000 systems in it, and the calculation would parallelise itself over those.

For some ensembles, these settings are not independent. In that case, the optional control.ens_corrs cell array may be specified.

control.ens_corrs Effect
'rho_match' Spinach will assume that the list of source-target pairs and the list of dirfts are not combinatorial, but aligned: each source-target pair on the list is to be matched with the corresponding operator in the drift list.
'rho_drift' Spinach will assume that an ensemble exists over drifts, power levels, and offsets, and each member of that ensemble has its own source-target pair provided by the user.

Other combinations may be implemented in request - send us an email if you need those.

## Stage 6: call optimcon.m constructor

After the dust has settled on the above, your control structure would looke something like:

    % Define control parameters
control.drifts={{H}};                           % Drift
control.operators={LxH,LyH,LxC,LyC,LxF,LyF};    % Controls
control.rho_init={rho_init};                    % Starting state
control.rho_targ={rho_targ};                    % Destination state
control.pwr_levels=2*pi*1e3;                    % Pulse power
control.pulse_dur=10e-3;                        % Pulse duration
control.pulse_nsteps=50;                        % Time points
control.penalties={'SNS'};                      % Penalty
control.p_weights=100;                          % Penalty weight
control.method='lbfgs';                         % Optimisation method
control.phase_cycle=[0  0  0  0*pi/2  -0*pi/2;
0  0  0  1*pi/2  -1*pi/2;
0  0  0  2*pi/2  -2*pi/2;
0  0  0  3*pi/2  -3*pi/2]; % Phase cycle

% Plots during optimisation
control.plotting={'correlation_order','local_each_spin','xy_controls'};


To validate and absorb this information, call the optimcon.m function:

    % Spinach housekeeping
spin_system=optimcon(spin_system,control);


It will update the spin_system data structure and print either the diagnostic summary or an informative error message to the console.

## Stage 7: specify the initial guess

Random numbers are usually fine. If the optimisation is carried out point-by-point (i.e. control.basis option is not specified), the dimension of the initial guess should be [n_controls x pulse_nsteps], for example:

    % Take a random guess for the waveform
guess=randn(numel(control.operators),control.pulse_nsteps)/3;


If the basis is specified, the dimension of the initial guess should be [n_controls x n_functions]. The initial guess is specified as a fraction of the amplifier power output, i.e. it should normally be within [-1,+1] range.

When an optimisation is to be carried out with respect to just the phase of the control sequence, the number of rows in the guess waveform should be equal to half the number of control operators, e.g.

    % Take a random guess for the phases
guess=randn(numel(control.operators)/2,control.pulse_nsteps)/3;


## Stage 8: run the optimization

The following functions return the fidelity, its gradient and its Hessian:

grape_xy.m
Optimisation of waveforms expressed as Lx, Ly coefficients.
grape_phase.m
Optimisation of waveforms expressed in amplitude-phase coordinates with respect to phase.
grape_coop.m
Cooperative control (experimental).

In principle, the derivatives these functions return can be supplied to any optimisation function, including those supplied with Matlab. Spinach does, however, include optimisation tools that are designed specifically for optimal control settings:

fminnewton.m
Gradient descent, quasi-Newton and Newton-Raphson methods.

For the purposes of optimal control optimisations, the following sets of settings provide a quick way to get started. The best balance between calculation speed and convergence rate is struck by LBFGS-GRAPE. To run Cartesian coordinate optimisation, use:

    % Run the optimization
fminnewton(spin_system,@grape_xy,guess);


To run phase optimisation, do not forget to specify the amplitude profile above (same dimensions as your phase guess) and call grape_phase.m instead:

    % Run the optimization
fminnewton(spin_system,@grape_phase,guess);


To use the Newton-Raphson method, specify 'newton' instead of 'lbfgs' in the control.method setting above. This is generally a good idea with fewer than 100 points in the waveform.