Principal component analysis balancing prediction and approximation accuracy for spatial data

Let $Y\in\mathbb{R}^{n\times p}$ represent $p$-dimensional spatially-structured variables of interest measured at $n$ locations. Examples include the concentrations of a mixture of $p$ pollutants at $n$ monitoring sites, or the expression of $p$ genes at $n$ spots on tissues. In addition, we observe the spatial coordinates $\{(s_{i1},s_{i2})\}_{i=1}^{n}$ of the $n$ locations, along with (potential) observations of $d$ covariates, $X\in\mathbb{R}^{n\times d}$. These could be, for example, population density and land use information for air pollution studies, or histology information from images for spatial transcriptomics.

The ultimate goal is to conduct prediction, clustering or other modeling tasks on $Y$ based on the spatial coordinates and covariates. Since the columns of $Y$ may be strongly correlated and/or noisy (Rao et al., 2021), dimension reduction can be used to to extract scientifically meaningful signals from the original data. The classical PCA achieves this by maximizing the proportion of variability in $Y$ that is explained by the lower dimensional PCs. We briefly introduce the classical PCA algorithm under the formulation of Shen and Huang (2008) to highlight its connection with related techniques.

Suppose the columns of $Y$ are centered and scaled. To find an $r$-dimensional representation of $Y$, classical PCA can be expressed as a sequence of rank-1 optimization problems for $l=1,\ldots,r$ (Shen and Huang, 2008):

where the 1-dimensional PC score $u_{l}$ is an $n$-vector, the loading $v_{l}$ is $p\times 1$, and $\lVert\cdot\rVert_{F}$ and $\lVert\cdot\rVert_{2}$ represent Frobenius norm for matrices and $\ell_{2}$ norm for vectors, respectively. $Y^{(l)}$ is the residual from approximation after each iteration: denoting $\tilde{u}_{l},\tilde{v}_{l}$ as the solution to (1), we define $Y^{(l)}=Y^{(l-1)}-\tilde{u}_{l-1}\tilde{v}_{l-1}^{\top}$ for $2\leq l\leq r$ and $Y^{(1)}=Y$.

While each of the $u_{l}$’s explain the greatest variability in $Y$, there is no guarantee that they are scientifically interpretable, and they do not explicitly account for the spatial structures underlying $Y$. Recognizing this, predictive PCA (Jandarov et al., 2017) builds upon the expression in (1), but constrains each PC score $u_{l}$ to be in a model space of the spatial coordinates and covariates $X$, leading to the optimiztion problem,

for each $l=1,\ldots,r$. Here, $Z=[X\ B]$, $B\in\mathbb{R}^{n\times m}$ contains $m$ thin-plate regression spline basis functions capturing the spatial effects, and in contrast to (1), normalization is done on the term $Z{\alpha_{l}}$ instead of $v_{l}$. As before, $Y^{(l)}$ is the residual after approximation, i.e. $Y^{(l)}=Y^{(l-1)}-Y^{(l-1)}\frac{\tilde{v}_{l-1}\tilde{v}_{l-1}^{\top}}{\lVert\tilde{v}_{l-1}\rVert_{2}}$ for $2\leq l\leq r$ and $Y^{(1)}=Y$. Note that the quantity $Y^{(l-1)}\tilde{v}_{l-1}$ corresponds to $\tilde{u}_{l-1}$ in the expression of PCA. Compared with (1), this formulation restricts each PC score $u_{l}=\frac{Z\alpha_{l}}{\lVert Z\alpha_{l}\rVert_{2}}$ to fall within the space spanned by the columns of $Z$. This ensures the spatial smoothness and the inclusion of information from covariates $X$ in the PC score $u_{l}$, and thus enforces predictability.

After solving either (1) or (2) for all $r$ PCs, the loadings are $\ell_{2}$-normalized ($\tilde{v}_{r}\leftarrow\tilde{v}_{r}/\lVert\tilde{v}_{r}\rVert_{2}$) if not already done, and then concatenated as $\tilde{V}_{p\times r}:=[\tilde{v}_{1}\ \ldots\ \tilde{v}_{r}]$. The PC scores are then defined as $\tilde{U}_{n\times r}:=Y\tilde{V}$.

