renew-dots, renew-matrix, nullify-dots

# Dynamic Learning of Correlation Potentials for a Time-Dependent Kohn-Sham
System

Harish S. Bhat<EMAIL_ADDRESS>
Department of Applied Mathematics  University of California Merced  Kevin
Collins<EMAIL_ADDRESS>
Department of Physics  University of California Merced  Prachi Gupta
<EMAIL_ADDRESS>
Department of Applied Mathematics  Department of Chemistry and Biochemistry
University of California Merced  Christine M. Isborn<EMAIL_ADDRESS>
Department of Chemistry and Biochemistry  University of California Merced

###### Abstract

We develop methods to learn the correlation potential for a time-dependent
Kohn-Sham (TDKS) system in one spatial dimension. We start from a low-
dimensional two-electron system for which we can numerically solve the time-
dependent Schrödinger equation; this yields electron densities suitable for
training models of the correlation potential. We frame the learning problem as
one of optimizing a least-squares objective subject to the constraint that the
dynamics obey the TDKS equation. Applying adjoints, we develop efficient
methods to compute gradients and thereby learn models of the correlation
potential. Our results show that it is possible to learn values of the
correlation potential such that the resulting electron densities match ground
truth densities. We also show how to learn correlation potential functionals
with memory, demonstrating one such model that yields reasonable results for
trajectories outside the training set.

###### keywords:

Physics-constrained learning, adjoint methods, quantum dynamics, TDDFT.

##  1 Introduction

The time-dependent Schrödinger equation (TDSE) governs the behavior of $N$
quantum particles,

$i\partial_{t}\Psi(r_{1},r_{2},...,r_{N},t)=\hat{H}(r_{1},r_{2},...,r_{N},t)\Psi(r_{1},r_{2},...,r_{N},t),$
(1)

where $\hat{H}$ is the Hamiltonian and $\Psi$ is the many-body wave function.
In $d$-dimensional space, the many-body Coulomb interaction in the potential
term of $\hat{H}$ leads to a coupled system of partial differential equations
(PDE) in $dN+1$ variables. Hence (1) can only be solved for simple model
problems, such as for one electron in three dimensions or two electrons in one
dimension. To simulate electron dynamics in molecules and materials, a widely
used approach is time-dependent density functional theory (TDDFT), in which
the many-body wave function $\Psi$ is replaced with the Kohn-Sham wave
function $\Phi(r)$ to give the time-dependent Kohn-Sham (TDKS) equation
(Maitra, 2016; Ullrich, 2011):

$i\partial_{t}\Phi(r,t)=\sum_{i=1}^{N}[-(1/2)\nabla^{2}_{i}+v^{\text{ext}}(r_{i},t)+v^{H}[n](r_{i},t)+v^{XC}[n,\Psi_{0},\Phi_{0}](r_{i},t)]\Phi(r,t).$
(2)

Because $\Phi(r)$ is constructed as a product of non-interacting single-
particle orbitals $\phi(r_{i})$, (2) decouples into $N$ separate evolution
equations in $3+1$ variables. Assuming all terms in (2) are specified, one can
use (2) to simulate molecular systems for which numerical simulation of (1) is
intractable.

In (2), the many-body Coulomb interaction between electrons is replaced by
known classical Hartree $v^{H}$ and unknown exchange-correlation $v^{XC}$
single-particle potentials, with the latter incorporating many-body effects.
TDDFT is formally an exact theory, as the Runge-Gross and Van Leeuwen theorems
proved the existence of a time-dependent electronic potential and the unique
mapping to the time-dependent electron density, which is generated from the KS
orbitals of the TDKS equation (Runge and Gross, 1984; van Leeuwen, 1999).

The challenge in TDDFT is to construct $v^{XC}$ potentials that yield an
electron density $n$ that is identical to the exact time-dependent many-body
electron density generated from the TDSE. Previous work has shown that the
unknown $v^{XC}$ formally depends on the initial many-body wave function
$\Psi_{0}$, the initial KS state $\Phi_{0}$, and the electron density at all
points in time $n(r,s<t)$ (Maitra et al., 2002). Although the development of
$v^{XC}$ for electrons is a very active area of research, almost all $v^{XC}$
make use of the so-called “adiabatic approximation” that only takes into
account the instantaneous electron density, leading to significant
inaccuracies in electron dynamics due to the lack of memory in $v^{XC}$. _The
desire for more accurate electron dynamics leads to a natural question: can we
learn $v^{XC}$ from time series data?_ Note that this is an entirely different
problem than the problem of learning static, ground state potentials from the
exact ground state electron density in _time-independent_ density functional
theory (DFT) (Nagai et al., 2018; Kalita et al., 2021).

For machine learning of $v^{XC}$ to proceed in the _time-dependent_ context
(TDDFT), a first obstacle is formulating a tractable learning problem. In
recent work, Suzuki et al. (2020) works with a _spatially one-dimensional_
electron-hydrogen scattering problem. For this model problem, one can solve
(1) numerically; from the solution, one can compute the electron density
$n(x,t)$ on spatial/temporal grids. In this problem, we know both the
functional form of $v^{X}$ and that $v^{XC}=v^{X}+v^{C}$. Furthermore, the
one-dimensionality enables one to solve for exact values of $v^{C}$ (Elliott
et al., 2012), again on spatial/temporal grids. With grid-based values of both
$v^{C}$ and $n$, the task of learning $v^{C}[n]$ becomes a static, supervised
learning problem, which Suzuki et al. (2020) solves using neural network
models. To our knowledge, this is the only prior work on learning $v^{XC}$ for
TDDFT.

We revisit the electron-hydrogen scattering model problem and develop methods
to learn $v^{C}[n]$ that do not require us to solve for grid-based values of
$v^{C}$ beforehand. _In short, we view the $v^{C}$ functional as a control
that guides TDKS propagation._ We formulate the learning problem as an optimal
control problem: find $v^{C}$ that minimizes the squared error between TDKS
electron densities $n$ and reference electron densities $\tilde{n}$. Implicit
in this formulation is the dynamical constraint that electron densities $n$
evolve forward in time via the TDKS equation with the model $v^{C}$. The
adjoint or costate method is often used to handle constraints of this kind
(Bryson and Ho, 1975; Hasdorff, 1976). To our knowledge, the derivations and
applications of the adjoint method, _to learn $v^{C}$ models with memory for
the TDKS equation_, are considered here for the first time111See Section 5.7
in the Appendix for further context.. We derive adjoint systems for two
settings: (i) to learn pointwise values of $v^{C}$ on a grid, and (ii) to
learn the functional dependence of $v^{C}$ on the electron density at two
points in time. We apply our methods to train both types of models, and study
their training and test performance. In particular, we train a neural network
model of $v^{C}[n]$ with memory that, when used to solve the TDKS equations
for initial conditions outside the training set, yields qualitatively accurate
predictions of electron density.

##  2 Methods

To formulate the problem of learning $v^{C}$ from time series, we first work
in continuous space and time. Later, to derive numerical algorithms to solve
this problem, we discretize.

#### Continuous Problem.

Define the 1D electron density created from KS orbitals

$n(x,t)=2|\phi(x,t)|^{2},$ (3)

and the soft-Coulomb external $v^{\text{ext}}$ and interaction $W^{ee}$
potentials

$\displaystyle v^{\text{ext}}(x)$ $\displaystyle=-((x+10)^{2}+1)^{-1/2},$ (4a)
$\displaystyle W^{ee}(x^{\prime},x)$
$\displaystyle=((x^{\prime}-x)^{2}+1)^{-1/2}.$ (4b)

The potentials (4a) and (4b) specify that we are working with the spatially
one-dimensional electron-hydrogen scattering problem considered by several
previous authors. For this problem, we know that $v^{XC}=v^{X}+v^{C}$. Let
$\phi$ and $n$ stand for $\phi(x,t)$ and $n(x,t)$. Then in one spatial
dimension and expressed in atomic units (a.u.), the TDKS system (2) becomes:

$\displaystyle i\partial_{t}\phi$
$\displaystyle=-\frac{1}{2}\partial_{xx}\phi+v^{\text{ext}}(x,t)\phi+v^{H}[n](x,t)\phi+v^{X}[n](x,t)\phi+v^{C}[\phi](x,t)\phi,$
(5a) $\displaystyle v^{H}[n](x,t)$
$\displaystyle=\int_{x^{\prime}}W^{ee}(x^{\prime},x)n(x^{\prime},t)\,dx^{\prime},\qquad\qquad
v^{X}[n](x,t)=-\frac{1}{2}v^{H}[n](x,t).$ (5b)

_In ( 5), the term that we are trying to learn (e.g., the control) is
$v^{C}[\phi]$._ Prior first principles work has shown that at time $t$,
$v^{C}$ should depend functionally on the electron density $n(x,s)$ for $s\leq
t$, the initial Kohn-Sham state $\phi(x,0)$ and the initial Schrödinger wave
function $\Psi(x,0)$ (Maitra et al., 2002; Wagner et al., 2012). In this work,
we ignore the dependence of $v^{C}$ on the initial states $\phi(x,0)$ and
$\Psi(x,0)$, and focus on modeling the dependence on present and past electron
densities. By (3), dependence on $n$ is equivalent to a particular type of
dependence on $\phi$; we use the notation $v^{C}[\phi]$ to refer to models
that depend on $\phi$ either directly or through $n$.

_For the sake of intuition, let us formulate the control problem in continuous
time and space._ Assume that for $t\in[0,T]$, we have access to a reference
electron density trajectory $\tilde{n}(x,t)$. Suppose that our model
$v^{C}[\phi;\boldsymbol{\theta}]$ is parameterized by $\boldsymbol{\theta}$.
Then we seek to minimize the squared loss

$\mathcal{J}(\boldsymbol{\theta})=\frac{1}{2}\int_{x=-\infty}^{\infty}\int_{t=0}^{T}(n(x,t)-\tilde{n}(x,t))^{2}\,dt\,dx,$
(6)

