File size: 5,250 Bytes
c059069 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """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()
|