CIUQ

Bayesian neural network for convective initiation forecast and uncertainty estimate.

Overview

This repository contains TensorFlow/TensorFlow Probability pipelines for training and evaluating Bayesian ResNet variants (including MOPED initialization) for convective initiation prediction. The core Bayesian building blocks live in src/ciuq/models/bayesian.py, and the executable experiment scripts live in src/ciuq/pipelines/.

Related paper: https://arxiv.org/abs/2507.16219

Bayesian Deep Learning for Convective Initiation Nowcasting Uncertainty Estimation
Da Fan, David John Gagne II, Steven J. Greybush, Eugene E. Clothiaux, John S. Schreck, Chaopeng Shen

Directory layout

• src/ciuq/models/: Model-building utilities and shared Bayesian helpers.
  • bayesian.py: Shared builders for Bayesian ResNet, deterministic ResNet, and MOPED variants.
• src/ciuq/pipelines/: End-to-end training and inference scripts.
  • Bayesian_ResNet_best_training.py: Train a Bayesian ResNet with ELBO loss.
  • Bayesian_ResNet_best_run.py: Run inference using a Bayesian ResNet checkpoint.
  • Bayesian_ResNet_MOPED_best_training.py: Train a MOPED-initialized Bayesian ResNet.
  • Bayesian_MOPED_ResNet_best_run.py: Run inference using a MOPED Bayesian ResNet.
• src/ciuq/evaluation/: Metrics and evaluation helpers.

Core helpers and their roles

The main shared utilities are defined in src/ciuq/models/bayesian.py:

• get_kernel_posterior_fn: Builds posterior distributions for Flipout layers.
• get_kernel_divergence_fn: Scales KL divergence by the training set size.
• build_elbo_loss: ELBO loss helper that combines data loss and KL loss.
• build_bayesian_resnet: Creates the Flipout-based Bayesian ResNet.
• build_det_resnet: Deterministic ResNet used to seed MOPED priors.
• build_moped_bayesian_resnet: Bayesian ResNet initialized from deterministic weights.

Usage

Typical usage is to run a pipeline script with a lead time argument. For example:

python src/ciuq/pipelines/Bayesian_ResNet_best_training.py --lead_time 10

For inference with a trained Bayesian checkpoint:

python src/ciuq/pipelines/Bayesian_ResNet_best_run.py --lead_time 10

MOPED variants are executed with the analogous MOPED scripts:

python src/ciuq/pipelines/Bayesian_ResNet_MOPED_best_training.py --lead_time 10
python src/ciuq/pipelines/Bayesian_MOPED_ResNet_best_run.py --lead_time 10