subject to the constraint that $n(x,t)$ is computed via (3), with $\phi(x,t)$
evolving on the interval $0\leq t\leq T$ according to the TDKS system (5). In
this TDKS system, we identify $v^{C}$ with our model
$v^{C}[\phi;\boldsymbol{\theta}]$. _In short, we seek $\boldsymbol{\theta}$
such that the resulting $v^{C}[\phi;\boldsymbol{\theta}]$ functional guides
the TDKS system to yield a solution $\phi(x,t)$ such that $n=2|\phi|^{2}$
matches the reference trajectory $\tilde{n}(x,t)$._

#### Direct and Adjoint Methods.

In a _direct method_ to minimize the loss (6), we compute gradients by
applying $\nabla_{\boldsymbol{\theta}}$ to both sides of (6). This will yield
an expression for $\nabla_{\boldsymbol{\theta}}\mathcal{J}$ that involves
$\nabla_{\boldsymbol{\theta}}\phi$. To compute this latter quantity, we
numerically solve an evolution equation derived by taking
$\nabla_{\boldsymbol{\theta}}$ of both sides of (5a). At each iteration of our
gradient-based optimizer, we would carry out this procedure to compute
$\nabla_{\boldsymbol{\theta}}\mathcal{J}$, which is then used to update
$\boldsymbol{\theta}$. In practice, this direct method suffers from one major
problem: if we discretize $\phi$ in space using $J+1$ grid points, and if
$\boldsymbol{\theta}$ has dimension $B$, then at each point in time,
$\nabla_{\boldsymbol{\theta}}\phi$ will have dimension $(J+1)B$. In our work,
$B$ can exceed $10^{7}$, while $J\geq 600$ is required for sufficient spatial
accuracy. Solving the evolution equation for
$\nabla_{\boldsymbol{\theta}}\phi$ in $(J+1)B$-dimensional space thus incurs
huge computational expense at each optimization step.

In this paper, we pursue the _adjoint method_ , which enables us to compute
all required gradients without computing or even storing any
$(J+1)B$-dimensional objects, thus dramatically reducing computational costs
relative to the direct method. Within the space of adjoint methods, there are
two broad approaches: (i) to use the continuous-time loss and constraints to
derive differential equations for continuous-time adjoint variables, and (ii)
to first discretize the loss and constraints, and then derive numerical
schemes for discrete-time adjoint variables. In approach (i), we must still
discretize the adjoint differential equations in order to solve them; the
choice of discretization can lead to subtle issues (Sanz-Serna, 2016). We
choose approach (ii) for its relative simplicity.

In the discrete adjoint method, we incorporate a discretized version of the
dynamical system (5) as a constraint using time-dependent Lagrange multipliers
$\boldsymbol{\lambda}(t)$. In this approach, we derive and numerically solve a
backward-in-time evolution equation for $\boldsymbol{\lambda}(t)$, from which
we compute required gradients. Importantly, $\boldsymbol{\lambda}(t)$ has the
same dimension as the state variables $\boldsymbol{\phi}(t)$ defined below; in
our implementation, both quantities are $(J+1)$-dimensional. We obtain the
gradients of the discretized loss at a computational cost that is proportional
to that of computing the loss itself.

#### Discretized Problem.

To keep this paper focused on the learning/control problem, we have moved
details of the numerical solution of the TDKS system (5) to Section 5.1 of the
Appendix. Here we include only the most important concepts. First, we
discretize the Kohn-Sham state by introducing
$\boldsymbol{\phi}(t_{k})=[\phi(x_{0},t_{k}),\ldots,\phi(x_{J},t_{k})]^{T}$.
The spatial domain is $x\in[L_{\text{min}},L_{\text{max}}]$. With $\Delta
x=(L_{\text{max}}-L_{\text{min}})/J>0$, our spatial grid is
$x_{j}=L_{\text{min}}+j\Delta x$. Our temporal grid is $t_{k}=k\Delta t$, with
$\Delta t=T/K$. The positive integers $J$ and $K$ are user-defined parameters
that control the accuracy of the discretization. Second, by applying finite
differences, Simpson’s quadrature rule, and operator splitting, we can derive
the following evolution equation for the discretized state $\boldsymbol{\phi}$
defined above:

$\boldsymbol{\phi}(t_{k+1})=\exp(-i\mathcal{K}\Delta
t/2)\exp(-iV(\boldsymbol{\phi}(t_{k}),\boldsymbol{v}^{C}_{k})\Delta
t)\exp(-i\mathcal{K}\Delta t/2)\boldsymbol{\phi}(t_{k}).$ (7)

Here $\mathcal{K}$ is a constant $(J+1)\times(J+1)$ matrix, while $V$ is a
diagonal $(J+1)\times(J+1)$ matrix that depends functionally on both the state
$\boldsymbol{\phi}$ and on $\boldsymbol{v}^{C}$, our spatially discretized
model of the correlation potential $v^{C}$ from (5). Detailed descriptions of
$\mathcal{K}$ and $V$ are provided in Section 5.1.

Evolving $\boldsymbol{\phi}$ according to (7) generates a numerical
approximation to the solution $\phi(x,t)$ of (5). This approximation has a
truncation error of $O(\Delta t^{2})$ in time and $O(\Delta x^{4})$ in space.

#### First Adjoint Method: Learning $v^{C}$ Pointwise.

Assume we have access to observed values of electron density on the grid—we
denote these observed or reference values by $\tilde{n}(x_{j},t_{k})$. The
first problem we consider is to learn $v^{C}(x_{j},t_{k})$ on the same grid.
Suppose we start from an initial condition $\boldsymbol{\phi}(0)$ and an
estimate $\boldsymbol{v}^{C}$. We iterate (7) forward in time and obtain a
trajectory $\boldsymbol{\phi}(t_{k})$ for $0\leq k\leq K$. We then form
$n(x_{j},t_{k})=|\phi(x_{j},t_{k})|^{2}$. In this subsection, $\phi$ and $n$
are the predicted wave function and density when we use the estimated
correlation potential $\boldsymbol{v}^{C}$. Let
$\mathcal{P}_{\mathcal{K}}=\exp(-i\mathcal{K}\Delta t/2)$ and abbreviate
$\boldsymbol{\phi}_{k}=\boldsymbol{\phi}(t_{k})$,
$\boldsymbol{v}^{C}_{k}=\boldsymbol{v}^{C}(t_{k})$. Define the discrete-time
propagator

$\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi},\boldsymbol{v}^{C})=\mathcal{P}_{\mathcal{K}}\exp(-iV(\boldsymbol{\phi},\boldsymbol{v}^{C})\Delta
t)\mathcal{P}_{\mathcal{K}}\boldsymbol{\phi},$ (8)

so that (7) can be written as the discrete-time system
$\boldsymbol{\phi}_{k+1}=\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi}_{k},\boldsymbol{v}^{C}_{k})$. Both sides of this system
are complex-valued. In order to form a real-valued Lagrangian and take real
variations, we split both $\boldsymbol{\phi}$ and $\boldsymbol{F}$ into real
and imaginary parts:
$\boldsymbol{\phi}=\boldsymbol{\phi}^{R}+i\boldsymbol{\phi}^{I}$ and
$\boldsymbol{F}_{\Delta t}=\boldsymbol{F}_{\Delta
t}^{R}+i\boldsymbol{F}_{\Delta t}^{I}$. Superscript $R$ and $I$ denote,
respectively, the real and imaginary parts of a complex quantity. Let the
uppercase $\boldsymbol{\Phi}$, $\boldsymbol{\Lambda}$, and
$\boldsymbol{V}^{C}$ denote the collections of all corresponding lowercase
$\boldsymbol{\phi}_{k}$, $\boldsymbol{\lambda}_{k}$, and
$\boldsymbol{v}^{C}_{k}$ for all $k$. Then we form a real-variable Lagrangian
that consists of the discretized squared loss with the constraint that
$\boldsymbol{\phi}$ evolves via (7).

$\mathscr{L}(\boldsymbol{\Phi}^{R},\boldsymbol{\Phi}^{I},\boldsymbol{\Lambda}^{R},\boldsymbol{\Lambda}^{I},\boldsymbol{v}^{C})=\frac{1}{2}\sum_{k=0}^{K}\sum_{j=0}^{J}(2\phi^{R}(x_{j},t_{k})^{2}+2\phi^{I}(x_{j},t_{k})^{2}-\tilde{n}(x_{j},t_{k}))^{2}\\\
-\sum_{k=0}^{K-1}[\boldsymbol{\lambda}_{k+1}^{R}]^{T}(\boldsymbol{\phi}_{k+1}^{R}-\boldsymbol{F}^{R}_{\Delta
t}(\boldsymbol{\phi}_{k}^{R},\boldsymbol{\phi}_{k}^{I},\boldsymbol{v}^{C}_{k}))+[\boldsymbol{\lambda}_{k+1}^{I}]^{T}(\boldsymbol{\phi}_{k+1}^{I}-\boldsymbol{F}^{I}_{\Delta
t}(\boldsymbol{\phi}_{k}^{R},\boldsymbol{\phi}_{k}^{I},\boldsymbol{v}^{C}_{k})).$
(9)

Setting $\delta\mathscr{L}=0$ for all variations
$\delta\boldsymbol{\phi}_{k}^{R}$ and $\delta\boldsymbol{\phi}_{k}^{I}$ for
$k\geq 1$, we obtain

$\displaystyle\boldsymbol{\lambda}_{K}$
$\displaystyle=4\left[(2|\boldsymbol{\phi}_{K}|^{2}-\tilde{\boldsymbol{n}}_{K})\circ\boldsymbol{\phi}_{K}\right]$
(10a) $\displaystyle\begin{bmatrix}\boldsymbol{\lambda}_{k}^{R}\\\
\boldsymbol{\lambda}_{k}^{I}\end{bmatrix}^{T}$
$\displaystyle=4(2|\boldsymbol{\phi}_{k}|^{2}-\tilde{\boldsymbol{n}}_{k})\circ\begin{bmatrix}\boldsymbol{\phi}_{k}^{R}\\\
\boldsymbol{\phi}_{k}^{I}\end{bmatrix}^{T}+\begin{bmatrix}\boldsymbol{\lambda}_{k+1}^{R}\\\
\boldsymbol{\lambda}_{k+1}^{I}\end{bmatrix}^{T}\boldsymbol{J}_{\boldsymbol{\phi}}\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi}_{k}^{R},\boldsymbol{\phi}_{k}^{I},\boldsymbol{v}^{C}_{k}).$
(10b)

