1from __future__ import annotations23import json4from pathlib import Path56import numpy as np7import torch89from ..models.model import MRPT, MRPTConfig101112def save_model(model: MRPT, path: Path, extra: dict | None = None) -> None:13 path.parent.mkdir(parents=True, exist_ok=True)14 torch.save({"cfg": model.cfg.__dict__, "model": model.state_dict(), "extra": extra or {}}, path)151617def load_model(path: Path, map_location="cpu") -> MRPT:18 ck = torch.load(path, map_location=map_location, weights_only=False)19 m = MRPT(MRPTConfig(**ck["cfg"]))20 m.load_state_dict(ck["model"])21 return m222324def save_predictions(res: dict[str, np.ndarray], path: Path, drop_samples: bool = True) -> None:25 path.parent.mkdir(parents=True, exist_ok=True)26 keep = {k: v for k, v in res.items() if not (drop_samples and ("samples" in k or k == "paths"))}27 np.savez_compressed(path, **keep)282930def write_json(obj, path: Path) -> None:31 path.parent.mkdir(parents=True, exist_ok=True)32 path.write_text(json.dumps(obj, indent=1, default=lambda o: float(o) if isinstance(o, (np.floating,)) else33 int(o) if isinstance(o, (np.integer,)) else o.tolist() if isinstance(o, np.ndarray) else str(o)))34