\section{Introduction}
Robust and high-fidelity control is a central prerequisite for quantum information processing \cite{TrainingSC}. The problem has been addressed differently by adiabatic techniques and composite pulses. The former is based on the adiabatic theorem, which mathematically requires an infinitely slow dynamics or infinitely large coupling strengths.
In practical applications, necessitating fast dynamics and low coupling, adiabatic passage is applied with finite times and is only approximate.
Methods to accelerate it and to make it accurate while preserving its relative robustness have been proposed.
One can mention (i) parallel adiabatic passage (PAP) \cite{parallel1,parallel2,parallel3,parallel4}, which allows the design of the shape of the control parameters minimizing the non-adiabatic transfer; (ii) counterdiabatic driving \cite{CD1,CD2,CD3} (included into a more general class of techniques usually referred to as shortcut to adiabaticity (STA) \cite{STA}), which is based on the addition of a coupling term that compensates the non-adiabatic coupling, and (iii)  local optimization \cite{localadiabatic}, referred to as FAQUAD \cite{FAQUAD}, based on the control of the dynamics via a single parameter through the gap of an avoided crossing. Composite pulses, on the other hand, use a sequence of precise $\pi$-pulse transfers and an error resilience by (static) phases, see e.g. \cite{Vitanov}. More general time-dependent phases can be used, as in the single-shot shaped pulse (SSSP) method, \cite{SSSP,Hybrid}, or more generally STA based on reverse engineering techniques \cite{STA}, which offer faster dynamical routes featuring both robustness and high-fidelity.

Methods to determine the fastest (or least energetic) processes under given conditions can be determined by optimal control,
through the minimization of a cost functional, such as, e.g., Pontryagin maximum principle \cite{Pontryagin}, which provides a global (absolute) optimum for simple systems. Numerical optimization, such as, e.g., gradient-based optimization algorithm (GRAPE algorithm \cite{GRAPE}), provide only local optimal solutions.
Only recently, techniques merging robustness and optimality have been proposed \cite{Extended,Barnes}, including a method based on inverse optimization (RIO) \cite{Dridi,Xavier}. In particular, time-optimal robust processes have been shown to feature driving pulses of constant amplitude, i.e. square pulses. However such square wave forms are difficult to generate in practice and can produce inaccuracies at the discontinuities. It is thus desirable to develop techniques that can systematically produce smooth control pulses close to ideal optimal square pulses.

In this paper, we investigate the construction of robust quantum control using smooth quasi-square pulses in the two regimes, optimal and adiabatic: (i) robust optimization \cite{Dridi,Xavier} and (ii) rescaled adiabatic passage with a contracted time \cite{TR}.
In both regimes, we use a formulation allowing a design of the pulse approaching square pulses. We choose a hyper-Gaussian shape (of high order, i.e. featuring fast ramps) as a smooth quasi-square pulse approaching the ideal time-optimal square pulse.

We compare the two protocols in terms of dynamics and robustness.
It is shown these two regimes feature very different dynamical process and that, for an identical pulse area, RIO offers a more efficient robustness compared to a time-contracted adiabatic passage, where adiabaticity is merely approximate. By increasing the pulse area over 50$\%$, the latter can achieve performance comparable to the RIO  protocol.

The paper is organized as follows. Section 2 introduces the model and the methods offering fast and close to time-optimal robust processes, namely RIO and time-contracted adiabatic passage. Section 3 presents the results with the use of a hyper-Gaussian pulse as a quasi-square pulse. Other quasi-square pulses are investigated for time-contracted adiabatic passage made on an expansion combining a linear and sine terms of the rescaled function. Section 4 concludes.
\section{Model and methods}
We consider the model
\begin{equation}
	H_0=\frac{\hbar}{2}\left[
	\begin{array}{cc}
		-\Delta & \Omega\\
		\Omega & \Delta
	\end{array}
	\right]
	\label{H},
\end{equation}
corresponding to a two-level system driven by a quasi-resonant pulse in rotate-wave approximation,
where $\Omega$ and $\Delta=\omega_{2}-\omega_{1}-\omega_{0}-\dot{\phi}$ denote the (time-dependent) Rabi frequency and detuning, respectively. Here $\omega_{21}=\omega_{2}-\omega_{1}$ is frequency difference between the two levels and $\omega_{0}+\dot{\phi}$ is the time-dependent instantaneous frequency of the laser field (with $\omega_0$ its mean frequency). Under resonance condition $\omega_0=\omega_2-\omega_1$, the phase of the field reads: $\phi(t)=\int_{t_i}^{t}\Delta(s)ds$.

The solution of the time dependent Schr\"odinger equation (TDSE)
$i\hbar\frac{\partial}{\partial t}\vert\phi_0(t)\rangle=H_0\vert\phi_0(t)\rangle$ is conveniently
parameterized with three angles: the mixing angle
$\theta\equiv\theta(t)\in[0,\pi]$, the internal (or relative) phase $\varphi\equiv\varphi(t)\in[-\pi,\pi]$
and a global phase $\gamma\equiv\gamma(t)\in[0,2\pi]$ as
\begin{equation}
	\label{phizero}
	\vert\phi_0(t)\rangle=\left[\begin{matrix}
		e^{i\varphi/2}\cos(\theta/2)\\
		e^{-i\varphi/2}\sin(\theta/2)
	\end{matrix}\right] e^{-i\gamma/2}.
\end{equation}
Inserting it into the TDSE, we obtain
\begin{subequations}
	\label{dotpara}
\begin{align}
		\label{dottheta}
		\dot\theta&=\Omega\sin\varphi,\\
		\label{dotphi}
		\dot\varphi&=\Delta+\Omega\cos\varphi\cot\theta,\\
		\label{dotgamma}
		\dot\gamma&
		=\Omega \frac{\cos\varphi}{\sin\theta}=\dot\theta \frac{\cot\varphi}{\sin\theta},
\end{align}
\end{subequations}
where the dot represents the derivation with respect to time $t$.

One can inversely determine the controls from the dynamical angles:
\begin{subequations}
\begin{eqnarray}
        \label{delta}
		&\Delta=\dot{\varphi}-\dot{\gamma}\cos\theta,\\
		&\Omega=\sqrt{\dot{\theta}^{2}+\dot{\gamma}^{2}\sin^{2}\theta}=|\dot{\gamma}|\sqrt{\bigl(\dot{\widetilde{\theta}}\bigr)^{2}+\sin^{2}\tilde{\theta}},
		\label{omega}
\end{eqnarray}
\end{subequations}
with $\dot{\tilde{\theta}}\equiv d\tilde{\theta}/d\gamma$, where $\theta$ as a function of $\gamma$ is denoted as $\widetilde\theta(\gamma)$.

\subsection{Robust inverse optimization}
To take into account robustness, we consider the perturbed Hamiltonian of the form $H_{\lambda}=H_0+\lambda V$,
\begin{equation}
	\label{Halpdelt}
	H_{\alpha,\beta,\delta}=\frac{\hbar}{2}\left[\begin{array}{cc} -\Delta & \Omega\\ \Omega & \Delta\end{array}\right]
	+\frac{\hbar}{2}\left[\begin{array}{cc} -\delta & \alpha\Omega + \beta \\ \alpha\Omega + \beta & \delta\end{array}\right],
\end{equation}
where $\alpha$ is a coefficient modifying the Rabi field amplitude (pulse inhomogeneities), $\delta$ features inhomogeneous broadening or a slow stochastic noise in the energy levels of the qubit (i.e. considered in a quasi-static representation), and $\beta$ (considered real) a slow stochastic transverse noise.
We consider the family of the state solutions $\vert\phi_{0}(t_f)\rangle$
that reach exactly the target state $|\phi_T\rangle$ at the end of the process, at $t=t_f$, without perturbation, from a given initial state at $t=t_i$. Among these solutions, we look for those that reach it in a robust way as defined below.

The SSSP method \cite{SSSP} consists in expanding the projection of the given state solution $\vert\phi_{\lambda}(t_f)\rangle$ on the given target state $|\phi_T\rangle$ at the end of the process with respect to $\lambda$, giving the probability amplitude:
\begin{equation}
	\label{expansion}