Here $\boldsymbol{J}_{\boldsymbol{\phi}}\boldsymbol{F}_{\Delta t}$ denotes the
Jacobian of $\boldsymbol{F}$ with respect to $\boldsymbol{\phi}$. We use (10a)
as a final condition and iterate (10b) backward in time for $k=K-1,\ldots,1$.
Having computed $\boldsymbol{\Lambda}$ from (10), we return to (9) and compute
the gradient with respect to $\boldsymbol{v}^{C}_{\ell}$:

$\nabla_{\boldsymbol{v}^{C}_{\ell}}\mathscr{L}=\begin{bmatrix}\boldsymbol{\lambda}_{\ell+1}^{R}\\\
\boldsymbol{\lambda}_{\ell+1}^{I}\end{bmatrix}^{T}\nabla_{\boldsymbol{v}^{C}_{\ell}}\begin{bmatrix}\boldsymbol{F}_{\Delta
t}^{R}\\\ \boldsymbol{F}_{\Delta
t}^{I}\end{bmatrix}(\boldsymbol{\phi}_{\ell}^{R},\boldsymbol{\phi}_{\ell}^{I},\boldsymbol{v}^{C}_{\ell}).$
(11)

Given a candidate $\boldsymbol{v}^{C}$, we solve the forward problem to obtain
$\boldsymbol{\Phi}$. We then solve the adjoint system to obtain
$\boldsymbol{\Lambda}$. This provides everything required to evaluate (11) for
each $\ell$. The variations, the block matrix form of the Jacobian
$\boldsymbol{J}_{\boldsymbol{\phi}}\boldsymbol{F}_{\Delta t}$, and the
gradients of the discrete-time propagator $\boldsymbol{F}$ can be found in
Sections 5.3 and 5.4 of the preprint Appendix.

#### Second Adjoint Method: Learning $v^{C}$ Functionals.

Here we rederive the adjoint method to enable learning the _functional
dependence_ of $v^{C}[\phi](x,t)$ on $\phi(x,t)$ and $\phi(x,t-\Delta t)$. We
take as our model
$v^{C}[\phi]=v^{C}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})$.
The parameters $\boldsymbol{\theta}$ determine a particular functional
dependence of $v^{C}$ on the present and previous Kohn-Sham states
$\boldsymbol{\phi}$ and $\boldsymbol{\phi}^{\prime}$. At spatial grid location
$x_{j}$ and time $t_{k}$, the model $v^{C}$ is

$v^{C}[\phi](x_{j},t_{k})=[\boldsymbol{v}^{C}(\boldsymbol{\phi}_{k},\boldsymbol{\phi}_{k-1};\boldsymbol{\theta})]_{j}.$
(12)

In short, we intend $\boldsymbol{\phi}^{\prime}$ to be the Kohn-Sham state at
the time step _prior_ to the time step that corresponds to
$\boldsymbol{\phi}$. Our goal is to learn $\boldsymbol{\theta}$. This requires
redefining the following quantities:

$\displaystyle V(\boldsymbol{\phi};\boldsymbol{\theta})$
$\displaystyle=\operatorname{diag}(\mathbf{v}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta}))$
$\displaystyle\boldsymbol{v}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})$
$\displaystyle=-((\boldsymbol{x}+10)^{2}+1)^{-1/2}+W(|\boldsymbol{\phi}|^{2}\circ\boldsymbol{w})+\boldsymbol{v}^{C}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})$
$\displaystyle\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})$
$\displaystyle=\mathcal{P}_{\mathcal{K}}\exp(-iV(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})\Delta
t)\mathcal{P}_{\mathcal{K}}\boldsymbol{\phi}.$

The Lagrangian still has the form of an objective function together with a
dynamical constraint:

$\mathscr{L}(\boldsymbol{\Phi}^{R},\boldsymbol{\Phi}^{I},\boldsymbol{\Lambda}^{R},\boldsymbol{\Lambda}^{I},\boldsymbol{\theta})=\frac{1}{2}\sum_{k=0}^{K}\sum_{j=0}^{J}(2\phi^{R}(x_{j},t_{k})^{2}+2\phi^{I}(x_{j},t_{k})^{2}-\tilde{n}(x_{j},t_{k}))^{2}-\sum_{k=1}^{K-1}[\boldsymbol{\lambda}_{k+1}^{R}]^{T}\\\
(\boldsymbol{\phi}_{k+1}^{R}-\boldsymbol{F}^{R}_{\Delta
t}(\boldsymbol{\phi}_{k}^{R},\boldsymbol{\phi}_{k-1}^{R},\boldsymbol{\phi}_{k}^{I},\boldsymbol{\phi}_{k-1}^{I};\boldsymbol{\theta}))+[\boldsymbol{\lambda}_{k+1}^{I}]^{T}(\boldsymbol{\phi}_{k+1}^{I}-\boldsymbol{F}^{I}_{\Delta
t}(\boldsymbol{\phi}_{k}^{R},\boldsymbol{\phi}_{k-1}^{R},\boldsymbol{\phi}_{k}^{I},\boldsymbol{\phi}_{k-1}^{I};\boldsymbol{\theta}))$
(13)

Setting $\delta\mathscr{L}=0$ for all variations $\boldsymbol{\phi}^{R}_{k}$
and $\boldsymbol{\phi}^{I}_{k}$ for $k\geq 1$, we obtain the following adjoint
system:

$\displaystyle\boldsymbol{\lambda}_{K}$
$\displaystyle=4\left[(2|\boldsymbol{\phi}_{K}|^{2}-\tilde{\boldsymbol{n}}_{K})\circ\boldsymbol{\phi}_{K}\right]$
(14a) $\displaystyle\boldsymbol{\lambda}_{K-1}$
$\displaystyle=4\left[(2|\boldsymbol{\phi}_{K-1}|^{2}-\tilde{\boldsymbol{n}}_{K-1})\circ\boldsymbol{\phi}_{K-1}\right]$
(14b)
$\displaystyle\qquad+[\boldsymbol{\lambda}_{K}^{R}]^{T}\nabla_{\boldsymbol{\phi}}\boldsymbol{F}_{\Delta
t}^{R}(\boldsymbol{\phi}_{K},\boldsymbol{\phi}_{K-1};\boldsymbol{\theta})+[\boldsymbol{\lambda}_{K}^{I}]^{T}\nabla_{\boldsymbol{\phi}}\boldsymbol{F}_{\Delta
t}^{I}(\boldsymbol{\phi}_{K},\boldsymbol{\phi}_{K-1};\boldsymbol{\theta})$
$\displaystyle\begin{bmatrix}\boldsymbol{\lambda}_{k}^{R}\\\
\boldsymbol{\lambda}_{k}^{I}\end{bmatrix}^{T}\\!\\!\\!\\!$
$\displaystyle=4(2|\boldsymbol{\phi}_{k}|^{2}\\!\\!-\\!\tilde{\boldsymbol{n}}_{k})\\!\circ\\!\begin{bmatrix}\boldsymbol{\phi}_{k}^{R}\\\
\boldsymbol{\phi}_{k}^{I}\end{bmatrix}^{T}\\!\\!\\!\\!\\!+\\!\begin{bmatrix}\boldsymbol{\lambda}_{k+1}^{R}\\\
\boldsymbol{\lambda}_{k+1}^{I}\end{bmatrix}^{T}\\!\\!\\!\boldsymbol{J}_{\boldsymbol{\phi}}\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi}_{k},\boldsymbol{\phi}_{k-1};\boldsymbol{\theta})+\begin{bmatrix}\boldsymbol{\lambda}_{k+2}^{R}\\\
\boldsymbol{\lambda}_{k+2}^{I}\end{bmatrix}^{T}\\!\\!\\!\boldsymbol{J}_{\boldsymbol{\phi}^{\prime}}\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi}_{k+1},\boldsymbol{\phi}_{k};\boldsymbol{\theta})$ (14c)

The key difference between (14c) and (10b) is that the right-hand side of
(14c) involves $\boldsymbol{\lambda}$ at two points in time. The adjoint
system is now a linear delay difference equation with time-dependent
coefficients. Additionally, the derivatives of $F_{\Delta t}$ needed to
evaluate (14-15) are different—see Section 5.4 of the preprint Appendix. For a
candidate value of $\boldsymbol{\theta}$, we solve the forward problem to
obtain $\boldsymbol{\phi}$ on our spatial and temporal grid. Then, to compute
gradients, we begin with the final conditions (14a-14b) and iterate (14c)
backwards in time from $k=K-2$ to $k=1$. Having solved the adjoint system, we
compute the gradient of $\mathscr{L}$ with respect to $\boldsymbol{\theta}$
via

$\nabla_{\boldsymbol{\theta}}\mathscr{L}=\sum_{k=1}^{K-1}\begin{bmatrix}\boldsymbol{\lambda}_{k+1}^{R}\\\
\boldsymbol{\lambda}_{k+1}^{I}\end{bmatrix}^{T}\nabla_{\boldsymbol{\theta}}\begin{bmatrix}\boldsymbol{F}_{\Delta
t}^{R}\\\ \boldsymbol{F}_{\Delta
t}^{I}\end{bmatrix}(\boldsymbol{\phi}_{k},\boldsymbol{\phi}_{k-1};\boldsymbol{\theta}).$
(15)

Figure 1: Training results for the problem of learning pointwise values of
$\boldsymbol{V}^{C}$ on a grid consisting of $K=30000$ points in time and
$J=600$ points in space. The adjoint method succeeds in producing
$\boldsymbol{V}^{C}$ values that yield TDKS solutions such that the
corresponding electron densities (red) match those computed from the 2D
Schrödinger equation (black).

##  3 Modeling and Implementation Details

#### Modeling Correlation Functionals.

