OneScience-Group/RF-ClimParam

1"""Pure NumPy joint multi-output random forests for the SAM parameterization."""2 3from dataclasses import dataclass4 5import numpy as np6 7 8FORMAT_VERSION = "rf_climparam_v2"9MODEL_NAME = "RF-ClimParam"10SCALES = ("x4", "x8", "x16", "x32")11TEND_INPUT_NAMES = tuple([f"T_{i:02d}" for i in range(48)] +12                         [f"qT_{i:02d}" for i in range(48)] +13                         [f"qp_{i:02d}" for i in range(48)] + ["abs_y"])14TEND_OUTPUT_NAMES = tuple([f"hL_tend_{i:02d}" for i in range(48)] +15                          [f"qT_tend_{i:02d}" for i in range(48)] +16                          [f"qp_tend_{i:02d}" for i in range(48)])17DIFF_INPUT_NAMES = tuple([f"T_low_{i:02d}" for i in range(15)] +18                         [f"qT_low_{i:02d}" for i in range(15)] +19                         [f"u_low_{i:02d}" for i in range(15)] +20                         [f"v_nh_low_{i:02d}" for i in range(15)] +21                         ["windsurf", "abs_y"])22DIFF_OUTPUT_NAMES = tuple([f"Dbar_{i:02d}" for i in range(15)] +23                          ["hL_surface_flux", "qT_surface_flux"])24 25 26def _check_array(name, value, features):27    value = np.asarray(value)28    if value.ndim != 2 or value.shape[1] != features:29        raise ValueError(f"{name} must have shape [N,{features}], got {value.shape}")30    if value.dtype not in (np.float32, np.float64) or not np.isfinite(value).all():31        raise ValueError(f"{name} must be finite float32/float64")32    return value.astype(np.float32, copy=False)33 34 35@dataclass36class TreeConfig:37    max_depth: int = 538    min_samples_leaf: int = 339    max_features: object = "sqrt"40    split_candidates: int = 841 42 43class ExtraRandomRegressionTree:44    """Randomized recursive tree whose leaves hold one joint output vector."""45 46    def __init__(self, config, seed=0):47        self.config = config48        self.rng = np.random.default_rng(seed)49        self.nodes = []50 51    def fit(self, x, y):52        x, y = np.asarray(x, np.float32), np.asarray(y, np.float32)53        self.nodes = []54        self._grow(x, y, np.arange(len(x)), 0)55        return self56 57    def _feature_count(self, total):58        value = self.config.max_features59        if value == "sqrt":60            return max(1, int(np.sqrt(total)))61        if value == "log2":62            return max(1, int(np.log2(total)))63        if isinstance(value, float):64            return max(1, min(total, int(np.ceil(value * total))))65        return max(1, min(total, int(value)))66 67    def _grow(self, x, y, indices, depth):68        node_id = len(self.nodes)69        self.nodes.append(None)70        leaf_value = y[indices].mean(axis=0).astype(np.float32)71        minimum = int(self.config.min_samples_leaf)72        if depth >= int(self.config.max_depth) or len(indices) < 2 * minimum:73            self.nodes[node_id] = {"value": leaf_value}74            return node_id75        features = self.rng.choice(x.shape[1], self._feature_count(x.shape[1]), replace=False)76        best = None77        parent_sse = float(np.square(y[indices] - leaf_value).sum())78        for feature in features:79            values = x[indices, feature]80            low, high = float(values.min()), float(values.max())81            if not low < high:82                continue83            thresholds = self.rng.uniform(low, high, int(self.config.split_candidates))84            for threshold in thresholds:85                mask = values <= threshold86                left, right = indices[mask], indices[~mask]87                if len(left) < minimum or len(right) < minimum:88                    continue89                left_mean, right_mean = y[left].mean(0), y[right].mean(0)90                loss = float(np.square(y[left] - left_mean).sum() +91                             np.square(y[right] - right_mean).sum())92                if best is None or loss < best[0]:93                    best = (loss, int(feature), float(threshold), left, right)94        if best is None or best[0] >= parent_sse - 1e-10:95            self.nodes[node_id] = {"value": leaf_value}96            return node_id97        _, feature, threshold, left, right = best98        self.nodes[node_id] = {"feature": feature, "threshold": threshold,99                               "left": self._grow(x, y, left, depth + 1),100                               "right": self._grow(x, y, right, depth + 1)}101        return node_id102 103    def predict(self, x):104        outputs = []105        for row in np.asarray(x):106            node = self.nodes[0]107            while "value" not in node:108                node = self.nodes[node["left"] if row[node["feature"]] <= node["threshold"] else node["right"]]109            outputs.append(node["value"])110        return np.asarray(outputs, dtype=np.float32)111 112    def state_dict(self):113        return {"config": vars(self.config), "nodes": self.nodes}114 115    @classmethod116    def from_state_dict(cls, state):117        tree = cls(TreeConfig(**state["config"]))118        tree.nodes = state["nodes"]119        return tree120 121 122class JointRandomForestRegressor:123    """Bootstrap ensemble retaining inseparable multi-output leaf predictions."""124 125    def __init__(self, n_trees=2, seed=0, **tree_options):126        self.n_trees, self.seed = int(n_trees), int(seed)127        self.tree_config = TreeConfig(**tree_options)128        self.trees = []129 130    def fit(self, x, y):131        x, y = np.asarray(x, np.float32), np.asarray(y, np.float32)132        rng = np.random.default_rng(self.seed)133        self.trees = []134        for index in range(self.n_trees):135            bootstrap = rng.integers(0, len(x), size=len(x))136            tree = ExtraRandomRegressionTree(self.tree_config, self.seed + 1009 * (index + 1))137            self.trees.append(tree.fit(x[bootstrap], y[bootstrap]))138        return self139 140    def predict(self, x):141        if not self.trees:142            raise RuntimeError("forest is not fitted")143        return np.mean([tree.predict(x) for tree in self.trees], axis=0, dtype=np.float32)144 145    def state_dict(self):146        return {"n_trees": self.n_trees, "seed": self.seed,147                "tree_config": vars(self.tree_config),148                "trees": [tree.state_dict() for tree in self.trees]}149 150    @classmethod151    def from_state_dict(cls, state):152        forest = cls(state["n_trees"], state["seed"], **state["tree_config"])153        forest.trees = [ExtraRandomRegressionTree.from_state_dict(item) for item in state["trees"]]154        return forest155 156 157class StandardizedForest:158    """Block-standardized wrapper; one scalar mean/std is used per variable block."""159 160    def __init__(self, forest, input_slices, output_slices, nonnegative_slice=None):161        self.forest = forest162        self.input_slices, self.output_slices = input_slices, output_slices163        self.nonnegative_slice = nonnegative_slice164        self.statistics = {}165 166    @staticmethod167    def _statistics(array, slices):168        means, stds = np.zeros(array.shape[1], np.float32), np.ones(array.shape[1], np.float32)169        for start, stop in slices:170            mean = float(array[:, start:stop].mean())171            std = max(float(array[:, start:stop].std()), 1e-6)172            means[start:stop], stds[start:stop] = mean, std173        return means, stds174 175    def fit(self, x, y):176        x = _check_array("inputs", x, self.input_slices[-1][1])177        y = _check_array("targets", y, self.output_slices[-1][1])178        x_mean, x_std = self._statistics(x, self.input_slices)179        y_mean, y_std = self._statistics(y, self.output_slices)180        self.statistics = {"input_mean": x_mean, "input_std": x_std,181                           "output_mean": y_mean, "output_std": y_std}182        self.forest.fit((x - x_mean) / x_std, (y - y_mean) / y_std)183        return self184 185    def predict(self, x):186        x = _check_array("inputs", x, len(self.statistics["input_mean"]))187        prediction = self.forest.predict((x - self.statistics["input_mean"]) / self.statistics["input_std"])188        prediction = prediction * self.statistics["output_std"] + self.statistics["output_mean"]189        if self.nonnegative_slice is not None:190            prediction[:, self.nonnegative_slice[0]:self.nonnegative_slice[1]] = np.maximum(191                prediction[:, self.nonnegative_slice[0]:self.nonnegative_slice[1]], 0.0)192        return prediction.astype(np.float32)193 194    def state_dict(self):195        return {"forest": self.forest.state_dict(), "input_slices": self.input_slices,196                "output_slices": self.output_slices, "nonnegative_slice": self.nonnegative_slice,197                "statistics": self.statistics}198 199    @classmethod200    def from_state_dict(cls, state):201        model = cls(JointRandomForestRegressor.from_state_dict(state["forest"]),202                    state["input_slices"], state["output_slices"], state["nonnegative_slice"])203        model.statistics = state["statistics"]204        return model205 206 207def build_pair(config, seed):208    options = config["engineering"]209    common = {"n_trees": options["trees"], "max_depth": options["max_depth"],210              "min_samples_leaf": options["min_samples_leaf"],211              "max_features": options["max_features"], "split_candidates": options["split_candidates"]}212    tend = StandardizedForest(JointRandomForestRegressor(seed=seed, **common),213                              [(0, 48), (48, 96), (96, 144), (144, 145)],214                              [(0, 48), (48, 96), (96, 144)])215    diff = StandardizedForest(JointRandomForestRegressor(seed=seed + 1, **common),216                              [(0, 15), (15, 30), (30, 45), (45, 60), (60, 61), (61, 62)],217                              [(0, 15), (15, 16), (16, 17)], nonnegative_slice=(0, 15))218    return {"rf_tend": tend, "rf_diff": diff}219 220 221def load_models(checkpoint):222    if checkpoint.get("format_version") != FORMAT_VERSION or checkpoint.get("model_name") != MODEL_NAME:223        raise ValueError("incompatible checkpoint model/format_version")224    if not isinstance(checkpoint.get("model"), dict) or set(checkpoint["model"]) != set(SCALES):225        raise ValueError("checkpoint model must contain all four scale forest states")226    expected = {"rf_tend_input": 145, "rf_tend_output": 144,227                "rf_diff_input": 62, "rf_diff_output": 17}228    if checkpoint.get("model_config", {}).get("dimensions") != expected:229        raise ValueError("checkpoint model_config dimensions are incompatible")230    return {scale: {name: StandardizedForest.from_state_dict(state)231                    for name, state in pair.items()}232            for scale, pair in checkpoint["model"].items()}233