\langle\phi_T|\phi_{\lambda}(t_f)\rangle=1+O_1+O_2+O_3+\cdots,
\end{equation}
where $O_n$ denotes the term of total order $n$: $O_n\equiv O(\lambda^n)$
\begin{align}
	\label{Errors}
		O_1&=-i\int_{t_i}^{t_f}\langle\phi_0(t)|V(t)|\phi_0(t)\rangle dt\equiv
		-i\int_{t_i}^{t_f}e(t)dt,\\
		O_2&=(-i)^2\int_{t_i}^{t_f}dt\int_{t_i}^{t}dt'[e(t)e(t')
		+f(t)\bar f(t')],
		\label{ErrorsO2}

\end{align}
with
\begin{align}
		e

		&=-\frac{1}{2}(\delta\cos\theta-\alpha\dot\gamma\sin^2\theta - \beta\sin\theta\cos\varphi),\\

		f

		&=\frac{1}{2}\Bigl[\delta\sin\theta
		+\alpha\Bigl(\frac{1}{2}\dot\gamma\sin2\theta-i\dot\theta\Bigr)
		\nonumber\\
		&\qquad\qquad{}+\beta(\cos\varphi\cos\theta-i\sin\varphi)\Bigr]e^{i\gamma}.

\end{align}
We will consider in this paper robustness with respect to pulse inhomogeneities ($\alpha$-robustness), i.e. $\delta=\beta=0$, $\lambda\equiv\alpha$, for population transfer at second order: $P_T(t_f)=1-(O_2+\bar O_2+\bar O_1O_1)=1-\vert \int_{t_i}^{t_f} f(t) dt\vert ^2$.
Robustness up to second order corresponds thus to the two constraints
\begin{eqnarray}
	\label{constraints}
	\int_{t_i}^{t_f} {\cal R}[f(t)] dt=0,\qquad \int_{t_i}^{t_f} {\cal I}[f(t)] dt=0.
\end{eqnarray}
Optimal trajectory of the dynamical angles with respect to a given cost, \textit{e.g.} the pulse area $\int_{t_i}^{t_f} \vert \Omega (t)\vert dt $, the pulse energy $\int_{t_i}^{t_f} \vert \Omega (t)\vert^2 dt $, or the pulse duration (time-optimization), that satisfy the constraints \eqref{constraints} can be obtained via an Euler-Lagrange formulation, and the controls are determined from this trajectory via \eqref{delta} and \eqref{omega}. This protocol has been named robust inverse optimization (RIO).

We can analyze the problem of complete or partial population transfer from the ground state ($\theta(t_i)\equiv \theta_i=0$, $\gamma(t_i)\equiv \gamma_i=\varphi(t_i)\equiv \varphi_i=\pi/2$),
which is optimal with respect to the pulse area, by fixing the final targeted mixing angle $\theta(t_f)\equiv \theta_f=\theta_0$ and internal phase $\varphi (t_f)\equiv\varphi_f=\varphi_0$ (the final value of $\gamma$ being irrelevant for population transfer).
In this paper, we consider the complete population transfer problem, i.e. $\theta_f=\pi$.
It corresponds to the trajectory of $\theta$ as a function of $\gamma$, denoted as $\widetilde\theta(\gamma)$,
displayed in Fig. \ref{figTraject}, as determined in \cite{Dridi}.
Time-optimal robust trajectory (for a given peak amplitude $\Omega_0$) corresponds to a special time-parametrization of the dynamical angle $\gamma(t)$ given by the constant of motion \begin{eqnarray}
	\dot\theta^2+\dot\gamma^2\sin^2\theta=\Omega_0^2,
\end{eqnarray}
with a constant pulse $\Omega=\Omega_0$, i.e. a square pulse, of duration $T_{\min}=t_f-t_i$:
\begin{eqnarray}
	\label{Tmin}
	T_{\min}=\frac{1}{\Omega_0}\int_{\gamma_i}^{\gamma_f}d\gamma\,\sqrt{\bigl(\dot{\widetilde\theta}\bigr)^2+\sin^2\widetilde\theta}\approx 5.84/\Omega_0.
\end{eqnarray}

\begin{figure}[t]
	\begin{center}
		\includegraphics[scale=0.75]{figpi.eps}
		\caption{Optimal $\alpha$-robust geodesic $\widetilde\theta(\gamma)$ in the dynamical variable space $(\gamma,\theta)$ leading to complete population transfer ($\theta_f=\pi$) from the ground state ($\theta_i=0$) (leading to $\gamma_f=5\pi/3$ and $\varphi_f=\pi/2$).}
		\label{figTraject}
	\end{center}
\end{figure}

The trajectory in Fig. \ref{figTraject} offers thus infinitely many optimal solutions  (robust with respect to the pulse area) depending on the time parametrization of $\gamma(t)$. In order to approach the ideal time-optimal robust dynamics, we impose a parametrization given by a smooth quasi-square pulse of same peak amplitude and same area as that of the square pulse. This is modeled by a hyper-Gaussian pulse of high-order even $n$ and width $\sigma$:
\begin{eqnarray}
	\Omega(t)=\Omega_{0}{\rm{exp}}(-(t/\sigma)^{n}),
\end{eqnarray}
such that $\int \Omega(t) dt=2\Omega_0\sigma \Gamma(1/n)/n=\Omega_0 T_{\min}$, i.e.
\begin{eqnarray}
	\sigma=\frac{n T_{\min}}{2\Gamma(1/n)}.
\end{eqnarray}
One could choose in principle any even parameter $n$; we expect a longer process for a smaller $n$ deviating more from the ideal (optimized) squared pulse.

The parametrization of $\gamma(t)$ is more precisely determined from the identity of the partial pulse areas:
\begin{eqnarray}
	\Omega_0\int_{t_{i}}^{t}{\rm {exp}} (-(t/\sigma)^{n})dt =\int_{\gamma_i}^{\gamma(t)}d\gamma\,\sqrt{\bigl(\dot{\widetilde\theta}\bigr)^2+\sin^2\widetilde\theta}.
\end{eqnarray}

\subsection{Time-contracted adiabatic passage}
Adiabatic passage means a dynamics which follows approximately the instantaneous eigenvector continuously connected to the initial condition (when $\Omega=0$).
The eigenvectors read
\begin{eqnarray}
		&\vert \psi_{+}\rangle=\cos(\Theta/2)|1\rangle + \sin(\Theta/2)|2\rangle,\\
		&\vert \psi_{-}\rangle= \cos(\Theta/2)|2\rangle-\sin(\Theta/2)|1\rangle,

\end{eqnarray}
with the mixing angle $\Theta=\arctan(-\Omega/\Delta)$, $0\le\Theta\le\pi$, and the corresponding eigenenergy $\lambda_{\pm}=\pm\hbar\Omega_{a}/2$ with $\Omega_{a}=\sqrt{\Delta^{2}+\Omega^{2}}$ .

To achieve an efficient and robust population transfer by adiabatic passage, we consider the parallel adiabatic passage technique \cite{parallel1}, where we define an analytic shape $\Lambda(t)$ of the Rabi frequency: $\Omega(t)=\widetilde\Omega_0\Lambda(t)$, from which we design the detuning $\Delta(t)$ such that the distance between the eigenvalues remain constant at all times:
\begin{align}
	\label{pulses}
		\Omega(t)&=\widetilde\Omega_{0}\Lambda(t/T),\\
		\Delta(t)&=\widetilde\Omega_{0} \,{\rm {sign}}(t/T) \sqrt{1-\Lambda^2(t/T)},
\end{align}
giving $\Omega_a=\widetilde\Omega_0$.
The parameter $T$ characterizes the width of the pulse.

Renormalizing the time $s=t/T$ in the Schr\"odinger equation
shows that the dynamics in the adiabatic basis $\vert\psi_a\rangle=T_a^{-1}\vert\psi\rangle$ with $T_a=[\vert\psi_-\rangle,\vert\psi_+\rangle]$, $D_a=T_a^{\dagger}HT_a={\rm{diag}}(\lambda_-,\lambda_+)$, given by
\begin{eqnarray}
	&i\hbar\frac{\partial}{\partial s}\vert\psi_a\rangle=T\widetilde\Omega_0\left[\begin{array}{cc}
		-1 & 0\\
		0 & 1 \end{array}
	\right]\vert\psi_a\rangle -i\hbar T_a^{\dagger} \dot T_a \vert\psi_a\rangle,
\end{eqnarray}
is better approximated by the diagonal Hamiltonian $D_a$ when the dimensionless quantity characterizing the pulse area is large $T\widetilde\Omega_0\gg 1$. It also shows that accelerating the dynamics with a shorter time $T'<T$ gives trivially the same (adiabatic) dynamics but for a larger peak Rabi frequency $\widetilde\Omega'_0=\widetilde\Omega_0 T/T'$ (with the same pulse area).

The method proposed in \cite{TR} shows a method to define, from a given adiabatic dynamics, a rescaled adiabatic passage featuring the same transfer and the same robustness as the original one. We will apply it with a time contraction and will refer it to as time-contracted adiabatic passage (TCAP). It consists in defining a rescaled time $\tau$, defined by the rescaled function $g(\tau)$ such that $t=g(\tau)$, in which the dynamics realizes exactly the same final state but in a shorter time. This is achieved when $\tau_f-\tau_i < t_f-t_i$,
and $\tau_i=g^{-1}(t_i)$ and $\tau_f=g^{-1}(t_f)$ the initial and final times in the new frame, respectively, such that the state functions are (i) the same ones at initial time, (ii) the same ones at final times, and that the Hamiltonians are (iii) the same ones initially and (iv) the same ones finally. More explicitly, the rescaled dynamics corresponds to the Schr\"odinger equation
\begin{eqnarray}
	i \hbar \frac{\partial}{\partial \tau} \psi_r(\tau)= H_r(\tau) \psi_r(\tau), \ H_r(\tau)=g'(\tau) H(g(\tau)),
\end{eqnarray}
with the rescaled state and Hamiltonian, respectively,
\begin{eqnarray}
	\psi_r(\tau)=\psi(g(\tau)),\quad H_r(\tau)=g'(\tau) H(g(\tau)).
\end{eqnarray}
The above conditions are satisfied when
\begin{eqnarray}
	\label{def_TR}
	g_i\equiv g(\tau_i)=t_i,\ g(\tau_f)=t_f,\ g'(\tau_i)= g'(\tau_f)=1.
\end{eqnarray}
The corresponding rescaled controls can be then expressed as
\begin{align}
		\label{pulse_TR}
		\Omega_c(\tau)&=\widetilde\Omega_{0}g^{\prime}(\tau)\Lambda(g(\tau)/T),\\
		\label{det_TR}
		\Delta_c(\tau)&=\widetilde\Omega_{0}g^{\prime}(\tau)  \,{\rm{sign}}(g(\tau)/T) \sqrt{1-\Lambda^2(g(\tau)/T)}.
\end{align}
We can notice that the pulse area is invariant: $\int \Omega_r(\tau) d\tau =\widetilde\Omega_{0} \int \Lambda(g/T) dg  =\int \Omega(t)dt$.

We consider the time contraction defined as
\begin{eqnarray}
	\label{contraction}
	\tau_i=t_i/a, \quad \tau_f=t_f/a.
\end{eqnarray}
which contracts the duration of the process by a factor $a>1$.
Our goal is to design, for a given time-contraction $a$, the rescaled function $g(t)$ such that the pulse features a quasi-square shape.

We search for quasi-square pulse following two strategies: (i) we start from an expansion of $g$ with free coefficients satisfying \eqref{def_TR} and \eqref{contraction},  and determine the optimal coefficients that minimize the pulse amplitude \eqref{pulse_TR}, (ii) we impose the rescaled pulse \eqref{pulse_TR} to be an hyper Gaussian pulse (of high order). In both cases, we infer the corresponding detuning \eqref{det_TR}.

We remark that we focus in this paper more specifically on the parallel adiabatic passage with a Gaussian pulse $\Lambda(t)=e^{-t^2}$, but that any other adiabatic passage design can be used.

\section{Robust control with quasi-square pulse}
\subsection{Robust inverse optimization with hyper-Gaussian pulse}

Fig. \ref{fig2} presents the dynamics driven by RIO with a hyper-Gaussian pulse (of high order $n=14$ and $\sigma\approx 1.095T$,
corresponding to the same pulse as in the TCAP, see  Fig. \ref{figAdiabPulses}, i.e. $\Omega_0\approx 2.77/T$) replacing the ideal optimal square pulse (of same area 5.84). It will be referred in short to as hyper-Gaussian RIO (hG-RIO).  As expected the corresponding detunings are very similar.
The detuning associated to the hyper-Gaussian pulse also appears smooth (i.e. regularized at the beginning and at the end of the process).

\begin{figure}
	\begin{center}
		\includegraphics[scale=0.75]{dynamicsop.eps}
		\caption{(a) The hG-RIO population dynamics (ground state: red line, excited state: blue line). (b) The quasi-square hyper-Gaussian pulse (solid lines) and  the time-optimal ideal square pulse (dashed line) with their corresponding oscillating detunings (almost undistinguishable at the scale of the figure, except at the beginning and at the end).

			}
		\label{fig2}
	\end{center}
\end{figure}

\subsection{Time-contracted adiabatic passage with linear-sine expansion}
We analyze TCAP using a $g$ function combing a linear term and a sine expansion with $N$ terms:
\begin{equation}
	g(\tau)=
	a\tau+\sum_{n=1}^N C_{n}\sin\left[\frac{2n\pi a}{t_{f}-t_{i}}\left(\tau-\frac{t_{i}}{a}\right)\right].
	\label{f0}
\end{equation}
This ensures in a simple way that \eqref{def_TR} with \eqref{contraction} are satisfied, while offering free parameters to design the shape of the pulse \eqref{pulse_TR}. More specifically, $g(t_i/a)=t_i$ and  $g(t_f/a)=t_f$ are satisfied by the linear term $a\tau$.
The derivatives give
\begin{equation}
	g'(t_i/a)=g'(t_f/a)=
	a+ \frac{2\pi a}{t_{f}-t_{i}} \sum_{n=1}^N n C_{n},
\end{equation}
which imposes a condition on the sum
\begin{equation}
	\label{constrCn}
	\sum_{n=1}^N nC_{n}=\frac{(1-a)  (t_{f}-t_{i})}{2\pi a}.
\end{equation}
This means that, for a given $N$, we can choose freely $N-1$ coefficients and the last one is constrained by \eqref{constrCn}.
We follow the protocol: starting with a Gaussian pulse, we determine numerically for a given $N$ the coefficients $C_{n<N}$ with \eqref{constrCn}, leading to the smallest pulse peak amplitude.

We consider the expansion \eqref{f0} of $g$ for various $N$.
Figure \ref{figAdiabPulses} shows the resulting shape of the pulses for $a=3$. We can notice as expected that, for larger $N$, the pulse tends to a quasi-square pulse with a decreasing peak amplitude.

\begin{figure}
	\begin{center}
		\includegraphics[scale=0.73]{APCshapes.eps}
		\caption{Time-contracted pulse shapes for $N=2,3,6$ where the coefficients $C_n$'s have been optimized to minimize the peak of the pulse and the hyper-Gaussian pulse for $n=14$ (labelled as hG) from the original Gaussian pulse (dashed line) that has been time-contracted.}
		\label{figAdiabPulses}
	\end{center}
\end{figure}

\subsection{Time-contracted adiabatic passage with hyper-Gaussian pulse}
In this section, we impose the time-contracted pulse to be a hyper-Gaussian pulse of high-order (even) $n$:
\begin{eqnarray}
	\label{pulse_TR_hG}
	\Omega_c(\tau)=\widetilde\Omega_{0}g^{\prime}(\tau)\Lambda(g(\tau)/T)

	\equiv \widehat\Omega_0\exp\bigl(-(\tau/\sigma)^n\bigr),
\end{eqnarray}
starting from a Gaussian pulse: $\Lambda(t)=e^{-t^2}$. It is referred to as hG-TCAP.
We derive the corresponding $g(t)$ function by integration.

We obtain (using $\tau_i=-\infty$ and the condition $g_i=t_i=-\infty$) for $\tau\ge0$
\begin{eqnarray}
	{\rm{erf}}\biggl(\frac{g}{T}\biggr)= \frac{2\sigma \widehat\Omega_{0}}{\sqrt{\pi}T\widetilde\Omega_{0}n}\biggl[ 2\Gamma\biggl(\frac{1}{n}\biggr)- \Gamma\biggl(\frac{1}{n},\frac{\tau^n}{\sigma^n}\biggr)\biggr]-1,
\end{eqnarray}

and for $\tau\le0$
\begin{eqnarray}
	{\rm{erf}}\biggl(\frac{g}{T}\biggr)=\frac{2\sigma \widehat\Omega_{0}}{\sqrt{\pi}T\widetilde\Omega_{0}n}\Gamma\biggl(\frac{1}{n},\frac{\tau^n}{\sigma^n}\biggr)-1.
\end{eqnarray}
Using an antisymmetric $g$, we have $g(0)=0$, i.e.
\begin{eqnarray}
	\sigma \widehat\Omega_{0}=\frac{\sqrt{\pi}T\widetilde\Omega_{0}n}{2 \Gamma\bigl(\frac{1}{n}\bigr)},
\end{eqnarray}
leading to
\begin{align}
	{\rm{erf}}\biggl(\frac{g}{T}\biggr)&= \frac{1}{\Gamma\bigl(\frac{1}{n}\bigr)}\biggl[ 2\Gamma\biggl(\frac{1}{n}\biggr)- \Gamma\biggl(\frac{1}{n},\frac{\tau^n}{\sigma^n}\biggr)\biggr]-1,\quad \tau\ge0,\nonumber\\
	{\rm{erf}}\biggl(\frac{g}{T}\biggr)&=\frac{1}{\Gamma\bigl(\frac{1}{n}\bigr)}\Gamma\biggl(\frac{1}{n},\frac{\tau^n}{\sigma^n}\biggr)-1,\quad \tau\le0.
\end{align}
From the derivative evaluated at $\tau_f=t_f/a$ for which $g(\tau_f)=t_f$ and $g'(\tau_f)=1$, we obtain
\begin{eqnarray}
	1=\frac{\sqrt{\pi}Tn}{2\sigma \Gamma\bigl(\frac{1}{n}\bigr)} \frac{e^{-(t_f/(a\sigma))^n}}{e^{-(t_f/T)^2}} ,
\end{eqnarray}
i.e.
\begin{eqnarray}
	\sigma = \frac{t_f}{a[\ln(\widehat\Omega_{0}/\widetilde\Omega_{0})+(t_f/T)^2 ]^{1/n}}\to  \frac{t_f}{a} (T/t_f)^{2/n}.
\end{eqnarray}
This shows that the width $\sigma$ tends to infinity for an infinite $t_f$, with an amplitude $\widehat\Omega_{0}$ going to 0.
In practice, we have to fix $t_f$ (and symmetrically $t_i$) to a finite value such that $\Lambda(t_f)=\epsilon$ is negligibly small.

\begin{figure}
	\begin{center}
		\includegraphics[scale=0.75]{dynamicsfap.eps}
		\caption{(a) The hG-TCAP population dynamics (ground state: red line, excited state: blue line) induced by the hyper-Gaussian pulse (b) (taken from Fig. \ref{figAdiabPulses}) and the corresponding detuning (c).}
		\label{fig4}
	\end{center}
\end{figure}

Figure \ref{figAdiabPulses} shows the resulting hyper-Gaussian pulse shape for $n=14$ generated from $a=3$ and $t_f=4T$, i.e. $\widehat\Omega_{0}\approx 0.84\widetilde\Omega_0=\Omega_0=2.77/T$ and $\sigma\approx 1.095T$.
Figure \ref{fig4} shows the hG-TCAP dynamics leading to the population transfer resulting from the hyper-Gaussian pulse and the detuning [determined from \eqref{det_TR}]. The latter features early and late high-amplitude bumps inducing far-from-resonance dynamics, which allows the fast ramps of the hyper-Gaussian pulse to satisfy adiabaticity. We highlight that these bumps result directly from the time-contraction of the adiabatic process. Between the two bumps, the detuning shows a slow quasi-linear chirp crossing the resonance during the quasi-plateau of the hyper-Gaussian pulse, which ensures adiabatic transfer.

\begin{figure}
	\begin{center}
		\includegraphics[scale=0.75]{Fidelitysum.eps}

		\caption{(a) Fidelity (dimensionless) of the transfer (in a base 10 logarithmic scale) as a function of the relative deviation $\alpha$ in the pulse amplitude $\Omega_0(1+\alpha)$ for hG-RIO (solid red line) and hG-TCAP (dashed red line) with the same hyper-Gaussian pulse shape and area (5.84). (b) Fidelity of the hG-TCAP transfer (in a base 10 logarithmic scale) with a pulse area more than 50\% larger (9.22) that leads to a fidelity comparable to that of hG-RIO (upper frame).}
		\label{fig5}
	\end{center}
\end{figure}
\subsection{Robustness}

Figure \ref{fig5}, upper frame, compares the robustness profile against a relative deviation of the pulse amplitude for hG-RIO and hG-TCAP, respectively, using the same pulse, via the fidelity of the transfer ${\cal F}=1-P$, where $P$ is the final population in the upper state. It shows that, as expected, hG-RIO achieves much better performance. The relative lack of robustness of the adiabatic process can be explained by the relatively low value of the pulse area $5.84$ (less than $2\pi$) imposed by the RIO protocol.

We recover a comparable robustness (on the left part of the profile) using hG-TCAP of pulse area more than 50\% larger (with the same hyper-Gaussian shape), as shown in Fig. \ref{fig5}, lower frame.

\section{Conclusion and Discussion}

We have analyzed fast robust control, very close to the time-optimal one, using hyper-Gaussian pulses as smooth quasi-square pulses replacing the ideal square pulse. Complete population transfer has been targeted. We have used two protocols: (i) the RIO protocol implementing a smooth pulse, and (ii) TCAP. In TCAP, we have first used an expansion combining a linear term with sine components with free coefficients. We have numerically shown that the time-contraction with an optimization of the free coefficients minimizing its peak induces the pulse to approach a square pulse. We have alternatively designed the contraction to precisely shape a hyper-Gaussian pulse.
We have compared the two protocols using the same hyper-Gaussian control pulse, namely hG-RIO and hG-TCAP. The dynamics in the two cases is very different: in hG-RIO, the dynamics takes much time and oscillate near the half superposition (see Fig. \ref{fig2}, upper frame), which is associated to the multi-valued function $\gamma$ in function of $\theta$ around $\pi/2$ in Fig. \ref{figTraject}.
In hG-TCAP, the dynamics monotonically reaches the transfer with quasi-linear evolution (except near the beginning and the end, see Fig. \ref{fig4}, upper frame).

We have numerically shown that, as expected, the robustness profile is much better for hG-RIO since this method is precisely constructed to produce an optimal robust dynamics and that the resulting pulse area is relatively too small to achieve an efficient adiabatic passage. However, hG-TCAP is of interest as it allows (i) a fast adiabatic passage (for a given pulse area) and (ii) a profile of robustness comparable to that of RIO using a 50\% larger pulse area.

The main drawback of adiabatic passage in such a two-level system is the lack of control of the final phase of the state, since it is given by the (dynamical) time-integral of the eigenvalue associated to the eigenstate involved in the passage, which is sensitive to the pulse amplitude. This is an important issue to produce quantum gate.
On the other hand, the RIO technique can be easily adapted to feature robustness including the phase of the state \cite{Dridi,Xavier}.

We highlight that, in addition of being optimal, the RIO protocol features a very simple (oscillatory) shape for the detuning. However, this process, as constructed, is not robust with respect to the detuning; it should be thus preferred when the frequencies are well controlled. Otherwise, the technique has to be extended to be resistant to inhomogeneous broadening [via the parameter $\delta$ in Eq. \eqref{Halpdelt}] or adiabatic passage can be considered if the internal phase is unimportant.

The model presents several practical limitations that are discussed below.

The use of short pulses might involve more than two levels. This would be typically the case in a superconducting qubits system that features a multi-level ladder of relatively weak anharmonicity \cite{superconductor}. In this situation, one can take the additional excited states into account at the lowest order, using perturbation theory. This would add dynamical Stark shift which can be easily treated and compensated by an appropriate detuning design, as a generalization of Ref. \cite{prl103}. Short pulses might also involve a coupling term beyond the RWA. It can be treated via Floquet theory, which, at the lowest order, will also give rise to a dynamical Stark shift, see, e.g., in \cite{parallel1}.

If the system is subjected to noise, such as decoherence due to environmental effects (e.g. dephasing noise or spontaneous emission), one can analyze the following typical situations. If the noise affects permanently the upper state (spontaneous emission), then one has to drive the system in a time much shorter than the typical decoherence time to avoid losses. A typical situation is when the loss drives the dynamics outside the two-state system, it is effectively modeled by a negative imaginary term on the diagonal of the Hamiltonian. This corresponds to resonance coalescence (exceptional points) with their specific adiabatic conditions \cite{EP}. The determination of a robust process could be determined for a given total loss. This has been shown in a different situation with a $\Lambda$-system, featuring a lossy upper state \cite{Xavier}, where robust pulse shapes have been derived for a given total loss. A slow stochastic noise in the energy levels of the qubit (corresponding to a systematic frequency errors  in a quasi-static representation) or a quasi-static transverse noise is represented by the static parameter $\delta$ and $\beta$, respectively, in the model \eqref{Halpdelt}. This will lead to different types of robustness integrals which have to be considered as new constraints in the Euler-Lagrange formulation.
A fast noise can be modeled by a Lindblad equation (in the Markov approximation) and to an effective Schr\"odinger equation in the case of pure dephasing \cite{pra75dephasing}. The Euler-Lagrange formulation can be applied to such linear (lossy) systems. It has been shown in \cite{whitenoise} that the flat $\pi$-pulse is the least sensitive protocol to phase noise since it is the fastest process. Allowing a given loss would allow here again the derivation of robust protocols in the same way as in \cite{Xavier}.

The implementation of hyper-Gaussian smooth quasi square pulses depends on the timescale considered. In the nano- and sub nano-second regimes, such waveforms can be produced directly in the time-domain by electronics. Below this regime, typically in the femto- and pico-second timescale, the principle of ultrashort pulse shaping is based on the spectral and amplitude modulation of the pulse spectral components \cite{Weiner}, including the liquid crystal programmable spatial light modulator, the acoustic-optic modulator programmable spatial light modulator and the acousto-optic programmable dispersive filter. Such techniques can been used to produce quasi square pulses \cite{Cicaldi}.

\section*{Acknowledgement}
This work was supported by the EUR-EIPHI Graduate School (17-EURE-0002). We also acknowledge support from the European Union's Horizon 2020 research and innovation program under the Marie Sklodowska-Curie grant agreement No. 765075 (LIMQUET). \XC{XC acknowledges the Ramon y Cajal program (Grant No. RYC-2017-22482).
JZ acknowledges additional support from the
CSC (China Scholarship Council).\\
This article is dedicated to the memory of Bruce W. Shore, who has been a great inspiration for SG for so many years, in particular concerning adiabatic and robust quantum control addressed in the present paper.}


\section{Introduction}\label{sec:intro}

Causal effects of spatially-varying exposures on spatially-varying outcomes may be subject to \emph{non-local confounding} (NLC), which occurs when the treatments and outcomes for a given unit are affected by \emph{covariates} of other nearby units \citep{cohen2008obesity, florax1992specification, chaix2010neighborhood, elhorst2010applied}. In simple cases, NLC can be resolved using simple summaries of non-local data, such as the averages of the covariates over pre-specified neighborhoods. But in many realistic settings, NLC is caused by the complex interaction of spatial factors, and thus it cannot be resolved using simple \emph{ad hoc} summaries of neighboring covariates. For such scenarios, we propose {\em weather2vec}\xspace, a framework that uses a U-net \citep{ronneberger2015u} to learn representations that encode NLC information and can be used in conjunction with standard causal inference tools. The method is broadly applicable to settings where the covariates are available over a grid of spatial units, and where the outcome and treatment are observed in some subset of the grid.

The name {\em weather2vec}\xspace stems from its motivation to address limitations in current methods for estimating causal effects in environmental studies where meteorological processes are known confounders, aiming to contribute to the development of new flexible machine learning tools to assess the effect of policies and climate-related events on health-relevant outcomes: a task which has been recently identified by \citet{rolnick2022tackling} as a pressing outstanding challenge for tackling the effects of climate change.

Two applications will be discussed in detail. The first application follows an earlier analysis by \citet{papadogeorgou2019adjusting}, who estimated the air quality impact of power plant emissions controls. This case study evaluates the method's ability to reduce NLC under sparsely observed treatments (in combination with with propensity matching methods \citep{rubin2005causal}).
The second example is an application to the problem of meteorological detrending \cite{wells2021improved}, and uses {\em weather2vec}\xspace to deconvolve climate variability from policy changes when characterizing long-term air quality trends.
These two examples are accompanied by a simulation study comparing alternative adjustments to account for NLC.

In summary, this article has three aims:
\begin{enumerate}[itemsep=0pt, parsep=0pt]
  \item Provide a rigorous characterization of NLC using the potential outcomes framework \citep{rubin2005causal}, clarifying some connections with interference and related methods \citep{tchetgen2012causal, forastiere2021identification, sobel2006randomized}.
  \item Expand the library of NN methods in causal inference by proposing a U-net \citep{ronneberger2015u} as a viable model to account for NLC in conjunction with standard causal inference tools.
   \item Establish a promising research direction for addressing NLC in scientific studies of air pollution exposure -- in which NLC is a common problem (driven by meteorology) for which widely applicable tools are lacking.
\end{enumerate}

We investigate two mechanisms to obtain the representations: one supervised, and one self-supervised.  The supervised one formally links the representation of NLC to the balancing property of propensity (and prognostic) scores in the causal inference literature \citep{rubin_for_2008,hansen2008prognostic}. This approach requires that the outcome and treatment are densely available throughout the covariates' grid. By contrast, the self-supervised approach first learns representations encoding neighboring covariate information into a low-dimensional vector, which can subsequently be included as confounders in downstream causal analyses when the outcomes and treatments are sparsely observed on the grid.

\paragraph{Related work}
Previous research has investigated NNs for the (non-spatial) estimation of balancing scores \citep{keller2015neural, westreich2010propensity, setoguchi2008evaluating} and counterfactual estimation  \citep{shalit2017estimating, johansson2016learning, shi2019adapting}.
None of these works, however, specifically consider NLC.

Relevant applications of U-nets in environmental studies include forecasting \citep{larraondo2019data, sadeghi2020improving}, estimating spatial data distributions from satellite images \citep{hanna2021multitask,fan2021resolving}, indicating that U-nets are powerful tools to manipulate rasterized weather data. Also relevant, \citet{lu2005meteorologically} give a specific application of NNs for meteorological detrending, although without considering adjusting for neighboring covariates.

Approaches to learn summaries of neighboring covariates for regression-based causal inference have been investigated in the econometrics literature. For example, WX-regression models \citep{elhorst2010applied} formulate the outcome as a linear function of the treatment and the covariates of some pre-specified neighborhood.
Similarly, CRAE \citep{blier2020encoding} uses an autoencoder to encode pre-extracted patches of regional census data into a lower-dimensional vector that is fed into an econometric regression, somewhat analogous to the dimensionality reduction step in the self-supervised formulation of {\em weather2vec}\xspace, but requiring to pre-process the data into patches.
In contrast to these regression-based approaches, {\em weather2vec}\xspace aims at learning balancing scores, which have known benefits that include the ability to empirically assess the threat of residual confounding and the offer of protection against model misspecification that arises when modeling outcomes directly \citep{rubin_for_2008}.

There is also a maturing literature on adjusting for unobserved spatially-varying confounding \citep{reich2021review, veitch2019using, papadogeorgou2019adjusting}.
Spatial random effect methods are popular in practice, although \citet{khan2020restricted} have highlighted their sensitivity to misspecification for the purposes of confounding adjustment. As an alternative, the distance adjusted propensity score matching (DAPSm) \citep{papadogeorgou2019adjusting} matches units based jointly on estimated propensity scores and spatial proximity under the rationale that spatial proximity can serve as a proxy for similarity in spatially-varying covariates. We use this model as a baseline in one of our applications.  \citet{veitch2019using} take a related approach in a context where network proximity is viewed analogously to spatial proximity, and show that, under certain regularity conditions, network proximity can be used as a proxy for a network-level unobserved confounder. They propose a mechanism to learn embeddings that capture confounding information and used them together with augmented inverse probability weighting to obtain unbiased causal estimates. Importantly, they only consider the ``pure homophily'' case \citep{shalizi2011homophily}, where the entirety of the confounding is assumed to be encoded by relative position in the network. While some of these methods could be useful for NLC, they all primarily target settings where confounding is local.

Finally, NLC is distinct from, but notionally similar to, \emph{causal interference} \citep{tchetgen2012causal, forastiere2021identification, sobel2006randomized, zigler2021bipartite, ogburn2014causal,bhattacharya2020causal}. Both interference and NLC arise from spatial (or network) interaction, and they both impose limitations on standard causal inference methods. Various works in this literature have discussed the role of conditional ignorability given neighboring covariates; for instance, in \citet{vansteelandt2007confounding} and \citet{ forastiere2021identification}. However, to the best of our knowledge, flexible statistical methods specifically addressing NLC by learning the dependencies with respect to neighboring covariates have been ignored.

\section{Potential outcomes and NLC}\label{sec:rubin}

The potential outcomes framework, also known as the Rubin Causal Model (RCM) \citep{rubin_for_2008}, distinguishes between the observed outcome $Y_s$ at unit $s$ and those that would be observed under counterfactual (potential) treatments $Y_s(a)$ (formally defined below). We start with some notaton.  The assigned treatment is denoted $A_s$. It is assumed to be binary for ease of presentation, although the ideas generalize to more general treatments. Next, ${\mathbb{S}}$ is the set where the outcome and treatment are measured; ${\mathbb{G}} \supset {\mathbb{S}}$ is a grid that contains the rasterized covariates $\{{\bm{X}}_s \in \mathbb{R}^d \colon s\in{\mathbb{G}}\}$; denoting ${\tens{X}}_{B}$ where $B\subset{\mathbb{G}}$ is a set means ${\tens{X}}_{B}= \{{\bm{X}}_s \mid s\in B\}$; and $X \perp \!\!\!\! \perp Y \mid Z$ means that $X$ and $Y$ are conditionally independent given $Z$.  Throughout ${\bm{X}}_s$ is assumed to consist of pre-treatment covariates only, meaning they are not affected by the treatment or outcome. Finally, we use the generic notation $p(U)$ to denote the density or probability function of a random variable $U$.

\begin{mydef}[Potential outcomes]\label{def:potential-outcomes} The potential outcome $Y_s({\bm{a}})$ is the outcome value that \textit{would} be observed at location $s$ under the global treatment assignment ${\bm{a}}=(a_1,\hdots, a_{|{\mathbb{S}}|})$.
\end{mydef}

For $Y_s({\bm{a}})$ to depend only on $a_s$, the RCM needs an additional condition called the \textit{stable unit treatment value assumption}, widely known as SUTVA, and encompassing notions of {\it consistency} and ruling out {\it interference}.

\begin{assumption}[SUTVA] \label{def:sutva} (1) Consistency: there is only one version of the treatment. (2) No interference: the potential outcomes for one location do not depend on treatments of other locations. Together, these conditions imply that $Y_s({\bm{a}})=Y_s(a_s)$ for any assignment vector ${\bm{a}}\in\{0,1\}^{|{\mathbb{S}}|}$, and that the observed outcome is the potential outcome for the observed treatment, i.e., $Y_s=Y_s(A_s)$.
\end{assumption}

The potential outcomes and SUTVA allow to define an important estimand of interest: the average treatment effect.

\begin{mydef}[ATE]\label{def:ate} The average treatment effect (ATE) is the quantity
$\tau_\text{ATE} = |{\mathbb{S}}|^{-1}\textstyle{\sum_{s\in{\mathbb{S}}}} \left\{Y_s(1) - Y_s(0)\right\}$.
\end{mydef}

\begin{figure}[htb]
    \centering
    \begin{subfigure}[t]{.3\linewidth}\centering
    \includegraphics[width=.85\linewidth]{img/causal_diagram_local.png}
    \subcaption{ Local Confounding.}
    \label{fig:confounding-types:local}
    \end{subfigure}
    \begin{subfigure}[t]{.3\linewidth}\centering
    \includegraphics[width=.85\linewidth]{img/causal_diagram_nonlocal.png}
    \subcaption{ NLC, no interference.}
    \label{fig:confounding-types:nlc}
    \end{subfigure}
    \begin{subfigure}[t]{.3\linewidth}\centering
    \includegraphics[width=.85\linewidth]{img/causal_diagram_interf.png}
    \subcaption{ Interference, no NLC.}
    \label{fig:confounding-types:interf}
    \end{subfigure}
    \caption{Confounding types.}
    \label{fig:confounding-types}
\end{figure}

One cannot estimate the ATE directly since one never simultaneously observes $Y_s(0)$ and $Y_s(1)$. The next assumption in the RCM formalizes conditions for estimating the ATE, (or other causal estimands) with observed data by stating that any  observed association between $A_s$ and $Y_s$ is not due to an unobserved factor.

\begin{assumption}[Treatment Ignorability]\label{def:ignorability} The treatment $A_s$ is ignorable with respect to some vector of controls ${\bm{L}}_s$ if and only if $Y_s(1), Y_s(0) \perp \!\!\!\! \perp A_s \mid {\bm{L}}_s$.
\end{assumption}

For the sake of brevity, we will say that ${\bm{L}}_s$ is \emph{sufficient} to mean that the treatment is ignorable conditional on ${\bm{L}}_s$. NLC occurs when local covariates are not sufficient. It is formally stated as follows:

\begin{mydef}[Non-local confounding] We say there is non-local confounding (NLC) when there exist neighborhoods $\{{\mathcal{N}}_s \subset {\mathbb{G}} \mid s\in{\mathbb{S}} \}$ such that ${\bm{L}}_s={\tens{X}}_{{\mathcal{N}}_s}$ is sufficient and the neighborhoods are necessarily non-trivial (${\mathcal{N}}_s\neq\{s\}$).
\end{mydef}

Figures \ref{fig:confounding-types:local} and \ref{fig:confounding-types:nlc} show a graphical representation of local confounding and NLC respectively.  Horizontal dotted lines emphasize that there may be spatial correlations in the covariate, treatment and outcome processes that do not result in confounding. For contrast, figure \ref{fig:confounding-types:interf} shows the distinct phenomenon of (direct) interference, in which $A_{s'}$ affects $A_{s}$ \citep{ogburn2014causal}. We remark that this depiction of is only on of the forms that interference can take. For instance, it may also happen through contagion \citep{ogburn2014causal}. But they key point is that \ref{fig:confounding-types:nlc} shows a scenario where the neighboring covariates suffice for ignorability and SUTVA still holds, whereas this is generally not the case with interference, which usually requires other causal estimands and implies the violation of SUTVA \citep{ogburn2014causal, forastiere2021identification}. Additional discussion of interference is in section \ref{sec:interference}.

Subsequent discussion of the size of the NLC neighborhood, ${\mathcal{N}}_s$, will make use of the following proposition stating that a neighborhood containing sufficient confounders can be enlarged without sacrificing the sufficiency.

\begin{proposition}\label{prop:pre-treatment}
Let ${\bm{L}}_s$ be a sufficient set of controls including only pre-treatment covariates. and let ${\bm{L}}_s'$ be another set of controls satisfying ${\bm{L}}_s' \supset {\bm{L}}_s$. Then, ${\bm{L}}_s'$ is also sufficient.
\end{proposition}

All the proofs are in Appendix \ref{appendix:proofs}. We conclude this section with a classic result stating that any sufficient ${\bm{L}}_s$ can be used to estimate the ATE from quantities and relations in the observed data.

\begin{proposition}\label{prop:unbiasedness} Assume SUTVA holds and that ${\bm{L}}_s$ is sufficient. Then
\begin{equation}\label{eq:ate}
  \begin{aligned}
      \mathbb{E}\left[\mathbb{E}[Y_s \mid {\bm{L}}_s, A_s=1] - \mathbb{E}[Y_s \mid {\bm{L}}_s, A_s=0]\right],
  \end{aligned}
\end{equation}
is an unbiased estimator of $ \tau_\text{ATE}$ (where $s$ is taken uniformly at random from ${\mathbb{S}}$).
\end{proposition}

\section{Adjustment for NLC with {\em weather2vec}\xspace}

Accounting for NLC would be fairly straightforward provided infinite data and the right set of confounders.  By virtue of Proposition \ref{prop:pre-treatment}, one could, in principle, specify a non-linear regression $Y_s \approx f(A_s, {\tens{X}}_{\mathbb{G}}, s)$ that includes every non-local covariate $
{\bm{X}}_{s'}\in {\tens{X}}_{{\mathbb{G}}}$ as part of the regressors. With large model capacity and infinite repeated samples per location, this regression would perfectly estimate $\mathbb{E}[\mathbb{E}[Y_s \mid {\bm{L}}_s, A_s=a]$ and thus be able to estimate the ATE using Proposition \ref{prop:unbiasedness}. But this scenario is far from realistic. Most commonly, there will be only be one observation for each $s$, and ${\mathbb{S}}$ can also be small, requiring additional structure to enable statistical estimation.
Thus, we consider the question: what kind of statistical and functional model (e.g., to predict the probability of treatment) reflects the causal structure of NLC and allows for flexible statistical models under such restrictions?

One desirable statistical property to consider is \emph{spatial stationarity}. Intuitively, it entails that the distributions of $Y_s$ and $A_s$ with respect to a neighboring covariate ${\bm{X}}_{s'}$ should only depend on $s-s'$ (their relative position). Formally, it requires that for any set $B\subset {\mathbb{G}}$, displacement vector $\delta$, and $s\in {\mathbb{G}}$, the following identity holds $p(A_s, Y_s \mid {\tens{X}}_B=x)\overset{}{=} p(A_{s + \delta}, Y_{s + \delta} \mid {\tens{X}}_{B + \delta}=x)$.

This paper propose to use U-nets \citep{ronneberger2015u} to learn representations of NLC, since they are spatially stationary and adhere to the causal structure of NLC shown in figure \ref{fig:confounding-types:nlc}. An overview of U-nets is provided in the next section for completeness. A key property is that a U-net $f_\theta$ can transform the input covariates ${\tens{X}}_{\mathbb{G}}$ onto an output grid ${\tens{Z}}_{\theta,{\mathbb{G}}}:=f_\theta({\tens{X}}_{\mathbb{G}})$ of same spatial dimensions that each scalar or vector ${\bm{Z}}_{\theta, s}\in {\tens{Z}}_{\theta,{\mathbb{G}}}$ localizes contextual spatial information from the input grid.  Notice that U-nets are not the only NN architecture with these properties. For instance, residual architectures \citep{he2016deep} could also be considered.
Nonetheless, our focus here is not on choosing an optimal architecture, but to highlight some fundamental statistical properties that an architecture must satisfy to be useful in accounting for NLC. The essence of {\em weather2vec}\xspace is to define appropriate learning tasks to obtain the NN weights $\theta$. Two such tasks are considered, summarized below and described in detail in subsequent sections.

\begin{enumerate}[itemsep=0pt, itemindent=0pt, parsep=0pt]
  \item (\textbf{Supervised)}\quad Assuming the treatment and outcome are densely available over ${\mathbb{G}}$, regress $A_s$ on ${\bm{Z}}_{\theta, s}$ (propensity score) or, for a subset of untreated units, regress $Y_s$ on ${\bm{Z}}_{\theta, s}$ (prognostic score).
  \item (\textbf{Self-supervised)}\quad If the treatment and outcome are not densely available over ${\mathbb{G}}$, then learn $\theta$ so that ${\bm{Z}}_{\theta, s}$ is highly predictive of ${\bm{X}}_{s'}$ for any $s'$ within a specified radius of $s$.
\end{enumerate}

\subsection{An overview of the U-net for summarizing NLC information}

In this section, we provide a brief description of the U-net's functional form, referring the reader to \citet{ronneberger2015u} for additional details. The U-net transformation involves two parts: a \emph{contractive} stage and an symmetric \emph{expansive stage}. Both of these steps use convolutions with learnable parameters and non-linear functions to aggregate information from the input grid spatially and create rich high-level features. The convolutions in the contractive path duplicate the number of latent features at each layer. Then, these intermediate outputs go through \textit{pooling} layers which halve the spatial dimensions. Together, these operations augment the dimensionality of each point of the grid, combining information at many spatial points to richer information contained at fewer points. Convolutions propagate information spatially, and the deeper they are in the contractive path, the larger their propagation reach (in the original scale of the input grid).
The expansive path, on the other hand, uses \textit{up-sampling} to progressively interpolate the deep higher-level features back to a finer spatial lattice, and then uses convolutions to reduce back the latent dimensionality at each grid point; with the characteristic that, in contrast to the input grid, every point now localizes spatial information. The output vector can have any arbitrary dimension after possibly applying an additional linear or convolutional layer followings the expansive path (or before the contractive path, or both). The unknown weights $\theta$ dictate the size -- the ``radius of influence" -- and what non-local information is summarized by ${\bm{Z}}_{\theta, s}$.
The functional model is spatially stationary because all of its components are (convolutions, pointwise activations, pooling, and upsampling). Figure \ref{fig:architecture} provides a visual example of the U-net architecture.

\subsection{Learning NLC representations via supervision}\label{sec:supervised}

The supervised approach links the proposed representation learning to the procedure of learning a balancing score \citep{rubin2005causal, hansen2008prognostic} in causal inference. We recall the definition here for completeness.

  \begin{mydef}[Balancing score]\label{def:balancing} $b({\bm{L}}_s)$ is a {\em balancing score} iff $A_s \perp \!\!\!\! \perp {\bm{L}}_s \mid b({\bm{L}}_s)$. The coarsest balancing score is $b({\bm{L}}_s) := p(A_s=1 \mid {\bm{L}}_s)$, widely known as the \emph{propensity score}.
  \end{mydef}

  \begin{mydef}[Prognostic score]\label{def:prognostic}  $b({\bm{L}}_s)$ is a {\em prognostic score} iff $Y_s(0) \perp \!\!\!\! \perp {\bm{L}}_s \mid b({\bm{L}}_s)$. The coarsest prognostic score is $b({\bm{L}}_s) := \mathbb{E}[Y_s(0) \mid {\bm{L}}_s]$.
  \end{mydef}

  The propensity score blocks confounding through the treatment \citep{rubin2005causal}; prognostic scores do so through the outcome \citep{hansen2008prognostic}. The importance of these definitions is summarized by the next well-known result.

  \begin{proposition}\label{prop:bs-reduction}
  If $b({\bm{L}}_s)$ is a balancing score, then ${\bm{L}}_s$ suffices to control for confounding iff $b({\bm{L}}_s)$ does. The same result holds for the prognostic score under the additional assumption of no effect modification.
  \end{proposition}

  Thus, using one-dimensional latent dimension output for ${\bm{Z}}_{\theta,s}$, the straightforward regression loss functions for the propensity score and prognostic score regression are given respectively by

\begin{equation}\label{eq:sup-prop}
\begin{aligned}
{\mathcal{L}}_\text{sup}^\text{prop}(\theta) =
\sum_{s\in {\mathbb{S}}} \text{CrossEnt}(A_s, {\bm{Z}}_{\theta,s})
\end{aligned}
\end{equation}
\begin{equation}\label{eq:sup-prog}
\begin{aligned}
{\mathcal{L}}_\text{sup}^\text{prog}(\theta) =
\sum_{s\in {\mathbb{S}}\colon A_s = 0} (Y_s - {\bm{Z}}_{\theta,s})^2
\end{aligned}
\end{equation}
where CrossEnt is the binary cross-entropy loss. Algorithm \ref{alg:sup} summarizes the supervised {\em weather2vec}\xspace approach.

\begin{algorithm}[tb]
\caption{Supervised representation learning and causal estimation}
\label{alg:sup}
\textbf{Input}: Covariates ${\bm{X}}_{{\mathbb{G}}}$; outcome and treatment $Y_s$ and $A_s$ at a dense subset ${\mathbb{S}}\subset {\mathbb{G}}$; and U-net model $f_\theta$ with 1-dimensional outputs.
\begin{algorithmic}[1] %
\STATE Obtain neural network weights $\hat{\theta}$ minimizing the supervised loss (equations \ref{eq:sup-prop} or \ref{eq:sup-prog}) and compute output balancing score grid from optimal weights $\hat{{\tens{Z}}}_{{\mathbb{G}}}=f_{\hat{\theta}}({\bm{X}}_{{\mathbb{G}}})$.
\STATE Use standard causal inference methods (e.g., IPTW in Appendix \ref{app:iptw}) to obtain $\hat{\tau}_\text{ATE}$ adjusting for $\hat{{\bm{Z}}}_{s}$ and other relevant local confounders  at each unit $s$.
\end{algorithmic}
\end{algorithm}

Learning $\theta$ through supervision results in an efficient scalar ${\bm{Z}}_{\theta, s}$ compressing NLC information, allowing for $\theta$ to just attend to relevant neighboring covariate information that pertains to confounding. However, supervision may not be possible with small-data studies where $Y_s$ and $A_s$ are only measured sparsely. In such cases, the supervised model will likely overfit to the data. For example, in application 1 in section \ref{sec:applications}, ${\mathbb{S}}$ consists only of measurements at 473 power plants, while the size of ${\mathbb{G}}$ is $128 \times 256$. An over-fitted propensity score would result in insufficient ``overlap'' \citep{stuart2010matching} by assigning zero probability to the unobserved treatment, resulting in causal inferences that would rely on model extrapolation to areas where covariate information is not represented in  both treatment groups.
To avoid this, the self-supervised approach
targets scenarios with sparse ${\mathbb{S}}$.

\subsection{Representations via self-supervised dimensionality reduction}\label{sec:self-supervised}

Self-supervision frames the representation learning problem as dimension reduction without reference to the treatment or outcome. The representations are then used to learn a balancing score for causal effect estimation in a second analysis stage.  This approach requires specification of a fixed neighborhood ${\mathcal{N}}_s$ (parameterized by a radius $R$) and latent dimension $k$, resulting on different representations for different hyper-parameter choices, which can be selected using standard model selection techniques (such as AIC) in the second stage. The dimension reduction's objective is that ${\bm{Z}}_{\theta, s}$ encodes predictive information of any ${\bm{X}}_{s + \delta}$ for $(s + \delta) \in {\mathcal{N}}_s$.
A simple predictive model ${\bm{X}}_{s + \delta} \approx g_\phi({\bm{Z}}_{\theta, s}, \delta)$ is proposed. First, let ${\bm{\Gamma}}_{\phi}(\cdot)$ be a function taking an offset $\delta$ as an input and yielding a $k\times k$ matrix, and let $h_\psi(\cdot)\colon \mathbb{R}^k \to \mathbb{R}^d$ be a decoder with output values in the covariate space. The idea is to consider ${\bm{\Gamma}}_{\phi}(\delta)$ as a selection operator acting on ${\bm{Z}}_{\theta, s}$. The task loss function can be written succinctly as

\begin{equation}\label{eq:self}
\begin{aligned}
{\mathcal{L}}_\text{self}(\theta, \phi,\psi\mid R) &= \\
\sum_{s\in {\mathbb{G}}}\sum_{ \{\delta\colon \lVert \delta \rVert \leq R\}} & \left({\bm{X}}_{s + \delta} - h_\psi(  {\bm{\Gamma}}_\phi(\delta){\bm{Z}}_{\theta, s})\right)^2.
\end{aligned}
\end{equation}
 Algorithm \ref{alg:self} summarizes the causal estimation procedure. See Appendix \ref{appendix:motivation-selfsup} for additional intuition. A connection with PCA is described in Appendix \ref{appendix:pca}.

\begin{algorithm}[tb]
\caption{Self-supervised representation learning and causal estimation}
\label{alg:self}
\textbf{Input}: Covariates ${\bm{X}}_{{\mathbb{G}}}$; outcome and treatment $Y_s$ and $A_s$ at any subset ${\mathbb{S}}\subset {\mathbb{G}}$; U-net model $f_\theta$ with $k$-dimensional outputs; and grid of candidate radii ${\mathcal{R}}=\{R_1,\hdots,R_\text{max}\}$.
\begin{algorithmic}[1] %
\STATE For each $R \in {\mathcal{R}}$, obtain neural network weights $\hat{\theta}^R$ minimizing the self-supervised loss (equations \ref{eq:self}) and compute $\hat{{\tens{Z}}}_{R, {\mathbb{G}}}=f_{\hat{\theta}^R}({\bm{X}}_{{\mathbb{G}}})$.
\STATE Use a propensity/prognostic score model to choose the optimal $r$. For example, use a logistic regression taking inputs $\hat{{\bm{Z}}}_{R,s}$ (and other relevant local confounders) and compute the Akaike information criterion (AIC). Then choose $\hat{r}$ that minimizes the AIC.
\STATE Use standard causal inference methods (e.g., IPTW in Appendix \ref{app:iptw}) to obtain $\hat{\tau}_\text{ATE}$ adjusting for $\hat{{\bm{Z}}}_{\hat{R}, s}$ and other local confounders at each unit $s$.
\end{algorithmic}
\end{algorithm}

\section{NLC and Interference}\label{sec:interference}

Section \ref{sec:intro} briefly contrasted NLC with the related problem of interference, a topic that we expand here. We first formalize the concept of interference. To ground the discussion, we follow closely the form of interference considered in \citet{forastiere2021identification}, which replaces SUTVA with the following neighborhood-level assumption, termed the \textit{stable unit neighborhood treatment value assignment} (SUTNVA).

\begin{assumption}[SUTNVA] \label{def:sutnva} (1) Consistency: there is only one version the treatment. (2) Neighborhood-level interference: for each location $s$, there is a neighborhood ${\mathcal{N}}_s$ such that the potential outcomes depend only on the treatments at ${\mathcal{N}}_s$. Together, these conditions imply that $Y_s({\bm{a}})=Y_s({\bm{a}}_{{\mathcal{N}}_s})$ for any assignment vector ${\bm{a}}\in\{0,1\}^{|{\mathbb{S}}|}$, and that the observed outcome is the potential outcome for the observed treatment, i.e., $Y_s=Y_s({\tens{A}}_{{\mathcal{N}}_s})$.
\end{assumption}
This definition of interference only considers \emph{direct} interference \citep{ogburn2014causal}, leaving aside indirect mechanisms such as contagion \citep{ogburn2018challenges, shalizi2011homophily}. Investigating the role of NLC in such scenarios is left for future work. We now describe one generalization of the ATE for this type of direct interference.  The statement uses potential outcomes of the form  $Y_s(a_s=a, {\tens{A}}_{{\mathcal{N}}_s\setminus \{s\}})$ -- a short-hand notation for the potential outcome that assigns the treatments of all the neighbors of $s$ to their observed treatments in the data.

\begin{mydef}[DATE] The direct average treatment effect (DATE) is the quantity $
\tau_\text{DATE} =  |{\mathbb{S}}|^{-1}\textstyle{\sum_{s\in{\mathbb{S}}}} \{Y_s(a_s=1, {\tens{A}}_{{\mathcal{N}}_s\setminus \{s\}})  -Y_s(a_s=0, {\tens{A}}_{{\mathcal{N}}_s\setminus \{s\}}) \}.$
\end{mydef}

Notice that $\tau_\text{DATE}$ is not the only estimand of interest. For instance, the literature often considers \emph{spill-over} effects \citep{ogburn2018challenges}. But for this work, we only focus on $\tau_\text{DATE}$, leaving other estimands for future work. Now consider the question: how does NLC affect the estimation of $\tau_\text{DATE}$ based on \eqref{eq:ate}? We provide a partial answer based on the following result by \citet{forastiere2021identification}, which states two conditions under which \eqref{eq:ate} is an unbiased estimator of $\tau_\text{DATE}$.

\begin{proposition}\label{prop:interference} Assume SUTNVA.
If (1) ${\tens{A}}_{{\mathcal{N}}_s} \perp \!\!\!\! \perp (Y_s({\bm{a}}))_{{\bm{a}}\in\{0,1\}^{|{\mathcal{N}}_s|}} \mid {\bm{L}}_s$  and (2) $A_s \perp \!\!\!\! \perp A_{s'} \mid {\bm{L}}_s$ for all $s\in{\mathbb{S}}, s'\in {\mathcal{N}}_s$. Then
\eqref{eq:ate} is an unbiased estimator of $\tau_\text{DATE}.$
\end{proposition}

\begin{figure}
    \centering
    \includegraphics[width=.3\linewidth]{img/causal_diagram_interf_nlc.png}
    \caption{ Interference + NLC.}
    \vspace*{-10pt}
    \label{fig:interf2}
\end{figure}

Conditions (1) and (2) correspond, respectively, to the notions of neighborhood-level ignorability and conditional independence of the neighboring treatments. It turns out that
when NLC is present (the arrows from $X_{s'}$ in figure \ref{fig:confounding-types:nlc})), conditions (1) and (2) in the proposition can be violated. To see this, consider Figure \ref{fig:interf2} representing the co-occurrence of interference and NLC. Adjusting only for local covariates would violate neighborhood ignorability condition (1) with a spurious correlation between $Y_s$ and $A_{s'}$ (through the backdoor path $Y_s \leftarrow {\bm{X}}_{s'} \to A_{s'}$). Similarly, a spurious correlation between $A_s$ and $A_{s'}$ would persist (via the the path $A_s \leftarrow {\bm{X}}_{s'} \to A_{s'}$). For such cases, {\em weather2vec}\xspace can play an important role in satisfying (1) and (2) since, after controlling for NLC (consisting in Figure \ref{fig:interf2} of adjusting for both ${\bm{X}}_s$ and ${\bm{X}}_{s'}$ and blocking the incoming arrows from neighboring covariates into one's treatments and outcomes), the residual dependencies would more closely resemble those of Figure \ref{fig:confounding-types:interf}. In summary, adjusting for NLC with {\em weather2vec}\xspace can aid satisfaction of the conditional independencies required to estimate causal effects with the same estimator used to estimate the ATE absent interference.

\begin{figure*}[htb]
\centering
\includegraphics[width=0.9\textwidth]{img/simstudy_results.png}
\caption{Simulation study results. (a) the basic linear task. (b) the non-linear task. Both tasks use outcomes and treatments densely available on the grid (${\mathbb{S}}={\mathbb{G}}$). (c) and (d) are the same tasks respectively, but only 500 random locations are in ${\mathbb{S}}$.}
\label{fig:simulation}
\end{figure*}

\begin{figure*}[htb]
  \centering
  \begin{subfigure}[t]{.31\textwidth}
    \centering
    \includegraphics[width=0.99\textwidth]{img/explained_variance.pdf}
   \caption{Self-supervision, NARR data}
   \label{fig:app1-self-supervision}
  \end{subfigure}%
  \begin{subfigure}[t]{.36\textwidth}
    \centering
    \includegraphics[width=0.99\textwidth]{img/ex1-metrics.png}
    \caption{Fit metrics, propensity score model}
    \label{fig:app1-fit-metrics}
  \end{subfigure}%
  \begin{subfigure}[t]{.27\textwidth} \centering
    \includegraphics[width=0.99\textwidth]{img/dapsm_estimates.png}
    \caption{Estimated causal effects}
    \label{fig:app1-effects}
  \end{subfigure} %
  \caption{Application 1: The effectiveness of catalytic devices to reduce power plant ozone emissions.
  }
  \label{fig:app1}
\end{figure*}

\section{Simulation study}\label{sec:simulation}

We conduct a simulation study that roughly mimics a dataset where pollution is dispersed in accordance with non-local meteorological covariates as in our applications. We briefly describe the setup here and the results. Appendix \ref{appendix:simulation} contains a detailed explanation and additional visualizations.

The covariates (simulating wind vectors) are generated from the gradient field of a random spatial process. The treatment probability and the outcome (simulating air pollution) are non-local functions of the covariates such that areas with lower outcomes have a higher probability of treatment, with a fixed treatment effect of 0.1. Two varying factors are considered: whether ${\mathbb{S}}$ is dense or sparse; and whether the simulated data is linear or non-linear on the covariates. We then use the self-supervised (W2V-self) and supervised (W2V-sup) variants in a propensity score model for causal estimation using IPTW (Appendix \ref{app:iptw}). Five additional baselines are considered: no controls (Unadjusted); controlling for local covariates (Local only); controlling for local and averages of neighboring covariates (Local$+$Averages); spatial random effects (spatial RE only), and an ablation of the supervised variant adding spatial random effects (W2V-sup$+$spatial RE). A total of 10 experiments are conducted for each configuration. The results are shown in Figure \ref{fig:simulation}. When ${\mathbb{S}}$ is dense, W2V-sup outperforms all other methods (panels a and b), exhibiting near zero bias in the linear case and a small amount of finite-sample bias for the non-linear task. When ${\mathbb{S}}$ is sparse (panels c and d), W2V-self gives the best results, noting further the poor performance of W2V-sup due to overfitting. Alone, spatial RE outperforms Local and Local$+$Averages, but hurts in W2V-sup$+$spatial RE.

\section{Applications in Air Pollution and Climate}\label{sec:applications}

\paragraph{Application 1: Quantifying the impact of power plant emission reduction technologies}

The study aims to quantify the impact of SCR/SNCR catalytic devices \citep{muzio2002overview} to reduce emissions among coal-fired power plants in the U.S \citep{DVN/M3D2NR_2016}. Appendix \ref{appendix:app1} provides a description of the dataset. Since air quality regulations are inherently regional and power plants are concentrated in regions with similar weather and economic demand factors, regional weather correlates with the assignment of the intervention. Further, weather patterns (such as wind vectors, precipitation and humidity) dictate regional differences in the formation and dispersion of ambient air pollution. Thus, the weather is a potential confounding factor which cannot be entirely characterized by local measurements.
s

\begin{figure*}[htb]
  \centering
  \begin{subfigure}[t]{.3\textwidth}
    \centering
    \includegraphics[width=0.99\textwidth]{img/prognostic_score.pdf}
   \caption{Prognostic score fit averaged over the entire grid ${\mathbb{G}}$.}
   \label{fig:app2-prognostic}
  \end{subfigure}\hfill%
  \begin{subfigure}[t]{.68\textwidth}
    \centering
    \includegraphics[width=0.5\textwidth]{img/app2_yearly.png}\hfill
    \includegraphics[width=0.5\textwidth]{img/app2_monthly.png}
    \caption{Detrended series at ${\mathbb{S}}^*$ resembles power plant emissions. (\emph{Left}) Yearly trend $\delta_{\textrm{year}(t)}$. (\emph{Right}) Monthly trend $\delta_{\text{year}(t)} + \gamma_{\text{month}(t)}$}
    \label{fig:app2-trends}
  \end{subfigure} %
  \caption{Application 2: Meteorological detrending of $\textrm{SO}_4$.
  }
  \label{fig:app2}
\end{figure*}

\emph{Self-supervised features from NARR.} We downloaded monthly NARR data \citep{mesinger2006north} containing averages of gridded atmospheric covariates across mainland U.S. for the period 2000-2014. We considered 5 covariates: temperature at 2m, relative humidity, total precipitation, and north-sound and east-west wind vector components. For each variable we also include its year-to-year average. Each grid cell covers roughly a $32\times 32$ km area and the lattice size is $128 \times 256$. We implemented the self-supervised {\em weather2vec}\xspace with a lightweight U-net of depth 2, 32 hidden units, and only one convolution per level. See Appendix \ref{appendix:app1} for more details and schematic of the U-net architecture. To measure the quality of the encoding, Figure \ref{fig:app1-self-supervision} shows the percentage of variance explained ($R^2$), comparing with neighbor averaging and local values. The results shows that the 32-dimensional self-supervised features provide a better reconstruction than averaging and using the local values. For instance, the 300km averages only capture 82\% of the variance, while the self-supervised {\em weather2vec}\xspace features capture 95\%. See Appendix \ref{appendix:app2} for details on the calculation of the $R^2$ and neural network architecture.

\emph{Estimated pollution reduction.}  We evaluate different propensity score models for different neighborhood sizes of the June 2004 NARR {\em weather2vec}\xspace-learned features with the same logistic model and other covariates as in DAPSm, augmented with the self-supervised features. We selected the representation using features within a 300km radius on the basis of its accuracy, recall, and AIC in the propensity score model relative to other considered neighborhood sizes (Figure \ref{fig:app1-fit-metrics}).
The causal effects are then obtained by performing 1:1 nearest neighbor matching on the estimated propensity score as in DAPSm.
Figure \ref{fig:app1-effects} compares treatment effect estimates for different estimation procedures. Overall, standard (naive) matching using the self-supervised features is comparable to DAPSm, but without requiring the additional spatial adjustments introduced by DAPSm. The same conclusion does not hold when using local weather only
, which (as in the most naive adjustment) provides the scientifically un-credible result that emissions reduction systems significantly {\it increase} ozone pollution. Do notice the wide confidence intervals which are constructed using conditional linear models fit to the matched data sets \citep{ho2007matching}. Thus, while the mean estimate shows a clear improvement, the intervals shows substantial overlap, warranting  caution.

\paragraph{Application 2: Meteorological detrending of sulfate}\label{sec:detrending}

We investigate meteorological detrending of the U.S. sulfate ($\textrm{SO}_4$) time series with the goal (common to the regulatory policy and atmospheric science literature) of adjusting long-term pollution trends by factoring out meteorologically-induced changes and isolating impacts of emission reduction policies \citep{wells2021improved}. We focus on $\textrm{SO}_4$ because it is known that its predominant source in the U.S. is $\textrm{SO}_2$ emissions from coal-fired power plants, on which observed data are available for comparison. Thus, we hypothesize that an effectively detrended $\textrm{SO}_4$ time series will closely resemble that of the power plant emissions.

\emph{Prognostic score.} We obtained gridded $\textrm{SO}_4$ concentration data publicly available from the Atmospheric Composition Analysis Group \citep{so4data, van2021monthly}, consisting of average monthly value for each raster cell in the mainland U.S. for the period of study 2000--2014. The data is aggregated into 32km-by-32km grids to match the resolution of atmospheric covariates. The model uses a U-net with quadratic loss for the (log) concentrations of $\text{SO}_4$.
Since the prognostic score is defined based on outcome data in the absence of treatment, we leverage the fact that the power plant emissions were relatively constant for the period 2000-2005 and using 2006 as test data -- regarding this period as absent of treatment. The model predictions, aggregated by all points in the grid is shown in Figure \ref{fig:app2-prognostic}.
The difference between the red line (the prognostic score fit) and the black dotted line (the $\textrm{SO}_4$) observations during 2000 - 2006 is a proxy for the meteorology-induced changes in the absence of treatment.

\emph{Trend estimation.} For comparability we adhere to the recommended detrending model by \citep{wells2021improved}.  Accordingly, we specify a regression with a year and seasonal fixed-effect term. Rather than pursue an entirely new methodology for detrending, we intentionally adhere to standard best practices and merely aim to evaluate whether augmenting this approach with the {\em weather2vec}\xspace representation of the prognostic score offers improvement. The outcome $\log(Y_{s,t})$ for untreated units is regressed using the predictive model
\begin{equation}\label{eq:detrending}
\begin{aligned}
\mu_{s,t} &=\alpha + \delta_{\text{year}(t)} + \gamma_{\text{month}(t)} + \textstyle{\sum_{j=1}^p} \beta_{p} X_{st}^p\\
 \end{aligned}
\end{equation}
for all $s\in{\mathbb{S}}^*$ and $t=1,\hdots, T$; and where $\delta_\ell$ is the year effect for $\ell=2000,\hdots,2014$; $\gamma_\kappa$ is the seasonal (monthly) effect for $\kappa=1,\hdots,12$; ${\mathbb{S}}^* \subset {\mathbb{S}}$ are the locations of the power plants; and $X_{st}^p$ are the controls with linear coefficients $\beta_{s,p}$. These controls are obtained from a B-spline basis of degree 3 using: 1) local weather only, and 2) local weather plus the {\em weather2vec}\xspace prognostic score. The model is fitted using Bayesian inference with a Gibbs sampler. Figure \ref{fig:app2-trends} shows the fitted (posterior median) yearly and monthly trends, which resemble the power plant emissions trends much more closely than the predicted trends from models that include local or neighborhood average weather. Notice the ``double peak'' per year in the monthly power plant emissions (owing to seasonal power demand), which is only captured by the detrended {\em weather2vec}\xspace series.

\section{Discussion and Future Work}

While notions of NLC have been acknowledged in causal inference (most explicitly in spatial econometrics but also alluded to in literature on spatial confounding and interference), potential-outcomes formalization of NLC and flexible tools to address it are lacking. We offer such a formalization, along with a flexible representation learning approach to account for NLC with gridded covariates and treatments and outcomes measured (possibly sparsely) on the same grid.  Our proposal is most closely tailored to problems in air pollution and climate science, where key relationships may be confounded by meteorological features, and promising results from two case studies evidence the potential of {\em weather2vec}\xspace to improve causal analyses over those with more typical accounts of local weather. A limitation of the approach is that the learned {\em weather2vec}\xspace representations are not as interpretable as direct weather covariates and using them could impede transparency when incorporated in policy decisions. Future work could explore new methods for interpretability. Other extensions could include additional data domains, such as graphs and longitudinal data with high temporal resolution. The links to causal interference explored in Section \ref{sec:interference} also offer clear directions for future work to formally account for NLC in the context of estimating causal effects with interference and spill-over.