In this work, all models of the form (12) consist of dense, feedforward neural
networks. For models of the form
$\boldsymbol{v}^{C}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})$,
we treat the real and imaginary parts of $\boldsymbol{\phi}$ and
$\boldsymbol{\phi}^{\prime}$ as real vectors each of length $J+1$. Hence for
$J=600$, we have an input layer of size $4(J+1)$. We follow this with three
hidden layers each with $256$ units and a scaled exponential linear unit
activation function (Klambauer et al., 2017). The output layer has $J+1$ units
to correspond to the vector-valued output $\boldsymbol{v}^{C}$. For models in
which $\boldsymbol{v}^{C}$ depend on $\boldsymbol{n}$ and
$\boldsymbol{n}^{\prime}$, we take the real and imaginary parts of
$\boldsymbol{\phi}$ and $\boldsymbol{\phi}^{\prime}$ as inputs and use them to
immediately compute $\boldsymbol{n}$ and $\boldsymbol{n}^{\prime}$, which we
then concatenate and feed into an input layer with $2(J+1)$ units. The
remainder of the network is as above. We started with smaller networks (fewer
layers, less units per layer) and increased the network size until we obtained
reasonable training results; no other architecture search or hyperparameter
tuning was carried out. We experimented with other activation functions and
convolutional layers—none of these models produced satisfactory results during
training.

#### Generation of Training Data.

To generate training data, we solve (1) for a model system consisting of $N=2$
electrons: a one-dimensional electron scattering off a one-dimensional
hydrogen atom. Hence (1) becomes a partial differential equation (PDE) for a
wave function $\Psi(x_{1},x_{2},t)$. We discretize this PDE using finite
differences on an equispaced grid in $(x_{1},x_{2})$ space with $J=1201$
points along each axis. Here $-80\leq x_{1},x_{2}\leq 40$, so that $\Delta
x=0.1$. After discretizing the kinetic and potential operators in space, we
propagate forward in time until $T=0.72$ fs, using second-order operator
splitting with $\Delta t=2.4\times 10^{-5}$ fs (or, in a.u., $\Delta t\approx
9.99219\times 10^{-4}$). Note that this is $1/100$-th the time step used by
Suzuki et al. (2020). For further details, consult Section 5.2. After
discretization, the wave function $\Psi(x_{1},x_{2},t)$ at time step $k$ is a
complex vector $\boldsymbol{\psi}_{k}$ of dimension $(J+1)^{2}$. For the
initial vector $\boldsymbol{\psi}_{0}$, we follow Suzuki et al. (2020) and use
a Gaussian wave packet that represents an electron initially centered at
$x=10$ a.u., approaching the H-atom localized at $x=-10$ a.u., with momentum
$p$. We generate training/test data by numerically solving the Schrödinger
system for initial conditions with $p\in\\{-1.0,-1.2,-1.4,-1.5,-1.6,-1.8\\}$.
From the resulting time series of wave functions, we compute the time-
dependent one-electron density $n(x,t)$; below, _we refer to this as the TDSE
electron density_.

##  4 Results

Figure 2: We use $300$ time steps (corresponding to $0.72$ fs) of the $p=-1.5$
data together with the adjoint method to train a neural network model of
$v^{C}$ that depends on the current and previous $\phi$. Using the learned
$v^{C}$, we propagate (5) for $60$ additional time steps and plot the test set
results (in red) against the reference electron density (in blue).

#### Pointwise Results.

Our first result concerns learning the pointwise values of
$\boldsymbol{V}^{C}$. Here we use the same fine time step $\Delta t=2.4\times
10^{-5}$ used to generate the training data. However, we increase $\Delta x$
by a factor of $2$, taking $J=600$ and sampling the initial condition
$\boldsymbol{\phi}_{0}$ at every other grid point. We retain this subsampling
in space in all training sets/results that follow. Still, our unknown
$\boldsymbol{V}^{C}$ consists of a total of $30000\cdot 601$ values.

We learn $\boldsymbol{V}^{C}$ by optimizing an objective function that
consists of the first line of (9) together with a regularization term. The
regularization consists of a finite-difference approximation of
$\mu\sum_{k}\sum_{j}(\partial_{x}v^{C}(x_{j},t_{k}))^{2}$, with $\mu=10^{-5}$.
This regularization is analogous to the $\int(f^{\prime\prime})^{2}dx$ penalty
used in smoothing splines (Hastie et al., 2009). We penalize the square of the
first (rather than second) derivative as we find this is sufficient to smooth
$v^{C}$ in space. The precise value of $\mu$ is unimportant; taking
$\mu\in[10^{-6},10^{-4}]$ yields similar results. For training data, we use
only the TDSE one-electron densities computed from the $p=-1.5$ initial
condition. To optimize, we use the quasi-Newton L-BFGS-B method, with
gradients $\nabla_{\boldsymbol{V}^{C}}\mathscr{L}$ computed via the procedure
described just below (11). We initialize the optimizer with
$\boldsymbol{V}^{C}\equiv 0$ and use default tolerances of $10^{-6}$.

In Figure 1, we present the results of this approach. Each panel shows a
snapshot of both the training electron density (in black, computed from TDSE
data) and the electron density $n=2|\phi|^{2}$ (in red) obtained by solving
TDKS (5) using the learned $\boldsymbol{V}^{C}$ values. Note the close
quantitative agreement between the black and red curves. The overall mean-
squared error (MSE) across all points in space and time is $2.035\times
10^{-6}$. Note that no exact $\boldsymbol{V}^{C}$ data was used; the learned
$\boldsymbol{V}^{C}$ does not match the exact $\boldsymbol{V}^{C}$
quantitatively, but does have some of the same qualitative features.

This problem suits the adjoint method well: regardless of the dimensionality
of $\boldsymbol{V}^{C}$, the dimensionality of the adjoint system is the same
as that of the discretized TDKS system. Note that, for this one-dimensional
TDKS problem (5), it is possible to solve for $\boldsymbol{V}^{C}$ on a grid
(Elliott et al., 2012). If we encounter solutions of _higher-dimensional,
multi-electron_ ($d\geq 2$ and $N\geq 2$) Schrödinger systems from which we
seek to learn $\boldsymbol{V}^{C}$, we will not be able to employ an exact
procedure. In this case, the adjoint-based method may yield numerical values
$\boldsymbol{V}^{C}$, with which we can pursue supervised learning of a
functional from electron densities $\boldsymbol{n}$ to correlation potentials
$\boldsymbol{V}^{C}$.

Figure 3: We plot training and test results at time $t=0.432$ fs for the
adjoint method, applied to estimating neural network models $v^{C}[\phi]$
that, at time $t$, depend on both $\phi(x,t)$ and $\phi(x,t-\Delta t)$.
Propagating TDKS (5) with the learned $v^{C}$ yields the red curves.

#### Functional Results.

Next we present results in which we learn $v^{C}$ functionals. In preliminary
work, we sought to model $v^{C}[\phi](x,t)$ as purely a function of
$\phi(x,t)$, a model without memory. These models did not yield satisfactory
training set results, and hence were abandoned. We focus first on models
$v^{C}[\phi](x,t)$ that allow for arbitrary dependence on the real and
imaginary parts of both $\phi(x,t)$ and $\phi(x,t-\Delta t)$. The TDDFT
literature emphasizes that $v^{C}$ should depend on $\phi$ through
present/past electron densities $n$, where $n=2|\phi|^{2}$. How important is
it to incorporate such physics-based constraints into our $v^{C}$ model? Let
us see how well a direct neural network model of $v^{C}[\phi]$ captures the
dynamics. The input layer is of dimension $4(J+1)$—see Section 3.

To train such a model, we again apply the L-BFGS-B optimizer with objective
function given by the first line of (13) and gradients computed with the
adjoint system (14-15). We initialize neural network parameters
$\boldsymbol{\theta}$ by sampling a mean-zero normal distribution with
standard deviation $\sigma=0.01$. For training data, we subsample the $p=-1.5$
TDSE electron density time series by a factor of $100$ in time, so that
$\Delta t=2.4\times 10^{-3}$ fs and the entire training trajectory consists of
$K=301$ time steps. We retain this time step in all training sets and results
that follow.

We omit the training set results here as they show excellent agreement between
training and model-predicted electron densities—see Section 5.6. The overall
training set mean-squared error (MSE) is $7.668\times 10^{-6}$. In Figure 2,
we display test set results obtained by propagating for $60$ additional time
steps beyond the end of the training data. On this test set, we see close
quantitative agreement near $t=0.72$ fs, which slowly degrades. Still, the
learned $v^{C}$ leads to TDKS electron densities that capture essential
features of the reference trajectory. Note that no regularization was used
during training of the $v^{C}$ functional, leading to a learned $v^{C}$ that
is not particularly smooth in space. We hypothesize that, with careful and
perhaps physically motivated regularization, the learned $v^{C}$ will yield
improved test set results over longer time intervals.

In the next set of results, we retrain our model using TDSE electron densities
with initial momenta equal to $p=-1.0$ and $p=-1.8$. We train two models: a
$v^{C}[\phi](x,t)$ model that depends on $\phi$ at times $t$ and $t-\Delta t$,
and a $v^{C}[n](x,t)$ model that depends on $n$ at times $t$ and $t-\Delta t$.
This latter model incorporates the physics-based constraints mentioned above.
We view the $v^{C}[n]$ model as more constrained because its the first hidden
layer can depend on $\phi(x,t)$ and $\phi(x,t-\Delta t)$ _only through_ the
electron densities $n(x,t)$ and $n(x,t-\Delta t)$. We keep all other details
of training the same. The final training set MSE values are $4.645\times
10^{-5}$ for the $v^{C}[\phi]$ model and $8.098\times 10^{-5}$ for the
$v^{C}[n]$ model.

In Figures 3 and 4, we plot both training and test set results for these
models. Here we have chosen a particular time ($t=0.432$ fs) and plotted the
electron density at this time for six different trajectories, each with a
different initial momentum $p$. We have chosen this time to highlight the
large, obvious differences between the $p=-1.0$ and $p=-1.8$ curves. The
$p=-1.0$ and $p=-1.8$ panels contain training set results; here the TDKS
electron densities (in red, produced using the learned $v^{C}$) lie closer to
the ground truth TDSE electron densities (in black).

