"""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")