| """Run ensemble inference for all held-out synthetic daily maps.""" |
|
|
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from model.precipdd import ensemble_predict, load_config, load_ensemble, validate_archive |
|
|
|
|
| def main(): |
| config = load_config(ROOT / "conf/config.yaml") |
| device = torch.device("cuda" if torch.cuda.is_available() and config["runtime"]["device"] != "cpu" else "cpu") |
| data = np.load(ROOT / config["paths"]["data"]) |
| validate_archive(data) |
| models, checkpoint = load_ensemble(ROOT / config["paths"]["checkpoint"], device) |
| mask = data["split"] == 2 |
| fields = torch.from_numpy(data["precipitation"][mask]).float().to(device) |
| prediction = ensemble_predict(models, fields, config["training"]["batch_size"]).cpu().numpy() |
| output = ROOT / config["paths"]["predictions"] |
| output.parent.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed(output, format_version=data["format_version"], prediction=prediction, target=data["agmt"][mask], |
| precipitation=data["precipitation"][mask], year=data["year"][mask], day_of_year=data["day_of_year"][mask], |
| latitude=data["latitude"], longitude=data["longitude"], ensemble_members=np.array(len(models)), |
| checkpoint_world_size=np.array(checkpoint["world_size"])) |
| print(f"predictions={output.relative_to(ROOT)} days={len(prediction)} members={len(models)}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
