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