| """MaskDiT-style denoiser for native packed radar wedges.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
| import json |
| from pathlib import Path |
| import numpy as np |
| import yaml |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
| from torch.utils.data import Dataset |
|
|
|
|
| def load_config(path): |
| with open(path, encoding="utf-8") as handle: |
| return yaml.safe_load(handle) |
|
|
|
|
| def geometry_mask(config): |
| data = config["data"] |
| mask = np.zeros((data["frames"], 4, 366, 180), dtype=bool) |
| for elevation, (azimuths, ranges) in enumerate(zip(data["azimuth_counts"], data["range_bins"])): |
| mask[:, elevation, :azimuths, :ranges] = True |
| return mask |
|
|
|
|
| def make_sample(config, seed): |
| rng = np.random.default_rng(seed); data = config["data"] |
| volume = np.zeros(tuple(data["storage_shape"]), dtype=np.float32) |
| phase = rng.uniform(0, 2*np.pi) |
| for t in range(data["frames"]): |
| for e, (na, nr) in enumerate(zip(data["azimuth_counts"], data["range_bins"])): |
| az=np.arange(na,dtype=np.float32)[:,None]; radius=np.arange(nr,dtype=np.float32)[None] |
| center=(80+12*t+8*e+7*np.sin(phase+t/3))%na; center_r=min(nr-15,42+4*t-5*e) |
| da=np.minimum(abs(az-center),na-abs(az-center)); core=np.exp(-.5*((da/(14+2*e))**2+((radius-center_r)/18)**2)) |
| volume[t,e,:na,:nr]=np.clip(58*np.exp(-.22*e)*core+2*np.sin(radius/8+az/30+phase+t/2),0,70) |
| valid=geometry_mask(config); observed=valid.copy() |
| for t in range(data["history"]): |
| for e,(na,nr) in enumerate(zip(data["azimuth_counts"],data["range_bins"])): |
| width=max(1,int(round(na*data["missing_ratio"]))); start=int(rng.integers(0,na-width+1)); observed[t,e,start:start+width,:nr]=False |
| packed=volume.copy();packed[~observed]=0 |
| return {"values":packed,"truth":volume,"validity":valid,"observed":observed} |
|
|
|
|
| class PackedRadarDataset(Dataset): |
| def __init__(self,directory,split): |
| self.files=sorted(Path(directory).glob(f"{split}_*.npz")) |
| if not self.files: raise FileNotFoundError(f"no {split} samples; run fake_data.py") |
| def __len__(self): return len(self.files) |
| def __getitem__(self,index): |
| with np.load(self.files[index]) as sample: |
| return {key:torch.from_numpy(sample[key].copy()) for key in ("values","truth","validity","observed")} |
|
|
|
|
| class DiffusionSchedule: |
| def __init__(self,steps,beta_start,beta_end,device): |
| self.steps=steps;self.betas=torch.linspace(beta_start,beta_end,steps,device=device);self.alphas=1-self.betas;self.alpha_bars=torch.cumprod(self.alphas,0) |
| def add_noise(self,clean,timestep): |
| noise=torch.randn_like(clean);alpha=self.alpha_bars[timestep].view(-1,1,1,1,1);return alpha.sqrt()*clean+(1-alpha).sqrt()*noise,noise |
| @torch.no_grad() |
| def sample(self,model,history,observed): |
| target=torch.randn(history.shape[0],5,*history.shape[2:],device=history.device) |
| for step in reversed(range(self.steps)): |
| t=torch.full((history.shape[0],),step,device=history.device,dtype=torch.long);clean=model(history,target,observed,t,None) |
| target=clean if step==0 else self.alpha_bars[step-1].sqrt()*clean+(1-self.alpha_bars[step-1]).sqrt()*torch.randn_like(clean) |
| return target |
|
|
|
|
| def evaluate_ensemble(ensemble,truth,valid,thresholds): |
| mean=ensemble.mean(0);error=mean-truth;result={"mae":float(abs(error)[valid].mean()),"rmse":float(np.sqrt((error**2)[valid].mean()))} |
| result["crps"]=float(abs(ensemble-truth[None])[:,valid].mean()-.5*abs(ensemble[:,None]-ensemble[None,:])[:,:,valid].mean()) |
| result["categorical"]={} |
| for threshold in thresholds: |
| forecast,observed=mean>=threshold,truth>=threshold;hits=int((forecast&observed&valid).sum());miss=int((~forecast&observed&valid).sum());false=int((forecast&~observed&valid).sum()) |
| result["categorical"][str(threshold)]={"csi":hits/max(hits+miss+false,1),"far":false/max(hits+false,1),"pod":hits/max(hits+miss,1)} |
| result["coverage_ratio"]=[[float(np.isfinite(ensemble).all(0)[lead,e][valid[lead,e]].mean()) for e in range(4)] for lead in range(5)] |
| return result |
|
|
|
|
| class DiTBlock(nn.Module): |
| """Linear-cost latent attention block suitable for 24k native wedge tokens.""" |
|
|
| def __init__(self, hidden: int, heads: int, latent_tokens: int): |
| super().__init__() |
| self.latents = nn.Parameter(torch.randn(1, latent_tokens, hidden) / math.sqrt(hidden)) |
| self.read = nn.MultiheadAttention(hidden, heads, batch_first=True) |
| self.write = nn.MultiheadAttention(hidden, heads, batch_first=True) |
| self.norm1 = nn.LayerNorm(hidden) |
| self.norm2 = nn.LayerNorm(hidden) |
| self.mlp = nn.Sequential(nn.Linear(hidden, hidden * 4), nn.GELU(), nn.Linear(hidden * 4, hidden)) |
|
|
| def forward(self, tokens: torch.Tensor, conditioning: torch.Tensor) -> torch.Tensor: |
| latent = self.latents.expand(tokens.shape[0], -1, -1) + conditioning[:, None] |
| latent = latent + self.read(latent, self.norm1(tokens), self.norm1(tokens), need_weights=False)[0] |
| tokens = tokens + self.write(self.norm1(tokens), latent, latent, need_weights=False)[0] |
| return tokens + self.mlp(self.norm2(tokens + conditioning[:, None])) |
|
|
|
|
| class EchoCast3D(nn.Module): |
| def __init__(self, config: dict): |
| super().__init__() |
| self.config = config |
| model = config["model"] |
| self.patch = tuple(model["patch_size"]) |
| patch_area = self.patch[0] * self.patch[1] |
| self.embed = nn.Linear(config["data"]["frames"] * patch_area, model["hidden_dim"]) |
| self.time_embed = nn.Embedding(config["diffusion"]["steps"], model["hidden_dim"]) |
| self.elevation_embed = nn.Embedding(4, model["hidden_dim"]) |
| self.mask_token = nn.Parameter(torch.zeros(1, 1, model["hidden_dim"])) |
| self.encoder = nn.ModuleList( |
| [DiTBlock(model["hidden_dim"], model["heads"], model["latent_tokens"]) for _ in range(model["depth"])] |
| ) |
| self.decoder = nn.ModuleList( |
| [DiTBlock(model["hidden_dim"], model["heads"], model["latent_tokens"]) for _ in range(model["depth"])] |
| ) |
| self.head = nn.Linear(model["hidden_dim"], config["data"]["target"] * patch_area) |
| self.layouts = [] |
| offset = 0 |
| for e, (na, nr) in enumerate(zip(config["data"]["azimuth_counts"], config["data"]["range_bins"])): |
| rows, cols = na // self.patch[0], nr // self.patch[1] |
| count = rows * cols |
| self.layouts.append((e, na, nr, rows, cols, offset, offset + count)) |
| offset += count |
| self.token_count = offset |
|
|
| def patchify(self, volume: torch.Tensor) -> torch.Tensor: |
| parts = [] |
| for e, na, nr, _, _, _, _ in self.layouts: |
| wedge = volume[:, :, e, :na, :nr] |
| patches = F.unfold(wedge, kernel_size=self.patch, stride=self.patch).transpose(1, 2) |
| parts.append(patches) |
| return torch.cat(parts, dim=1) |
|
|
| def token_validity(self, validity: torch.Tensor) -> torch.Tensor: |
| parts = [] |
| for e, na, nr, _, _, _, _ in self.layouts: |
| wedge = validity[:, :, e, :na, :nr].float() |
| patch = F.unfold(wedge, kernel_size=self.patch, stride=self.patch) |
| parts.append((patch > 0.5).all(dim=1)) |
| return torch.cat(parts, dim=1) |
|
|
| def unpatchify(self, patches: torch.Tensor) -> torch.Tensor: |
| batch = patches.shape[0] |
| output = patches.new_zeros((batch, 5, 4, 366, 180)) |
| for e, na, nr, rows, cols, start, end in self.layouts: |
| wedge = F.fold( |
| patches[:, start:end].transpose(1, 2), |
| output_size=(na, nr), |
| kernel_size=self.patch, |
| stride=self.patch, |
| ) |
| output[:, :, e, :na, :nr] = wedge |
| return output |
|
|
| def forward( |
| self, |
| history: torch.Tensor, |
| noisy_target: torch.Tensor, |
| observed: torch.Tensor, |
| timestep: torch.Tensor, |
| token_mask: torch.Tensor | None, |
| ) -> torch.Tensor: |
| full = torch.cat((history, noisy_target), dim=1) |
| tokens = self.embed(self.patchify(full)) |
| elevation_ids = torch.cat( |
| [torch.full((end - start,), e, device=full.device) for e, _, _, _, _, start, end in self.layouts] |
| ) |
| tokens = tokens + self.elevation_embed(elevation_ids)[None] |
| if token_mask is not None: |
| tokens = torch.where(token_mask[..., None], self.mask_token.expand_as(tokens), tokens) |
| condition = self.time_embed(timestep) |
| for block in self.encoder: |
| tokens = block(tokens, condition) |
| for block in self.decoder: |
| tokens = block(tokens, condition) |
| return self.unpatchify(self.head(tokens)) |
|
|
|
|
| def masked_objective( |
| model: EchoCast3D, |
| prediction: torch.Tensor, |
| truth: torch.Tensor, |
| validity: torch.Tensor, |
| token_mask: torch.Tensor, |
| reconstruction_weight: float, |
| ) -> tuple[torch.Tensor, dict[str, float]]: |
| target_valid = validity[:, 3:] |
| token_pixels = token_mask[:, :, None].expand(-1, -1, 5 * 9) |
| masked_pixels = model.unpatchify(token_pixels.float()).bool() |
| squared = (prediction - truth) ** 2 |
| unmasked = target_valid & ~masked_pixels |
| masked = target_valid & masked_pixels |
| score = squared[unmasked].mean() |
| reconstruction = squared[masked].mean() |
| loss = score + reconstruction_weight * reconstruction |
| return loss, {"score": score.item(), "reconstruction": reconstruction.item()} |
|
|
