README.md

---
license: apache-2.0
language:
- en
tags:
- OneScience
- Earth Science
- Data Assimilation
- Mass Conservation
frameworks: PyTorch
---

<p align="center">
  <strong><span style="font-size: 30px;">MassConservingCNN</span></strong>
</p>

# Model Introduction

MassConservingCNN corrects mass-conservation violations introduced by localization in ensemble Kalman filter data assimilation. Given an unconstrained analysis and radar-observation locations, it generates an analysis field with nonnegative rain and improved mass conservation for research on physically constrained data assimilation and analysis postprocessing.

Paper: Training a convolutional neural network to conserve mass in data assimilation  
https://doi.org/10.5194/npg-28-111-2021

# Model Description

MassConservingCNN was proposed by researchers from the Meteorological Institute of Ludwig-Maximilians-Universität München and ClimateAi. The paper trains and validates the model with EnKF unconstrained analyses, QPEns constrained analyses, and radar-observation locations generated by twin experiments with a one-dimensional modified shallow-water model. The model is suitable for mass-conserving data-assimilation correction, rain non-negativity constraints, and physically consistent analysis generation.

# Use Cases

| Use Case | Description |
| :---: | :--- |
| Data-assimilation correction | Predict a QPEns-style analysis from `X^a` and a radar-location indicator. |
| Mass-aware training | Train with the paper Equation 6 error and Equation 7 mass penalty. |
| ModelScope/OneCode execution | Validate training, inference, evaluation, visualization, and checkpoint workflows in ModelScope or OneCode. |
| Multi-GPU training | Launch distributed data-parallel training with `torchrun`. |

# Usage Instructions

## 1.OneCode

Experience intelligent, one-click AI4S programming through the OneCode online environment:

[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

## 2.Download and Installation

```bash
hf download OneScience-Group/MassConservingCNN --local-dir ./MassConservingCNN
cd MassConservingCNN
```

### Environment Dependencies

**Hardware Requirements**

- A GPU or DCU is recommended.
- A CPU can be used for connectivity validation with the default small-sample configuration.
- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.

**DCU Environment**

```bash
# Activate DTK and Conda first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
```

**GPU Environment**

```bash
# Activate Conda first
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
conda activate onescience311
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
```

### Training Data

The default generator creates 24 training samples and 12 validation samples while preserving `[B,4,250]` inputs and `[B,3,250]` targets. The data combine periodic waves, smooth convective cells, nonnegative rain related to velocity convergence, rainy-region radar masks, and smooth EnKF-style errors. They validate the engineering workflow only and are not equivalent to the paper's 48,000-sample QPEns datasets.

```bash
python scripts/fake_data.py --force
```

### Training

For single-device training, use:

```bash
python scripts/train.py
```

For multi-GPU training, use:

```bash
torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
```

Training checks the data version, shapes, dtypes, finite values, binary radar masks, and rain non-negativity. The checkpoint stores model parameters, optimizer state, model configuration, normalization statistics, variable order, data version, `eta`, epoch, and seed. Outputs are written to:

```text
result/checkpoints/massconservingcnn.pt
result/training/metrics.json
```

### Trained Weights

This repository does not include weights under `weight/`. The paper does not provide a confirmed official checkpoint, and the current engineering checkpoint is not claimed to be compatible with external weights.

### Inference

```bash
python scripts/inference.py
```

Inference restores the model parameters produced by training and generates mass-corrected analysis fields from unconstrained analyses and radar-location indicators. Results include input analyses, target analyses, model predictions, radar-observation locations, and the corresponding physical and normalization information, and are saved to:

```text
result/output/predictions.npz
```

### Evaluation and Visualization

```bash
python scripts/result.py
```

Evaluation results include core prediction-error, mass-conservation-error, and relative-improvement results, together with an input, target, and prediction comparison figure. Structured metrics and the auxiliary figure are saved to the paths below; synthetic-data results validate the engineering workflow only and do not represent paper performance.

```text
result/evaluation/metrics.json
result/evaluation/input_target_prediction.png
```

# Official OneScience Information

| Platform | OneScience Main Repository | Skills Repository |
| --- | --- | --- |
| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |

# Citation and License

This repository is an independent engineering reproduction of the public MassConservingCNN specifications.

Use of this repository's code, official model weights, and data remains subject to the licenses and terms of their respective projects.