Upload 4 files

unified_graph_weather_dataloader.py

@@ -0,0 +1,415 @@
+# graph_weather_dataloader.py
+from __future__ import annotations
+from pathlib import Path
+from typing import Optional, Tuple, List, Dict
+
+import re
+import json
+import math
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader
+import pytorch_lightning as pl
+
+
+# ==========================
+# Utils
+# ==========================
+def _parse_ym_from_name(p: Path):
+    m = re.search(r"(\d{4})-(\d{2})\.pt$", p.name)
+    if not m:
+        return None, None
+    return int(m.group(1)), int(m.group(2))
+
+
+def _sincos_pos2d_nodes(lat_deg: np.ndarray, lon_deg: np.ndarray) -> np.ndarray:
+    """
+    По вузлах графа -> [4, N]: sinφ, cosφ, sin(λ·cosφ̄), cos(λ·cosφ̄).
+    φ̄ — середня широта домену (стабільно, як у твоєму build-скрипті).
+    """
+    lat = np.deg2rad(lat_deg.astype(np.float32, copy=False))
+    lon = np.deg2rad(lon_deg.astype(np.float32, copy=False))
+    lat0 = float(lat.mean())
+    lon_adj = lon * np.cos(lat0)
+    pos = np.stack([np.sin(lat), np.cos(lat), np.sin(lon_adj), np.cos(lon_adj)], axis=0)
+    return pos.astype(np.float32, copy=False)
+
+
+def _load_graph_static(npz_path: Path):
+    """
+    Зчитує статичний мультискейл граф:
+      обов'язково: coord_lon[N], coord_lat[N], static_geo_mesh[N,4]
+      опційно: edge_index_horiz_base[2,Eh] або edge_index_horiz[2, E_total] + horiz_edge_ptr[L+1]
+      (fallback на edges_h_src/edges_h_dst або edges_horiz)
+    """
+    d = np.load(npz_path, allow_pickle=True)
+
+    # --- coords
+    lon = d["coord_lon"].astype(np.float32, copy=False)
+    lat = d["coord_lat"].astype(np.float32, copy=False)
+    N = int(lon.shape[0])
+
+    # --- static geo (4 канали, вже нормалізовані по зонах)
+    if "static_geo_mesh" not in d:
+        raise KeyError(f"{npz_path} must contain 'static_geo_mesh'")
+    static_geo = d["static_geo_mesh"].astype(np.float32, copy=False)  # [N,4]
+    if static_geo.shape[0] != N:
+        raise ValueError("static_geo_mesh first dim must match coord arrays")
+
+    # --- pos2d: беремо з файла, якщо є; інакше — будуємо
+    if "static_pos2d_mesh" in d:
+        pos2d = d["static_pos2d_mesh"].astype(np.float32, copy=False).T  # [4,N]
+    else:
+        pos2d = _sincos_pos2d_nodes(lat, lon)  # [4,N]
+
+    # --- levels/meta (не обов'язково, але корисно)
+    levels = [str(x) for x in (d["levels"].tolist() if "levels" in d else [])]
+    level_offsets = d["level_offsets"].astype(np.int32, copy=False) if "level_offsets" in d else None
+    mesh_hash = None
+    if "mesh_hash" in d:
+        mh = d["mesh_hash"].item()
+        mesh_hash = (mh.decode() if isinstance(mh, (bytes, np.bytes_)) else str(mh))
+
+    # --- edge_index_base: кілька шляхів
+    edge_index = None
+
+    # 1) ідеально — вже готове поле
+    if "edge_index_horiz_base" in d:
+        ei = d["edge_index_horiz_base"].astype(np.int64, copy=False)
+        edge_index = torch.from_numpy(ei)
+
+    # 2) зібране по рівнях: edge_index_horiz + horiz_edge_ptr → беремо рівень 0
+    elif "edge_index_horiz" in d and "horiz_edge_ptr" in d:
+        e_all = d["edge_index_horiz"]   # [2, E_total] або [E_total, 2]
+        if e_all.ndim == 2 and e_all.shape[0] == 2:
+            e_all = e_all
+        elif e_all.ndim == 2 and e_all.shape[1] == 2:
+            e_all = e_all.T
+        else:
+            raise ValueError("edge_index_horiz must be [2, E] or [E, 2]")
+
+        ptr = d["horiz_edge_ptr"].astype(np.int64, copy=False)  # [L+1]
+        if ptr.size < 2:
+            raise ValueError("horiz_edge_ptr must have at least 2 entries")
+        e0 = e_all[:, int(ptr[0]):int(ptr[1])].astype(np.int64, copy=False)  # slice рівня 0
+
+        # якщо чомусь індекси зсунуто — прибираємо офсет рівня 0
+        if level_offsets is not None and int(level_offsets[0]) != 0:
+            e0 = e0 - int(level_offsets[0])
+
+        if e0.min() < 0 or e0.max() >= N:
+            raise RuntimeError(f"edge_index_base out of range after slice: min={e0.min()}, max={e0.max()}, N={N}")
+
+        edge_index = torch.from_numpy(e0.copy())
+
+    # 3) старі ключі: edges_h_src/edges_h_dst або edges_horiz
+    elif "edges_h_src" in d and "edges_h_dst" in d:
+        src = torch.from_numpy(d["edges_h_src"].astype(np.int64, copy=False))
+        dst = torch.from_numpy(d["edges_h_dst"].astype(np.int64, copy=False))
+        edge_index = torch.stack([src, dst], dim=0)
+    elif "edges_horiz" in d:
+        eh = d["edges_horiz"]
+        if eh.ndim == 2 and eh.shape[1] == 2:
+            edge_index = torch.from_numpy(eh.astype(np.int64, copy=False)).t().contiguous()
+        elif eh.ndim == 2 and eh.shape[0] == 2:
+            edge_index = torch.from_numpy(eh.astype(np.int64, copy=False))
+
+    # --- упакуємо відповіді
+    meta = {
+        "N_mesh": N,
+        "levels": levels,
+        "level_offsets": level_offsets,
+        "mesh_hash": mesh_hash,
+    }
+    geo = static_geo.T  # [4,N]
+    return {
+        "lon": lon, "lat": lat, "geo": geo, "pos2d": pos2d,      # [4,N] і [4,N]
+        "edge_index_base": edge_index,                            # torch.Long[2,Eh] або None
+        "meta": meta,
+    }
+
+
+# ==========================
+# Dataset
+# ==========================
+class GraphShardDataset(Dataset):
+    """
+    Зчитує графові шарди і віддає семпли для GAT-моделі.
+
+    __getitem__ -> ((x_all, geo, pos2d), y_all)
+      x_all: [Tin, C, N]
+      y_all: [Tout, C, N]
+      geo:   [G,   N]  (4)
+      pos2d: [P,   N]  (4)
+    """
+    def __init__(self,
+                 shards_root: str,
+                 years: range,
+                 static_graph_npz: str = "./data/graph/graph_static_from_stats_vgeo_aligned.npz",
+                 allowed_months: Optional[List[int]] = None,  # None → всі 1..12
+                 force_tin_tout: Optional[Tuple[int, int]] = None,  # None → брати з шару
+                 dtype_out: str = "float32",
+                 shuffle_within_file: bool = True):
+        super().__init__()
+        self.root = Path(shards_root)
+        self.years = set(int(y) for y in years)
+        self.allowed_months = None if allowed_months is None else {int(m) for m in allowed_months}
+        self.force_tin_tout = force_tin_tout
+        assert dtype_out in {"float16", "float32"}
+        self.dtype_out = torch.float16 if dtype_out == "float16" else torch.float32
+        self.shuffle_within_file = bool(shuffle_within_file)
+
+        # ---- зчитуємо статичні фічі графа один раз ----
+        sg = _load_graph_static(Path(static_graph_npz))
+        self.geo   = torch.from_numpy(sg["geo"].copy())     # [4,N]
+        self.pos2d = torch.from_numpy(sg["pos2d"].copy())   # [4,N]
+        self.N = int(sg["geo"].shape[1])
+        self.edge_index_base = sg["edge_index_base"]  # може бути None
+        self.static_meta = sg["meta"]
+
+        # ---- зібрати список файлів і entries ----
+        files = []
+        for p in sorted(self.root.glob("graph_*.pt")):
+            y, m = _parse_ym_from_name(p)
+            if y is None:
+                continue
+            if (y in self.years) and (self.allowed_months is None or m in self.allowed_months):
+                files.append(p)
+        if not files:
+            raise FileNotFoundError(f"No graph_YYYY-MM.pt shards in {self.root}")
+
+        self.entries: List[Tuple[str, int]] = []  # (path, s_start)
+        self._file_meta: Dict[str, Dict] = {}     # path -> {T,C,N,Tin,Tout,stride,S}
+        self._file_counts: Dict[str, int] = {}
+
+        for p in files:
+            d = torch.load(p, map_location="cpu")
+            values = d["values"]         # [T,C,N]
+            T, C, N = values.shape
+            if N != self.N:
+                raise ValueError(f"{p.name}: N({N}) != static N({self.N})")
+
+            Tin  = int(d["Tin"])
+            Tout = int(d["Tout"])
+            stride = int(d["stride"])
+            starts = d["sample_starts"].cpu().numpy().astype(np.int64)  # [S]
+            if self.force_tin_tout is not None:
+                # Перерахувати sample_starts під інші Tin/Tout (залишимо stride=1)
+                fTin, fTout = self.force_tin_tout
+                max_start = T - (fTin + fTout)
+                if max_start < 0:
+                    continue
+                starts = np.arange(0, max_start + 1, 1, dtype=np.int64)
+                Tin, Tout = int(fTin), int(fTout)
+
+            if self.shuffle_within_file:
+                rng = np.random.default_rng(int(y) * 100 + int(m))
+                rng.shuffle(starts)
+
+            for s in starts:
+                self.entries.append((str(p), int(s)))
+
+            self._file_meta[str(p)] = {
+                "T": T, "C": C, "N": N, "Tin": Tin, "Tout": Tout, "stride": stride, "S": int(len(starts))
+            }
+            self._file_counts[str(p)] = int(len(starts))
+
+        # cache last shard
+        self._last_path: Optional[str] = None
+        self._last_data = None
+
+        # зручні поля
+        self.C = next(iter(self._file_meta.values()))["C"]
+        self.Tin = next(iter(self._file_meta.values()))["Tin"]
+        self.Tout = next(iter(self._file_meta.values()))["Tout"]
+
+    def __len__(self) -> int:
+        return len(self.entries)
+
+    def _load_shard(self, path: str):
+        if self._last_path != path:
+            self._last_data = torch.load(path, map_location="cpu")
+            self._last_path = path
+        return self._last_data
+
+    def __getitem__(self, i: int):
+        path, s = self.entries[i]
+        meta = self._file_meta[path]
+        Tin, Tout = meta["Tin"], meta["Tout"]
+
+        d = self._load_shard(path)
+        vals = d["values"]  # torch Tensor [T,C,N], може бути float16
+        # збираємо x/y як float32 (рекомендовано для стабільності)
+        x_all = vals[s:s+Tin].to(dtype=self.dtype_out)        # [Tin,C,N]
+        y_all = vals[s+Tin:s+Tin+Tout].to(dtype=self.dtype_out)  # [Tout,C,N]
+
+        # статичні — копії не потрібні, вони read-only, але distinct тензори бажані
+        geo   = self.geo.to(dtype=self.dtype_out)     # [4,N]
+        pos2d = self.pos2d.to(dtype=self.dtype_out)   # [4,N]
+
+        return (x_all, geo, pos2d), y_all
+
+    # ---------- diagnostics ----------
+    @property
+    def file_counts(self) -> Dict[str, int]:
+        return dict(self._file_counts)
+
+    @property
+    def N_mesh(self) -> int:
+        return self.N
+
+
+# ==========================
+# Prefetch (GPU)
+# ==========================
+class PrefetchLoader:
+    """GPU-prefetch з відомим __len__."""
+    def __init__(self, loader: DataLoader, device: str = 'cuda'):
+        self.loader = loader
+        self.device = device
+
+    def __len__(self):
+        return len(self.loader)
+
+    def __iter__(self):
+        if self.device.startswith("cuda"):
+            stream = torch.cuda.Stream()
+            first = True
+            for (x_all, geo, pos2d), y_all in self.loader:
+                with torch.cuda.stream(stream):
+                    x_all = x_all.to(self.device, non_blocking=True)
+                    geo   = geo.to(self.device, non_blocking=True)
+                    pos2d = pos2d.to(self.device, non_blocking=True)
+                    y_all = y_all.to(self.device, non_blocking=True)
+                    next_in, next_tgt = (x_all, geo, pos2d), y_all
+                if not first:
+                    yield cur_in, cur_tgt
+                else:
+                    first = False
+                torch.cuda.current_stream().wait_stream(stream)
+                cur_in, cur_tgt = next_in, next_tgt
+            yield cur_in, cur_tgt
+        else:
+            # CPU → прозора прокладка
+            for batch in self.loader:
+                yield batch
+
+
+# ==========================
+# Lightning DataModule
+# ==========================
+class GraphWeatherDataModule(pl.LightningDataModule):
+    """
+    Створює train/val лоадери для графових шард-файлів.
+    Виносить edge_index_base і N_mesh як властивості, щоб інітити модель.
+
+    Batch: ((x_all, geo, pos2d), y_all)
+      x_all: [B, Tin, C, N]
+      y_all: [B, Tout, C, N]
+      geo:   [B, 4,   N]
+      pos2d: [B, 4,   N]
+    """
+    def __init__(self,
+                 shards_root_train: str = "./data/shards_graph",
+                 shards_root_val: str   = "./data/shards_graph",
+                 years_train: range = range(1980, 2011),
+                 years_val:   range = range(2011, 2013),
+                 static_graph_npz: str = "./data/graph/graph_static_from_stats_vgeo_aligned.npz",
+                 allowed_months: Optional[List[int]] = None,
+                 force_tin_tout: Optional[Tuple[int,int]] = None,
+                 dtype_out: str = "float32",
+                 shuffle_within_file_train: bool = True,
+                 shuffle_within_file_val:   bool = False,
+                 batch_size_train: int = 1,
+                 batch_size_val:   int = 1,
+                 num_workers: Optional[int] = None,
+                 pin_memory: bool = True,
+                 persistent_workers: bool = True,
+                 device: str = "cuda"):
+        super().__init__()
+        self.paths = {"train": shards_root_train, "val": shards_root_val}
+        self.years = {"train": years_train, "val": years_val}
+        self.kw_ds = dict(
+            static_graph_npz=static_graph_npz,
+            allowed_months=allowed_months,
+            force_tin_tout=force_tin_tout,
+            dtype_out=dtype_out,
+        )
+        self.shuffle = {"train": shuffle_within_file_train, "val": shuffle_within_file_val}
+        self.bs = {"train": batch_size_train, "val": batch_size_val}
+        self.num_workers = num_workers if num_workers is not None else max((torch.get_num_threads() or 1) - 1, 0)
+        self.pin_memory = pin_memory
+        self.persistent_workers = persistent_workers
+        self.device = device
+
+        self._train_ds = None
+        self._val_ds   = None
+        self._train_loader = None
+        self._val_loader   = None
+
+        # public: для ініціалізації моделі
+        self.edge_index_base: Optional[torch.Tensor] = None  # Long[2,Eh] або None
+        self.N_mesh: Optional[int] = None
+        self.C: Optional[int] = None
+        self.Tin: Optional[int] = None
+        self.Tout: Optional[int] = None
+
+    def setup(self, stage=None):
+        self._train_ds = GraphShardDataset(
+            shards_root=self.paths["train"],
+            years=self.years["train"],
+            shuffle_within_file=self.shuffle["train"],
+            **self.kw_ds
+        )
+        self._val_ds = GraphShardDataset(
+            shards_root=self.paths["val"],
+            years=self.years["val"],
+            shuffle_within_file=self.shuffle["val"],
+            **self.kw_ds
+        )
+
+        # заповнимо публічні поля із train DS (ідентичні у val при тих самих даних)
+        self.N_mesh = self._train_ds.N_mesh
+        self.edge_index_base = self._train_ds.edge_index_base   # може бути None
+        self.C = self._train_ds.C
+        self.Tin = self._train_ds.Tin
+        self.Tout = self._train_ds.Tout
+
+    def _make_loader(self, ds: Dataset, bs: int, shuffle: bool):
+        base = DataLoader(
+            ds, batch_size=bs, shuffle=shuffle,
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+            persistent_workers=self.persistent_workers,
+            drop_last=False,
+        )
+        return PrefetchLoader(base, device=self.device)
+
+    def train_dataloader(self):
+        self._train_loader = self._make_loader(self._train_ds, self.bs["train"], shuffle=True)
+        return self._train_loader
+
+    def val_dataloader(self):
+        self._val_loader = self._make_loader(self._val_ds, self.bs["val"], shuffle=False)
+        return self._val_loader
+
+    # handy stats
+    @property
+    def train_samples(self):
+        return len(self._train_ds) if self._train_ds is not None else None
+
+    @property
+    def val_samples(self):
+        return len(self._val_ds) if self._val_ds is not None else None
+
+    @property
+    def train_steps_per_epoch(self):
+        if self._train_ds is None:
+            return None
+        return math.ceil(len(self._train_ds) / self.bs["train"])
+
+    @property
+    def val_steps(self):
+        if self._val_ds is None:
+            return None
+        return math.ceil(len(self._val_ds) / self.bs["val"])

