OneScience-Group/ScaleAdaptiveCM

1"""Compact consistency model for full-grid precipitation downscaling."""2import json3from pathlib import Path4 5import numpy as np6import torch7from torch import nn8import torch.nn.functional as F9import yaml10 11 12def load_config(root):13    return yaml.safe_load((Path(root) / "conf/config.yaml").read_text())14 15 16class TimeBlock(nn.Module):17    def __init__(self, cin, cout, time_dim):18        super().__init__()19        groups = min(4, cout)20        self.conv = nn.Conv2d(cin, cout, 3, padding=1)21        self.norm = nn.GroupNorm(groups, cout)22        self.time = nn.Linear(time_dim, cout)23 24    def forward(self, x, embedding):25        return F.silu(self.norm(self.conv(x)) + self.time(embedding)[:, :, None, None])26 27 28class ScaleAdaptiveCM(nn.Module):29    def __init__(self, channels=(4, 8, 16), time_dim=16, sigma_data=0.5):30        super().__init__()31        self.sigma_data = float(sigma_data)32        self.time_dim = int(time_dim)33        self.time_mlp = nn.Sequential(nn.Linear(time_dim, time_dim), nn.SiLU(), nn.Linear(time_dim, time_dim))34        self.enc1 = TimeBlock(1, channels[0], time_dim)35        self.enc2 = TimeBlock(channels[0], channels[1], time_dim)36        self.mid = TimeBlock(channels[1], channels[2], time_dim)37        self.dec2 = TimeBlock(channels[2] + channels[1], channels[1], time_dim)38        self.dec1 = TimeBlock(channels[1] + channels[0], channels[0], time_dim)39        self.out = nn.Conv2d(channels[0], 1, 1)40        self.model_config = {"channels": list(channels), "time_dim": time_dim, "sigma_data": sigma_data}41 42    def embed_time(self, t):43        half = self.time_dim // 244        freq = torch.exp(torch.linspace(0, -7, half, device=t.device))45        emb = torch.cat((torch.sin(t[:, None] * freq), torch.cos(t[:, None] * freq)), 1)46        return self.time_mlp(emb)47 48    def forward(self, noisy, t):49        if noisy.ndim != 4 or noisy.shape[1] != 1:50            raise ValueError("expected [B,1,H,W]")51        emb = self.embed_time(t.float())52        e1 = self.enc1(noisy, emb)53        e2 = self.enc2(F.avg_pool2d(e1, 2), emb)54        mid = self.mid(F.avg_pool2d(e2, 2), emb)55        d2 = self.dec2(torch.cat((F.interpolate(mid, e2.shape[-2:], mode="bilinear", align_corners=False), e2), 1), emb)56        d1 = self.dec1(torch.cat((F.interpolate(d2, e1.shape[-2:], mode="bilinear", align_corners=False), e1), 1), emb)57        raw = self.out(d1)58        sigma2 = self.sigma_data ** 259        cskip = sigma2 / ((t[:, None, None, None] - 0.002).square() + sigma2)60        cout = self.sigma_data * t[:, None, None, None] / torch.sqrt(t[:, None, None, None].square() + sigma2)61        return cskip * noisy + cout * raw62 63 64def structured_fields(samples, high_h, high_w, seed):65    rng = np.random.default_rng(seed)66    yy, xx = np.mgrid[-1:1:complex(high_h), -1:1:complex(high_w)]67    fields = []68    for i in range(samples):69        phase = 2 * np.pi * i / max(samples, 4)70        itcz = 9 * np.exp(-((yy - .12 * np.sin(phase)) / .16) ** 2)71        storms = 18 * np.exp(-((xx - .45 * np.cos(phase)) ** 2 + (yy - .3 * np.sin(phase)) ** 2) / .025)72        texture = 2 * np.maximum(0, np.sin(18 * xx + phase) * np.cos(13 * yy - phase))73        fields.append(np.maximum(0, itcz + storms + texture + rng.normal(0, .15, yy.shape)))74    return np.asarray(fields, np.float32)[:, None]75 76 77def radial_spectrum(field):78    power = np.abs(np.fft.fftshift(np.fft.fft2(field))) ** 279    y, x = np.indices(field.shape); r = np.sqrt((y-field.shape[0]/2)**2 + (x-field.shape[1]/2)**2).astype(int)80    return np.bincount(r.ravel(), power.ravel()) / np.maximum(np.bincount(r.ravel()), 1)81 82 83def write_json(path, value):84    path = Path(path); path.parent.mkdir(parents=True, exist_ok=True)85    path.write_text(json.dumps(value, indent=2) + "\n")86