Note that, despite the greater freedom enjoyed by the $v^{C}[\phi]$ model, its
generalization to trajectories _outside the training set_ ($-1.2\leq
p\leq-1.6$) is noticeably worse than that of the more constrained $v^{C}[n]$
model. In fact, the $v^{C}[n]$ model’s results (Figure 4, in red) show broad
qualitative agreement with the test set TDSE curves (in blue). The test set
MSE values are $9.363\times 10^{-4}$ for the $v^{C}[\phi]$ model and
$2.482\times 10^{-4}$ for the $v^{C}[n]$ model. Overall, these results support
the view that $v^{C}$ should depend on $\phi$ through $n$. Again, we
hypothesize that if we were to filter out short-wavelength oscillations in the
electron density—perhaps by regularizing the $v^{C}[n]$ model or by training
on a larger set of trajectories—the agreement could be improved.

Figure 4: We plot training and test results at time $t=0.432$ fs for the
adjoint method, applied to estimating a neural network model $v^{C}[n]$ that,
at time $t$, depends on both $n(x,t)$ and $n(x,t-\Delta t)$. Propagating TDKS
(5) with the learned $v^{C}[n]$ yields the red curves.

#### Conclusion.

For a low-dimensional model problem, we have developed adjoint-based methods
to learn the correlation potential $v^{C}$ using data from TDSE simulations.
The adjoint method can be used to _directly train_ $v^{C}[n]$ models,
sidestepping the need for either exact $v^{C}$ values or density-to-potential
inversion. Our work provides a foundation for learning models that depend on
present and past snapshots of the electron density. We find that our trained
$v^{C}[n]$ models (with memory) generalize well to trajectories outside the
training set. Further improvements to the model may be possible, e.g., by
incorporating known physics in the form of model constraints. Overall, the
results show the promise of learning $v^{C}$ via TDKS-constrained
optimization.

This work was supported by the U.S. Department of Energy, Office of Science,
Basic Energy Sciences under Award Number DE-SC0020203. This research used
resources of the National Energy Research Scientific Computing Center (NERSC),
a U.S. Department of Energy Office of Science User Facility located at
Lawrence Berkeley National Laboratory, operated under Contract No. DE-
AC02-05CH11231 using NERSC award BES-m2530 for 2021. We acknowledge
computational time on the Pinnacles cluster at UC Merced (supported by NSF
OAC-2019144). We also acknowledge computational time on the Nautilus cluster,
supported by the Pacific Research Platform (NSF ACI-1541349), CHASE-CI (NSF
CNS-1730158), and Towards a National Research Platform (NSF OAC-1826967).
Additional funding for Nautilus has been supplied by the University of
California Office of the President.

##  5 Appendix

###  5.1 Solving the Spatially One-Dimensional TDKS System

To solve (5), we use a finite-difference discretization of (5) on the spatial
domain $x\in[L_{\text{min}},L_{\text{max}}]$ and temporal domain $t\in[0,T]$.
Fix spatial and temporal grid spacings $\Delta x>0$ and $\Delta t>0$; let
$x_{j}=L_{\text{min}}+j\Delta x$ and $t_{k}=k\Delta t$. Then our spatial grid
is $\\{x_{j}\\}_{j=0}^{j=J}$ with $J\Delta x=L_{\text{max}}-L_{\text{min}}$,
and our temporal grid is $\\{t_{k}\\}_{k=0}^{K}$ with $K\Delta t=T$.

Suppose the correlation functional has been specified in one of two ways: (i)
for a particular trajectory, we have access to the _values_ of
$v^{C}(x_{j},t_{k})$ at all spatial and temporal grid points $x_{j}$ and
$t_{k}$, or (ii) we have a model $v^{C}[\phi]$ that takes as input $\phi(x,s)$
for $s\leq t_{k}$ and produces as output $v^{C}(x_{j},t_{k})$ for all $j$.
Then, given an initial condition $\phi(x,0)$, the _forward problem_ is to
solve (5) numerically on the grids defined above.

Our first step is to discretize (5) in space. Let $\boldsymbol{\phi}(t)$ be
the $(J+1)\times 1$ column vector

$\boldsymbol{\phi}(t)=[\phi(x_{0},t),\ldots,\phi(x_{J},t)]^{T},$ (16)

where T denotes transpose. We use the integers $0,\ldots,J$ to index this
vector, so that $[\boldsymbol{\phi}(t)]_{j}=\phi(x_{j},t)$. We discretize
$\partial_{xx}$ with a fourth-order Laplacian matrix $\Delta$, defined in
(19), such that

$\partial_{xx}\phi(x,t)\bigr{|}_{x=x_{j}}=[\Delta\boldsymbol{\phi}(t)]_{j}+O(\Delta
x^{4}).$ (17)

Besides the $\partial_{xx}$ term, the remaining terms on the right-hand side
of (5) result in a diagonal matrix multiplied by $\boldsymbol{\phi}$. The only
term that requires further numerical approximation is the integral term. For
this purpose, we define the symmetric matrix
$W_{j,j^{\prime}}=((x_{j^{\prime}}-x_{j})^{2}+1)^{-1/2}\Delta x$. Let $\circ$
denote the entry-wise product of vectors and let $|\boldsymbol{\phi}(t)|^{2}$
be the vector of entry-wise magnitudes $|\phi(x_{j},t)|^{2}$. With quadrature
weights $w_{j}$ given by Simpson’s rule, the integral at $x=x_{j}$ becomes

$\displaystyle\int_{x^{\prime}}((x^{\prime}-x_{j})^{2}+1)^{-1/2}|\phi(x^{\prime},t)|^{2}dx^{\prime}$
$\displaystyle=\sum_{j^{\prime}}W_{j,j^{\prime}}\phi(x_{j^{\prime}},t)\phi^{\ast}(x_{j^{\prime}},t)w_{j^{\prime}}+O(\Delta
x^{4})$
$\displaystyle=\left[W(|\boldsymbol{\phi}(t)|^{2}\circ\boldsymbol{w})\right]_{j}+O(\Delta
x^{4})$ (18)

Let $\mathcal{K}=-(1/2)\Delta$ be the spatially discretized kinetic operator,
with $\Delta$ defined by the following fourth-order discrete Laplacian:

$\Delta=\frac{1}{12\Delta x^{2}}\begin{bmatrix}-30&16&-1&&\\\
16&\ddots&\ddots&\ddots&\\\ -1&\ddots&\ddots&\ddots&\ddots\\\
&\ddots&\ddots&\ddots&\ddots&-1\\\ &&\ddots&\ddots&\ddots&16\\\
&&&-1&16&-30\end{bmatrix}.$ (19)

Evaluating both sides of (5) at $x=x_{j}$ for all $j$ at once, we arrive at
the following nonlinear system of ordinary differential equations (ODE) for
$\boldsymbol{\phi}=\boldsymbol{\phi}(t)$:

$i\frac{d}{dt}\boldsymbol{\phi}=\mathcal{K}\boldsymbol{\phi}+V(\boldsymbol{\phi},\boldsymbol{v}^{C})\boldsymbol{\phi},$
(20)

where $V(\boldsymbol{\phi},\boldsymbol{v}^{C})$ is a diagonal matrix whose
diagonal is the vector
$\boldsymbol{v}(\boldsymbol{\phi},\boldsymbol{v}^{C})=-((\boldsymbol{x}+10)^{2}+1)^{-1/2}+W(|\boldsymbol{\phi}|^{2}\circ\boldsymbol{w})+\boldsymbol{v}^{C}$.
Here $\boldsymbol{x}$ is the vector whose $j$-th entry is $x_{j}$ and
operations involving $\boldsymbol{x}$ should be interpreted entry-wise. We
have deliberately kept $\boldsymbol{v}^{C}$ general to encompass both the
cases where (i) $\boldsymbol{v}^{C}(t)$ is a vector of time-dependent
parameters whose $j$-th entry is $v^{C}[\phi](x_{j},t)$, and (ii)
$\boldsymbol{v}^{C}$ is a function that takes as input, e.g.,
$\\{\boldsymbol{\phi}_{k},\boldsymbol{\phi}_{k-1},\ldots\boldsymbol{\phi}_{0}\\}$
and produces as output the values $v^{C}[\phi](x_{j},t_{k})$.

To solve (20), we apply operator splitting (Castro et al., 2004), resulting in
the fully discretized propagation equation

$\boldsymbol{\phi}(t_{k+1})=\exp(-i\mathcal{K}\Delta
t/2)\exp(-iV(\boldsymbol{\phi}(t_{k}),\boldsymbol{v}^{C}_{k})\Delta
t)\exp(-i\mathcal{K}\Delta t/2)\boldsymbol{\phi}(t_{k}).$ (21)

We choose this method for two reasons. First, the propagator is unitary and
hence preserves the normalization of $\boldsymbol{\phi}$ over long times.
Second, as $V$ is diagonal, both the matrix exponential of $-iV\Delta t$ and
its Jacobian with respect to $\boldsymbol{v}^{C}$ are simple to calculate. The
ease with which we can compute derivatives of the right-hand side of (21)
balances its second-order accuracy in time.

Note that $\mathcal{K}$ is time-independent and symmetric. For small systems,
$\exp(-i\mathcal{K}\Delta t/2)$ can be computed by diagonalizing
$\mathcal{K}$. If $\mathcal{K}=SDS^{-1}$, then $\exp(-i\mathcal{K}\Delta
t/2)=S\exp(-iD\Delta t/2)S^{-1}$. During the initial part of our codes, we
compute this kinetic propagator once and store it for future use. The adjoint
derivations below can be extended straightforwardly to higher-order version of
operator splitting, as long as the discrete propagation scheme involves
alternating products of kinetic and potential propagation terms as in (21),
with $V$ and $\boldsymbol{v}^{C}$ exponentiated diagonally.

###  5.2 Solving the Spatially Two-Dimensional (2D) Schrödinger System