As noted earlier, classical PCA aims to achieve optimal representability of the PC scores $\tilde{u}$, but these $\tilde{u}$’s may not retain meaningful signals in $X$ or important spatial patterns to be well predictable. In contrast, predictive PCA and its variants (Jandarov et al., 2017; Bose et al., 2018) minimize the approximation gap, which will be formally defined in Section 2.2, after constraining each $\tilde{u}$ to lie within a specific model space; but when these PC scores $\tilde{u}$ are predicted at new locations and transformed back to the higher dimensional space of $Y$, closeness in the space of $\tilde{u}$ may not necessarily translate to the original space of observations $Y$ due to reduced approximation accuracy. Our proposed method, introduced in Section 3, is based on an interpolation between classical and predictive PCA, and encourages $\tilde{u}$ to be close to, but not exactly within, the specified model space while balancing the quality of representation.

Before introducing our proposal, it is helpful to formalize different aspects of dimension reduction performances, such as representability and predictability, and to characterize the overall balance between them.

Suppose the PC scores $\tilde{U}_{\text{trn}}=[\tilde{u}_{1}\ \ldots\tilde{u}_{r}]$ and loadings $\tilde{V}$ are obtained on a set of training data with size $n_{\text{trn}}$, and $\tilde{u}_{1},\ldots,\tilde{u}_{r}$ are then predicted at $n_{\text{tst}}$ test locations with a user-specified spatial prediction model as $\hat{U}_{\text{tst}}:=[\hat{u}_{1}\ \ldots\ \hat{u}_{r}]$. We define the mean-squared prediction error (MSPE) of this procedure as

where $U^{*}=\arg\min_{U}\lVert Y_{\text{tst}}-U\tilde{V}^{\top}\rVert_{F}^{2}=Y_{\text{tst}}\tilde{V}$. Thus, $U^{*}$ is what the actual PC scores on the test set would be according to the loadings $\tilde{V}$, if $Y_{\text{tst}}$ were known. Consequently, MSPE characterizes the gap between the predicted and true PC scores and reflects the predictability of the PCs.

The mean-squared representation error (MSRE) is defined as

which measures the gap between the true data, $Y_{\text{tst}}$, and the closest approximation possible based on $\tilde{V}$. Since the quality of representation alone, without considering predictive performance, is of more interest for the training than the test data, we instead examine MSRE-trn as the metric for representation: $\text{MSRE-trn}={n_{\text{trn}}}^{-1}\lVert Y_{\text{trn}}-\tilde{U}_{\text{trn}}\tilde{V}^{\top}\rVert_{F}^{2}$. It can be seen that MSRE-trn exactly matches the objective function for PCA in (1).

Given any outcome data $Y$, loadings $\tilde{V}$ from a dimension reduction procedure and predicted PC scores $\hat{U}$, we define the total mean-squared error (TMSE) resulting from both dimension reduction and prediction as $\text{TMSE}=n^{-1}\left\lVert Y-\hat{U}\tilde{V}^{\top}\right\rVert_{F}^{2}$. This quantity measures the discrepancy between the true data $Y$ and the predicted scores that are transformed back to the original, higher-dimensional space. In the specific training-test setting, the expression of TMSE becomes

MSPE and MSRE can be viewed informally as a decomposition of the overall TMSE, where the prediction and approximation gaps, i.e., $(U^{*}-\hat{U}_{\text{tst}})\tilde{V}^{\top}$ and $Y_{\text{tst}}-U^{*}\tilde{V}^{\top}$, are components of the overall error $Y_{\text{tst}}-\hat{U}_{\text{tst}}\tilde{V}^{\top}$.

We first consider the setting with 3 well-predictable PCs and 3 PCs mainly consisting of structured and unstructured Gaussian errors. More specifically, for the mean model in (5), we let $f_{l}(X)=X\beta_{l}$ where the entries of $\beta_{l}$ are drawn independently from Uniform$[-1,1]$, and $\sigma_{l}^{2}=0$, for $l=1,2,3$; we let $f_{l}=0$ and $\sigma_{l}^{2}=1$ for $l=4,5,6$. $M$ is set to be $\Lambda\tilde{M}$, where $\Lambda$ is diagonal with $\Lambda_{ll}=1/\text{sd}\big{(}\text{PC}^{(l)}\big{)}$ and the entries of $\tilde{M}$ are independent Uniform$[-1,1]$. In other words, we first scale the 6 PCs to have equal variability, and then assign random but overall comparable weights to them.

Comparing the two existing methods, it is expected that in this scenario predictive PCA would outperform classical PCA in prediction without severely compromising approximation accuracy, since the predictable PCs explain a similar amount of variability compared to the unpredictable ones. Consequently, predictive PCA would achieve better overall performance as reflected by TMSE. RapPCA flexibly interpolates between classical and predictive PCA, and is expected to have comparable or better performance than predictive PCA.

