This repository contains the modularized implementation of PI-RAFT, a physics-informed deep learning model for retrieving Atmospheric Motion Vectors (AMVs) [wind speed and direction vectors estimated by tracking cloud features in sequential satellite images] using geostationary satellite data.
The primary goal of this research project is to estimate high-resolution wind velocity fields directly from sequential satellite images. Rather than treating this purely as an optical flow [the pattern of apparent motion of objects or surfaces in a visual scene] problem, the model incorporates meteorological physical laws (such as fluid dynamics and pressure heights) into the neural network optimization objective to ensure that the retrieved winds are physically realistic.
The model is structured around a recurrent refinement network inspired by RAFT:
- Feature and Context Encoders: Custom convolutional architectures that downsample raw multi-spectral grids to extract spatial representations.
- All-Pairs Correlation Volume: A 4D tensor representing the dot product similarity between features across successive frames.
- ConvGRU [Convolutional Gated Recurrent Unit]: An iterative solver that updates the displacement vector fields over multiple steps.
- Multi-Head Prediction: Outputs horizontal wind vectors
[u, v]and vertical Cloud-Top Pressure (CTP) [the vertical atmospheric pressure level at the top of a cloud, measured in hectopascals (hPa)]. - Physics-Informed Loss: Couples standard data-matching loss with Brightness Constancy [the physical assumption that a moving cloud feature retains its brightness/temperature value over short intervals], spatial gradient constancy, and fluid smoothness regularization.
PI-RAFT-INSAT/
│
├── configs/ # YAML configuration files for train, eval, and experiments
├── datasets/ # Dataset loaders, preprocessing, calibration, and geolocation
├── models/ # Model architectures (encoders, correlation, ConvGRU, pressure head)
├── losses/ # Physics-informed loss functions and warping utilities
├── evaluation/ # Metrics, consistency checks, and visualization scripts
├── export/ # NetCDF4 exporting code
├── scripts/ # Execution entrypoints for training and inference
├── tests/ # Unit and integration tests
├── notebooks/ # Orchestration notebooks for experiments and Colab execution
├── checkpoints/ # Serialized model and optimizer weights
└── outputs/ # Local cache and exported NetCDF files
- Modular model components (
ResidualBlock,FeatureEncoder,ConvGRUCell). - Differentiable spatial warping and multi-constraint physics loss framework.
- CF-compliant NetCDF4 exporter.
- GOES-16 proxy data streaming from public S3 buckets.
- Establishing a local validation pipeline with quantitative metrics (Mean Vector Difference).
- Scaffolding INSAT-3DS HDF5 data loading interfaces.
- Designing temporal frame sequence builders for 4-frame tracking.
- Optimization of the cloud-top pressure head via external reference data.
- Integration of actual topography (Digital Elevation Model) inputs.
- Resolution of the correlation lookup mapping (replacing the random placeholder).
Clone the repository and install it in editable mode:
git clone https://github.com/ice-user/PI-RAFT.git
cd PI-RAFT-INSAT
pip install -r requirements.txt
pip install -e .To start training the model using the default configuration parameters:
python scripts/train.py --config configs/train.yamlTo execute inference, evaluate iteration scaling, and export wind vectors:
python scripts/inference.py --config configs/inference.yamlFor Colab environments, import the orchestration notebook located at baseline.ipynb. It is structured to run setting commands and packages directly, install the repository, and execute model training headless.