unified_graph_weather_model.py

@@ -0,0 +1,366 @@
+# unified_graph_weather_model_gat.py
+# Графова версія: Vertical (C) attention + Horizontal GAT over edges
+from __future__ import annotations
+from typing import Tuple, Optional, Literal
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pytorch_lightning as pl
+
+
+# -----------------------------
+# SDPA helpers (без Flash)
+# -----------------------------
+def _sdpa_no_flash_ctx():
+    try:
+        from torch.nn.attention import sdpa_kernel, SDPBackend
+        return sdpa_kernel(backends=[SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION])
+    except Exception:
+        return torch.backends.cuda.sdp_kernel(
+            enable_flash=False, enable_math=True, enable_mem_efficient=True
+        )
+
+
+class EfficientScaledDotProductAttention(nn.Module):
+    """Self-attention по послідовності (torch SDPA)."""
+    def __init__(self, embed_dim: int, num_heads: int):
+        super().__init__()
+        assert embed_dim % num_heads == 0
+        self.num_heads = num_heads
+        self.embed_dim = embed_dim
+        self.head_dim  = embed_dim // num_heads
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+        self.norm = nn.LayerNorm(embed_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B, N, E]
+        B, N, _ = x.shape
+        q = self.q_proj(x).view(B, N, self.num_heads, self.head_dim).transpose(1, 2)  # [B,h,N,d]
+        k = self.k_proj(x).view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
+        with _sdpa_no_flash_ctx():
+            attn = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False)
+        attn = attn.transpose(1, 2).contiguous().view(B, N, self.embed_dim)                 # [B,N,E]
+        out = self.out_proj(attn)
+        return self.norm(out + x)
+
+
+# -----------------------------
+# GAT mixer (multi-head, без сторонніх бібліотек)
+# -----------------------------
+class GraphGATMixer(nn.Module):
+    """
+    Multi-head GAT по горизонтальних ребрах.
+    Працює окремо для кожного вертикального токена (каналу) і кожного елемента батчу.
+
+    Вхід:  H [B, C, N, E]
+    Вихід: H' того ж розміру
+
+    edge_index: Long[2, Eh] (індекси 0..N-1)
+    Порада: бажано додати самопетлі в edge_index поза моделлю (i->i), або використовується skip-резідюал.
+    """
+    def __init__(self, embed_dim: int, num_heads: int = 4, attn_drop: float = 0.0, neg_slope: float = 0.2):
+        super().__init__()
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+        self.E = embed_dim
+        self.H = num_heads
+        self.Dh = embed_dim // num_heads
+        self.attn_drop = attn_drop
+        self.leaky = nn.LeakyReLU(neg_slope)
+
+        # Лінійна проєкція у head-простір
+        self.lin = nn.Linear(self.E, self.E, bias=False)
+
+        # a_src / a_dst параметри на кожну голову: [H, Dh]
+        self.a_src = nn.Parameter(torch.empty(self.H, self.Dh))
+        self.a_dst = nn.Parameter(torch.empty(self.H, self.Dh))
+
+        # вихідна проєкція + нормалізація
+        self.out_proj = nn.Linear(self.E, self.E, bias=True)
+        self.norm = nn.LayerNorm(self.E)
+        self.act = nn.GELU()
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.xavier_uniform_(self.lin.weight)
+        nn.init.xavier_uniform_(self.out_proj.weight)
+        nn.init.zeros_(self.out_proj.bias)
+        nn.init.xavier_uniform_(self.a_src.unsqueeze(-1))
+        nn.init.xavier_uniform_(self.a_dst.unsqueeze(-1))
+
+    def forward(self, H: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor:
+        """
+        H: [B, C, N, E]
+        edge_index: [2, Eh]
+        """
+        B, C, N, E = H.shape
+        src = edge_index[0].to(H.device, non_blocking=True)
+        dst = edge_index[1].to(H.device, non_blocking=True)
+        Eh = src.numel()
+
+        # Перетворення у [B*C, N, H, Dh]
+        X = H.view(B * C, N, E)
+        X = self.lin(X)                                 # [BC, N, E]
+        X = X.view(B * C, N, self.H, self.Dh)           # [BC, N, H, Dh]
+
+        # За один прохід пройдемо BC-семпли
+        out = torch.zeros(B * C, N, self.H, self.Dh, device=H.device, dtype=H.dtype)
+
+        # Попередньо підготуємо a_src/dst у потрібному dtype/пристрої
+        a_src = self.a_src.to(device=H.device, dtype=H.dtype)
+        a_dst = self.a_dst.to(device=H.device, dtype=H.dtype)
+
+        for bc in range(B * C):
+            Xbc = X[bc]                                 # [N, H, Dh]
+            # Витягуємо ознаки джерел і приймачів по ребрах
+            x_i = Xbc.index_select(0, src)              # [Eh, H, Dh] (source features)
+            x_j = Xbc.index_select(0, dst)              # [Eh, H, Dh] (target features)
+
+            # e_ij = Leaky(a_src·x_i + a_dst·x_j), по головах
+            e_src = torch.einsum("ehd,hd->eh", x_i, a_src)   # [Eh, H]
+            e_dst = torch.einsum("ehd,hd->eh", x_j, a_dst)   # [Eh, H]
+            e = self.leaky(e_src + e_dst)                    # [Eh, H]
+
+            # стабілізуємо softmax: обмежимо значення
+            e = torch.clamp(e, -10.0, 10.0)
+            exp_e = torch.exp(e)                             # [Eh, H]
+            if self.attn_drop > 0:
+                exp_e = F.dropout(exp_e, p=self.attn_drop, training=self.training)
+
+            # ∑_{k in N(dst)} exp_e; нормалізація по вхідних ребрах кожного dst (per head)
+            # Використовуємо index_add по другій осі (уздовж вузлів)
+            sum_dst = torch.zeros(self.H, N, device=H.device, dtype=H.dtype)  # [H, N]
+            # transpose -> [H, Eh], add along dim=1 at indices 'dst'
+            sum_dst.index_add_(1, dst, exp_e.transpose(0, 1))
+            # вибрати суму для кожного ребра/голови
+            denom = sum_dst[:, dst].transpose(0, 1) + 1e-12                     # [Eh, H]
+            alpha = exp_e / denom                                               # [Eh, H]
+
+            # повідомлення від джерела до приймача: alpha * x_i (можна x_j; беремо x_i як у GAT v1)
+            m = alpha.unsqueeze(-1) * x_i                                       # [Eh, H, Dh]
+
+            # агрегуємо у вузли-приймачі (per head)
+            # робимо head-цикл — H невелике, зате пам’ять економна
+            for h in range(self.H):
+                out[bc, :, h, :].index_add_(0, dst, m[:, h, :])                 # [N, Dh]
+
+        # concat голови → [BC, N, E]
+        out = out.reshape(B * C, N, self.H * self.Dh)
+        out = self.out_proj(out)                                                # [BC, N, E]
+        out = self.act(out)
+        # назад у [B, E, C, N]
+        out = out.view(B, C, N, E).permute(0, 3, 1, 2).contiguous()             # [B, E, C, N]
+        # skip connection + norm опрацьовуємо у блоці
+        return out
+
+
+# -----------------------------
+# Один блок: Vertical SDPA + GAT + FFN
+# -----------------------------
+class GraphAxialGATBlock(nn.Module):
+    """
+    Працює у представленні [B, E, C, N].
+    1) Vertical attention (по C у кожній вершині)
+    2) GAT по графу (по горизонталі) окремо для кожного каналу
+    3) FFN (1x1 Conv2d) + residual + norm
+    """
+    def __init__(self, E: int, heads_v: int = 4, heads_xy: int = 4, attn_drop_xy: float = 0.0):
+        super().__init__()
+        self.attn_v = EfficientScaledDotProductAttention(E, heads_v)
+        self.gat    = GraphGATMixer(E, num_heads=heads_xy, attn_drop=attn_drop_xy)
+        self.ffn = nn.Sequential(
+            nn.Conv2d(E, 2 * E, kernel_size=1), nn.GELU(),
+            nn.Conv2d(2 * E, E, kernel_size=1),
+        )
+        self.norm = nn.GroupNorm(8, E)
+
+    def forward(self, F: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor:
+        # (1) Vertical attention
+        B, E, C, N = F.shape
+        V = F.permute(0, 3, 2, 1).reshape(B * N, C, E)    # [B*N, C, E]
+        V = self.attn_v(V)                                # [B*N, C, E]
+        V = V.reshape(B, N, C, E).permute(0, 3, 2, 1)     # [B, E, C, N]
+
+        # (2) GAT mixing (per channel)
+        G_in  = V.permute(0, 2, 3, 1).contiguous()        # [B, C, N, E]
+        G_out = self.gat(G_in, edge_index=edge_index)     # [B, E, C, N]
+
+        S = V + G_out                                     # residual після GAT
+
+        # (3) FFN + residual + norm
+        U = self.ffn(S)
+        return self.norm(U + S)
+
+
+# -----------------------------
+# Core network
+# -----------------------------
+class UnifiedGraphSpatioVerticalGATNet(nn.Module):
+    """
+    input:  (x_all, geo, pos2d), diffusion_step(optional)
+            x_all: [B, Tin, C, N]
+            geo:   [B, G,   N]
+            pos2d: [B, P,   N]
+    output: [B, Tout, C, N]
+
+    Потрібно подати edge_index_base: Long[2, Eh] (горизонтальні ребра базового рівня).
+    """
+    def __init__(self,
+                 C: int,
+                 Tin: int,
+                 Tout: int,
+                 edge_index_base: torch.Tensor,
+                 N_mesh: int,
+                 static_geo_ch: int = 4,
+                 static_pos2d_ch: int = 4,
+                 embed_dim: int = 192,
+                 blocks: int = 4,
+                 heads_v: int = 4,
+                 heads_xy: int = 4,
+                 attn_drop_xy: float = 0.0):
+        super().__init__()
+        self.C = int(C)
+        self.Tin = int(Tin)
+        self.Tout = int(Tout)
+        self.N = int(N_mesh)
+        self.static_in = int(static_geo_ch + static_pos2d_ch)
+        E = int(embed_dim)
+
+        # зафіксуємо ребра як буфери
+        if edge_index_base.dtype != torch.long:
+            edge_index_base = edge_index_base.long()
+        self.register_buffer("edge_src", edge_index_base[0].contiguous())
+        self.register_buffer("edge_dst", edge_index_base[1].contiguous())
+
+        # diffusion-step embedding (опц.)
+        self.diffusion_embed = nn.Sequential(
+            nn.Linear(1, E), nn.ReLU(inplace=True),
+            nn.Linear(E, E)
+        )
+
+        # Temporal projection: [B, Tin, C, N] -> [B, E, C, N]
+        self.temporal_proj = nn.Conv2d(self.Tin, E, kernel_size=1)
+
+        # Static encoder: [B, static_in, 1, N] -> [B, E, 1, N] -> broadcast по C
+        self.static_encoder = nn.Sequential(
+            nn.Conv2d(self.static_in, 64, kernel_size=1),
+            nn.GroupNorm(8, 64), nn.ReLU(inplace=True),
+            nn.Conv2d(64, E, kernel_size=1),
+            nn.GroupNorm(8, E), nn.ReLU(inplace=True),
+        )
+
+        # Blocks
+        self.blocks = nn.ModuleList([
+            GraphAxialGATBlock(E, heads_v=heads_v, heads_xy=heads_xy, attn_drop_xy=attn_drop_xy)
+            for _ in range(blocks)
+        ])
+
+        # Decoder: [B, E, C, N] -> [B, Tout, C, N]
+        self.decoder = nn.Sequential(
+            nn.Conv2d(E, E, kernel_size=1), nn.GELU(),
+            nn.Conv2d(E, self.Tout, kernel_size=1)
+        )
+
+    def forward(self,
+                inputs: Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+                diffusion_step: Optional[torch.Tensor] = None) -> torch.Tensor:
+        x_all, geo, pos2d = inputs
+        x_all = x_all.float()  # [B, Tin, C, N]
+        geo   = geo.float()    # [B, G,   N]
+        pos2d = pos2d.float()  # [B, P,   N]
+
+        B, Tin, C, N = x_all.shape
+        assert Tin == self.Tin and C == self.C and N == self.N, \
+            f"Shape mismatch: x_all={x_all.shape} expected Tin={self.Tin}, C={self.C}, N={self.N}"
+
+        # (1) Temporal projection
+        F = self.temporal_proj(x_all)   # [B, E, C, N]
+
+        # (2) Diffusion step embedding (optional)
+        if diffusion_step is None:
+            diffusion_step = torch.zeros(B, device=x_all.device, dtype=x_all.dtype)
+        d_emb = self.diffusion_embed(diffusion_step.unsqueeze(-1))          # [B, E]
+        d_map = d_emb.view(B, -1, 1, 1).expand(-1, -1, C, N)                # [B, E, C, N]
+        F = F + d_map
+
+        # (3) Static features (broadcast per channel)
+        static = torch.cat([geo, pos2d], dim=1)          # [B, static_in, N]
+        static_e = self.static_encoder(static.unsqueeze(2))     # [B, E, 1, N]
+        F = F + static_e.expand(-1, -1, C, -1)                  # [B, E, C, N]
+
+        # (4) Axial GAT blocks
+        edge_index = torch.stack([self.edge_src, self.edge_dst], dim=0)
+        for blk in self.blocks:
+            F = blk(F, edge_index=edge_index)
+
+        # (5) Decode to Tout
+        y = self.decoder(F)  # [B, Tout, C, N]
+        return y
+
+
+# -----------------------------
+# Lightning wrapper
+# -----------------------------
+class UnifiedGraphWeatherGATLightning(pl.LightningModule):
+    """
+    Очікує batches у форматі:
+        ((x_all, geo, pos2d), y_all)
+    де:
+        x_all: [B, Tin, C, N]
+        y_all: [B, Tout, C, N]
+    """
+    def __init__(self,
+                 C: int,
+                 Tin: int,
+                 Tout: int,
+                 edge_index_base: torch.Tensor,
+                 N_mesh: int,
+                 embed_dim: int = 192,
+                 blocks: int = 4,
+                 heads_v: int = 4,
+                 heads_xy: int = 4,
+                 attn_drop_xy: float = 0.0,
+                 static_geo_ch: int = 4,
+                 static_pos2d_ch: int = 4,
+                 lr: float = 3e-4,
+                 weight_decay: float = 1e-4,
+                 loss: Literal["mse", "huber"] = "mse"):
+        super().__init__()
+        self.save_hyperparameters(ignore=["edge_index_base"])
+        self.net = UnifiedGraphSpatioVerticalGATNet(
+            C=C, Tin=Tin, Tout=Tout, N_mesh=N_mesh,
+            edge_index_base=edge_index_base,
+            static_geo_ch=static_geo_ch, static_pos2d_ch=static_pos2d_ch,
+            embed_dim=embed_dim, blocks=blocks,
+            heads_v=heads_v, heads_xy=heads_xy, attn_drop_xy=attn_drop_xy,
+        )
+        self.criterion = nn.MSELoss() if loss == "mse" else nn.HuberLoss()
+        self.lr = lr
+        self.weight_decay = weight_decay
+
+    def forward(self, x_tuple, diffusion_step: Optional[torch.Tensor] = None):
+        return self.net(x_tuple, diffusion_step=diffusion_step)
+
+    def _step(self, batch, stage: str):
+        (x_all, geo, pos2d), y_true = batch
+        B = x_all.size(0)
+        t = torch.zeros(B, device=x_all.device, dtype=x_all.dtype)  # нульовий diffusion-step
+        y_pred = self.forward((x_all, geo, pos2d), diffusion_step=t)
+        loss = self.criterion(y_pred, y_true)
+        self.log(f"{stage}_loss", loss, prog_bar=True, on_step=(stage=="train"), on_epoch=True, batch_size=B)
+        return loss
+
+    def training_step(self, batch, batch_idx):
+        return self._step(batch, "train")
+
+    def validation_step(self, batch, batch_idx):
+        self._step(batch, "val")
+
+    def configure_optimizers(self):
+        opt = torch.optim.AdamW(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
+        sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=self.trainer.max_epochs)
+        return {"optimizer": opt, "lr_scheduler": {"scheduler": sch, "interval": "epoch"}}

unified_graph_weather_model_tp.py

@@ -0,0 +1,272 @@
+# model/unified_graph_weather_model_tp.py
+import math
+from typing import Optional, Sequence, Tuple
+
+import torch
+import torch.nn as nn
+import pytorch_lightning as pl
+
+# базова загальна мережа
+from model.unified_graph_weather_model import UnifiedGraphSpatioVerticalGATNet
+
+
+# ----------------------- допоміжні утиліти -----------------------
+def _coerce_list_str(x) -> list[str]:
+    if isinstance(x, (list, tuple)):
+        return [str(v.decode("utf-8")) if isinstance(v, (bytes, bytearray)) else str(v) for v in x]
+    if isinstance(x, str):
+        return [x]
+    return [str(x)]
+
+
+def _find_tp_index(order: Sequence[str]) -> int:
+    """
+    Знаходимо індекс каналу опадів у shard order.
+    Пробуємо 'tp_log@sfc', 'tp@sfc', 'tp_log', 'tp' (в порядку пріоритету).
+    """
+    order = [str(v) for v in order]
+    low   = [s.lower() for s in order]
+    for cand in ("tp_log@sfc", "tp@sfc", "tp_log", "tp"):
+        if cand.lower() in low:
+            return low.index(cand.lower())
+    raise KeyError(f"Не знайшов 'tp' каналу у channel_order. Приклади: {order[:6]} ...")
+
+
+def _stats_key_for_tp(name: str) -> str:
+    # у статистиці ми використовуємо лог-опади
+    return "tp_log"
+
+
+# ---------- зручні перетворення (нормалізований лог ↔ фізичні мм) ----------
+class TPDenormHelper:
+    """
+    Конвертації між нормалізованим 'tp_log' і фізичними мм.
+    Використовує ROI-статистику: mean/std для 'tp_log'.
+    """
+    def __init__(self, stats_order: Sequence[str], mu: torch.Tensor, sd: torch.Tensor):
+        if isinstance(mu, (list, tuple)):
+            mu = torch.tensor(mu, dtype=torch.float32)
+        if isinstance(sd, (list, tuple)):
+            sd = torch.tensor(sd, dtype=torch.float32)
+
+        key = _stats_key_for_tp("tp")
+        idx = _coerce_list_str(stats_order).index(key)
+        self.mu = float(mu[idx])
+        self.sd = float(sd[idx])
+
+    def lognorm_to_mm(self, arr_norm: torch.Tensor) -> torch.Tensor:
+        # arr_norm → arr_log → мм
+        arr_log = arr_norm * self.sd + self.mu
+        mm = torch.expm1(arr_log)
+        return torch.clamp(mm, min=0.0)
+
+    def mm_to_lognorm(self, mm: torch.Tensor) -> torch.Tensor:
+        arr_log = torch.log1p(torch.clamp(mm, min=0.0))
+        return (arr_log - self.mu) / max(self.sd, 1e-6)
+
+
+# ----------------------- лосс під опади -----------------------
+class SparsePrecipLoss(nn.Module):
+    """
+    Комбінований лосс для опадів:
+      L = huber(log_norm) + α * w * |mm_pred - mm_true|
+    де w підсилює "мокрі" пікселі (mm_true >= wet_thr_mm), а для дуже малих опадів
+    можна тримати меншу вагу.
+
+    Параметри:
+      huber_delta: δ для Huber в нормалізованому лог-просторі
+      wet_thr_mm:  поріг "мокро" в міліметрах
+      wet_weight:  множник ваги для мокрих пікселів
+      dry_weight:  вага для сухих/майже сухих пікселів
+      alpha_mm:    коефіцієнт при терміні в мм
+      gamma_focal: опційно фокалізація для мокрих (підсилює великі значення)
+    """
+    def __init__(self,
+                 huber_delta: float = 0.5,
+                 wet_thr_mm: float = 0.1,
+                 wet_weight: float = 5.0,
+                 dry_weight: float = 0.5,
+                 alpha_mm: float = 0.5,
+                 gamma_focal: float = 0.0):
+        super().__init__()
+        self.delta = float(huber_delta)
+        self.wet_thr = float(wet_thr_mm)
+        self.w_wet = float(wet_weight)
+        self.w_dry = float(dry_weight)
+        self.alpha = float(alpha_mm)
+        self.gamma = float(gamma_focal)
+        self.huber = nn.SmoothL1Loss(beta=self.delta, reduction="none")  # PyTorch Huber v2
+
+    def forward(self,
+                y_pred_norm: torch.Tensor,   # [B, N]  (tp_log нормалізований)
+                y_true_norm: torch.Tensor,   # [B, N]
+                denorm_helper: TPDenormHelper) -> torch.Tensor:
+
+        # 1) Huber у нормалізованому лог-просторі
+        L_h = self.huber(y_pred_norm, y_true_norm)  # [B, N]
+
+        if self.alpha <= 0.0:
+            return L_h.mean()
+
+        # 2) Додатковий термін у фізичних мм
+        mm_pred = denorm_helper.lognorm_to_mm(y_pred_norm)  # [B,N]
+        mm_true = denorm_helper.lognorm_to_mm(y_true_norm)  # [B,N]
+        L_mm = torch.abs(mm_pred - mm_true)
+
+        # Ваги: мокре/сухе
+        wet_mask = (mm_true >= self.wet_thr).float()
+        dry_mask = 1.0 - wet_mask
+        w = self.w_wet * wet_mask + self.w_dry * dry_mask
+
+        if self.gamma > 0.0:
+            # фокалізація: підсилити великі значення (не обов’язково)
+            w = w * torch.pow(1.0 + mm_true, self.gamma)
+
+        L = L_h + self.alpha * (w * L_mm)
+        return L.mean()
+
+
+# ------------------- TP-адаптер над загальною моделлю -------------------
+class UnifiedGraphSpatioVerticalGATNetTP(nn.Module):
+    """
+    Обгортка над загальною моделлю, що повертає тільки 1 канал (tp_log@sfc).
+    Внутрішня вартість майже як у повної, але head узятий з бази (ми лише обрізаємо вихід).
+    """
+    def __init__(self,
+                 C: int, Tin: int, Tout: int,
+                 edge_index_base: torch.Tensor,
+                 N_mesh: Optional[int] = None,
+                 embed_dim: int = 192, blocks: int = 4, heads_v: int = 4, heads_xy: int = 4,
+                 attn_drop_xy: float = 0.0,
+                 channel_order: Optional[Sequence[str]] = None):
+        super().__init__()
+        self.base = UnifiedGraphSpatioVerticalGATNet(
+            C=C, Tin=Tin, Tout=Tout,
+            edge_index_base=edge_index_base, N_mesh=N_mesh,
+            embed_dim=embed_dim, blocks=blocks, heads_v=heads_v, heads_xy=heads_xy,
+            attn_drop_xy=attn_drop_xy,
+        )
+        self.channel_order = list(channel_order) if channel_order is not None else None
+        self.idx_tp = None
+        if self.channel_order is not None:
+            self.idx_tp = _find_tp_index(self.channel_order)
+
+    def set_channel_order(self, order: Sequence[str]):
+        self.channel_order = list(order)
+        self.idx_tp = _find_tp_index(self.channel_order)
+
+    def forward(self, x_tuple, diffusion_step: Optional[torch.Tensor] = None,
+                channel_order: Optional[Sequence[str]] = None):
+        """
+        x_tuple: (x_all, geo, pos2d)
+          x_all: [B, Tin, C, N]
+          geo:   [B, 4,   N]
+          pos2d: [B, 4,   N]
+        return: [B, Tout, 1, N]  (tp_log нормалізований)
+        """
+        if (self.idx_tp is None) and (channel_order is not None):
+            self.set_channel_order(channel_order)
+
+        y_full = self.base(x_tuple, diffusion_step=diffusion_step)  # [B,Tout,C,N]
+        if self.idx_tp is None:
+            raise RuntimeError("Не визначено індекс tp у channel_order. Передай order у forward або set_channel_order().")
+        tp = y_full[:, :, self.idx_tp].unsqueeze(2)  # [B, Tout, 1, N]
+        return tp
+
+
+# ------------------------- Lightning під опади -------------------------
+class PrecipTPLightning(pl.LightningModule):
+    """
+    1-кроковий TP (tp_log@sfc), лосс адаптований до розрідженості.
+    Підтримує файнтюн з загального чекпойнта (через .net.base.load_state_dict()).
+    """
+    def __init__(self,
+                 C: int, Tin: int, Tout: int,
+                 edge_index_base: torch.Tensor,
+                 N_mesh: Optional[int] = None,
+                 embed_dim: int = 192, blocks: int = 4, heads_v: int = 4, heads_xy: int = 4,
+                 attn_drop_xy: float = 0.0,
+                 lr: float = 3e-4, weight_decay: float = 1e-4,
+                 # лосс-параметри
+                 huber_delta: float = 0.5,
+                 wet_thr_mm: float = 0.1,
+                 wet_weight: float = 5.0,
+                 dry_weight: float = 0.5,
+                 alpha_mm: float = 0.5,
+                 gamma_focal: float = 0.0,
+                 diffusion_max_step: int = 0,
+                 channel_order: Optional[Sequence[str]] = None,
+                 stats_order: Optional[Sequence[str]] = None,
+                 stats_mu: Optional[Sequence[float]] = None,
+                 stats_sd: Optional[Sequence[float]] = None):
+        super().__init__()
+        self.save_hyperparameters(ignore=["edge_index_base", "channel_order", "stats_order", "stats_mu", "stats_sd"])
+        self.net = UnifiedGraphSpatioVerticalGATNetTP(
+            C=C, Tin=Tin, Tout=Tout,
+            edge_index_base=edge_index_base, N_mesh=N_mesh,
+            embed_dim=embed_dim, blocks=blocks, heads_v=heads_v, heads_xy=heads_xy,
+            attn_drop_xy=attn_drop_xy,
+            channel_order=channel_order,
+        )
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.diffusion_max_step = int(diffusion_max_step)
+
+        self.crit = SparsePrecipLoss(
+            huber_delta=huber_delta,
+            wet_thr_mm=wet_thr_mm,
+            wet_weight=wet_weight,
+            dry_weight=dry_weight,
+            alpha_mm=alpha_mm,
+            gamma_focal=gamma_focal,
+        )
+
+        # для другого терміну лоссу потрібні ROI-статистики
+        if (stats_order is None) or (stats_mu is None) or (stats_sd is None):
+            self.denorm_helper = None
+        else:
+            # збережемо в torch.Tensors для зручності на GPU
+            self.register_buffer("_stats_mu", torch.tensor(stats_mu, dtype=torch.float32))
+            self.register_buffer("_stats_sd", torch.tensor(stats_sd, dtype=torch.float32))
+            self.stats_order = list(stats_order)
+            self.denorm_helper = TPDenormHelper(self.stats_order, self._stats_mu, self._stats_sd)
+
+    # ----------- shared step -----------
+    def _shared_step(self, batch, stage: str):
+        (x_all, geo, pos2d), y_full = batch            # y_full: [B, Tout, C, N]
+        B = x_all.size(0)
+        # дифузійний крок як у загальній
+        steps = torch.zeros((B,), device=x_all.device, dtype=torch.float32)
+        if self.diffusion_max_step > 0:
+            steps = torch.randint(0, self.diffusion_max_step, (B,), device=x_all.device, dtype=torch.float32)
+
+        # канал-ордер беремо з DM
+        order = getattr(self.trainer.datamodule, "channel_order", None)
+        if order is None:
+            raise RuntimeError("DataModule не заповнив channel_order.")
+        # прогноз [B,Tout,1,N]
+        y_hat = self.net((x_all, geo, pos2d), diffusion_step=steps, channel_order=order)
+        y_true_tp = y_full[:, :, self.net.idx_tp].unsqueeze(2)  # [B,Tout,1,N]
+
+        y_pred = y_hat[:, 0, 0]   # [B,N] нормалізований tp_log
+        y_true = y_true_tp[:, 0, 0]
+
+        if self.denorm_helper is None:
+            loss = nn.SmoothL1Loss(beta=0.5)(y_pred, y_true)  # fallback, якщо нема стат
+        else:
+            loss = self.crit(y_pred, y_true, self.denorm_helper)
+
+        self.log_dict({f"{stage}/loss": loss}, prog_bar=True, on_step=(stage == "train"), on_epoch=True)
+        return loss
+
+    def training_step(self, batch, batch_idx):
+        return self._shared_step(batch, "train")
+
+    def validation_step(self, batch, batch_idx):
+        self._shared_step(batch, "val")
+
+    # ----------- optim -----------
+    def configure_optimizers(self):
+        opt = torch.optim.AdamW(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
+        sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(self.trainer.max_epochs, 1))
+        return {"optimizer": opt, "lr_scheduler": {"scheduler": sch, "interval": "epoch"}}

unified_graph_weather_model_uv10.py

@@ -0,0 +1,229 @@
+# model_unified_graph_uv10.py
+import math
+from typing import Optional, Tuple, Sequence
+
+import torch
+import torch.nn as nn
+import pytorch_lightning as pl
+
+# ВАЖЛИВО: імпортуємо вашу загальну модель
+from model.unified_graph_weather_model import UnifiedGraphSpatioVerticalGATNet
+
+
+# ------------------------- Допоміжний векторний лосс -------------------------
+class WindVectorLoss(nn.Module):
+    """
+    Лосс для вітру: компоненти + (опційно) модуль і напрям.
+    y_pred, y_true: [B, 2, N] (u, v) у НОРМАЛІЗОВАНОМУ масштабі (як у тренінгу).
+    Зазвичай достатньо component MSE. Якщо хочеться трохи "фізики", вмикаємо модуль/кут.
+    """
+    def __init__(self,
+                 w_components: float = 1.0,
+                 w_magnitude: float = 0.0,
+                 w_direction: float = 0.0,
+                 robust: bool = False,
+                 eps: float = 1e-6):
+        super().__init__()
+        self.wc = float(w_components)
+        self.wm = float(w_magnitude)
+        self.wd = float(w_direction)
+        self.eps = float(eps)
+        self.comp_loss = nn.SmoothL1Loss(reduction="mean") if robust else nn.MSELoss(reduction="mean")
+        self.mag_loss  = nn.SmoothL1Loss(reduction="mean") if robust else nn.MSELoss(reduction="mean")
+
+    def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
+        # компоненти
+        Lc = self.comp_loss(y_pred, y_true)
+
+        if (self.wm == 0.0) and (self.wd == 0.0):
+            return self.wc * Lc
+
+        # модуль
+        mag_p = torch.sqrt(y_pred[:, 0]**2 + y_pred[:, 1]**2 + self.eps)  # [B, N]
+        mag_t = torch.sqrt(y_true[:, 0]**2 + y_true[:, 1]**2 + self.eps)
+        Lm = self.mag_loss(mag_p, mag_t)
+
+        # напрям (1 - cos(Δθ)) через косинусну схожість
+        dot = (y_pred[:, 0]*y_true[:, 0] + y_pred[:, 1]*y_true[:, 1])  # [B, N]
+        cos_sim = dot / (mag_p * mag_t + self.eps)
+        # clamp щоб уникнути NaN при рідкісних виродженнях
+        cos_sim = torch.clamp(cos_sim, -1.0, 1.0)
+        Ld = (1.0 - cos_sim).mean()
+
+        return self.wc * Lc + self.wm * Lm + self.wd * Ld
+
+
+# ----------------- UV-адаптер над загальною Graph GAT моделлю ----------------
+class UnifiedGraphSpatioVerticalGATNetUV(nn.Module):
+    """
+    Тонкий адаптер: всередині — Ваша загальна мережа.
+    На виході повертаємо лише 2 канали (u10, v10).
+    Обчислювальні витрати майже ті самі (всі внутрішні блоки працюють як і раніше),
+    економія — лише на невеликій «голові». Це компроміс заради сумісності з чекпойнтом.
+
+    Параметри:
+      - chan_u_name, chan_v_name: як канали називаються у shard'ах (звично "u10@sfc", "v10@sfc")
+      - channel_order: список з шарду для пошуку індексів
+    """
+    def __init__(self,
+                 C: int, Tin: int, Tout: int,
+                 edge_index_base: torch.Tensor,
+                 N_mesh: Optional[int] = None,
+                 embed_dim: int = 192, blocks: int = 4, heads_v: int = 4, heads_xy: int = 4,
+                 attn_drop_xy: float = 0.0,
+                 chan_u_name: str = "u10@sfc",
+                 chan_v_name: str = "v10@sfc",
+                 channel_order: Optional[Sequence[str]] = None):
+        super().__init__()
+        self.chan_u_name = chan_u_name
+        self.chan_v_name = chan_v_name
+        self.channel_order = list(channel_order) if channel_order is not None else None
+
+        # базова мережа як у загальній моделі (C_out = C)
+        self.base = UnifiedGraphSpatioVerticalGATNet(
+            C=C, Tin=Tin, Tout=Tout,
+            edge_index_base=edge_index_base, N_mesh=N_mesh,
+            embed_dim=embed_dim, blocks=blocks, heads_v=heads_v, heads_xy=heads_xy,
+            attn_drop_xy=attn_drop_xy,
+        )
+
+        # індекси каналів у повному виході
+        if self.channel_order is not None:
+            self.idx_u, self.idx_v = self._find_uv_indices(self.channel_order)
+        else:
+            # якщо не дали під час ініт — довизначимо при першому forward (по наданому channel_order)
+            self.idx_u, self.idx_v = None, None
+
+    @staticmethod
+    def _find_uv_indices(order: Sequence[str],
+                         u_candidates=("u10@sfc", "u10"),
+                         v_candidates=("v10@sfc", "v10")) -> Tuple[int, int]:
+        lower = [str(x).lower() for x in order]
+        iu, iv = None, None
+        for cand in u_candidates:
+            if cand.lower() in lower:
+                iu = lower.index(cand.lower()); break
+        for cand in v_candidates:
+            if cand.lower() in lower:
+                iv = lower.index(cand.lower()); break
+        if iu is None or iv is None:
+            raise KeyError(f"Не знайшов індекси u/v у channel_order. order приклад: {order[:5]} ...")
+        return int(iu), int(iv)
+
+    def set_channel_order(self, order: Sequence[str]):
+        self.channel_order = list(order)
+        self.idx_u, self.idx_v = self._find_uv_indices(self.channel_order)
+
+    @torch.no_grad()
+    def _lazy_indices_from_order(self, order: Sequence[str]):
+        if (self.idx_u is None) or (self.idx_v is None):
+            self.set_channel_order(order)
+
+    def forward(self, x_tuple, diffusion_step: Optional[torch.Tensor] = None,
+                channel_order: Optional[Sequence[str]] = None):
+        """
+        x_tuple: (x_all, geo, pos2d)
+          x_all: [B, Tin, C, N]
+          geo:   [B, 4,   N]
+          pos2d: [B, 4,   N]
+        return: [B, Tout, 2, N]  (u, v)
+        """
+        if (self.idx_u is None or self.idx_v is None) and (channel_order is not None):
+            self._lazy_indices_from_order(channel_order)
+
+        y_full = self.base(x_tuple, diffusion_step=diffusion_step)  # [B, Tout, C, N]
+        if self.idx_u is None or self.idx_v is None:
+            raise RuntimeError("UV індекси не визначені (передайте channel_order у forward або set_channel_order()).")
+        uv = torch.stack([y_full[:, :, self.idx_u], y_full[:, :, self.idx_v]], dim=2)  # [B, Tout, 2, N]
+        return uv
+
+
+# ----------------------- Lightning: UV-файнтюнінг ----------------------------
+class WindUV10Lightning(pl.LightningModule):
+    """
+    Lightning-клас для файнтюнінгу під (u10, v10).
+    - Повертає лише 2-канальний вихід
+    - Лосс: компоненти + (опц.) модуль/напрям
+    - Може підвантажувати «знання» з загальної моделі (ckpt)
+    """
+    def __init__(self,
+                 C: int, Tin: int, Tout: int,
+                 edge_index_base: torch.Tensor,
+                 N_mesh: Optional[int] = None,
+                 embed_dim: int = 192, blocks: int = 4, heads_v: int = 4, heads_xy: int = 4,
+                 attn_drop_xy: float = 0.0,
+                 lr: float = 3e-4, weight_decay: float = 1e-4,
+                 loss_components: float = 1.0,
+                 loss_magnitude: float = 0.0,
+                 loss_direction: float = 0.0,
+                 robust_loss: bool = False,
+                 chan_u_name: str = "u10@sfc",
+                 chan_v_name: str = "v10@sfc",
+                 channel_order: Optional[Sequence[str]] = None,
+                 diffusion_max_step: int = 100):
+        super().__init__()
+        self.save_hyperparameters(ignore=["edge_index_base", "channel_order"])
+        self.net = UnifiedGraphSpatioVerticalGATNetUV(
+            C=C, Tin=Tin, Tout=Tout,
+            edge_index_base=edge_index_base, N_mesh=N_mesh,
+            embed_dim=embed_dim, blocks=blocks, heads_v=heads_v, heads_xy=heads_xy,
+            attn_drop_xy=attn_drop_xy,
+            chan_u_name=chan_u_name, chan_v_name=chan_v_name,
+            channel_order=channel_order,
+        )
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.diffusion_max_step = int(diffusion_max_step)
+        self.criterion = WindVectorLoss(
+            w_components=loss_components,
+            w_magnitude=loss_magnitude,
+            w_direction=loss_direction,
+            robust=robust_loss,
+        )
+
+    # -------- helper: вирізати GT (u,v) із повного y ----------
+    def _select_uv_from_full(self, y_full: torch.Tensor, order: Sequence[str]) -> torch.Tensor:
+        """
+        y_full: [B, Tout, C, N] -> [B, Tout, 2, N]
+        """
+        self.net._lazy_indices_from_order(order)
+        iu, iv = self.net.idx_u, self.net.idx_v
+        return torch.stack([y_full[:, :, iu], y_full[:, :, iv]], dim=2)
+
+    # ----------------------- train/val step -----------------------
+    def _shared_step(self, batch, stage: str):
+        (x_all, geo, pos2d), y_full = batch                  # y_full: [B, Tout, C, N]
+        # дифузійний крок (як у загальній моделі)
+        steps = torch.randint(
+            low=0, high=max(self.diffusion_max_step, 1),
+            size=(x_all.size(0),), device=x_all.device, dtype=torch.float32
+        )
+        # forward -> [B, Tout, 2, N]
+        y_hat = self.net((x_all, geo, pos2d), diffusion_step=steps,
+                         channel_order=getattr(self.trainer.datamodule, "channel_order", None))
+
+        # GT тільки UV
+        order = getattr(self.trainer.datamodule, "channel_order", None)
+        if order is None:
+            raise RuntimeError("DataModule не надав channel_order.")
+        y_uv = self._select_uv_from_full(y_full, order)     # [B, Tout, 2, N]
+
+        # lead-1 (Tout=1)
+        y_pred = y_hat[:, 0]                                # [B, 2, N]
+        y_true = y_uv[:, 0]                                 # [B, 2, N]
+
+        loss = self.criterion(y_pred, y_true)
+        self.log_dict({f"{stage}/loss": loss}, prog_bar=True, on_step=(stage=="train"), on_epoch=True)
+        return loss
+
+    def training_step(self, batch, batch_idx):
+        return self._shared_step(batch, "train")
+
+    def validation_step(self, batch, batch_idx):
+        self._shared_step(batch, "val")
+
+    # ----------------------- optim -----------------------
+    def configure_optimizers(self):
+        opt = torch.optim.AdamW(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
+        sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(self.trainer.max_epochs, 1))
+        return {"optimizer": opt, "lr_scheduler": {"scheduler": sch, "interval": "epoch"}}