HoqueMahmudul/convective-cloud-denoising-diffusion-probabilistic-model-3d

3D Convective-Cloud Denoising Diffusion Probabilistic Model

An unconditional 3D Denoising Diffusion Probabilistic Model (DDPM) that generates synthetic 64×64×64 convective-cloud cells (vertical-velocity volumes). Roughly 190M parameters.

What it is for: producing synthetic 3D cloud cells to augment scarce atmospheric-simulation data — for example as additional training data for 3D cloud segmentation.

What it is not: this is not a weather-forecasting model and not a physical simulation. It samples volumes resembling its training distribution; it does not solve atmospheric dynamics.

Status: research artifact from work with an associated manuscript in preparation. Methodology, training details, and evaluation are deliberately not described here and will appear in the publication.

Requirements

This is a MONAI model, not a transformers model. You need monai-generative (which provides the generative package) to build the network before loading the weights:

pip install monai-generative==0.2.3 monai==1.5.2 torch safetensors huggingface_hub

A requirements.txt with the verified-working versions is included in this repo:

pip install -r <(curl -sL https://huggingface.co/HoqueMahmudul/convective-cloud-denoising-diffusion-probabilistic-model-3d/raw/main/requirements.txt)

Use monai-generative, not MONAI core. MONAI core (>= 1.4) ships its own DiffusionModelUNet, DDPMScheduler and DiffusionInferer, and they look like drop-in replacements. They are not weight-compatible with this checkpoint.

	monai-generative 0.2.3	MONAI core 1.5
Attention keys	...attentions.0.to_q.weight	...attentions.0.attn.to_q.weight
Constructor arg	num_channels=[...]	channels=[...]
Strict load here	works	52 missing / 52 unexpected keys

Both classes have the identical parameter count (189,815,297), so if you hit the key error and "fix" it with strict=False, you will get a model with 52 randomly-initialised attention tensors. It will run and produce plausible-looking output, and it will be silently wrong.

Never load these weights with strict=False. If a strict load fails, the environment is wrong — fix the environment.

Files

File	What it is
best_model.safetensors	The model weights (a plain PyTorch state_dict)
requirements.txt	Verified-working dependency versions

How to load and generate a sample

The configuration below is required — the weights are a bare state_dict, so the network must be constructed exactly this way before loading.

import torch, numpy as np
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from generative.networks.nets import DiffusionModelUNet
from generative.networks.schedulers import DDPMScheduler
from generative.inferers import DiffusionInferer

REPO = "HoqueMahmudul/convective-cloud-denoising-diffusion-probabilistic-model-3d"
device = "cuda"

# 1. build the network with exactly this configuration
model = DiffusionModelUNet(
    spatial_dims=3, in_channels=1, out_channels=1,
    num_channels=[256, 256, 512],
    attention_levels=[False, False, True],
    num_head_channels=[0, 0, 512],
    num_res_blocks=2,
)

# 2. load the weights -- strict=True is deliberate; see the warning above
weights = hf_hub_download(REPO, "best_model.safetensors")
model.load_state_dict(load_file(weights), strict=True)
model.to(device).eval()

# 3. the sampler must use these settings
scheduler = DDPMScheduler(
    num_train_timesteps=1000, schedule="scaled_linear_beta",
    beta_start=0.0005, beta_end=0.0195,
)
scheduler.set_timesteps(1000)
inferer = DiffusionInferer(scheduler)

# 4. sample one volume (1000 reverse steps -- use a GPU)
noise = torch.randn((1, 1, 64, 64, 64), device=device)
with torch.no_grad(), torch.amp.autocast("cuda"):
    x = inferer.sample(input_noise=noise, diffusion_model=model, scheduler=scheduler)

# 5. convert the [0,1] output back to raw units
#    clip the LOWER bound at 1, not 0: the model's background output is very
#    slightly negative, which would otherwise denormalise to 0 and leave you
#    with a background value the format does not use.
blob = x[0, 0].cpu().numpy()
blob = np.clip(blob * (3100 - 1) + 1, 1, 3100).astype(np.float32)

# 6. enforce the value gap the data format requires (see below)
gap = (blob > 1) & (blob < 300)
blob[gap] = np.where(blob[gap] < 150, 1, 300)

blob = np.round(blob).astype(np.uint16)   # 64^3 synthetic cloud cell
mask = blob >= 300                        # binary cloud mask

Understanding the output

You need these conventions to interpret a generated volume:

Property	Value
Shape / dtype	64 × 64 × 64, uint16
Background	1
Cloud voxels	300 – 3022 (vertical velocity, scaled units)
Value gap	nothing lies strictly between 1 and 300
Voxel size	0.5 km isotropic (so 20 voxels = 10 km)
Normalisation	(raw - 1) / (3100 - 1), clipped to [0, 1]

The model operates in the normalised [0, 1] space, so step 5 above is required to return to raw units, and step 6 restores the value gap that real volumes have.

Post-processing — extracting a single cloud cell

Steps 5 and 6 of the snippet give you a valid volume, but a raw sample often contains several disconnected pieces (or none). To get a usable single cloud cell you need to threshold, find connected components, keep the largest, and blank the rest. Requires scipy:

import numpy as np
from scipy import ndimage

def extract_cloud(blob, min_voxels=6, max_voxels=None):
    """Keep only the largest 26-connected cloud component; blank everything else.

    Returns (cleaned, mask, size).
      cleaned : uint16 volume containing at most one cloud component
      mask    : boolean cloud mask for `cleaned`
      size    : voxel count of the kept component; 0 means "discard this sample"
    """
    mask = blob >= 300
    if not mask.any():
        return blob.copy(), mask, 0                  # empty sample

    # 26-connectivity: faces, edges AND corners count as connected
    labels, n = ndimage.label(mask, structure=np.ones((3, 3, 3)))
    sizes = ndimage.sum(mask, labels, range(1, n + 1))
    keep = int(np.argmax(sizes)) + 1
    largest = int(sizes.max())

    cleaned = blob.copy()
    cleaned[labels != keep] = 1                      # fragments -> background

    if largest < min_voxels:
        return cleaned, cleaned >= 300, 0            # noise only -> discard
    if max_voxels is not None and largest > max_voxels:
        return cleaned, cleaned >= 300, largest      # implausibly large -> review
    return cleaned, cleaned >= 300, largest


cleaned, mask, size = extract_cloud(blob)
if size == 0:
    print("empty or noise-only sample -- generate another")
else:
    print(f"single cloud cell, {size} voxels "
          f"({size * 0.5**3:.1f} km^3 at 0.5 km voxels)")

min_voxels rejects noise-only samples and max_voxels (optional) flags blobs that fill an implausible share of the volume. Both are yours to choose — pick them from the size range your application needs. Expect to discard a large fraction of samples; see the note below.

Practical notes before you use this

• Sampling is variable and low-yield. Only a minority of raw samples are usable single cloud cells; the rest come out empty, as noise fragments, or occasionally as a blob filling the whole volume. Plan on generating a surplus and filtering with extract_cloud() above, discarding samples where it returns size == 0. This is the intended workflow, not a sign of a defect.
• The size distribution is skewed. Generated cells tend to be smaller and weaker than those in the source data, so a raw synthetic set is not a distribution-matched substitute for real data.
• One source dataset. Behaviour outside that simulation regime is unknown.
• Cost. Each sample takes 1000 reverse diffusion steps. A GPU is strongly recommended.

Intended use and limitations

Intended for research and educational use — synthetic data augmentation and generative-modelling experiments on volumetric atmospheric data.

No perceptual or physical-fidelity metrics are published for this model, so no fidelity or accuracy claims are made. Do not use it for operational forecasting, hazard assessment, or any decision-making application.

Data and credit

Trained on convective-cloud cells extracted from atmospheric-simulation output. The dataset is not included or redistributed here.

The underlying simulations come from Dr. Xiaowen Li's lab at Morgan State University; please credit that lab as the source of the simulation data if you use this model. The diffusion model in this repository was developed by Mahmudul Hoque.

License

CC BY-NC 4.0 — Creative Commons Attribution-NonCommercial 4.0 International.

You may use, share, and adapt these weights for non-commercial purposes with attribution. Commercial use is not permitted. Please also credit the source of the simulation data as described above.

References

• Ho, Jain & Abbeel, Denoising Diffusion Probabilistic Models, NeurIPS 2020 (arXiv:2006.11239)
• MONAI GenerativeModels

Contact

Mahmudul Hoque, Morgan State University — mahoq1@morgan.edu