README.md

---
datasets:
- OneScience/ERA5
frameworks:
- ""
language:
- en
license: mit
tags:
- OneScience
- Earth Science
- Weather Forecasting
- ERA5
- Regional Forecasting
tasks: []
---
<p align="center">
  <strong>
    <span style="font-size: 30px;">ClimaX</span>
  </strong>
</p>

# Model Overview

ClimaX is the first general-purpose foundation model designed for weather and climate science. It was jointly developed by researchers at Microsoft Research and the University of California, Los Angeles (UCLA), and was presented at ICML 2023.

Paper: *ClimaX: A Foundation Model for Weather and Climate*

https://arxiv.org/abs/2301.10343


# Model Description

ClimaX is built on the Vision Transformer architecture. It was the first model to unify pre-training on heterogeneous climate datasets and adaptation to downstream tasks, such as regional forecasting and climate projections, within a single framework.


# Use Cases

| Use Case | Description |
| :---: | :--- |
| Weather forecasting training | Train ClimaX on ERA5 data in HDF5 format. |
| Quick local validation | Use synthetic data to validate data loading, model training and inference, and visualization of inference results. |
| ModelScope/OneCode execution | Download the standalone model package, install its dependencies, and run the included scripts directly. |
| Multi-GPU training | Launch multi-process training with `torchrun`. |


# Usage

## 1. Using OneCode

Use the OneCode online environment for an intelligent, one-click AI4S development experience:

[Try one-click AI4S development with OneCode](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

## 2. Manual Setup

**Hardware Requirements**

- A GPU or DCU is recommended.
- A CPU can be used for import checks and connectivity validation with a minimal configuration, but full training and inference will be slow.
- DCU users must install DTK in advance. DTK 25.04.2 or later is recommended; alternatively, use the OneScience-recommended version compatible with your cluster.

### Download the Model Package

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

### Set Up the Runtime Environment

**DCU Environment**

```bash
# Activate DTK and conda first.
conda create -n onescience311 python=3.11 -y
conda activate onescience311
# Installation with uv is also supported.
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
# Installation with uv is also supported.
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
```


### Training Data

The OneScience community provides ERA5 data for training. Because of file-size constraints, the repository currently contains a self-contained data slice. Download the data with the following command and ensure that the data path in `conf/config.yaml` is configured correctly:

```bash
hf download --dataset OneScience-Group/ERA5 --local-dir ./data
```

### Training

Single GPU:

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

Multiple GPUs:

```bash
torchrun --nproc_per_node=8 scripts/train.py
```

Training saves the `model_bak.pth` checkpoint under `data/checkpoints/`.

### Pre-trained Weights

This repository will provide weights trained on ERA5 data in the `weights/` directory. The weight files are being prepared and will be uploaded soon.

### Inference

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

Inference results are saved to `result/output/`.

### Evaluation and Visualization

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


# Official OneScience Resources

| 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 a reproduction of the original ClimaX paper and is licensed under the MIT License.