Panel A in Figure 1 presents the breakdown of MSPE, MSRE-trn and TMSE for the first PC by $\gamma$, in comparison to PCA and predictive PCA for this scenario. We only present this breakdown for the first PC but not the subsequent ones because starting from the second PC, different PCA algorithms will have different residuals $Y^{(l)}$ (recall Equations 1, 2 and 4), and the PC-specific breakdown of metrics is consequently no longer comparable. We observe that MSPE decreases and training MSRE increases as we increase $\gamma$ for RapPCA, since this imposes a greater penalty on prediction errors in optimization.

Figure 2 compares the prediction MSEs for each PC, as well as the overall TMSE, MSPE and MSRE-trn for different methods. As expected for this scenario, we observe in Panel A that predictive PCA achieves lower prediction errors than PCA for each PC, as well as lower overall MSPE. This advantage does not cost a higher approximation gap (MSRE), since the predictable components in the outcome $Y$ explain a similar amount of variability as the unpredictable ones. RapPCA is able to adjust the penalties on the covariate versus spatial effect terms adaptively, especially when there are a large number of spatial basis terms (recall the separate tuning parameters $\lambda_{1}$ and $\lambda_{2}$ in Equation 4), as opposed to predictive PCA which is restricted to a low-dimensional combination of covariates and spatial basis functions (recall the term $Z$ in Equation 2). Consequently, predictive PCA achieves better overall performance (reflected by smaller TMSE) than PCA, while RapPCA further improves the prediction and overall performance due to its flexibility to data-adaptively capture the covariate and spatial effects.

We now consider a scenario where the predictable PCs explain a higher proportion of variability than the unpredictable PCs. In particular, the simulation setting is similar to Scenario 1, except that the entries of $\tilde{M}$ are first drawn from independent Uniform$[-1,1]$ distributions, and then the rows are scaled to have decaying norm (recall that the first 3 PCs are more predictable than the last 3 PCs). We expect in this case that PCA and predictive PCA would have comparable prediction, approximation and overall accuracy, since the well-predictable components of $Y$ coincide with those that explain the greatest amount of variance.

For the same reason, adjusting $\gamma$ in this case would not lead to steep changes in MSPE, MSRE or TMSE for RapPCA, as reflected by the flat curves of each metric in Panel B of Figure 1. Despite the similar behaviors of PCA and predictive PCA, Panel B of Figure 2 reveals that RapPCA is still able to improve the prediction and overall performance by more flexibly exploiting the covariate and spatial information when identifying the PCs.

In our last simulation, we investigate the setting where the relationship between the true PCs and the covariates is non-linear. We modify the setting in Scenario 1 so that

for $l=1,2,3$, where $\odot$ and $\odot^{2}$ denote element-wise multiplication and square, and $\lfloor x\rfloor$ denotes the largest integer not exceeding $x$. In other words, we specify the mean function to be a combination of squared and interaction terms of the covariates, where interactions exist between pairs of consecutive covariates ($X_{\cdot 1}$ and $X_{\cdot 2}$; $X_{\cdot 3}$ and $X_{\cdot 4}$, etc).

Panel C of Figure 1 indicates that by flexibly handling non-linear effects, RapPCA gains a clear advantage in prediction accuracy compared to both predictive and classical PCA in this scenario, when we impose a large enough penalty $\gamma$ on prediction errors in optimization. It is not surprising that putting an emphasis on prediction leads to some loss in approximation accuracy; however, we observe from Panel C of Figure 2 that by balancing the two objectives data-adaptively, RapPCA achieves meaningful improvement in MSPE while maintaining comparable, if not superior, overall performance.

The performance of RapPCA in comparison with common existing methods, including classical and predictive PCA, is first illustrated with the multivariate traffic-related air pollution (TRAP) data in Seattle (Blanco et al., 2022). The study leverages a mobile monitoring campaign where a vehicle equipped with air monitors repeatedly collected two-minute samples at $n=309$ stationary roadside sites in the greater Seattle area. Repeated samples for the concentrations of 6 types of pollutants were obtained, including ultrafine particles (UFP), black carbon (BC), nitrogen dioxide (NO₂), carbon monoxide (CO), carbon dioxide (CO₂) and fine particulate matter (PM_2.5). Following Blanco et al. (2022), we focus on UFP, BC and NO₂ in our analysis due to data quality considerations. For UFP and BC, the concentrations were measured with multiple instruments corresponding to different measurement ranges or units. In particular, NanoScan instrument with 13 different bin sizes, as well as DiSCmini and PTRAK 8525 (with and without diffusion screen) instruments measured the concentration of UFP in terms of the counts of particles with different size ranges, along with the median size and overall count of particles. Black carbon was assessed by micro-aethalometer (MA200), which measures the concentration of particles with different light absorbing properties, corresponding to 5 different ranges of wavelengths.