To generate training data, we numerically solve the 2D Schrödinger model
system with Hamiltonian $\hat{H}=\hat{K}+\hat{V}$. Here the kinetic operator
is $\hat{K}=-(1/2)\nabla^{2}$, and $\nabla^{2}$ is the two-dimensional
Laplacian. The electronic potential $\hat{V}$ consists of a sum of electron-
nuclear and electron-electron terms:

$\hat{V}(x_{1},x_{2})=v^{\text{ext}}(x_{1})+v^{\text{ext}}(x_{2})+W^{ee}(x_{1},x_{2}),$
(22)

with $v^{\text{ext}}$ and $W^{ee}$ defined via the soft-Coulomb potentials
(4a) and (4b), respectively.

For forward time-evolution of the 2D Schrödinger model system, we use a
discrete Laplacian $\Delta_{2}$ that consists of

$\Delta_{2}=\Delta\otimes I+I\otimes\Delta,$

where $I$ is the $(J+1)\times(J+1)$ identity matrix and $\otimes$ denotes the
Kronecker product. The resulting $\Delta_{2}$ is a fourth-order approximation
to the two-dimensional Laplacian
$\nabla^{2}=\partial_{x_{1}x_{1}}+\partial_{x_{2}x_{2}}$. In the Schrödinger
system, spatially discretizing the kinetic operator yields the matrix
$\mathcal{K}=-(1/2)\Delta_{2}$, which is of dimension
$(J+1)^{2}\times(J+1)^{2}$ with $J=1200$. To compute the kinetic portion of
the propagator,

$\mathcal{P}_{\mathcal{K}}=\exp(-i\mathcal{K}\Delta t/2),$

we used a straightforward fourth-order series expansion of the matrix
exponential:

$\mathcal{P}_{\mathcal{K}}\approx\sum_{j=0}^{4}\frac{(-i\mathcal{K}\Delta
t/2)^{j}}{j!}.$

With $\Delta t=2.4\times 10^{-5}$ fs (or, in a.u., $\Delta t\approx
9.99219\times 10^{-4}$), this series approximation of the matrix exponential
incurs negligible error.

The potential portion of the propagator, $\mathcal{P}_{V}=\exp(-iV\Delta t)$,
is a purely diagonal matrix—the $(J+1)^{2}$ entries along its diagonal consist
of a flattened version of the $(J+1)\times(J+1)$ matrix obtained by evaluating
(22) on our finite-difference spatial grid.

Equipped with $\mathcal{P}_{\mathcal{K}}$ and $\mathcal{P}_{V}$, both of which
are time-independent, we propagate forward using second-order operator
splitting as in (8):

$\boldsymbol{\psi}_{k+1}=\mathcal{P}_{\mathcal{K}}\mathcal{P}_{V}\mathcal{P}_{\mathcal{K}}\boldsymbol{\psi}_{k}.$
(23)

As a numerical method for the TDSE (1), operator splitting goes back at least
to the work of Fleck Jr et al. (1976) and Feit et al. (1982). Starting from an
initial condition $\boldsymbol{\psi}_{0}$ represented as a complex vector of
dimension $(J+1)^{2}$, we iterate for $K=30000$ steps until we reach a final
time of $0.72$ fs. We have implemented the above Schrödinger solver using
sparse linear algebra and CuPy.

#### Initializing TDKS Simulations with Memory.

When we solve (5) with a correlation potential with memory, _e.g._ ,
$v^{C}[\phi]$ that depends on both $\phi(x,t)$ and $\phi(x,t-\Delta t)$, how
do we initialize the simulation? Our solution is to start with the wave
function data generated by solving the TDSE (as above). With this data, we
apply the methods from Ullrich (2011, Appendix E) to compute exact Kohn-Sham
states $\boldsymbol{\phi}_{k}$ corresponding to $k=0$ and $k=1$ with $\Delta
t=2.4\times 10^{-3}$ fs. We use the exact $\boldsymbol{\phi}_{0}$ and
$\boldsymbol{\phi}_{1}$ to initialize our TDKS simulations when we use a
$v^{C}$ model with memory.

###  5.3 Derivation of the Adjoint System

Variations of (9) with respect to $\boldsymbol{\Lambda}^{R}$ and
$\boldsymbol{\Lambda}^{I}$ give the real and imaginary parts of the equality
constraint (21). For the variation with respect to $\boldsymbol{\Phi}^{R}$, we
obtain

$\displaystyle\delta\mathscr{L}=\frac{d}{d\epsilon}\biggr{|}_{\epsilon=0}\mathscr{L}(\boldsymbol{\Phi}^{R}+\epsilon\delta\boldsymbol{\Phi}^{R},\boldsymbol{\Phi}^{I},\boldsymbol{\Lambda}^{R},\boldsymbol{\Lambda}^{I},\boldsymbol{v}^{C})$
$\displaystyle=\sum_{k=0}^{K}\sum_{j=0}^{J}(2\phi^{R}(x_{j},t_{k})^{2}+2\phi^{I}(x_{j},t_{k})^{2}-\tilde{n}(x_{j},t_{k}))(4\phi^{R}(x_{j},t_{k})\delta\phi^{R}(x_{j},t_{k}))$
$\displaystyle\quad-\sum_{k=0}^{K-1}[\boldsymbol{\lambda}_{k+1}^{R}]^{T}\delta\boldsymbol{\phi}_{k+1}^{R}$
$\displaystyle\quad+\sum_{k=0}^{K-1}[\boldsymbol{\lambda}_{k+1}^{R}]^{T}\nabla_{\boldsymbol{\phi}^{R}}\boldsymbol{F}^{R}_{\Delta
t}(\boldsymbol{\phi}_{k}^{R},\boldsymbol{\phi}_{k}^{I},\boldsymbol{v}^{C}_{k})\delta\boldsymbol{\phi}^{R}_{k}+[\boldsymbol{\lambda}_{k+1}^{I}]^{T}\nabla_{\boldsymbol{\phi}^{R}}\boldsymbol{F}^{I}_{\Delta
t}(\boldsymbol{\phi}_{k}^{R},\boldsymbol{\phi}_{k}^{I},\boldsymbol{v}^{C}_{k})\delta\boldsymbol{\phi}^{R}_{k}$
$\displaystyle=\sum_{k=0}^{K}4\left[(2|\boldsymbol{\phi}_{k}|^{2}-\tilde{\boldsymbol{n}}_{k})\circ\boldsymbol{\phi}_{k}^{R}\right]^{T}\delta\boldsymbol{\phi}^{R}_{k}-\sum_{k=1}^{K}[\boldsymbol{\lambda}_{k}^{R}]^{T}\delta\boldsymbol{\phi}_{k}^{R}+\sum_{k=0}^{K-1}\begin{bmatrix}\boldsymbol{\lambda}_{k+1}^{R}\\\
\boldsymbol{\lambda}_{k+1}^{I}\end{bmatrix}^{T}\nabla_{\boldsymbol{\phi}^{R}}\begin{bmatrix}\boldsymbol{F}_{\Delta
t}^{R}\\\ \boldsymbol{F}_{\Delta
t}^{I}\end{bmatrix}\delta\boldsymbol{\phi}_{k}^{R}.$

Analogously, for the variation with respect to $\boldsymbol{\Phi}^{I}$, we
obtain

$\delta\mathscr{L}=\sum_{k=0}^{K}4\left[(2|\boldsymbol{\phi}_{k}|^{2}-\tilde{\boldsymbol{n}}_{k})\circ\boldsymbol{\phi}_{k}^{I}\right]^{T}\delta\boldsymbol{\phi}^{I}_{k}-\sum_{k=1}^{K}[\boldsymbol{\lambda}_{k}^{I}]^{T}\delta\boldsymbol{\phi}_{k}^{I}+\sum_{k=0}^{K-1}\begin{bmatrix}\boldsymbol{\lambda}_{k+1}^{R}\\\
\boldsymbol{\lambda}_{k+1}^{I}\end{bmatrix}^{T}\nabla_{\boldsymbol{\phi}^{I}}\begin{bmatrix}\boldsymbol{F}_{\Delta
t}^{R}\\\ \boldsymbol{F}_{\Delta
t}^{I}\end{bmatrix}\delta\boldsymbol{\phi}_{k}^{I}.$

Setting $\delta\mathscr{L}=0$ for all variations
$\delta\boldsymbol{\phi}_{k}^{R}$ and $\delta\boldsymbol{\phi}_{k}^{I}$ for
$k\geq 1$, we obtain the following backward-in-time system for
$\boldsymbol{\Lambda}$:

$\displaystyle\boldsymbol{\lambda}_{K}$
$\displaystyle=4\left[(2|\boldsymbol{\phi}_{K}|^{2}-\tilde{\boldsymbol{n}}_{K})\circ\boldsymbol{\phi}_{K}\right]$
(24a) $\displaystyle[\boldsymbol{\lambda}_{k}^{R}]^{T}$
$\displaystyle=4\left[(2|\boldsymbol{\phi}_{k}|^{2}-\tilde{\boldsymbol{n}}_{k})\circ\boldsymbol{\phi}_{k}^{R}\right]^{T}+\begin{bmatrix}\boldsymbol{\lambda}_{k+1}^{R}\\\
\boldsymbol{\lambda}_{k+1}^{I}\end{bmatrix}^{T}\nabla_{\boldsymbol{\phi}^{R}}\begin{bmatrix}\boldsymbol{F}_{\Delta
t}^{R}\\\ \boldsymbol{F}_{\Delta
t}^{I}\end{bmatrix}(\boldsymbol{\phi}_{k}^{R},\boldsymbol{\phi}_{k}^{I},\boldsymbol{v}^{C}_{k})$
(24b) $\displaystyle[\boldsymbol{\lambda}_{k}^{I}]^{T}$
$\displaystyle=4\left[(2|\boldsymbol{\phi}_{k}|^{2}-\tilde{\boldsymbol{n}}_{k})\circ\boldsymbol{\phi}_{k}^{I}\right]^{T}+\begin{bmatrix}\boldsymbol{\lambda}_{k+1}^{R}\\\
\boldsymbol{\lambda}_{k+1}^{I}\end{bmatrix}^{T}\nabla_{\boldsymbol{\phi}^{I}}\begin{bmatrix}\boldsymbol{F}_{\Delta
t}^{R}\\\ \boldsymbol{F}_{\Delta
t}^{I}\end{bmatrix}(\boldsymbol{\phi}_{k}^{R},\boldsymbol{\phi}_{k}^{I},\boldsymbol{v}^{C}_{k})$
(24c)

