WIT-SMSR-PUNJ

Interactive soil-moisture maps: WIT Spatial Maps

WIT-SMSR-PUNJ is a PyTorch pipeline for downscaling SMAP surface soil moisture from 9 km to 1 km over Punjab, Pakistan. It supports model training, fine-tuning, evaluation, interpretation, Punjab-wide prediction, and in-situ and spatial validation.

The main GeoTemporalTransformerSR_Residual model combines 1 km spatial predictors and geographic/calendar metadata with a temporal sequence of native SMAP observations. It learns a residual correction to the coarse input, while the loss constrains 9 × 9 block means of the prediction against native 9 km SMAP and penalizes physically invalid values.

Architecture	Role
GeoTemporalTransformerSR_Residual	Main spatial-temporal Transformer
GeoTransformerSR_Residual	Transformer ablation without the temporal stream
SRCNN_Residual	Parameter-comparable residual CNN baseline

All models use the same spatial predictors and per-tile metadata. The configured temporal_eligible cohort provides one-to-one samples for model comparison.

Setup

The code is currently tested with Python 3.14.6 and PyTorch 2.13 in a Conda environment named geo.

git clone https://github.com/LUMS-WIT/WIT-SMSR-PUNJ.git
cd WIT-SMSR-PUNJ

conda create -n geo python=3.14
conda activate geo
conda install -c conda-forge numpy pandas scipy rasterio scikit-image \
  scikit-learn matplotlib netcdf4 pyyaml pyproj pillow tqdm einops \
  shapely cartopy

Install the appropriate CPU, CUDA, or MPS build of PyTorch using the official installation selector.

Data and configuration

Data, checkpoints, normalization statistics, and generated outputs are not tracked by Git. Their locations are set in configs/config.yaml.

training/temporal/forward/
├── train/
│   ├── 1km-r-m-la/    # 11-band 1 km predictors
│   ├── 9km/           # native SMAP
│   └── 1km/           # 1 km reference targets
└── test/
    ├── 1km-r-m-la/
    ├── 9km/
    └── 1km/

training/SMAP9km/      # temporal native-SMAP archive
checkpoints/           # statistics and model checkpoints
Output/                # figures and regional predictions

Matching rasters must preserve date/tile identifiers, CRS, transform, extent, resolution, nodata convention, and band order. The 11 raw predictor bands are encoded into 21 model channels; LULC and HSG are one-hot encoded.

Before running, check device, model.arch, the paths and artifact names, and the temporal window under data.temporal. Statistics are calculated from the training split only and reused when the configured JSON already exists.

Run the pipeline

Run commands from the repository root:

# Train
conda run -n geo python main.py --stage train --config configs/config.yaml

# Fine-tune
conda run -n geo python main.py --stage finetune --config configs/finetune_config.yaml

# Test
conda run -n geo python main.py --stage test --config configs/config.yaml

# Qualitative held-out inference
conda run -n geo python main.py --stage infer --config configs/config.yaml

# Punjab-wide prediction
conda run -n geo python main.py --stage predict --config configs/config.yaml

Testing reports PSNR, MSE, RMSE, bias, ubRMSE, Pearson correlation, sample count, and the coarse-input baseline. For a controlled architecture comparison, keep the data, split, temporal eligibility, statistics, and evaluation settings fixed; change only model.arch and its matching checkpoint.

Additional analyses

Input-corruption ablation:

conda run -n geo python ablation_updated.py --config configs/config.yaml

Quick feature-attribution check:

conda run -n geo python -m SHAP.main \
  --config configs/config.yaml --quick --workflow

The complete WITSMS workflow is configured in configs/validations_config.yaml. Begin by checking station-file matching without making changes:

conda run -n geo python -m sensor_validations.update_id --dry_run

Spatial validation is provided under spatial_validations/. Irrigation and seasonal applications are documented in spatio_temporal_application/README.md. Files named *_legacy.py, the older ablation.py, and Extras/unit_tests/ are not used by the active main.py pipeline.

Reproducibility

• The scale factor is fixed at 9 (9 km to 1 km).
• Train/validation splitting is deterministic and tile-grouped by default.
• Training statistics and the eligible evaluation cohort are shared across comparison models.
• Metrics and exported GeoTIFFs use volumetric soil moisture in m³/m³.

Citation

The article citation and DOI will be added when available.