"""Pure NumPy multi-output random forest for Beijing air-quality normalization.""" from dataclasses import dataclass import pickle import numpy as np FORMAT_VERSION = "meteonorm_rf_v1" MODEL_NAME = "MeteoNorm-RF" POLLUTANTS = ("PM2.5", "PM10", "NO2", "SO2", "O3", "CO") STATIONS = ("dongsi", "tiantan", "guanyuan", "wanshouxigong", "aotizhongxin", "nongzhanguan", "wanliu", "beibuxinqu", "zhiwuyuan", "fengtaihuayuan", "yungang", "gucheng") BASE_FEATURES = ("ttrend", "day_of_year", "weekend", "hour", "wind_speed", "wind_direction", "pressure", "temperature", "relative_humidity") FEATURE_NAMES = BASE_FEATURES + tuple(f"station={name}" for name in STATIONS) def encode_features(data): """Encode station as an explicit 12-column one-hot categorical feature.""" station = np.asarray(data["station_id"], dtype=np.int64) if np.any((station < 0) | (station >= len(STATIONS))): raise ValueError("station_id must be in [0, 11]") numeric = np.column_stack([np.asarray(data[name], dtype=np.float32) for name in BASE_FEATURES]) one_hot = np.eye(len(STATIONS), dtype=np.float32)[station] result = np.concatenate([numeric, one_hot], axis=1) if result.shape[1] != len(FEATURE_NAMES) or not np.isfinite(result).all(): raise ValueError("features must be finite and follow the 21-column protocol") return result.astype(np.float32) @dataclass class TreeConfig: max_depth: int = 8 min_samples_leaf: int = 8 max_features: object = "sqrt" split_candidates: int = 12 class RegressionTree: def __init__(self, config, seed=0): self.config = config self.rng = np.random.default_rng(seed) self.nodes = [] def _feature_count(self, total): value = self.config.max_features if value == "sqrt": return max(1, int(np.sqrt(total))) if isinstance(value, float): return max(1, min(total, int(np.ceil(total * value)))) return max(1, min(total, int(value))) def fit(self, x, y): self.nodes = [] self._grow(np.asarray(x, np.float32), np.asarray(y, np.float32), np.arange(len(x)), 0) return self def _grow(self, x, y, indices, depth): node_id = len(self.nodes) self.nodes.append(None) value = y[indices].mean(0).astype(np.float32) minimum = int(self.config.min_samples_leaf) if depth >= self.config.max_depth or len(indices) < 2 * minimum: self.nodes[node_id] = {"value": value} return node_id best = None features = self.rng.choice(x.shape[1], self._feature_count(x.shape[1]), replace=False) for feature in features: values = x[indices, feature] low, high = float(values.min()), float(values.max()) if low >= high: continue for threshold in self.rng.uniform(low, high, self.config.split_candidates): mask = values <= threshold left, right = indices[mask], indices[~mask] if len(left) < minimum or len(right) < minimum: continue loss = (np.square(y[left] - y[left].mean(0)).sum() + np.square(y[right] - y[right].mean(0)).sum()) if best is None or loss < best[0]: best = (float(loss), int(feature), float(threshold), left, right) if best is None: self.nodes[node_id] = {"value": value} return node_id _, feature, threshold, left, right = best self.nodes[node_id] = {"feature": feature, "threshold": threshold, "left": self._grow(x, y, left, depth + 1), "right": self._grow(x, y, right, depth + 1)} return node_id def predict(self, x): output = [] for row in np.asarray(x, np.float32): node = self.nodes[0] while "value" not in node: branch = "left" if row[node["feature"]] <= node["threshold"] else "right" node = self.nodes[node[branch]] output.append(node["value"]) return np.asarray(output, dtype=np.float32) class MultiOutputRandomForest: def __init__(self, n_trees=12, seed=2019, **tree_options): self.n_trees, self.seed = int(n_trees), int(seed) self.tree_config = TreeConfig(**tree_options) self.trees = [] def fit(self, x, y, tree_indices=None): x, y = np.asarray(x, np.float32), np.asarray(y, np.float32) indices = range(self.n_trees) if tree_indices is None else tree_indices self.trees = [] for index in indices: rng = np.random.default_rng(self.seed + 7919 * (index + 1)) sample = rng.integers(0, len(x), len(x)) tree = RegressionTree(self.tree_config, self.seed + 104729 * (index + 1)) self.trees.append((int(index), tree.fit(x[sample], y[sample]))) return self def predict(self, x): if not self.trees: raise RuntimeError("forest is not fitted") return np.mean([tree.predict(x) for _, tree in self.trees], axis=0, dtype=np.float32) def state_dict(self): return {"n_trees": self.n_trees, "seed": self.seed, "tree_config": vars(self.tree_config), "trees": [(index, tree.nodes) for index, tree in self.trees]} @classmethod def from_state_dict(cls, state): model = cls(state["n_trees"], state["seed"], **state["tree_config"]) for index, nodes in state["trees"]: tree = RegressionTree(model.tree_config) tree.nodes = nodes model.trees.append((index, tree)) return model class MeteoNormRF: def __init__(self, forest): self.forest = forest self.statistics = {} def fit(self, x, y, tree_indices=None): x, y = np.asarray(x, np.float32), np.asarray(y, np.float32) self.statistics = {"x_mean": x.mean(0), "x_std": np.maximum(x.std(0), 1e-6), "y_mean": y.mean(0), "y_std": np.maximum(y.std(0), 1e-6)} self.forest.fit((x - self.statistics["x_mean"]) / self.statistics["x_std"], (y - self.statistics["y_mean"]) / self.statistics["y_std"], tree_indices) return self def predict(self, x): x = np.asarray(x, np.float32) value = self.forest.predict((x - self.statistics["x_mean"]) / self.statistics["x_std"]) return np.maximum(value * self.statistics["y_std"] + self.statistics["y_mean"], 0.0) def state_dict(self): return {"forest": self.forest.state_dict(), "statistics": self.statistics} @classmethod def from_state_dict(cls, state): model = cls(MultiOutputRandomForest.from_state_dict(state["forest"])) model.statistics = state["statistics"] return model def merge_states(states): base = states[0] trees = [item for state in states for item in state["forest"]["trees"]] base["forest"]["trees"] = sorted(trees, key=lambda item: item[0]) return base def save_checkpoint(path, model, metadata): with open(path, "wb") as stream: pickle.dump({"format_version": FORMAT_VERSION, "model_name": MODEL_NAME, "model": model.state_dict(), "metadata": metadata}, stream) def load_checkpoint(path): with open(path, "rb") as stream: state = pickle.load(stream) if state.get("format_version") != FORMAT_VERSION or state.get("model_name") != MODEL_NAME: raise ValueError("incompatible checkpoint") return MeteoNormRF.from_state_dict(state["model"]), state["metadata"]