OneScience-Group/ML-MODIS

1#!/usr/bin/env python32"""Run all serialized trees and retain ensemble and per-tree predictions."""3 4from __future__ import annotations5 6import argparse7import sys8from pathlib import Path9 10import numpy as np11import torch12import yaml13 14ROOT = Path(__file__).resolve().parents[1]15sys.path.insert(0, str(ROOT / "model"))16from ml_modis import BootstrapRandomForestRegressor, validate_multimodal_keys17 18 19def main() -> None:20    parser = argparse.ArgumentParser()21    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))22    parser.add_argument("--data", default=None)23    parser.add_argument("--checkpoint", default=None)24    parser.add_argument("--output", default=None)25    args = parser.parse_args()26    config = yaml.safe_load(Path(args.config).read_text())27    with np.load(ROOT / (args.data or config["data"]["path"])) as archive:28        data = {key: archive[key] for key in archive.files}29    validate_multimodal_keys(data)30    checkpoint = torch.load(ROOT / (args.checkpoint or config["paths"]["checkpoint"]), map_location="cpu", weights_only=False)31    if checkpoint.get("format_version") != config["format_version"]:32        raise ValueError("Checkpoint format_version does not match configuration")33    targets = list(checkpoint["model_config"]["targets"])34    tree_count = len(next(iter(checkpoint["model"].values()))["state"]["trees"])35    tree_predictions = np.full((data["X"].shape[0], len(targets), tree_count), np.nan, dtype=np.float32)36    for month in checkpoint["model_config"]["months"]:37        mask = data["month"] == month38        for target_index, target in enumerate(targets):39            model = BootstrapRandomForestRegressor.from_state_dict(checkpoint["model"][f"{month}:{target}"]["state"])40            tree_predictions[mask, target_index, :] = model.predict_trees(data["X"][mask])41    prediction = tree_predictions.mean(axis=2)42    safe_prediction = np.where(np.abs(prediction) > 1e-8, prediction, np.nan)43    ratio = data["Y"] / safe_prediction44    if not np.isfinite(prediction).all() or not np.isfinite(ratio).all():45        raise FloatingPointError("Inference produced non-finite values")46    output = ROOT / (args.output or config["paths"]["predictions"])47    output.parent.mkdir(parents=True, exist_ok=True)48    np.savez_compressed(output, pred=prediction, pred_trees=tree_predictions, obs=data["Y"],49                        obs_over_pred=ratio, relative_response=ratio - 1.0,50                        year=data["year"], month=data["month"], platform=data["platform"],51                        latitude=data["latitude"], longitude=data["longitude"],52                        target_names=np.asarray(targets))53    print(f"output={output.relative_to(ROOT)} samples={prediction.shape[0]} "54          f"targets={targets} trees_per_prediction={tree_count}")55 56 57if __name__ == "__main__":58    main()59