We can write the Jacobian as a block matrix:

$\boldsymbol{J}_{\boldsymbol{\phi}}\boldsymbol{F}_{\Delta
t}=\begin{bmatrix}\nabla_{\boldsymbol{\phi}^{R}}\boldsymbol{F}^{R}_{\Delta
t}&\nabla_{\boldsymbol{\phi}^{I}}\boldsymbol{F}^{R}_{\Delta t}\\\
\nabla_{\boldsymbol{\phi}^{R}}\boldsymbol{F}^{I}_{\Delta
t}&\nabla_{\boldsymbol{\phi}^{I}}\boldsymbol{F}^{I}_{\Delta t}\end{bmatrix}.$
(25)

###  5.4 Gradients of the TDKS Propagator

Here we consider gradients of the propagator $\boldsymbol{F}_{\Delta t}$
defined in (8). Note that $\boldsymbol{F}_{\Delta t}$ also satisfies

$\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi},\boldsymbol{v}^{C})=\boldsymbol{F}^{R}_{\Delta
t}(\boldsymbol{\phi}^{R},\boldsymbol{\phi}^{I},\boldsymbol{v}^{C})+i\boldsymbol{F}^{I}_{\Delta
t}(\boldsymbol{\phi}^{R},\boldsymbol{\phi}^{I},\boldsymbol{v}^{C}).$

#### Gradients of $\boldsymbol{F}$ when we seek pointwise values of
$\boldsymbol{v}^{C}$.

Because $V$ is diagonal, the $\ell$-th element of $\boldsymbol{F}$ is

$\left[\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi},\boldsymbol{v}^{C})\right]_{\ell}=\sum_{q,r}\mathcal{P}_{\mathcal{K};\ell,q}\exp(-iV_{qq}(\boldsymbol{\phi},\boldsymbol{v}^{C})\Delta
t)\mathcal{P}_{\mathcal{K};q,r}\phi_{r}.$

First let us compute the derivative of this $\ell$-th element with respect to
$\phi_{m}^{R}$. We obtain

$\frac{\partial\left[\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi},\boldsymbol{v}^{C})\right]_{\ell}}{\partial\phi_{m}^{R}}=\sum_{q,r}\mathcal{P}_{\mathcal{K};\ell,q}\exp(-iV_{qq}(\boldsymbol{\phi},\boldsymbol{v}^{C})\Delta
t)(-i\Delta t)\frac{\partial
V_{qq}}{\partial\phi_{m}^{R}}\mathcal{P}_{\mathcal{K};q,r}\phi_{r}\\\
+\sum_{q}\mathcal{P}_{\mathcal{K};\ell,q}\exp(-iV_{qq}(\boldsymbol{\phi},\boldsymbol{v}^{C})\Delta
t)\mathcal{P}_{\mathcal{K};q,m},$ (26)

with

$\frac{\partial
V_{qq}}{\partial\phi_{m}^{R}}=\frac{\partial}{\partial\phi_{m}^{R}}\sum_{s}W_{q,s}w_{s}\phi_{s}\phi_{s}^{\ast}=2W_{q,m}w_{m}\phi_{m}^{R}.$

The derivative with respect to $\phi_{m}^{I}$ is similar:

$\frac{\partial\left[\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi},\boldsymbol{v}^{C})\right]_{\ell}}{\partial\phi_{m}^{I}}=\sum_{q,r}\mathcal{P}_{\mathcal{K};\ell,q}\exp(-iV_{qq}(\boldsymbol{\phi},\boldsymbol{v}^{C})\Delta
t)(-i\Delta t)\frac{\partial
V_{qq}}{\partial\phi_{m}^{I}}\mathcal{P}_{\mathcal{K};q,r}\phi_{r}\\\
+\sum_{q}\mathcal{P}_{\mathcal{K};\ell,q}\exp(-iV_{qq}(\boldsymbol{\phi},\boldsymbol{v}^{C})\Delta
t)\mathcal{P}_{\mathcal{K};q,m}i,$ (27)

with

$\frac{\partial
V_{qq}}{\partial\phi_{m}^{I}}=\frac{\partial}{\partial\phi_{m}^{I}}\sum_{s}W_{q,s}w_{s}\phi_{s}\phi_{s}^{\ast}=2W_{q,m}w_{m}\phi_{m}^{I}.$

Taking the real and imaginary parts of (26-27), we obtain all necessary
elements of the block Jacobian (25). Next we compute the derivative of the
$\ell$-th element of $\boldsymbol{F}$ with respect to the $m$-th element of
$\boldsymbol{v}^{C}$:

$\frac{\partial\left[\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi},\boldsymbol{v}^{C})\right]_{\ell}}{\partial
v^{C}_{m}}=\sum_{r}\mathcal{P}_{\mathcal{K};\ell,m}\exp(-iV_{mm}(\boldsymbol{\phi},\boldsymbol{v}^{C})\Delta
t)(-i\Delta t)\mathcal{P}_{\mathcal{K};m,r}\phi_{r},$

which follows from

$\frac{\partial V_{qq}}{\partial v^{C}_{m}}=\delta_{qm}=\begin{cases}1&q=m\\\
0&q\neq m.\end{cases}$

#### Gradients of $\boldsymbol{F}$ when we model the functional dependence of
$\boldsymbol{v}^{C}$ on present and past states.

The $\ell$-th element of $\boldsymbol{F}$ is now

$\left[\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})\right]_{\ell}=\sum_{q,r}\mathcal{P}_{\mathcal{K};\ell,q}\exp(-iV_{qq}(\boldsymbol{\phi};\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})\Delta
t)\mathcal{P}_{\mathcal{K};q,r}\phi_{r}.$

Using the new expression for $V$, we derive

$\frac{\partial\left[\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})\right]_{\ell}}{\partial\phi_{m}^{R}}=\sum_{q,r}\mathcal{P}_{\mathcal{K};\ell,q}\exp(-iV_{qq}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})\Delta
t)(-i\Delta t)\frac{\partial
V_{qq}}{\partial\phi_{m}^{R}}\mathcal{P}_{\mathcal{K};q,r}\phi_{r}\\\
+\sum_{q}\mathcal{P}_{\mathcal{K};\ell,q}\exp(-iV_{qq}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})\Delta
t)\mathcal{P}_{\mathcal{K};q,m},$ (28)

with

$\frac{\partial
V_{qq}}{\partial\phi_{m}^{R}}=2W_{q,m}w_{m}\phi_{m}^{R}+\frac{\partial
v^{C}_{q}}{\partial\phi_{m}^{R}}.$

The derivative with respect to $\phi_{m}^{I}$ is similar:

$\frac{\partial\left[\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})\right]_{\ell}}{\partial\phi_{m}^{I}}=\sum_{q,r}\mathcal{P}_{\mathcal{K};\ell,q}\exp(-iV_{qq}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})\Delta
t)(-i\Delta t)\frac{\partial
V_{qq}}{\partial\phi_{m}^{I}}\mathcal{P}_{\mathcal{K};q,r}\phi_{r}\\\
+\sum_{q}\mathcal{P}_{\mathcal{K};\ell,q}\exp(-iV_{qq}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})\Delta
t)\mathcal{P}_{\mathcal{K};q,m}i,$ (29)

with

$\frac{\partial
V_{qq}}{\partial\phi_{m}^{I}}=2W_{q,m}w_{m}\phi_{m}^{I}+\frac{\partial
v^{C}_{q}}{\partial\phi_{m}^{I}}.$

The derivatives with respect to the past state $\boldsymbol{\phi}^{\prime}$
are

$\displaystyle\frac{\partial\left[\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})\right]_{\ell}}{\partial\phi_{m}^{{}^{\prime},R}}$
$\displaystyle=\sum_{q,r}\mathcal{P}_{\mathcal{K};\ell,q}\exp(-iV_{qq}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})\Delta
t)(-i\Delta t)\frac{\partial
v^{C}_{q}}{\partial\phi_{m}^{{}^{\prime},R}}\mathcal{P}_{\mathcal{K};q,r}\phi_{r}$
(30a) $\displaystyle\frac{\partial\left[\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})\right]_{\ell}}{\partial\phi_{m}^{{}^{\prime},I}}$
$\displaystyle=\sum_{q,r}\mathcal{P}_{\mathcal{K};\ell,q}\exp(-iV_{qq}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})\Delta
t)(-i\Delta t)\frac{\partial
v^{C}_{q}}{\partial\phi_{m}^{{}^{\prime},I}}\mathcal{P}_{\mathcal{K};q,r}\phi_{r}$
(30b)

Taking the real and imaginary parts of (28-29-30), we obtain all necessary
elements of both block Jacobians in (14c).

Finally, we need

$\frac{\partial\left[\boldsymbol{F}_{\Delta
t}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})\right]_{\ell}}{\partial\theta_{m}}=\sum_{q,r}\mathcal{P}_{\mathcal{K};\ell,q}\exp(-iV_{qq}(\boldsymbol{\phi},\boldsymbol{\phi}^{\prime};\boldsymbol{\theta})\Delta
t)(-i\Delta t)\frac{\partial
V_{qq}}{\partial\theta_{m}}\mathcal{P}_{\mathcal{K};q,r}\phi_{r}.$

###  5.5 Further Implementation Details

We implemented the adjoint method in JAX. Derivatives of the $v^{C}$ model are
computed via automatic differentiation. XLA compilation enables us to run the
code on GPUs. To optimize via L-BFGS-B, we use scipy.optimize. All source code
is available upon request.

###  5.6 Training Set Results

