model/pbl_emulator.py

from __future__ import annotations

import json
import math
import os
import random
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import torch
from torch import nn
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, TensorDataset


INPUT_NAMES = ["Q2", "T2", "U10", "V10", "GRDFLX", "SWDOWN", "GLW", "LH", "HFX", "PBLH", "UST", "TSK", "TSLB", "SMOIS", "Ug", "Vg"]
OUTPUT_NAMES = ["U", "V", "W", "tk", "QVAPOR"]
PARAMETER_COUNTS = {"FFN": 10693, "HPC": 16597, "HAC": 26197}
CHECKPOINT_FORMAT_VERSION = "1.0"


def load_yaml(path: str | Path) -> dict[str, Any]:
    try:
        import yaml
    except ImportError as exc:
        raise RuntimeError("PyYAML is required to read conf/config.yaml") from exc
    with open(path, "r", encoding="utf-8") as handle:
        return yaml.safe_load(handle)


class FFN(nn.Module):
    def __init__(self, width: int = 16, levels: int = 17, variables: int = 5):
        super().__init__()
        layers: list[nn.Module] = []
        in_features = 16
        for _ in range(34):
            layers.extend((nn.Linear(in_features, width), nn.ReLU()))
            in_features = width
        self.hidden = nn.Sequential(*layers)
        self.output = nn.Linear(width, levels * variables)
        self.levels, self.variables = levels, variables

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.output(self.hidden(x)).reshape(-1, self.levels, self.variables)


class HierarchicalNetwork(nn.Module):
    def __init__(self, mode: str, width: int = 16, levels: int = 17, variables: int = 5):
        super().__init__()
        self.mode, self.levels, self.variables = mode, levels, variables
        blocks = []
        for level in range(levels):
            conditioned_outputs = variables * (level if mode == "HAC" else min(level, 1))
            blocks.append(nn.Sequential(
                nn.Linear(16 + conditioned_outputs, width), nn.ReLU(),
                nn.Linear(width, width), nn.ReLU(),
                nn.Linear(width, width), nn.ReLU(),
                nn.Linear(width, variables),
            ))
        self.blocks = nn.ModuleList(blocks)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        outputs = []
        for block in self.blocks:
            if not outputs:
                conditioned = x
            elif self.mode == "HPC":
                conditioned = torch.cat((x, outputs[-1]), dim=-1)
            else:
                conditioned = torch.cat((x, *outputs), dim=-1)
            outputs.append(block(conditioned))
        return torch.stack(outputs, dim=1)


class HAC(HierarchicalNetwork):
    def __init__(self, width: int = 16, levels: int = 17, variables: int = 5):
        super().__init__("HAC", width, levels, variables)


class PBLEmulator(nn.Module):
    def __init__(self, architecture: str = "HAC", width: int = 16, levels: int = 17, variables: int = 5):
        super().__init__()
        self.architecture = architecture.upper()
        if self.architecture == "FFN":
            self.model = FFN(width, levels, variables)
        elif self.architecture == "HPC":
            self.model = HierarchicalNetwork(self.architecture, width, levels, variables)
        elif self.architecture == "HAC":
            self.model = HAC(width, levels, variables)
        else:
            raise ValueError(f"Unknown architecture: {architecture}")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.model(x)


def build_model(architecture: str = "HAC", width: int = 16, levels: int = 17, variables: int = 5) -> PBLEmulator:
    model = PBLEmulator(architecture, width, levels, variables)
    count = sum(parameter.numel() for parameter in model.parameters())
    if width == 16 and levels == 17 and variables == 5:
        assert count == PARAMETER_COUNTS[model.architecture], f"{model.architecture}: expected {PARAMETER_COUNTS[model.architecture]}, got {count}"
    return model


@dataclass
class ColumnScaler:
    mean: np.ndarray
    scale: np.ndarray
    minimum: np.ndarray
    span: np.ndarray

    @classmethod
    def fit(cls, values: np.ndarray) -> "ColumnScaler":
        flat = np.asarray(values, dtype=np.float64).reshape(len(values), -1)
        mean = flat.mean(axis=0)
        scale = flat.std(axis=0)
        scale[scale < 1e-12] = 1.0
        standardized = (flat - mean) / scale
        minimum = standardized.min(axis=0)
        span = standardized.max(axis=0) - minimum
        span[span < 1e-12] = 1.0
        return cls(mean, scale, minimum, span)

    def transform(self, values: np.ndarray) -> np.ndarray:
        shape = values.shape
        flat = np.asarray(values, dtype=np.float64).reshape(len(values), -1)
        return (((flat - self.mean) / self.scale - self.minimum) / self.span).reshape(shape).astype(np.float32)

    def inverse_transform(self, values: np.ndarray) -> np.ndarray:
        shape = values.shape
        flat = np.asarray(values, dtype=np.float64).reshape(len(values), -1)
        return ((flat * self.span + self.minimum) * self.scale + self.mean).reshape(shape).astype(np.float32)

    def state_dict(self) -> dict[str, np.ndarray]:
        return {"mean": self.mean, "scale": self.scale, "minimum": self.minimum, "span": self.span}

    @classmethod
    def from_state_dict(cls, state: dict[str, Any]) -> "ColumnScaler":
        return cls(*(np.asarray(state[key]) for key in ("mean", "scale", "minimum", "span")))