The median 2-minute visit concentrations were winsorized at the site level such that concentrations below the 5^th and above the 95^th quantile for a given site were set to those thresholds, respectively. This was done to reduce the influence of outliers in the measurements. These winsorized medians were then averaged for each site, leading to $p=23$ annual average measurements of pollutant concentrations in total, including 17 for UFP, 5 for BC, and one for NO₂. We use the concentration of NO₂ to normalize all other variables except the median size measurements for UFP from the DiSCmini. All variables are centered and scaled before running each PCA algorithm.

We are interested in extrapolating the pollutant measurements to unobserved locations in the same study region. Because of the autocorrelation between these measurements, predicting each of them separately might not always lead to sensible results; instead, dimension reduction is needed before building spatial prediction models. Dimension reduction also allows us to understand the health effects of pollutant mixtures rather than separate pollutants alone. As in Section 4, we compare the performance of RapPCA with classical and predictive PCA, each extracting the top 3 PCs from the 23 measurements of air pollutants. Due to the high dimensionality of geographical covariates, we first ran PCA on these 189 covariates and used the top 15 PCs as predictors for predictive PCA. PCA and RapPCA were applied without this additional step. We examine the overall metrics as well as the individual prediction MSEs for each PC via 10-fold cross-validation following the same spatial prediction approach as Section 4, namely, a random forest model followed by spatial smoothing via TPRS on the residuals (Wai et al., 2020).

Table 1 compares each dimension reduction algorithm in terms of cross-validated TMSE, MSPE and MSRE-trn as well as prediction MSEs for each PC. Consistent with the findings from simulation studies, RapPCA achieves the lowest prediction and overall errors, while PCA has the smallest approximation gap on the training data. The advantage in predictability of RapPCA is reflected in both the overall MSPE and the individual MSEs for each PC. In this particular data analysis task, predictive PCA has good predictive performance for the second PC but is otherwise similar to or outperformed by both classical PCA and RapPCA. More generally, in settings with high-dimensional predictors (covariates and/or spatial basis), predictive PCA is not guaranteed to show an advantage in predictability since it requires a separate processing step on these predictors instead of selecting the information to use for prediction data-adaptively as RapPCA does.

We then run dimension reduction with each of the PCA algorithms on the whole dataset to assess the spatial distribution of top 3 PC scores, where for RapPCA we select the tuning parameters $\gamma,\lambda_{1}$ and $\lambda_{2}$ that lead to the optimal TMSE. Next, we train spatial prediction models for each of the PC scores, and evaluate these models on a finer grid of 5,040 locations over the study region to obtain smoothed plots of each PC score. Figure S4 in Appendix B.3 visualizes the smoothed, finer-grain PC scores obtained by each method across the study region. Figure 3 presents the PC loadings reflecting the contribution of each pollutant on the PC scores. Roughly categorizing various UFP particle sizes into small, moderate and large groups, one key difference we observe from Figure 3 between RapPCA and the two benchmark algorithms is that RapPCA identifies one size group to be the main contributor to each of the 3 PCs, while the other two groups have negligible loadings: small particles contributing to the top PC, moderately-sized particles to the second PC, and large ones to the third PC. In addition, BC is mainly contributing to the second PC according to RapPCA. This observation is consistent with findings from other studies, where aircraft emissions (primarily composed of small UFPs) and road traffic, especially diesel exhaust (comprising moderately-sized UFPs and BC), were identified as two key sources of traffic-related air pollution (Austin et al., 2021; Nie et al., 2022). These findings reflect the benefits in downstream model interpretability brought by dimension reduction that well preserves scientific and spatial information. In contrast, neither predictive nor classical PCA effectively distinguishes the contributions of UFPs with different sizes or the contribution of BC.

In this section, we present another use case of RapPCA by analyzing spatial transcriptomics data. While dimension reduction is also a common data handling step in these applications, there are a variety of downstream modeling tasks where spatial prediction, as in the Seattle TRAP study, may not be the primary focus. We will nevertheless illustrate that optimizing the balance between predictability and representability with RapPCA still leads to desirable performance in these tasks.

We analyze the human epidermal growth factor receptor 2 positive (HER2+) breast tumor data, which includes expression measures of genes in HER2+ tumors from eight individuals (Andersson et al., 2021). We focus on the first tumor section of the last individual, coded as sample H1 in Andersson et al. (2021), where the data consists of expression counts of $15,030$ genes on $613$ tissue locations. Following the same filtering steps as in Shang and Zhou (2022), we removed genes with non-zero expression at less than 20 locations and those confounded by technical artifacts (Andersson et al., 2021), as well as locations with non-zero expressions for less than 20 genes. The filtered set of data contains $p=10,053$ genes measured on $n=607$ spots, which is then normalized via regularized negative binomial regression as implemented by the Seurat R package (Hafemeister and Satija, 2019).

