MAC2STI: Mamba network with autoregressive clustering for two-stage spatio-temporal imputation

Soil moisture (SM) data on the Qinghai–Tibet Plateau (QTP) is essential for environmental monitoring but often suffers from incompleteness due to sensor malfunctions or harsh conditions. To address this issue, we propose a two-stage spatio-temporal imputation framework, the Mamba network with autoregressive clustering (MAC2STI), utilizing a benchmark QTP SM dataset for imputation. In stage 1, we introduce an innovative autoregressive model for fine-grained clustering based on similar spatio-temporal patterns, regardless of geographical proximity. This method outperforms traditional clustering techniques that focus solely on temporal or spatial dimensions, providing a more refined understanding of the heterogeneous SM distribution. Temporal features are captured using deep echo state networks, while spatial features are extracted through graph message propagation. A novel bidirectional recurrent neural network-based decoder is then developed to reconstruct the benchmark SM data, with the average hidden states serving as clustering features to segment the QTP into meaningful sub-regions. In stage 2, a Mamba-based imputation architecture is applied, incorporating clustering features into the state transition matrix of selective state space models to ensure precise imputation within each cluster. Experimental results demonstrate that the proposed clustering method effectively captures the spatio-temporal heterogeneity of the QTP SM with excellent interpretability. Comprehensive experiments demonstrate that MAC2STI reduces mean absolute error (MAE) by 62% to 0.16%, achieves a mean absolute percentage error (MAPE) of 0.78% and a root mean squared error (RMSE) of 0.35%. Moreover, it converges at 24.43 iterations/s, which is 40% faster than four baselines—including the previous state-of-the-art method, PriSTI. Finally, ablation studies confirm the indispensable roles of clustering and feature integration, while the imputed dataset faithfully preserves cluster-specific dynamic patterns.