Publish MeteoNorm-RF reproduction

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
 *.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.png binary

conf/config.yaml

@@ -0,0 +1,29 @@
+seed: 2019
+data:
+  path: data/meteonorm_rf_hourly.npz
+  synthetic_days: 30
+  start_day_of_year: 1
+  stations: 12
+  pollutants: [PM2.5, PM10, NO2, SO2, O3, CO]
+  frequency: hourly
+  train_fraction: 0.70
+model:
+  engineering:
+    trees: 12
+    max_depth: 8
+    min_samples_leaf: 8
+    max_features: sqrt
+    split_candidates: 12
+paper_protocol:
+  trees: hundreds
+  normalization_resamples: 1000
+normalization:
+  engineering_resamples: 32
+  season_window_days: 14
+  match_hour: true
+  preserve_time_features: [ttrend, day_of_year, weekend, hour]
+paths:
+  checkpoint: result/checkpoints/meteonorm_rf.pkl
+  training_metrics: result/training/metrics.json
+  inference: result/output/predictions.npz
+  evaluation_dir: result/evaluation

config.json

@@ -0,0 +1 @@
+{"model_name":"MeteoNorm-RF","model_type":"meteonorm_rf","architectures":["MeteoNormRF"],"framework":"NumPy/PyTorch","domain":"air-quality","task":"meteorological-normalization","implementation":{"entry_point":"model/meteonorm_rf.py","scope":"core-method full-feature reduced-sample engineering reproduction","train_script":"scripts/train.py","inference_script":"scripts/inference.py","evaluation_script":"scripts/result.py","synthetic_data_script":"scripts/fake_data.py"},"architecture":{"input_features":21,"outputs":6,"engineering_trees":12,"paper_trees":"hundreds"},"data":{"dataset":"Beijing CNEM hourly air quality and meteorology","stations":12,"pollutants":["PM2.5","PM10","NO2","SO2","O3","CO"],"frequency":"hourly","synthetic":true},"configuration_sources":["conf/config.yaml","model/meteonorm_rf.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/meteonorm_rf.py

@@ -0,0 +1,189 @@
+"""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"]

scripts/fake_data.py

@@ -0,0 +1,60 @@
+"""Generate compact hourly data retaining all 12 stations and six pollutants."""
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.meteonorm_rf import BASE_FEATURES, FORMAT_VERSION, POLLUTANTS, STATIONS
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    rng = np.random.default_rng(int(config["seed"]))
+    days = int(config["data"]["synthetic_days"])
+    hours = np.tile(np.arange(days * 24), len(STATIONS))
+    station = np.repeat(np.arange(len(STATIONS)), days * 24)
+    day = hours // 24
+    hour = hours % 24
+    doy = day + int(config["data"]["start_day_of_year"])
+    weekend = ((day + 1) % 7 >= 5).astype(np.float32)
+    seasonal = np.sin(2 * np.pi * (doy - 172) / 365.25)
+    diurnal = np.sin(2 * np.pi * (hour - 14) / 24)
+    synoptic = np.sin(2 * np.pi * day / 6.5 + station * 0.17)
+    wind_speed = np.maximum(0.2, 2.7 + 1.1 * synoptic + rng.normal(0, .55, len(hours)))
+    wind_direction = np.mod(185 + 75 * synoptic + station * 9 + rng.normal(0, 28, len(hours)), 360)
+    temperature = 13 + 12 * seasonal + 5.5 * diurnal + rng.normal(0, 1.4, len(hours))
+    rh = np.clip(52 - 18 * seasonal - 10 * diurnal - 3 * wind_speed + rng.normal(0, 7, len(hours)), 15, 98)
+    pressure = 1013 - .32 * temperature + 4 * np.cos(2 * np.pi * day / 8) + rng.normal(0, 1.8, len(hours))
+    ttrend = 2013.0 + hours / (24 * 365.25)
+    emission = 1.0 - .22 * day / max(days - 1, 1)
+    stagnation = np.exp(-wind_speed / 3.2) * (0.65 + rh / 150)
+    rush = ((hour >= 7) & (hour <= 10)) | ((hour >= 17) & (hour <= 21))
+    station_factor = 0.82 + station * .025
+    pm25 = (22 + 78 * stagnation + 16 * rush) * emission * station_factor
+    pm10 = 28 + 1.22 * pm25 + 7 * wind_speed
+    no2 = (15 + 37 * stagnation + 19 * rush) * emission * station_factor
+    so2 = (4 + 25 * stagnation) * emission ** 1.8 * (1.1 - .15 * seasonal)
+    o3 = np.maximum(3, 30 + 29 * np.maximum(diurnal, -.4) + 1.2 * temperature - .28 * no2)
+    co = (0.25 + .012 * pm25 + .008 * no2) * emission
+    pollution = np.column_stack([pm25, pm10, no2, so2, o3, co])
+    pollution *= rng.lognormal(0, .09, pollution.shape)
+    output = ROOT / config["data"]["path"]
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, format_version=np.array(FORMAT_VERSION),
+                        station_names=np.asarray(STATIONS), pollutant_names=np.asarray(POLLUTANTS),
+                        station_id=station.astype(np.int16), timestamp_hour=hours.astype(np.int32),
+                        ttrend=ttrend.astype(np.float32), day_of_year=doy.astype(np.float32),
+                        weekend=weekend, hour=hour.astype(np.float32),
+                        wind_speed=wind_speed.astype(np.float32), wind_direction=wind_direction.astype(np.float32),
+                        pressure=pressure.astype(np.float32), temperature=temperature.astype(np.float32),
+                        relative_humidity=rh.astype(np.float32), pollution=pollution.astype(np.float32),
+                        feature_names=np.asarray(BASE_FEATURES))
+    print(f"data={output.relative_to(ROOT)} rows={len(hours)} stations=12 pollutants=6 hourly_days={days}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,53 @@
+"""Predict the test set and perform hour/season-window weather normalization."""
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.meteonorm_rf import BASE_FEATURES, encode_features, load_checkpoint
+
+
+def circular_day_distance(a, b):
+    distance = np.abs(a - b)
+    return np.minimum(distance, 365 - distance)
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    data = np.load(ROOT / config["data"]["path"])
+    model, metadata = load_checkpoint(ROOT / config["paths"]["checkpoint"])
+    x = encode_features(data)
+    test = np.asarray(metadata["test_indices"], dtype=np.int64)
+    observed = data["pollution"][test]
+    predicted = model.predict(x[test])
+    repeats = int(config["normalization"]["engineering_resamples"])
+    window = int(config["normalization"]["season_window_days"])
+    rng = np.random.default_rng(int(config["seed"]) + 17)
+    normalized = np.empty_like(predicted)
+    weather_columns = [BASE_FEATURES.index(name) for name in
+                       ("wind_speed", "wind_direction", "pressure", "temperature", "relative_humidity")]
+    all_hour, all_doy = data["hour"], data["day_of_year"]
+    for output_index, row_index in enumerate(test):
+        candidates = np.flatnonzero((all_hour == all_hour[row_index]) &
+                                    (circular_day_distance(all_doy, all_doy[row_index]) <= window))
+        draws = rng.choice(candidates, repeats, replace=True)
+        replicated = np.repeat(x[row_index:row_index + 1], repeats, axis=0)
+        replicated[:, weather_columns] = x[draws][:, weather_columns]
+        # ttrend, day-of-year, weekend, hour, and station remain at the target time/station.
+        normalized[output_index] = model.predict(replicated).mean(0)
+    output = ROOT / config["paths"]["inference"]
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, test_indices=test, observed=observed, predicted=predicted,
+                        normalized=normalized, ttrend=data["ttrend"][test],
+                        day_of_year=data["day_of_year"][test], hour=data["hour"][test],
+                        station_id=data["station_id"][test], pollutant_names=data["pollutant_names"],
+                        resamples=np.array(repeats), season_window_days=np.array(window))
+    print(f"predictions={output.relative_to(ROOT)} test_rows={len(test)} normalization_resamples={repeats}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,85 @@
+"""Compute paper metrics, Theil-Sen trends, and the comparison figure."""
+
+import json
+from pathlib import Path
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def metrics(observed, predicted):
+    error = predicted - observed
+    denominator = np.maximum(np.sum(observed, axis=0), 1e-12)
+    r2_denominator = np.maximum(np.sum((observed - observed.mean(0)) ** 2, axis=0), 1e-12)
+    ioa_denominator = np.maximum(np.sum((np.abs(predicted - observed.mean(0)) +
+                                         np.abs(observed - observed.mean(0))) ** 2, axis=0), 1e-12)
+    return {"RMSE": np.sqrt(np.mean(error ** 2, axis=0)),
+            "r2": 1 - np.sum(error ** 2, axis=0) / r2_denominator,
+            "FAC2": np.mean((predicted / np.maximum(observed, 1e-12) >= .5) &
+                            (predicted / np.maximum(observed, 1e-12) <= 2), axis=0),
+            "MB": np.mean(error, axis=0), "MGE": np.mean(np.abs(error), axis=0),
+            "NMB": np.sum(error, axis=0) / denominator,
+            "NMGE": np.sum(np.abs(error), axis=0) / denominator,
+            "COE": 1 - np.sum(np.abs(error), axis=0) /
+                   np.maximum(np.sum(np.abs(observed - observed.mean(0)), axis=0), 1e-12),
+            "IOA": 1 - np.sum(error ** 2, axis=0) / ioa_denominator}
+
+
+def theil_sen(x, y, max_pairs=200000, seed=0):
+    """Median pairwise slope, with deterministic subsampling for large series."""
+    x, y = np.asarray(x, float), np.asarray(y, float)
+    total = len(x) * (len(x) - 1) // 2
+    if total <= max_pairs:
+        slopes = [(y[j] - y[i]) / (x[j] - x[i]) for i in range(len(x) - 1)
+                  for j in range(i + 1, len(x)) if x[j] != x[i]]
+    else:
+        rng = np.random.default_rng(seed)
+        left = rng.integers(0, len(x) - 1, max_pairs)
+        right = rng.integers(left + 1, len(x), max_pairs)
+        valid = x[right] != x[left]
+        slopes = (y[right[valid]] - y[left[valid]]) / (x[right[valid]] - x[left[valid]])
+    slope = float(np.median(slopes))
+    return slope, float(np.median(y - slope * x))
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    data = np.load(ROOT / config["paths"]["inference"])
+    names = [str(value) for value in data["pollutant_names"]]
+    scores = metrics(data["observed"], data["predicted"])
+    report = {"metrics": {name: {metric: float(values[i]) for metric, values in scores.items()}
+                           for i, name in enumerate(names)}, "theil_sen_per_year": {}}
+    trends = []
+    for i, name in enumerate(names):
+        raw = theil_sen(data["ttrend"], data["observed"][:, i], seed=i)
+        normalized = theil_sen(data["ttrend"], data["normalized"][:, i], seed=100 + i)
+        report["theil_sen_per_year"][name] = {"observed_slope": raw[0],
+                                                       "normalized_slope": normalized[0]}
+        trends.append((raw[0], normalized[0]))
+    output = ROOT / config["paths"]["evaluation_dir"]
+    output.mkdir(parents=True, exist_ok=True)
+    (output / "metrics.json").write_text(json.dumps(report, indent=2) + "\n")
+    figure, axes = plt.subplots(2, 1, figsize=(11, 8), gridspec_kw={"height_ratios": [2, 1]})
+    order = np.argsort(data["ttrend"])
+    sample = order[::max(1, len(order) // 600)]
+    axes[0].plot(data["ttrend"][sample], data["observed"][sample, 0], ".", alpha=.25, label="observed PM2.5")
+    axes[0].plot(data["ttrend"][sample], data["normalized"][sample, 0], ".", alpha=.4, label="weather-normalized PM2.5")
+    axes[0].set(ylabel="Concentration (ug m-3)", title="Observed and meteorologically normalized test samples")
+    axes[0].legend()
+    positions = np.arange(len(names)); width = .36
+    axes[1].bar(positions - width / 2, [v[0] for v in trends], width, label="observed")
+    axes[1].bar(positions + width / 2, [v[1] for v in trends], width, label="normalized")
+    axes[1].axhline(0, color="black", linewidth=.7); axes[1].set_xticks(positions, names)
+    axes[1].set(ylabel="Theil-Sen slope per year", title="Robust concentration trends")
+    axes[1].legend(); figure.tight_layout(); figure.savefig(output / "comparison.png", dpi=160)
+    plt.close(figure)
+    print(f"evaluation={output.relative_to(ROOT)} pollutants={len(names)}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,71 @@
+"""Train a pure NumPy RF; torchrun ranks build disjoint tree subsets."""
+
+import json
+import os
+import pickle
+import sys
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.meteonorm_rf import (FEATURE_NAMES, FORMAT_VERSION, MODEL_NAME,
+                                  MultiOutputRandomForest, MeteoNormRF,
+                                  encode_features, merge_states, save_checkpoint)
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    seed = int(config["seed"])
+    data = np.load(ROOT / config["data"]["path"])
+    x, y = encode_features(data), data["pollution"].astype(np.float32)
+    rng = np.random.default_rng(seed)
+    order = rng.permutation(len(x))
+    cut = int(float(config["data"]["train_fraction"]) * len(order))
+    train_indices, test_indices = order[:cut], order[cut:]
+    world, rank = int(os.environ.get("WORLD_SIZE", "1")), int(os.environ.get("RANK", "0"))
+    options = config["model"]["engineering"]
+    trees = int(options["trees"])
+    assigned = list(range(rank, trees, world))
+    forest = MultiOutputRandomForest(n_trees=trees, seed=seed,
+                                     max_depth=int(options["max_depth"]),
+                                     min_samples_leaf=int(options["min_samples_leaf"]),
+                                     max_features=options["max_features"],
+                                     split_candidates=int(options["split_candidates"]))
+    model = MeteoNormRF(forest).fit(x[train_indices], y[train_indices], assigned)
+    checkpoint = ROOT / config["paths"]["checkpoint"]
+    checkpoint.parent.mkdir(parents=True, exist_ok=True)
+    shard = checkpoint.with_suffix(f".rank{rank}.pkl")
+    with open(shard, "wb") as stream:
+        pickle.dump(model.state_dict(), stream)
+    if world > 1:
+        import torch
+        torch.distributed.init_process_group("gloo")
+        torch.distributed.barrier()
+    if rank == 0:
+        states = []
+        for item in range(world):
+            with open(checkpoint.with_suffix(f".rank{item}.pkl"), "rb") as stream:
+                states.append(pickle.load(stream))
+        merged = MeteoNormRF.from_state_dict(merge_states(states))
+        metadata = {"feature_names": FEATURE_NAMES, "train_indices": train_indices,
+                    "test_indices": test_indices, "split": "seeded random 70/30",
+                    "paper_trees": config["paper_protocol"]["trees"], "engineering_trees": trees}
+        save_checkpoint(checkpoint, merged, metadata)
+        metrics = ROOT / config["paths"]["training_metrics"]
+        metrics.parent.mkdir(parents=True, exist_ok=True)
+        metrics.write_text(json.dumps({"format_version": FORMAT_VERSION, "model": MODEL_NAME,
+                                       "rows": len(x), "train_rows": cut, "test_rows": len(x) - cut,
+                                       "trees": trees, "world_size": world}, indent=2) + "\n")
+        for item in range(world):
+            checkpoint.with_suffix(f".rank{item}.pkl").unlink()
+        print(f"checkpoint={checkpoint.relative_to(ROOT)} trees={trees} train={cut} test={len(x)-cut}")
+    if world > 1:
+        torch.distributed.barrier()
+        torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()