Due to the large number of genes along with the noise present in the expression measurements (Rao et al., 2021), dimension reduction is commonly performed before the main analytic tasks on genomic data to extract biologically meaningful signals as well as to facilitate computation; see Sun et al. (2019); Zhao et al. (2021); Shang and Zhou (2022) for related examples. We focus on two modeling tasks: spatial extrapolation and domain detection. In addition to the three PCA methods (PCA, predictive PCA, RapPCA) that were investigated in previous sections, we also include spatial PCA in our comparison. Spatial PCA is a probabilistic PCA algorithm for spatial transcriptomics data which incorporates spatial proximity information when modeling PC scores (Shang and Zhou, 2022).

The first application is more similar to the analysis in Section 5.1, where we are interested in predicting the gene expression PC scores at tissue locations with no measurements. This prediction problem is involved, for example, in the reconstruction of high-resolution spatial maps of gene expression (Gryglewski et al., 2018; Shang and Zhou, 2022). For each dimension reduction method, we extract the top 20 PCs from the expression of all genes, and conduct 10-fold cross-validation to assess TMSE, MSPE and MSRE-trn, where spatial prediction is done by random forest followed by TPRS smoothing as in previous sections. For spatial PCA, we also include results from its built-in spatial prediction function (using the SpatialPCA R package) for completeness. Table 2 summarizes the overall metrics along with individual prediction MSEs for the top 3 PCs. RapPCA demonstrates significant advantage in overall prediction accuracy (MSPE) compared to all other methods, followed by predictive PCA which also has the lowest per-PC MSEs for the top 3 PCs. Also, consistent with findings from previous sections, RapPCA achieves the optimal trade-off between prediction and approximation gaps, as reflected by TMSE.

Figure 4 visualizes the top 3 PC scores from the expressions of all genes. We observe that spatial PCA and predictive PCA produce spatially smoother surfaces of PC scores compared to the other two methods. This is natural for predictive PCA since it prioritizes spatial predictability of the PC scores, which results in stronger spatial smoothing; for spatial PCA, the prediction-representation balance is implicit and less straightforward to interpret. The gap in the prediction performance of spatial or predictive PCA compared with RapPCA may be the consequence of over-smoothing, which also indicates the effectiveness of explicit optimization for the prediction-representation balance achieved by (4). Figure S5 in Appendix B.3 visualizes the smoothed high-resolution maps of PC scores based on predictions at unmeasured locations.

In the second application, we investigate the problem of domain detection following the dimension reduction step. In domain detection, different sections of tissues are identified as various spatially coherent and functionally distinct regions (Dong and Zhang, 2022; Shang and Zhou, 2022), providing helpful insights into the biological function of tissues. To this end, we conduct domain detection via the walktrap clustering algorithm (Pons and Latapy, 2005) using the top 20 PCs extracted by each algorithm. The number of clusters is set to align with the “ground truth” labels based on pathologist annotations of tissue regions in the original study (Andersson et al., 2021). The hyperparameters $\gamma,\lambda_{1},\lambda_{2}$ are chosen based on silhouette value measuring the difference between nearest-cluster versus intra-cluster distances (Rousseeuw, 1987).

Figure 5 visualizes the detected spatial domains based on PCs extracted from different dimension reduction procedures, compared to the ground truth labels. In general, the relative performance of each algorithm differs across spatial domains. For example, PCA fails to detect the adipose tissue region and produces noisier results for regions surrounding invasive cancer cells. Predictive PCA and RapPCA both mis-classify part of the cancer in situ region as invasive cancer, where predictive PCA results are noisier for the top left region, and RapPCA shows larger uncertainty regarding adipose tissue. Spatial PCA, on the other hand, infers part of the connective tissue to be cancerous.

To make a more in-depth comparison, we examine the breakdown of clustering accuracy by ground truth labels; that is, whether or not each approach correctly classifies the spots belonging to each spatial domain, with the undetermined cluster removed from comparison. Figure 6 presents this breakdown for each method in terms of precision, recall and F1 score. As observed from the spatial domain plots, the relative accuracy for each method differs by region, and in particular, RapPCA has the best performance in both precision and recall on the identification of invasive cancer regions, whereas spatial PCA shows an advantage in recall for cancer in situ.