Restore model card as README

README.md

@@ -1,278 +1,82 @@
-<div align="center">
-
-# ConvGRU-Ensemble
-
-**Ensemble precipitation nowcasting using Convolutional GRU networks**
-
-*Pretrained model for Italy:* ***IRENE*** — **I**talian **R**adar **E**nsemble **N**owcasting **E**xperiment
-
-[![CI](https://github.com/DSIP-FBK/ConvGRU-Ensemble/actions/workflows/ci.yml/badge.svg)](https://github.com/DSIP-FBK/ConvGRU-Ensemble/actions)
-[![License: BSD-2](https://img.shields.io/badge/license-BSD--2-blue.svg)](LICENSE)
-[![Python 3.13+](https://img.shields.io/badge/python-3.13%2B-blue.svg)](https://python.org)
-[![HuggingFace](https://img.shields.io/badge/%F0%9F%A4%97-Model-yellow)](https://huggingface.co/it4lia/irene)
+---
+license: bsd-2-clause
+language: en
+tags:
+  - weather
+  - nowcasting
+  - radar
+  - precipitation
+  - ensemble-forecasting
+  - convgru
+  - earth-observation
+library_name: pytorch
+pipeline_tag: image-to-image
+---
 
-<br>
+# IRENE — Italian Radar Ensemble Nowcasting Experiment
 
-<a href="https://www.fbk.eu"><img src="https://webvalley.fbk.eu/static/img/logos/fbk-logo-blue.png" height="55" alt="Fondazione Bruno Kessler"></a>
-&nbsp;&nbsp;&nbsp;&nbsp;
-<a href="https://it4lia-aifactory.eu"><img src="https://it4lia-aifactory.eu/wp-content/uploads/2025/05/logo-IT4LIA-AI-factory.svg" height="55" alt="IT4LIA AI-Factory"></a>
-&nbsp;&nbsp;&nbsp;&nbsp;
-<a href="https://www.italiameteo.eu"><img src="https://it4lia-aifactory.eu/wp-content/uploads/2025/08/logo-italiameteo.svg" height="55" alt="ItaliaMeteo"></a>
+**IRENE** is a ConvGRU encoder-decoder model for short-term precipitation forecasting (nowcasting) from radar data. The model generates probabilistic ensemble forecasts, producing multiple plausible future scenarios from a single input sequence.
 
-<br><br>
+## Model Description
 
-The model encodes past radar frames into multi-scale hidden states and decodes them into an **ensemble of probabilistic forecasts** by running the decoder multiple times with different noise inputs, trained with **CRPS loss**.
+- **Architecture**: ConvGRU encoder-decoder with PixelShuffle/PixelUnshuffle for spatial scaling
+- **Input**: Sequence of past radar rain rate fields (T, H, W) in mm/h
+- **Output**: Ensemble of future rain rate forecasts (E, T, H, W) in mm/h
+- **Temporal resolution**: 5 minutes per timestep
+- **Training loss**: Continuous Ranked Probability Score (CRPS) with temporal consistency regularization
 
-</div>
+The model encodes past radar observations into multi-scale hidden states using stacked ConvGRU blocks with PixelUnshuffle downsampling. The decoder generates forecasts by unrolling with different random noise inputs, producing diverse ensemble members that capture forecast uncertainty.
 
----
+## Intended Uses
 
-## Quick Start
+- Short-term precipitation forecasting (0-60 min ahead) from radar reflectivity data
+- Probabilistic nowcasting with uncertainty quantification via ensemble spread
+- Research on deep learning for weather prediction
+- Fine-tuning on regional radar datasets
 
-<details open>
-<summary><b>Load from HuggingFace Hub</b></summary>
+## How to Use
 
 ```python
 from convgru_ensemble import RadarLightningModel
 
+# Load from HuggingFace Hub
 model = RadarLightningModel.from_pretrained("it4lia/irene")
 
+# Run inference on past radar data (rain rate in mm/h)
 import numpy as np
-past = np.load("past_radar.npy")  # rain rate in mm/h, shape (T_past, H, W)
+past = np.random.rand(6, 256, 256).astype(np.float32)  # 6 past timesteps
 forecasts = model.predict(past, forecast_steps=12, ensemble_size=10)
-# forecasts.shape = (10, 12, H, W) — 10 members, 12 future steps, mm/h
-```
-
-</details>
-
-<details>
-<summary><b>CLI Inference</b></summary>
-
-```bash
-convgru-ensemble predict \
-    --input examples/sample_data.nc \
-    --hub-repo it4lia/irene \
-    --forecast-steps 12 \
-    --ensemble-size 10 \
-    --output predictions.nc
-```
-
-</details>
-
-<details>
-<summary><b>Serve via API</b></summary>
-
-```bash
-# With Docker
-docker compose up
-
-# Or directly
-pip install convgru-ensemble[serve]
-convgru-ensemble serve --hub-repo it4lia/irene --port 8000
-```
-
-**Submit a forecast request:**
-
-```bash
-# 4-hour forecast (4 steps × 1h) with 5 ensemble members
-curl -X POST "http://localhost:8000/predict?forecast_steps=4&ensemble_size=5" \
-    -F "file=@examples/sample_data.nc" \
-    -o predictions.nc
-
-# Use default settings (12 steps, 10 members)
-curl -X POST http://localhost:8000/predict \
-    -F "file=@examples/sample_data.nc" -o predictions.nc
-```
-
-**Read the predictions:**
-
-```python
-import xarray as xr
-
-ds = xr.open_dataset("predictions.nc")
-print(ds.precipitation_forecast.shape)
-# (5, 4, 1400, 1200) — ensemble_member, forecast_step, y, x
-```
-
-| Endpoint | Method | Description |
-|---|---|---|
-| `/health` | GET | Health check |
-| `/model/info` | GET | Model metadata and hyperparameters |
-| `/predict` | POST | Upload NetCDF, get ensemble forecast as NetCDF |
-
-**`/predict` query parameters:**
-
-| Parameter | Default | Description |
-|---|---|---|
-| `variable` | `RR` | Name of the rain rate variable in the NetCDF |
-| `forecast_steps` | `12` | Number of future 5-min steps (1–48, i.e. max 4h) |
-| `ensemble_size` | `10` | Number of ensemble members (1–10) |
-
-The input NetCDF must contain a 3D variable `(T, H, W)` with rain rate in mm/h and at least 2 timesteps.
-
-</details>
-
-<details>
-<summary><b>Fine-tune on your data</b></summary>
-
-```bash
-pip install convgru-ensemble
-# See "Training" section below
-```
-
-</details>
-
-## Setup
-
-Requires Python >= 3.13. Uses [uv](https://github.com/astral-sh/uv) for dependency management.
-
-```bash
-uv sync                    # core dependencies
-uv sync --extra serve      # + FastAPI serving
-```
-
-## Data Preparation
-
-The training pipeline expects a Zarr dataset with a rain rate variable `RR` indexed by `(time, x, y)`.
-
-<details>
-<summary><b>1. Filter valid datacubes</b></summary>
-
-Scan the Zarr and find all space-time datacubes with fewer than `n_nan` NaN values:
-
-```bash
-cd importance_sampler
-uv run python filter_nan.py path/to/dataset.zarr \
-    --start_date 2021-01-01 --end_date 2025-12-11 \
-    --Dt 24 --w 256 --h 256 \
-    --step_T 3 --step_X 16 --step_Y 16 \
-    --n_nan 10000 --n_workers 8
+# forecasts.shape = (10, 12, 256, 256) — 10 ensemble members, 12 future steps
 ```
 
-</details>
+## Training Data
 
-<details>
-<summary><b>2. Importance sampling</b></summary>
+Trained on the Italian DPC (Dipartimento della Protezione Civile) radar mosaic surface rain intensity (SRI) dataset, covering the Italian territory at ~1 km resolution with 5-minute temporal resolution.
 
-Sample valid datacubes with higher probability for rainier events:
+## Training Procedure
 
-```bash
-uv run python sample_valid_datacubes.py path/to/dataset.zarr valid_datacubes_*.csv \
-    --q_min 1e-4 --m 0.1 --n_workers 8
-```
-
-A pre-sampled CSV is provided in [`importance_sampler/output/`](importance_sampler/output/).
-
-</details>
-
-## Training
-
-Training is configured via [Fiddle](https://github.com/google/fiddle). Run with defaults:
+- **Optimizer**: Adam (lr=1e-4)
+- **Loss**: CRPS with temporal consistency penalty (lambda=0.01)
+- **Batch size**: 16
+- **Ensemble size during training**: 2 members
+- **Input window**: 6 past timesteps (30 min)
+- **Forecast horizon**: 12 future timesteps (60 min)
+- **Data augmentation**: Random rotations and flips
+- **NaN handling**: Masked loss for missing radar data
 
-```bash
-uv run python -m convgru_ensemble.train
-```
-
-Override parameters from the command line:
-
-```bash
-uv run python -m convgru_ensemble.train \
-    --config config:experiment \
-    --config set:model.num_blocks=5 \
-    --config set:model.forecast_steps=12 \
-    --config set:model.loss_class=crps \
-    --config set:model.ensemble_size=2 \
-    --config set:datamodule.batch_size=16 \
-    --config set:trainer.max_epochs=100
-```
+## Limitations
 
-Monitor with TensorBoard: `uv run tensorboard --logdir logs/`
-
-| Parameter | Description | Default |
-|---|---|---|
-| `model.num_blocks` | Encoder/decoder depth | `5` |
-| `model.forecast_steps` | Future steps to predict | `12` |
-| `model.ensemble_size` | Ensemble members during training | `2` |
-| `model.loss_class` | Loss function (`mse`, `mae`, `crps`, `afcrps`) | `crps` |
-| `model.masked_loss` | Mask NaN regions in loss | `True` |
-| `datamodule.steps` | Total timesteps per sample (past + future) | `18` |
-| `datamodule.batch_size` | Batch size | `16` |
-
-## Architecture
-
-```
-Input (B, T_past, 1, H, W)
-    |
-    v
-+--------------------------+
-|        Encoder           |  ConvGRU + PixelUnshuffle (x num_blocks)
-|  Spatial dims halve at   |  Channels: 1 -> 4 -> 16 -> 64 -> 256 -> 1024
-|  each block              |
-+----------+---------------+
-           | hidden states
-           v
-+--------------------------+
-|        Decoder           |  ConvGRU + PixelShuffle (x num_blocks)
-|  Noise input (x M runs)  |  Each run produces one ensemble member
-|  for ensemble generation |
-+----------+---------------+
-           |
-           v
-Output (B, T_future, M, H, W)
-```
-
-## Docker
-
-```bash
-docker build -t convgru-ensemble .
-
-# Run with local checkpoint
-docker run -p 8000:8000 -v ./checkpoints:/app/checkpoints \
-    -e MODEL_CHECKPOINT=/app/checkpoints/model.ckpt convgru-ensemble
-
-# Run with HuggingFace Hub
-docker run -p 8000:8000 -e HF_REPO_ID=it4lia/irene convgru-ensemble
-```
-
-## Project Structure
-
-```
-ConvGRU-Ensemble/
-+-- convgru_ensemble/          # Python package
-|   +-- model.py               # ConvGRU encoder-decoder architecture
-|   +-- losses.py              # CRPS, afCRPS, masked loss wrappers
-|   +-- lightning_model.py     # PyTorch Lightning training module
-|   +-- datamodule.py          # Dataset and data loading
-|   +-- train.py               # Training entry point (Fiddle config)
-|   +-- utils.py               # Rain rate <-> reflectivity conversions
-|   +-- hub.py                 # HuggingFace Hub upload/download
-|   +-- cli.py                 # CLI for inference and serving
-|   +-- serve.py               # FastAPI inference server
-+-- examples/                  # Sample data for testing
-+-- importance_sampler/        # Data preparation scripts
-+-- notebooks/                 # Example notebooks
-+-- scripts/                   # Utility scripts (e.g., upload to Hub)
-+-- tests/                     # Test suite
-+-- Dockerfile                 # Container for serving API
-+-- MODEL_CARD.md              # HuggingFace model card template
-```
+- Trained on Italian radar data; performance may degrade on other domains without fine-tuning
+- 5-minute temporal resolution only
+- Best suited for convective and stratiform precipitation; extreme events may be underrepresented
+- Ensemble spread is generated via noisy decoder inputs, not a full Bayesian approach
 
 ## Acknowledgements
 
-<div align="center">
-
-Developed at **Fondazione Bruno Kessler (FBK)**, Trento, Italy, as part of the **Italian AI-Factory (IT4LIA)**, an EU-funded initiative supporting AI adoption across SMEs, academia, and public/private sectors. This work showcases capabilities in the **Earth (weather and climate) vertical domain**.
-
-<br>
-
-<a href="https://www.fbk.eu"><img src="https://webvalley.fbk.eu/static/img/logos/fbk-logo-blue.png" height="45" alt="Fondazione Bruno Kessler"></a>
-&nbsp;&nbsp;&nbsp;&nbsp;
-<a href="https://it4lia-aifactory.eu"><img src="https://it4lia-aifactory.eu/wp-content/uploads/2025/05/logo-IT4LIA-AI-factory.svg" height="45" alt="IT4LIA AI-Factory"></a>
-&nbsp;&nbsp;&nbsp;&nbsp;
-<a href="https://www.italiameteo.eu"><img src="https://it4lia-aifactory.eu/wp-content/uploads/2025/08/logo-italiameteo.svg" height="45" alt="ItaliaMeteo"></a>
+This model was developed as part of the **Italian AI-Factory** (IT4LIA), an EU-funded initiative supporting the adoption of AI across SMEs, academia, and public/private sectors. The AI-Factory provides free HPC compute, consultancy, and AI-ready datasets. This work showcases capabilities in the **Earth (weather and climate) vertical domain**.
 
-</div>
+Developed at **Fondazione Bruno Kessler (FBK)**, Trento, Italy.
 
 ## License
 
-BSD 2-Clause — see [LICENSE](LICENSE).
+BSD 2-Clause License