TF-ZIC: Training-Free Zero Inflation Correction

Energy-Efficient Training-Free Zero Inflation Correction for Rainfall Forecasting with Time Series Foundation Models

This repository contains the official implementation of TF-ZIC, a novel, training-free, and energy-efficient framework designed to correct the zero-inflation bias in precipitation forecasts generated by general-purpose Time Series Foundation Models (TSFMs).

📄 Abstract

Time Series Foundation Models (TSFMs) have demonstrated remarkable zero-shot forecasting capabilities. However, when applied to precipitation forecasting, they often struggle with the "zero-inflation" problem—the prevalence of zero values (no rain) in the data—leading to biased forecasts that underestimate rainfall occurrence and intensity.

TF-ZIC addresses this challenge without the need for expensive fine-tuning. It employs a three-stage distribution-free correction mechanism:

1. Zero Calibration: Adjusts the probability of precipitation occurrence using a climatology-aware mixing strategy.
2. Bulk Alignment: Corrects the intensity of non-extreme precipitation using quantile mapping.
3. Tail Correction: Refines extreme value predictions using Generalized Pareto Distribution (GPD) fitting.

This approach significantly improves forecasting accuracy while maintaining the computational efficiency of pre-trained models.

🚀 Features

• Training-Free: No gradient updates or fine-tuning required.
• Model-Agnostic: Compatible with any TSFM (TimesFM, Chronos, Lag-Llama, etc.).
• Energy-Efficient: Minimal computational overhead compared to retraining.
• Modular Design: Easy to integrate into existing forecasting pipelines.

🛠️ Installation

git clone https://github.com/yourusername/TF_ZIC.git
cd TF_ZIC
pip install -e .

To install dependencies for specific foundation models:

pip install -e ".[models]"

💻 Usage

Basic Example

import numpy as np
from tf_zic import TFZICorrector

# 1. Prepare Data
# Historical data (climatology)
climatology = np.random.exponential(scale=2, size=1000)
# Add zeros (zero-inflation)
climatology[np.random.rand(1000) < 0.7] = 0

# 2. Initialize Corrector
corrector = TFZICorrector(
    climatology_data=climatology,
    eta=0.5,  # Weight for climatology vs local window
    window_size=100
)

# 3. Fit (Calibration)
# Recent predictions and observations
recent_preds = np.random.exponential(scale=1.5, size=100)
recent_obs = np.random.exponential(scale=2, size=100)
recent_obs[np.random.rand(100) < 0.7] = 0

corrector.fit(recent_preds, recent_obs)

# 4. Predict (Correction)
# New raw forecast from Foundation Model
raw_forecast = np.array([0.5, 1.2, 0.0, 5.5, 0.1])
corrected_forecast = corrector.predict(raw_forecast)

print("Raw:", raw_forecast)
print("Corrected:", corrected_forecast)

Using with Foundation Models

from tf_zic import TimesFMWrapper, TFZICorrector

# Load Model
model = TimesFMWrapper(repo_id="google/timesfm-1.0-200m-pytorch")

# ... (Load data and fit corrector as above) ...

# Forecast
raw_preds = model.predict(context_data, horizon=24)
corrected_preds = corrector.predict(raw_preds)

📊 Methodology

The TF-ZIC framework operates in three stages:

Stage 1: Zero Calibration

Corrects the binary rain/no-rain decision. It estimates the target zero probability $\pi^*$ by mixing the local zero rate $\pi_{FM}$ with the climatological zero rate $\pi_{clim}$: $\pi^* = \eta \pi_{clim} + (1-\eta) \pi_{FM}$

Stage 2: Bulk Alignment

Adjusts the distribution of positive values using Empirical Quantile Mapping (EQM). This ensures that the bulk of the forecasted distribution matches the observed distribution.

Stage 3: Tail Correction

Models the extreme tail of the distribution using a Generalized Pareto Distribution (GPD). This is crucial for capturing heavy-tailed precipitation events that are often underestimated by standard models.

📂 Project Structure

TF_ZIC/
├── src/
│   └── tf_zic/
│       ├── core/       # Core correction logic (TF-ZIC)
│       └── models/     # Wrappers for Foundation Models
├── examples/           # Usage examples
├── tests/              # Unit tests
├── setup.py            # Installation script
└── README.md           # This file

📝 License

This project is licensed under the MIT License - see the LICENSE file for details.