def _distributed() -> tuple[bool, int, int, int]:
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    rank = int(os.environ.get("RANK", "0"))
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    if world_size > 1 and not torch.distributed.is_initialized():
        torch.distributed.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
    return world_size > 1, rank, local_rank, world_size


def _reduce_loss(total: float, count: int, device: torch.device) -> float:
    pair = torch.tensor([total, count], dtype=torch.float64, device=device)
    if torch.distributed.is_initialized():
        torch.distributed.all_reduce(pair, op=torch.distributed.ReduceOp.SUM)
    return float(pair[0] / pair[1].clamp_min(1))


def _epoch(model: nn.Module, loader: DataLoader, device: torch.device, optimizer: torch.optim.Optimizer | None) -> float:
    model.train(optimizer is not None)
    total, count = 0.0, 0
    context = torch.enable_grad() if optimizer is not None else torch.no_grad()
    with context:
        for x_batch, y_batch in loader:
            x_batch, y_batch = x_batch.to(device), y_batch.to(device)
            if optimizer is not None:
                optimizer.zero_grad(set_to_none=True)
            loss = torch.mean((model(x_batch) - y_batch) ** 2)
            if optimizer is not None:
                loss.backward()
                optimizer.step()
            total += float(loss.detach()) * len(x_batch)
            count += len(x_batch)
    return _reduce_loss(total, count, device)