Here we consider training a model $v^{C}[\phi](x,t)$ that allows for arbitrary
dependence on the real and imaginary parts of both $\phi(x,t)$ and
$\phi(x,t-\Delta t)$. Here the input layer is of dimension $4(J+1)$—see
Section 3.

To train such a model, we apply the L-BFGS-B optimizer with objective function
given by the first line of (13) and gradients computed with the adjoint system
(14-15). We initialize neural network parameters $\boldsymbol{\theta}$ by
sampling a mean-zero normal distribution with standard deviation
$\sigma=0.01$. For training data, we subsample the $p=-1.5$ TDSE electron
density time series by a factor of $100$ in time, so that $\Delta t=2.4\times
10^{-3}$ fs and the entire training trajectory consists of $K=301$ time steps.
We retain this time step in all training sets and results that follow.

In Figure 5, we show the resulting model’s results on the training set. The
trained $v^{C}[\phi]$ functional, when used to solve the TDKS equation (5),
yields electron densities that agree closely with the reference TDSE electron
densities.

Figure 5: We use $300$ time steps (corresponding to $0.72$ fs) of the $p=-1.5$
data (black) together with the adjoint method to train a neural network model
$v^{C}[\phi]$ that depends on the current and previous $\phi$. We plot in red
the results of using the learned $v^{C}$ to propagate (5) from $t=0$ to
$t=0.72$ fs. Note the close agreement.

###  5.7 Relationship to Existing Literature on Optimal Control for TDKS
Systems

Here we contrast our work with prior work on optimal control for TDKS systems,
specifically work that involves the adjoint method.

First let us view our work through the lens of optimal control: we generate
reference data by first solving the two-dimensional TDSE. Our cost function is
then the mismatch between (i) electron densities computed from TDKS, and (ii)
electron densities computed from the time-dependent wave functions obtained
from TDSE, all on a discrete temporal grid. We view $v^{C}$ as a control that,
properly chosen, guides TDKS to produce the same electron densities that would
have been produced by solving TDSE.

A common feature of both present and prior work is the idea of incorporating
the TDKS equation as a time-dependent constraint—see Eq. (8) in Castro et al.
(2012), Eq. (43) in Castro and Gross (2013), and Eq. (3.4) in Sprengel et al.
(2018). Upon taking functional derivatives, this leads naturally to adjoint
systems, which have been analyzed in detail for TDKS systems (Borzì, 2012;
Sprengel et al., 2017). In particular, Sprengel et al. (2017) and Sprengel et
al. (2018) develop and analyze optimal control problems for multidimensional
TDKS systems. In all prior work we have seen, the correlation potential
$v^{C}$ is taken as _adiabatic_ with fixed functional form throughout the
solution of the optimal control problem.

In prior work, the control $u$ is distinct from $v^{C}$. In Castro et al.
(2012), the control $u$ governs the Fourier spectrum of the amplitudes of an
applied electric field. In Sprengel et al. (2018), the control $u$ influences
the system through potentials such as $V_{u}(x)=x^{2}$ and $V_{u}(x)=x\cdot
p$, modeling the control of a quantum dot.

In Castro et al. (2012), the objective is to balance (i) maximization of
charge transfer from one potential well to a neighboring potential well with
(ii) minimization of the intensity of the applied field. The resulting cost
function models both parts of this physical objective. In Sprengel et al.
(2018), the authors do include in their cost function the $L^{2}$ distance
between the electron density computed from TDKS and a reference electron
density, all in continuous time. They apply this to the problem of guiding
TDKS towards a target trajectory that itself was computed by solving TDKS.

Viewed in this context, the distinguishing features of the present work are as
follows: (i) treating $v^{C}$ itself as the control, (ii) allowing $v^{C}$ to
be _non-adiabatic_ in the sense that it depends on both present and past
electron densities, and (iii) applying this method to match electron density
trajectories computed from TDSE. By treating $v^{C}$ as the object of
interest, and by guiding TDKS trajectories to match TDSE trajectories, the
present work addresses the system identification problem of learning $v^{C}$
from data.

## References

  * Borzì (2012) Alfio Borzì.  Quantum optimal control using the adjoint method.  _Nanoscale Systems: Mathematical Modeling, Theory and Applications_ , 1:93–111, 2012.  URL http://eudml.org/doc/266625.
  * Bryson and Ho (1975) A. E. Bryson and Y.-C. Ho.  _Applied Optimal Control: Optimization, Estimation and Control_.  Halsted Press Book. Taylor & Francis, 1975.  Revised printing.
  * Castro et al. (2012) A. Castro, J. Werschnik, and E. K. U. Gross.  Controlling the dynamics of many-electron systems from first principles: A combination of optimal control and time-dependent density-functional theory.  _Phys. Rev. Lett._ , 109:153603, Oct 2012.  10.1103/PhysRevLett.109.153603.  URL https://link.aps.org/doi/10.1103/PhysRevLett.109.153603.
  * Castro and Gross (2013) Alberto Castro and E. K. U. Gross.  Optimal control theory for quantum-classical systems: Ehrenfest molecular dynamics based on time-dependent density-functional theory.  _Journal of Physics A: Mathematical and Theoretical_ , 47(2):025204, 2013.
  * Castro et al. (2004) Alberto Castro, Miguel A. L. Marques, and Angel Rubio.  Propagators for the time-dependent Kohn-Sham equations.  _The Journal of Chemical Physics_ , 121(8):3425–3433, 2004.
  * Elliott et al. (2012) Peter Elliott, Johanna I Fuks, Angel Rubio, and Neepa T Maitra.  Universal dynamical steps in the exact time-dependent exchange-correlation potential.  _Physical Review Letters_ , 109(26):266404, 2012\.
  * Feit et al. (1982) M. D. Feit, J. A. Fleck Jr, and A. Steiger.  Solution of the Schrödinger equation by a spectral method.  _Journal of Computational Physics_ , 47(3):412–433, 1982.
  * Fleck Jr et al. (1976) J. A. Fleck Jr, J. R. Morris, and M. D. Feit.  Time-dependent propagation of high energy laser beams through the atmosphere.  _Applied Physics_ , 10(2):129–160, 1976.
  * Hasdorff (1976) L. Hasdorff.  _Gradient Optimization and Nonlinear Control_.  Wiley, 1976.
  * Hastie et al. (2009) Trevor Hastie, Robert Tibshirani, Jerome H Friedman, and Jerome H Friedman.  _The Elements of Statistical Learning: Data Mining, Inference, and Prediction_.  Springer, second edition, 2009.
  * Kalita et al. (2021) Bhupalee Kalita, Li Li, Ryan J. McCarty, and Kieron Burke.  Learning to approximate density functionals.  _Accounts of Chemical Research_ , 54(4):818–826, 2021.  10.1021/acs.accounts.0c00742.  URL https://doi.org/10.1021/acs.accounts.0c00742.
  * Klambauer et al. (2017) Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter.  Self-Normalizing Neural Networks.  In _Proceedings of the 31st International Conference on Neural Information Processing Systems_ , NIPS’17, pages 972–981, 2017.
  * Maitra (2016) Neepa T. Maitra.  Perspective: Fundamental aspects of time-dependent density functional theory.  _The Journal of Chemical Physics_ , 144(22):220901, 2016.  10.1063/1.4953039.  URL https://doi.org/10.1063/1.4953039.
  * Maitra et al. (2002) Neepa T. Maitra, Kieron Burke, and Chris Woodward.  Memory in time-dependent density functional theory.  _Phys. Rev. Lett._ , 89:023002, Jun 2002.  10.1103/PhysRevLett.89.023002.  URL https://link.aps.org/doi/10.1103/PhysRevLett.89.023002.
  * Nagai et al. (2018) Ryo Nagai, Ryosuke Akashi, Shu Sasaki, and Shinji Tsuneyuki.  Neural-network Kohn-Sham exchange-correlation potential and its out-of-training transferability.  _The Journal of Chemical Physics_ , 148(24):241737, 2018.
  * Runge and Gross (1984) Erich Runge and E. K. U. Gross.  Density-functional theory for time-dependent systems.  _Phys. Rev. Lett._ , 52:997–1000, Mar 1984.  10.1103/PhysRevLett.52.997.  URL https://link.aps.org/doi/10.1103/PhysRevLett.52.997.
  * Sanz-Serna (2016) J. M. Sanz-Serna.  Symplectic Runge–Kutta schemes for adjoint equations, automatic differentiation, optimal control, and more.  _SIAM Review_ , 58(1):3–33, 2016.  10.1137/151002769.
  * Sprengel et al. (2017) Martin Sprengel, Gabriele Ciaramella, and Alfio Borzì.  A Theoretical Investigation of Time-Dependent Kohn–Sham Equations.  _SIAM Journal on Mathematical Analysis_ , 49(3):1681–1704, 2017.
  * Sprengel et al. (2018) Martin Sprengel, Gabriele Ciaramella, and Alfio Borzì.  Investigation of optimal control problems governed by a time-dependent Kohn-Sham model.  _Journal of Dynamical and Control Systems_ , 24(4):657–679, 2018.  https://arxiv.org/abs/1701.02679.
  * Suzuki et al. (2020) Yasumitsu Suzuki, Ryo Nagai, and Jun Haruyama.  Machine learning exchange-correlation potential in time-dependent density-functional theory.  _Physical Review A_ , 101(5):050501, 2020.
  * Ullrich (2011) Carsten A. Ullrich.  _Time-Dependent Density-Functional Theory: Concepts and Applications_.  Oxford Graduate Texts. Oxford University Press, Oxford, 2011.  10.1093/acprof:oso/9780199563029.001.0001.
  * van Leeuwen (1999) Robert van Leeuwen.  Mapping from densities to potentials in time-dependent density-functional theory.  _Phys. Rev. Lett._ , 82:3863–3866, May 1999.  10.1103/PhysRevLett.82.3863.  URL https://link.aps.org/doi/10.1103/PhysRevLett.82.3863.
  * Wagner et al. (2012) Lucas O Wagner, Zeng-hui Yang, and Kieron Burke.  Exact conditions and their relevance in TDDFT.  In _Fundamentals of Time-Dependent Density Functional Theory_ , pages 101–123. Springer, 2012.