
\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.