def train_model(data_path: str | Path, checkpoint_path: str | Path, metrics_path: str | Path, settings: dict[str, Any]) -> dict[str, Any]:
    distributed, rank, local_rank, world_size = _distributed()
    seed = int(settings.get("seed", 19))
    random.seed(seed + rank); np.random.seed(seed + rank); torch.manual_seed(seed + rank)
    if torch.cuda.is_available():
        torch.cuda.set_device(local_rank)
        device = torch.device("cuda", local_rank)
    else:
        device = torch.device("cpu")
    raw = np.load(data_path)
    x_train, y_train = raw["x_train"], raw["y_train"]
    x_val, y_val = raw["x_val"], raw["y_val"]
    assert x_train.shape[1:] == (16,) and y_train.shape[1:] == (17, 5)
    model_config = {
        "architecture": settings.get("architecture", "HAC"),
        "width": int(settings.get("width", 16)),
        "levels": int(settings.get("levels", 17)),
        "output_variables": int(settings.get("output_variables", 5)),
    }
    start_epoch, history, best_loss = 0, [], math.inf
    checkpoint_path = Path(checkpoint_path)
    saved = None
    if settings.get("resume") and checkpoint_path.exists():
        saved = torch.load(checkpoint_path, map_location=device, weights_only=False)
        if saved.get("format_version") != CHECKPOINT_FORMAT_VERSION:
            raise ValueError(f"Unsupported checkpoint format_version: {saved.get('format_version')!r}; expected {CHECKPOINT_FORMAT_VERSION!r}")
        model_config = saved["model_config"]
        x_scaler = ColumnScaler.from_state_dict(saved["x_scaler"])
        y_scaler = ColumnScaler.from_state_dict(saved["y_scaler"])
    else:
        x_scaler, y_scaler = ColumnScaler.fit(x_train), ColumnScaler.fit(y_train)
    x_train, y_train = x_scaler.transform(x_train), y_scaler.transform(y_train)
    x_val, y_val = x_scaler.transform(x_val), y_scaler.transform(y_val)
    model = build_model(model_config["architecture"], model_config["width"], model_config["levels"], model_config["output_variables"]).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=float(settings.get("learning_rate", 0.001)))
    best_state = None
    if saved is not None:
        model.load_state_dict(saved.get("last_model", saved["model"]))
        optimizer.load_state_dict(saved["optimizer_state"])
        start_epoch, history, best_loss = saved["epoch"] + 1, saved["history"], saved["best_val_loss"]
        best_state = {key: value.detach().cpu().clone() for key, value in saved["model"].items()}
    sampler = torch.utils.data.distributed.DistributedSampler(TensorDataset(torch.from_numpy(x_train), torch.from_numpy(y_train)), shuffle=True) if distributed else None
    train_set = sampler.dataset if sampler else TensorDataset(torch.from_numpy(x_train), torch.from_numpy(y_train))
    train_loader = DataLoader(train_set, batch_size=int(settings.get("batch_size", 64)), sampler=sampler, shuffle=sampler is None, num_workers=int(settings.get("num_workers", 0)))
    val_loader = DataLoader(TensorDataset(torch.from_numpy(x_val), torch.from_numpy(y_val)), batch_size=int(settings.get("batch_size", 64)), shuffle=False)
    if distributed:
        model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
    patience, stale = int(settings.get("early_stopping_patience", 10)), 0
    for epoch in range(start_epoch, int(settings.get("epochs", 6))):
        if sampler is not None:
            sampler.set_epoch(epoch)
        train_loss = _epoch(model, train_loader, device, optimizer)
        val_loss = _epoch(model, val_loader, device, None)
        history.append({"epoch": epoch, "train_mse": train_loss, "val_mse": val_loss})
        if val_loss < best_loss:
            best_loss, stale = val_loss, 0
            best_state = {key: value.detach().cpu().clone() for key, value in (model.module if distributed else model).state_dict().items()}
        else:
            stale += 1
        if stale >= patience:
            break
    if rank == 0:
        checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
        final_model = model.module if distributed else model
        last_state = {key: value.detach().cpu().clone() for key, value in final_model.state_dict().items()}
        if best_state is not None:
            final_model.load_state_dict(best_state)
        checkpoint = {
            "format_version": CHECKPOINT_FORMAT_VERSION,
            "model_config": model_config,
            "model": best_state or last_state,
            "last_model": last_state,
            "optimizer_state": optimizer.state_dict(), "epoch": history[-1]["epoch"],
            "best_val_loss": best_loss, "history": history, "settings": settings, "x_scaler": x_scaler.state_dict(), "y_scaler": y_scaler.state_dict(),
            "random_state": {"python": random.getstate(), "numpy": np.random.get_state(), "torch": torch.get_rng_state()},
            "world_size": world_size,
        }
        torch.save(checkpoint, checkpoint_path)
        metrics_path = Path(metrics_path); metrics_path.parent.mkdir(parents=True, exist_ok=True)
        metrics_path.write_text(json.dumps({"architecture": model_config["architecture"], "parameters": sum(p.numel() for p in final_model.parameters()), "best_val_mse_normalized": best_loss, "epochs_completed": len(history), "history": history, "world_size": world_size}, indent=2), encoding="utf-8")
    if distributed:
        torch.distributed.barrier(); torch.distributed.destroy_process_group()
    return {"best_val_mse_normalized": best_loss, "epochs_completed": len(history)}


def run_inference(data_path: str | Path, checkpoint_path: str | Path, output_path: str | Path) -> dict[str, Any]:
    data = np.load(data_path)
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    if checkpoint.get("format_version") != CHECKPOINT_FORMAT_VERSION:
        raise ValueError(f"Unsupported checkpoint format_version: {checkpoint.get('format_version')!r}; expected {CHECKPOINT_FORMAT_VERSION!r}")
    settings = checkpoint["model_config"]
    model = build_model(settings["architecture"], int(settings["width"]), int(settings["levels"]), int(settings["output_variables"]))
    model.load_state_dict(checkpoint["model"]); model.eval()
    x_scaler, y_scaler = ColumnScaler.from_state_dict(checkpoint["x_scaler"]), ColumnScaler.from_state_dict(checkpoint["y_scaler"])
    with torch.no_grad():
        prediction_scaled = model(torch.from_numpy(x_scaler.transform(data["x_test"]))).numpy()
    prediction = y_scaler.inverse_transform(prediction_scaled)
    output_path = Path(output_path); output_path.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output_path, inputs=data["x_test"], targets=data["y_test"], predictions=prediction,
                        timestamps=data["time_test"], heights_m=data["heights_m"], pblh_m=data["pblh_test"],
                        input_names=np.asarray(INPUT_NAMES), output_names=np.asarray(OUTPUT_NAMES), architecture=np.asarray(settings["architecture"]))
    return {"samples": len(prediction), "shape": list(prediction.shape), "architecture": settings["architecture"]}


