zhangrenchao commited on
Commit
87f2bd3
·
verified ·
1 Parent(s): 1b38903

Upload folder using huggingface_hub

Browse files
.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
- *.pth filter=lfs diff=lfs merge=lfs -text
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