krotov_pseudocode.tex

\documentclass[aps,pra,twocolumn,notitlepage,letterpaper]{revtex4}

\def\Author{Michael Goerz}
\def\Title{Pseudocode for Krotov's Method}

\usepackage[utf8]{inputenc}
\usepackage{amsmath}
\usepackage{url}
\usepackage{xcolor}
\usepackage[
  pdftitle={\Title},
  pdfauthor={\Author},
  colorlinks=true, linkcolor=black, urlcolor=black, citecolor=black,
  bookmarksopen=false, breaklinks=true, plainpages=false, pdfpagelabels
]{hyperref}

\usepackage{braket}
\usepackage{listings}
\usepackage{algorithm}
\usepackage{algpseudocode}

\newcommand{\tr}[0]{\operatorname{tr}}
\newcommand{\diag}[0]{\operatorname{diag}}
\newcommand{\abs}[0]{\operatorname{abs}}
\newcommand{\pop}[0]{\operatorname{pop}}
\newcommand{\aux}[0]{\text{aux}}
\newcommand{\opt}[0]{\text{opt}}
\newcommand{\tgt}[0]{\text{tgt}}
\newcommand{\init}[0]{\text{init}}
\newcommand{\lab}[0]{\text{lab}}
\newcommand{\rwa}[0]{\text{rwa}}
\renewcommand{\Braket}[2]{\left\langle{}#1\vphantom{#2}\mid{}#2\vphantom{#1}\right\rangle}
\newcommand{\op}[1]{\hat{#1}}
\newcommand{\Op}[1]{\hat{#1}}
\newcommand{\dd}[0]{\,\text{d}}
\newcommand{\Liouville}[0]{\mathcal{L}}
\newcommand{\DynMap}[0]{\mathcal{E}}
\newcommand{\identity}[0]{\mathbf{1}}
\newcommand{\norm}[1]{\lVert#1\rVert}
\newcommand{\Norm}[1]{\left\lVert#1\right\rVert}
\newcommand{\Abs}[1]{\left\vert#1\right\vert}
\newcommand{\avg}[1]{\langle#1\rangle}
\newcommand{\Avg}[1]{\left\langle#1\right\rangle}
\newcommand{\AbsSq}[1]{\left\vert#1\right\vert^2}
\renewcommand{\Re}[0]{\operatorname{Re}}
\renewcommand{\Im}[0]{\operatorname{Im}}

\definecolor{yaleblue}{rgb}{0.06, 0.3, 0.57}
\definecolor{tropicalrainforest}{rgb}{0.0, 0.46, 0.37}

% algorithmic notation and color coding
\algrenewcommand\algorithmicrequire{\textbf{Input:}}
\algrenewcommand\algorithmicensure{\textbf{Output:}}
\algnewcommand{\IIf}[1]{\State\algorithmicif\ #1\ \algorithmicthen}
\floatname{algorithm}{Algorithm}
\definecolor{clVarScalar}{rgb}{0.06, 0.3, 0.57} % blue
\definecolor{clVarState}{rgb}{0.57, 0.06, 0.43} % purple
\definecolor{clVarOperator}{rgb}{0.09, 0.4, 0.04} % green
\definecolor{clVarArray}{rgb}{0.67, 0.06, 0.11} % red
\newcommand{\VarScalar}[1]{{\color{clVarScalar}#1}}
\newcommand{\VarState}[1]{\ensuremath{\textcolor{clVarState}{#1}}}
\newcommand{\PropAnnotation}[1]{\textcolor{black!30}{#1}}
\newcommand{\VarPropState}[4]{\ensuremath{\VarState{#1}_{\textcolor{clVarState}{#2}}^{\PropAnnotation{#3}}\PropAnnotation{(#4)}}}  % State with propagation-annotations: \VarPropState{symbol}{subscript}{iteration index superscript}{time argument}
\newcommand{\VarOperator}[2]{{\color{clVarOperator}#1_{#2}}} % \VarOperator{symbol}{subscripts} (no hat)
\newcommand{\VarArray}[1]{{\color{clVarArray}#1}}
\newcommand{\Forall}{\boldsymbol{\forall}}


\lstset{%   % style for source code listings
  basicstyle=\ttfamily\footnotesize,
  frame=single,
  framerule=0.5pt,
  rulecolor=\color{lightgray},
  numbersep=5pt,
  xleftmargin=0.5cm,
  xrightmargin=0.5cm,
  numberstyle=\tiny\textsf,
  showstringspaces=false,
  numberfirstline=true,
  backgroundcolor=\color{lightgray!20},
  commentstyle=\color{tropicalrainforest},    % comment style
  keywordstyle=\color{yaleblue},       % keyword style
  stringstyle=\color{tropicalrainforest},     % string literal style
  inputencoding=utf8,
}

\begin{document}

\title{\Title}
\author{\Author}
\date{\today}

\maketitle

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
For reference, Algorithm~\ref{al:KrotovsMethod} shows the complete pseudocode of
an optimization with Krotov's method, as implemented in the
\texttt{krotov} package (\url{https://github.com/qucontrol/krotov}).

Variables are color coded.
Scalars are set in blue, e.g.\ $\VarScalar{\epsilon_{ln}^{(0)}}$.
States  (Hilbert space states or vectorized density matrices) are set in purple,
e.g.\ $\VarState{\phi_k^{\text{init}}}$.
They may be annotated with light gray superscripts to indicate the
iteration-index $\VarScalar{i}$ of the control under which state was propagated,
and with light gray time arguments.
These annotations serve only to connect the variables to the equations of
motion: $\VarPropState{\phi}{k}{(0)}{t_n}$ and
$\VarPropState{\phi}{k}{(0)}{t_{n-1}}$ are the same variable
$\VarState{\phi_k}$.
Operators acting on states are set in green, e.g.\ $\VarOperator{\mu}{lkn}$.
These may be implemented as a sparse matrix or implicitly as a function that
returns the result of applying the operator to a state.
Lastly, storage arrays are set in red, e.g.\ $\VarArray{\Phi_0}$.
Each element of a storage array is a state.

The Python implementation groups several of the algorithm's input parameters by
introducing a list of $N$ ``objectives''.
The objectives are indexed by $\VarScalar{k}$, and each objective contains the
initial state $\VarState{\phi_k^{\text{init}}}$, the Hamiltonian or Liouvillian
$H_k$ to be used by the propagator $U$ and for the operators
$\VarOperator{\mu}{lkn}$, and possibly a ``target'' to be taken into account by
the function $\chi$.
In many applications, $H_k \equiv H$ is the same in all objectives, and
$\VarOperator{\mu}{lkn} \equiv \VarOperator{\mu}{l}$ if $H$ is linear in the
controls in addition.
The subscript $n$ and the superscript $(i-1)$ for
$\VarOperator{\mu^{\PropAnnotation{(i-1)}}}{lkn}$ in
lines~\ref{ln:pulse_update},~\ref{ln:pulse_update2} comes into play only if $H$
is \emph{not} linear in the control.
Mathematically, $\VarOperator{\mu}{lkn}$ would then have to be evaluated using
the \emph{updated} control.
Since the update is not yet known, the \emph{guess} control may be used as an
approximation (valid for sufficiently large $\lambda_{a,l}$).

The CPU resources required for the optimization are dominated by the time
propagation (calls to the function $U$ in
lines~\ref{ln:prop1},~\ref{ln:prop2}~\ref{ln:prop3}). This is under the
assumption that evaluating $U$ dominates the application of the operator
$\VarOperator{\mu^{\PropAnnotation{(i-1)}}}{lkn}$
to the state $\VarPropState{\phi}{k}{(i)}{t_{n-1}}$ and
the evaluation of the inner product of two states,
lines~\ref{ln:pulse_update},~\ref{ln:pulse_update2}.
This condition is fulfilled for any non-trivial Hilbert space dimension.

Loops over the index $\VarScalar{k}$ are parallelizable, in particular in a
shared-memory (multi-threaded) parallelization environment like OpenMP. In a
(multi-process) method-passing environment like MPI, some care must be taken to
minimize communication overhead from passing large state vectors.
For some (but not all) functionals, inter-process communication can be reduced
to only the scalar values constituting the sum over $\VarScalar{k}$ in
lines~\ref{ln:pulse_update},~\ref{ln:pulse_update2}.

The memory requirements of the algorithm are dominated by the storage arrays
$\VarArray{\Phi_0}$, $\VarArray{\Phi_1}$, and $\VarArray{X}$.
Each of these must store $N (N_T + 1)$ full state vectors (a full time
propagation for each of the $N$ objectives). Each state vector is typically an
array of double-precision complex numbers.
For a Hilbert space dimension $d$, a state vector thus requires $16 d$ bytes of
memory, or $16 d^2$ bytes for a density matrix.
Under certain conditions, the use of $\VarArray{\Phi_0}$ and $\VarArray{\Phi_1}$
can be avoided: both are required only when the second order update is used
($\sigma(t) \neq 0$). When the first order update is sufficient,
$\VarArray{\Phi_1}$ may overwrite $\VarArray{\Phi_0}$ so that the two collapse
into a single forward-storage $\VarArray{\Phi}$.
The states stored in $\VarArray{\Phi}$ are only used for the inhomogeneity
$\partial g_b / \partial \bra{\phi_k}$ in Eq.~\eqref{eq:bw_eqm}, and no storage
$\VarArray{\Phi}$ of forward-propagated states at all is required if $g_b \equiv
0$.
Thus, in most examples, only the storage $\VarArray{X}$ of the
backward-propagated states remains.
In principle, if the time propagation $U$ is unitary (i.e., invertible), the
states stored in $\VarArray{X}$ could be recovered by forward-propagation of
$\{\VarPropState{\chi}{k}{(i-1)}{t=0}\}$, eliminating $\VarArray{X}$ at the
(considerable) runtime cost of an additional time propagation.

\onecolumngrid

\begin{algorithm}
\caption{{\sc Krotov's Method for Quantum Optimal Control}
\label{al:KrotovsMethod} }
  \begin{algorithmic}[1]
    \Require{%
      ~\par
      \begin{enumerate}
      \item
      list of guess control values $\{\VarScalar{\epsilon^{(0)}_{ln}}\}$ where
      $\VarScalar{\epsilon^{(0)}_{ln}}$ is the value of the $l$'th control field
      on the $n$'th interval of the propagation time grid ($t_0 = 0, \dots,
      t_{N_T} = T$),
      i.e., $\VarScalar{\epsilon^{(0)}_{ln}} \equiv
      \epsilon^{(0)}_l(\tilde{t}_{n-1})$ with
      $n \in [1, N_T]$ and $\tilde{t}_n \equiv (t_{n} + t_{n+1})/2$
      \item
      list of update-shape values $\{\VarScalar{S_{ln}}\}$ with each
      $\VarScalar{S_{ln}} \in [0,1]$
      \item
      list of update step size values $\{\VarScalar{\lambda_{a,l}}\}$
      \item
      list of $N$ initial states $\{\VarState{\phi_{k}^{\text{init}}}\}$ at
      $t=t_0=0$
      \item
      propagator function $U$ that in ``forward mode'' receives a state
      $\VarPropState{\phi}{k}{}{t_n}$ and a list of control values
      $\{\VarScalar{\epsilon_{ln}}\}$ and
      returns $\VarPropState{\phi}{k}{}{t_{n+1}}$ by solving the differential
      equation~\eqref{eq:fw_eqm}, respectively in ``backward mode'' (indicated
      as $U^\dagger$) receives a state $\VarPropState{\chi}{k}{}{t_{n}}$ and returns
      $\VarPropState{\chi}{k}{}{t_{n-1}}$ by solving the differential
      equation~\eqref{eq:bw_eqm}
      \item
      list of operators
      $\VarOperator{\mu}{lkn}
      = \frac{\partial H_k}{\partial \epsilon_{l}(t)}\big\vert
        _{\epsilon_{ln}}$,
      where $H_k$ is the right-hand-side of the
      equation of motion of $\VarPropState{\phi}{k}{}{t}$, up to a factor of
      $(-\mathrm{i}/\hbar)$, cf.~Eq.~\eqref{eq:fw_eqm}
      \item
      function $\chi$ that receives a list of states $\{\VarPropState{\phi}{k}{}{T}\}$ and
      returns a list of states $\{\VarPropState{\chi}{k}{}{T}\}$ according to
      Eq.~\eqref{eq:chi_boundary}
      \item
      optionally, if  a second order construction of the
      pulse update is necessary: function $\sigma(t)$
      \end{enumerate}
    }
    \Ensure{%
      optimized control values $\{\VarScalar{\epsilon^{(\text{opt})}_{ln}}\}$,
      such that $J[\{\VarScalar{\epsilon_{ln}^{(\text{opt})}}\}]
      \le J[\{\VarScalar{\epsilon_{ln}^{(0)}}\}]$,
      with $J$ defined in Eq.~\eqref{eq:functional}.
    }
    \vspace{6pt}
    \Procedure{KrotovOptimization}{%
      $\{\VarScalar{\epsilon^{(0)}_{ln}}\}$,
      $\{\VarScalar{S_{ln}}\}$,
      $\{\VarScalar{\lambda_{a,l}}\}$,
      $\{\VarState{\phi_k^{\text{init}}}\}$,
      $U$,
      $\{\VarOperator{\mu}{lkn}\}$,
      $\chi$,
      $\sigma$
    }
      \State $\VarScalar{i} \gets 0$ \Comment{iteration number}
      \State allocate forward storage array $\VarArray{\Phi_0}[1\dots N, 0\dots N_T]$
      \For{$\VarScalar{k} \gets 1, \dots, N$} \Comment{initial forward-propagation}
         \State $\VarArray{\Phi_0}[\VarScalar{k},0] \gets \VarPropState{\phi}{k}{(0)}{t_0} \gets \VarState{\phi_k^{\text{init}}}$
         \For{$\VarScalar{n} \gets 1, \dots, N_T$}
           \State $\VarArray{\Phi_0}[\VarScalar{k}, \VarScalar{n}] \gets \VarPropState{\phi}{k}{(0)}{t_n} \gets U(\VarPropState{\phi}{k}{(0)}{t_{n-1}}, \{\VarScalar{\epsilon^{(0)}_{ln}}\})$ \Comment{propagate and store} \label{ln:prop1}
         \EndFor
      \EndFor
      \While{not converged} \Comment{optimization loop}
         \State $\VarScalar{i} \gets \VarScalar{i} + 1$
         \State $\VarArray{\Phi_1}, \{\VarScalar{\epsilon_{ln}^{(i)}}\} \gets \text{\textsc{KrotovIteration}}(\VarArray{\Phi_0}, \{\VarScalar{\epsilon^{(i-1)}_{ln}}\}, \dots)$
         \State $\VarArray{\Phi_0} \gets \VarArray{\Phi_1}$
      \EndWhile
      \State $\Forall \VarScalar{l}, \Forall \VarScalar{n}: \VarScalar{\epsilon_{ln}^{(\text{opt})}} \gets  \VarScalar{\epsilon_{ln}^{(i)}}$ \Comment{final optimized controls}
    \EndProcedure
    \algstore{krotovalg}
  \end{algorithmic}
    \vspace{6pt}

%\begin{algorithm}
  \begin{algorithmic}[1]
   \algrestore{krotovalg}
   \Procedure{KrotovIteration}{%
     $\VarArray{\Phi_0}$,
     $\{\VarScalar{\epsilon^{(i-1)}_{ln}}\}$,
     $\{\VarScalar{S_{ln}}\}$,
     $\{\VarScalar{\lambda_{a,l}}\}$,
     $\{\VarState{\phi_k^{\text{init}}}\}$,
     $U$,
     $\{\VarOperator{\mu}{lkn}\}$,
     $\chi$,
     $\sigma$
   }
      \State $\Forall \VarScalar{k}: \VarPropState{\phi}{k}{(i-1)}{T} \gets \VarArray{\Phi_0}[\VarScalar{k}, N_T]$
      \State $\{\VarPropState{\chi}{k}{(i-1)}{T}\} \gets \chi(\{\VarPropState{\phi}{k}{(i-1)}{T}\})$ \Comment{backward boundary condition} \label{ln:chi_boundary}
      \State allocate backward storage array $\VarArray{X}[1 \dots N, 0 \dots N_T]$.
      \For{$\VarScalar{k} \gets 1, \dots, N$}
        \State $\VarArray{X}[\VarScalar{k}, N_T] \gets \VarPropState{\chi}{k}{(i-1)}{T}$
         \For{$\VarScalar{n} \gets N_T, \dots, 1$} \Comment{backward-propagate and store}
           \State $\VarArray{X}[\VarScalar{k}, \VarScalar{n}-1] \gets \VarPropState{\chi}{k}{(i-1)}{t_{n-1}} \gets U^{\dagger}(\VarPropState{\chi}{k}{(i-1)}{t_{n}}, \{\VarScalar{\epsilon^{(i-1)}_{ln}}\}, \VarArray{\Phi_0})$  \label{ln:prop2}
         \EndFor
      \EndFor
      \State allocate forward storage array $\VarArray{\Phi_1}[1\dots N, 0\dots N_T]$
      \State $\Forall \VarScalar{k}: \VarArray{\Phi_1}[\VarScalar{k},0] \gets \VarPropState{\phi}{k}{(i)}{t_0} \gets \VarState{\phi_k^{\text{init}}}$
      \For{$\VarScalar{n} \gets 1, \dots, N_T$} \Comment{sequential update loop}
         \State $\Forall \VarScalar{k}: \VarPropState{\chi}{k}{(i-1)}{t_{n-1}} \gets \VarArray{X}[\VarScalar{k}, \VarScalar{n}-1]$
         \State $\Forall \VarScalar{l}: \VarScalar{\Delta \epsilon_{ln}} \gets \frac{\VarScalar{S_{ln}}}{\VarScalar{\lambda_{a, l}}} \Im \sum_{\VarScalar{k}} \big\langle \VarPropState{\chi}{k}{(i-1)}{t_{n-1}} \big\vert \VarOperator{\mu^{\PropAnnotation{(i-1)}}}{lkn} \big\vert \VarPropState{\phi}{k}{(i)}{t_{n-1}} \big\rangle$ \Comment{first order} \label{ln:pulse_update}
         \If{$\sigma(t) \neq 0$} \Comment{second order}
           \State $\Forall \VarScalar{k}: \VarPropState{\Delta\phi}{k}{(i)}{t_{n-1}} \gets \VarPropState{\phi}{k}{(i)}{t_{n-1}} - \VarArray{\Phi_0}[\VarScalar{k},\VarScalar{n}-1]$ \label{ln:delta_phi}
           \State $\Forall \VarScalar{l}: \VarScalar{\Delta \epsilon_{ln}} \gets \VarScalar{\Delta \epsilon_{ln}} + \frac{\VarScalar{S_{ln}}}{\VarScalar{\lambda_{a, l}}} \Im \sum_{\VarScalar{k}} \frac{1}{2} \sigma(\tilde{t}_n) \big\langle \VarPropState{\Delta\phi}{k}{(i)}{t_{n-1}} \big\vert \VarOperator{\mu^{\PropAnnotation{(i-1)}}}{lkn} \big\vert \VarPropState{\phi}{k}{(i)}{t_{n-1}} \big\rangle$ \label{ln:pulse_update2}
         \EndIf
         \State $\Forall \VarScalar{l}: \VarScalar{\epsilon_{ln}^{(i)}} \gets \VarScalar{\epsilon_{ln}^{(i-1)}} + \VarScalar{\Delta \epsilon_{ln}}$ \Comment{apply update}
         \State $\Forall \VarScalar{k}: \VarArray{\Phi_1}[\VarScalar{k},\VarScalar{n}] \gets \VarPropState{\phi}{k}{(i)}{t_n} \gets U(\VarPropState{\phi}{k}{(i)}{t_{n-1}}, \{\VarScalar{\epsilon_{ln}^{(i)}}\})$ \Comment{propagate and store} \label{ln:prop3}
      \EndFor
      \If{$\sigma(t) \neq 0$}
        \State Update internal parameters of $\sigma(t)$ \label{ln:sigma_update} if necessary
      \EndIf
    \EndProcedure
  \end{algorithmic}
\end{algorithm}


\section*{Optimization Functional and Equations of Motion}

\vspace{-5mm}
\begin{equation}
  \label{eq:functional}
  J[\{\ket{\phi_k^{(i)}(t)}\}, \{\epsilon_l^{(i)}(t)\}]
    = J_T(\{\ket{\phi_k^{(i)}(T)}\})
        + \sum_l \int_0^T g_a(\epsilon_l^{(i)}(t)) \dd t
        + \int_0^T g_b(\{\phi^{(i)}_k(t)\}) \dd t
\end{equation}

\begin{equation}%
  \label{eq:fw_eqm}
  \frac{\partial}{\partial t} \Ket{\phi_k^{(i)}(t)} =
  -\frac{\mathrm{i}}{\hbar} \Op{H}^{(i)}(t) \Ket{\phi_k^{(i)}(t)}
\end{equation}

\begin{equation}%
  \label{eq:bw_eqm}
  \frac{\partial}{\partial t} \Ket{\chi_k^{(i-1)}(t)}
  = -\frac{\mathrm{i}}{\hbar} \Op{H}^{\dagger\,(i-1)}(t) \Ket{\chi_k^{(i-1)}(t)}
    + \left.\frac{\partial g_b}{\partial \Bra{\phi_k}}\right\vert_{(i-1)}\,
\end{equation}

\vspace{-3mm}
\begin{equation}%
  \label{eq:chi_boundary}
  \text{with}\qquad
  \Ket{\chi_k^{(i-1)}(T)}
  = - \left.\frac{\partial J_T}{\partial \Bra{\phi_k(T)}}
    \right\vert_{(i-1)}
\end{equation}


\vspace{5pt} \textbf{Notes:}
  \begin{itemize}
    \item The index $\VarScalar{k}$ numbers the independent states to be
      propagated, respectively the independent ``objectives'' (see text for
      details), $\VarScalar{l}$ numbers the independent control fields, and
      $\VarScalar{n}$ numbers the intervals on the time grid. All of these
      indices start at 1.
    \item The optimization loop may be stopped if the optimization functional or
      the change of functional falls below a pre-defined threshold, a maximum
      number of iterations is reached, or any other criterion.
    \item The braket notation in line~\ref{ln:pulse_update} indicates the
      (Hilbert-Schmidt) inner product of the state
      $\VarPropState{\chi}{k}{(i-1)}{t_n-1}$ and the state
      resulting from applying $\VarOperator{\mu^{\PropAnnotation{(i-1)}}}{lkn}$
      to $\VarPropState{\phi}{k}{(i)}{t_{n-1}}$.
      In Hilbert space, this is the standard braket.
      In Liouville space, it is $\tr\left(\VarState{\chi_k}^\dagger
      \;\VarOperator{\mu}{lkn}[\VarState{\phi_k}]\right)$ with density
      matrices $\VarState{\chi_k}$, $\VarState{\phi_k}$ and a super-operator
      $\VarOperator{\mu}{lkn}$.
    \item For numerical stability, the states $\VarPropState{\chi}{k}{(i-1)}{T}$
      in line~\ref{ln:chi_boundary} may be normalized. This norm then has to
      taken into account in the pulse update, line~\ref{ln:pulse_update}.
    \item In line~\ref{ln:prop2}, the storage array $\VarArray{\Phi_0}$ is
      passed to $U^\dagger$ only to account for the inhomogeneity due to
      a possible state-dependent constraint, $\partial g_b / \partial
      \bra{\phi_k}$ in Eq.~\eqref{eq:bw_eqm}.
      If $g_b \equiv 0$, the parameter can be omitted.
  \end{itemize}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\end{document}