Upload folder using huggingface_hub
Browse files- .gitattributes +1 -34
- README.md +149 -0
- conf/config.yaml +30 -0
- config.json +34 -0
- configuration.json +12 -0
- model/climemu_s2l.py +168 -0
- scripts/fake_data.py +88 -0
- scripts/inference.py +60 -0
- scripts/result.py +103 -0
- scripts/train.py +130 -0
- weight/.gitkeep +0 -0
.gitattributes
CHANGED
|
@@ -1,35 +1,2 @@
|
|
| 1 |
-
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
-
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
-
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
-
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
-
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
-
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
-
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
-
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
-
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
-
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
-
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
-
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
-
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
-
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
-
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
-
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
-
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
-
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
-
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
-
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
-
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
-
*.
|
| 24 |
-
*.rar filter=lfs diff=lfs merge=lfs -text
|
| 25 |
-
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
-
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
| 27 |
-
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
| 28 |
-
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
-
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
-
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
-
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
-
*.xz filter=lfs diff=lfs merge=lfs -text
|
| 33 |
-
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
-
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
-
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 2 |
+
*.npz filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
README.md
ADDED
|
@@ -0,0 +1,149 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: cc-by-4.0
|
| 3 |
+
language:
|
| 4 |
+
- en
|
| 5 |
+
tags:
|
| 6 |
+
- OneScience
|
| 7 |
+
- Earth Science
|
| 8 |
+
- Climate Prediction
|
| 9 |
+
- Regression
|
| 10 |
+
frameworks: PyTorch
|
| 11 |
+
---
|
| 12 |
+
|
| 13 |
+
<p align="center">
|
| 14 |
+
<strong><span style="font-size: 30px;">ClimEmu-S2L</span></strong>
|
| 15 |
+
</p>
|
| 16 |
+
|
| 17 |
+
# Model Introduction
|
| 18 |
+
|
| 19 |
+
ClimEmu-S2L predicts long-term climate-change patterns from the first 10 years of global near-surface temperature responses to reduce the cost of multi-scenario long-term climate simulations and support regional temperature-response analysis.
|
| 20 |
+
|
| 21 |
+
Paper: Predicting global patterns of long-term climate change from short-term simulations using machine learning
|
| 22 |
+
https://doi.org/10.1038/s41612-020-00148-5
|
| 23 |
+
|
| 24 |
+
# Model Description
|
| 25 |
+
|
| 26 |
+
The method reproduced by ClimEmu-S2L was proposed by research teams from Imperial College London, the University of Reading, the University of East Anglia, the University of Warwick, the Technical University of Crete, and collaborating institutions. The paper uses 21 HadGEM3 climate-forcing scenarios from projects including PDRMIP, ECLIPSE, and Kasoar. The model performs the task of predicting long-term spatial patterns from short-term global near-surface temperature responses and evaluates global and regional climate responses.
|
| 27 |
+
|
| 28 |
+
# Use Cases
|
| 29 |
+
|
| 30 |
+
| Use Case | Description |
|
| 31 |
+
| :---: | :--- |
|
| 32 |
+
| Long-term climate response prediction | Predict global long-term patterns after year 70 from the first 10 years of temperature response. |
|
| 33 |
+
| Ridge regression validation | Select regularization with inner cross-validation and execute 21-fold LOSO prediction. |
|
| 34 |
+
| GPR regression validation | Execute 21-fold LOSO prediction with shared-kernel Gaussian process regression. |
|
| 35 |
+
| ModelScope/OneCode execution | Execute full-grid training, inference, regional evaluation, and visualization in ModelScope or OneCode environments. |
|
| 36 |
+
| Multi-GPU training | Distribute the 21 LOSO folds across processes with `torchrun`. |
|
| 37 |
+
|
| 38 |
+
# Usage Instructions
|
| 39 |
+
|
| 40 |
+
## 1.OneCode
|
| 41 |
+
|
| 42 |
+
Experience intelligent, one-click AI4S programming through the OneCode online environment:
|
| 43 |
+
|
| 44 |
+
[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
|
| 45 |
+
|
| 46 |
+
## 2. Download and Installation
|
| 47 |
+
|
| 48 |
+
```bash
|
| 49 |
+
hf download OneScience-Group/ClimEmu-S2L --local-dir ./ClimEmu-S2L
|
| 50 |
+
cd ClimEmu-S2L
|
| 51 |
+
```
|
| 52 |
+
|
| 53 |
+
### Environment Dependencies
|
| 54 |
+
|
| 55 |
+
**Hardware Requirements**
|
| 56 |
+
|
| 57 |
+
- A GPU or DCU is recommended.
|
| 58 |
+
- A CPU can be used for connectivity validation with the default small-sample configuration.
|
| 59 |
+
- DCU users must install DTK first. DTK 25.04.2 or later, or the OneScience-recommended version matching the cluster, is recommended.
|
| 60 |
+
|
| 61 |
+
**DCU Environment**
|
| 62 |
+
|
| 63 |
+
```bash
|
| 64 |
+
# Activate DTK and Conda first
|
| 65 |
+
conda create -n onescience311 python=3.11 -y
|
| 66 |
+
conda activate onescience311
|
| 67 |
+
pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
|
| 68 |
+
```
|
| 69 |
+
|
| 70 |
+
**GPU Environment**
|
| 71 |
+
|
| 72 |
+
```bash
|
| 73 |
+
# Activate Conda first
|
| 74 |
+
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
|
| 75 |
+
conda activate onescience311
|
| 76 |
+
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
|
| 77 |
+
```
|
| 78 |
+
|
| 79 |
+
### Training Data
|
| 80 |
+
|
| 81 |
+
This repository uses a small set of structured synthetic climate-response samples to validate the engineering workflow, comprising 21 forcing scenarios of single-channel near-surface temperature anomalies on the complete `145×192` grid in `degC`. Inputs are the mean responses during the first 10 years of each scenario, and targets are the same-scenario long-term mean responses after year 70. The synthetic data preserve the scenario count, spatial grid, and time windows while reducing GPR to a shared non-ARD kernel structure and fewer kernel-optimization iterations to control engineering-validation cost; they do not represent the official HadGEM3 data distribution or training scale.
|
| 82 |
+
|
| 83 |
+
```bash
|
| 84 |
+
python scripts/fake_data.py
|
| 85 |
+
```
|
| 86 |
+
|
| 87 |
+
### Training
|
| 88 |
+
|
| 89 |
+
For single-device training, use:
|
| 90 |
+
|
| 91 |
+
```bash
|
| 92 |
+
python scripts/train.py
|
| 93 |
+
```
|
| 94 |
+
|
| 95 |
+
For multi-GPU training, use:
|
| 96 |
+
|
| 97 |
+
```bash
|
| 98 |
+
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
|
| 99 |
+
```
|
| 100 |
+
|
| 101 |
+
Training fits Ridge and GPR models over all 21 leave-one-scenario-out (LOSO) folds; Ridge selects regularization through inner three-fold cross-validation, while GPR optimizes a shared non-ARD kernel. Standard training artifacts are saved to:
|
| 102 |
+
|
| 103 |
+
```text
|
| 104 |
+
result/checkpoints/climemu_s2l.pt
|
| 105 |
+
result/training/metrics.json
|
| 106 |
+
```
|
| 107 |
+
|
| 108 |
+
### Trained Weights
|
| 109 |
+
|
| 110 |
+
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 contain the paper's formal model parameters.
|
| 111 |
+
|
| 112 |
+
### Inference
|
| 113 |
+
|
| 114 |
+
```bash
|
| 115 |
+
python scripts/inference.py
|
| 116 |
+
```
|
| 117 |
+
|
| 118 |
+
Inference loads the training checkpoint and generates Ridge and GPR long-term response fields for all 21 scenarios. Complete numerical results are saved to:
|
| 119 |
+
|
| 120 |
+
```text
|
| 121 |
+
result/output/predictions.npz
|
| 122 |
+
```
|
| 123 |
+
|
| 124 |
+
### Evaluation and Visualization
|
| 125 |
+
|
| 126 |
+
```bash
|
| 127 |
+
python scripts/result.py
|
| 128 |
+
```
|
| 129 |
+
|
| 130 |
+
Evaluation reports area-weighted full-grid RMSE, global mean absolute error, and regional mean absolute error by method, per scenario, and per region; this is neither a multi-step nor classification task, so no per-step or per-class results are provided. Visualizations include target and prediction spatial maps and cross-scenario error boxplots, with structured results and figures saved to the paths below. Synthetic-data results validate the engineering workflow only and do not represent formal paper performance.
|
| 131 |
+
|
| 132 |
+
```text
|
| 133 |
+
result/evaluation/metrics.json
|
| 134 |
+
result/evaluation/spatial_fields.png
|
| 135 |
+
result/evaluation/error_boxplots.png
|
| 136 |
+
```
|
| 137 |
+
|
| 138 |
+
# Official OneScience Information
|
| 139 |
+
|
| 140 |
+
| Platform | OneScience Main Repository | Skills Repository |
|
| 141 |
+
| --- | --- | --- |
|
| 142 |
+
| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
|
| 143 |
+
| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
|
| 144 |
+
|
| 145 |
+
# Citation and License
|
| 146 |
+
|
| 147 |
+
This repository is an independent engineering reproduction of the public ClimEmu-S2L specifications.
|
| 148 |
+
|
| 149 |
+
Use of this repository's code, official model weights, and data remains subject to the licenses and terms of their respective projects.
|
conf/config.yaml
ADDED
|
@@ -0,0 +1,30 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
seed: 2020
|
| 2 |
+
data:
|
| 3 |
+
path: data/hadgem3_fake.npz
|
| 4 |
+
scenario_count: 21
|
| 5 |
+
grid: [145, 192]
|
| 6 |
+
channels: 1
|
| 7 |
+
model:
|
| 8 |
+
methods: [ridge, gpr]
|
| 9 |
+
ridge:
|
| 10 |
+
alphas: [0.001, 0.01, 0.1, 1.0, 10.0, 100.0]
|
| 11 |
+
inner_folds: 3
|
| 12 |
+
gpr:
|
| 13 |
+
kernel_mode: paper_rbf
|
| 14 |
+
supported_kernel_modes: [paper_rbf, official_rbf_linear]
|
| 15 |
+
optimizer_steps: 8
|
| 16 |
+
learning_rate: 0.08
|
| 17 |
+
jitter: 0.00001
|
| 18 |
+
ard: false
|
| 19 |
+
ard_gap: "The paper/official GPy workflow can use per-input ARD; this implementation deliberately uses one shared non-ARD lengthscale because p=27840 and N=20."
|
| 20 |
+
runtime:
|
| 21 |
+
device: cpu
|
| 22 |
+
dtype: float64
|
| 23 |
+
ddp_backend: gloo
|
| 24 |
+
paths:
|
| 25 |
+
checkpoint: result/checkpoints/climemu_s2l.pt
|
| 26 |
+
training_metrics: result/training/metrics.json
|
| 27 |
+
inference: result/output/predictions.npz
|
| 28 |
+
evaluation_dir: result/evaluation
|
| 29 |
+
evaluation:
|
| 30 |
+
plot_scenario_id: PDRMIP_2xCO2
|
config.json
ADDED
|
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_name": "ClimEmu-S2L",
|
| 3 |
+
"model_type": "climemu_s2l",
|
| 4 |
+
"architectures": ["DualRidge", "SharedKernelGPR"],
|
| 5 |
+
"framework": "PyTorch",
|
| 6 |
+
"domain": "earth-science",
|
| 7 |
+
"task": "short-to-long-term-climate-response",
|
| 8 |
+
"implementation": {
|
| 9 |
+
"entry_point": "model/climemu_s2l.py",
|
| 10 |
+
"train_script": "scripts/train.py",
|
| 11 |
+
"inference_script": "scripts/inference.py",
|
| 12 |
+
"evaluation_script": "scripts/result.py",
|
| 13 |
+
"synthetic_data_script": "scripts/fake_data.py"
|
| 14 |
+
},
|
| 15 |
+
"data": {
|
| 16 |
+
"scenarios": 21,
|
| 17 |
+
"grid": [145, 192],
|
| 18 |
+
"features_per_scenario": 27840,
|
| 19 |
+
"input_channels": 1,
|
| 20 |
+
"outer_validation": "21-fold leave-one-scenario-out"
|
| 21 |
+
},
|
| 22 |
+
"methods": {
|
| 23 |
+
"ridge": "multi-output sample-space dual closed form with internal 3-fold alpha selection",
|
| 24 |
+
"gpr": "multi-output posterior mean with one shared non-ARD kernel and optimized marginal likelihood",
|
| 25 |
+
"kernel_modes": ["paper_rbf", "official_rbf_linear"]
|
| 26 |
+
},
|
| 27 |
+
"paper": {
|
| 28 |
+
"title": "Predicting global patterns of long-term climate change from short-term simulations using machine learning",
|
| 29 |
+
"doi": "10.1038/s41612-020-00148-5",
|
| 30 |
+
"data_doi": "10.5281/zenodo.3971024",
|
| 31 |
+
"license": "CC BY 4.0"
|
| 32 |
+
},
|
| 33 |
+
"configuration_sources": ["conf/config.yaml", "model/climemu_s2l.py", "scripts/fake_data.py", "scripts/train.py", "scripts/inference.py", "scripts/result.py"]
|
| 34 |
+
}
|
configuration.json
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"framework": "PyTorch",
|
| 3 |
+
"task": "short-to-long-term-climate-response",
|
| 4 |
+
"model": "ClimEmu-S2L",
|
| 5 |
+
"input_format": "NPZ with short_response and long_response shaped [21,145,192]",
|
| 6 |
+
"protocol": "full-grid 21-fold LOSO dual Ridge and shared-kernel GPR",
|
| 7 |
+
"default_config": "conf/config.yaml",
|
| 8 |
+
"training": "scripts/train.py",
|
| 9 |
+
"inference": "scripts/inference.py",
|
| 10 |
+
"evaluation": "scripts/result.py",
|
| 11 |
+
"visualization": "scripts/result.py"
|
| 12 |
+
}
|
model/climemu_s2l.py
ADDED
|
@@ -0,0 +1,168 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Full-grid sample-space regressors for HadGEM3 short-to-long responses."""
|
| 2 |
+
|
| 3 |
+
from dataclasses import dataclass
|
| 4 |
+
|
| 5 |
+
import numpy as np
|
| 6 |
+
import torch
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
MODEL_NAME = "ClimEmu-S2L"
|
| 10 |
+
FORMAT_VERSION = "climemu_s2l_v1"
|
| 11 |
+
GRID_SHAPE = (145, 192)
|
| 12 |
+
SCENARIO_COUNT = 21
|
| 13 |
+
FEATURE_COUNT = 27840
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
def area_weights(latitude_deg, nx=192):
|
| 17 |
+
weights = np.cos(np.deg2rad(np.asarray(latitude_deg, dtype=np.float64)))[:, None]
|
| 18 |
+
weights = np.broadcast_to(np.maximum(weights, 0.0), (len(latitude_deg), nx)).copy()
|
| 19 |
+
return weights / weights.sum()
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def weighted_rmse(target, prediction, weights):
|
| 23 |
+
residual = np.asarray(target) - np.asarray(prediction)
|
| 24 |
+
return float(np.sqrt(np.sum(weights * residual.reshape(weights.shape) ** 2)))
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def _normalise_inputs(x):
|
| 28 |
+
mean = x.mean(dim=0, keepdim=True)
|
| 29 |
+
scale = torch.sqrt(torch.mean((x - mean) ** 2)).clamp_min(1e-8)
|
| 30 |
+
return (x - mean) / scale, mean, scale
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
@dataclass
|
| 34 |
+
class DualRidge:
|
| 35 |
+
alpha: float
|
| 36 |
+
x_train: torch.Tensor = None
|
| 37 |
+
dual_coef: torch.Tensor = None
|
| 38 |
+
x_mean: torch.Tensor = None
|
| 39 |
+
x_scale: torch.Tensor = None
|
| 40 |
+
y_mean: torch.Tensor = None
|
| 41 |
+
|
| 42 |
+
def fit(self, x, y):
|
| 43 |
+
x = torch.as_tensor(x)
|
| 44 |
+
y = torch.as_tensor(y, dtype=x.dtype, device=x.device)
|
| 45 |
+
x_scaled, self.x_mean, self.x_scale = _normalise_inputs(x)
|
| 46 |
+
self.y_mean = y.mean(dim=0, keepdim=True)
|
| 47 |
+
yc = y - self.y_mean
|
| 48 |
+
gram = x_scaled @ x_scaled.T
|
| 49 |
+
system = gram + self.alpha * torch.eye(len(x), dtype=x.dtype, device=x.device)
|
| 50 |
+
self.x_train = x_scaled
|
| 51 |
+
self.dual_coef = torch.linalg.solve(system, yc)
|
| 52 |
+
return self
|
| 53 |
+
|
| 54 |
+
def predict(self, x):
|
| 55 |
+
x = torch.as_tensor(x, dtype=self.x_train.dtype, device=self.x_train.device)
|
| 56 |
+
return ((x - self.x_mean) / self.x_scale) @ self.x_train.T @ self.dual_coef + self.y_mean
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def select_ridge_alpha(x, y, alphas, folds, weights):
|
| 60 |
+
"""Paper-style internal CV using area-weighted grid RMSE."""
|
| 61 |
+
x = torch.as_tensor(x)
|
| 62 |
+
y = torch.as_tensor(y, dtype=x.dtype)
|
| 63 |
+
split = np.array_split(np.arange(len(x)), folds)
|
| 64 |
+
scores = {}
|
| 65 |
+
flat_weights = torch.as_tensor(weights.reshape(-1), dtype=x.dtype)
|
| 66 |
+
for alpha in alphas:
|
| 67 |
+
fold_scores = []
|
| 68 |
+
for validation in split:
|
| 69 |
+
training = np.setdiff1d(np.arange(len(x)), validation)
|
| 70 |
+
prediction = DualRidge(float(alpha)).fit(x[training], y[training]).predict(x[validation])
|
| 71 |
+
mse = torch.sum((prediction - y[validation]) ** 2 * flat_weights, dim=1)
|
| 72 |
+
fold_scores.append(float(torch.sqrt(mse).mean()))
|
| 73 |
+
scores[float(alpha)] = float(np.mean(fold_scores))
|
| 74 |
+
selected = min(scores, key=scores.get)
|
| 75 |
+
return selected, scores
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
class SharedKernelGPR(torch.nn.Module):
|
| 79 |
+
"""One scenario-space covariance shared by all 27,840 output cells."""
|
| 80 |
+
|
| 81 |
+
def __init__(self, kernel_mode="paper_rbf", jitter=1e-5, dtype=torch.float64):
|
| 82 |
+
super().__init__()
|
| 83 |
+
if kernel_mode not in {"paper_rbf", "official_rbf_linear"}:
|
| 84 |
+
raise ValueError(f"unknown GPR kernel mode: {kernel_mode}")
|
| 85 |
+
self.kernel_mode = kernel_mode
|
| 86 |
+
self.jitter = float(jitter)
|
| 87 |
+
self.log_lengthscale = torch.nn.Parameter(torch.tensor(0.0, dtype=dtype))
|
| 88 |
+
self.log_outputscale = torch.nn.Parameter(torch.tensor(0.0, dtype=dtype))
|
| 89 |
+
self.log_noise = torch.nn.Parameter(torch.tensor(-2.5, dtype=dtype))
|
| 90 |
+
self.log_linear_scale = torch.nn.Parameter(torch.tensor(-1.0, dtype=dtype))
|
| 91 |
+
|
| 92 |
+
@staticmethod
|
| 93 |
+
def geometry(x1, x2):
|
| 94 |
+
feature_count = x1.shape[1]
|
| 95 |
+
distance = ((x1[:, None, :] - x2[None, :, :]) ** 2).mean(dim=2)
|
| 96 |
+
linear = (x1 @ x2.T) / feature_count
|
| 97 |
+
return distance, linear
|
| 98 |
+
|
| 99 |
+
def kernel_from_geometry(self, distance, linear):
|
| 100 |
+
length2 = torch.exp(2.0 * self.log_lengthscale).clamp_min(1e-10)
|
| 101 |
+
kernel = torch.exp(2.0 * self.log_outputscale) * torch.exp(-0.5 * distance / length2)
|
| 102 |
+
if self.kernel_mode == "official_rbf_linear":
|
| 103 |
+
kernel = kernel + torch.exp(2.0 * self.log_linear_scale) * linear
|
| 104 |
+
return kernel
|
| 105 |
+
|
| 106 |
+
def fit(self, x, y, steps=8, learning_rate=0.08):
|
| 107 |
+
x = torch.as_tensor(x, dtype=self.log_lengthscale.dtype)
|
| 108 |
+
y = torch.as_tensor(y, dtype=x.dtype)
|
| 109 |
+
self.x_train, self.x_mean, self.x_scale = _normalise_inputs(x)
|
| 110 |
+
self.y_mean = y.mean(dim=0, keepdim=True)
|
| 111 |
+
self.y_scale = torch.sqrt(torch.mean((y - self.y_mean) ** 2)).clamp_min(1e-8)
|
| 112 |
+
self.y_train = (y - self.y_mean) / self.y_scale
|
| 113 |
+
distance, linear = self.geometry(self.x_train, self.x_train)
|
| 114 |
+
identity = torch.eye(len(x), dtype=x.dtype)
|
| 115 |
+
optimizer = torch.optim.Adam(self.parameters(), lr=float(learning_rate))
|
| 116 |
+
initial = {name: value.detach().clone() for name, value in self.named_parameters()}
|
| 117 |
+
losses, gradient_seen = [], False
|
| 118 |
+
for _ in range(int(steps)):
|
| 119 |
+
optimizer.zero_grad()
|
| 120 |
+
kernel = self.kernel_from_geometry(distance, linear)
|
| 121 |
+
noise = torch.exp(2.0 * self.log_noise) + self.jitter
|
| 122 |
+
chol = torch.linalg.cholesky(kernel + noise * identity)
|
| 123 |
+
solved = torch.cholesky_solve(self.y_train, chol)
|
| 124 |
+
data_fit = 0.5 * torch.mean(torch.sum(self.y_train * solved, dim=0))
|
| 125 |
+
loss = data_fit + torch.log(torch.diagonal(chol)).sum()
|
| 126 |
+
loss.backward()
|
| 127 |
+
gradient_seen = gradient_seen or any(p.grad is not None and torch.isfinite(p.grad).all() for p in self.parameters())
|
| 128 |
+
torch.nn.utils.clip_grad_norm_(self.parameters(), 20.0)
|
| 129 |
+
optimizer.step()
|
| 130 |
+
losses.append(float(loss.detach()))
|
| 131 |
+
parameter_updated = any(not torch.equal(initial[name], value.detach()) for name, value in self.named_parameters())
|
| 132 |
+
if not gradient_seen or not parameter_updated:
|
| 133 |
+
raise RuntimeError("GPR marginal-likelihood backward/update check failed")
|
| 134 |
+
with torch.no_grad():
|
| 135 |
+
kernel = self.kernel_from_geometry(distance, linear)
|
| 136 |
+
noise = torch.exp(2.0 * self.log_noise) + self.jitter
|
| 137 |
+
self.posterior_coef = torch.linalg.solve(kernel + noise * identity, self.y_train)
|
| 138 |
+
return {"loss": losses, "gradient_seen": gradient_seen, "parameter_updated": parameter_updated}
|
| 139 |
+
|
| 140 |
+
def predict(self, x):
|
| 141 |
+
x = torch.as_tensor(x, dtype=self.x_train.dtype)
|
| 142 |
+
scaled = (x - self.x_mean) / self.x_scale
|
| 143 |
+
distance, linear = self.geometry(scaled, self.x_train)
|
| 144 |
+
kernel = self.kernel_from_geometry(distance, linear)
|
| 145 |
+
return self.y_mean + self.y_scale * (kernel @ self.posterior_coef)
|
| 146 |
+
|
| 147 |
+
def hyperparameters(self):
|
| 148 |
+
return {name: float(value.detach()) for name, value in self.named_parameters()}
|
| 149 |
+
|
| 150 |
+
def load_hyperparameters(self, values):
|
| 151 |
+
with torch.no_grad():
|
| 152 |
+
for name, value in values.items():
|
| 153 |
+
getattr(self, name).copy_(torch.tensor(value, dtype=getattr(self, name).dtype))
|
| 154 |
+
return self
|
| 155 |
+
|
| 156 |
+
def restore_posterior(self, x, y):
|
| 157 |
+
x = torch.as_tensor(x, dtype=self.log_lengthscale.dtype)
|
| 158 |
+
y = torch.as_tensor(y, dtype=x.dtype)
|
| 159 |
+
self.x_train, self.x_mean, self.x_scale = _normalise_inputs(x)
|
| 160 |
+
self.y_mean = y.mean(dim=0, keepdim=True)
|
| 161 |
+
self.y_scale = torch.sqrt(torch.mean((y - self.y_mean) ** 2)).clamp_min(1e-8)
|
| 162 |
+
self.y_train = (y - self.y_mean) / self.y_scale
|
| 163 |
+
distance, linear = self.geometry(self.x_train, self.x_train)
|
| 164 |
+
kernel = self.kernel_from_geometry(distance, linear)
|
| 165 |
+
identity = torch.eye(len(x), dtype=x.dtype)
|
| 166 |
+
noise = torch.exp(2.0 * self.log_noise) + self.jitter
|
| 167 |
+
self.posterior_coef = torch.linalg.solve(kernel + noise * identity, self.y_train)
|
| 168 |
+
return self
|
scripts/fake_data.py
ADDED
|
@@ -0,0 +1,88 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Generate 21 structured, physically interpretable full-grid response pairs."""
|
| 2 |
+
|
| 3 |
+
import sys
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
import numpy as np
|
| 7 |
+
import yaml
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 11 |
+
sys.path.insert(0, str(ROOT))
|
| 12 |
+
from model.climemu_s2l import FORMAT_VERSION, GRID_SHAPE
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
SCENARIOS = [
|
| 16 |
+
("PDRMIP_2xCO2", 2.5, "global", 0, 0),
|
| 17 |
+
("PDRMIP_3xCH4", 1.5, "global", 0, 0),
|
| 18 |
+
("PDRMIP_10xCFC12", 1.1, "global", 0, 0),
|
| 19 |
+
("PDRMIP_solar", 1.0, "global", 0, 0),
|
| 20 |
+
("PDRMIP_5xSO4", -1.6, "global", 0, 0),
|
| 21 |
+
("PDRMIP_10xBC", 1.2, "global", 0, 0),
|
| 22 |
+
("PDRMIP_10xSO4_Europe", -1.0, "regional", 52, 15),
|
| 23 |
+
("PDRMIP_10xSO4_Asia", -1.1, "regional", 35, 105),
|
| 24 |
+
("PDRMIP_preindustrial_SO4", 0.8, "global", 0, 0),
|
| 25 |
+
("ECLIPSE_CH4_minus20", -0.7, "global", 0, 0),
|
| 26 |
+
("ECLIPSE_2xCO2", 2.3, "global", 0, 0),
|
| 27 |
+
("ECLIPSE_BC_minus100", -0.6, "global", 0, 0),
|
| 28 |
+
("ECLIPSE_SO2_minus100", 1.0, "global", 0, 0),
|
| 29 |
+
("ECLIPSE_CO_minus100", -0.45, "global", 0, 0),
|
| 30 |
+
("KASOAR_SO2_NHML_minus100", 0.75, "regional", 42, 30),
|
| 31 |
+
("KASOAR_BC_NHML_minus100", -0.55, "regional", 42, 30),
|
| 32 |
+
("KASOAR_SO2_China_minus100", 0.70, "regional", 34, 105),
|
| 33 |
+
("KASOAR_SO2_EastAsia_minus100", 0.65, "regional", 40, 125),
|
| 34 |
+
("KASOAR_SO2_Europe_minus100", 0.62, "regional", 52, 15),
|
| 35 |
+
("KASOAR_SO2_US_minus100", 0.58, "regional", 40, 260),
|
| 36 |
+
("PDRMIP_SO4_Asia_alt", -0.72, "regional", 25, 80),
|
| 37 |
+
]
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def wrapped_distance(longitude, centre):
|
| 41 |
+
return (longitude - centre + 180.0) % 360.0 - 180.0
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
def main():
|
| 45 |
+
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
|
| 46 |
+
ny, nx = map(int, config["data"]["grid"])
|
| 47 |
+
if (ny, nx) != GRID_SHAPE or len(SCENARIOS) != 21:
|
| 48 |
+
raise ValueError("protocol requires exactly 21 scenarios on a 145x192 grid")
|
| 49 |
+
latitude = np.linspace(-90.0, 90.0, ny, dtype=np.float64)
|
| 50 |
+
longitude = np.linspace(0.0, 360.0, nx, endpoint=False, dtype=np.float64)
|
| 51 |
+
lat, lon = np.meshgrid(latitude, longitude, indexing="ij")
|
| 52 |
+
polar = np.sin(np.deg2rad(lat)) ** 4
|
| 53 |
+
land_wave = np.cos(np.deg2rad(2.0 * lon - 0.6 * lat)) * np.cos(np.deg2rad(lat)) ** 2
|
| 54 |
+
short_fields, long_fields = [], []
|
| 55 |
+
for index, (scenario_id, amplitude, forcing_type, centre_lat, centre_lon) in enumerate(SCENARIOS):
|
| 56 |
+
local = np.exp(-0.5 * ((lat - centre_lat) / 16.0) ** 2
|
| 57 |
+
-0.5 * (wrapped_distance(lon, centre_lon) / 28.0) ** 2)
|
| 58 |
+
if forcing_type == "global":
|
| 59 |
+
local = 0.25 * np.cos(np.deg2rad(lat)) ** 2
|
| 60 |
+
phase = 2.0 * np.pi * index / len(SCENARIOS)
|
| 61 |
+
forcing_region = amplitude * local
|
| 62 |
+
global_warming = amplitude * (0.52 + 0.34 * polar)
|
| 63 |
+
circulation = 0.13 * amplitude * np.cos(np.deg2rad(lat * 2.0) + phase) * np.sin(np.deg2rad(lon) - phase)
|
| 64 |
+
short = global_warming + 0.72 * forcing_region + 0.10 * amplitude * land_wave + circulation
|
| 65 |
+
zonal = short.mean(axis=1, keepdims=True)
|
| 66 |
+
remote_east = np.roll(short, nx // 5, axis=1)
|
| 67 |
+
cross_equatorial = np.flip(zonal, axis=0)
|
| 68 |
+
planetary_wave = np.cos(np.deg2rad(lon * 2.0 + centre_lon)) * np.cos(np.deg2rad(lat))
|
| 69 |
+
long = (1.42 * short + 0.28 * remote_east + 0.23 * cross_equatorial
|
| 70 |
+
+ short.mean() * (0.48 + 0.62 * polar) + 0.09 * amplitude * planetary_wave)
|
| 71 |
+
short_fields.append(short.astype(np.float32))
|
| 72 |
+
long_fields.append(long.astype(np.float32))
|
| 73 |
+
short_response = np.stack(short_fields)
|
| 74 |
+
long_response = np.stack(long_fields)
|
| 75 |
+
if short_response.shape != (21, 145, 192) or not np.isfinite(long_response).all():
|
| 76 |
+
raise ValueError("invalid synthetic response fields")
|
| 77 |
+
output = ROOT / config["data"]["path"]
|
| 78 |
+
output.parent.mkdir(parents=True, exist_ok=True)
|
| 79 |
+
np.savez_compressed(output, format_version=np.array(FORMAT_VERSION),
|
| 80 |
+
scenario_ids=np.asarray([item[0] for item in SCENARIOS]),
|
| 81 |
+
forcing_amplitude=np.asarray([item[1] for item in SCENARIOS], dtype=np.float32),
|
| 82 |
+
latitude_deg=latitude.astype(np.float32), longitude_deg=longitude.astype(np.float32),
|
| 83 |
+
short_response=short_response, long_response=long_response)
|
| 84 |
+
print(f"data={output.relative_to(ROOT)} shape={short_response.shape} features={ny * nx}")
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
if __name__ == "__main__":
|
| 88 |
+
main()
|
scripts/inference.py
ADDED
|
@@ -0,0 +1,60 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Restore every LOSO fold and emit both methods' 21 complete response maps."""
|
| 2 |
+
|
| 3 |
+
import sys
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
import numpy as np
|
| 7 |
+
import torch
|
| 8 |
+
import yaml
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 12 |
+
sys.path.insert(0, str(ROOT))
|
| 13 |
+
from model.climemu_s2l import FEATURE_COUNT, FORMAT_VERSION, GRID_SHAPE, DualRidge, SharedKernelGPR
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
def load_checkpoint(path):
|
| 17 |
+
try:
|
| 18 |
+
return torch.load(path, map_location="cpu", weights_only=False)
|
| 19 |
+
except TypeError:
|
| 20 |
+
return torch.load(path, map_location="cpu")
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
def main():
|
| 24 |
+
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
|
| 25 |
+
checkpoint = load_checkpoint(ROOT / config["paths"]["checkpoint"])
|
| 26 |
+
if checkpoint["format_version"] != FORMAT_VERSION or tuple(checkpoint["grid_shape"]) != GRID_SHAPE:
|
| 27 |
+
raise ValueError("checkpoint protocol mismatch")
|
| 28 |
+
x = checkpoint["short_response"].double()
|
| 29 |
+
y = checkpoint["long_response"].double()
|
| 30 |
+
required = {"model", "model_config", "format_version"}
|
| 31 |
+
if not required.issubset(checkpoint):
|
| 32 |
+
raise ValueError(f"checkpoint is missing standard fields: {sorted(required - checkpoint.keys())}")
|
| 33 |
+
predictions = {"ridge": [], "gpr": []}
|
| 34 |
+
for expected_fold, state in enumerate(checkpoint["model"]["folds"]):
|
| 35 |
+
if state["fold"] != expected_fold or state["held_out_scenario_id"] != checkpoint["scenario_ids"][expected_fold]:
|
| 36 |
+
raise ValueError("fold/scenario identity mismatch")
|
| 37 |
+
training = state["train_indices"]
|
| 38 |
+
ridge = DualRidge(float(state["ridge_alpha"])).fit(x[training], y[training])
|
| 39 |
+
gpr = SharedKernelGPR(state["gpr_kernel_mode"], config["model"]["gpr"]["jitter"])
|
| 40 |
+
gpr.load_hyperparameters(state["gpr_hyperparameters"]).restore_posterior(x[training], y[training])
|
| 41 |
+
predictions["ridge"].append(ridge.predict(x[expected_fold:expected_fold + 1]).detach().numpy()[0])
|
| 42 |
+
predictions["gpr"].append(gpr.predict(x[expected_fold:expected_fold + 1]).detach().numpy()[0])
|
| 43 |
+
ridge = np.asarray(predictions["ridge"], dtype=np.float32).reshape(21, *GRID_SHAPE)
|
| 44 |
+
gpr = np.asarray(predictions["gpr"], dtype=np.float32).reshape(21, *GRID_SHAPE)
|
| 45 |
+
if ridge.shape != (21, 145, 192) or not np.isfinite(gpr).all():
|
| 46 |
+
raise ValueError("inference did not produce 21 finite full-grid fields")
|
| 47 |
+
source = np.load(ROOT / config["data"]["path"])
|
| 48 |
+
output = ROOT / config["paths"]["inference"]
|
| 49 |
+
output.parent.mkdir(parents=True, exist_ok=True)
|
| 50 |
+
np.savez_compressed(output, format_version=np.array(FORMAT_VERSION),
|
| 51 |
+
scenario_ids=np.asarray(checkpoint["scenario_ids"]),
|
| 52 |
+
latitude_deg=source["latitude_deg"], longitude_deg=source["longitude_deg"],
|
| 53 |
+
short_response=x.numpy().reshape(21, *GRID_SHAPE).astype(np.float32),
|
| 54 |
+
long_response=y.numpy().reshape(21, *GRID_SHAPE).astype(np.float32),
|
| 55 |
+
ridge_prediction=ridge, gpr_prediction=gpr)
|
| 56 |
+
print(f"predictions={output.relative_to(ROOT)} methods=2 shape={ridge.shape}")
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
if __name__ == "__main__":
|
| 60 |
+
main()
|
scripts/result.py
ADDED
|
@@ -0,0 +1,103 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Compute paper metrics for every scenario and draw maps and error boxes."""
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
import sys
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
|
| 7 |
+
import matplotlib
|
| 8 |
+
matplotlib.use("Agg")
|
| 9 |
+
import matplotlib.pyplot as plt
|
| 10 |
+
import numpy as np
|
| 11 |
+
import yaml
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 15 |
+
sys.path.insert(0, str(ROOT))
|
| 16 |
+
from model.climemu_s2l import GRID_SHAPE, area_weights, weighted_rmse
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
REGIONS = {
|
| 20 |
+
"North America": (15, 75, 190, 310), "Europe": (35, 70, 350, 45),
|
| 21 |
+
"South Asia": (5, 35, 60, 100), "East Asia": (20, 55, 100, 150),
|
| 22 |
+
"Arctic": (66, 90, 0, 360), "Northwest Asia": (50, 75, 45, 120),
|
| 23 |
+
"Northern Africa": (15, 35, 340, 55), "Southern Africa": (-35, 0, 10, 55),
|
| 24 |
+
"South America": (-60, 15, 275, 330), "Australia": (-45, -10, 110, 155),
|
| 25 |
+
}
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
def region_mask(latitude, longitude, bounds):
|
| 29 |
+
south, north, west, east = bounds
|
| 30 |
+
latitude_mask = (latitude[:, None] >= south) & (latitude[:, None] <= north)
|
| 31 |
+
longitude_mask = ((longitude[None, :] >= west) & (longitude[None, :] <= east)
|
| 32 |
+
if west <= east else (longitude[None, :] >= west) | (longitude[None, :] <= east))
|
| 33 |
+
return latitude_mask & longitude_mask
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
def weighted_mean(field, weights, mask=None):
|
| 37 |
+
selected = np.ones(field.shape, dtype=bool) if mask is None else mask
|
| 38 |
+
local_weights = weights[selected]
|
| 39 |
+
return float(np.sum(field[selected] * local_weights) / np.sum(local_weights))
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
def main():
|
| 43 |
+
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
|
| 44 |
+
data = np.load(ROOT / config["paths"]["inference"])
|
| 45 |
+
latitude, longitude = data["latitude_deg"], data["longitude_deg"]
|
| 46 |
+
weights = area_weights(latitude, GRID_SHAPE[1])
|
| 47 |
+
targets = data["long_response"]
|
| 48 |
+
scenario_ids = [str(item) for item in data["scenario_ids"]]
|
| 49 |
+
masks = {name: region_mask(latitude, longitude, bounds) for name, bounds in REGIONS.items()}
|
| 50 |
+
report = {"grid": list(GRID_SHAPE), "scenario_count": 21, "regions": list(REGIONS), "methods": {}}
|
| 51 |
+
box_values = {}
|
| 52 |
+
for method in ("ridge", "gpr"):
|
| 53 |
+
predictions = data[f"{method}_prediction"]
|
| 54 |
+
scenarios = []
|
| 55 |
+
for index, scenario_id in enumerate(scenario_ids):
|
| 56 |
+
target, prediction = targets[index], predictions[index]
|
| 57 |
+
regional = {name: abs(weighted_mean(prediction, weights, mask) - weighted_mean(target, weights, mask))
|
| 58 |
+
for name, mask in masks.items()}
|
| 59 |
+
scenarios.append({"scenario_id": scenario_id,
|
| 60 |
+
"area_weighted_grid_rmse": weighted_rmse(target, prediction, weights),
|
| 61 |
+
"global_mean_absolute_error": abs(weighted_mean(prediction, weights) - weighted_mean(target, weights)),
|
| 62 |
+
"regional_mean_absolute_error": regional})
|
| 63 |
+
grid_errors = [item["area_weighted_grid_rmse"] for item in scenarios]
|
| 64 |
+
global_errors = [item["global_mean_absolute_error"] for item in scenarios]
|
| 65 |
+
report["methods"][method] = {"scenarios": scenarios,
|
| 66 |
+
"summary": {"mean_grid_rmse": float(np.mean(grid_errors)), "median_grid_rmse": float(np.median(grid_errors)),
|
| 67 |
+
"mean_global_absolute_error": float(np.mean(global_errors)),
|
| 68 |
+
"mean_regional_absolute_error": {name: float(np.mean([item["regional_mean_absolute_error"][name]
|
| 69 |
+
for item in scenarios])) for name in REGIONS}}}
|
| 70 |
+
box_values[method] = [grid_errors, global_errors] + [[item["regional_mean_absolute_error"][name] for item in scenarios]
|
| 71 |
+
for name in REGIONS]
|
| 72 |
+
numeric = [value for method in report["methods"].values() for value in
|
| 73 |
+
(method["summary"]["mean_grid_rmse"], method["summary"]["mean_global_absolute_error"])]
|
| 74 |
+
if not np.isfinite(numeric).all():
|
| 75 |
+
raise FloatingPointError("evaluation metrics are not finite")
|
| 76 |
+
output = ROOT / config["paths"]["evaluation_dir"]
|
| 77 |
+
output.mkdir(parents=True, exist_ok=True)
|
| 78 |
+
(output / "metrics.json").write_text(json.dumps(report, indent=2) + "\n")
|
| 79 |
+
plot_index = scenario_ids.index(config["evaluation"]["plot_scenario_id"])
|
| 80 |
+
fields = [targets[plot_index], data["ridge_prediction"][plot_index], data["gpr_prediction"][plot_index],
|
| 81 |
+
data["ridge_prediction"][plot_index] - targets[plot_index], data["gpr_prediction"][plot_index] - targets[plot_index]]
|
| 82 |
+
titles = ["HadGEM3 target", "Dual Ridge", "Shared-kernel GPR", "Ridge error", "GPR error"]
|
| 83 |
+
figure, axes = plt.subplots(2, 3, figsize=(14, 7), constrained_layout=True)
|
| 84 |
+
for axis, field, title in zip(axes.flat, fields, titles):
|
| 85 |
+
limit = max(abs(np.percentile(field, 1)), abs(np.percentile(field, 99)))
|
| 86 |
+
image = axis.imshow(field, origin="lower", extent=(0, 360, -90, 90), cmap="RdBu_r", vmin=-limit, vmax=limit, aspect="auto")
|
| 87 |
+
axis.set(title=title, xlabel="Longitude", ylabel="Latitude"); figure.colorbar(image, ax=axis, shrink=0.75)
|
| 88 |
+
axes.flat[-1].axis("off")
|
| 89 |
+
figure.suptitle(config["evaluation"]["plot_scenario_id"] + " full 145x192 fields")
|
| 90 |
+
figure.savefig(output / "spatial_fields.png", dpi=150); plt.close(figure)
|
| 91 |
+
labels = ["Grid RMSE", "Global"] + list(REGIONS)
|
| 92 |
+
figure, axes = plt.subplots(1, 2, figsize=(15, 5), sharey=False, constrained_layout=True)
|
| 93 |
+
for axis, method, color in zip(axes, ("ridge", "gpr"), ("#2a6f97", "#c45d35")):
|
| 94 |
+
boxes = axis.boxplot(box_values[method], patch_artist=True, showmeans=True)
|
| 95 |
+
for patch in boxes["boxes"]: patch.set_facecolor(color); patch.set_alpha(0.65)
|
| 96 |
+
axis.set_xticks(range(1, len(labels) + 1), labels, rotation=55, ha="right")
|
| 97 |
+
axis.set(title=method.upper(), ylabel="Absolute error / RMSE (deg C)"); axis.grid(axis="y", alpha=0.25)
|
| 98 |
+
figure.savefig(output / "error_boxplots.png", dpi=150); plt.close(figure)
|
| 99 |
+
print(f"evaluation={output.relative_to(ROOT)} scenarios=21 regions=10 methods=2")
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
if __name__ == "__main__":
|
| 103 |
+
main()
|
scripts/train.py
ADDED
|
@@ -0,0 +1,130 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Train all 21 LOSO folds; DDP ranks independently shard outer folds."""
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
import argparse
|
| 5 |
+
import os
|
| 6 |
+
import random
|
| 7 |
+
import sys
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
|
| 10 |
+
import numpy as np
|
| 11 |
+
import torch
|
| 12 |
+
import yaml
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 16 |
+
sys.path.insert(0, str(ROOT))
|
| 17 |
+
from model.climemu_s2l import (FEATURE_COUNT, FORMAT_VERSION, GRID_SHAPE,
|
| 18 |
+
MODEL_NAME, DualRidge, SharedKernelGPR,
|
| 19 |
+
area_weights, select_ridge_alpha, weighted_rmse)
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def load_data(path):
|
| 23 |
+
data = np.load(path)
|
| 24 |
+
short, long = data["short_response"], data["long_response"]
|
| 25 |
+
if str(data["format_version"]) != FORMAT_VERSION:
|
| 26 |
+
raise ValueError("data format version mismatch")
|
| 27 |
+
if short.shape != (21, *GRID_SHAPE) or long.shape != short.shape:
|
| 28 |
+
raise ValueError("responses must both have shape [21,145,192]")
|
| 29 |
+
ids = [str(value) for value in data["scenario_ids"]]
|
| 30 |
+
if len(set(ids)) != 21 or not np.isfinite(short).all() or not np.isfinite(long).all():
|
| 31 |
+
raise ValueError("scenario IDs must be unique and fields finite")
|
| 32 |
+
return data, short.reshape(21, FEATURE_COUNT), long.reshape(21, FEATURE_COUNT), ids
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def train_fold(fold, x, y, scenario_ids, config, weights):
|
| 36 |
+
training = np.asarray([index for index in range(21) if index != fold])
|
| 37 |
+
alpha, cv_scores = select_ridge_alpha(
|
| 38 |
+
x[training], y[training], config["model"]["ridge"]["alphas"],
|
| 39 |
+
int(config["model"]["ridge"]["inner_folds"]), weights)
|
| 40 |
+
ridge = DualRidge(alpha).fit(x[training], y[training])
|
| 41 |
+
ridge_prediction = ridge.predict(x[fold:fold + 1]).detach().cpu().numpy()[0]
|
| 42 |
+
gpr_config = config["model"]["gpr"]
|
| 43 |
+
gpr = SharedKernelGPR(gpr_config["kernel_mode"], gpr_config["jitter"])
|
| 44 |
+
optimizer_trace = gpr.fit(x[training], y[training], gpr_config["optimizer_steps"], gpr_config["learning_rate"])
|
| 45 |
+
gpr_prediction = gpr.predict(x[fold:fold + 1]).detach().cpu().numpy()[0]
|
| 46 |
+
state = {"fold": fold, "held_out_scenario_id": scenario_ids[fold], "train_indices": training.tolist(),
|
| 47 |
+
"ridge_alpha": alpha, "ridge_cv_scores": cv_scores,
|
| 48 |
+
"gpr_kernel_mode": gpr.kernel_mode, "gpr_hyperparameters": gpr.hyperparameters(),
|
| 49 |
+
"gpr_optimizer": optimizer_trace}
|
| 50 |
+
metrics = {"scenario_id": scenario_ids[fold],
|
| 51 |
+
"ridge_weighted_rmse": weighted_rmse(y[fold], ridge_prediction, weights),
|
| 52 |
+
"gpr_weighted_rmse": weighted_rmse(y[fold], gpr_prediction, weights),
|
| 53 |
+
"selected_alpha": alpha, "gpr_final_loss": optimizer_trace["loss"][-1],
|
| 54 |
+
"gpr_backward": optimizer_trace["gradient_seen"],
|
| 55 |
+
"gpr_parameter_update": optimizer_trace["parameter_updated"]}
|
| 56 |
+
return state, metrics
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def main():
|
| 60 |
+
parser = argparse.ArgumentParser()
|
| 61 |
+
parser.add_argument("--resume", action="store_true", help="resume completed LOSO folds from the checkpoint")
|
| 62 |
+
args = parser.parse_args()
|
| 63 |
+
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
|
| 64 |
+
seed = int(config["seed"])
|
| 65 |
+
random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
|
| 66 |
+
torch.set_default_dtype(torch.float64)
|
| 67 |
+
distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
|
| 68 |
+
if distributed:
|
| 69 |
+
torch.distributed.init_process_group(config["runtime"]["ddp_backend"])
|
| 70 |
+
rank = torch.distributed.get_rank() if distributed else 0
|
| 71 |
+
world = torch.distributed.get_world_size() if distributed else 1
|
| 72 |
+
data, x_numpy, y_numpy, scenario_ids = load_data(ROOT / config["data"]["path"])
|
| 73 |
+
x, y = torch.from_numpy(x_numpy).double(), torch.from_numpy(y_numpy).double()
|
| 74 |
+
weights = area_weights(data["latitude_deg"], GRID_SHAPE[1])
|
| 75 |
+
checkpoint_path = ROOT / config["paths"]["checkpoint"]
|
| 76 |
+
resumed_states, resumed_metrics = {}, {}
|
| 77 |
+
if args.resume and checkpoint_path.exists():
|
| 78 |
+
try:
|
| 79 |
+
previous = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
|
| 80 |
+
except TypeError:
|
| 81 |
+
previous = torch.load(checkpoint_path, map_location="cpu")
|
| 82 |
+
if previous.get("format_version") != FORMAT_VERSION or tuple(previous.get("grid_shape", ())) != GRID_SHAPE:
|
| 83 |
+
raise ValueError("checkpoint protocol mismatch")
|
| 84 |
+
restored_folds = previous.get("model", {}).get("folds", previous.get("folds", []))
|
| 85 |
+
resumed_states = {state["fold"]: state for state in restored_folds}
|
| 86 |
+
resumed_metrics = {index: value for index, value in enumerate(previous.get("training_metrics", []))}
|
| 87 |
+
local_states, local_metrics = {}, {}
|
| 88 |
+
for fold in range(rank, 21, world):
|
| 89 |
+
if fold in resumed_states and fold in resumed_metrics:
|
| 90 |
+
continue
|
| 91 |
+
state, metrics = train_fold(fold, x, y, scenario_ids, config, weights)
|
| 92 |
+
local_states[fold], local_metrics[fold] = state, metrics
|
| 93 |
+
print(f"rank={rank} fold={fold:02d} scenario={scenario_ids[fold]} alpha={state['ridge_alpha']}")
|
| 94 |
+
if distributed:
|
| 95 |
+
gathered_states, gathered_metrics = [None] * world, [None] * world
|
| 96 |
+
torch.distributed.all_gather_object(gathered_states, local_states)
|
| 97 |
+
torch.distributed.all_gather_object(gathered_metrics, local_metrics)
|
| 98 |
+
states = {**resumed_states, **{key: value for item in gathered_states for key, value in item.items()}}
|
| 99 |
+
metrics = {**resumed_metrics, **{key: value for item in gathered_metrics for key, value in item.items()}}
|
| 100 |
+
else:
|
| 101 |
+
states, metrics = {**resumed_states, **local_states}, {**resumed_metrics, **local_metrics}
|
| 102 |
+
if rank == 0:
|
| 103 |
+
if sorted(states) != list(range(21)):
|
| 104 |
+
raise RuntimeError("all 21 LOSO folds must be trained")
|
| 105 |
+
fold_states = [states[index] for index in range(21)]
|
| 106 |
+
fold_metrics = [metrics[index] for index in range(21)]
|
| 107 |
+
checkpoint = {"epoch": 1,
|
| 108 |
+
"model": {"folds": fold_states},
|
| 109 |
+
"model_config": config["model"],
|
| 110 |
+
"optimizer_state_dict": {"folds": [state["gpr_optimizer"] for state in fold_states]},
|
| 111 |
+
"loss": float(np.mean([item["gpr_final_loss"] for item in fold_metrics])),
|
| 112 |
+
"config": config,
|
| 113 |
+
"model_name": MODEL_NAME, "format_version": FORMAT_VERSION,
|
| 114 |
+
"grid_shape": GRID_SHAPE, "scenario_ids": scenario_ids,
|
| 115 |
+
"short_response": torch.from_numpy(x_numpy), "long_response": torch.from_numpy(y_numpy),
|
| 116 |
+
"training_metrics": fold_metrics,
|
| 117 |
+
"gpr_ard": False, "gpr_ard_gap": config["model"]["gpr"]["ard_gap"]}
|
| 118 |
+
checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
|
| 119 |
+
torch.save(checkpoint, checkpoint_path)
|
| 120 |
+
metrics_path = ROOT / config["paths"]["training_metrics"]
|
| 121 |
+
metrics_path.parent.mkdir(parents=True, exist_ok=True)
|
| 122 |
+
metrics_path.write_text(json.dumps({"fold_count": 21, "world_size": world,
|
| 123 |
+
"folds": fold_metrics}, indent=2) + "\n")
|
| 124 |
+
print(f"checkpoint={checkpoint_path.relative_to(ROOT)} folds=21 world_size={world} resumed={len(resumed_states)}")
|
| 125 |
+
if distributed:
|
| 126 |
+
torch.distributed.barrier(); torch.distributed.destroy_process_group()
|
| 127 |
+
|
| 128 |
+
|
| 129 |
+
if __name__ == "__main__":
|
| 130 |
+
main()
|
weight/.gitkeep
ADDED
|
File without changes
|