| |
| """Run all serialized trees and retain ensemble and per-tree predictions.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import yaml |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT / "model")) |
| from ml_modis import BootstrapRandomForestRegressor, validate_multimodal_keys |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config", default=str(ROOT / "conf/config.yaml")) |
| parser.add_argument("--data", default=None) |
| parser.add_argument("--checkpoint", default=None) |
| parser.add_argument("--output", default=None) |
| args = parser.parse_args() |
| config = yaml.safe_load(Path(args.config).read_text()) |
| with np.load(ROOT / (args.data or config["data"]["path"])) as archive: |
| data = {key: archive[key] for key in archive.files} |
| validate_multimodal_keys(data) |
| checkpoint = torch.load(ROOT / (args.checkpoint or config["paths"]["checkpoint"]), map_location="cpu", weights_only=False) |
| if checkpoint.get("format_version") != config["format_version"]: |
| raise ValueError("Checkpoint format_version does not match configuration") |
| targets = list(checkpoint["model_config"]["targets"]) |
| tree_count = len(next(iter(checkpoint["model"].values()))["state"]["trees"]) |
| tree_predictions = np.full((data["X"].shape[0], len(targets), tree_count), np.nan, dtype=np.float32) |
| for month in checkpoint["model_config"]["months"]: |
| mask = data["month"] == month |
| for target_index, target in enumerate(targets): |
| model = BootstrapRandomForestRegressor.from_state_dict(checkpoint["model"][f"{month}:{target}"]["state"]) |
| tree_predictions[mask, target_index, :] = model.predict_trees(data["X"][mask]) |
| prediction = tree_predictions.mean(axis=2) |
| safe_prediction = np.where(np.abs(prediction) > 1e-8, prediction, np.nan) |
| ratio = data["Y"] / safe_prediction |
| if not np.isfinite(prediction).all() or not np.isfinite(ratio).all(): |
| raise FloatingPointError("Inference produced non-finite values") |
| output = ROOT / (args.output or config["paths"]["predictions"]) |
| output.parent.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed(output, pred=prediction, pred_trees=tree_predictions, obs=data["Y"], |
| obs_over_pred=ratio, relative_response=ratio - 1.0, |
| year=data["year"], month=data["month"], platform=data["platform"], |
| latitude=data["latitude"], longitude=data["longitude"], |
| target_names=np.asarray(targets)) |
| print(f"output={output.relative_to(ROOT)} samples={prediction.shape[0]} " |
| f"targets={targets} trees_per_prediction={tree_count}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