def _pearson(a: np.ndarray, b: np.ndarray) -> float:
    a, b = a.ravel(), b.ravel()
    if len(a) < 2 or np.std(a) < 1e-12 or np.std(b) < 1e-12:
        return 0.0
    return float(np.corrcoef(a, b)[0, 1])


def _scores(target: np.ndarray, prediction: np.ndarray) -> dict[str, float]:
    return {"rmse": float(np.sqrt(np.mean((prediction - target) ** 2))), "pearson": _pearson(target, prediction)}


def evaluate(predictions_path: str | Path, metrics_path: str | Path, figure_path: str | Path) -> dict[str, Any]:
    data = np.load(predictions_path)
    target, prediction = data["targets"], data["predictions"]
    heights, pblh = data["heights_m"], data["pblh_m"]
    mask = heights[None, :] <= pblh[:, None]
    metrics: dict[str, Any] = {
        "physical_scale": True,
        "primary_protocol": "synthetic virtual-level height <= synthetic PBLH mask",
        "primary": {"standardized_rmse_by_variable": {}},
        "full_17_level_diagnostics": {"by_variable": {}, "by_level_and_variable": {}},
    }
    for index, name in enumerate(OUTPUT_NAMES):
        masked_target, masked_prediction = target[:, :, index][mask], prediction[:, :, index][mask]
        scale = max(float(np.std(masked_target)), 1e-12)
        metrics["primary"]["standardized_rmse_by_variable"][name] = float(np.sqrt(np.mean(((masked_prediction - masked_target) / scale) ** 2)))
        metrics["full_17_level_diagnostics"]["by_variable"][name] = _scores(target[:, :, index], prediction[:, :, index])
    primary_values = metrics["primary"]["standardized_rmse_by_variable"].values()
    metrics["primary"]["macro_mean_standardized_rmse"] = float(np.mean(list(primary_values)))
    metrics["primary"]["definition"] = "Unweighted mean of per-variable RMSE divided by that variable's target standard deviation within the synthetic PBLH mask; no physical units are mixed."
    for level, height in enumerate(heights):
        metrics["full_17_level_diagnostics"]["by_level_and_variable"][str(level)] = {
            "height_m": float(height),
            "by_variable": {name: _scores(target[:, level, index], prediction[:, level, index]) for index, name in enumerate(OUTPUT_NAMES)},
        }
    speed_true = np.hypot(target[:, :, 0], target[:, :, 1]); speed_pred = np.hypot(prediction[:, :, 0], prediction[:, :, 1])
    direction_true = np.mod(1.5 * np.pi - np.arctan2(target[:, :, 1], target[:, :, 0]), 2 * np.pi)
    direction_pred = np.mod(1.5 * np.pi - np.arctan2(prediction[:, :, 1], prediction[:, :, 0]), 2 * np.pi)
    delta = np.arctan2(np.sin(direction_pred - direction_true), np.cos(direction_pred - direction_true))
    metrics["wind_speed"] = _scores(speed_true, speed_pred)
    metrics["wind_direction"] = {"convention": "meteorological direction from: 0 degrees from north, increasing clockwise", "circular_rmse_degrees": float(np.degrees(np.sqrt(np.mean(delta ** 2)))), "mean_absolute_circular_error_degrees": float(np.degrees(np.mean(np.abs(delta)))), "circular_correlation_cosine": float(np.mean(np.cos(delta)))}
    metrics["synthetic_pblh_mask"] = {"synthetic": True, "definition": "virtual level height <= synthetic PBLH; not a paper or real-WRF PBL mask", "sample_level_pairs": int(mask.sum())}
    metrics_path = Path(metrics_path); metrics_path.parent.mkdir(parents=True, exist_ok=True)
    metrics_path.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    try:
        import matplotlib.pyplot as plt
        figure_path = Path(figure_path); figure_path.parent.mkdir(parents=True, exist_ok=True)
        fig, axes = plt.subplots(1, 5, figsize=(15, 4), sharey=True)
        for index, (axis, name) in enumerate(zip(axes, OUTPUT_NAMES)):
            axis.plot(target[:, :, index].mean(0), heights, label="target")
            axis.plot(prediction[:, :, index].mean(0), heights, "--", label="prediction")
            axis.set_title(name); axis.grid(alpha=0.25)
        axes[0].set_ylabel("synthetic height (m)"); axes[-1].legend()
        fig.tight_layout(); fig.savefig(figure_path, dpi=150); plt.close(fig)
    except ImportError:
        metrics["figure_note"] = "matplotlib unavailable; numerical evaluation completed"
    return metrics