Code and Data for "Explainable physics-guided deep learning for spatiotemporal groundwater level prediction"

Code and Data for "Explainable physics-guided deep learning for spatiotemporal groundwater level prediction"

This repository contains the Python implementation of a multi-stage machine learning framework designed to predict groundwater levels by coupling temporal deep learning with spatial physical constraints.

Model Hierarchy and Workflow

The research utilizes a two-stage approach to ensure both temporal continuity and spatial physical consistency:

Stage 1: Preliminary Temporal Prediction (ConvGRU)

Core Logic: Uses a Convolutional Gated Recurrent Unit (ConvGRU) with residual connections to process historical groundwater data.

Function: Captures the temporal evolution of the groundwater head to generate an initial estimate for the next time step.

Stage 2: Physics-Informed Refinement (Parallel Architectures)

SegNet-GRU: An encoder-decoder architecture that integrates static geological parameters (K, Ss, Sy, etc.) and dynamic hydrological factors (precipitation, ET, recharge) with the ConvGRU output.

ResNet-GRU: A residual network based on ResNet-18 that uses the ConvGRU prediction as a base and learns the physics-based residuals through deep spatial features.

Comparison: These two models run in parallel to refine the preliminary Stage 1 results. Their outputs are compared against observed data (OBS) to determine the most effective architecture for specific hydrogeological conditions.

File Structure and Descriptions

File
Description

convgru.py
Training script for the Stage 1 model. Includes data normalization and masked MSE loss to handle invalid geological values (-999).

segnet_gru.py
Implementation of the SegNet refinement model. It couples 19 physical channels (static + dynamic) with the ConvGRU prediction.

resnet_gru.py
Implementation of the ResNet-18 refinement model. It uses a residual structure where the final output is the sum of learned features and the initial prediction.

data_comparasion.py
Validation script that extracts results from MODFLOW and all AI models at specific monitoring well coordinates (I, J) for final performance metrics.

Technical Features

1. Physics-Informed Inputs

Both Stage 2 models integrate a comprehensive set of hydrogeological parameters:

Static Grids: BOT1 (Bottom elevation), K (Hydraulic conductivity), Ss (Specific storage), Sy (Specific yield), TOP1 (Top elevation).

Dynamic Grids: Precipitation (PREC), Evapotranspiration (ET), Recharge (RECH), Surface Runoff (SURQ), and Boundary conditions.

2. Explainable AI (XAI)

The framework includes integrated SHAP (SHapley Additive exPlanations) analysis to quantify the contribution of each physical factor:

Global Importance: Bar charts showing which features (e.g., ConvGRU prediction vs. Precipitation) dominate the model output.

Spatial SHAP: Heatmaps showing where specific parameters like Hydraulic Conductivity influence the prediction most strongly across the study area.

3. Partial Dependence Plots (PDP)

resnet_gru.py contains functionality to generate PDPs, allowing researchers to visualize the marginal effect of physical variables on predicted water levels across different hydrogeological zones.