model/ml_modis.py

"""Pure NumPy bootstrap random-forest regression used by ML-MODIS."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Sequence, Tuple

import numpy as np


TARGETS = ("Nd", "reff", "LWP", "CF")
PRESSURE_VARIABLES = ("temperature", "specific_humidity", "relative_humidity", "u_wind", "v_wind", "omega", "geopotential", "cloud_liquid", "cloud_fraction")
PRESSURE_LEVELS = (1000, 950, 900, 850, 800, 750, 700, 650, 600, 550)
SINGLE_FEATURES = (
    "sst", "surface_pressure", "mslp", "skin_temperature", "t2m", "d2m",
    "u10", "v10", "surface_solar_radiation", "surface_thermal_radiation",
    "latent_heat_flux", "sensible_heat_flux", "boundary_layer_height",
    "total_column_water_vapour", "total_column_cloud_liquid", "cape", "cin",
    "low_cloud_cover", "sea_ice_fraction", "precipitation", "cos_sza",
    "latitude", "longitude", "platform_hour",
)


def feature_names() -> List[str]:
    names = [f"{variable}_{level}hPa" for variable in PRESSURE_VARIABLES for level in PRESSURE_LEVELS]
    names.extend(SINGLE_FEATURES)
    if len(names) != 114:
        raise RuntimeError("The ERA5 predictor ledger must contain exactly 114 features")
    return names


def regression_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> Dict[str, float]:
    mask = np.isfinite(y_true) & np.isfinite(y_pred)
    if mask.sum() < 2:
        return {"n": int(mask.sum()), "mse": float("nan"), "r2": float("nan"), "pearson": float("nan")}
    y = np.asarray(y_true[mask], dtype=np.float64)
    p = np.asarray(y_pred[mask], dtype=np.float64)
    mse = float(np.mean((y - p) ** 2))
    variance = float(np.sum((y - y.mean()) ** 2))
    r2 = float(1.0 - np.sum((y - p) ** 2) / variance) if variance > 0 else float("nan")
    pearson = float(np.corrcoef(y, p)[0, 1]) if np.std(y) > 0 and np.std(p) > 0 else float("nan")
    return {"n": int(mask.sum()), "mse": mse, "r2": r2, "pearson": pearson}


@dataclass
class TreeConfig:
    min_leaf: int = 7
    max_features: int = 38
    max_depth: Optional[int] = None
    split_candidates: int = 12


class RandomRegressionTree:
    """CART regressor with random feature subsets and compact array state."""

    def __init__(self, config: TreeConfig, seed: int):
        self.config = config
        self.seed = int(seed)
        self.feature: List[int] = []
        self.threshold: List[float] = []
        self.left: List[int] = []
        self.right: List[int] = []
        self.value: List[float] = []

    def fit(self, x: np.ndarray, y: np.ndarray) -> "RandomRegressionTree":
        x = np.asarray(x, dtype=np.float32)
        y = np.asarray(y, dtype=np.float64)
        rng = np.random.default_rng(self.seed)

        def build(indices: np.ndarray, depth: int) -> int:
            node = len(self.value)
            self.feature.append(-1)
            self.threshold.append(np.nan)
            self.left.append(-1)
            self.right.append(-1)
            self.value.append(float(y[indices].mean()))
            if indices.size < 2 * self.config.min_leaf:
                return node
            if self.config.max_depth is not None and depth >= self.config.max_depth:
                return node
            parent_sse = float(np.sum((y[indices] - y[indices].mean()) ** 2))
            if parent_sse <= 1e-12:
                return node
            n_features = min(self.config.max_features, x.shape[1])
            candidates = rng.choice(x.shape[1], size=n_features, replace=False)
            best: Optional[Tuple[float, int, float, np.ndarray]] = None
            quantiles = np.linspace(0.05, 0.95, self.config.split_candidates)
            for feature in candidates:
                values = x[indices, feature]
                thresholds = np.unique(np.quantile(values, quantiles))
                for threshold in thresholds:
                    is_left = values <= threshold
                    nl = int(is_left.sum())
                    nr = indices.size - nl
                    if nl < self.config.min_leaf or nr < self.config.min_leaf:
                        continue
                    yl, yr = y[indices[is_left]], y[indices[~is_left]]
                    score = float(np.sum((yl - yl.mean()) ** 2) + np.sum((yr - yr.mean()) ** 2))
                    if best is None or score < best[0]:
                        best = (score, int(feature), float(threshold), is_left.copy())
            if best is None or best[0] >= parent_sse - 1e-12:
                return node
            _, split_feature, split_threshold, is_left = best
            self.feature[node] = split_feature
            self.threshold[node] = split_threshold
            self.left[node] = build(indices[is_left], depth + 1)
            self.right[node] = build(indices[~is_left], depth + 1)
            return node

        build(np.arange(y.size, dtype=np.int64), 0)
        return self

    def predict(self, x: np.ndarray) -> np.ndarray:
        x = np.asarray(x, dtype=np.float32)
        output = np.empty(x.shape[0], dtype=np.float32)
        for row in range(x.shape[0]):
            node = 0
            while self.feature[node] >= 0:
                node = self.left[node] if x[row, self.feature[node]] <= self.threshold[node] else self.right[node]
            output[row] = self.value[node]
        return output

    def state_dict(self) -> Dict[str, Any]:
        return {
            "seed": self.seed,
            "config": self.config.__dict__.copy(),
            "feature": np.asarray(self.feature, dtype=np.int32),
            "threshold": np.asarray(self.threshold, dtype=np.float32),
            "left": np.asarray(self.left, dtype=np.int32),
            "right": np.asarray(self.right, dtype=np.int32),
            "value": np.asarray(self.value, dtype=np.float32),
        }

    @classmethod
    def from_state_dict(cls, state: Dict[str, Any]) -> "RandomRegressionTree":
        tree = cls(TreeConfig(**state["config"]), int(state["seed"]))
        for name in ("feature", "threshold", "left", "right", "value"):
            setattr(tree, name, np.asarray(state[name]).tolist())
        return tree


class BootstrapRandomForestRegressor:
    """Regression forest with explicit approximately 60% bootstrap and OOB state."""

    def __init__(self, n_trees: int = 100, min_leaf: int = 7, max_features: int = 38,
                 bootstrap_fraction: float = 0.6, max_depth: Optional[int] = None,
                 split_candidates: int = 12, seed: int = 0):
        if n_trees < 1 or min_leaf < 1 or not 0 < bootstrap_fraction <= 1:
            raise ValueError("Invalid forest configuration")
        self.n_trees = int(n_trees)
        self.bootstrap_fraction = float(bootstrap_fraction)
        self.seed = int(seed)
        self.tree_config = TreeConfig(int(min_leaf), int(max_features), max_depth, int(split_candidates))
        self.trees: List[RandomRegressionTree] = []
        self.oob_indices: List[np.ndarray] = []

    def fit(self, x: np.ndarray, y: np.ndarray) -> "BootstrapRandomForestRegressor":
        x = np.asarray(x, dtype=np.float32)
        y = np.asarray(y, dtype=np.float32)
        if x.ndim != 2 or x.shape[1] != 114 or y.shape != (x.shape[0],):
            raise ValueError(f"Expected X [N,114] and y [N], got {x.shape} and {y.shape}")
        rng = np.random.default_rng(self.seed)
        draw_size = max(2 * self.tree_config.min_leaf, int(round(self.bootstrap_fraction * x.shape[0])))
        self.trees, self.oob_indices = [], []
        for _ in range(self.n_trees):
            bootstrap = rng.integers(0, x.shape[0], size=draw_size)
            used = np.zeros(x.shape[0], dtype=bool)
            used[np.unique(bootstrap)] = True
            oob = np.flatnonzero(~used)
            tree_seed = int(rng.integers(0, 2**31 - 1))
            self.trees.append(RandomRegressionTree(self.tree_config, tree_seed).fit(x[bootstrap], y[bootstrap]))
            self.oob_indices.append(oob.astype(np.int32))
        return self

    def predict_trees(self, x: np.ndarray) -> np.ndarray:
        if not self.trees:
            raise RuntimeError("Forest is not fitted")
        return np.stack([tree.predict(x) for tree in self.trees], axis=1)

    def predict(self, x: np.ndarray) -> np.ndarray:
        return self.predict_trees(x).mean(axis=1)

    def oob_predict(self, x: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        sums = np.zeros(x.shape[0], dtype=np.float64)
        counts = np.zeros(x.shape[0], dtype=np.int32)
        for tree, indices in zip(self.trees, self.oob_indices):
            if indices.size:
                sums[indices] += tree.predict(x[indices])
                counts[indices] += 1
        prediction = np.full(x.shape[0], np.nan, dtype=np.float32)
        valid = counts > 0
        prediction[valid] = (sums[valid] / counts[valid]).astype(np.float32)
        return prediction, counts

    def permutation_importance(self, x: np.ndarray, y: np.ndarray, seed: int = 0) -> np.ndarray:
        """Breiman OOB permuted-predictor delta MSE, averaged over eligible trees."""
        rng = np.random.default_rng(seed)
        deltas = np.zeros(x.shape[1], dtype=np.float64)
        counts = np.zeros(x.shape[1], dtype=np.int32)
        for tree, indices in zip(self.trees, self.oob_indices):
            if indices.size < 2:
                continue
            xo = np.asarray(x[indices], dtype=np.float32)
            yo = np.asarray(y[indices], dtype=np.float32)
            baseline = float(np.mean((yo - tree.predict(xo)) ** 2))
            for feature in range(x.shape[1]):
                changed = xo.copy()
                changed[:, feature] = changed[rng.permutation(indices.size), feature]
                deltas[feature] += float(np.mean((yo - tree.predict(changed)) ** 2)) - baseline
                counts[feature] += 1
        return np.divide(deltas, counts, out=np.zeros_like(deltas), where=counts > 0).astype(np.float32)

    def state_dict(self) -> Dict[str, Any]:
        return {
            "n_trees": self.n_trees,
            "bootstrap_fraction": self.bootstrap_fraction,
            "seed": self.seed,
            "tree_config": self.tree_config.__dict__.copy(),
            "trees": [tree.state_dict() for tree in self.trees],
            "oob_indices": self.oob_indices,
        }

    @classmethod
    def from_state_dict(cls, state: Dict[str, Any]) -> "BootstrapRandomForestRegressor":
        config = state["tree_config"]
        forest = cls(state["n_trees"], config["min_leaf"], config["max_features"],
                     state["bootstrap_fraction"], config["max_depth"],
                     config["split_candidates"], state["seed"])
        forest.trees = [RandomRegressionTree.from_state_dict(item) for item in state["trees"]]
        forest.oob_indices = [np.asarray(item, dtype=np.int32) for item in state["oob_indices"]]
        return forest


def validate_multimodal_keys(data: Dict[str, np.ndarray]) -> None:
    required = ("year", "month", "platform", "latitude", "longitude", "X", "Y")
    missing = [key for key in required if key not in data]
    if missing:
        raise ValueError(f"Missing aligned arrays: {missing}")
    n = data["X"].shape[0]
    if data["X"].shape[1] != 114 or data["Y"].shape != (n, 4):
        raise ValueError("Predictors must be [N,114] and targets [N,4]")
    if any(np.asarray(data[key]).shape[0] != n for key in required[:-2]):
        raise ValueError("Year/month/platform/coordinates are not row-aligned")
    keys = list(zip(data["year"].tolist(), data["month"].tolist(), data["platform"].tolist(),
                    np.round(data["latitude"], 4).tolist(), np.round(data["longitude"], 4).tolist()))
    if len(set(keys)) != n:
        raise ValueError("Multimodal year-month-platform-latitude-longitude keys are not unique")