| """Train four scale-specific pairs of joint multi-output random forests.""" |
|
|
| import json |
| import os |
| import random |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from model.rf_climparam import FORMAT_VERSION, MODEL_NAME, build_pair |
|
|
|
|
| def load_scale(path, scale): |
| data = np.load(path) |
| expected = {"tend_inputs": 145, "tend_targets": 144, "diff_inputs": 62, "diff_targets": 17} |
| if str(data["format_version"]) != FORMAT_VERSION or str(data["scale"]) != scale: |
| raise ValueError(f"invalid metadata in {path}") |
| count = len(data["tend_inputs"]) |
| ny, nx = map(int, data["grid_shape"]) |
| if count != ny * nx or int(data["snapshot_count"]) != 1: |
| raise ValueError(f"{scale} must contain one complete [{ny},{nx}] snapshot") |
| linear = data["grid_row"].astype(np.int64) * nx + data["grid_column"].astype(np.int64) |
| if not np.array_equal(linear, np.arange(count)): |
| raise ValueError(f"{scale} grid cannot be reversibly flattened") |
| for name, width in expected.items(): |
| value = data[name] |
| if value.shape != (count, width) or value.dtype != np.float32 or not np.isfinite(value).all(): |
| raise ValueError(f"{scale}/{name} requires finite float32 [{count},{width}]") |
| if np.any(data["diff_targets"][:, :15] < 0): |
| raise ValueError("Dbar training targets must be nonnegative") |
| return data |
|
|
|
|
| def sample_training_columns(data, columns_per_latitude, seed): |
| ny, nx = map(int, data["grid_shape"]) |
| if not 1 <= columns_per_latitude <= nx: |
| raise ValueError("train_columns_per_latitude must be between 1 and grid width") |
| rng = np.random.default_rng(seed) |
| selected = [] |
| for row in range(ny): |
| selected.extend(row * nx + rng.choice(nx, columns_per_latitude, replace=False)) |
| return np.asarray(selected, dtype=np.int64) |
|
|
|
|
| def main(): |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| seed = int(config["seed"]) |
| random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) |
| distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1 |
| local_rank = int(os.environ.get("LOCAL_RANK", "0")) |
| if distributed: |
| torch.distributed.init_process_group("gloo") |
| rank = torch.distributed.get_rank() if distributed else 0 |
| world = torch.distributed.get_world_size() if distributed else 1 |
| scales = list(config["data"]["scales"]) |
| assigned = [scale for i, scale in enumerate(scales) if i % world == rank] |
| local_states, local_metrics = {}, {} |
| for scale_index, scale in enumerate(assigned): |
| data = load_scale(ROOT / config["data"]["root"] / f"{scale}.npz", scale) |
| selected = sample_training_columns( |
| data, int(config["data"]["train_columns_per_latitude"]), seed + scales.index(scale)) |
| pair = build_pair(config["model"], seed + scales.index(scale) * 100) |
| pair["rf_tend"].fit(data["tend_inputs"][selected], data["tend_targets"][selected]) |
| pair["rf_diff"].fit(data["diff_inputs"][selected], data["diff_targets"][selected]) |
| local_states[scale] = {name: model.state_dict() for name, model in pair.items()} |
| local_metrics[scale] = {"complete_grid_points": len(data["tend_inputs"]), |
| "train_columns": len(selected), |
| "columns_per_latitude": int(config["data"]["train_columns_per_latitude"])} |
| print(f"rank={rank} trained={scale} sampled_columns={len(selected)}") |
| if distributed: |
| gathered_states, gathered_metrics = [None] * world, [None] * world |
| torch.distributed.all_gather_object(gathered_states, local_states) |
| torch.distributed.all_gather_object(gathered_metrics, local_metrics) |
| states = {key: value for item in gathered_states for key, value in item.items()} |
| metrics = {key: value for item in gathered_metrics for key, value in item.items()} |
| else: |
| states, metrics = local_states, local_metrics |
| if rank == 0: |
| if set(states) != set(scales): |
| raise RuntimeError("not all scales were trained") |
| checkpoint = { |
| "model": states, |
| "model_name": MODEL_NAME, |
| "model_config": {"dimensions": {"rf_tend_input": 145, "rf_tend_output": 144, |
| "rf_diff_input": 62, "rf_diff_output": 17}, |
| "engineering": config["model"]["engineering"], |
| "paper_model": config["paper_model"], "scales": scales}, |
| "format_version": FORMAT_VERSION, "seed": seed} |
| path = ROOT / config["paths"]["checkpoint"] |
| path.parent.mkdir(parents=True, exist_ok=True) |
| torch.save(checkpoint, path) |
| metrics_path = ROOT / config["paths"]["training_metrics"] |
| metrics_path.parent.mkdir(parents=True, exist_ok=True) |
| metrics_path.write_text(json.dumps({"format_version": FORMAT_VERSION, "scales": metrics}, indent=2) + "\n") |
| print(f"checkpoint={path.relative_to(ROOT)} scales={len(states)}") |
| if distributed: |
| torch.distributed.barrier(); torch.distributed.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
