Initial commit

.cursorignore

@@ -0,0 +1,17 @@
+# Python cache
+__pycache__/
+*.pyc
+
+# Experiment artifacts
+cache/
+logs/
+wandb/
+figs/
+checkpoints/*.pth
+checkpoints/*.bin
+
+# Data files
+examples/data/*.p
+*.pkl
+*.pickle
+

.gitattributes

@@ -0,0 +1,3 @@
+# Auto detect text files and perform LF normalization
+* text=auto
+*.bin filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,76 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyTorch
+*.pth
+*.pt
+*.bin
+*.ckpt
+!checkpoints/**/*.pth
+!checkpoints/**/*.bin
+!checkpoints/**/*.json
+!LWMTemporal/models/config.json
+
+# Data
+*.p
+*.pkl
+*.pickle
+*.h5
+*.hdf5
+cache/
+data/
+!examples/data/
+!examples/data/*.p
+!examples/data/README.md
+
+# Experiments
+logs/
+figs/
+wandb/
+outputs/
+# checkpoints/
+runs/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+.DS_Store
+
+# Testing
+.pytest_cache/
+.coverage
+htmlcov/
+.tox/
+
+# Environment
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+

LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2025 Sadjad Alikhani
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+

LWMTemporal/__init__.py

@@ -0,0 +1,10 @@
+"""LWM Temporal Model package."""
+
+import warnings
+
+warnings.filterwarnings("ignore")
+
+from .models.lwm import LWMConfig, LWMModel, LWMBackbone
+
+__version__ = "0.1.0"
+__all__ = ["LWMConfig", "LWMModel", "LWMBackbone", "__version__"]

LWMTemporal/cli/__init__.py

@@ -0,0 +1 @@
+"""Command line entrypoints for the LWM foundation package."""

LWMTemporal/cli/channel_prediction.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Optional, Sequence
+import torch
+
+from ..tasks.channel_prediction import (
+    ChannelPredictionArgs,
+    ChannelPredictionTrainer,
+    DatasetArgs,
+    ModelArgs,
+    PredictionArgs,
+    TrainingArgs,
+)
+from ..utils.logging import setup_logging
+
+
+def parse_args(argv: Optional[Sequence[str]] = None) -> ChannelPredictionArgs:
+    parser = argparse.ArgumentParser(description="Channel prediction trainer")
+    parser.add_argument("--data_path", type=Path, required=True)
+    parser.add_argument("--keep_percentage", type=float, default=0.25)
+    parser.add_argument("--normalize", type=str, default="global_rms", choices=["global_rms", "per_sample_rms", "none"])
+    parser.add_argument("--cache_dir", type=Path, default=Path("cache"))
+    parser.add_argument("--no_cache", action="store_true")
+    parser.add_argument("--overwrite_cache", action="store_true")
+    parser.add_argument("--snr_db", type=float, default=None)
+    parser.add_argument("--noise_seed", type=int, default=None)
+    parser.add_argument("--max_time_steps", type=int, default=None)
+    parser.add_argument("--train_limit", type=int, default=500)
+    parser.add_argument("--val_limit", type=int, default=100)
+    parser.add_argument("--seed", type=int, default=42)
+
+    parser.add_argument("--patch_size", type=int, nargs=2, default=(1, 1))
+    parser.add_argument("--phase_mode", type=str, default="real_imag", choices=["real_imag", "mag_phase"])
+    parser.add_argument("--embed_dim", type=int, default=32)
+    parser.add_argument("--depth", type=int, default=12)
+    parser.add_argument("--num_heads", type=int, default=8)
+    parser.add_argument("--mlp_ratio", type=float, default=4.0)
+    parser.add_argument("--same_frame_window", type=int, default=2)
+    parser.add_argument("--temporal_offsets", type=int, nargs="*", default=[-1, -2, -3, -4, -5, -6, -7])
+    parser.add_argument("--temporal_spatial_window", type=int, default=2)
+    parser.add_argument("--temporal_drift_h", type=int, default=1)
+    parser.add_argument("--temporal_drift_w", type=int, default=1)
+    parser.add_argument("--routing_topk_enable", action="store_true", default=True)
+    parser.add_argument("--routing_topk_fraction", type=float, default=0.2)
+    parser.add_argument("--routing_topk_min", type=int, default=8)
+    parser.add_argument("--routing_topk_max", type=int, default=32)
+    parser.add_argument("--topk_per_head", action="store_true", default=True)
+    parser.add_argument("--posenc", type=str, default="learned", choices=["learned", "rope_sincos"])
+    parser.add_argument("--rope_base", type=float, default=10000.0)
+    parser.add_argument("--global_cls", action="store_true")
+    parser.add_argument("--pretrained", type=Path, default=None)
+    parser.add_argument("--finetune_last_n", type=int, default=0)
+    parser.add_argument("--train_head_only", action="store_true")
+
+    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--epochs", type=int, default=3)
+    parser.add_argument("--batch_size", type=int, default=16)
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--weight_decay", type=float, default=1e-4)
+    parser.add_argument("--warmup_ratio", type=float, default=0.1)
+    parser.add_argument("--loss", type=str, default="nmse", choices=["nmse", "mse"])
+    parser.add_argument("--use_dataparallel", action="store_true")
+    parser.add_argument("--grad_clip", type=float, default=1.0)
+    parser.add_argument("--log_interval", type=int, default=10)
+    parser.add_argument("--save_dir", type=Path, default=Path("models"))
+    parser.add_argument("--save_prefix", type=str, default="channel_prediction")
+    parser.add_argument("--inference_only", action="store_true")
+    parser.add_argument("--inference_split", type=str, default="val", choices=["train", "val", "all"])
+    parser.add_argument("--verbose_inference", action="store_true")
+    parser.add_argument("--log_dir", type=Path, default=Path("logs"))
+    parser.add_argument("--use_wandb", action="store_true")
+    parser.add_argument("--wandb_project", type=str, default=None)
+    parser.add_argument("--wandb_entity", type=str, default=None)
+    parser.add_argument("--wandb_run_name", type=str, default=None)
+
+    parser.add_argument("--Tpast", type=int, default=10)
+    parser.add_argument("--horizon", type=int, default=1)
+    parser.add_argument("--num_visual_samples", type=int, default=4)
+    parser.add_argument("--viz_dir", type=Path, default=Path("figs/predictions"))
+
+    ns = parser.parse_args(argv)
+
+    dataset = DatasetArgs(
+        data_path=ns.data_path,
+        keep_percentage=ns.keep_percentage,
+        normalize=ns.normalize,
+        cache_dir=ns.cache_dir,
+        use_cache=not ns.no_cache,
+        overwrite_cache=ns.overwrite_cache,
+        snr_db=ns.snr_db,
+        noise_seed=ns.noise_seed,
+        max_time_steps=ns.max_time_steps,
+        train_limit=ns.train_limit,
+        val_limit=ns.val_limit,
+        seed=ns.seed,
+    )
+
+    model = ModelArgs(
+        patch_size=tuple(ns.patch_size),
+        phase_mode=ns.phase_mode,
+        embed_dim=ns.embed_dim,
+        depth=ns.depth,
+        num_heads=ns.num_heads,
+        mlp_ratio=ns.mlp_ratio,
+        same_frame_window=ns.same_frame_window,
+        temporal_offsets=tuple(ns.temporal_offsets),
+        temporal_spatial_window=ns.temporal_spatial_window,
+        temporal_drift_h=ns.temporal_drift_h,
+        temporal_drift_w=ns.temporal_drift_w,
+        routing_topk_enable=ns.routing_topk_enable,
+        routing_topk_fraction=ns.routing_topk_fraction,
+        routing_topk_min=ns.routing_topk_min,
+        routing_topk_max=ns.routing_topk_max,
+        topk_per_head=ns.topk_per_head,
+        posenc=ns.posenc,
+        rope_base=ns.rope_base,
+        global_cls=ns.global_cls,
+        pretrained=ns.pretrained,
+        finetune_last_n=ns.finetune_last_n,
+        train_head_only=ns.train_head_only,
+    )
+
+    training = TrainingArgs(
+        device=ns.device,
+        epochs=ns.epochs,
+        batch_size=ns.batch_size,
+        lr=ns.lr,
+        weight_decay=ns.weight_decay,
+        warmup_ratio=ns.warmup_ratio,
+        loss=ns.loss,
+        use_dataparallel=ns.use_dataparallel,
+        grad_clip=ns.grad_clip,
+        log_interval=ns.log_interval,
+        save_dir=ns.save_dir,
+        save_prefix=ns.save_prefix,
+        inference_only=ns.inference_only,
+        inference_split=ns.inference_split,
+        verbose_inference=ns.verbose_inference,
+        log_dir=ns.log_dir,
+        use_wandb=ns.use_wandb,
+        wandb_project=ns.wandb_project,
+        wandb_entity=ns.wandb_entity,
+        wandb_run_name=ns.wandb_run_name,
+    )
+
+    prediction = PredictionArgs(
+        Tpast=ns.Tpast,
+        horizon=ns.horizon,
+        num_visual_samples=ns.num_visual_samples,
+        viz_dir=ns.viz_dir,
+    )
+
+    return ChannelPredictionArgs(dataset=dataset, model=model, training=training, prediction=prediction)
+
+
+def main(argv: Optional[Sequence[str]] = None) -> None:
+    args = parse_args(argv)
+    logger = setup_logging("LWMTemporal.channel_prediction", args.training.log_dir)
+    logger.info(
+        "Starting channel prediction run | device=%s inference_only=%s use_wandb=%s",
+        args.training.device,
+        args.training.inference_only,
+        args.training.use_wandb,
+    )
+    trainer = ChannelPredictionTrainer(args, logger=logger)
+    trainer.train()
+
+
+__all__ = ["parse_args", "main"]
+
+
+if __name__ == "__main__":
+    main()

LWMTemporal/cli/pretrain.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from typing import Optional, Sequence
+
+from ..tasks.pretraining import build_parser, build_pretraining_args, PretrainingTrainer
+from ..utils.logging import setup_logging
+
+
+def main(argv: Optional[Sequence[str]] = None) -> None:
+    parser = build_parser()
+    args_ns = parser.parse_args(args=list(argv) if argv is not None else None)
+    args = build_pretraining_args(args_ns)
+    logger = setup_logging("LWMTemporal.pretraining", args.logging.log_dir)
+    logger.info(
+        "Starting pretraining run | device=%s epochs=%d batch_size=%d use_wandb=%s",
+        args.optim.device,
+        args.optim.epochs,
+        args.optim.batch_size,
+        args.logging.use_wandb,
+    )
+    trainer = PretrainingTrainer(args, logger=logger)
+    trainer.train()
+
+
+__all__ = ["main"]

LWMTemporal/models/__init__.py

@@ -0,0 +1,3 @@
+from .lwm import LWMConfig, LWMModel, LWMBackbone
+
+__all__ = ["LWMConfig", "LWMModel", "LWMBackbone"]

LWMTemporal/models/config.json

@@ -0,0 +1,37 @@
+{
+    "patch_size": [1, 1],
+    "phase_mode": "real_imag",
+    "embed_dim": 32,
+    "depth": 12,
+    "num_heads": 8,
+    "mlp_ratio": 4.0,
+    "same_frame_window": 2,
+    "same_frame_window_h": null,
+    "same_frame_window_w": null,
+    "same_frame_dilation_h": 1,
+    "same_frame_dilation_w": 1,
+    "temporal_offsets": [-4, -3, -2, -1, 1, 2, 3],
+    "temporal_spatial_window": 2,
+    "temporal_spatial_window_h": null,
+    "temporal_spatial_window_w": null,
+    "temporal_spatial_dilation_h": 1,
+    "temporal_spatial_dilation_w": 1,
+    "temporal_drift_h": 1,
+    "temporal_drift_w": 1,
+    "spatial_only": false,
+    "routing_topk_enable": true,
+    "routing_topk_fraction": 0.2,
+    "routing_topk_min": 8,
+    "routing_topk_max": 32,
+    "routing_topk_per_head": true,
+    "topk_neighbors": null,
+    "topk_per_head": true,
+    "global_cls": false,
+    "posenc": "learned",
+    "rope_base": 10000.0,
+    "rope_mode": "flat",
+    "rope_base_t": null,
+    "rope_base_h": null,
+    "rope_base_w": null,
+    "max_seq_len": null
+  }

LWMTemporal/models/lwm.py

@@ -0,0 +1,576 @@
+from __future__ import annotations
+
+import json
+import math
+from dataclasses import dataclass, asdict, fields
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.utils.data import Dataset
+
+# -----------------------------------------------------------------------------
+# Tokenization
+# -----------------------------------------------------------------------------
+class ComplexPatchTokenizer:
+    def __init__(self, phase_mode: str = "real_imag") -> None:
+        if phase_mode not in {"real_imag", "mag_phase"}:
+            raise ValueError("phase_mode must be 'real_imag' or 'mag_phase'")
+        self.phase_mode = phase_mode
+
+    def _split_channels(self, tensor: Tensor) -> Tensor:
+        if self.phase_mode == "real_imag":
+            real = tensor.real.unsqueeze(-1)
+            imag = tensor.imag.unsqueeze(-1)
+            return torch.cat([real, imag], dim=-1)
+        magnitude = tensor.abs().unsqueeze(-1)
+        phase = torch.angle(tensor).unsqueeze(-1)
+        return torch.cat([magnitude, phase], dim=-1)
+
+    def __call__(self, seq: Tensor, patch_size: Tuple[int, int]) -> Tuple[Tensor, Tensor]:
+        if not torch.is_complex(seq):
+            raise TypeError("expected complex tensor shaped (B, T, N, M)")
+        ph, pw = patch_size
+        if seq.size(2) % ph != 0 or seq.size(3) % pw != 0:
+            raise ValueError("patch_size must evenly divide channel dimensions")
+        channels = self._split_channels(seq)
+        b, t, n, m, c = channels.shape
+        h = n // ph
+        w = m // pw
+        channels = channels.view(b, t, h, ph, w, pw, c)
+        channels = channels.permute(0, 1, 2, 4, 3, 5, 6).contiguous()
+        tokens = channels.view(b, t * h * w, ph * pw * c)
+        mask = torch.zeros((b, tokens.size(1)), dtype=torch.bool, device=tokens.device)
+        return tokens, mask
+
+
+# -----------------------------------------------------------------------------
+# Sparse spatio-temporal attention
+# -----------------------------------------------------------------------------
+@dataclass(frozen=True)
+class AttentionCacheKey:
+    temporal: int
+    height: int
+    width: int
+    same_frame_window: int
+    same_frame_window_h: Optional[int]
+    same_frame_window_w: Optional[int]
+    same_frame_dilation_h: int
+    same_frame_dilation_w: int
+    temporal_offsets: Tuple[int, ...]
+    temporal_spatial_window: int
+    temporal_spatial_window_h: Optional[int]
+    temporal_spatial_window_w: Optional[int]
+    temporal_spatial_dilation_h: int
+    temporal_spatial_dilation_w: int
+    temporal_drift_h: int
+    temporal_drift_w: int
+    include_cls: bool
+
+
+class NeighborIndexer:
+    def __init__(self) -> None:
+        self._cache: Dict[Tuple[int, int, int, AttentionCacheKey], Tensor] = {}
+
+    def get(self, T: int, H: int, W: int, include_cls: bool, config: "LWMConfig", device: torch.device) -> Tensor:
+        key = (
+            T,
+            H,
+            W,
+            AttentionCacheKey(
+                temporal=T,
+                height=H,
+                width=W,
+                same_frame_window=config.same_frame_window,
+                same_frame_window_h=config.same_frame_window_h,
+                same_frame_window_w=config.same_frame_window_w,
+                same_frame_dilation_h=config.same_frame_dilation_h,
+                same_frame_dilation_w=config.same_frame_dilation_w,
+                temporal_offsets=config.temporal_offsets,
+                temporal_spatial_window=config.temporal_spatial_window,
+                temporal_spatial_window_h=config.temporal_spatial_window_h,
+                temporal_spatial_window_w=config.temporal_spatial_window_w,
+                temporal_spatial_dilation_h=config.temporal_spatial_dilation_h,
+                temporal_spatial_dilation_w=config.temporal_spatial_dilation_w,
+                temporal_drift_h=config.temporal_drift_h,
+                temporal_drift_w=config.temporal_drift_w,
+                include_cls=include_cls,
+            ),
+        )
+        if key in self._cache:
+            tensor = self._cache[key]
+            return tensor if tensor.device == device else tensor.to(device)
+        indices = self._build_indices(T, H, W, include_cls, config)
+        if indices:
+            max_len = max(len(neighbors) for neighbors in indices)
+            if any(len(neighbors) != max_len for neighbors in indices):
+                padded = []
+                for neighbors in indices:
+                    if len(neighbors) < max_len:
+                        neighbors = neighbors + [-1] * (max_len - len(neighbors))
+                    padded.append(neighbors)
+                indices = padded
+        tensor = torch.as_tensor(indices, dtype=torch.long, device=device)
+        self._cache[key] = tensor
+        return tensor
+
+    def _build_indices(self, T: int, H: int, W: int, include_cls: bool, config: "LWMConfig") -> List[List[int]]:
+        neighbors: List[List[int]] = []
+        same_h = config.same_frame_window if config.same_frame_window_h is None else config.same_frame_window_h
+        same_w = config.same_frame_window if config.same_frame_window_w is None else config.same_frame_window_w
+
+        def frame_base(frame: int) -> int:
+            return frame * H * W
+
+        for t_idx in range(T):
+            base = frame_base(t_idx)
+            for h_idx in range(H):
+                for w_idx in range(W):
+                    current = base + h_idx * W + w_idx
+                    local: List[int] = []
+                    if config.same_frame_window < 0:
+                        local.extend(range(base, base + H * W))
+                    else:
+                        for dh in range(-same_h, same_h + 1, config.same_frame_dilation_h):
+                            for dw in range(-same_w, same_w + 1, config.same_frame_dilation_w):
+                                nh = h_idx + dh
+                                nw = w_idx + dw
+                                if 0 <= nh < H and 0 <= nw < W:
+                                    local.append(base + nh * W + nw)
+                    if not config.spatial_only:
+                        for dt in config.temporal_offsets:
+                            other_t = t_idx + dt
+                            if other_t < 0 or other_t >= T:
+                                continue
+                            other_base = frame_base(other_t)
+                            drift_h = config.temporal_spatial_window if config.temporal_drift_h == 0 else min(config.temporal_spatial_window, abs(dt) * config.temporal_drift_h)
+                            drift_w = config.temporal_spatial_window if config.temporal_drift_w == 0 else min(config.temporal_spatial_window, abs(dt) * config.temporal_drift_w)
+                            window_h = config.temporal_spatial_window if config.temporal_spatial_window_h is None else config.temporal_spatial_window_h
+                            window_w = config.temporal_spatial_window if config.temporal_spatial_window_w is None else config.temporal_spatial_window_w
+                            for dh in range(-min(window_h, drift_h), min(window_h, drift_h) + 1, config.temporal_spatial_dilation_h):
+                                for dw in range(-min(window_w, drift_w), min(window_w, drift_w) + 1, config.temporal_spatial_dilation_w):
+                                    nh = max(0, min(H - 1, h_idx + dh))
+                                    nw = max(0, min(W - 1, w_idx + dw))
+                                    local.append(other_base + nh * W + nw)
+                    if include_cls:
+                        local.append(T * H * W)
+                    if not local:
+                        local.append(current)
+                    neighbors.append(sorted(set(local)))
+        if include_cls:
+            neighbors.append(list(range(T * H * W)))
+        return neighbors
+
+
+class SparseSpatioTemporalAttention(nn.Module):
+    def __init__(self, config: "LWMConfig", embed_dim: int, num_heads: int) -> None:
+        super().__init__()
+        self.config = config
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        if self.head_dim * num_heads != embed_dim:
+            raise ValueError("embed_dim must be divisible by num_heads")
+        self.scale = self.head_dim ** -0.5
+        self.qkv = nn.Linear(embed_dim, 3 * embed_dim, bias=False)
+        self.proj = nn.Linear(embed_dim, embed_dim)
+        self.indexer = NeighborIndexer()
+
+    def _apply_rope(self, x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
+        x1 = x[..., ::2]
+        x2 = x[..., 1::2]
+        rotated_first = x1 * cos - x2 * sin
+        rotated_second = x1 * sin + x2 * cos
+        return torch.stack([rotated_first, rotated_second], dim=-1).flatten(-2)
+
+    def _rope_factors(self, S: int, device: torch.device) -> Tuple[Tensor, Tensor]:
+        half = self.head_dim // 2
+        inv_freq = 1.0 / (self.config.rope_base ** (torch.arange(0, half, dtype=torch.float32, device=device) / max(1, half)))
+        positions = torch.arange(S, dtype=torch.float32, device=device)
+        angles = positions[:, None] * inv_freq[None, :]
+        return torch.cos(angles)[None, None, :, :], torch.sin(angles)[None, None, :, :]
+
+    def forward(self, hidden_states: Tensor, T: int, H: int, W: int, include_cls: bool) -> Tensor:
+        bsz, seq_len, _ = hidden_states.shape
+        neighbors = self.indexer.get(T, H, W, include_cls, self.config, hidden_states.device)
+        qkv = self.qkv(hidden_states)
+        qkv = qkv.view(bsz, seq_len, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        if self.config.posenc == "rope_sincos":
+            cos, sin = self._rope_factors(seq_len, hidden_states.device)
+            q = self._apply_rope(q, cos, sin)
+            k = self._apply_rope(k, cos, sin)
+
+        gather_idx = neighbors.clamp_min(0)
+        valid_mask = neighbors >= 0
+        k = k[:, :, gather_idx, :]
+        v = v[:, :, gather_idx, :]
+        scores = torch.einsum("bhqd,bhqkd->bhqk", q, k) * self.scale
+        scores = scores.masked_fill(~valid_mask.unsqueeze(0).unsqueeze(0), float("-inf"))
+
+        if self.config.routing_topk_enable:
+            K = scores.size(-1)
+            keep = min(self.config.routing_topk_max, max(self.config.routing_topk_min, int(self.config.routing_topk_fraction * K)))
+            if self.config.routing_topk_per_head:
+                _, idx = torch.topk(scores, keep, dim=-1)
+                topk_mask = torch.zeros_like(scores, dtype=torch.bool)
+                topk_mask.scatter_(-1, idx, True)
+            else:
+                avg_scores = scores.mean(dim=1, keepdim=True)
+                _, idx = torch.topk(avg_scores, keep, dim=-1)
+                topk_mask = torch.zeros_like(scores, dtype=torch.bool)
+                topk_mask.scatter_(-1, idx.expand_as(scores), True)
+            scores = scores.masked_fill(~topk_mask, float("-inf"))
+        elif self.config.topk_neighbors is not None:
+            keep = min(self.config.topk_neighbors, scores.size(-1))
+            if self.config.topk_per_head:
+                _, idx = torch.topk(scores, keep, dim=-1)
+                topk_mask = torch.zeros_like(scores, dtype=torch.bool)
+                topk_mask.scatter_(-1, idx, True)
+            else:
+                avg_scores = scores.mean(dim=1, keepdim=True)
+                _, idx = torch.topk(avg_scores, keep, dim=-1)
+                topk_mask = torch.zeros_like(scores, dtype=torch.bool)
+                topk_mask.scatter_(-1, idx.expand_as(scores), True)
+            scores = scores.masked_fill(~topk_mask, float("-inf"))
+
+        attn = torch.softmax(scores, dim=-1)
+        attn = attn.masked_fill(~valid_mask.unsqueeze(0).unsqueeze(0), 0.0)
+        context = torch.einsum("bhqk,bhqkd->bhqd", attn, v)
+        context = context.transpose(1, 2).contiguous().view(bsz, seq_len, self.embed_dim)
+        return self.proj(context)
+
+
+class LWMEncoderLayer(nn.Module):
+    def __init__(self, config: "LWMConfig") -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(config.embed_dim)
+        self.attn = SparseSpatioTemporalAttention(config, config.embed_dim, config.num_heads)
+        self.norm2 = nn.LayerNorm(config.embed_dim)
+        hidden_dim = int(config.embed_dim * config.mlp_ratio)
+        self.mlp = nn.Sequential(
+            nn.Linear(config.embed_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, config.embed_dim),
+        )
+
+    def forward(self, x: Tensor, T: int, H: int, W: int, include_cls: bool) -> Tensor:
+        x = x + self.attn(self.norm1(x), T, H, W, include_cls)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class LWMEncoder(nn.Module):
+    def __init__(self, config: "LWMConfig") -> None:
+        super().__init__()
+        self.layers = nn.ModuleList([LWMEncoderLayer(config) for _ in range(config.depth)])
+        self.norm = nn.LayerNorm(config.embed_dim)
+
+    def forward(self, x: Tensor, T: int, H: int, W: int, include_cls: bool) -> Tensor:
+        for layer in self.layers:
+            x = layer(x, T, H, W, include_cls)
+        return self.norm(x)
+
+
+# -----------------------------------------------------------------------------
+# Hugging Face configuration and model definitions
+# -----------------------------------------------------------------------------
+@dataclass
+class LWMConfig:
+    patch_size: Tuple[int, int] = (1, 1)
+    phase_mode: str = "real_imag"
+    embed_dim: int = 32
+    depth: int = 12
+    num_heads: int = 8
+    mlp_ratio: float = 4.0
+    same_frame_window: int = 2
+    same_frame_window_h: Optional[int] = None
+    same_frame_window_w: Optional[int] = None
+    same_frame_dilation_h: int = 1
+    same_frame_dilation_w: int = 1
+    temporal_offsets: Tuple[int, ...] = (-4, -3, -2, -1, 1, 2, 3)
+    temporal_spatial_window: int = 2
+    temporal_spatial_window_h: Optional[int] = None
+    temporal_spatial_window_w: Optional[int] = None
+    temporal_spatial_dilation_h: int = 1
+    temporal_spatial_dilation_w: int = 1
+    temporal_drift_h: int = 1
+    temporal_drift_w: int = 1
+    spatial_only: bool = False
+    routing_topk_enable: bool = True
+    routing_topk_fraction: float = 0.2
+    routing_topk_min: int = 8
+    routing_topk_max: int = 32
+    routing_topk_per_head: bool = True
+    topk_neighbors: Optional[int] = None
+    topk_per_head: bool = True
+    global_cls: bool = False
+    posenc: str = "learned"
+    rope_base: float = 10000.0
+    rope_mode: str = "flat"
+    rope_base_t: Optional[float] = None
+    rope_base_h: Optional[float] = None
+    rope_base_w: Optional[float] = None
+    max_seq_len: Optional[int] = None
+
+    def __post_init__(self) -> None:
+        self.patch_size = (int(self.patch_size[0]), int(self.patch_size[1]))
+        self.temporal_offsets = tuple(int(o) for o in self.temporal_offsets)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "LWMConfig":
+        return cls(**data)
+
+
+class LWMModel(nn.Module):
+    def __init__(self, config: LWMConfig) -> None:
+        super().__init__()
+        self.config = config
+        patch_dim = config.patch_size[0] * config.patch_size[1] * 2
+        self.tokenizer = ComplexPatchTokenizer(config.phase_mode)
+        self.patch_embed = nn.Linear(patch_dim, config.embed_dim)
+        self.global_cls = config.global_cls
+        pos_len = (config.max_seq_len or 0) + (1 if self.global_cls else 0)
+        if pos_len == 0:
+            pos_len = 1
+        if config.posenc == "learned":
+            self.pos_embed = nn.Parameter(torch.zeros(1, pos_len, config.embed_dim))
+            nn.init.trunc_normal_(self.pos_embed, std=0.02)
+        else:
+            self.register_buffer("pos_embed", torch.zeros(1, pos_len, config.embed_dim), persistent=False)
+        if self.global_cls:
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dim))
+            nn.init.trunc_normal_(self.cls_token, std=0.02)
+        self.encoder = LWMEncoder(config)
+        self.head = nn.Linear(config.embed_dim, patch_dim)
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.trunc_normal_(module.weight, std=0.02)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, nn.LayerNorm):
+                nn.init.ones_(module.weight)
+                nn.init.zeros_(module.bias)
+
+    def _add_positional(self, tokens: Tensor) -> Tensor:
+        if self.config.posenc == "learned":
+            return tokens + self.pos_embed[:, : tokens.size(1)]
+        return tokens
+
+    def forward_tokens(
+        self,
+        tokens: Tensor,
+        mask: Tensor,
+        T: int,
+        H: int,
+        W: int,
+        *,
+        return_cls: bool = False,
+    ) -> Dict[str, Optional[Tensor]]:
+        embeddings = self.patch_embed(tokens)
+        include_cls = self.global_cls
+        if include_cls:
+            cls_tokens = self.cls_token.expand(embeddings.size(0), -1, -1)
+            embeddings = torch.cat([embeddings, cls_tokens], dim=1)
+            cls_mask = torch.zeros((embeddings.size(0), 1), dtype=torch.bool, device=embeddings.device)
+            mask = torch.cat([mask, cls_mask], dim=1)
+        # Add positional embeddings BEFORE masking (matching original implementation)
+        embeddings = self._add_positional(embeddings)
+        # Then mask embeddings (zeros out both token embedding AND positional embedding)
+        embeddings = embeddings.masked_fill(mask.unsqueeze(-1), 0.0)
+        encoded = self.encoder(embeddings, T, H, W, include_cls)
+        if include_cls:
+            reconstruction = self.head(encoded[:, :-1, :])
+            cls = encoded[:, -1, :]
+        else:
+            reconstruction = self.head(encoded)
+            cls = None
+        return {"reconstruction": reconstruction, "cls": cls if return_cls else None}
+
+    def forward(self, seq: Tensor, mask: Optional[Tensor] = None, *, return_cls: bool = False) -> Dict[str, Optional[Tensor]]:
+        tokens, base_mask = self.tokenizer(seq, self.config.patch_size)
+        total_mask = base_mask if mask is None else mask
+        ph, pw = self.config.patch_size
+        T = seq.size(1)
+        H = seq.size(2) // ph
+        W = seq.size(3) // pw
+        return self.forward_tokens(tokens, total_mask, T, H, W, return_cls=return_cls)
+
+    @torch.no_grad()
+    def forward_features(self, seq: Tensor) -> Tuple[Tensor, Optional[Tensor]]:
+        outputs = self.forward(seq, return_cls=True)
+        return outputs["reconstruction"], outputs["cls"]
+
+
+class LWMBackbone(LWMModel):
+    """Minor alias kept for backwards compatibility with legacy scripts."""
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: str | Path,
+        *model_args: Any,
+        config: Optional[LWMConfig] = None,
+        map_location: str | torch.device = "cpu",
+        **kwargs: Any,
+    ) -> "LWMBackbone":
+        path = Path(pretrained_model_name_or_path)
+        state: Dict[str, Tensor]
+        checkpoint_config: Optional[Dict[str, Any]] = None
+
+        if path.is_dir():
+            directory = path
+            state_path = directory / "pytorch_model.bin"
+            if not state_path.exists():
+                raise FileNotFoundError(f"Pretrained weights not found at {state_path}")
+            raw = torch.load(state_path, map_location=map_location)
+            if isinstance(raw, dict) and any(isinstance(v, torch.Tensor) for v in raw.values()):
+                state = {k: v for k, v in raw.items() if isinstance(v, torch.Tensor)}
+            else:
+                raise ValueError(f"Unexpected checkpoint format at {state_path}")
+            # Always try to load checkpoint config first, then merge with provided config
+            checkpoint_config_dict = None
+            config_path = directory / "config.json"
+            if config_path.exists():
+                with config_path.open("r") as handle:
+                    checkpoint_config_dict = json.load(handle)
+                    checkpoint_config = LWMConfig.from_dict(checkpoint_config_dict)
+                    if config is None:
+                        config = checkpoint_config
+                    else:
+                        # Merge: use checkpoint config as base, override with provided config
+                        checkpoint_dict = checkpoint_config.to_dict()
+                        provided_dict = config.to_dict()
+                        merged_dict = {**checkpoint_dict, **provided_dict}
+                        config = LWMConfig.from_dict(merged_dict)
+        else:
+            if not path.exists():
+                raise FileNotFoundError(f"Pretrained weights not found at {path}")
+            raw = torch.load(path, map_location=map_location)
+            if isinstance(raw, dict) and "model_state_dict" in raw:
+                state = raw["model_state_dict"]
+                checkpoint_config = raw.get("config")
+            elif isinstance(raw, dict):
+                state = {k: v for k, v in raw.items() if isinstance(v, torch.Tensor)}
+            else:
+                raise ValueError("Unsupported checkpoint format; expected a state_dict or training checkpoint.")
+
+        if config is None and checkpoint_config is not None:
+            config = cls._config_from_checkpoint(checkpoint_config)
+        if config is None:
+            config = LWMConfig()
+
+        if config.max_seq_len is None and "pos_embed" in state:
+            pos_len = int(state["pos_embed"].shape[1])
+            cls_tokens = 1 if config.global_cls else 0
+            inferred = max(0, pos_len - cls_tokens)
+            if inferred > 0:
+                config.max_seq_len = inferred
+        remapped_state = cls._remap_state_dict(state)
+        model = cls(config, *model_args, **kwargs)
+        model.load_state_dict(remapped_state, strict=False)
+        return model
+
+    def save_pretrained(self, save_directory: str | Path, **kwargs: Any) -> None:
+        directory = Path(save_directory)
+        directory.mkdir(parents=True, exist_ok=True)
+        config_path = directory / "config.json"
+        with config_path.open("w") as handle:
+            json.dump(self.config.to_dict(), handle, indent=2)
+        state_path = directory / "pytorch_model.bin"
+        torch.save(self.state_dict(), state_path)
+
+    @staticmethod
+    def _config_from_checkpoint(data: Any) -> Optional[LWMConfig]:
+        if not isinstance(data, dict):
+            return None
+        model_cfg = data.get("model", data)
+        if not isinstance(model_cfg, dict):
+            return None
+        allowed = {field.name for field in fields(LWMConfig)}
+        kwargs: Dict[str, Any] = {}
+        for key, value in model_cfg.items():
+            if key not in allowed:
+                continue
+            if key == "patch_size" and isinstance(value, (list, tuple)):
+                value = tuple(int(v) for v in value)
+            if key == "temporal_offsets" and isinstance(value, (list, tuple)):
+                value = tuple(int(v) for v in value)
+            kwargs[key] = value
+        if not kwargs:
+            return None
+        return LWMConfig(**kwargs)
+
+    @staticmethod
+    def _remap_state_dict(state: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        remapped: Dict[str, Tensor] = {}
+        for key, value in state.items():
+            new_key = key
+            if key.startswith("embed."):
+                new_key = key.replace("embed", "patch_embed", 1)
+            elif key.startswith("blocks."):
+                new_key = key.replace("blocks", "encoder.layers", 1)
+            elif key.startswith("norm."):
+                new_key = key.replace("norm", "encoder.norm", 1)
+            remapped[new_key] = value
+        return remapped
+
+
+def compute_nmse(pred: Tensor, target: Tensor, mask: Tensor) -> float:
+    """
+    Compute NMSE per sample, then average across batch (matching original implementation).
+    For each sample: nmse_b = sum((pred-target)^2 [mask]) / sum(target^2 [mask])
+    """
+    B = pred.size(0)
+    nmse_vals = []
+    for b in range(B):
+        m = mask[b]
+        if m.sum() == 0:
+            continue
+        se = (pred[b][m] - target[b][m]).pow(2).sum()
+        sp = target[b][m].pow(2).sum().clamp_min(1e-12)
+        nmse_vals.append((se / sp).item())
+    if not nmse_vals:
+        return float('nan')
+    return sum(nmse_vals) / len(nmse_vals)
+
+
+def masked_nmse_loss(pred: Tensor, target: Tensor, mask: Tensor) -> Tensor:
+    diff = (pred - target).abs() ** 2
+    power = target.abs() ** 2
+    mask_f = mask.float()
+    diff_sum = (diff.sum(-1) * mask_f).sum(-1)
+    power_sum = (power.sum(-1) * mask_f).sum(-1).clamp_min(1e-12)
+    nmse = diff_sum / power_sum
+    valid = mask.sum(-1) > 0
+    return nmse[valid].mean() if valid.any() else nmse.mean()
+
+
+def masked_mse_loss(pred: Tensor, target: Tensor, mask: Tensor) -> Tensor:
+    diff = (pred - target).abs() ** 2
+    mask_f = mask.float()
+    num = (diff.sum(-1) * mask_f).sum()
+    denom = mask_f.sum().clamp_min(1.0)
+    return num / denom
+
+
+__all__ = [
+    "ComplexPatchTokenizer",
+    "LWMConfig",
+    "LWMModel",
+    "LWMBackbone",
+    "compute_nmse",
+    "masked_nmse_loss",
+    "masked_mse_loss",
+]

LWMTemporal/tasks/channel_prediction.py

@@ -0,0 +1,641 @@
+from __future__ import annotations
+
+import dataclasses
+import logging
+import math
+from pathlib import Path
+from typing import Dict, List, Optional, Sequence, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.cuda.amp import GradScaler, autocast
+from torch.utils.data import DataLoader, Subset
+
+from ..data import AngleDelayDatasetConfig, AngleDelaySequenceDataset
+from ..models import LWMBackbone, LWMConfig
+from ..models.lwm import masked_mse_loss, masked_nmse_loss, compute_nmse
+
+try:
+    import wandb  # type: ignore
+except ImportError:  # pragma: no cover
+    wandb = None  # type: ignore
+
+
+@dataclasses.dataclass
+class DatasetArgs:
+    data_path: Path
+    keep_percentage: float = 0.25
+    normalize: str = "global_rms"
+    cache_dir: Path = Path("cache")
+    use_cache: bool = True
+    overwrite_cache: bool = False
+    snr_db: Optional[float] = None
+    noise_seed: Optional[int] = None
+    max_time_steps: Optional[int] = None
+    train_limit: int = 500
+    val_limit: int = 1000
+    seed: int = 42
+
+
+@dataclasses.dataclass
+class ModelArgs:
+    patch_size: Tuple[int, int] = (1, 1)
+    phase_mode: str = "real_imag"
+    embed_dim: int = 32
+    depth: int = 12
+    num_heads: int = 8
+    mlp_ratio: float = 4.0
+    same_frame_window: int = 2
+    temporal_offsets: Sequence[int] = dataclasses.field(default_factory=lambda: (-1, -2, -3, -4, -5, -6, -7))
+    temporal_spatial_window: int = 2
+    temporal_drift_h: int = 1
+    temporal_drift_w: int = 1
+    routing_topk_enable: bool = True
+    routing_topk_fraction: float = 0.2
+    routing_topk_min: int = 8
+    routing_topk_max: int = 32
+    topk_per_head: bool = True
+    posenc: str = "learned"
+    rope_base: float = 10000.0
+    global_cls: bool = False
+    pretrained: Optional[Path] = None
+    finetune_last_n: int = 0
+    train_head_only: bool = False
+
+
+@dataclasses.dataclass
+class TrainingArgs:
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    epochs: int = 3
+    batch_size: int = 16
+    lr: float = 1e-4
+    weight_decay: float = 1e-4
+    warmup_ratio: float = 0.1
+    loss: str = "nmse"
+    use_dataparallel: bool = False
+    grad_clip: float = 1.0
+    log_interval: int = 10
+    save_dir: Path = Path("models")
+    save_prefix: str = "channel_prediction"
+    inference_only: bool = False
+    inference_split: str = "val"
+    verbose_inference: bool = False
+    log_dir: Path = Path("logs")
+    use_wandb: bool = False
+    wandb_project: Optional[str] = None
+    wandb_entity: Optional[str] = None
+    wandb_run_name: Optional[str] = None
+
+
+@dataclasses.dataclass
+class PredictionArgs:
+    Tpast: int = 10
+    horizon: int = 1
+    num_visual_samples: int = 4
+    viz_dir: Path = Path("figs/predictions")
+
+
+@dataclasses.dataclass
+class ChannelPredictionArgs:
+    dataset: DatasetArgs
+    model: ModelArgs
+    training: TrainingArgs
+    prediction: PredictionArgs
+
+
+class ChannelPredictionDataModule:
+    def __init__(self, args: DatasetArgs, patch_size: Tuple[int, int], phase_mode: str) -> None:
+        cfg = AngleDelayDatasetConfig(
+            raw_path=args.data_path,
+            keep_percentage=args.keep_percentage,
+            normalize=args.normalize,
+            cache_dir=args.cache_dir,
+            use_cache=args.use_cache,
+            overwrite_cache=args.overwrite_cache,
+            snr_db=args.snr_db,
+            noise_seed=args.noise_seed,
+            max_time_steps=args.max_time_steps,
+            patch_size=patch_size,
+            phase_mode=phase_mode,
+        )
+        self.dataset = AngleDelaySequenceDataset(cfg)
+        generator = torch.Generator().manual_seed(args.seed)
+        indices = torch.randperm(len(self.dataset), generator=generator).tolist()
+        train_len = min(args.train_limit, len(indices))
+        val_len = min(args.val_limit, max(0, len(indices) - train_len))
+        self.train_indices = indices[:train_len]
+        self.val_indices = indices[train_len:train_len + val_len]
+        self.patch_size = patch_size
+        self.phase_mode = phase_mode
+
+    def train_loader(self, batch_size: int, drop_last: bool = True) -> DataLoader:
+        subset = Subset(self.dataset, self.train_indices)
+        return DataLoader(subset, batch_size=batch_size, shuffle=True, drop_last=drop_last)
+
+    def val_loader(self, batch_size: int, drop_last: bool = False) -> Optional[DataLoader]:
+        if not self.val_indices:
+            return None
+        subset = Subset(self.dataset, self.val_indices)
+        return DataLoader(subset, batch_size=batch_size, shuffle=False, drop_last=drop_last)
+
+
+class AutoregressiveEngine:
+    def __init__(self, patch_size: Tuple[int, int], phase_mode: str) -> None:
+        self.patch_size = patch_size
+        self.phase_mode = phase_mode
+
+    def detokenize(self, tokens: Tensor, T: int, H: int, W: int) -> Tensor:
+        B = tokens.size(0)
+        ph, pw = self.patch_size
+        patches = tokens.view(B, T, H, W, ph * pw * 2)
+        patches = patches.view(B, T, H, W, ph, pw, 2)
+        patches = patches.permute(0, 1, 2, 4, 3, 5, 6).contiguous()
+        recon = patches.view(B, T, H * ph, W * pw, 2)
+        if self.phase_mode == "real_imag":
+            real = recon[..., 0]
+            imag = recon[..., 1]
+            return torch.complex(real, imag)
+        magnitude = recon[..., 0]
+        phase = recon[..., 1]
+        real = magnitude * torch.cos(phase)
+        imag = magnitude * torch.sin(phase)
+        return torch.complex(real, imag)
+
+    def autoregressive_rollout(
+        self,
+        model: LWMBackbone,
+        tokens: Tensor,
+        Tpast: int,
+        horizon: int,
+        H: int,
+        W: int,
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        B, S_full, D = tokens.shape
+        S_per_time = H * W
+        if S_full % S_per_time != 0:
+            raise ValueError("Token sequence length incompatible with H and W")
+        T_total = S_full // S_per_time
+        S_per_time = H * W
+        window_tokens = Tpast + 1
+        if T_total < Tpast + horizon:
+            raise ValueError("sequence shorter than Tpast + horizon")
+        mask_window = torch.zeros((window_tokens, H, W), dtype=torch.bool, device=tokens.device)
+        mask_window[Tpast, :, :] = True
+        mask_window = mask_window.view(window_tokens * S_per_time)
+        mask_future = torch.zeros((T_total, H, W), dtype=torch.bool, device=tokens.device)
+        mask_future[Tpast:Tpast + horizon, :, :] = True
+        mask_flat = mask_future.view(1, T_total * S_per_time).expand(B, -1)
+
+        source_tokens = tokens.clone()
+        pred_tokens = torch.zeros_like(tokens)
+
+        for step in range(horizon):
+            start_time = step
+            end_time = step + window_tokens
+            abs_start = start_time * S_per_time
+            abs_end = end_time * S_per_time
+            window_slice = source_tokens[:, abs_start:abs_end, :].clone()  # Clone to avoid in-place modification
+            mask_slice = mask_window.unsqueeze(0).expand(B, -1)
+            # Zero masked tokens before model forward (matching original implementation)
+            window_slice = window_slice.masked_fill(mask_slice.unsqueeze(-1), 0.0)
+            outputs = model.forward_tokens(window_slice, mask_slice, window_tokens, H, W, return_cls=False)
+            predicted_window = outputs["reconstruction"]
+            # Extract predictions for the last time position in the window using slicing
+            win_last_start = Tpast * S_per_time
+            win_last_end = (Tpast + 1) * S_per_time
+            step_pred_last = predicted_window[:, win_last_start:win_last_end, :]
+            # Write back into absolute position
+            target_range_start = (Tpast + step) * S_per_time
+            target_range_end = target_range_start + S_per_time
+            source_tokens[:, target_range_start:target_range_end, :] = step_pred_last
+            pred_tokens[:, target_range_start:target_range_end, :] = step_pred_last
+
+        target_tokens = tokens
+        return pred_tokens, target_tokens, mask_flat
+
+
+class PredictionVisualizer:
+    def __init__(self, engine: AutoregressiveEngine, save_dir: Path, num_samples: int) -> None:
+        self.engine = engine
+        self.save_dir = save_dir
+        self.num_samples = num_samples
+        self.save_dir.mkdir(parents=True, exist_ok=True)
+
+    def save(self, model: LWMBackbone, tokens: Tensor, H: int, W: int, args: PredictionArgs) -> None:
+        model.eval()
+        with torch.no_grad():
+            preds, tgt, mask = self.engine.autoregressive_rollout(
+                model,
+                tokens,
+                args.Tpast,
+                args.horizon,
+                H,
+                W,
+            )
+        tokens_per_time = H * W
+        T_total = tokens.size(1) // tokens_per_time
+        B = tokens.size(0)
+        for idx in range(min(B, self.num_samples)):
+            pred_seq = preds[idx].view(T_total, tokens_per_time, -1)
+            tgt_seq = tgt[idx].view(T_total, tokens_per_time, -1)
+            pred_complex = self.engine.detokenize(pred_seq.unsqueeze(0), T_total, H, W)[0]
+            tgt_complex = self.engine.detokenize(tgt_seq.unsqueeze(0), T_total, H, W)[0]
+            self._plot_sample(pred_complex, tgt_complex, args, sample_idx=idx)
+
+    def _plot_sample(self, pred: Tensor, tgt: Tensor, args: PredictionArgs, sample_idx: int) -> None:
+        import matplotlib.pyplot as plt
+
+        fig, axes = plt.subplots(args.horizon, 2, figsize=(8, 3 * args.horizon), squeeze=False)
+        for step in range(args.horizon):
+            t_idx = args.Tpast + step
+            gt_mag = tgt[t_idx].abs().cpu().numpy()
+            pred_mag = pred[t_idx].abs().cpu().numpy()
+            ax_gt, ax_pred = axes[step]
+            im0 = ax_gt.imshow(gt_mag, cmap="viridis", aspect="auto")
+            im1 = ax_pred.imshow(pred_mag, cmap="viridis", aspect="auto")
+            ax_gt.set_title(f"GT t={t_idx}")
+            ax_pred.set_title(f"Pred t={t_idx}")
+            for ax in (ax_gt, ax_pred):
+                ax.set_xticks([])
+                ax.set_yticks([])
+            fig.colorbar(im0, ax=ax_gt, fraction=0.046, pad=0.04)
+            fig.colorbar(im1, ax=ax_pred, fraction=0.046, pad=0.04)
+        fig.tight_layout()
+        out_path = self.save_dir / f"sample_{sample_idx}.png"
+        fig.savefig(out_path)
+        plt.close(fig)
+
+
+class ChannelPredictionTrainer:
+    def __init__(self, args: ChannelPredictionArgs, *, logger: Optional[logging.Logger] = None) -> None:
+        self.args = args
+        torch.manual_seed(args.dataset.seed)
+        np.random.seed(args.dataset.seed)
+        self.device = torch.device(args.training.device)
+        self.engine = AutoregressiveEngine(args.model.patch_size, args.model.phase_mode)
+        self.data = ChannelPredictionDataModule(args.dataset, args.model.patch_size, args.model.phase_mode)
+        self.model = self._build_model().to(self.device)
+        self.model.eval()  # Set to eval mode immediately after loading
+        if args.training.use_dataparallel and torch.cuda.device_count() > 1:
+            self.model = nn.DataParallel(self.model)
+            if hasattr(self.model, 'module'):
+                self.model.module.eval()
+        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=args.training.lr, weight_decay=args.training.weight_decay)
+        self.scheduler = self._build_scheduler()
+        self.scaler = GradScaler()
+        self.viz = PredictionVisualizer(self.engine, args.prediction.viz_dir, args.prediction.num_visual_samples)
+        self.logger = logger or logging.getLogger(__name__)
+        self.global_step = 0
+        self._wandb_run = self._maybe_init_wandb()
+
+    def _wandb_enabled(self) -> bool:
+        return self._wandb_run is not None
+
+    def _maybe_init_wandb(self) -> Optional["wandb.sdk.wandb_run.Run"]:
+        training = self.args.training
+        if not training.use_wandb:
+            return None
+        if wandb is None:
+            self.logger.warning("Weights & Biases not installed; disabling wandb logging.")
+            return None
+        config = {
+            "dataset": dataclasses.asdict(self.args.dataset),
+            "model": dataclasses.asdict(self.args.model),
+            "training": dataclasses.asdict(self.args.training),
+            "prediction": dataclasses.asdict(self.args.prediction),
+        }
+        run = wandb.init(
+            project=training.wandb_project,
+            entity=training.wandb_entity,
+            name=training.wandb_run_name,
+            config=config,
+        )
+        wandb.watch(self.model, log="all", log_freq=self.args.training.log_interval)
+        self.logger.info("Initialized Weights & Biases run: %s", run.name)
+        return run
+
+    def _wandb_log(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
+        if not self._wandb_enabled():
+            return
+        wandb.log(metrics, step=step)
+
+    def _finish_wandb(self) -> None:
+        if self._wandb_enabled():
+            wandb.finish()
+
+    def _build_model(self) -> LWMBackbone:
+        # Calculate max_seq_len based on window size (matching original implementation)
+        # This is critical for channel prediction with autoregressive rollout
+        sample_batch = next(iter(self.data.val_loader(1) or self.data.train_loader(1)))
+        _, _, H, W = self._prepare_batch(sample_batch)
+        max_seq_len = (self.args.prediction.Tpast + 1) * H * W
+        
+        cfg = LWMConfig(
+            patch_size=self.args.model.patch_size,
+            phase_mode=self.args.model.phase_mode,
+            embed_dim=self.args.model.embed_dim,
+            depth=self.args.model.depth,
+            num_heads=self.args.model.num_heads,
+            mlp_ratio=self.args.model.mlp_ratio,
+            same_frame_window=self.args.model.same_frame_window,
+            temporal_offsets=self.args.model.temporal_offsets,
+            temporal_spatial_window=self.args.model.temporal_spatial_window,
+            temporal_drift_h=self.args.model.temporal_drift_h,
+            temporal_drift_w=self.args.model.temporal_drift_w,
+            routing_topk_enable=self.args.model.routing_topk_enable,
+            routing_topk_fraction=self.args.model.routing_topk_fraction,
+            routing_topk_min=self.args.model.routing_topk_min,
+            routing_topk_max=self.args.model.routing_topk_max,
+            topk_per_head=self.args.model.topk_per_head,
+            posenc=self.args.model.posenc,
+            rope_base=self.args.model.rope_base,
+            global_cls=self.args.model.global_cls,
+            max_seq_len=max_seq_len,
+        )
+        model = LWMBackbone(cfg)
+        if self.args.model.pretrained is not None and self.args.model.pretrained.exists():
+            model = LWMBackbone.from_pretrained(self.args.model.pretrained, config=cfg)
+        if self.args.model.train_head_only:
+            for param in model.parameters():
+                param.requires_grad = False
+            for param in model.head.parameters():
+                param.requires_grad = True
+        elif self.args.model.finetune_last_n > 0:
+            model.freeze_backbone()
+            if hasattr(model, "encoder"):
+                layers = model.encoder.layers
+                for layer in layers[-self.args.model.finetune_last_n:]:
+                    for param in layer.parameters():
+                        param.requires_grad = True
+            for param in model.head.parameters():
+                param.requires_grad = True
+        return model
+
+    def _build_scheduler(self) -> torch.optim.lr_scheduler.LambdaLR:
+        train_loader = self.data.train_loader(self.args.training.batch_size)
+        steps_per_epoch = max(1, len(train_loader))
+        total_steps = steps_per_epoch * max(1, self.args.training.epochs)
+        warmup_steps = int(self.args.training.warmup_ratio * total_steps)
+
+        def schedule(step: int) -> float:
+            if step < warmup_steps:
+                return float(step) / max(1, warmup_steps)
+            progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
+            return 0.5 * (1.0 + math.cos(math.pi * progress))
+
+        return torch.optim.lr_scheduler.LambdaLR(self.optimizer, schedule)
+
+    def _prepare_batch(self, batch: Dict[str, Tensor]) -> Tuple[Tensor, Tensor, int, int]:
+        tokens = batch["tokens"].to(self.device)
+        base_mask = batch["base_mask"].to(self.device)
+        shapes = batch["shape"]
+        if not isinstance(shapes, torch.Tensor):
+            shapes = torch.tensor(shapes)
+        if shapes.dim() == 1:
+            shapes = shapes.unsqueeze(0)
+        ref_shape = shapes[0]
+        if not torch.all(shapes.eq(ref_shape)):
+            raise ValueError("Mixed sequence shapes within the same batch are not supported")
+        T = int(ref_shape[0].item())
+        H = int(ref_shape[1].item())
+        W = int(ref_shape[2].item())
+        T_needed = self.args.prediction.Tpast + self.args.prediction.horizon
+        if T < T_needed:
+            raise ValueError("Sequence shorter than required Tpast+horizon frames")
+        S_per_time = H * W
+        tokens = tokens[:, : T_needed * S_per_time, :]
+        mask = base_mask[:, : T_needed * S_per_time]
+        return tokens, mask, H, W
+
+    def _compute_loss(self, pred: Tensor, tgt: Tensor, mask: Tensor) -> Tensor:
+        if self.args.training.loss == "mse":
+            return masked_mse_loss(pred, tgt, mask)
+        return masked_nmse_loss(pred, tgt, mask)
+
+    def train(self) -> None:
+        if self.args.training.inference_only:
+            # self.logger.info(
+            #     "Running inference-only evaluation on split '%s'", self.args.training.inference_split
+            # )
+            self.evaluate(split=self.args.training.inference_split)
+            self._finish_wandb()
+            return
+        train_loader = self.data.train_loader(self.args.training.batch_size)
+        val_loader = self.data.val_loader(self.args.training.batch_size)
+        for epoch in range(1, self.args.training.epochs + 1):
+            self.model.train()
+            running_loss = 0.0
+            running_nmse: List[float] = []
+            loader_len = len(train_loader)
+            for step, batch in enumerate(train_loader, start=1):
+                tokens, _, H, W = self._prepare_batch(batch)
+                with autocast():
+                    preds, target, mask = self.engine.autoregressive_rollout(
+                        self.model, tokens, self.args.prediction.Tpast, self.args.prediction.horizon, H, W
+                    )
+                    loss = self._compute_loss(preds, target, mask)
+                self.scaler.scale(loss).backward()
+                self.scaler.unscale_(self.optimizer)
+                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.training.grad_clip)
+                self.scaler.step(self.optimizer)
+                self.scaler.update()
+                self.optimizer.zero_grad()
+                self.scheduler.step()
+                running_loss += loss.item()
+                running_nmse.append(compute_nmse(preds, target, mask))
+                self.global_step += 1
+                if step % self.args.training.log_interval == 0:
+                    avg_loss = running_loss / step
+                    avg_nmse = float(np.mean(running_nmse)) if running_nmse else float("nan")
+                    lr_cur = self.optimizer.param_groups[0]["lr"]
+                    self.logger.info(
+                        "Train: [%d/%d][%d/%d] loss=%0.6f nmse=%0.6f lr=%0.2e",
+                        epoch,
+                        self.args.training.epochs,
+                        step,
+                        loader_len,
+                        avg_loss,
+                        avg_nmse,
+                        lr_cur,
+                    )
+                    self._wandb_log(
+                        {
+                            "train/loss": avg_loss,
+                            "train/nmse": avg_nmse,
+                            "train/lr": lr_cur,
+                        },
+                        step=self.global_step,
+                    )
+            avg_train_loss = running_loss / max(1, len(train_loader))
+            avg_train_nmse = float(np.mean(running_nmse)) if running_nmse else float("nan")
+            self.logger.info(
+                "Train Epoch %d/%d Summary: loss=%0.6f nmse=%0.6f",
+                epoch,
+                self.args.training.epochs,
+                avg_train_loss,
+                avg_train_nmse,
+            )
+            self._wandb_log(
+                {
+                    "train/epoch_loss": avg_train_loss,
+                    "train/epoch_nmse": avg_train_nmse,
+                },
+                step=self.global_step,
+            )
+            if val_loader is not None:
+                self.evaluate(loader=val_loader, split="val", epoch=epoch)
+            first_batch = next(iter(train_loader))
+            tokens_vis, _, H_vis, W_vis = self._prepare_batch(first_batch)
+            self.viz.save(self.model, tokens_vis, H_vis, W_vis, self.args.prediction)
+        self._finish_wandb()
+
+    def evaluate(
+        self,
+        loader: Optional[DataLoader] = None,
+        split: str = "val",
+        epoch: Optional[int] = None,
+    ) -> None:
+        if loader is None:
+            if split == "train":
+                subset = Subset(self.data.dataset, self.data.train_indices)
+                loader = DataLoader(subset, batch_size=self.args.training.batch_size, shuffle=False, drop_last=False)
+            elif split == "val":
+                subset = Subset(self.data.dataset, self.data.val_indices)
+                loader = DataLoader(subset, batch_size=self.args.training.batch_size, shuffle=False, drop_last=False)
+            else:
+                loader = DataLoader(self.data.dataset, batch_size=self.args.training.batch_size, shuffle=False)
+        if loader is None:
+            self.logger.warning("No %s loader available", split)
+            return
+        self.model.eval()
+        losses: List[float] = []
+        nmses: List[float] = []
+        per_step_nmses: List[List[float]] = []  # List of lists: [batch][step]
+        with torch.no_grad():
+            total_steps = len(loader)
+            for step, batch in enumerate(loader, start=1):
+                tokens, _, H, W = self._prepare_batch(batch)
+                preds, target, mask = self.engine.autoregressive_rollout(
+                    self.model, tokens, self.args.prediction.Tpast, self.args.prediction.horizon, H, W
+                )
+                loss = self._compute_loss(preds, target, mask)
+                batch_loss = loss.item()
+                batch_nmse = compute_nmse(preds, target, mask)
+                losses.append(batch_loss)
+                nmses.append(batch_nmse)
+                
+                # Compute per-step NMSE for this batch
+                S_per_time = H * W
+                Tpast = self.args.prediction.Tpast
+                horizon = self.args.prediction.horizon
+                step_nmses = []
+                for h in range(horizon):
+                    t_idx = Tpast + h
+                    step_start = t_idx * S_per_time
+                    step_end = step_start + S_per_time
+                    step_mask = mask[:, step_start:step_end]
+                    if step_mask.sum() > 0:
+                        step_pred = preds[:, step_start:step_end, :]
+                        step_target = target[:, step_start:step_end, :]
+                        step_nmse = compute_nmse(step_pred, step_target, step_mask)
+                        step_nmses.append(step_nmse)
+                    else:
+                        step_nmses.append(float('nan'))
+                per_step_nmses.append(step_nmses)
+                
+                # Report per-step NMSE for this batch (matching original package format)
+                per_step_strs = []
+                for h, step_nmse in enumerate(step_nmses):
+                    if not math.isnan(step_nmse):
+                        t = Tpast + h + 1  # t=11, 12, ... (1-indexed)
+                        nmse_db = 10.0 * math.log10(max(step_nmse, 1e-12))
+                        per_step_strs.append(f"t={t}: {nmse_db:.3f} dB")
+                if per_step_strs:
+                    self.logger.info(
+                        "[%s] per-step NMSE dB: %s",
+                        split,
+                        ", ".join(per_step_strs),
+                    )
+                
+                if self.args.training.verbose_inference:
+                    tag = split.upper()
+                    nmse_db = 10.0 * math.log10(max(batch_nmse, 1e-12))
+                    self.logger.info(
+                        "%s: [%d/%d] loss=%0.6f nmse=%0.6f (%0.2f dB)",
+                        tag,
+                        step,
+                        total_steps,
+                        batch_loss,
+                        batch_nmse,
+                        nmse_db,
+                    )
+        
+        avg_loss = float(np.mean(losses)) if losses else float("nan")
+        avg_nmse = float(np.mean(nmses)) if nmses else float("nan")
+        tag = f"[{split}]" if epoch is None else f"Epoch {epoch} [{split}]"
+        avg_nmse_db = 10.0 * math.log10(max(avg_nmse, 1e-12))
+        self.logger.info(
+            "Inference [%s] NMSE=%e (%0.3f dB) over %d batches",
+            split,
+            avg_nmse,
+            avg_nmse_db,
+            len(losses),
+        )
+        
+        # Compute per-step average in dB scale (matching original implementation)
+        if per_step_nmses:
+            horizon = len(per_step_nmses[0])
+            per_step_avg_db = []
+            Tpast = self.args.prediction.Tpast
+            for h in range(horizon):
+                # Average dB values (not linear values!)
+                step_dbs = []
+                for batch_nmses in per_step_nmses:
+                    if not math.isnan(batch_nmses[h]):
+                        step_db = 10.0 * math.log10(max(batch_nmses[h], 1e-12))
+                        step_dbs.append(step_db)
+                if step_dbs:
+                    avg_db = float(np.mean(step_dbs))
+                    per_step_avg_db.append(f"t={Tpast + h + 1}: {avg_db:.3f} dB")
+            if per_step_avg_db:
+                self.logger.info(
+                    "Inference [%s] per-step average NMSE dB: %s",
+                    split,
+                    ", ".join(per_step_avg_db),
+                )
+        
+        metrics = {
+            f"{split}/loss": avg_loss,
+            f"{split}/nmse": avg_nmse,
+            f"{split}/nmse_db": avg_nmse_db,
+        }
+        self._wandb_log(metrics, step=self.global_step)
+
+    def _save_checkpoint(self, epoch: int, metric: float) -> None:
+        self.args.training.save_dir.mkdir(parents=True, exist_ok=True)
+        filename = f"{self.args.training.save_prefix}_epoch{epoch:02d}.pth"
+        path = self.args.training.save_dir / filename
+        state = {
+            "epoch": epoch,
+            "metric": metric,
+            "model_state_dict": self.model.state_dict(),
+            "optimizer_state_dict": self.optimizer.state_dict(),
+            "scheduler_state_dict": self.scheduler.state_dict(),
+            "config": dataclasses.asdict(self.args),
+        }
+        torch.save(state, path)
+        print(f"Saved checkpoint to {path}")
+
+
+__all__ = [
+    "DatasetArgs",
+    "ModelArgs",
+    "TrainingArgs",
+    "PredictionArgs",
+    "ChannelPredictionArgs",
+    "ChannelPredictionDataModule",
+    "AutoregressiveEngine",
+    "PredictionVisualizer",
+    "ChannelPredictionTrainer",
+]

LWMTemporal/tasks/pretraining.py

@@ -0,0 +1,684 @@
+from __future__ import annotations
+
+import argparse
+import dataclasses
+import logging
+import math
+import pickle
+import random
+from pathlib import Path
+from typing import Dict, List, Optional, Sequence, Tuple
+
+import numpy as np
+import torch
+from torch import Tensor
+from torch.cuda.amp import GradScaler, autocast
+from torch.utils.data import DataLoader, Dataset
+
+from ..data.angle_delay import AngleDelayConfig, AngleDelayProcessor
+from ..models import LWMBackbone, LWMConfig
+from ..models.lwm import ComplexPatchTokenizer, masked_nmse_loss
+
+try:
+    import wandb  # type: ignore
+except ImportError:  # pragma: no cover
+    wandb = None  # type: ignore
+
+@dataclasses.dataclass
+class DataArgs:
+    data_dir: Path
+    keep_percentage: float = 0.25
+    normalize: str = "global_rms"
+    max_time_steps: Optional[int] = None
+
+
+@dataclasses.dataclass
+class MaskArgs:
+    mask_ratio: float = 0.75
+    mask_mode: str = "auto"
+    random_fraction: float = 0.2
+
+
+@dataclasses.dataclass
+class CurriculumArgs:
+    strategy: str = "mask"
+    warmup_epochs: int = 4
+    min_mask_ratio: float = 0.3
+    max_mask_ratio: float = 0.75
+
+
+@dataclasses.dataclass
+class AugmentationArgs:
+    phase_p: float = 0.0
+    amp_p: float = 0.0
+    amp_min: float = 0.7
+    amp_max: float = 1.3
+    awgn_p: float = 0.0
+    awgn_snr_min: float = 20.0
+    awgn_snr_max: float = 30.0
+
+
+@dataclasses.dataclass
+class LoggingArgs:
+    log_dir: Path = Path("logs")
+    use_wandb: bool = False
+    wandb_project: Optional[str] = None
+    wandb_entity: Optional[str] = None
+    wandb_run_name: Optional[str] = None
+
+
+@dataclasses.dataclass
+class OptimizationArgs:
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    epochs: int = 20
+    batch_size: int = 32
+    lr: float = 2e-4
+    weight_decay: float = 1e-4
+    warmup_ratio: float = 0.1
+    grad_clip: float = 1.0
+    log_interval: int = 1
+    save_dir: Path = Path("models")
+    save_prefix: str = "lwm_pretrain"
+    resume_from: Optional[Path] = None
+
+
+@dataclasses.dataclass
+class ModelArgs:
+    patch_size: Tuple[int, int] = (1, 1)
+    phase_mode: str = "real_imag"
+    embed_dim: int = 32
+    depth: int = 12
+    num_heads: int = 8
+    mlp_ratio: float = 4.0
+    same_frame_window: int = 2
+    temporal_offsets: Sequence[int] = dataclasses.field(default_factory=lambda: (-4, -3, -2, -1, 1, 2, 3))
+    temporal_spatial_window: int = 2
+    temporal_drift_h: int = 1
+    temporal_drift_w: int = 1
+    routing_topk_enable: bool = True
+    routing_topk_fraction: float = 0.2
+    routing_topk_min: int = 8
+    routing_topk_max: int = 32
+    topk_per_head: bool = True
+    posenc: str = "learned"
+    rope_base: float = 10000.0
+    global_cls: bool = False
+
+
+@dataclasses.dataclass
+class PretrainingArgs:
+    data: DataArgs
+    mask: MaskArgs
+    curriculum: CurriculumArgs
+    augment: AugmentationArgs
+    optim: OptimizationArgs
+    model: ModelArgs
+    logging: LoggingArgs
+
+
+class PretrainingDataset(Dataset):
+    def __init__(
+        self,
+        args: DataArgs,
+        tokenizer: ComplexPatchTokenizer,
+        augmenter: Augmenter,
+        masker: MaskGenerator,
+        patch_size: Tuple[int, int],
+    ) -> None:
+        self.args = args
+        self.tokenizer = tokenizer
+        self.augmenter = augmenter
+        self.masker = masker
+        self.patch_size = patch_size
+        self.samples = self._load_sequences()
+        if args.normalize != "none":
+            self.samples = [self._normalize(sample, args.normalize) for sample in self.samples]
+
+    def _load_sequences(self) -> List[Tensor]:
+        processor = AngleDelayProcessor(AngleDelayConfig(keep_percentage=self.args.keep_percentage))
+        samples: List[Tensor] = []
+        for path in sorted(self.args.data_dir.glob("*.p")):
+            with path.open("rb") as handle:
+                payload = pickle.load(handle)
+            if isinstance(payload, dict) and "channel" in payload:
+                tensor = torch.as_tensor(payload["channel"], dtype=torch.complex64)
+            else:
+                tensor = torch.as_tensor(payload, dtype=torch.complex64)
+            if tensor.ndim == 3:
+                tensor = tensor.unsqueeze(0)
+            for seq in tensor:
+                ad = processor.forward(seq)
+                truncated, _ = processor.truncate_delay_bins(ad)
+                if self.args.max_time_steps is not None and truncated.size(0) > self.args.max_time_steps:
+                    truncated = truncated[: self.args.max_time_steps]
+                samples.append(truncated)
+        return samples
+
+    def _normalize(self, tensor: Tensor, mode: str) -> Tensor:
+        if mode == "global_rms":
+            rms = torch.sqrt((tensor.real.float() ** 2 + tensor.imag.float() ** 2).mean()).clamp_min(1e-8)
+            return tensor / rms.to(tensor.dtype)
+        if mode == "per_sample_rms":
+            rms = torch.sqrt((tensor.real.float() ** 2 + tensor.imag.float() ** 2).mean()).clamp_min(1e-8)
+            return tensor / rms.to(tensor.dtype)
+        return tensor
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, index: int) -> Dict[str, Tensor]:
+        sample = self.samples[index]
+        if self.augmenter is not None:
+            sample = self.augmenter(sample)
+        tokens, _ = self.tokenizer(sample.unsqueeze(0), self.patch_size)
+        tokens = tokens.squeeze(0)
+        ph, pw = self.patch_size
+        T, N, M = sample.shape
+        H = N // ph
+        W = M // pw
+        mask = self.masker(T, H, W, device=tokens.device).view(-1)
+        shape = torch.tensor([T, H, W], dtype=torch.long)
+        return {
+            "tokens": tokens,
+            "mask": mask,
+            "shape": shape,
+        }
+
+
+class MaskGenerator:
+    def __init__(self, args: MaskArgs) -> None:
+        self.args = args
+
+    def __call__(self, T: int, H: int, W: int, device: torch.device) -> torch.BoolTensor:
+        if self.args.mask_mode == "random" or (self.args.mask_mode == "auto" and random.random() < self.args.random_fraction):
+            return self.random_mask(T, H, W, device)
+        if self.args.mask_mode in {"rect", "auto"} and random.random() < 0.33:
+            return self.rect_mask(T, H, W, device)
+        if self.args.mask_mode in {"tube", "auto"} and random.random() < 0.33:
+            return self.tube_mask(T, H, W, device)
+        return self.comb_mask(T, H, W, device)
+ 
+    def random_mask(self, T: int, H: int, W: int, device: torch.device) -> torch.BoolTensor:
+        total = T * H * W
+        num_mask = int(self.args.mask_ratio * total)
+        mask = torch.zeros(total, dtype=torch.bool, device=device)
+        idx = torch.randperm(total, device=device)[:num_mask]
+        mask[idx] = True
+        return mask.view(T, H, W)
+
+    def rect_mask(self, T: int, H: int, W: int, device: torch.device) -> torch.BoolTensor:
+        mask = torch.zeros((T, H, W), dtype=torch.bool, device=device)
+        blocks = max(1, int(self.args.mask_ratio * T))
+        for _ in range(blocks):
+            t = random.randrange(T)
+            h_size = random.randint(1, max(1, H // 2))
+            w_size = random.randint(1, max(1, W // 2))
+            h0 = random.randint(0, H - h_size)
+            w0 = random.randint(0, W - w_size)
+            mask[t, h0:h0 + h_size, w0:w0 + w_size] = True
+        return mask
+
+    def tube_mask(self, T: int, H: int, W: int, device: torch.device) -> torch.BoolTensor:
+        mask = torch.zeros((T, H, W), dtype=torch.bool, device=device)
+        start_t = random.randrange(T)
+        h = random.randrange(H)
+        w = random.randrange(W)
+        length = random.randint(max(1, T // 2), T)
+        for k in range(length):
+            t_idx = (start_t + k) % T
+            mask[t_idx, max(0, h - 1):min(H, h + 2), max(0, w - 1):min(W, w + 2)] = True
+            h = max(0, min(H - 1, h + random.randint(-1, 1)))
+            w = max(0, min(W - 1, w + random.randint(-1, 1)))
+        return mask
+
+    def comb_mask(self, T: int, H: int, W: int, device: torch.device) -> torch.BoolTensor:
+        mask = torch.zeros((T, H, W), dtype=torch.bool, device=device)
+        stride_t = random.choice([2, 3]) if T >= 2 else 1
+        stride_w = random.choice([3, 4, 6]) if W >= 3 else 1
+        offset_t = random.randrange(stride_t)
+        offset_w = random.randrange(stride_w)
+        for t in range(T):
+            for w in range(W):
+                visible = (t % stride_t == offset_t) and (w % stride_w == offset_w)
+                if not visible:
+                    mask[t, :, w] = True
+        return mask
+
+
+class Augmenter:
+    def __init__(self, args: AugmentationArgs) -> None:
+        self.args = args
+
+    def __call__(self, tensor: Tensor) -> Tensor:
+        x = tensor.clone()
+        if torch.rand(()) < self.args.phase_p:
+            theta = (torch.rand((), device=x.device) * 2 * math.pi) - math.pi
+            rotation = torch.cos(theta) + 1j * torch.sin(theta)
+            x = x * rotation
+        if torch.rand(()) < self.args.amp_p:
+            scale = self.args.amp_min + (self.args.amp_max - self.args.amp_min) * torch.rand((), device=x.device)
+            x = x * scale
+        if torch.rand(()) < self.args.awgn_p:
+            snr_db = torch.empty((), device=x.device).uniform_(self.args.awgn_snr_min, self.args.awgn_snr_max)
+            snr_lin = 10 ** (snr_db / 10.0)
+            power = (x.real.float().pow(2) + x.imag.float().pow(2)).mean().item()
+            if power > 0:
+                noise_var = power / snr_lin
+                std = math.sqrt(noise_var / 2.0)
+                noise_real = torch.randn_like(x.real.float()) * std
+                noise_imag = torch.randn_like(x.imag.float()) * std
+                noise = torch.complex(noise_real.to(x.dtype), noise_imag.to(x.dtype))
+                x = x + noise
+        return x
+
+
+class PretrainingTrainer:
+    def __init__(self, args: PretrainingArgs, *, logger: Optional[logging.Logger] = None) -> None:
+        self.args = args
+        self.logger = logger or logging.getLogger(__name__)
+        self.device = torch.device(args.optim.device)
+        self.tokenizer = ComplexPatchTokenizer(args.model.phase_mode)
+        self.masker = MaskGenerator(args.mask)
+        self.augmenter = Augmenter(args.augment)
+        self.dataset = PretrainingDataset(args.data, self.tokenizer, self.augmenter, self.masker, args.model.patch_size)
+        self.dataloader = DataLoader(self.dataset, batch_size=args.optim.batch_size, shuffle=True, drop_last=True)
+        self.model = self._build_model().to(self.device)
+        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=args.optim.lr, weight_decay=args.optim.weight_decay)
+        self.scheduler = self._build_scheduler()
+        self.scaler = GradScaler()
+        self.global_step = 0
+        self._wandb_run = self._maybe_init_wandb()
+        if self.args.optim.resume_from is not None:
+            self._load_checkpoint(self.args.optim.resume_from)
+
+    def _wandb_enabled(self) -> bool:
+        return self._wandb_run is not None
+
+    def _maybe_init_wandb(self) -> Optional["wandb.sdk.wandb_run.Run"]:
+        logging_args = self.args.logging
+        if not logging_args.use_wandb:
+            return None
+        if wandb is None:
+            self.logger.warning("Weights & Biases not installed; disabling wandb logging.")
+            return None
+        config = dataclasses.asdict(self.args)
+        run = wandb.init(
+            project=logging_args.wandb_project,
+            entity=logging_args.wandb_entity,
+            name=logging_args.wandb_run_name,
+            config=config,
+        )
+        wandb.watch(self.model, log="all", log_freq=self.args.optim.log_interval)
+        self.logger.info("Initialized Weights & Biases run: %s", run.name)
+        return run
+
+    def _wandb_log(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
+        if not self._wandb_enabled():
+            return
+        wandb.log(metrics, step=step)
+
+    def _finish_wandb(self) -> None:
+        if self._wandb_enabled():
+            wandb.finish()
+
+    def _build_model(self) -> LWMBackbone:
+        cfg = LWMConfig(
+            patch_size=self.args.model.patch_size,
+            phase_mode=self.args.model.phase_mode,
+            embed_dim=self.args.model.embed_dim,
+            depth=self.args.model.depth,
+            num_heads=self.args.model.num_heads,
+            mlp_ratio=self.args.model.mlp_ratio,
+            same_frame_window=self.args.model.same_frame_window,
+            temporal_offsets=self.args.model.temporal_offsets,
+            temporal_spatial_window=self.args.model.temporal_spatial_window,
+            temporal_drift_h=self.args.model.temporal_drift_h,
+            temporal_drift_w=self.args.model.temporal_drift_w,
+            routing_topk_enable=self.args.model.routing_topk_enable,
+            routing_topk_fraction=self.args.model.routing_topk_fraction,
+            routing_topk_min=self.args.model.routing_topk_min,
+            routing_topk_max=self.args.model.routing_topk_max,
+            topk_per_head=self.args.model.topk_per_head,
+            posenc=self.args.model.posenc,
+            rope_base=self.args.model.rope_base,
+            global_cls=self.args.model.global_cls,
+        )
+        return LWMBackbone(cfg)
+
+    def _build_scheduler(self) -> torch.optim.lr_scheduler.LambdaLR:
+        steps_per_epoch = max(1, len(self.dataloader))
+        total_steps = steps_per_epoch * max(1, self.args.optim.epochs)
+        warmup_steps = int(self.args.optim.warmup_ratio * total_steps)
+
+        def schedule(step: int) -> float:
+            if step < warmup_steps:
+                return float(step) / max(1, warmup_steps)
+            progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
+            return 0.5 * (1.0 + math.cos(math.pi * progress))
+
+        return torch.optim.lr_scheduler.LambdaLR(self.optimizer, schedule)
+
+    def _adjust_curriculum(self, epoch: int) -> None:
+        if self.args.curriculum.strategy == "mask" and epoch <= self.args.curriculum.warmup_epochs:
+            ratio = np.interp(
+                epoch,
+                [0, self.args.curriculum.warmup_epochs],
+                [self.args.curriculum.min_mask_ratio, self.args.curriculum.max_mask_ratio],
+            )
+            self.masker.args.mask_ratio = float(ratio)
+            self.logger.info(
+                "Curriculum update | epoch=%d/%d mask_ratio=%0.2f",
+                epoch,
+                self.args.optim.epochs,
+                self.masker.args.mask_ratio,
+            )
+            self._wandb_log(
+                {"train/curriculum_mask_ratio": self.masker.args.mask_ratio},
+                step=self.global_step,
+            )
+            self.logger.info(
+                "Curriculum update | epoch=%d/%d mask_ratio=%0.2f",
+                epoch,
+                self.args.optim.epochs,
+                self.masker.args.mask_ratio,
+            )
+
+    def train(self) -> None:
+        for epoch in range(1, self.args.optim.epochs + 1):
+            self._adjust_curriculum(epoch)
+            running_loss = 0.0
+            loader_len = len(self.dataloader)
+            for step, batch in enumerate(self.dataloader, start=1):
+                tokens = batch["tokens"].to(self.device)
+                mask_tokens = batch["mask"].to(self.device)
+                shapes = batch["shape"]
+                if not isinstance(shapes, torch.Tensor):
+                    shapes = torch.tensor(shapes)
+                if shapes.dim() == 1:
+                    shapes = shapes.unsqueeze(0)
+                ref_shape = shapes[0]
+                if not torch.all(shapes.eq(ref_shape)):
+                    raise ValueError("Mixed sequence shapes within the same batch are not supported")
+                T = int(ref_shape[0].item())
+                H = int(ref_shape[1].item())
+                W = int(ref_shape[2].item())
+                with autocast():
+                    outputs = self.model.forward_tokens(tokens, mask_tokens, T, H, W, return_cls=False)
+                    preds = outputs["reconstruction"]
+                    loss = masked_nmse_loss(preds, tokens, mask_tokens)
+                self.scaler.scale(loss).backward()
+                self.scaler.unscale_(self.optimizer)
+                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.optim.grad_clip)
+                self.scaler.step(self.optimizer)
+                self.scaler.update()
+                self.optimizer.zero_grad()
+                self.scheduler.step()
+                running_loss += loss.item()
+                self.global_step += 1
+                if step % self.args.optim.log_interval == 0:
+                    avg_loss = running_loss / step
+                    lr_cur = self.optimizer.param_groups[0]["lr"]
+                    self.logger.info(
+                        "Train: [%d/%d][%d/%d] loss=%0.6f mask=%0.2f lr=%0.2e",
+                        epoch,
+                        self.args.optim.epochs,
+                        step,
+                        loader_len,
+                        avg_loss,
+                        self.masker.args.mask_ratio,
+                        lr_cur,
+                    )
+                    self._wandb_log(
+                        {
+                            "train/loss": avg_loss,
+                            "train/mask_ratio": self.masker.args.mask_ratio,
+                            "train/lr": lr_cur,
+                        },
+                        step=self.global_step,
+                    )
+            avg_epoch_loss = running_loss / max(1, len(self.dataloader))
+            self.logger.info(
+                "Train Epoch %d/%d Summary: loss=%0.6f",
+                epoch,
+                self.args.optim.epochs,
+                avg_epoch_loss,
+            )
+            self._wandb_log(
+                {
+                    "train/epoch_loss": avg_epoch_loss,
+                },
+                step=self.global_step,
+            )
+            self._save_checkpoint(epoch, avg_epoch_loss)
+        self._finish_wandb()
+
+    def _save_checkpoint(self, epoch: int, metric: float) -> None:
+        self.args.optim.save_dir.mkdir(parents=True, exist_ok=True)
+        save_prefix = Path(self.args.optim.save_prefix)
+        suffix = save_prefix.suffix or ".pth"
+        stem = save_prefix.stem if save_prefix.suffix else save_prefix.name
+        filename = f"{stem}_epoch{epoch:03d}{suffix}"
+        path = self.args.optim.save_dir / filename
+        suffix = path.suffix.lower()
+        if suffix == ".bin":
+            torch.save(self.model.state_dict(), path)
+            self.logger.info("Saved weights-only checkpoint to %s", path)
+        else:
+            torch.save(
+                {
+                    "epoch": epoch,
+                    "metric": metric,
+                    "model_state_dict": self.model.state_dict(),
+                    "optimizer_state_dict": self.optimizer.state_dict(),
+                    "scheduler_state_dict": self.scheduler.state_dict(),
+                    "config": dataclasses.asdict(self.args),
+                },
+                path,
+            )
+            self.logger.info("Saved checkpoint to %s", path)
+        if self._wandb_enabled():
+            wandb.save(str(path))
+
+    def _load_checkpoint(self, checkpoint_path: Path) -> None:
+        if not checkpoint_path.exists():
+            raise FileNotFoundError(f"Checkpoint not found at {checkpoint_path}")
+        self.logger.info("Resuming from checkpoint %s", checkpoint_path)
+
+        payload = torch.load(checkpoint_path, map_location=self.device)
+        suffix = checkpoint_path.suffix.lower()
+
+        if suffix == ".bin":
+            model_state = payload
+            missing, unexpected = self.model.load_state_dict(model_state, strict=False)
+            if missing:
+                self.logger.warning("Missing keys when loading model: %s", missing)
+            if unexpected:
+                self.logger.warning("Unexpected keys when loading model: %s", unexpected)
+            self.logger.info("Loaded weights-only checkpoint.")
+            return
+
+        model_state = payload.get("model_state_dict")
+        if model_state is not None:
+            missing, unexpected = self.model.load_state_dict(model_state, strict=False)
+            if missing:
+                self.logger.warning("Missing keys when loading model: %s", missing)
+            if unexpected:
+                self.logger.warning("Unexpected keys when loading model: %s", unexpected)
+
+        opt_state = payload.get("optimizer_state_dict")
+        if opt_state is not None:
+            try:
+                self.optimizer.load_state_dict(opt_state)
+            except Exception as exc:
+                self.logger.warning("Failed to load optimizer state: %s", exc)
+
+        sched_state = payload.get("scheduler_state_dict")
+        if sched_state is not None:
+            try:
+                self.scheduler.load_state_dict(sched_state)
+            except Exception as exc:
+                self.logger.warning("Failed to load scheduler state: %s", exc)
+
+        epoch = payload.get("epoch", 0)
+        metric = payload.get("metric")
+        self.logger.info("Checkpoint contained epoch=%s metric=%s", epoch, metric)
+
+
+def build_pretraining_args(ns: argparse.Namespace) -> PretrainingArgs:
+    data_args = DataArgs(
+        data_dir=ns.data_dir,
+        keep_percentage=ns.keep_percentage,
+        normalize=ns.normalize,
+        max_time_steps=ns.max_time_steps,
+    )
+    mask_args = MaskArgs(mask_ratio=ns.mask_ratio, mask_mode=ns.mask_mode, random_fraction=ns.mask_random_fraction)
+    curriculum_args = CurriculumArgs(
+        strategy=ns.curriculum_strategy,
+        warmup_epochs=ns.curriculum_warmup_epochs,
+        min_mask_ratio=ns.curriculum_min_mask,
+        max_mask_ratio=ns.curriculum_max_mask,
+    )
+    augment_args = AugmentationArgs(
+        phase_p=ns.aug_phase_p,
+        amp_p=ns.aug_amp_p,
+        amp_min=ns.aug_amp_min,
+        amp_max=ns.aug_amp_max,
+        awgn_p=ns.aug_awgn_p,
+        awgn_snr_min=ns.aug_awgn_snr_min,
+        awgn_snr_max=ns.aug_awgn_snr_max,
+    )
+    optim_args = OptimizationArgs(
+        device=ns.device,
+        epochs=ns.epochs,
+        batch_size=ns.batch_size,
+        lr=ns.lr,
+        weight_decay=ns.weight_decay,
+        warmup_ratio=ns.warmup_ratio,
+        grad_clip=ns.grad_clip,
+        log_interval=ns.log_interval,
+        save_dir=ns.save_dir,
+        save_prefix=ns.save_prefix,
+        resume_from=ns.resume_from,
+    )
+    logging_args = LoggingArgs(
+        log_dir=ns.log_dir,
+        use_wandb=ns.use_wandb,
+        wandb_project=ns.wandb_project,
+        wandb_entity=ns.wandb_entity,
+        wandb_run_name=ns.wandb_run_name,
+    )
+    model_args = ModelArgs(
+        patch_size=tuple(ns.patch_size),
+        phase_mode=ns.phase_mode,
+        embed_dim=ns.embed_dim,
+        depth=ns.depth,
+        num_heads=ns.num_heads,
+        mlp_ratio=ns.mlp_ratio,
+        same_frame_window=ns.same_frame_window,
+        temporal_offsets=tuple(ns.temporal_offsets),
+        temporal_spatial_window=ns.temporal_spatial_window,
+        temporal_drift_h=ns.temporal_drift_h,
+        temporal_drift_w=ns.temporal_drift_w,
+        routing_topk_enable=ns.routing_topk_enable,
+        routing_topk_fraction=ns.routing_topk_fraction,
+        routing_topk_min=ns.routing_topk_min,
+        routing_topk_max=ns.routing_topk_max,
+        topk_per_head=ns.topk_per_head,
+        posenc=ns.posenc,
+        rope_base=ns.rope_base,
+    )
+    return PretrainingArgs(
+        data=data_args,
+        mask=mask_args,
+        curriculum=curriculum_args,
+        augment=augment_args,
+        optim=optim_args,
+        model=model_args,
+        logging=logging_args,
+    )
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Pretrain LWM foundation model")
+    parser.add_argument("--data_dir", type=Path, required=True)
+    parser.add_argument("--keep_percentage", type=float, default=0.25)
+    parser.add_argument("--normalize", type=str, default="global_rms", choices=["global_rms", "per_sample_rms", "none"])
+    parser.add_argument("--max_time_steps", type=int, default=None)
+
+    parser.add_argument("--mask_ratio", type=float, default=0.60)
+    parser.add_argument("--mask_mode", type=str, default="auto", choices=["auto", "random", "rect", "tube", "comb"])
+    parser.add_argument("--mask_random_fraction", type=float, default=0.2)
+
+    parser.add_argument("--curriculum_strategy", type=str, default="mask", choices=["none", "mask"])
+    parser.add_argument("--curriculum_warmup_epochs", type=int, default=4)
+    parser.add_argument("--curriculum_min_mask", type=float, default=0.3)
+    parser.add_argument("--curriculum_max_mask", type=float, default=0.75)
+    parser.add_argument("--log_dir", type=Path, default=Path("logs"))
+    parser.add_argument("--use_wandb", action="store_true")
+    parser.add_argument("--wandb_project", type=str, default=None)
+    parser.add_argument("--wandb_entity", type=str, default=None)
+    parser.add_argument("--wandb_run_name", type=str, default=None)
+
+    parser.add_argument("--phase_mode", type=str, default="real_imag", choices=["real_imag", "mag_phase"])
+    parser.add_argument("--patch_size", type=int, nargs=2, default=(1, 1))
+    parser.add_argument("--embed_dim", type=int, default=32)
+    parser.add_argument("--depth", type=int, default=12)
+    parser.add_argument("--num_heads", type=int, default=8)
+    parser.add_argument("--mlp_ratio", type=float, default=4.0)
+    parser.add_argument("--same_frame_window", type=int, default=2)
+    parser.add_argument("--temporal_offsets", type=int, nargs="*", default=[-4, -3, -2, -1, 1, 2, 3])
+    parser.add_argument("--temporal_spatial_window", type=int, default=2)
+    parser.add_argument("--temporal_drift_h", type=int, default=1)
+    parser.add_argument("--temporal_drift_w", type=int, default=1)
+    parser.add_argument("--routing_topk_enable", action="store_true", default=True)
+    parser.add_argument("--routing_topk_fraction", type=float, default=0.2)
+    parser.add_argument("--routing_topk_min", type=int, default=8)
+    parser.add_argument("--routing_topk_max", type=int, default=32)
+    parser.add_argument("--topk_per_head", action="store_true", default=True)
+    parser.add_argument("--posenc", type=str, default="learned", choices=["learned", "rope_sincos"])
+    parser.add_argument("--rope_base", type=float, default=10000.0)
+
+    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--epochs", type=int, default=20)
+    parser.add_argument("--batch_size", type=int, default=32)
+    parser.add_argument("--lr", type=float, default=2e-4)
+    parser.add_argument("--weight_decay", type=float, default=1e-4)
+    parser.add_argument("--warmup_ratio", type=float, default=0.1)
+    parser.add_argument("--grad_clip", type=float, default=1.0)
+    parser.add_argument("--log_interval", type=int, default=1)
+    parser.add_argument("--save_dir", type=Path, default=Path("models"))
+    parser.add_argument("--save_prefix", type=str, default="lwm_pretrain")
+    parser.add_argument("--resume_from", type=Path, default=None, help="Path to checkpoint to resume from")
+
+    parser.add_argument("--aug_phase_p", type=float, default=0.0)
+    parser.add_argument("--aug_amp_p", type=float, default=0.0)
+    parser.add_argument("--aug_amp_min", type=float, default=0.7)
+    parser.add_argument("--aug_amp_max", type=float, default=1.3)
+    parser.add_argument("--aug_awgn_p", type=float, default=0.0)
+    parser.add_argument("--aug_awgn_snr_min", type=float, default=20.0)
+    parser.add_argument("--aug_awgn_snr_max", type=float, default=30.0)
+
+    return parser
+
+
+def main(argv: Optional[Sequence[str]] = None) -> None:
+    ns = build_parser().parse_args(args=list(argv) if argv is not None else None)
+    args = build_pretraining_args(ns)
+    trainer = PretrainingTrainer(args)
+    trainer.train()
+
+
+__all__ = [
+    "DataArgs",
+    "MaskArgs",
+    "CurriculumArgs",
+    "AugmentationArgs",
+    "OptimizationArgs",
+    "ModelArgs",
+    "PretrainingArgs",
+    "PretrainingDataset",
+    "PretrainingTrainer",
+    "build_pretraining_args",
+    "build_parser",
+    "main",
+]

LWMTemporal/training/__init__.py

@@ -0,0 +1 @@
+"""Training utilities for LWM foundation models."""

LWMTemporal/utils/logging.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from pathlib import Path
+import logging
+
+LOG_FORMAT = "[%(asctime)s,%(msecs)03d %(levelname)s %(name)s line %(lineno)d %(process)d] %(message)s"
+DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
+
+
+def setup_logging(
+    name: str = "LWMTemporal",
+    log_dir: Path | None = None,
+    level: int = logging.INFO,
+) -> logging.Logger:
+    """Configure and return a logger using the original package's style."""
+    logger = logging.getLogger(name)
+    logger.setLevel(level)
+
+    # Avoid duplicating handlers when called multiple times
+    if logger.hasHandlers():
+        logger.handlers.clear()
+
+    formatter = logging.Formatter(LOG_FORMAT, DATE_FORMAT)
+
+    if log_dir is not None:
+        log_dir = Path(log_dir)
+        log_dir.mkdir(parents=True, exist_ok=True)
+        file_handler = logging.FileHandler(log_dir / f"{name}.log")
+        file_handler.setFormatter(formatter)
+        file_handler.setLevel(level)
+        logger.addHandler(file_handler)
+
+    stream_handler = logging.StreamHandler()
+    stream_handler.setFormatter(formatter)
+    stream_handler.setLevel(level)
+    logger.addHandler(stream_handler)
+
+    return logger
+

MANIFEST.in

@@ -0,0 +1,14 @@
+include README.md
+include LICENSE
+include requirements.txt
+include LWMTemporal/models/config.json
+recursive-exclude * __pycache__
+recursive-exclude * *.py[co]
+recursive-exclude * .DS_Store
+exclude cache
+exclude logs
+exclude figs
+exclude wandb
+exclude checkpoints
+exclude test.py
+

README.md

@@ -0,0 +1,353 @@
+# LWMTemporal
+
+Large Wireless Model (LWM) with sparse spatio-temporal attention for wireless channel prediction and forecasting.
+
+This package provides a transformer-based model for spatio-temporal wireless channel prediction with support for both pretraining and fine-tuning tasks. It follows Hugging Face conventions for model checkpoints and configurations.
+
+---
+
+## Installation
+
+### From PyPI (Recommended)
+
+```bash
+pip install lwm-temporal
+```
+
+### From Source
+
+```bash
+git clone https://github.com/yourusername/lwm-temporal.git
+cd lwm-temporal
+pip install -e .
+```
+
+### Optional Dependencies
+
+For Weights & Biases logging:
+```bash
+pip install lwm-temporal[wandb]
+```
+
+For development:
+```bash
+pip install lwm-temporal[dev]
+```
+
+---
+
+## Quick Start
+
+### Python API
+
+```python
+from pathlib import Path
+from LWMTemporal import LWMBackbone, LWMConfig
+
+# Load pretrained model
+model = LWMBackbone.from_pretrained("checkpoints/m18_cp.pth")
+model.eval()
+
+# Use for inference - see examples/ for complete scripts
+```
+
+### Command Line Interface
+
+```bash
+# Run channel prediction inference
+python -m LWMTemporal.cli.channel_prediction \
+  --data_path examples/data/city_8_tempe_3p5_20_32_32.p \
+  --pretrained checkpoints/m18_cp.pth \
+  --inference_only \
+  --device cpu
+```
+
+See `examples/` directory for more detailed usage examples.
+
+---
+
+## 1. Environment Setup
+
+### Requirements
+
+- Python >= 3.9
+- PyTorch >= 2.0.0
+- NumPy >= 1.21.0
+- Matplotlib >= 3.5.0
+
+Verify that your PyTorch build matches your hardware (CPU vs CUDA). Mixed-precision (AMP) is optional; on CPU it will automatically disable itself.
+
+---
+
+## 2. Repository Layout
+
+```
+LWMTemporal/          # Main package
+  cli/                # Command-line entry points
+  data/               # Dataset loaders & preprocessing utilities
+  models/             # LWM model + backbone + configs
+  tasks/              # High-level training/inference orchestration
+  utils/              # Logging and helper utilities
+examples/             # Example scripts and sample data
+  data/               # Example datasets
+checkpoints/          # Pretrained model checkpoints
+cache/                # Optional dataset cache (auto-created)
+figs/predictions/     # Visualization output (auto-created)
+logs/                 # Training logs (auto-created)
+```
+
+---
+
+## 3. Checkpoint Format (Hugging Face Compatible)
+
+The code supports two checkpoint formats:
+
+**Format 1: Directory (Hugging Face style)**
+```
+checkpoints/my_model/
+  config.json           # Model configuration
+  pytorch_model.bin     # Model weights
+```
+
+**Format 2: Single file**
+```
+checkpoints/my_model.pth  # Contains both weights and optional config
+```
+
+The package automatically detects and loads either format.
+
+> **Tip:** If you only have a single file (e.g. `model_best.pth`), move it to a directory and rename to `pytorch_model.bin`. Copy or recreate a matching `config.json`. The loader infers `max_seq_len` when it sees a longer positional embedding in the checkpoint, so older weights continue to work.
+
+The directory can be uploaded to Hugging Face Hub as-is and loaded via `AutoModel.from_pretrained` if you create a thin wrapper.
+
+---
+
+## 4. Dataset Preparation
+
+- The pipeline consumes pickle (`.p`) payloads with a `channel` key (complex tensor) and optional metadata (`pos`, `dt`).
+- `AngleDelaySequenceDataset` normalizes, truncates, and caches angle-delay representations on demand.
+- Configure preprocessing through `DatasetArgs`:
+  - `keep_percentage` – fraction of strongest taps to keep.
+  - `normalize` – `global_rms`, `per_sample_rms`, or `none`.
+  - `cache_dir`, `use_cache`, `overwrite_cache` – caching behavior.
+  - `snr_db`, `noise_seed` – synthetic AWGN injection.
+  - `max_time_steps` – optional temporal truncation.
+
+Cached tensors are stored under `cache/adseq_<stem>_keepXX_<normalize>.pt`.
+
+---
+
+## 5. Command-Line Usage
+
+The CLI mirrors the Hugging Face workflow (`python -m package.cli ...`).
+
+### 5.1 Inference / Evaluation
+
+```bash
+python -m LWMTemporal.cli.channel_prediction \
+  --data_path examples/data/parow.p \
+  --pretrained checkpoints/m18_cp.pth \
+  --inference_only \
+  --inference_split val \
+  --Tpast 10 \
+  --horizon 1
+```
+
+- `--Tpast` / `--horizon` define the autoregressive roll-out window.
+- `--inference_split` selects which subset to score (`train`, `val`, `all`).
+- Visualizations are written to `figs/predictions/`.
+
+### 5.2 Training / Fine-Tuning
+
+Remove `--inference_only` to launch training:
+
+```bash
+python -m LWMTemporal.cli.channel_prediction \
+  --data_path examples/data/parow.p \
+  --save_dir models/finetune_run \
+  --epochs 5 \
+  --batch_size 8 \
+  --lr 3e-4 \
+  --Tpast 10 \
+  --horizon 2
+```
+
+Notable flags:
+
+- `--pretrained` – resume from existing weights.
+- `--train_head_only` – freeze encoder, train output head.
+- `--finetune_last_n` – unfreeze last N transformer blocks.
+- `--global_cls` – enable CLS token for global prediction heads.
+- `--routing_topk_enable`, `--topk_per_head`, etc. – sparse attention controls.
+- `--temporal_offsets` – defaults to `[-4, -3, -2, -1]` so the attention only reaches the previous four frames.
+- `--use_wandb` (with `--wandb_project`, `--wandb_run_name`, `--wandb_entity`) – stream training/eval metrics to Weights & Biases.
+
+Checkpoints are saved as `save_dir/<prefix>_epochXX.pth` together with optimizer state. Use `ChannelPredictionTrainer._save_checkpoint` for custom logic.
+
+---
+
+## 6. Python API Usage
+
+Construct arguments programmatically and drive training/evaluation via `ChannelPredictionTrainer`:
+
+```python
+from pathlib import Path
+from LWMTemporal.tasks.channel_prediction import (
+    ChannelPredictionArgs,
+    DatasetArgs,
+    ModelArgs,
+    TrainingArgs,
+    PredictionArgs,
+    ChannelPredictionTrainer,
+)
+
+args = ChannelPredictionArgs(
+    dataset=DatasetArgs(
+        data_path=Path("examples/data/parow.p"),
+        keep_percentage=0.25,
+        train_limit=500,
+        val_limit=1000,
+    ),
+    model=ModelArgs(
+        patch_size=(1, 1),
+        phase_mode="real_imag",
+        pretrained=Path("checkpoints/m18_cp.pth"),
+    ),
+    training=TrainingArgs(
+        inference_only=True,
+        device="cpu",
+        batch_size=4,
+    ),
+    prediction=PredictionArgs(Tpast=10, horizon=1),
+)
+
+trainer = ChannelPredictionTrainer(args)
+trainer.train()  # runs evaluate() because inference_only=True
+```
+
+From here you can:
+
+- Access `trainer.model` (an `LWMBackbone`) for custom forward passes.
+- Call `trainer.data.train_loader(...)` / `val_loader(...)` for raw dataloaders.
+- Use `trainer.engine.autoregressive_rollout(...)` to obtain `(pred_tokens, target_tokens, mask)` tensors for downstream metrics.
+- Generate visualizations with `trainer.viz.save(...)`.
+
+---
+
+## 7. Working With `LWMBackbone`
+
+- Instantiate from scratch: `LWMBackbone(LWMConfig(...))`.
+- Load checkpoints:
+  ```python
+  from LWMTemporal.models.lwm import LWMBackbone, LWMConfig
+
+  cfg = LWMConfig(patch_size=(1, 1), embed_dim=32, max_seq_len=2816)
+  model = LWMBackbone.from_pretrained("checkpoints/m18_cp.pth", config=cfg)
+  ```
+- Save checkpoints: `model.save_pretrained("path/to/output")`.
+- The loader automatically adjusts `config.max_seq_len` when the checkpoint’s positional embedding is longer than the provided config.
+
+`LWMModel.forward(seq, mask=None, return_cls=False)` accepts complex tensors shaped `(B, T, N, M)` and returns reconstruction tokens along with an optional CLS embedding when enabled.
+
+---
+
+## 8. Visualization & Metrics
+
+- `PredictionVisualizer` renders magnitude plots (`|H|`) for predicted vs. ground-truth angle-delay grids.
+- Metrics:
+  - `masked_nmse_loss` / `compute_nmse` – Normalized MSE over valid tokens.
+  - `masked_mse_loss` – standard MSE with masking support.
+
+Configure masking via token `mask` tensors (boolean) where `True` indicates dropped tokens.
+
+---
+
+## 9. Advanced Configuration
+
+- **Sparse Attention Windows:** Control spatial/temporal neighborhoods via `same_frame_window`, `temporal_offsets`, `temporal_spatial_window`, `temporal_drift_*`, and dilation parameters. The default `temporal_offsets = (-4, -3, -2, -1)` limits attention to the previous four frames.
+- **Routing & Top-k Pruning:** Enable dynamic neighbour pruning with `routing_topk_enable`, `routing_topk_fraction`, `routing_topk_min/max`, or fallback to static `topk_neighbors`.
+- **Positional Encoding:** `posenc` supports `learned` or `rope_sincos`. Additional `rope_base_*` parameters adjust RoPE scaling.
+- **CLS Token:** Toggle `global_cls`; the autoregressive rollout handles CLS automatically when present.
+- **Detokenization:** `AutoregressiveEngine.detokenize` converts predicted tokens back to complex-valued channel coefficients.
+
+---
+
+## 10. Troubleshooting
+
+- **Circular Imports:** The project avoids cross-imports by keeping tokenizers in `models`. Ensure you are on the latest code if you encounter import errors.
+- **Checkpoint Shape Mismatch:** Confirm `patch_size`, `phase_mode`, and positional embedding lengths match between config and weights.
+- **Neighbor Padding Errors:** Patched `NeighborIndexer` pads ragged neighbour lists with `-1`, so any older ValueError is resolved once you update to the current code.
+- **AMP Warnings:** On CPU you may see `GradScaler` warnings; they are benign because AMP disables itself.
+- **Data Shape Mismatch:** Sequences must have consistent `(T, H, W)` dimensions within each batch. The trainer raises a descriptive error otherwise.
+
+---
+
+## 11. Hugging Face Integration
+
+- Because checkpoints follow the standard `config.json` + `pytorch_model.bin` scheme, you can do:
+  ```python
+  from transformers import AutoConfig, AutoModel
+
+  cfg = AutoConfig.from_pretrained("path/to/model_best")
+  model = AutoModel.from_pretrained("path/to/model_best", config=cfg)
+  ```
+- Wrap `LWMBackbone` in a custom `transformers.PreTrainedModel` subclass if you need full pipeline compatibility.
+- Use the same directory structure when publishing to the Hugging Face Hub.
+
+---
+
+## 12. Reproducibility Checklist
+
+- Seed control: `DatasetArgs.seed` (for train/val splits); manual seeding via `torch.manual_seed` and `np.random.seed` happens inside the trainer.
+- Log frequency: `TrainingArgs.log_interval`.
+- Gradient clipping: `TrainingArgs.grad_clip` (defaults to 1.0).
+- Warm-up / Scheduler: cosine decay after a configurable warm-up fraction (`TrainingArgs.warmup_ratio`).
+
+---
+
+## 13. Getting Help
+
+- Issues related to data format, training instabilities, or new features can be logged on your preferred tracking system or discussed with collaborators.
+- For general transformer best practices, refer to the Hugging Face BERT documentation and friends ([link](https://huggingface.co/docs/transformers/en/model_doc/bert?usage=Pipeline)). The workflow above mirrors that style for LWMTemporal.
+
+Happy experimenting!
+
+## Citation
+
+If you use LWMTemporal in your research, please cite:
+
+```bibtex
+@article{lwmtemporal2025,
+  title={Large Wireless Model for Spatio-Temporal Channel Prediction},
+  author={Alikhani, Sadjad and others},
+  journal={arXiv preprint arXiv:XXXX.XXXXX},
+  year={2025}
+}
+```
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
+
+## Acknowledgments
+
+- Built with PyTorch
+- Inspired by Vision Transformer architectures
+- Supports Hugging Face model hub integration
+
+## Contact
+
+For questions or issues, please:
+- Open an issue on GitHub
+- Contact: sadjad.alikhani@asu.edu
+
+## Contributing
+
+Contributions are welcome! Please:
+1. Fork the repository
+2. Create a feature branch
+3. Make your changes
+4. Submit a pull request
+
+For major changes, please open an issue first to discuss the proposed changes.
+

checkpoints/README.md

@@ -0,0 +1,58 @@
+# Model Checkpoints
+
+This directory contains pretrained model checkpoints.
+
+## Available Checkpoints
+
+### `m18_cp.pth`
+- **Task**: Channel prediction (fine-tuned)
+- **Architecture**: 12-layer transformer with 32-dim embeddings
+- **Temporal Attention**: Causal (attends to past 7 frames)
+- **Performance**: ~-20 dB NMSE on validation set
+
+### `pytorch_model.bin`
+- **Task**: Pretrained backbone
+- **Architecture**: Same as above
+- **Temporal Attention**: Bidirectional
+
+## Loading Checkpoints
+
+### Python API
+
+```python
+from pathlib import Path
+from LWMTemporal import LWMBackbone, LWMConfig
+
+# Load with default config from checkpoint
+model = LWMBackbone.from_pretrained("checkpoints/m18_cp.pth")
+
+# Or override config for fine-tuning
+cfg = LWMConfig(
+    temporal_offsets=(-1, -2, -3, -4),  # Override for different task
+)
+model = LWMBackbone.from_pretrained("checkpoints/m18_cp.pth", config=cfg)
+```
+
+### CLI
+
+```bash
+python -m LWMTemporal.cli.channel_prediction \
+  --pretrained checkpoints/m18_cp.pth \
+  --data_path examples/data/parow.p \
+  --inference_only
+```
+
+## Hosting on Hugging Face Hub (Recommended)
+
+For production use, upload checkpoints to Hugging Face Hub:
+
+```bash
+huggingface-cli login
+huggingface-cli upload your-username/lwm-temporal checkpoints/m18_cp.pth
+```
+
+Then load directly from the hub:
+```python
+model = LWMBackbone.from_pretrained("your-username/lwm-temporal")
+```
+

checkpoints/config.json

@@ -0,0 +1,37 @@
+{
+    "patch_size": [1, 1],
+    "phase_mode": "real_imag",
+    "embed_dim": 32,
+    "depth": 12,
+    "num_heads": 8,
+    "mlp_ratio": 4.0,
+    "same_frame_window": 2,
+    "same_frame_window_h": null,
+    "same_frame_window_w": null,
+    "same_frame_dilation_h": 1,
+    "same_frame_dilation_w": 1,
+    "temporal_offsets": [-4, -3, -2, -1, 1, 2, 3],
+    "temporal_spatial_window": 2,
+    "temporal_spatial_window_h": null,
+    "temporal_spatial_window_w": null,
+    "temporal_spatial_dilation_h": 1,
+    "temporal_spatial_dilation_w": 1,
+    "temporal_drift_h": 1,
+    "temporal_drift_w": 1,
+    "spatial_only": false,
+    "routing_topk_enable": true,
+    "routing_topk_fraction": 0.2,
+    "routing_topk_min": 8,
+    "routing_topk_max": 32,
+    "routing_topk_per_head": true,
+    "topk_neighbors": null,
+    "topk_per_head": true,
+    "global_cls": false,
+    "posenc": "learned",
+    "rope_base": 10000.0,
+    "rope_mode": "flat",
+    "rope_base_t": null,
+    "rope_base_h": null,
+    "rope_base_w": null,
+    "max_seq_len": null
+  }

checkpoints/hist/config.json

@@ -0,0 +1,37 @@
+{
+    "patch_size": [1, 1],
+    "phase_mode": "real_imag",
+    "embed_dim": 32,
+    "depth": 12,
+    "num_heads": 8,
+    "mlp_ratio": 4.0,
+    "same_frame_window": 2,
+    "same_frame_window_h": null,
+    "same_frame_window_w": null,
+    "same_frame_dilation_h": 1,
+    "same_frame_dilation_w": 1,
+    "temporal_offsets": [-4, -3, -2, -1, 1, 2, 3],
+    "temporal_spatial_window": 2,
+    "temporal_spatial_window_h": null,
+    "temporal_spatial_window_w": null,
+    "temporal_spatial_dilation_h": 1,
+    "temporal_spatial_dilation_w": 1,
+    "temporal_drift_h": 1,
+    "temporal_drift_w": 1,
+    "spatial_only": false,
+    "routing_topk_enable": true,
+    "routing_topk_fraction": 0.2,
+    "routing_topk_min": 8,
+    "routing_topk_max": 32,
+    "routing_topk_per_head": true,
+    "topk_neighbors": null,
+    "topk_per_head": true,
+    "global_cls": false,
+    "posenc": "learned",
+    "rope_base": 10000.0,
+    "rope_mode": "flat",
+    "rope_base_t": null,
+    "rope_base_h": null,
+    "rope_base_w": null,
+    "max_seq_len": null
+  }

examples/README.md

@@ -0,0 +1,76 @@
+# LWMTemporal Examples
+
+This directory contains example scripts demonstrating how to use the LWMTemporal package.
+
+## Quick Start Examples
+
+### 1. Masked Reconstruction (`example_reconstruction.py`)
+
+Demonstrates how to:
+- Load wireless channel data
+- Tokenize complex channels
+- Mask random positions
+- Reconstruct using the pretrained model
+
+```bash
+python examples/example_reconstruction.py
+```
+
+### 2. Channel Prediction Inference (`inference_channel_prediction.py`)
+
+Run inference with a fine-tuned channel prediction model:
+
+```bash
+python examples/inference_channel_prediction.py
+```
+
+Expected output: Per-step NMSE around -20 dB
+
+### 3. Train Channel Prediction (`train_channel_prediction.py`)
+
+Fine-tune the model for channel prediction:
+
+```bash
+python examples/train_channel_prediction.py
+```
+
+This will:
+- Load pretrained weights
+- Fine-tune on your dataset
+- Save checkpoints to `models/`
+- Generate visualizations in `figs/predictions/`
+
+## Using the CLI
+
+The package also provides command-line interfaces:
+
+### Channel Prediction
+
+```bash
+python -m LWMTemporal.cli.channel_prediction \
+  --data_path examples/data/city_8_tempe_3p5_20_32_32.p \
+  --pretrained checkpoints/m18_cp.pth \
+  --inference_only \
+  --val_limit 100 \
+  --device cpu
+```
+
+### Pretraining
+
+```bash
+python -m LWMTemporal.cli.pretrain \
+  --data_dir examples/data/ \
+  --save_prefix models/pretrained \
+  --epochs 100 \
+  --batch_size 32 \
+  --device cuda
+```
+
+## Data Format
+
+Example data files are in `examples/data/`. See `examples/data/README.md` for details on the expected format.
+
+## Checkpoints
+
+Pretrained checkpoints are in `checkpoints/`. See `checkpoints/README.md` for available models and loading instructions.
+

examples/__init__.py

@@ -0,0 +1,2 @@
+"""LWMTemporal usage examples."""
+

examples/example_reconstruction.py

@@ -0,0 +1,138 @@
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+import torch
+import numpy as np
+
+from LWMTemporal.data.datasets import AngleDelayDatasetConfig, AngleDelaySequenceDataset
+from LWMTemporal.models.lwm import (
+    LWMBackbone,
+    LWMConfig,
+    ComplexPatchTokenizer,
+    masked_nmse_loss,
+)
+
+# ----- 1. Load one sequence (complex tensor) -----
+data_cfg = AngleDelayDatasetConfig(raw_path=Path("examples/data/parow.p"))
+dataset = AngleDelaySequenceDataset(data_cfg)
+
+sequence = dataset[0]["sequence"].unsqueeze(0)        # (1, T, N, M)
+sequence = sequence[:, :11]                          # keep only the first 11 time steps
+print("Sequence shape:", sequence.shape)             # expect (1, 11, 32, 8)
+
+# ----- 2. Tokenise and select tokens to mask -----
+tokenizer = ComplexPatchTokenizer(phase_mode="real_imag")
+tokens, base_mask = tokenizer(sequence, patch_size=(1, 1))  # tokens: (B, S, D)
+
+B, S, D = tokens.shape
+mask_ratio = 0.60                                        # choose the fraction to hide
+mask = base_mask.clone()
+
+# randomly choose the positions that will be hidden
+for b in range(B):
+    num_mask = int(mask_ratio * S)
+    masked_positions = torch.randperm(S)[:num_mask]
+    mask[b, masked_positions] = True
+
+# create the corrupted input by zeroing the masked tokens
+corrupted_tokens = tokens.clone()
+corrupted_tokens[mask] = 0.0
+
+# ----- 3. Load the pretrained backbone -----
+# Need max_seq_len >= S (here 11 * 32 * 8 = 2816)
+cfg = LWMConfig(
+    patch_size=(1, 1),
+    phase_mode="real_imag",
+    embed_dim=32,
+    depth=12,
+    num_heads=8,
+    mlp_ratio=4.0,
+    same_frame_window=2,
+    temporal_offsets=(-4, -3, -2, -1, 1, 2, 3),
+    temporal_spatial_window=2,
+    temporal_drift_h=1,
+    temporal_drift_w=1,
+    routing_topk_enable=True,
+    topk_per_head=True,
+    max_seq_len=2816, # 2816
+)
+
+backbone = LWMBackbone.from_pretrained(Path("checkpoints/m18_cp.pth"), config=cfg)
+backbone.eval()
+
+# ---- 4. Run reconstruction and compute NMSE on the masked positions -----
+with torch.no_grad():
+    # compute H, W from the sequence (N and M dimensions)
+    T = sequence.size(1)
+    H = sequence.size(2)
+    W = sequence.size(3)
+
+    outputs = backbone.forward_tokens(corrupted_tokens, mask, T, H, W, return_cls=False)
+    reconstructed = outputs["reconstruction"]
+
+    nmse = masked_nmse_loss(reconstructed, tokens, mask)
+    nmse_db = 10 * torch.log10(nmse)
+
+print(f"Masked {mask_ratio*100:.1f}% of tokens ({mask.sum().item()} / {S})")
+print(f"NMSE (linear): {nmse.item():.6f}")
+print(f"NMSE (dB):     {nmse_db.item():.2f} dB")
+
+
+
+
+
+# import torch
+# from pathlib import Path
+# from LWMTemporal.data.datasets import AngleDelayDatasetConfig, AngleDelaySequenceDataset
+# from LWMTemporal.models.lwm import (
+#     LWMBackbone,
+#     LWMConfig,
+#     ComplexPatchTokenizer,
+#     masked_nmse_loss,
+# )
+
+# # --- 1. Load one sample from the dataset and keep the first 11 frames ---
+# cfg = AngleDelayDatasetConfig(raw_path=Path("LWMTemporal/data/parow.p"))
+# dataset = AngleDelaySequenceDataset(cfg)
+# sequence = dataset[0]["sequence"].unsqueeze(0)[:, :11]  # (1, 11, 32, 8)
+
+# # --- 2. Tokenise and randomly mask 40% of the tokens ---
+# tokenizer = ComplexPatchTokenizer(phase_mode="real_imag")
+# tokens, base_mask = tokenizer(sequence, patch_size=(1, 1))
+# mask = base_mask.clone()
+
+# B, S, _ = tokens.shape
+# mask_fraction = 0.40
+
+# for b in range(B):
+#     num_mask = int(mask_fraction * S)
+#     masked_positions = torch.randperm(S)[:num_mask]
+#     mask[b, masked_positions] = True
+
+# corrupted_tokens = tokens.clone()
+# corrupted_tokens[mask] = 0.0
+
+# T = sequence.size(1)
+# H = sequence.size(2)
+# W = sequence.size(3)
+
+# # --- 3. Helper to run a model and report NMSE ---
+# def run_model(model: LWMBackbone, label: str) -> None:
+#     model.eval()
+#     with torch.no_grad():
+#         outputs = model.forward_tokens(corrupted_tokens, mask, T, H, W, return_cls=False)
+#         reconstructed = outputs["reconstruction"]
+#         nmse = masked_nmse_loss(reconstructed, tokens, mask)
+#         nmse_db = 10 * torch.log10(nmse)
+#     print(f"{label:>12}: NMSE = {nmse.item():.6f}   ({nmse_db.item():.2f} dB)")
+
+# # --- 4. Random-weights model ---
+# cfg_random = LWMConfig(max_seq_len=11 * sequence.size(2) * sequence.size(3))
+# model_random = LWMBackbone(cfg_random)
+# run_model(model_random, "random init")
+
+# # --- 5. Pretrained checkpoint ---
+# cfg_pretrained = LWMConfig(max_seq_len=11 * sequence.size(2) * sequence.size(3))
+# model_ckpt = LWMBackbone.from_pretrained(Path("LWMTemporal/models"), config=cfg_pretrained)
+# run_model(model_ckpt, "checkpoint")

examples/inference_channel_prediction.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python
+"""Example: Run inference with a trained channel prediction model."""
+
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from LWMTemporal.tasks.channel_prediction import (
+    ChannelPredictionArgs,
+    ChannelPredictionTrainer,
+    DatasetArgs,
+    ModelArgs,
+    TrainingArgs,
+    PredictionArgs,
+)
+from LWMTemporal.utils.logging import setup_logging
+
+# Setup logging
+logger = setup_logging("channel_prediction_inference", log_dir=Path("logs"))
+
+# Configure dataset
+dataset_args = DatasetArgs(
+    data_path=Path("examples/data/city_8_tempe_3p5_20_32_32.p"),
+    keep_percentage=0.25,
+    normalize="global_rms",
+    seed=0,
+    val_limit=100,
+)
+
+# Configure model
+model_args = ModelArgs(
+    patch_size=(1, 1),
+    phase_mode="real_imag",
+    embed_dim=32,
+    depth=12,
+    num_heads=8,
+    mlp_ratio=4.0,
+    same_frame_window=2,
+    temporal_offsets=(-1, -2, -3, -4, -5, -6, -7),  # Causal attention
+    temporal_spatial_window=2,
+    temporal_drift_h=1,
+    temporal_drift_w=1,
+    routing_topk_enable=True,
+    routing_topk_fraction=0.2,
+    routing_topk_max=32,
+    pretrained=Path("checkpoints/m18_cp.pth"),
+)
+
+# Configure training (inference only)
+training_args = TrainingArgs(
+    device="cpu",
+    batch_size=2,
+    inference_only=True,
+    inference_split="val",
+)
+
+# Configure prediction
+prediction_args = PredictionArgs(
+    Tpast=10,
+    horizon=1,
+)
+
+# Build full config
+args = ChannelPredictionArgs(
+    dataset=dataset_args,
+    model=model_args,
+    training=training_args,
+    prediction=prediction_args,
+)
+
+# Run inference
+trainer = ChannelPredictionTrainer(args, logger=logger)
+trainer.train()  # train() handles inference_only mode
+
+logger.info("Inference complete!")
+

examples/train_channel_prediction.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python
+"""Example: Train a channel prediction model."""
+
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+import torch
+from LWMTemporal.tasks.channel_prediction import (
+    ChannelPredictionArgs,
+    ChannelPredictionTrainer,
+    DatasetArgs,
+    ModelArgs,
+    TrainingArgs,
+    PredictionArgs,
+)
+from LWMTemporal.utils.logging import setup_logging
+
+# Setup logging
+logger = setup_logging("channel_prediction_example", log_dir=Path("logs"))
+
+# Configure dataset
+dataset_args = DatasetArgs(
+    data_path=Path("examples/data/city_8_tempe_3p5_20_32_32.p"),
+    keep_percentage=0.25,
+    normalize="global_rms",
+    seed=0,
+    train_limit=500,
+    val_limit=100,
+)
+
+# Configure model
+model_args = ModelArgs(
+    patch_size=(1, 1),
+    phase_mode="real_imag",
+    embed_dim=32,
+    depth=12,
+    num_heads=8,
+    mlp_ratio=4.0,
+    same_frame_window=2,
+    temporal_offsets=(-1, -2, -3, -4, -5, -6, -7),  # Causal attention
+    temporal_spatial_window=2,
+    temporal_drift_h=1,
+    temporal_drift_w=1,
+    routing_topk_enable=True,
+    routing_topk_fraction=0.2,
+    routing_topk_max=32,
+    pretrained=Path("checkpoints/m18_cp.pth"),  # Load pretrained weights
+)
+
+# Configure training
+training_args = TrainingArgs(
+    device="cuda" if torch.cuda.is_available() else "cpu",
+    epochs=10,
+    batch_size=16,
+    lr=1e-4,
+    weight_decay=1e-4,
+    warmup_ratio=0.1,
+    save_dir=Path("models"),
+    use_wandb=False,  # Set to True to enable Weights & Biases logging
+)
+
+# Configure prediction
+prediction_args = PredictionArgs(
+    Tpast=10,
+    horizon=1,
+    viz_dir=Path("figs/predictions"),
+)
+
+# Build full config
+args = ChannelPredictionArgs(
+    dataset=dataset_args,
+    model=model_args,
+    training=training_args,
+    prediction=prediction_args,
+)
+
+# Train
+trainer = ChannelPredictionTrainer(args, logger=logger)
+trainer.train()
+
+logger.info("Training complete!")
+

pyproject.toml

@@ -0,0 +1,60 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "lwm-temporal"
+version = "0.1.0"
+description = "Large Wireless Model (LWM) for spatio-temporal wireless channel representation learning"
+readme = "README.md"
+requires-python = ">=3.9"
+license = {text = "MIT"}
+authors = [
+    {name = "Sadjad Alikhani", email = "alikhani@asu.edu"}
+]
+keywords = ["wireless", "sparse-spatiotemporal-attention", "transformer", "deep-learning", "pytorch"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+dependencies = [
+    "torch>=2.0.0",
+    "numpy>=1.21.0",
+    "matplotlib>=3.5.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0",
+    "black>=22.0",
+    "flake8>=4.0",
+    "mypy>=0.950",
+]
+wandb = [
+    "wandb>=0.13.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/yourusername/lwm-temporal"
+Repository = "https://github.com/yourusername/lwm-temporal"
+Documentation = "https://github.com/yourusername/lwm-temporal#readme"
+
+[project.scripts]
+lwm-pretrain = "LWMTemporal.cli.pretrain:main"
+lwm-channel-prediction = "LWMTemporal.cli.channel_prediction:main"
+
+[tool.setuptools.packages.find]
+include = ["LWMTemporal*"]
+exclude = ["tests*", "examples*", "checkpoints*", "cache*", "logs*", "figs*", "wandb*"]
+
+[tool.setuptools.package-data]
+LWMTemporal = ["models/config.json"]
+

requirements.txt

@@ -0,0 +1,4 @@
+torch>=2.0.0
+numpy>=1.21.0
+matplotlib>=3.5.0
+

setup.py

@@ -0,0 +1,61 @@
+"""
+Setup script for LWMTemporal package.
+This is kept for backward compatibility; the package primarily uses pyproject.toml.
+"""
+
+from setuptools import setup, find_packages
+from pathlib import Path
+
+# Read the README
+this_directory = Path(__file__).parent
+long_description = (this_directory / "README.md").read_text(encoding="utf-8")
+
+setup(
+    name="lwm-temporal",
+    version="0.1.0",
+    author="Sadjad Alikhani",
+    author_email="alikhani@asu.edu",
+    description="Large Wireless Model (LWM) for spatio-temporal wireless channel prediction",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/yourusername/lwm-temporal",
+    packages=find_packages(include=["LWMTemporal", "LWMTemporal.*"]),
+    package_data={
+        "LWMTemporal": ["models/config.json"],
+    },
+    install_requires=[
+        "torch>=2.0.0",
+        "numpy>=1.21.0",
+        "matplotlib>=3.5.0",
+    ],
+    extras_require={
+        "dev": [
+            "pytest>=7.0",
+            "black>=22.0",
+            "flake8>=4.0",
+            "mypy>=0.950",
+        ],
+        "wandb": ["wandb>=0.13.0"],
+    },
+    entry_points={
+        "console_scripts": [
+            "lwm-pretrain=LWMTemporal.cli.pretrain:main",
+            "lwm-channel-prediction=LWMTemporal.cli.channel_prediction:main",
+        ],
+    },
+    classifiers=[
+        "Development Status :: 4 - Beta",
+        "Intended Audience :: Science/Research",
+        "License :: OSI Approved :: MIT License",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    ],
+    python_requires=">=3.9",
+    license="MIT",
+    keywords="wireless channel-prediction transformer deep-learning pytorch",
+)
+

test_package.py

@@ -0,0 +1,174 @@
+#!/usr/bin/env python
+"""
+Test script to verify the LWMTemporal package is properly structured and functional.
+Run this before releasing to ensure everything works.
+"""
+
+import sys
+from pathlib import Path
+
+def test_imports():
+    """Test that all core components can be imported."""
+    print("Testing imports...")
+    try:
+        from LWMTemporal import LWMBackbone, LWMConfig, LWMModel, __version__
+        from LWMTemporal.data import AngleDelaySequenceDataset, AngleDelayDatasetConfig
+        from LWMTemporal.tasks.channel_prediction import ChannelPredictionTrainer
+        from LWMTemporal.tasks.pretraining import PretrainingTrainer
+        print(f"  ✓ All imports successful (version {__version__})")
+        return True
+    except ImportError as e:
+        print(f"  ✗ Import failed: {e}")
+        return False
+
+def test_file_structure():
+    """Test that required files exist."""
+    print("\nTesting file structure...")
+    required_files = [
+        "README.md",
+        "LICENSE",
+        "pyproject.toml",
+        "setup.py",
+        "requirements.txt",
+        "MANIFEST.in",
+        ".gitignore",
+        "CHANGELOG.md",
+        "LWMTemporal/__init__.py",
+        "LWMTemporal/models/lwm.py",
+        "LWMTemporal/models/config.json",
+        "examples/README.md",
+        "checkpoints/README.md",
+    ]
+    
+    all_exist = True
+    for file in required_files:
+        path = Path(file)
+        if path.exists():
+            print(f"  ✓ {file}")
+        else:
+            print(f"  ✗ {file} NOT FOUND")
+            all_exist = False
+    
+    return all_exist
+
+def test_checkpoints():
+    """Test that checkpoints are accessible."""
+    print("\nTesting checkpoints...")
+    checkpoint_dir = Path("checkpoints")
+    
+    if not checkpoint_dir.exists():
+        print(f"  ✗ Checkpoints directory not found")
+        return False
+    
+    checkpoints = list(checkpoint_dir.glob("*.pth")) + list(checkpoint_dir.glob("*.bin"))
+    if checkpoints:
+        print(f"  ✓ Found {len(checkpoints)} checkpoint(s)")
+        for ckpt in checkpoints:
+            print(f"    - {ckpt.name}")
+        return True
+    else:
+        print(f"  ✗ No checkpoint files found")
+        return False
+
+def test_examples():
+    """Test that example files exist."""
+    print("\nTesting examples...")
+    examples_dir = Path("examples")
+    
+    if not examples_dir.exists():
+        print(f"  ✗ Examples directory not found")
+        return False
+    
+    py_files = list(examples_dir.glob("*.py"))
+    if py_files:
+        print(f"  ✓ Found {len(py_files)} example script(s)")
+        for script in py_files:
+            print(f"    - {script.name}")
+        return True
+    else:
+        print(f"  ✗ No example scripts found")
+        return False
+
+def test_data():
+    """Test that example data exists."""
+    print("\nTesting example data...")
+    data_dir = Path("examples/data")
+    
+    if not data_dir.exists():
+        print(f"  ✗ Example data directory not found")
+        return False
+    
+    data_files = list(data_dir.glob("*.p"))
+    if data_files:
+        print(f"  ✓ Found {len(data_files)} data file(s)")
+        for data_file in data_files:
+            size_mb = data_file.stat().st_size / (1024 * 1024)
+            print(f"    - {data_file.name} ({size_mb:.1f} MB)")
+        return True
+    else:
+        print(f"  ⚠ No example data files found (optional)")
+        return True  # Not critical
+
+def test_no_data_in_package():
+    """Test that data files are not in the main package."""
+    print("\nTesting package cleanliness...")
+    package_dir = Path("LWMTemporal")
+    
+    data_files = list(package_dir.rglob("*.p"))
+    checkpoints = list(package_dir.rglob("*.pth")) + list(package_dir.rglob("*.bin"))
+    
+    issues = []
+    if data_files:
+        issues.append(f"Found {len(data_files)} .p files in package (should be in examples/)")
+    if checkpoints:
+        # config.json is OK, but not checkpoints
+        checkpoint_files = [f for f in checkpoints if 'hist' not in str(f)]
+        if checkpoint_files:
+            issues.append(f"Found checkpoint files in package (should be in checkpoints/)")
+    
+    if issues:
+        for issue in issues:
+            print(f"  ⚠ {issue}")
+        return False
+    else:
+        print(f"  ✓ Package directory is clean")
+        return True
+
+def main():
+    """Run all tests."""
+    print("=" * 60)
+    print("LWMTemporal Package Structure Test")
+    print("=" * 60)
+    
+    results = []
+    results.append(("Imports", test_imports()))
+    results.append(("File Structure", test_file_structure()))
+    results.append(("Checkpoints", test_checkpoints()))
+    results.append(("Examples", test_examples()))
+    results.append(("Example Data", test_data()))
+    results.append(("Package Cleanliness", test_no_data_in_package()))
+    
+    print("\n" + "=" * 60)
+    print("SUMMARY")
+    print("=" * 60)
+    
+    passed = sum(1 for _, result in results if result)
+    total = len(results)
+    
+    for name, result in results:
+        status = "✓ PASS" if result else "✗ FAIL"
+        print(f"{status:8} | {name}")
+    
+    print("=" * 60)
+    print(f"Result: {passed}/{total} tests passed")
+    
+    if passed == total:
+        print("\n🎉 Package is ready for release!")
+        return 0
+    else:
+        print("\n⚠️  Some tests failed. Please review and fix.")
+        return 1
+
+if __name__ == "__main__":
+    sys.exit(main())
+