File size: 2,790 Bytes
41dd7e0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | from __future__ import annotations
import torch
from ..base import Method
from . import core
class CascadedSIMethod(Method):
def __init__(self, sigma_train, t_eps, sampling, beta_power=1,
eval_noise_scales=None, optimizer_cfg=None):
self.sigma_train = float(sigma_train)
self.t_eps = float(t_eps)
self.beta_power = int(beta_power)
self.n_sde_steps_per_seg = int(sampling["n_sde_steps_per_seg"])
self.n_segments = int(sampling.get("n_segments", 10))
self.noise_scale = float(sampling.get("noise_scale", 1.0))
# Segment boundaries — uniform partition of [0, 1] by default, matching
# the normalized physical forecast horizon.
if "t_values" in sampling:
self.t_values = list(sampling["t_values"])
else:
self.t_values = [k / self.n_segments for k in range(self.n_segments + 1)]
self.eval_noise_scales = ([float(s) for s in eval_noise_scales]
if eval_noise_scales is not None else [self.noise_scale])
self.optimizer_cfg = optimizer_cfg or {}
@classmethod
def from_config(cls, cfg):
if cfg.get("anchor_ic", False):
raise ValueError("This release supports the paper's SI models without retained initial-state inputs")
return cls(
sigma_train=cfg["sigma_train"],
t_eps=cfg.get("t_eps", 0.0),
sampling=cfg["sampling"],
beta_power=cfg.get("beta_power", 1),
eval_noise_scales=cfg.get("eval_noise_scales"),
optimizer_cfg=cfg.get("optimizer"),
)
def variants_for_eval(self):
return {f"noise_scale_{_fmt(s)}": {"noise_scale": s} for s in self.eval_noise_scales}
def loss(self, model, batch):
x_tk, x_tk1, t_k, t_k1 = batch
return core.local_si_loss(model, x_tk, x_tk1, t_k, t_k1,
sigma=self.sigma_train, t_eps=self.t_eps,
beta_power=self.beta_power)
def sample(self, model, x_0, n_samples, noise_scale=None):
if noise_scale is None:
noise_scale = self.noise_scale
if x_0.ndim == 1:
x_0 = x_0.unsqueeze(0)
B = x_0.shape[0]
out = []
for i in range(B):
cloud = core.rollout(model, x_0[i], n_samples=n_samples,
t_values=self.t_values,
n_sde_steps_per_seg=self.n_sde_steps_per_seg,
sigma=self.sigma_train,
noise_scale=noise_scale)
out.append(cloud)
return torch.stack(out, dim=0)
def _fmt(s):
if s == int(s):
return f"{int(s)}p0"
return f"{s}".replace(".", "p")
|