Release LWM Competition Package

README.md

@@ -0,0 +1,170 @@
+<div align="center">
+
+# ⚡ LARGE WIRELESS MODELS (LWMs) 2025 CHALLENGE
+
+**The goal is to improve performance across five wireless downstream tasks by optimizing a baseline LWM and/or designing new downstream models**
+
+[![Model Hub](https://img.shields.io/badge/🤗%20HuggingFace-Model%20Hub-orange?style=flat-square)](https://huggingface.co/wi-lab/lwm-v1.1)
+[![Tutorials](https://img.shields.io/badge/🎓%20Tutorials-Available-brightgreen?style=flat-square)](https://lwm-wireless.net/tutorials)
+[![Website](https://img.shields.io/badge/🌐%20Website-lwm--wireless.net-blue?style=flat-square)](https://lwm-wireless.net/)
+[![Contact](https://img.shields.io/badge/📬%20Contact-lwmwireless@gmail.com-red?style=flat-square)](mailto:lwmwireless@gmail.com)
+
+<p align="center">
+  <a href="#-challenge-overview">Challenge Overview</a> •
+  <a href="#-provided-materials">Provided Materials</a> •
+  <a href="#-getting-started">Getting Started</a> •
+  <a href="#-submission-process">Submission Process</a> •
+  <a href="#-tutorials">Tutorials</a> •
+  <a href="#-citation">Citation</a> •
+  <a href="#-community--support">Community & Support</a> •
+  <a href="#-team">Team</a>
+</p>
+
+&nbsp;
+
+<a target="_blank" href="https://huggingface.co/spaces/wi-lab/lwm-interactive-demo">
+  <img src="https://img.shields.io/badge/▶️%20Try%20Interactive%20Demo-HuggingFace%20Spaces-yellow?style=for-the-badge" height="36px" alt="Try Interactive Demo"/>
+</a>
+
+&nbsp;
+
+</div>
+
+# 📡 Large Wireless Model (LWM) Challenge
+
+Welcome to the official repository of the **LWM 2025 Challenge**, a competition designed to advance the state of foundation models in wireless communications and sensing. Participants are invited to optimize a provided baseline Large Wireless Model (LWM) and design downstream models to tackle five core wireless tasks with limited labeled data.
+
+---
+
+## 🧠 About LWM
+
+**Large Wireless Model (LWM) 1.1** is a Transformer-based foundation model pre-trained using self-supervised learning on over 1 million unlabeled wireless channel samples. It generates rich, task-agnostic embeddings that significantly outperform raw channel representations on downstream tasks—especially when data is scarce or noisy or downstream models need to be simple.
+
+---
+
+## 🏁 Challenge Overview
+
+Participants are given:
+
+- A pre-trained LWM 1.1 checkpoint
+- Baseline downstream task models
+- Training, validation, and public test sets for each task
+- Helper functions and templates
+
+Your goal is to improve the **Composite Generalization Score (CG-Score)** across these five tasks:
+
+1. **LoS/NLoS Classification** – F1-score
+2. **Sub-6 GHz Channel to mmWave Beam Prediction** – Top-1 Beam F1-score
+3. **Channel Interpolation** – Normalized MSE
+4. **Channel Estimation** – Normalized MSE
+5. **Localization** – Normalized Localization Error
+
+Final rankings are based on hidden test sets evaluated by the organizers.
+
+---
+
+## 📦 Provided Materials
+
+This repository contains:
+
+- `pretrained_model.py` — Loads the baseline or your refined LWM model
+- `train_heads.py` — The main script for training and evaluating all task-specific models. **This file must not be modified.** It is provided as a standardized template to ensure fairness and consistency across all teams. Participants must design their submissions to align with this script. The organizers will use an equivalent version of `train_heads.py` for final evaluation, and any deviation from the expected structure will result in automatic disqualification.
+- `train_heads_config.py` — Contains training configs and model head definitions
+- `train_lwm.py` — Contains LWM 1.1 pre-training and dataset reproducibility script
+- `utils.py` — Helper functions (training, scoring, data handling)
+- `task_{t}/` — Contains the training, validation, and public test sets for each downstream task. These datasets are used for jointly fine-tuning your refined LWM and training the corresponding task-specific models. While downstream training is restricted to the provided datasets, you are free to use any dataset for LWM pre-training. Participants are granted early access to the **DeepMIMO v4** dataset, which offers new, large-scale scenarios suitable for extended LWM refinement.
+- `requirements.yml` — Conda environment file for dependency setup
+
+---
+
+## 🚀 Getting Started
+
+### 📥 Clone the repo
+
+```bash
+git clone https://github.com/wireless-intelligence-lab/lwm-competition-2025.git
+cd lwm-competition-2025
+```
+
+### 🛠️ Set up the environment
+```bash
+conda env create -f requirements.yml
+conda activate lwm_env
+```
+
+### 🧪 Run baseline pipeline
+```bash
+python train_heads.py
+```
+This jointly finetunes LWM and trains downstream heads, evaluates on public test sets, and creates a submission ZIP file.
+
+
+### 🧩 Submission Process
+1. Refine your LWM or downstream heads
+2. Update `pretrained_model.py`, `train_heads_config.py`, and `utils.py`.
+3. Run:
+```bash
+python train_heads.py
+```
+4. Submit the generated ZIP file to the competition portal
+
+🛑 **Do not modify `train_heads.py`.** While you may adapt it for local development or experimentation, your final submission **must be fully compatible with the original, unmodified version** provided. The evaluation script used by the organizers assumes this exact structure—any deviation may result in disqualification.
+
+---
+
+## 📚 Tutorials
+
+Visit the official tutorials page:
+
+👉 [https://lwm-wireless.net/tutorials](https://lwm-wireless.net/tutorials)
+
+---
+
+## 🧪 Citation
+
+If you use the LWM model or its components, please cite:
+
+```bibtex
+@misc{alikhani2025largewirelessmodellwm,
+      title={Large Wireless Model (LWM): A Foundation Model for Wireless Channels}, 
+      author={Sadjad Alikhani and Gouranga Charan and Ahmed Alkhateeb},
+      year={2025},
+      eprint={2411.08872},
+      archivePrefix={arXiv},
+      primaryClass={cs.IT},
+      url={https://arxiv.org/abs/2411.08872}, 
+}
+```
+
+---
+
+## 👥 Community & Support
+
+* 💬 [Discussion Forum](https://huggingface.co/wi-lab/lwm-v1.1/discussions)
+* 📨 Contact: [lwmwireless@gmail.com](mailto:lwmwireless@gmail.com)
+
+---
+
+## 👨‍🔬 Team
+
+Developed by the [Wireless Intelligence Lab](https://wi-lab.net) at Arizona State University.
+
+<p align="center">
+  <a href="https://scholar.google.com/citations?user=PKjnTR4AAAAJ&hl=en" target="_blank">
+    <img src="https://img.shields.io/badge/👤 Sadjad Alikhani-Click to view Scholar profile-blue?style=for-the-badge" alt="Sadjad Alikhani">
+  </a>
+  &nbsp;
+  <a href="https://scholar.google.com/citations?user=MHKcvFMAAAAJ&hl=en" target="_blank">
+    <img src="https://img.shields.io/badge/👤 Gouranga Charan-Click to view Scholar profile-green?style=for-the-badge" alt="Gouranga Charan">
+  </a>
+  &nbsp;
+  <a href="https://scholar.google.com/citations?user=dLHw2qcAAAAJ&hl=en" target="_blank">
+    <img src="https://img.shields.io/badge/👤 Ahmed Alkhateeb-Click to view Scholar profile-orange?style=for-the-badge" alt="Ahmed Alkhateeb">
+  </a>
+</p>
+
+
+
+
+
+

model_checkpoint.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:485611f1a0f819f9c673827b8e613887b39672e97072bd7a412866b49d8dd40f
+size 9960738

pretrained_model.py

@@ -0,0 +1,207 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+class LayerNormalization(nn.Module):
+    """
+    Custom Layer Normalization module with learnable scale and bias.
+
+    Args:
+        d_model (int): Dimensionality of the input embeddings.
+        eps (float): A small constant added to variance to avoid division by zero.
+    """
+    def __init__(self, d_model: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.eps = eps
+        self.alpha = nn.Parameter(torch.ones(d_model))
+        self.bias = nn.Parameter(torch.zeros(d_model))
+
+    def forward(self, x):
+        mean = x.mean(dim=-1, keepdim=True)
+        std = x.std(dim=-1, keepdim=True)
+        return self.alpha * (x - mean) / (std + self.eps) + self.bias
+
+class Embedding(nn.Module):
+    """
+    Input embedding module with linear projection and positional encoding.
+
+    Args:
+        element_length (int): Length of each input element (e.g., patch size).
+        d_model (int): Output embedding dimension.
+        max_len (int): Maximum sequence length for positional embeddings.
+    """
+    def __init__(self, element_length, d_model, max_len=513):
+        super().__init__()
+        self.element_length = element_length
+        self.d_model = d_model
+        self.proj = nn.Linear(element_length, d_model)
+        self.pos_embed = nn.Embedding(max_len, d_model)
+        self.norm = LayerNormalization(d_model)
+
+    def forward(self, x):
+        seq_len = x.size(1)
+        pos = torch.arange(seq_len, dtype=torch.long, device=x.device)
+        pos_encodings = self.pos_embed(pos)
+        tok_emb = self.proj(x.float())
+        embedding = tok_emb + pos_encodings
+        return self.norm(embedding)
+
+class ScaledDotProductAttention(nn.Module):
+    """
+    Computes scaled dot-product attention.
+
+    Args:
+        d_k (int): Dimensionality of the key vectors.
+    """
+    def __init__(self, d_k):
+        super().__init__()
+        self.d_k = d_k
+
+    def forward(self, Q, K, V):
+        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(self.d_k)
+        attn = F.softmax(scores, dim=-1)
+        context = torch.matmul(attn, V)
+        return context, attn
+
+class MultiHeadAttention(nn.Module):
+    """
+    Multi-head attention mechanism.
+
+    Args:
+        d_model (int): Total input/output dimension.
+        n_heads (int): Number of attention heads.
+        dropout (float): Dropout probability applied after attention.
+    """
+    def __init__(self, d_model, n_heads, dropout):
+        super().__init__()
+        self.d_k = d_model // n_heads
+        self.d_v = d_model // n_heads
+        self.n_heads = n_heads
+        self.W_Q = nn.Linear(d_model, self.d_k * n_heads)
+        self.W_K = nn.Linear(d_model, self.d_k * n_heads)
+        self.W_V = nn.Linear(d_model, self.d_v * n_heads)
+        self.linear = nn.Linear(n_heads * self.d_v, d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.scaled_dot_attn = ScaledDotProductAttention(self.d_k)
+
+    def forward(self, Q, K, V):
+        residual, batch_size = Q, Q.size(0)
+        q_s = self.W_Q(Q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
+        k_s = self.W_K(K).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
+        v_s = self.W_V(V).view(batch_size, -1, self.n_heads, self.d_v).transpose(1, 2)
+
+        context, attn = self.scaled_dot_attn(q_s, k_s, v_s)
+        output = context.transpose(1, 2).contiguous().view(batch_size, -1, self.n_heads * self.d_v)
+        output = self.linear(output)
+        return residual + self.dropout(output), attn
+
+class PoswiseFeedForwardNet(nn.Module):
+    """
+    Position-wise feed-forward network applied to each token independently.
+
+    Args:
+        d_model (int): Input and output dimensionality.
+        d_ff (int): Hidden layer size in the feed-forward block.
+        dropout (float): Dropout rate applied between layers.
+    """
+    def __init__(self, d_model, d_ff, dropout):
+        super().__init__()
+        self.fc1 = nn.Linear(d_model, d_ff)
+        self.fc2 = nn.Linear(d_ff, d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        return self.fc2(self.dropout(F.relu(self.fc1(x))))
+
+class EncoderLayer(nn.Module):
+    """
+    Transformer encoder block composed of multi-head self-attention,
+    feed-forward network, and layer normalization.
+
+    Args:
+        d_model (int): Embedding dimension.
+        n_heads (int): Number of attention heads.
+        d_ff (int): Hidden size of the feed-forward subnetwork.
+        dropout (float): Dropout probability.
+    """
+    def __init__(self, d_model, n_heads, d_ff, dropout):
+        super().__init__()
+        self.enc_self_attn = MultiHeadAttention(d_model, n_heads, dropout)
+        self.pos_ffn = PoswiseFeedForwardNet(d_model, d_ff, dropout)
+        self.norm1 = LayerNormalization(d_model)
+        self.norm2 = LayerNormalization(d_model)
+
+    def forward(self, enc_inputs):
+        attn_outputs, attn = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs)
+        attn_outputs = self.norm1(enc_inputs + attn_outputs)
+        ff_outputs = self.pos_ffn(attn_outputs)
+        enc_outputs = self.norm2(attn_outputs + ff_outputs)
+        return enc_outputs, attn
+
+class lwm(nn.Module):
+    """
+    Large Wireless Model (LWM): A Transformer-based encoder model for 
+    extracting rich embeddings from wireless channel data.
+
+    Args:
+        element_length (int): Dimensionality of input tokens.
+        d_model (int): Embedding dimension used throughout the network.
+        n_layers (int): Number of Transformer encoder layers.
+        max_len (int): Maximum number of tokens (sequence length).
+        n_heads (int): Number of self-attention heads.
+        dropout (float): Dropout probability used across the model.
+    """
+    def __init__(self, element_length=32, d_model=128, n_layers=12, max_len=513, n_heads=8, dropout=0.1):
+        super().__init__()
+        
+        self.element_length = element_length
+        self.d_model        = d_model
+        self.n_layers       = n_layers
+        self.max_len        = max_len
+        self.n_heads        = n_heads
+        self.dropout        = dropout
+        
+        self.embedding = Embedding(element_length, d_model, max_len)
+        self.layers = nn.ModuleList(
+            [EncoderLayer(d_model, n_heads, d_model*4, dropout) for _ in range(n_layers)]
+        )
+        self.linear = nn.Linear(d_model, d_model)
+        self.norm = LayerNormalization(d_model)
+
+        embed_weight = self.embedding.proj.weight
+        _, n_dim = embed_weight.size()
+        self.decoder = nn.Linear(d_model, n_dim, bias=False)
+        self.decoder_bias = nn.Parameter(torch.zeros(n_dim))
+
+    def forward(self, input_ids, masked_pos=None):
+        """
+        Forward pass of the LWM model.
+
+        Args:
+            input_ids (torch.Tensor): Input tensor of shape (B, T, element_length), where
+                                      B is batch size, T is sequence length.
+            masked_pos (torch.Tensor, optional): Indices of masked positions for patch prediction.
+                                                 If provided, returns logits for these positions.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor] if masked_pos is provided:
+                - logits_lm: Predicted values for masked positions.
+                - output: Full contextualized embeddings for all tokens.
+
+            torch.Tensor if masked_pos is None:
+                - output: Full contextualized embeddings of shape (B, T, d_model).
+        """
+        output = self.embedding(input_ids)
+
+        for layer in self.layers:
+            output, attn = layer(output)
+
+        if masked_pos is not None:
+            masked_pos = masked_pos.long()[:, :, None].expand(-1, -1, output.size(-1))
+            h_masked = torch.gather(output, 1, masked_pos)
+            h_masked = self.norm(F.relu(self.linear(h_masked)))
+            logits_lm = self.decoder(h_masked) + self.decoder_bias
+            return logits_lm, output
+        else:
+            return output

requirements.yml

@@ -0,0 +1,183 @@
+name: lwm_env
+channels:
+  - pytorch
+  - nvidia
+  - conda-forge
+dependencies:
+  - asttokens=3.0.0=pyhd8ed1ab_1
+  - blas=1.0=mkl
+  - brotli-python=1.1.0=py312h275cf98_2
+  - bzip2=1.0.8=h2466b09_7
+  - ca-certificates=2025.4.26=h4c7d964_0
+  - certifi=2025.4.26=pyhd8ed1ab_0
+  - cffi=1.17.1=py312h4389bb4_0
+  - charset-normalizer=3.4.2=pyhd8ed1ab_0
+  - cloudpickle=3.1.1=pyhd8ed1ab_0
+  - colorama=0.4.6=pyhd8ed1ab_1
+  - comm=0.2.2=pyhd8ed1ab_1
+  - cpython=3.12.10=py312hd8ed1ab_0
+  - cuda-cccl=12.9.27=0
+  - cuda-cccl_win-64=12.9.27=0
+  - cuda-cudart=12.1.105=0
+  - cuda-cudart-dev=12.1.105=0
+  - cuda-cupti=12.1.105=0
+  - cuda-libraries=12.1.0=0
+  - cuda-libraries-dev=12.1.0=0
+  - cuda-nvrtc=12.1.105=0
+  - cuda-nvrtc-dev=12.1.105=0
+  - cuda-nvtx=12.1.105=0
+  - cuda-opencl=12.9.19=0
+  - cuda-opencl-dev=12.9.19=0
+  - cuda-profiler-api=12.9.19=0
+  - cuda-runtime=12.1.0=0
+  - cuda-version=12.9=3
+  - debugpy=1.8.14=py312h275cf98_0
+  - decorator=5.2.1=pyhd8ed1ab_0
+  - exceptiongroup=1.3.0=pyhd8ed1ab_0
+  - executing=2.2.0=pyhd8ed1ab_0
+  - filelock=3.18.0=pyhd8ed1ab_0
+  - freetype=2.13.3=h57928b3_1
+  - h2=4.2.0=pyhd8ed1ab_0
+  - hpack=4.1.0=pyhd8ed1ab_0
+  - hyperframe=6.1.0=pyhd8ed1ab_0
+  - idna=3.10=pyhd8ed1ab_1
+  - importlib-metadata=8.7.0=pyhe01879c_1
+  - intel-openmp=2025.1.0=h57928b3_980
+  - ipykernel=6.29.5=pyh4bbf305_0
+  - ipython=8.36.0=pyh9ab4c32_0
+  - jedi=0.19.2=pyhd8ed1ab_1
+  - jinja2=3.1.6=pyhd8ed1ab_0
+  - joblib=1.5.1=pyhd8ed1ab_0
+  - jupyter_client=8.6.3=pyhd8ed1ab_1
+  - jupyter_core=5.8.0=pyh5737063_0
+  - khronos-opencl-icd-loader=2024.10.24=h2466b09_1
+  - krb5=1.21.3=hdf4eb48_0
+  - lcms2=2.17=hbcf6048_0
+  - lerc=4.0.0=h6470a55_1
+  - libblas=3.9.0=1_h8933c1f_netlib
+  - libcblas=3.9.0=12_hb3dda5d_netlib
+  - libcublas=12.1.0.26=0
+  - libcublas-dev=12.1.0.26=0
+  - libcufft=11.0.2.4=0
+  - libcufft-dev=11.0.2.4=0
+  - libcurand=10.3.10.19=0
+  - libcurand-dev=10.3.10.19=0
+  - libcusolver=11.4.4.55=0
+  - libcusolver-dev=11.4.4.55=0
+  - libcusparse=12.0.2.55=0
+  - libcusparse-dev=12.0.2.55=0
+  - libdeflate=1.22=h2466b09_0
+  - libexpat=2.6.3=he0c23c2_0
+  - libffi=3.4.2=h8ffe710_5
+  - libfreetype=2.13.3=h57928b3_1
+  - libfreetype6=2.13.3=h0b5ce68_1
+  - libhwloc=2.11.2=default_hc8275d1_1000
+  - libiconv=1.18=h135ad9c_1
+  - libjpeg-turbo=3.1.0=h2466b09_0
+  - liblapack=3.9.0=12_h13b7882_netlib
+  - libnpp=12.0.2.50=0
+  - libnpp-dev=12.0.2.50=0
+  - libnvjitlink=12.1.105=0
+  - libnvjitlink-dev=12.1.105=0
+  - libnvjpeg=12.1.1.14=0
+  - libnvjpeg-dev=12.1.1.14=0
+  - libpng=1.6.47=h7a4582a_0
+  - libsodium=1.0.20=hc70643c_0
+  - libsqlite=3.46.1=h2466b09_0
+  - libtiff=4.7.0=hfc51747_1
+  - libuv=1.50.0=h2466b09_0
+  - libwebp=1.5.0=h3b0e114_0
+  - libwebp-base=1.5.0=h3b0e114_0
+  - libxcb=1.16=h013a479_1
+  - libxml2=2.13.8=h442d1da_0
+  - libzlib=1.3.1=h2466b09_1
+  - llvmlite=0.44.0=py312h1f7db74_1
+  - m2w64-gcc-libgfortran=5.3.0=6
+  - m2w64-gcc-libs=5.3.0=7
+  - m2w64-gcc-libs-core=5.3.0=7
+  - m2w64-gmp=6.1.0=2
+  - m2w64-libwinpthread-git=5.0.0.4634.697f757=2
+  - markupsafe=3.0.2=py312h31fea79_1
+  - matplotlib-inline=0.1.7=pyhd8ed1ab_1
+  - mkl=2023.1.0=h6a75c08_48682
+  - mpmath=1.3.0=pyhd8ed1ab_1
+  - msys2-conda-epoch=20160418=1
+  - nest-asyncio=1.6.0=pyhd8ed1ab_1
+  - networkx=3.4.2=pyh267e887_2
+  - numba=0.61.2=py312hcccf92d_0
+  - numpy=2.2.6=py312h3150e54_0
+  - opencl-headers=2024.10.24=he0c23c2_0
+  - openjpeg=2.5.3=h4d64b90_0
+  - openssl=3.5.0=ha4e3fda_1
+  - packaging=25.0=pyh29332c3_1
+  - parso=0.8.4=pyhd8ed1ab_1
+  - pickleshare=0.7.5=pyhd8ed1ab_1004
+  - pillow=10.4.0=py312h381445a_1
+  - pip=24.2=pyh8b19718_1
+  - platformdirs=4.3.8=pyhe01879c_0
+  - prompt-toolkit=3.0.51=pyha770c72_0
+  - psutil=7.0.0=py312h4389bb4_0
+  - pthread-stubs=0.4=hcd874cb_1001
+  - pthreads-win32=2.9.1=h2466b09_4
+  - pure_eval=0.2.3=pyhd8ed1ab_1
+  - pycparser=2.22=pyh29332c3_1
+  - pygments=2.19.1=pyhd8ed1ab_0
+  - pynndescent=0.5.13=pyhd8ed1ab_1
+  - pysocks=1.7.1=pyh09c184e_7
+  - python=3.12.6=hce54a09_1_cpython
+  - python-dateutil=2.9.0.post0=pyhff2d567_1
+  - python_abi=3.12=7_cp312
+  - pytorch=2.5.1=py3.12_cuda12.1_cudnn9_0
+  - pytorch-cuda=12.1=hde6ce7c_6
+  - pytorch-mutex=1.0=cuda
+  - pywin32=307=py312h275cf98_3
+  - pyyaml=6.0.2=py312h31fea79_2
+  - pyzmq=26.4.0=py312hd7027bb_0
+  - requests=2.32.3=pyhd8ed1ab_1
+  - scikit-learn=1.6.1=py312h816cc57_0
+  - setuptools=75.1.0=pyhd8ed1ab_0
+  - six=1.17.0=pyhd8ed1ab_0
+  - spyder-kernels=3.0.5=win_pyh7428d3b_0
+  - stack_data=0.6.3=pyhd8ed1ab_1
+  - sympy=1.14.0=pyh04b8f61_5
+  - tbb=2021.13.0=h62715c5_1
+  - threadpoolctl=3.6.0=pyhecae5ae_0
+  - tk=8.6.13=h5226925_1
+  - tornado=6.5.1=py312h4389bb4_0
+  - tqdm=4.67.1=pyhd8ed1ab_1
+  - traitlets=5.14.3=pyhd8ed1ab_1
+  - typing_extensions=4.13.2=pyh29332c3_0
+  - tzdata=2024a=h8827d51_1
+  - ucrt=10.0.22621.0=h57928b3_0
+  - umap-learn=0.5.7=py312h2e8e312_1
+  - urllib3=2.4.0=pyhd8ed1ab_0
+  - vc=14.3=h8a93ad2_21
+  - vc14_runtime=14.40.33810=ha82c5b3_21
+  - vs2015_runtime=14.40.33810=h3bf8584_21
+  - wcwidth=0.2.13=pyhd8ed1ab_1
+  - wheel=0.44.0=pyhd8ed1ab_0
+  - win_inet_pton=1.1.0=pyh7428d3b_8
+  - xorg-kbproto=1.0.7=hcd874cb_1002
+  - xorg-libx11=1.8.9=h0076a8d_1
+  - xorg-libxau=1.0.11=hcd874cb_0
+  - xorg-libxdmcp=1.1.3=hcd874cb_0
+  - xorg-xextproto=7.3.0=hcd874cb_1003
+  - xorg-xproto=7.0.31=hcd874cb_1007
+  - xz=5.2.6=h8d14728_0
+  - yaml=0.2.5=h8ffe710_2
+  - zeromq=4.3.5=ha9f60a1_7
+  - zipp=3.22.0=pyhd8ed1ab_0
+  - zstandard=0.23.0=py312h4389bb4_2
+  - zstd=1.5.7=hbeecb71_2
+  - pip:
+      - contourpy==1.3.2
+      - cycler==0.12.1
+      - fonttools==4.58.0
+      - kiwisolver==1.4.8
+      - matplotlib==3.10.3
+      - pyparsing==3.2.3
+      - scipy==1.15.3
+      - torchaudio==2.5.1
+      - torchvision==0.20.1
+      - deepmimo
+prefix: C:\Users\salikha4\AppData\Local\miniforge3\envs\lwm_env

task_1/config.json

@@ -0,0 +1,15 @@
+{
+  "task_id": 1,
+  "n_total_candidates": 85055,
+  "n_train_samples": 6,
+  "n_val_samples": 500,
+  "n_public_test_samples": 1500,
+  "n_antennas_bs": 8,
+  "n_subcarriers": 8,
+  "max_head_parameters": 500000,
+  "max_wrapper_parameters": 500000,
+  "report_loss_function": "CrossEntropyLoss",
+  "num_classes": 2,
+  "bounding_box_coord": null,
+  "snr": 10
+}

task_1/test_data.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3fcbe4647720f2c8b33a3928dc1535a5f1d1e974c76e589bfe2a6da3b3c2e956
+size 775460

task_1/train_data.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4db0f33461da632166ebc0a550f8e37e2968d47c27c5bd551da6f6d8ad22540a
+size 4586

task_1/val_data.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be18692b4527d205ad0db47a7c20d60b5fc4a71e30e039744febdb6131a51a85
+size 259486

task_2/config.json

@@ -0,0 +1,15 @@
+{
+  "task_id": 2,
+  "n_total_candidates": 67769,
+  "n_train_samples": 1000,
+  "n_val_samples": 500,
+  "n_public_test_samples": 1500,
+  "n_antennas_bs": 64,
+  "n_subcarriers": 16,
+  "max_head_parameters": 500000,
+  "max_wrapper_parameters": 5000000,
+  "report_loss_function": "CrossEntropyLoss",
+  "num_classes": 64,
+  "bounding_box_coord": null,
+  "snr": null
+}

task_2/test_data.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9bca2d3256e34727edc60368990d455ae63773a37b604d04aaac71aaac7d6dab
+size 12295460

task_2/train_data.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d7a7055d6ce7949424f16a356d1038e82a47acb97621e27ec651452c14c9cdf2
+size 8197482

task_2/val_data.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:931ff8ac11532dc0dba852e13c1aa94bd6b2613377ff6ddbf39035628647cfbf
+size 4099486

task_3/config.json

@@ -0,0 +1,15 @@
+{
+  "task_id": 3,
+  "n_total_candidates": 74226,
+  "n_train_samples": 300,
+  "n_val_samples": 500,
+  "n_public_test_samples": 1500,
+  "n_antennas_bs": 64,
+  "n_subcarriers": 64,
+  "max_head_parameters": 500000,
+  "max_wrapper_parameters": 5000000,
+  "report_loss_function": "MSELoss",
+  "num_classes": null,
+  "bounding_box_coord": null,
+  "snr": null
+}

task_3/train_data.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:60150d8975c362133c104a65ec7d036e48a16c947fff3a0eedcc934530759b81
+size 19662314

task_3/val_data.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d48d0644282d658ef2954f64b8fdce8b17056355c0f9b49674d9197acd3c1c1d
+size 32769502

task_4/config.json

@@ -0,0 +1,15 @@
+{
+  "task_id": 4,
+  "n_total_candidates": 75666,
+  "n_train_samples": 250,
+  "n_val_samples": 500,
+  "n_public_test_samples": 1500,
+  "n_antennas_bs": 32,
+  "n_subcarriers": 32,
+  "max_head_parameters": 500000,
+  "max_wrapper_parameters": 5000000,
+  "report_loss_function": "MSELoss",
+  "num_classes": null,
+  "bounding_box_coord": null,
+  "snr": 10
+}

task_4/test_data.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a6d9555fa19cd9c633c643029c5c1fc8b8984c352a731cac184c8529c772fe8
+size 24577508

task_4/train_data.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5210cd429217f85010d329b6f91837945d7bd7092f4d00b7c802a67c18f358ff
+size 4097514

task_4/val_data.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:34901cf96ac5acf22578c3622aeca9e815ede2b0032bfe15702de48ed72ad32e
+size 8193502

task_5/config.json

@@ -0,0 +1,24 @@
+{
+  "task_id": 5,
+  "n_total_candidates": 75044,
+  "n_train_samples": 450,
+  "n_val_samples": 500,
+  "n_public_test_samples": 1500,
+  "n_antennas_bs": 256,
+  "n_subcarriers": 32,
+  "max_head_parameters": 500000,
+  "max_wrapper_parameters": 5000000,
+  "report_loss_function": "MSELoss",
+  "num_classes": null,
+  "bounding_box_coord": [
+    [
+      -144.9029998779297,
+      -164.35699462890625
+    ],
+    [
+      145.0970001220703,
+      163.64300537109375
+    ]
+  ],
+  "snr": null
+}

task_5/train_data.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bec5d102cd89fa1b0e697db6f311821081422d702ca783c201c5ae68e155d174
+size 29496298

task_5/val_data.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:40c81e686ea6d5fbac4b3ea5102ebb92585193a0f4da5ad20b7440160aedd7ca
+size 32773470

train_heads.py

@@ -0,0 +1,777 @@
+import os
+import torch
+import torch.nn as nn
+import json
+import numpy as np
+from torch.utils.data import TensorDataset, DataLoader
+import shutil
+from tqdm import tqdm
+from sklearn.metrics import f1_score
+import matplotlib.pyplot as plt
+from typing import Optional, Tuple, List, Dict, Any
+import sys
+import warnings
+warnings.filterwarnings("ignore")
+from utils import embedding_space_visual, tokenizer, plot_radar_chart
+from pretrained_model import lwm
+import train_heads_config as thc
+
+# Set environment variable for CuBLAS deterministic behavior
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+# Ensure deterministic behavior
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+torch.use_deterministic_algorithms(True)
+
+# Worker initialization for DataLoader to ensure reproducible shuffling
+def worker_init_fn(worker_id):
+    np.random.seed(42 + worker_id)
+
+# List of TaskHeads
+task_heads = [
+    thc.LosNlosClassificationHead,
+    thc.BeamPredictionHead,
+    thc.ChannelInterpolationHead,
+    thc.ChannelEstimationHead,
+    thc.ChannelChartingHead
+]
+
+# Fine-tuning wrapper for LWM and downstream model
+class FineTuningWrapper(nn.Module):
+    def __init__(self, model, task_head, fine_tune_layers="full"):
+        """
+        Initialize the FineTuningWrapper to manage fine-tuning of a model with a task-specific head.
+
+        Args:
+            model (nn.Module): The base model (e.g., LWM) to be fine-tuned.
+            task_head (nn.Module): The task-specific head for downstream tasks.
+            fine_tune_layers (str or list, optional): Specifies which layers to fine-tune.
+                If "full", all model layers are unfrozen. If a list, only specified layers are unfrozen.
+                Defaults to "full".
+
+        Raises:
+            ValueError: If a specified layer in fine_tune_layers is not found in the model.
+        """
+        super().__init__()
+        self.model = model
+        self.task_head = task_head
+        
+        # Freeze all layers initially
+        for param in self.model.parameters():
+            param.requires_grad = False
+        
+        # Handle fine-tuning layers
+        if fine_tune_layers is not None:
+            if fine_tune_layers == "full":
+                # Unfreeze all layers if "full" is specified
+                for param in self.model.parameters():
+                    param.requires_grad = True
+            else:
+                # Get a list of all available layer names in the model
+                available_layers = [name for name, _ in self.model.named_parameters()]
+                
+                # Validate that specified layers exist in the model
+                for layer in fine_tune_layers:
+                    if not any(layer in lname for lname in available_layers):
+                        raise ValueError(
+                            f"Layer '{layer}' not found in the model. "
+                            f"Available layers: {available_layers}"
+                        )
+                
+                # Unfreeze only the specified layers
+                for name, param in self.model.named_parameters():
+                    if any(layer in name for layer in fine_tune_layers):
+                        param.requires_grad = True
+
+    def forward(self, x, input_type="cls_emb", selected_tokens=None):
+        """
+        Forward pass through the model and task head, processing input based on specified type.
+
+        Args:
+            x (torch.Tensor): Input tensor to the model.
+            input_type (str, optional): Type of embedding to extract from the model.
+                Options: "raw", "cls_emb", "channel_emb", "combined", "mean_pooled",
+                "arbitrary_concat", "arbitrary_meanPooled". Defaults to "cls_emb".
+            selected_tokens (list, optional): List of token indices for "arbitrary_concat"
+                or "arbitrary_meanPooled" input types. Defaults to None.
+
+        Returns:
+            torch.Tensor: Output of the task head after processing the input embeddings.
+        """
+        if input_type == "raw":
+            # Use the original raw channel input directly for the downstream task
+            task_input = x
+        
+        else:
+            # Pass input through the LWM model to obtain transformer embeddings
+            embeddings = self.model(x)
+        
+            if input_type == "cls_emb":
+                # Extract only the [CLS] token embedding (assumed to be at index 0)
+                task_input = embeddings[:, [0]]
+        
+            elif input_type == "channel_emb":
+                # Use all patch embeddings except the [CLS] token
+                task_input = embeddings[:, 1:]
+        
+            elif input_type == "combined":
+                # Concatenate [CLS] and patch embeddings for full representation
+                task_input = embeddings
+        
+            elif input_type == "mean_pooled":
+                # Compute the mean over all token embeddings and retain sequence dimension
+                task_input = torch.mean(embeddings, dim=1).unsqueeze(1)
+        
+            elif input_type == "arbitrary_concat":
+                # Concatenate a selected subset of token embeddings by index
+                # `selected_tokens` should be a list of token indices to include
+                task_input = embeddings[:, selected_tokens]
+        
+            elif input_type == "arbitrary_meanPooled":
+                # Compute mean-pooled embedding over a selected subset of tokens
+                # and add a singleton sequence dimension
+                task_input = torch.mean(embeddings[:, selected_tokens], dim=1).unsqueeze(1)
+
+        return self.task_head(task_input)
+
+def nmse(y_true, y_pred):
+    """
+    Calculate the Normalized Mean Squared Error (NMSE) between true and predicted values.
+
+    Args:
+        y_true (array-like): Ground truth values.
+        y_pred (array-like): Predicted values.
+
+    Returns:
+        float: The NMSE value, computed as the mean squared error divided by the mean
+               squared magnitude of the true values.
+    """
+    y_true = np.array(y_true)
+    y_pred = np.array(y_pred)
+    return np.mean(np.abs(y_true - y_pred)**2) / np.mean(np.abs(y_true)**2)
+
+def pow2db(nmse):
+    """
+    Convert a Normalized Mean Squared Error (NMSE) value to decibels (dB).
+
+    Args:
+        nmse (float): The NMSE value to convert.
+
+    Returns:
+        float: The NMSE value in decibels, calculated as 10 * log10(nmse).
+    """
+    return 10 * np.log10(nmse)
+
+def finetune(
+    base_model: nn.Module,
+    train_loader: DataLoader,
+    val_loader: Optional[DataLoader] = None,
+    test_loader: Optional[DataLoader] = None,
+    input_type: str = "cls_emb",
+    fine_tune_layers: Optional[str] = None,
+    optimizer_config: Optional[Dict[str, Any]] = None,
+    scheduler_config: Optional[Dict[str, Any]] = None,
+    epochs: int = 50,
+    device: str = "cuda",
+    task: Optional[str] = None,
+    d_model: Optional[int] = None,
+    sequence_length: Optional[int] = None,
+    selected_tokens: Optional[List[int]] = None,
+    bbox_coord: Optional[float] = None,
+    max_head_pars: int = 1e5,
+    max_wrapper_pars: int = 3e6,
+) -> Tuple[nn.Module, List[float], List[float], List[float], List[float], List[torch.Tensor], List[torch.Tensor]]:
+    """
+    Fine-tune a pre-trained base model with a task-specific head on a given dataset.
+
+    Args:
+        base_model (nn.Module): Pre-trained base model (e.g., LWM) to fine-tune.
+        train_loader (DataLoader): DataLoader for the training dataset.
+        val_loader (Optional[DataLoader]): DataLoader for the validation dataset. Defaults to None.
+        test_loader (Optional[DataLoader]): DataLoader for the test dataset. Defaults to None.
+        input_type (str): Type of input embedding to use. Options: 'cls_emb', 'mean_pooled',
+            'channel_emb', 'combined', 'arbitrary_meanPooled'. Defaults to 'cls_emb'.
+        fine_tune_layers (Optional[str]): Layers to fine-tune in the base model. If 'full', all
+            layers are fine-tuned; if a list, only specified layers are fine-tuned. Defaults to None.
+        optimizer_config (Optional[Dict[str, Any]]): Configuration for the optimizer.
+            Defaults to {'lr': 1e-3} if None.
+        scheduler_config (Optional[Dict[str, Any]]): Configuration for the learning rate scheduler.
+            Defaults to {'step_size': 1000, 'gamma': 0.99} if None.
+        epochs (int): Number of training epochs. Defaults to 50.
+        device (str): Device for training ('cuda' or 'cpu'). Defaults to 'cuda'.
+        task (Optional[str]): Task name. Options: 'LosNlosClassification', 'BeamPrediction',
+            'ChannelInterpolation', 'ChannelEstimation', 'ChannelCharting'. Defaults to None.
+        d_model (Optional[int]): Dimensionality of the model embeddings. Required.
+        sequence_length (Optional[int]): Length of the input sequence. Required for
+            'channel_emb' or 'combined' input types.
+        selected_tokens (Optional[List[int]]): List of token indices for 'arbitrary_meanPooled'
+            or 'arbitrary_concat' input types. Defaults to None.
+        bbox_coord (Optional[float]): Bounding box coordinate (not used in the function).
+            Defaults to None.
+        max_head_pars (int): Maximum allowed parameters in the task head. Defaults to 100,000.
+        max_wrapper_pars (int): Maximum allowed parameters in the wrapper. Defaults to 3,000,000.
+
+    Returns:
+        Tuple containing:
+            - nn.Module: Fine-tuned wrapper model.
+            - List[float]: Training losses per epoch.
+            - List[float]: Validation losses per epoch.
+            - List[float]: Test loss (single value) after training.
+            - List[float]: Task-specific score (e.g., F1-score or normalized score).
+            - List[torch.Tensor]: Ground truth labels from the test set.
+            - List[torch.Tensor]: Predictions from the test set.
+
+    Raises:
+        ValueError: If task, d_model, or input_type is invalid, or required parameters
+            (e.g., sequence_length, selected_tokens) are missing.
+    """
+    # Validate inputs
+    if task is None or d_model is None:
+        raise ValueError("Task and d_model must be provided.")
+    if input_type not in ["cls_emb", "mean_pooled", "channel_emb", "combined", "arbitrary_meanPooled"]:
+        raise ValueError(f"Invalid input_type: {input_type}")
+
+    # Determine number of patches based on input type
+    if input_type in ["cls_emb", "mean_pooled", "arbitrary_meanPooled"]:
+        n_patches = 1
+    elif input_type == "channel_emb":
+        if sequence_length is None:
+            raise ValueError("sequence_length must be provided for input_type 'channel_emb'.")
+        n_patches = sequence_length - 1
+    elif input_type == "combined":
+        if sequence_length is None:
+            raise ValueError("sequence_length must be provided for input_type 'combined'.")
+        n_patches = sequence_length
+    else:  # arbitrary_meanPooled
+        if selected_tokens is None:
+            raise ValueError("selected_tokens must be provided for input_type 'arbitrary_meanPooled'.")
+        n_patches = len(selected_tokens)
+
+    # Define input dimension
+    input_dim = (n_patches, d_model)
+
+    # Dynamically determine output_dim for regression tasks
+    output_dim = None
+    if task in ["ChannelInterpolation", "ChannelEstimation"]:
+        for batch in train_loader:
+            output_dim = batch[1].shape[1:]
+            break  # Use the first batch to determine output shape
+
+    # Handle DataParallel models
+    if isinstance(base_model, nn.DataParallel):
+        base_model = base_model.module
+
+    # Initialize task-specific head
+    if task == "LosNlosClassification":
+        task_head = thc.LosNlosClassificationHead(input_dim)
+    elif task == "BeamPrediction":
+        task_head = thc.BeamPredictionHead(input_dim)
+    elif task == "ChannelInterpolation":
+        if output_dim is None:
+            raise ValueError("output_dim could not be determined for ChannelInterpolation.")
+        task_head = thc.ChannelInterpolationHead(input_dim, output_dim)
+    elif task == "ChannelEstimation":
+        if output_dim is None:
+            raise ValueError("output_dim could not be determined for ChannelEstimation.")
+        task_head = thc.ChannelEstimationHead(input_dim, output_dim)
+    elif task == "ChannelCharting":
+        task_head = thc.ChannelChartingHead(input_dim)
+    else:
+        raise ValueError(f"Unsupported task: {task}")
+
+    # Set up loss criterion
+    if task in ["LosNlosClassification", "BeamPrediction"]:
+        criterion = nn.CrossEntropyLoss()
+    elif task in ["ChannelInterpolation", "ChannelEstimation", "ChannelCharting"]:
+        criterion = nn.MSELoss()
+
+    # Initialize the fine-tuning wrapper
+    fine_tune_layers_config = None if task == "LosNlosClassification" else fine_tune_layers
+    wrapper = FineTuningWrapper(
+        model=base_model,
+        task_head=task_head,
+        fine_tune_layers=fine_tune_layers_config
+    )
+    wrapper = wrapper.to(device)
+    
+    n_head_pars = count_parameters(wrapper.task_head)
+    n_wrapper_pars = count_parameters(wrapper)
+    print(f"\nNumber of head parameters:    {n_head_pars}")
+    print(f"Number of wrapper parameters: {n_wrapper_pars}\n")
+    if n_head_pars > max_head_pars or n_wrapper_pars > max_wrapper_pars:
+        reasons = []
+        if n_head_pars > max_head_pars:
+            reasons.append(
+                f"head parameters ({n_head_pars}) exceed maximum allowed ({max_head_pars})"
+            )
+        if n_wrapper_pars > max_wrapper_pars:
+            reasons.append(
+                f"wrapper parameters ({n_wrapper_pars}) exceed maximum allowed ({max_wrapper_pars})"
+            )
+        print("Stopping run because " + " and ".join(reasons))
+        sys.exit(1)
+        
+    # Save universal LWM weights
+    os.makedirs("submission", exist_ok=True)
+    torch.save(base_model.state_dict(), "submission/model_checkpoint.pth")
+    shutil.copy("pretrained_model.py", "submission/pretrained_model.py")
+    shutil.copy("utils.py", "submission/utils.py")
+    shutil.copy("train_heads_config.py", "submission/train_heads_config.py")
+    shutil.copy("train_heads.py", "submission/train_heads.py")
+    
+    # Set default optimizer config if not provided
+    if optimizer_config is None:
+        optimizer_config = {"lr": 1e-3}
+    optimizer = torch.optim.Adam(wrapper.parameters(), **optimizer_config)
+    
+    # Set up the scheduler
+    if scheduler_config is None:
+        scheduler_config = {"step_size": 1000, "gamma": 0.99}
+    scheduler = torch.optim.lr_scheduler.StepLR(
+        optimizer,
+        step_size=scheduler_config["step_size"],
+        gamma=scheduler_config["gamma"]
+    )
+    
+    # Initialize training utilities
+    train_losses, val_losses, f1_scores = [], [], []
+    predictions, ground_truth = [], []
+    
+    # Training loop
+    for epoch in range(epochs):
+        wrapper.train()
+        epoch_loss = 0.0
+        batch_count = 0
+        train_preds, train_targets = [], []
+        
+        # Prepare a single validation batch
+        val_batch = None
+        val_iterator = iter(val_loader) if val_loader else None
+        if val_iterator:
+            try:
+                val_batch = next(val_iterator)
+            except StopIteration:
+                val_iterator = None
+                
+        with tqdm(train_loader, desc=f"Task Epoch {epoch + 1}/{epochs}", leave=True) as progress_bar:
+            for batch in progress_bar:
+                input_data, targets = batch[0].to(device), batch[1].to(device)
+                optimizer.zero_grad()
+    
+                outputs = wrapper(input_data, 
+                                input_type=input_type,
+                                selected_tokens=selected_tokens)
+                if task in ["LosNlosClassification", "BeamPrediction"]:
+                    preds = torch.argmax(outputs, dim=1).cpu().numpy()
+                    train_preds.extend(preds)
+                    train_targets.extend(targets.cpu().numpy())
+                elif task in ["ChannelInterpolation", "ChannelEstimation"]:
+                    train_preds.extend(outputs.cpu().detach().numpy().flatten())
+                    train_targets.extend(targets.cpu().detach().numpy().flatten())
+                loss = criterion(outputs, targets)
+                
+                loss.backward()
+                optimizer.step()
+                
+                epoch_loss += loss.item()
+                batch_count += 1
+                running_avg_loss = epoch_loss / batch_count
+                
+                train_metric = None
+                if task in ["LosNlosClassification", "BeamPrediction"] and train_preds and train_targets:
+                    train_metric = f1_score(train_targets, train_preds, average="weighted")
+                elif task in ["ChannelInterpolation", "ChannelEstimation"] and train_preds and train_targets:
+                    train_metric = nmse(train_targets, train_preds)
+    
+                val_loss = 0.0
+                val_preds, val_targets = [], []
+                if val_batch:
+                    wrapper.eval()
+                    with torch.no_grad():
+                        val_input_data, val_targets_batch = val_batch[0].to(device), val_batch[1].to(device)
+                        val_outputs = wrapper(val_input_data, input_type=input_type)
+                        if task in ["LosNlosClassification", "BeamPrediction"]:
+                            val_preds = torch.argmax(val_outputs, dim=1).cpu().numpy()
+                            val_targets = val_targets_batch.cpu().numpy()
+                        elif task in ["ChannelInterpolation", "ChannelEstimation"]:
+                            val_preds = val_outputs.cpu().numpy().flatten()
+                            val_targets = val_targets_batch.cpu().numpy().flatten()
+                        val_loss = criterion(val_outputs, val_targets_batch).item()
+    
+                    avg_val_loss = val_loss if val_loss > 0 else None
+                    val_metric = None
+                    if task in ["LosNlosClassification", "BeamPrediction"] and len(val_preds) and len(val_targets):
+                        val_metric = f1_score(val_targets, val_preds, average="weighted")
+                    elif task in ["ChannelInterpolation", "ChannelEstimation"] and len(val_preds) and len(val_targets):
+                        val_metric = nmse(val_targets, val_preds)
+                    
+                    # Switch back to training mode for the next batch
+                    wrapper.train()
+    
+                postfix_dict = {
+                    "Batch Loss": f"{loss.item():.6f}",
+                    "Avg Train Loss": f"{running_avg_loss:.6f}",
+                }
+                if train_metric is not None:
+                    if task in ["LosNlosClassification", "BeamPrediction"]:
+                        postfix_dict["Train F1-Score"] = f"{train_metric:.4f}"
+                    elif task in ["ChannelInterpolation", "ChannelEstimation"]:
+                        postfix_dict["Train NMSE"] = f"{pow2db(train_metric):.6f}"
+                if avg_val_loss is not None:
+                    postfix_dict["Avg Val Loss"] = f"{avg_val_loss:.6f}"
+                if val_metric is not None:
+                    if task in ["LosNlosClassification", "BeamPrediction"]:
+                        postfix_dict["Val F1-Score"] = f"{val_metric:.4f}"
+                    elif task in ["ChannelInterpolation", "ChannelEstimation"]:
+                        postfix_dict["Val NMSE"] = f"{pow2db(val_metric):.6f}"
+    
+                progress_bar.set_postfix(postfix_dict)
+                progress_bar.refresh()
+    
+            avg_train_loss = epoch_loss / len(train_loader)
+            train_losses.append(avg_train_loss)
+    
+            train_metric = None
+            if task in ["LosNlosClassification", "BeamPrediction"] and train_preds and train_targets:
+                train_metric = f1_score(train_targets, train_preds, average="weighted")
+            elif task in ["ChannelInterpolation", "ChannelEstimation"] and train_preds and train_targets:
+                train_metric = nmse(train_targets, train_preds)
+    
+            val_loss = 0.0
+            val_preds, val_targets = [], []
+            if val_loader:
+                wrapper.eval()
+                with torch.no_grad():
+                    for batch in val_loader:
+                        input_data, targets = batch[0].to(device), batch[1].to(device)
+                        outputs = wrapper(input_data, input_type=input_type)
+                        if task in ["LosNlosClassification", "BeamPrediction"]:
+                            preds = torch.argmax(outputs, dim=1).cpu().numpy()
+                            val_preds.extend(preds)
+                            val_targets.extend(targets.cpu().numpy())
+                        elif task in ["ChannelInterpolation", "ChannelEstimation"]:
+                            val_preds.extend(outputs.cpu().numpy().flatten())
+                            val_targets.extend(targets.cpu().numpy().flatten())
+                        elif task == "ChannelCharting":
+                            val_preds.extend(outputs.cpu().numpy().flatten())
+                            val_targets.extend(targets.cpu().numpy().flatten())
+                        loss = criterion(outputs, targets)
+                        val_loss += loss.item()
+    
+                avg_val_loss = val_loss / len(val_loader)
+                val_losses.append(avg_val_loss)
+    
+                val_metric = None
+                if task in ["LosNlosClassification", "BeamPrediction"] and val_preds and val_targets:
+                    val_metric = f1_score(val_targets, val_preds, average="weighted")
+                    f1_scores.append(val_metric)
+                elif task in ["ChannelInterpolation", "ChannelEstimation"] and val_preds and val_targets:
+                    val_metric = nmse(val_targets, val_preds)
+                elif task == "ChannelCharting" and val_preds and val_targets:
+                    val_metric = np.mean(np.abs(np.array(val_targets) - np.array(val_preds)))
+    
+                if val_metric is not None and task == "ChannelCharting":
+                    print(f"Validation Prediction Error (meters) at epoch {epoch + 1}: {val_metric:.2f}")
+    
+            postfix_dict = {
+                "Avg Train Loss": f"{avg_train_loss:.6f}",
+            }
+            if train_metric is not None:
+                if task in ["LosNlosClassification", "BeamPrediction"]:
+                    postfix_dict["Train F1-Score"] = f"{train_metric:.4f}"
+                elif task in ["ChannelInterpolation", "ChannelEstimation"]:
+                    postfix_dict["Train NMSE"] = f"{pow2db(train_metric):.6f}"
+            if avg_val_loss is not None:
+                postfix_dict["Avg Val Loss"] = f"{avg_val_loss:.6f}"
+            if val_metric is not None:
+                if task in ["LosNlosClassification", "BeamPrediction"]:
+                    postfix_dict["Val F1-Score"] = f"{val_metric:.4f}"
+                elif task in ["ChannelInterpolation", "ChannelEstimation"]:
+                    postfix_dict["Val NMSE"] = f"{pow2db(val_metric):.6f}"
+    
+            progress_bar.set_postfix(postfix_dict)
+            progress_bar.refresh()
+    
+        scheduler.step()
+
+    # Test evaluation
+    test_loss = 0.0
+    test_preds, test_targets = [], []
+    if test_loader:
+        wrapper.eval()
+        with torch.no_grad():
+            for batch in test_loader:
+                input_data, targets = batch[0].to(device), batch[1].to(device)
+                outputs = wrapper(input_data, input_type=input_type)
+                if epoch == epochs - 1:
+                    predictions.append(outputs)
+                    ground_truth.append(targets)
+                if task in ["LosNlosClassification", "BeamPrediction"]:
+                    preds = torch.argmax(outputs, dim=1).cpu().numpy()
+                    test_preds.extend(preds)
+                    test_targets.extend(targets.cpu().numpy())
+                elif task in ["ChannelInterpolation", "ChannelEstimation"]:
+                    test_preds.extend(outputs.cpu().numpy().flatten())
+                    test_targets.extend(targets.cpu().numpy().flatten())
+                elif task == "ChannelCharting":
+                    test_preds.extend(outputs.cpu().numpy().flatten())
+                    test_targets.extend(targets.cpu().numpy().flatten())
+                loss = criterion(outputs, targets)
+                test_loss += loss.item()
+
+        avg_test_loss = test_loss / len(test_loader)
+        
+        test_metric = None
+        if task in ["LosNlosClassification", "BeamPrediction"] and test_preds and test_targets:
+            test_metric = f1_score(test_targets, test_preds, average="weighted")
+        elif task in ["ChannelInterpolation", "ChannelEstimation"] and test_preds and test_targets:
+            test_metric = nmse(test_targets, test_preds)
+        elif task == "ChannelCharting" and test_preds and test_targets:
+            test_metric = np.mean(np.abs(np.array(test_targets) - np.array(test_preds)))
+            
+        print(f"Test Loss: {avg_test_loss:.6f}")
+        if test_metric is not None:
+            if task in ["LosNlosClassification", "BeamPrediction"]:
+                print(f"Test F1-Score: {test_metric:.4f}")
+            elif task in ["ChannelInterpolation", "ChannelEstimation"]:
+                print(f"Test NMSE (dB): {pow2db(test_metric):.6f}")
+            elif task == "ChannelCharting":
+                print(f"Test Prediction Error (meters): {test_metric:.2f}")
+                
+    plt.figure(figsize=(10, 6), dpi=300)
+    plt.plot(range(1, epochs + 1), train_losses, label="Train Loss")
+    if val_losses:
+        plt.plot(range(1, epochs + 1), val_losses, label="Validation Loss")
+    plt.xlabel("Epoch")
+    plt.ylabel("Loss")
+    plt.title("Learning Curves")
+    plt.legend()
+    plt.grid(True)
+    plt.show()
+
+    test_losses = [avg_test_loss] if test_loader else []
+
+    if task in ["LosNlosClassification", "BeamPrediction"]:
+        score = test_metric
+    elif task in ["ChannelInterpolation", "ChannelEstimation"]:
+        db_value = pow2db(test_metric)
+        db_min, db_max = -20.0, 0.0
+        normalized = (db_value - db_min) / (db_max - db_min)
+        score = 1.0 - normalized
+        score = max(0.0, min(1.0, score))
+    elif task == "ChannelCharting":
+        localization_error = max(0.0, min(100.0, test_metric))
+        score = (100.0 - localization_error) / 100.0
+    
+    print("\n=============================================================")
+    print(f"The score for the {task} task is {score:.5f}")
+    print("=============================================================\n")
+    
+    return wrapper, train_losses, val_losses, test_losses, score, ground_truth, predictions
+
+def count_parameters(model):
+    """
+    Calculate the total number of learnable parameters in a PyTorch model.
+
+    Args:
+        model (nn.Module): The PyTorch model to count parameters for.
+
+    Returns:
+        int: The total number of parameters that require gradients.
+    """
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+# Set device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Process each task
+scores = []
+num_tasks = 5
+for t in range(1, num_tasks + 1):
+    # Set random seed for reproducibility
+    seed = thc.training_configs[t-1]["seed"]
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    torch.cuda.manual_seed_all(seed)  # For multi-GPU setups
+    
+    # Load universal LWM
+    pretrained_checkpoint_path = "model_checkpoint.pth"
+    universal_lwm = lwm().to(device)
+    checkpoint = torch.load(pretrained_checkpoint_path, map_location=device)
+    clean_state_dict = {k.replace("module.", ""): v for k, v in checkpoint.items()}
+    universal_lwm.load_state_dict(clean_state_dict)
+    
+    # Create task directory
+    task_dir = f"task_{t}"
+    os.makedirs(task_dir, exist_ok=True)
+    
+    # Load task configuration
+    with open(f"{task_dir}/config.json", "r") as f:
+        config = json.load(f)
+    
+    # Load data
+    train_data = torch.load(f"{task_dir}/train_data.pt", map_location="cpu")
+    val_data = torch.load(f"{task_dir}/val_data.pt", map_location="cpu") if os.path.exists(f"{task_dir}/val_data.pt") else None
+    test_data = torch.load(f"{task_dir}/test_data.pt", map_location="cpu") if os.path.exists(f"{task_dir}/test_data.pt") else None
+    
+    # Retrieve training configuration and task head
+    training_config = thc.training_configs[t-1]
+    TaskHead = task_heads[t-1]
+    
+    # Display task name
+    task_name = training_config['task']
+    title = f" Task {t}: {task_name} "
+    border = "+" + "-" * len(title) + "+"
+    print()
+    print(border)
+    print(f"|{title}|")
+    print(border)
+    print()
+    
+    # Extract channels and labels
+    train_channels = train_data["channels"]
+    val_channels = val_data["channels"] if val_data else None
+    test_channels = test_data["channels"] if test_data else None
+    if t <= 2:
+        train_labels = train_data["labels"].to(device).long()
+        val_labels = val_data["labels"].to(device).long() if val_data else None
+        test_labels = test_data["labels"].to(device).long() if test_data else None
+    else:
+        train_labels = train_data["labels"].to(device)
+        val_labels = val_data["labels"].to(device) if val_data else None
+        test_labels = test_data["labels"].to(device) if test_data else None
+    
+    # Tokenize input data
+    train_tokens = tokenizer(train_channels)
+    val_tokens = tokenizer(val_channels) if val_channels is not None else None
+    test_tokens = tokenizer(test_channels) if test_channels is not None else None
+    
+    # Determine sequence length
+    sequence_length = train_tokens.shape[1]
+    
+    # Create datasets and data loaders
+    train_dataset = TensorDataset(train_tokens, train_labels)
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=training_config["batch_size"],
+        shuffle=True,
+        worker_init_fn=worker_init_fn,
+        num_workers=0  # Single-threaded for reproducibility
+    )
+    if val_data:
+        val_dataset = TensorDataset(val_tokens, val_labels)
+        val_loader = DataLoader(
+            val_dataset,
+            batch_size=training_config["batch_size"],
+            shuffle=False,
+            worker_init_fn=worker_init_fn,
+            num_workers=0
+        )
+    else:
+        val_loader = None
+    if test_data:
+        test_dataset = TensorDataset(test_tokens, test_labels)
+        test_loader = DataLoader(
+            test_dataset,
+            batch_size=training_config["batch_size"],
+            shuffle=False,
+            worker_init_fn=worker_init_fn,
+            num_workers=0
+        )
+    else:
+        test_loader = None
+    
+    # Visualize embeddings before fine-tuning
+    embeddings = embedding_space_visual(
+        universal_lwm,
+        test_tokens,
+        input_type=training_config["input_type"],
+        batch_size=training_config["batch_size"],
+        selected_tokens=training_config["selected_tokens"],
+        task=training_config["task"],
+        labels=test_labels if t <= 2 or t == 5 else None,
+        visualization=True,
+        visualization_method="tsne",
+        device=device
+    )
+    
+    # Fine-tune the model
+    wrapper, train_losses, val_losses, test_losses, score, ground_truth, predictions = finetune(
+        base_model=universal_lwm,
+        train_loader=train_loader,
+        val_loader=val_loader,
+        test_loader=test_loader,
+        input_type=training_config["input_type"],
+        fine_tune_layers=training_config["fine_tune_layers"],
+        optimizer_config=training_config["optimizer_config"],
+        scheduler_config=training_config["scheduler"],
+        epochs=training_config["epochs"],
+        task=training_config["task"],
+        d_model=universal_lwm.d_model,
+        sequence_length=sequence_length,
+        selected_tokens=training_config["selected_tokens"],
+        bbox_coord=config["bounding_box_coord"] if t == 5 else None,
+        max_head_pars=config["max_head_parameters"],
+        max_wrapper_pars=config["max_wrapper_parameters"],
+        device=device
+    )
+    
+    # Visualize embeddings after fine-tuning
+    finetuned_embeddings = embedding_space_visual(
+        wrapper.model,
+        test_tokens,
+        input_type=training_config["input_type"],
+        batch_size=training_config["batch_size"],
+        selected_tokens=training_config["selected_tokens"],
+        task=training_config["task"],
+        labels=test_labels if t <= 2 or t == 5 else None,
+        visualization=True,
+        visualization_method="tsne",
+        device=device
+    )
+
+    # Create submission directory
+    task_dir = f"submission/task_{t}"
+    os.makedirs(task_dir, exist_ok=True)
+    
+    # Save fine-tuned wrapper model weights
+    wrapper_weights_path = os.path.join(task_dir, "wrapper.pt")
+    torch.save(wrapper.state_dict(), wrapper_weights_path)
+    print(f"Saved wrapper weights for task {t} to {wrapper_weights_path}")
+    
+    # Save ground truth and predictions
+    ground_truth_path = os.path.join(task_dir, "ground_truth.pt")
+    predictions_path = os.path.join(task_dir, "predictions.pt")
+    torch.save(ground_truth, ground_truth_path)
+    torch.save(predictions, predictions_path)
+    print(f"Saved ground truth and predictions for task {t}")
+    
+    # Save task score
+    score_path = os.path.join(task_dir, "score.json")
+    with open(score_path, "w") as f:
+        json.dump(float(score), f, indent=7)
+    print(f"Saved task score to {score_path}")
+    
+    scores.append(float(score))
+    
+# Calculate and save composite score
+composite_score = np.mean(scores)
+composite_score_path = os.path.join("submission", "composite_score.json")
+with open(composite_score_path, "w") as f:
+    json.dump(composite_score, f, indent=7)
+print("Saved composite score")
+    
+# Create zip archive
+shutil.make_archive("submission", format="zip", root_dir="submission")
+
+# Define task names and baseline scores
+task_names = ["LoS/NLoS\nClassification", "Beam\nPrediction", "Channel\nInterpolation", "Channel\nEstimation", "User\nLocalization"]
+baseline_scores = [
+    0.9396,
+    0.6137,
+    0.4165,
+    0.4576,
+    0.6711
+]
+plot_radar_chart(task_names, scores, baseline_scores)

train_heads_config.py

@@ -0,0 +1,213 @@
+import torch
+import torch.nn as nn
+from utils import patch_reconstructor
+
+# Define TaskHead for each task
+class LosNlosClassificationHead(nn.Module):
+    """
+    Task head for LoS/NLoS classification.
+
+    Takes flattened patch embeddings as input and outputs class logits for binary classification.
+
+    Args:
+        input_dim (tuple): (n_patches, d_model) — number of patches and feature dimension.
+    """
+    def __init__(self, input_dim):
+        super().__init__()
+        n_patches, d_model = input_dim
+        flattened_dim = n_patches * d_model 
+        self.classifier = nn.Sequential(
+            nn.Linear(flattened_dim, 8),  
+            nn.BatchNorm1d(8),       
+            nn.ReLU(),                 
+            nn.Dropout(0.1),           
+            nn.Linear(8, 2), 
+        )
+
+    def forward(self, x):
+        batch_size = x.size(0)
+        x = x.view(batch_size, -1)  
+        x = self.classifier(x)  
+        return x
+    
+class BeamPredictionHead(nn.Module):
+    """
+    Task head for mmWave beam index prediction.
+
+    Processes flattened patch embeddings and outputs logits over 64 possible beam indices.
+
+    Args:
+        input_dim (tuple): (n_patches, d_model) — number of patches and feature dimension.
+    """
+    def __init__(self, input_dim):
+        super().__init__()
+        n_patches, d_model = input_dim
+        flattened_dim = n_patches * d_model 
+        self.classifier = nn.Sequential(
+            nn.Linear(flattened_dim, 256),  
+            nn.BatchNorm1d(256),       
+            nn.ReLU(),                 
+            nn.Dropout(0.1),           
+            nn.Linear(256, 128),       
+            nn.BatchNorm1d(128),       
+            nn.ReLU(),                 
+            nn.Linear(128, 64)
+        )
+    def forward(self, x):
+        batch_size = x.size(0)
+        x = x.view(batch_size, -1)  
+        x = self.classifier(x)  
+        return x
+
+class ChannelInterpolationHead(nn.Module):
+    """
+    Task head for reconstructing missing channel values from patch embeddings.
+
+    Applies a linear layer to each patch and reconstructs the full channel using patch_reconstructor.
+
+    Args:
+        input_dim (tuple): (n_patches, d_model).
+        output_dim (tuple): (target_channels, n_rows, n_cols) — shape of the output channel matrix.
+    """
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+        n_patches, d_model = input_dim 
+        target_channels, self.n_rows, self.n_cols = output_dim 
+        self.fcn = nn.Sequential(
+            nn.Linear(d_model, 32)
+        )
+
+    def forward(self, x):
+        batch_size, n_patches, d_model = x.size()
+        x = x.reshape(batch_size * n_patches, d_model)
+        x = self.fcn(x)  
+        x = x.reshape(batch_size, n_patches, 32)
+        x = patch_reconstructor(x, self.n_rows, self.n_cols)
+        return x
+
+class ChannelEstimationHead(nn.Module):
+    """
+    Task head for full channel estimation from embeddings.
+
+    Similar to interpolation but typically used for denoising or noisy reconstruction.
+
+    Args:
+        input_dim (tuple): (n_patches, d_model).
+        output_dim (tuple): (target_channels, n_rows, n_cols) — shape of the target full-resolution channel.
+    """
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+        n_patches, d_model = input_dim 
+        target_channels, self.n_rows, self.n_cols = output_dim 
+        self.fcn = nn.Sequential(
+            nn.Linear(d_model, 32)
+        )
+        
+    def forward(self, x):
+        batch_size, n_patches, d_model = x.size()
+        x = x.reshape(batch_size * n_patches, d_model)
+        x = self.fcn(x)  
+        x = x.reshape(batch_size, n_patches, 32)
+        x = patch_reconstructor(x, self.n_rows, self.n_cols)
+        return x
+
+class ChannelChartingHead(nn.Module):
+    """
+    Task head for 2D channel charting (e.g., learning spatial topology).
+
+    Reduces the flattened embeddings into 2D coordinates.
+
+    Args:
+        input_dim (tuple): (n_patches, d_model).
+    """
+    def __init__(self, input_dim):
+        super().__init__()
+        n_patches, d_model = input_dim
+        flattened_dim = n_patches * d_model 
+        self.fcn = nn.Sequential(
+            nn.Linear(flattened_dim, 64),      
+            nn.ReLU(),                       
+            nn.Linear(64, 32),        
+            nn.ReLU(),                      
+            nn.Linear(32, 2)
+        )
+
+    def forward(self, x):
+        batch_size = x.size(0)
+        x = x.view(batch_size, -1)
+        x = self.fcn(x)  
+        return x
+
+# training_configs is a list of dictionaries, each specifying the setup for one downstream task.
+# Each entry includes:
+# - task: Name of the task.
+# - optimizer_config: Learning rate for the optimizer.
+# - scheduler: Step size and decay rate for the learning rate scheduler.
+# - epochs: Total number of training epochs.
+# - batch_size: Number of samples per batch.
+# - loss_function: Loss type ("CrossEntropyLoss" or "MSELoss").
+# - seed: Random seed for reproducibility.
+# - fine_tune_layers: Specifies which parts of the LWM model to fine-tune:
+#     • "full" means all layers are trainable
+#     • A list like ["layers.10", "layers.11"] specifies partial fine-tuning
+# - input_type: Type of embedding input used from LWM:
+#     • "cls_emb", "channel_emb", "mean_pooled", etc.
+# - selected_tokens: Specific tokens to select for input, used in some tasks.
+
+training_configs = [
+    {  # Task 1
+        "task": "LosNlosClassification",
+        "optimizer_config": {"lr": 1e-3},
+        "scheduler": {"step_size": 20, "gamma": 0.5},
+        "epochs": 200,
+        "batch_size": 128,
+        "seed": 42,
+        "fine_tune_layers": ["layers.9", "layers.10", "layers.11"],
+        "input_type": "cls_emb",
+        "selected_tokens": None
+    },
+    {  # Task 2
+        "task": "BeamPrediction",
+        "optimizer_config": {"lr": 1e-3},
+        "scheduler": {"step_size": 20, "gamma": 0.8},
+        "epochs": 70,
+        "batch_size": 128,
+        "seed": 42,
+        "fine_tune_layers": "full",
+        "input_type": "mean_pooled",
+        "selected_tokens": None
+    },
+    {  # Task 3
+        "task": "ChannelInterpolation",
+        "optimizer_config": {"lr": 1e-2},
+        "scheduler": {"step_size": 25, "gamma": 0.2},
+        "epochs": 100,
+        "batch_size": 128,
+        "seed": 42,
+        "fine_tune_layers": ["layers.10", "layers.11"],
+        "input_type": "channel_emb",
+        "selected_tokens": None
+    },
+    {  # Task 4
+        "task": "ChannelEstimation",
+        "optimizer_config": {"lr": 1e-2},
+        "scheduler": {"step_size": 50, "gamma": 0.3},
+        "epochs": 200,
+        "batch_size": 128,
+        "seed": 42,
+        "fine_tune_layers": ["layers.9", "layers.10", "layers.11"],
+        "input_type": "channel_emb",
+        "selected_tokens": None
+    },
+    {  # Task 5
+        "task": "ChannelCharting",
+        "optimizer_config": {"lr": 1e-3},
+        "scheduler": {"step_size": 40, "gamma": 0.6},
+        "epochs": 150,
+        "batch_size": 128,
+        "seed": 42,
+        "fine_tune_layers": ["layers.10", "layers.11"],
+        "input_type": "mean_pooled",
+        "selected_tokens": None
+    }
+]

train_lwm.py

@@ -0,0 +1,818 @@
+# =============================================================================
+# 1. IMPORTS AND WARNINGS SETUP
+#    - Load necessary PyTorch modules, utilities, and suppress UserWarnings
+# =============================================================================
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, random_split, TensorDataset
+import torch.optim as optim
+from utils import (generate_channels_and_labels, tokenizer_train, 
+                   create_train_dataloader, count_parameters, train_lwm)
+import numpy as np
+import pretrained_model  # Assuming this contains the LWM model definition
+from torch.optim.lr_scheduler import LambdaLR
+from torch.optim import AdamW
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# =============================================================================
+# 2. SCENARIO LIST DEFINITION
+#    - Define the list of scenario names to iterate over for data generation
+# =============================================================================
+def scenarios_list():
+    scen_list = np.array([
+        'city_0_newyork_3p5_lwm', 
+        'city_1_losangeles_3p5_lwm', 
+        'city_2_chicago_3p5_lwm', 
+        'city_3_houston_3p5_lwm', 
+        'city_4_phoenix_3p5_lwm', 
+        'city_5_philadelphia_3p5_lwm', 
+        'city_6_miami_3p5_lwm', 
+        'city_7_sandiego_3p5_lwm',
+        'city_8_dallas_3p5_lwm', 
+        'city_9_sanfrancisco_3p5_lwm', 
+        'city_10_austin_3p5_lwm', 
+        'city_11_santaclara_3p5_lwm', 
+        'city_12_fortworth_3p5_lwm', 
+        'city_13_columbus_3p5_lwm', 
+        'city_14_charlotte_3p5_lwm',
+        'city_15_indianapolis_3p5_lwm',
+        'city_16_sanfrancisco_3p5_lwm',  
+        'city_17_seattle_3p5_lwm', 
+        'city_18_denver_3p5_lwm', 
+        'city_19_oklahoma_3p5_lwm', 
+        'asu_campus_3p5',
+        'o1_3p5',
+        'boston5G_3p5'
+    ])
+    return scen_list  
+
+# =============================================================================
+# 3. SCENARIO PROPERTIES MAPPING
+#    - Map each scenario name to its corresponding rows, antenna count, and subcarrier count
+# =============================================================================
+ 
+def scenario_prop():
+    row_column_users = {
+    'city_0_newyork_3p5_lwm': {
+        'n_rows': 109,
+        'n_per_row': 291,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 32
+    },
+    'city_1_losangeles_3p5_lwm': {
+        'n_rows': 142,
+        'n_per_row': 201,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 64
+    },
+    'city_2_chicago_3p5_lwm': {
+        'n_rows': 139,
+        'n_per_row': 200,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 128
+    },
+    'city_3_houston_3p5_lwm': {
+        'n_rows': 154,
+        'n_per_row': 202,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 256
+    },
+    'city_4_phoenix_3p5_lwm': {
+        'n_rows': 198,
+        'n_per_row': 214,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 512
+    },
+    'city_5_philadelphia_3p5_lwm': {
+        'n_rows': 239,
+        'n_per_row': 164,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 1024
+    },
+    'city_6_miami_3p5_lwm': {
+        'n_rows': 199,
+        'n_per_row': 216,
+        'grid_idx': 1,
+        'n_ant_bs': 16,
+        'n_subcarriers': 32
+    },
+    'city_7_sandiego_3p5_lwm': {
+        'n_rows': 176,
+        'n_per_row': 207,
+        'grid_idx': 1,
+        'n_ant_bs': 16,
+        'n_subcarriers': 64
+    },
+    'city_8_dallas_3p5_lwm': {
+        'n_rows': 207,
+        'n_per_row': 190,
+        'grid_idx': 1,
+        'n_ant_bs': 16,
+        'n_subcarriers': 128
+    },
+    'city_9_sanfrancisco_3p5_lwm': {
+        'n_rows': 196,
+        'n_per_row': 206,
+        'grid_idx': 1,
+        'n_ant_bs': 16,
+        'n_subcarriers': 256
+    },
+    'city_10_austin_3p5_lwm': {
+        'n_rows': 255,
+        'n_per_row': 137,
+        'grid_idx': 1,
+        'n_ant_bs': 16,
+        'n_subcarriers': 512
+    },
+    'city_11_santaclara_3p5_lwm': {
+        'n_rows': 117,
+        'n_per_row': 285,
+        'grid_idx': 1,
+        'n_ant_bs': 32,
+        'n_subcarriers': 32
+    },
+    'city_12_fortworth_3p5_lwm': {
+        'n_rows': 214,
+        'n_per_row': 179,
+        'grid_idx': 1,
+        'n_ant_bs': 32,
+        'n_subcarriers': 64
+    },
+    'city_13_columbus_3p5_lwm': {
+        'n_rows': 178,
+        'n_per_row': 240,
+        'grid_idx': 1,
+        'n_ant_bs': 32,
+        'n_subcarriers': 128
+    },
+    'city_14_charlotte_3p5_lwm': {
+        'n_rows': 216,
+        'n_per_row': 177,
+        'grid_idx': 1,
+        'n_ant_bs': 32,
+        'n_subcarriers': 256
+    },
+    'city_15_indianapolis_3p5_lwm': {
+        'n_rows': 200,
+        'n_per_row': 196,
+        'grid_idx': 1,
+        'n_ant_bs': 64,
+        'n_subcarriers': 32
+    },
+    'city_16_sanfrancisco_3p5_lwm': {
+        'n_rows': 201,
+        'n_per_row': 208,
+        'grid_idx': 1,
+        'n_ant_bs': 64,
+        'n_subcarriers': 64
+    },
+    'city_17_seattle_3p5_lwm': {
+        'n_rows': 185,
+        'n_per_row': 205,
+        'grid_idx': 1,
+        'n_ant_bs': 64,
+        'n_subcarriers': 128
+    },
+    'city_18_denver_3p5_lwm': {
+        'n_rows': 212,
+        'n_per_row': 204,
+        'grid_idx': 1,
+        'n_ant_bs': 128,
+        'n_subcarriers': 32
+    },
+    'city_19_oklahoma_3p5_lwm': {
+        'n_rows': 204,
+        'n_per_row': 188,
+        'grid_idx': 1,
+        'n_ant_bs': 128,
+        'n_subcarriers': 64
+    },
+    'asu_campus_3p5_v1': {
+        'n_rows': [0, 1*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 32
+    },
+    'asu_campus_3p5_v2': {
+        'n_rows': [1*int(321/20), 2*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 64
+    },
+    'asu_campus_3p5_v3': {
+        'n_rows': [2*int(321/20), 3*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 128
+    },
+    'asu_campus_3p5_v4': {
+        'n_rows': [3*int(321/20), 4*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 256
+    },
+    'asu_campus_3p5_v5': {
+        'n_rows': [4*int(321/20), 5*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 512
+    },
+    'asu_campus_3p5_v6': {
+        'n_rows': [5*int(321/20), 6*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 1024
+    },
+    'asu_campus_3p5_v7': {
+        'n_rows': [6*int(321/20), 7*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 16,
+        'n_subcarriers': 32
+    },
+    'asu_campus_3p5_v8': {
+        'n_rows': [7*int(321/20), 8*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs':16,
+        'n_subcarriers': 64
+    },
+    'asu_campus_3p5_v9': {
+        'n_rows': [8*int(321/20), 9*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 16,
+        'n_subcarriers': 128
+    },
+    'asu_campus_3p5_v10': {
+        'n_rows': [9*int(321/20), 10*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 16,
+        'n_subcarriers': 256
+    },
+    'asu_campus_3p5_v11': {
+        'n_rows': [10*int(321/20), 11*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 16,
+        'n_subcarriers': 512
+    },
+    'asu_campus_3p5_v12': {
+        'n_rows': [11*int(321/20), 12*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 32,
+        'n_subcarriers': 32
+    },
+    'asu_campus_3p5_v13': {
+        'n_rows': [12*int(321/20), 13*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 32,
+        'n_subcarriers': 64
+    },
+    'asu_campus_3p5_v14': {
+        'n_rows': [13*int(321/20), 14*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 32,
+        'n_subcarriers': 128
+    },
+    'asu_campus_3p5_v15': {
+        'n_rows': [14*int(321/20), 15*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 32,
+        'n_subcarriers': 256
+    },
+    'asu_campus_3p5_v16': {
+        'n_rows': [15*int(321/20), 16*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 64,
+        'n_subcarriers': 32
+    },
+    'asu_campus_3p5_v17': {
+        'n_rows': [16*int(321/20), 17*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 64,
+        'n_subcarriers': 64 
+    },
+    'asu_campus_3p5_v18': {
+        'n_rows': [17*int(321/20), 18*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 64,
+        'n_subcarriers': 128
+    },
+    'asu_campus_3p5_v19': {
+        'n_rows': [18*int(321/20), 19*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 128,
+        'n_subcarriers': 32
+    },
+    'asu_campus_3p5_v20': {
+        'n_rows': [19*int(321/20), 20*int(321/20)],
+        'n_per_row': 411,
+        'grid_idx': 1,
+        'n_ant_bs': 128,
+        'n_subcarriers': 64
+    },
+    'boston5G_3p5_v1': {
+        'n_rows': [812 + 0, 812 + 1*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 8,
+        'n_subcarriers': 32
+    },
+    'boston5G_3p5_v2': {
+        'n_rows': [812 + 1*int((1622-812)/20), 812 + 2*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 8,
+        'n_subcarriers': 64
+    },
+    'boston5G_3p5_v3': {
+        'n_rows': [812 + 2*int((1622-812)/20), 812 + 3*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 8,
+        'n_subcarriers': 128
+    },
+    'boston5G_3p5_v4': {
+        'n_rows': [812 + 3*int((1622-812)/20), 812 + 4*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 8,
+        'n_subcarriers': 256
+    },
+    'boston5G_3p5_v5': {
+        'n_rows': [812 + 4*int((1622-812)/20), 812 + 5*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 8,
+        'n_subcarriers': 512
+    },
+    'boston5G_3p5_v6': {
+        'n_rows': [812 + 5*int((1622-812)/20), 812 + 6*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 8,
+        'n_subcarriers': 1024
+    },
+    'boston5G_3p5_v7': {
+        'n_rows': [812 + 6*int((1622-812)/20), 812 + 7*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 16,
+        'n_subcarriers': 32
+    },
+    'boston5G_3p5_v8': {
+        'n_rows': [812 + 7*int((1622-812)/20), 812 + 8*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs':16,
+        'n_subcarriers': 64
+    },
+    'boston5G_3p5_v9': {
+        'n_rows': [812 + 8*int((1622-812)/20), 812 + 9*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 16,
+        'n_subcarriers': 128
+    },
+    'boston5G_3p5_v10': {
+        'n_rows': [812 + 9*int((1622-812)/20), 812 + 10*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 16,
+        'n_subcarriers': 256
+    },
+    'boston5G_3p5_v11': {
+        'n_rows': [812 + 10*int((1622-812)/20), 812 + 11*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 16,
+        'n_subcarriers': 512
+    },
+    'boston5G_3p5_v12': {
+        'n_rows': [812 + 11*int((1622-812)/20), 812 + 12*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 32,
+        'n_subcarriers': 32
+    },
+    'boston5G_3p5_v13': {
+        'n_rows': [812 + 12*int((1622-812)/20), 812 + 13*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 32,
+        'n_subcarriers': 64
+    },
+    'boston5G_3p5_v14': {
+        'n_rows': [812 + 13*int((1622-812)/20), 812 + 14*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 32,
+        'n_subcarriers': 128
+    },
+    'boston5G_3p5_v15': {
+        'n_rows': [812 + 14*int((1622-812)/20), 812 + 15*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 32,
+        'n_subcarriers': 256
+    },
+    'boston5G_3p5_v16': {
+        'n_rows': [812 + 15*int((1622-812)/20), 812 + 16*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 64,
+        'n_subcarriers': 32
+    },
+    'boston5G_3p5_v17': {
+        'n_rows': [812 + 16*int((1622-812)/20), 812 + 17*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 64,
+        'n_subcarriers': 64 
+    },
+    'boston5G_3p5_v18': {
+        'n_rows': [812 + 17*int((1622-812)/20), 812 + 18*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 64,
+        'n_subcarriers': 128
+    },
+    'boston5G_3p5_v19': {
+        'n_rows': [812 + 18*int((1622-812)/20), 812 + 19*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 128,
+        'n_subcarriers': 32
+    },
+    'boston5G_3p5_v20': {
+        'n_rows': [812 + 19*int((1622-812)/20), 812 + 20*int((1622-812)/20)],
+        'n_per_row': 595,
+        'grid_idx': 2,
+        'n_ant_bs': 128,
+        'n_subcarriers': 64
+    },
+    'o1_3p5_v1': {
+        'n_rows': [0*int(3852/12), 1*int(3852/12)],
+        'n_per_row': 181,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 32
+    },
+    'o1_3p5_v2': {
+        'n_rows': [1*int(3852/12), 2*int(3852/12)],
+        'n_per_row': 181,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 64
+    },
+    'o1_3p5_v3': {
+        'n_rows': [2*int(3852/12), 3*int(3852/12)],
+        'n_per_row': 181,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 128
+    },
+    'o1_3p5_v4': {
+        'n_rows': [3*int(3852/12), 4*int(3852/12)],
+        'n_per_row': 181,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 256
+    },
+    'o1_3p5_v5': {
+        'n_rows': [4*int(3852/12), 5*int(3852/12)],
+        'n_per_row': 181,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 512
+    },
+    'o1_3p5_v6': {
+        'n_rows': [5*int(3852/12), 6*int(3852/12)],
+        'n_per_row': 181,
+        'grid_idx': 1,
+        'n_ant_bs': 8,
+        'n_subcarriers': 1024
+    },
+    'o1_3p5_v7': {
+        'n_rows': [6*int(3852/12), 7*int(3852/12)],
+        'n_per_row': 181,
+        'grid_idx': 1,
+        'n_ant_bs': 16,
+        'n_subcarriers': 32
+    },
+    'o1_3p5_v8': {
+        'n_rows': [7*int(3852/12), 8*int(3852/12)],
+        'n_per_row': 181,
+        'grid_idx': 1,
+        'n_ant_bs': 16,
+        'n_subcarriers': 64
+    },
+    'o1_3p5_v9': {
+        'n_rows': [8*int(3852/12), 2750],
+        'n_per_row': 181,
+        'grid_idx': 1,
+        'n_ant_bs': 16,
+        'n_subcarriers': 128
+    },
+    'o1_3p5_v10': {
+        'n_rows': [2751, 10*int(3852/12)],
+        'n_per_row': 181,
+        'grid_idx': 2,
+        'n_ant_bs': 16,
+        'n_subcarriers': 256
+    },
+    'o1_3p5_v11': {
+        'n_rows': [10*int(3852/12), 11*int(3852/12)],
+        'n_per_row': 181,
+        'grid_idx': 2,
+        'n_ant_bs': 16,
+        'n_subcarriers': 512
+    },
+    'o1_3p5_v12': {
+        'n_rows': [11*int(3852/12), 3851],
+        'n_per_row': 181,
+        'grid_idx': 2,
+        'n_ant_bs': 32,
+        'n_subcarriers': 32
+    },
+    'o1_3p5_v13': {
+        'n_rows': [3852, 12*int(3852/12)+1*int(1351/10)],
+        'n_per_row': 361,
+        'grid_idx': 3,
+        'n_ant_bs': 32,
+        'n_subcarriers': 64
+    },
+    'o1_3p5_v14': {
+        'n_rows': [12*int(3852/12)+1*int(1351/10), 12*int(3852/12)+2*int(1351/10)],
+        'n_per_row': 181,
+        'grid_idx': 3,
+        'n_ant_bs': 32,
+        'n_subcarriers': 128
+    },
+    'o1_3p5_v15': {
+        'n_rows': [12*int(3852/12)+2*int(1351/10), 12*int(3852/12)+3*int(1351/10)],
+        'n_per_row': 181,
+        'grid_idx': 3,
+        'n_ant_bs': 32,
+        'n_subcarriers': 256
+    },
+    'o1_3p5_v16': {
+        'n_rows': [12*int(3852/12)+3*int(1351/10), 12*int(3852/12)+4*int(1351/10)],
+        'n_per_row': 181,
+        'grid_idx': 3,
+        'n_ant_bs': 64,
+        'n_subcarriers': 32
+    },
+    'o1_3p5_v17': {
+        'n_rows': [12*int(3852/12)+4*int(1351/10), 12*int(3852/12)+5*int(1351/10)],
+        'n_per_row': 181,
+        'grid_idx': 3,
+        'n_ant_bs': 64,
+        'n_subcarriers': 64
+    },
+    'o1_3p5_v18': {
+        'n_rows': [12*int(3852/12)+5*int(1351/10), 12*int(3852/12)+6*int(1351/10)],
+        'n_per_row': 181,
+        'grid_idx': 3,
+        'n_ant_bs': 64,
+        'n_subcarriers': 128
+    },
+    'o1_3p5_v19': {
+        'n_rows': [12*int(3852/12)+6*int(1351/10), 12*int(3852/12)+7*int(1351/10)],
+        'n_per_row': 181,
+        'grid_idx': 3,
+        'n_ant_bs': 128,
+        'n_subcarriers': 32
+    },
+    'o1_3p5_v20': {
+        'n_rows': [12*int(3852/12)+7*int(1351/10), 12*int(3852/12)+8*int(1351/10)],
+        'n_per_row': 181,
+        'grid_idx': 3,
+        'n_ant_bs': 128,
+        'n_subcarriers': 64
+    }}
+    return row_column_users   
+
+# =============================================================================
+# 4. TRAINING PARAMETERS AND HYPERPARAMETERS
+#    - Set training epochs, batch sizes, learning rates, model dimensions, etc.
+# =============================================================================
+
+EPOCHS = 50
+BATCH_SIZE = 128
+VAL_BATCH_SIZE = 64
+WARMUP_EPOCHS = 5
+BASE_LR = 5e-4
+MIN_LR = 1e-8
+N_ROWS = 4
+N_COLUMNS = 4
+ELEMENT_LENGTH = N_ROWS * N_COLUMNS * 2
+D_MODEL = 128
+MAX_LEN = 513
+N_LAYERS = 12
+N_ANT_BS = 64
+N_SUBCARRIERS = 64
+device_idx = 0
+WEIGHT_DECAY = 0.05
+BETA1 = 0.9
+BETA2 = 0.999
+MASK_PERCENT = 0.4
+N_HEADS = 8
+DROPOUT = 0.1
+task = ["LosNlosClassification", 
+        "BeamPrediction", 
+        "ChannelInterpolation", 
+        "ChannelEstimation", 
+        "ChannelCharting",
+        None][-1]
+
+# =============================================================================
+# 5. DATA GENERATION LOOP
+#    - Iterate over scenarios and base station indices to generate channel samples and labels
+#    - Handle both full-scenario and zoned sub-scenarios for campus and Boston data
+# =============================================================================
+
+scenarios = scenarios_list()
+
+channels = []
+labels = []
+scenario_properties = scenario_prop()
+preprocessed_data = []
+ 
+for scenario in scenarios[:-3]: 
+    for bs_idx in range (1,4):
+        scenario_channels, scenario_labels = generate_channels_and_labels(
+            n_ant_bs=scenario_properties[scenario]["n_ant_bs"],
+            n_subcarriers=scenario_properties[scenario]["n_subcarriers"],
+            bs_idx=bs_idx,
+            scenario_name=scenario,
+            task=task,
+            n_beams=64
+        )
+        labels.extend(scenario_labels)
+        channels.append(scenario_channels)
+
+bs_idxs = [[1], [4, 15], [2]]
+for scenario_idx, scenario in enumerate(scenarios[-3:]):  
+    for bs_idx in bs_idxs[scenario_idx]:
+        for zone in range (20):
+            row_start = scenario_properties[scenario+f"_v{zone+1}"]["n_rows"][0]
+            row_end = scenario_properties[scenario+f"_v{zone+1}"]["n_rows"][1]
+            grid_idx = scenario_properties[scenario+f"_v{zone+1}"]["grid_idx"]-1
+            scenario_channels, scenario_labels = generate_channels_and_labels(
+                n_ant_bs=scenario_properties[scenario+f"_v{zone+1}"]["n_ant_bs"],
+                n_subcarriers=scenario_properties[scenario+f"_v{zone+1}"]["n_subcarriers"],
+                grid_idx=grid_idx,
+                bs_idx=bs_idx,
+                scenario_name=scenario,
+                rows=np.arange(row_start, row_end),
+                task=task,
+                n_beams=64
+            )
+            
+            if scenario_channels.numel() == 0:
+                print(f"No candidate user in zone {zone} for scenario {scenario} has a path to bs_idx {bs_idx} (All channels are zero)")
+                continue
+            
+            labels.extend(scenario_labels)
+            channels.append(scenario_channels)
+
+# =============================================================================
+# 6. DATA TOKENIZATION
+#    - Tokenize channel matrices into input sequences with masking for pretraining
+# =============================================================================
+
+preprocessed_data = tokenizer_train(
+    channels,
+    max_len=MAX_LEN,
+    masking_percent=MASK_PERCENT,
+    mask=True,
+    seed=42
+)
+
+# =============================================================================
+# 7. TRAIN/VALIDATION/TEST SPLIT
+#    - Split each tokenized dataset into train, validation, and test subsets with a fixed random seed
+# =============================================================================
+
+SEED = 42
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+train_ratio = 0.8
+val_ratio = 0.2
+train_data = {}
+val_data = {}
+test_data = {}
+
+for key, samples in preprocessed_data.items():
+    print(f"key: {key}")
+    total_samples = len(samples)
+    train_size = int(train_ratio * total_samples)
+    val_size = int(val_ratio * total_samples)
+    test_size = total_samples - train_size - val_size
+    
+    train_data[key], val_data[key], test_data[key] = random_split(
+        samples, [train_size, val_size, test_size]
+    )
+
+# =============================================================================
+# 8. DATALOADER CREATION
+#    - Build PyTorch DataLoader objects for batched training and validation
+# =============================================================================
+
+train_loaders = create_train_dataloader(train_data, batch_size=BATCH_SIZE, shuffle=True)
+val_loaders = create_train_dataloader(val_data, batch_size=VAL_BATCH_SIZE, shuffle=False)
+
+# =============================================================================
+# 9. MODEL INITIALIZATION
+#    - Instantiate the LWM transformer model and optionally load pre-trained weights
+#    - Wrap with DataParallel for multi-GPU support
+# =============================================================================
+
+gpu_ids = [1] # device_idx
+device = torch.device(f"cuda:{gpu_ids[0]}" if torch.cuda.is_available() else "cpu")
+model = pretrained_model.lwm(
+    element_length=ELEMENT_LENGTH,
+    d_model=D_MODEL,
+    n_layers=N_LAYERS,
+    max_len=MAX_LEN,
+    n_heads=N_HEADS,
+    dropout=DROPOUT
+).to(device)
+
+# Optional: Load pre-trained model
+load_model = False
+if load_model:
+    model.load_state_dict(torch.load("models/model_checkpoint.pth", map_location=device))
+    print("Pre-trained model loaded successfully.")
+
+# Use DataParallel for multi-GPU support
+model = nn.DataParallel(model, device_ids=gpu_ids)
+print(f"Model loaded successfully on GPU {device.index}")
+n_parameters = count_parameters(model)
+print(f"Number of trainable parameters: {n_parameters:,}")
+
+# =============================================================================
+# 10. OPTIMIZER AND LEARNING RATE SCHEDULER
+#     - Configure AdamW optimizer and a cosine-with-warmup LR schedule based on total steps
+# =============================================================================
+
+TOTAL_STEPS = sum(len(loader) for loader in train_loaders.values()) * EPOCHS
+WARMUP_STEPS = sum(len(loader) for loader in train_loaders.values()) * WARMUP_EPOCHS
+
+optimizer = AdamW(
+    model.parameters(),
+    lr=BASE_LR,
+    betas=(BETA1, BETA2),
+    weight_decay=WEIGHT_DECAY
+)
+
+def lr_lambda(current_step):
+    if current_step < WARMUP_STEPS:
+        return current_step / WARMUP_STEPS
+    else:
+        scaled_progress = (current_step - WARMUP_STEPS) / (TOTAL_STEPS - WARMUP_STEPS)
+        cosine_decay = 0.5 * (1 + np.cos(np.pi * scaled_progress))
+        return cosine_decay * (BASE_LR - MIN_LR) / BASE_LR + MIN_LR / BASE_LR
+
+scheduler = LambdaLR(optimizer, lr_lambda=lr_lambda)
+
+# =============================================================================
+# 11. PRE-TRAINING LOOP
+#     - Call the train_lwm utility to run the pre-training epochs, logging metrics and saving models
+# =============================================================================
+
+pretrained_model = train_lwm(
+    model,
+    train_loaders,
+    val_loaders,
+    optimizer,
+    scheduler,
+    EPOCHS,
+    device=device,
+    save_dir="pretrained_models",
+    log_file="training_log.csv"
+)

utils.py

@@ -0,0 +1,1119 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import os
+from tqdm import tqdm
+import pickle
+import deepmimo as dm
+from collections import defaultdict
+from torch.utils.data import TensorDataset, DataLoader
+import matplotlib.pyplot as plt
+from sklearn.decomposition import PCA
+from sklearn.manifold import TSNE
+import umap
+from math import pi
+import zipfile
+import shutil
+
+def generate_channels_and_labels(
+    n_ant_bs=16,
+    n_subcarriers=64,
+    bs_idx=3,
+    grid_idx=0,
+    scenario_name=0,
+    scenario_idx=None,
+    rows=None,
+    task="LosNlosClassification",
+    n_beams=64
+    ):
+    """
+    Generate wireless channel samples and task-specific labels for pre-training or evaluation.
+
+    Args:
+        n_ant_bs (int): Number of antennas at the base station. Defaults to 16.
+        n_subcarriers (int): Number of subcarriers per channel. Defaults to 64.
+        bs_idx (int): Index of the base station to generate data for. Defaults to 3.
+        scenario_idx (int): Index of the scenario to select from available scenarios. Defaults to 0.
+        task (str): Task for which to generate labels. Defaults to "LosNlosClassification".
+        n_beams (int): Number of beams for beam-related tasks. Defaults to 64.
+
+    Returns:
+        tuple: A tuple containing:
+            - channels (list of torch.Tensor): List of generated channel tensors.
+            - labels (list): Corresponding task-specific labels.
+    """
+    if scenario_idx:
+        scenario_name = dm.search({})[scenario_idx]
+        
+    channels, labels = dataset_generator(
+        n_ant_bs=n_ant_bs,
+        n_subcarriers=n_subcarriers,
+        grid_idx=grid_idx,
+        scenario_name=scenario_name,
+        rows=rows,
+        bs_idx=bs_idx,
+        task=task,
+        n_beams=n_beams
+    )
+
+    return channels, labels
+
+def dataset_generator(
+        n_ant_bs=32, 
+        n_subcarriers=32, 
+        scenario_name="city_0_newyork_3p5", 
+        grid_idx=0,
+        rows=None,
+        bs_idx=1, 
+        save_dir="data", 
+        task="LosNlosClassification", 
+        n_beams=64, 
+        snr=None, 
+        seed=42
+        ):
+    """
+    Generate wireless channel data and task-specific labels using DeepMIMO dataset.
+
+    Args:
+        n_ant_bs (int): Number of antennas at the base station. Defaults to 32.
+        n_subcarriers (int): Number of subcarriers per channel. Defaults to 32.
+        scenario_name (str): Name of the scenario for data generation. Defaults to "city_0_newyork_3p5".
+        bs_idx (int): Index of the base station to generate data for. Defaults to 1.
+        save_dir (str): Directory to save generated data. Defaults to "data".
+        task (str): Task for which to generate labels. Defaults to "LosNlosClassification".
+        n_beams (int): Number of beams for beam-related tasks. Defaults to 64.
+        snr (float, optional): Signal-to-noise ratio for adding Gaussian noise in robust beamforming task.
+            Defaults to None.
+        seed (int): Random seed for reproducibility in noise generation. Defaults to 42.
+
+    Returns:
+        tuple: A tuple containing:
+            - cleaned_deepmimo_data (torch.Tensor): Cleaned channel data tensor.
+            - labels (torch.Tensor or list): Task-specific labels or NaN-filled tensor if task is None.
+    """ 
+    os.makedirs(save_dir, exist_ok=True)
+    print(f"\nGenerating data for scenario: {scenario_name}, BS #{bs_idx}")
+    deepmimo_data = DeepMIMO_data_gen(scenario_name, 
+                                      n_ant_bs, 
+                                      1, n_subcarriers, 
+                                      bs_idx=bs_idx, 
+                                      row_indices=rows,
+                                      grid_idx=grid_idx)
+    if task is not None:
+        labels = label_gen(deepmimo_data, task, scenario_name, bs_idx=bs_idx, n_beams=n_beams)
+    else:
+        n_channels = len(np.where(deepmimo_data.los != -1)[0])
+        labels = np.nan * torch.ones(n_channels) 
+    cleaned_deepmimo_data = deepmimo_data_cleaning(deepmimo_data)
+    if snr is not None and task in["ChannelEstimation", "LosNlosClassification"]:
+        cleaned_deepmimo_data = generate_gaussian_noise(cleaned_deepmimo_data, snr, seed=seed)
+    return cleaned_deepmimo_data.squeeze(1), labels
+
+def label_gen(data, task, scenario_name, bs_idx=1, n_beams=64):
+    """
+    Generate task-specific labels for wireless channel data.
+
+    Args:
+        data (object): DeepMIMO data object containing channel, LOS status, and receiver positions.
+        task (str): Task for which to generate labels. Options: 'LosNlosClassification',
+            'BeamPrediction', 'ChannelCharting', 'ChannelEstimation', 'ChannelInterpolation'.
+        scenario_name (str): Name of the scenario for data generation.
+        bs_idx (int): Index of the base station. Defaults to 1.
+        n_beams (int): Number of beams for beam-related tasks. Defaults to 64.
+
+    Returns:
+        torch.Tensor: Task-specific labels for the valid data indices.
+    """
+    labels = 0
+    idxs = np.where(data.los != -1)[0]
+
+    if len(idxs) == 0: # Users with no path (do not need them for pre-training). You should set the task to None in train_lwm.py
+        labels = np.full(data.n_ue, np.nan, dtype=float)
+
+    else:        
+        if task == 'LosNlosClassification':
+            
+            labels = data.los[idxs].astype(int)
+            dm.plot_coverage(data.rx_pos[idxs], data.los[idxs], cbar_title='LoS status')
+            
+        elif task == 'BeamPrediction':
+            
+            parameters = get_parameters()
+            n_users = len(data.channel)
+            n_subbands = 1
+            fov = 180
+
+            # Setup Beamformers
+            beam_angles = np.around(np.arange(-fov/2, fov/2+.1, fov/(n_beams-1)), 2)
+
+            F1 = np.array([dm.steering_vec(data.ch_params.bs_antenna.shape, phi=azi).squeeze()
+                for azi in beam_angles])
+
+            full_dbm = np.zeros((n_beams, n_subbands, n_users), dtype=float)
+            for ue_idx in tqdm(range(n_users), desc='Computing the channel for each user'):
+                if data.los[ue_idx] == -1:
+                    full_dbm[:,:,ue_idx] = np.nan
+                else:
+                    chs = F1 @ data.channel[ue_idx]
+                    full_linear = np.abs(np.mean(chs.squeeze().reshape((n_beams, n_subbands, -1)), axis=-1))
+                    full_dbm[:,:,ue_idx] = np.around(20*np.log10(full_linear) + 30, 1)
+
+            best_beams = np.argmax(np.mean(full_dbm,axis=1), axis=0)
+            best_beams = best_beams.astype(float)
+            best_beams[np.isnan(full_dbm[0,0,:])] = np.nan
+            
+            dm.plot_coverage(data.rx_pos[idxs], best_beams[idxs], bs_pos=data.tx_pos, 
+                            bs_ori=parameters.bs_antenna.rotation*np.pi/180, 
+                            cbar_title='Best beam index')
+            
+            labels = best_beams[idxs].astype(int)
+            
+        elif task == 'ChannelCharting':
+            
+            labels = torch.tensor(data.rx_pos[:,:2][idxs]).to(dtype=torch.float32)
+            
+        elif task == 'ChannelEstimation':
+            
+            channels = torch.tensor(data.channel[idxs]*1e6, dtype=torch.complex64).squeeze(1)
+            labels = torch.stack((channels.real, channels.imag), dim=1) 
+            
+        elif task == 'ChannelInterpolation':
+            
+            channels = torch.tensor(data.channel[idxs]*1e6, dtype=torch.complex64).squeeze(1)
+            labels = torch.stack((channels.real, channels.imag), dim=1)
+    
+        labels = torch.tensor(labels)
+    
+    return labels        
+
+def generate_gaussian_noise(data, snr_db, seed=42):
+    """
+    Add complex Gaussian noise to channel data based on a specified signal-to-noise ratio (SNR).
+
+    Args:
+        data (torch.Tensor): Input complex-valued channel data with shape (n_samples, 1, n_ant, n_sc).
+        snr_db (float): Signal-to-noise ratio in decibels.
+        seed (int): Random seed for reproducibility of noise generation. Defaults to 42.
+
+    Returns:
+        torch.Tensor: Noisy channel data with the same shape as the input, with added complex Gaussian noise.
+    """
+    torch.manual_seed(seed) 
+    data = data.squeeze(1)  # Shape: (n_samples, n_ant, n_sc)
+    flat_data = data.view(data.size(0), -1)  
+    
+    # Compute signal power
+    signal_power = torch.mean(flat_data.abs() ** 2, dim=1, keepdim=True)  
+    snr_linear = 10 ** (snr_db / 10)  
+    noise_power = signal_power / snr_linear  
+
+    # Generate noise
+    noise_real = torch.randn_like(flat_data.real) * torch.sqrt(noise_power / 2)
+    noise_imag = torch.randn_like(flat_data.imag) * torch.sqrt(noise_power / 2)
+    noise = torch.complex(noise_real, noise_imag) 
+
+    # Reshape noise and add to data
+    noise = noise.view_as(data)
+    noisy_data = data + noise
+    noisy_data = noisy_data.unsqueeze(1) 
+
+    return noisy_data
+
+# REMOVE ZERO CHANNELS AND SCALE
+def deepmimo_data_cleaning(deepmimo_data):
+    """
+    Clean DeepMIMO channel data by removing invalid channels and scaling the valid ones.
+
+    Args:
+        deepmimo_data (object): DeepMIMO data object containing channel data and LOS status.
+
+    Returns:
+        torch.Tensor: Cleaned and scaled channel data as a complex-valued tensor with dtype torch.complex64.
+    """
+    idxs = np.where(deepmimo_data.los != -1)[0]
+    cleaned_deepmimo_data = deepmimo_data.channel[idxs]
+    return torch.tensor(cleaned_deepmimo_data, dtype=torch.complex64) * 1e6
+
+def manual_unzip_scenario(scenario_name):
+    """Manually unzip a downloaded scenario to avoid DeepMIMO corruption issues."""
+    scenarios_dir = os.path.join(os.getcwd(), "deepmimo_scenarios")
+    zip_path = os.path.join(scenarios_dir, f"{scenario_name}_downloaded.zip")
+    scenario_path = os.path.join(scenarios_dir, scenario_name)
+    
+    # Remove existing unzipped folder if it exists
+    if os.path.exists(scenario_path):
+        print(f"Removing existing scenario folder: {scenario_path}")
+        shutil.rmtree(scenario_path)
+    
+    # Manually unzip the downloaded file
+    if os.path.exists(zip_path):
+        print(f"Manually unzipping: {zip_path}")
+        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+            zip_ref.extractall(scenarios_dir)
+        print(f"Successfully extracted to: {scenario_path}")
+        return True
+    else:
+        print(f"Zip file not found: {zip_path}")
+        return False   
+
+# Data Generation 
+def DeepMIMO_data_gen(scenario, num_ant_hor, num_ant_vert, n_subcarriers, bs_idx, row_indices, grid_idx):
+    """
+    Generates wireless channel data for a specified DeepMIMO scenario and base station configuration.
+
+    This function downloads and loads the specified DeepMIMO scenario, filters the users (if row_indices are given),
+    and computes the MIMO channels using the provided antenna and subcarrier configuration. It also plots the
+    line-of-sight (LoS) coverage map for visualization.
+
+    Args:
+        scenario (str): Name of the DeepMIMO scenario (e.g., 'o1_3p5').
+        num_ant_hor (int): Number of horizontal antennas at the base station (for the UPA configuration).
+        num_ant_vert (int): Number of vertical antennas at the base station.
+        n_subcarriers (int): Number of subcarriers to simulate per user.
+        bs_idx (int): Index of the base station to extract channels from.
+        row_indices (list or None): List of user row indices to subset the dataset. If None, all users are included.
+
+    Returns:
+        object: A DeepMIMO data object containing:
+            - Computed channel matrices for selected users and base station.
+            - User metadata including location, LoS status, and other parameters.
+            - Scenario-specific attributes for analysis and visualization.
+    """
+    parameters = get_parameters(num_ant_hor, num_ant_vert, n_subcarriers)
+    dm.download(scenario)
+
+    # Manual unzip for LWM scenarios to avoid corruption
+    if "lwm" in scenario:
+        manual_unzip_scenario(scenario)
+
+    data = dm.load(scenario, tx_sets=[bs_idx], rx_sets=[grid_idx])
+    
+    if row_indices is not None:
+        if grid_idx == 0:
+            row_idxs = data.get_row_idxs(row_indices) # example: row_indices = np.arange(40,60)
+            data = data.subset(row_idxs)
+        elif grid_idx == 1:
+            if scenario == "o1_3p5":
+                col_idxs = data.get_col_idxs(row_indices-2751)
+                data = data.subset(col_idxs)
+            elif scenario == "boston5G_3p5":
+                row_idxs = data.get_row_idxs(row_indices-812)
+                data = data.subset(row_idxs)
+        elif grid_idx == 2:
+            if scenario == "o1_3p5":
+                col_idxs = data.get_col_idxs(row_indices-3852)
+            data = data.subset(col_idxs)
+    
+    # data.plot_coverage(data.los)
+    data.compute_channels(parameters)
+    
+    return data
+
+def get_parameters(num_ant_hor=32, num_ant_vert=1, n_subcarriers=32):
+    """
+    Generate channel parameters for DeepMIMO dataset generation.
+
+    Args:
+        num_ant_hor (int): Number of horizontal antennas at the base station. Defaults to 32.
+        num_ant_vert (int): Number of vertical antennas at the base station. Defaults to 1.
+        n_subcarriers (int): Number of subcarriers per channel. Defaults to 32.
+        bs_idx (int): Index of the base station. Defaults to 1.
+
+    Returns:
+        dm.ChannelGenParameters: Configured channel parameters object for DeepMIMO data generation.
+    """
+    # Create channel parameters with all options
+    ch_params = dm.ChannelParameters()
+
+    # Antenna parameters
+
+    # Base station antenna parameters
+    ch_params.bs_antenna.rotation = np.array([0, 0, -135])  # [az, el, pol] in degrees
+    ch_params.bs_antenna.fov = np.array([360, 180])      # [az, el] in degrees
+    ch_params.bs_antenna.shape = np.array([num_ant_hor, num_ant_vert])        # [horizontal, vertical] elements
+    ch_params.bs_antenna.spacing = 0.5                   # Element spacing in wavelengths
+
+    # User equipment antenna parameters
+    ch_params.ue_antenna.rotation = np.array([0, 0, 0])  # [az, el, pol] in degrees
+    ch_params.ue_antenna.fov = np.array([360, 180])      # [az, el] in degrees
+    ch_params.ue_antenna.shape = np.array([1, 1])        # [horizontal, vertical] elements
+    ch_params.ue_antenna.spacing = 0.5                   # Element spacing in wavelengths
+    
+    # Channel parameters
+    ch_params.freq_domain = True  # Whether to compute frequency domain channels
+    ch_params.num_paths = 20      # Number of paths
+
+    # OFDM parameters
+    subcarrier_spacing = 30e3                                 
+    ch_params.ofdm.subcarriers = n_subcarriers                       # Number of subcarriers
+    ch_params.ofdm.selected_subcarriers = np.arange(n_subcarriers)   # Which subcarriers to generate
+    ch_params.ofdm.bandwidth = subcarrier_spacing * n_subcarriers    # Bandwidth in Hz
+    ch_params.ofdm.rx_filter = 0     
+    
+    return ch_params
+
+def tokenizer_train(channels,
+                    max_len=513, 
+                    masking_percent=0.40, 
+                    mask=False, 
+                    seed=42):
+    """
+    Tokenize wireless channel data into patches and optionally apply masking.
+
+    Args:
+        channels (torch.Tensor or list): Input channel data to be tokenized.
+        max_len (int): Maximum sequence length for tokenized samples. Defaults to 513.
+        masking_percent (float): Percentage of patches to mask if mask is True. Defaults to 0.40.
+        mask (bool): Whether to apply masking to the tokenized samples. Defaults to False.
+        seed (int): Random seed for reproducibility in masking. Defaults to 42.
+
+    Returns:
+        dict or torch.Tensor: If mask is True, returns a dictionary mapping sequence lengths to lists
+            of tokenized samples. If mask is False, returns a tensor of stacked tokenized samples.
+    """
+    patches = [patch_maker(channel_set, patch_rows=4, patch_cols=4) for channel_set in channels]
+    patches = [patch for patch_list in patches for patch in patch_list]
+    print("\nTotal number of samples:", len(patches))
+    
+    grouped_data = defaultdict(list)  # Group samples by sequence length
+    grouped_data_2 = []
+    
+    for user_idx in tqdm(range(len(patches)), desc="Processing items"):
+        patch_size = patches[user_idx].shape[1]
+        n_patches = patches[user_idx].shape[0]
+        n_masks_half = int(masking_percent * n_patches)
+        
+        word2id = {
+            '[CLS]': 0.2 * np.ones((patch_size)),
+            '[MASK]': 0.1 * np.ones((patch_size))
+        }
+        
+        sample = make_sample(
+            user_idx, patches, word2id, n_patches, n_masks_half, patch_size, mask=mask, seed=seed
+        )
+        
+        if mask:
+            seq_length = len(sample[0]) 
+            grouped_data[seq_length].append(sample)
+        else:
+            grouped_data_2.append(sample)
+    
+    if mask:
+        normalized_grouped_data = {i: grouped_data[key] for i, key in enumerate(sorted(grouped_data.keys()))}
+    else: 
+        normalized_grouped_data = torch.stack(grouped_data_2, dim=0)
+    
+    return normalized_grouped_data 
+
+def tokenizer(channels,
+              max_len=513, 
+              masking_percent=0.40, 
+              mask=False, 
+              seed=42):
+    """
+    Tokenize wireless channel data into patches and optionally apply masking.
+
+    Args:
+        channels (torch.Tensor or list): Input channel data to be tokenized.
+        max_len (int): Maximum sequence length for tokenized samples. Defaults to 513.
+        masking_percent (float): Percentage of patches to mask if mask is True. Defaults to 0.40.
+        mask (bool): Whether to apply masking to the tokenized samples. Defaults to False.
+        seed (int): Random seed for reproducibility in masking. Defaults to 42.
+
+    Returns:
+        dict or torch.Tensor: If mask is True, returns a dictionary mapping sequence lengths to lists
+            of tokenized samples. If mask is False, returns a tensor of stacked tokenized samples.
+    """
+    patches = patch_maker(channels, patch_rows=4, patch_cols=4)
+    print("\nTotal number of samples:", len(patches))
+    
+    grouped_data = defaultdict(list)  # Group samples by sequence length
+    grouped_data_2 = []
+    
+    for user_idx in tqdm(range(len(patches)), desc="Processing items"):
+        patch_size = patches[user_idx].shape[1]
+        n_patches = patches[user_idx].shape[0]
+        n_masks_half = int(masking_percent * n_patches)
+        
+        word2id = {
+            '[CLS]': 0.2 * np.ones((patch_size)),
+            '[MASK]': 0.1 * np.ones((patch_size))
+        }
+        
+        sample = make_sample(
+            user_idx, patches, word2id, n_patches, n_masks_half, patch_size, mask=mask, seed=seed
+        )
+        
+        if mask:
+            seq_length = len(sample[0]) 
+            grouped_data[seq_length].append(sample)
+        else:
+            grouped_data_2.append(sample)
+    
+    if mask:
+        normalized_grouped_data = {i: grouped_data[key] for i, key in enumerate(sorted(grouped_data.keys()))}
+    else: 
+        normalized_grouped_data = torch.stack(grouped_data_2, dim=0)
+    
+    return normalized_grouped_data 
+
+def make_sample(user_idx, patch, word2id, n_patches, n_masks, patch_size, mask=True, seed=None):
+    """
+    Create a tokenized sample from patch data, optionally applying masking for a specific user.
+
+    Args:
+        user_idx (int): Index of the user whose patch data is to be processed.
+        patch (numpy.ndarray or torch.Tensor): Patch data for all users.
+        word2id (dict): Dictionary mapping special tokens ('[CLS]', '[MASK]') to their representations.
+        n_patches (int): Number of patches in the input data.
+        n_masks (int): Number of patches to mask if mask is True.
+        patch_size (int): Size of each patch.
+        mask (bool): Whether to apply masking to the sample. Defaults to True.
+        seed (int, optional): Random seed for reproducibility in masking. Defaults to None.
+
+    Returns:
+        torch.Tensor or list: If mask is False, returns a tensor of input IDs with [CLS] prepended.
+            If mask is True, returns a list containing input IDs, masked tokens, and masked positions.
+    """
+    if seed is not None:
+        np.random.seed(seed)  
+        
+    # Step 1: Retrieve tokens and prepend [CLS]
+    tokens = patch[user_idx]
+    input_ids = np.vstack((word2id['[CLS]'], tokens))
+
+    # Step 2: Mask real and imaginary patches
+    tokens_size = int(n_patches)  # int(n_patches / 2)
+    masked_pos = np.random.choice(range(1, tokens_size), size=n_masks, replace=False)
+
+    masked_tokens = []
+    for pos in masked_pos:
+        original_masked_tokens = input_ids[pos].copy()
+        masked_tokens.append(original_masked_tokens)
+        if mask:
+            rnd_num = np.random.rand()
+            if rnd_num < 0.1:
+                input_ids[pos] = np.random.rand(patch_size)  # Replace with random values
+            elif rnd_num < 0.9:
+                input_ids[pos] = word2id['[MASK]']  # Replace with [MASK]
+    
+    if not mask:
+        return torch.tensor(input_ids)
+    else:
+        return [input_ids, masked_tokens, masked_pos]
+    
+# Patch GENERATION
+def patch_maker(original_ch, patch_rows=4, patch_cols=4):
+    """
+    Converts complex-valued channel matrices into flattened, interleaved real-imaginary patch embeddings.
+
+    This function takes a batch of complex-valued 2D channel matrices (one per sample), splits the real 
+    and imaginary components, interleaves them along the last dimension, and divides the result into
+    non-overlapping patches of specified size. The output is a set of flattened patches per sample,
+    ready for use in models like Transformers.
+
+    Args:
+        original_ch (np.ndarray): Input array of shape (n_samples, n_rows, n_cols) with complex values.
+        patch_rows (int): Number of rows per patch. Default is 4.
+        patch_cols (int): Number of columns per patch. Default is 4.
+
+    Returns:
+        np.ndarray: Array of shape (n_samples, n_patches, patch_rows * patch_cols * 2), where each patch
+                    is flattened and contains interleaved real and imaginary parts.
+    """
+    # Step 1: Remove the singleton channel dimension
+    n_samples, n_rows, n_cols = original_ch.shape  # Unpack shape
+    # original_ch = original_ch[:, 0]  # Remove the singleton dimension
+
+    # Step 2: Split into real and imaginary parts and interleave them
+    flat_real = original_ch.real
+    flat_imag = original_ch.imag
+
+    # Interleave real and imaginary parts along the last axis
+    interleaved = np.empty((n_samples, n_rows, n_cols * 2), dtype=np.float32)
+    interleaved[:, :, 0::2] = flat_real
+    interleaved[:, :, 1::2] = flat_imag
+
+    # Step 3: Compute the number of patches along rows and columns
+    n_patches_rows = int(np.ceil(n_rows / patch_rows))
+    n_patches_cols = int(np.ceil(n_cols / patch_cols))
+
+    # Step 4: Pad the matrix if necessary to make it divisible by patch size
+    padded_rows = n_patches_rows * patch_rows - n_rows
+    padded_cols = n_patches_cols * patch_cols - n_cols
+    if padded_rows > 0 or padded_cols > 0:
+        interleaved = np.pad(
+            interleaved,
+            ((0, 0), (0, padded_rows), (0, padded_cols * 2)),  # Double padding for interleaved axis
+            mode='constant',
+            constant_values=0,
+        )
+
+    # Step 5: Create patches by dividing into blocks
+    n_samples, padded_rows, padded_cols = interleaved.shape
+    padded_cols //= 2  # Adjust for interleaving (real and imaginary parts count as one)
+    patches = []
+
+    for i in range(0, padded_rows, patch_rows):
+        for j in range(0, padded_cols, patch_cols):
+            patch = interleaved[:, i:i + patch_rows, j * 2:(j + patch_cols) * 2]
+            patches.append(patch.reshape(n_samples, -1))  # Flatten each patch
+
+    # Step 6: Stack patches to form the final array
+    patches = np.stack(patches, axis=1)  # Shape: (num_samples, n_patches, patch_rows * patch_cols * 2)
+
+    return patches
+
+def patch_reconstructor(patches, original_rows, original_cols, patch_rows=4, patch_cols=4):
+    """
+    Reconstructs the original channel matrix with real and imaginary parts as separate channels using PyTorch.
+
+    Args:
+        patches (torch.Tensor): Patches of shape (n_samples, n_patches, patch_rows * patch_cols * 2)
+        original_rows (int): Original number of rows (n_rows)
+        original_cols (int): Original number of columns (n_cols)
+        patch_rows (int): Number of rows per patch (default: 4)
+        patch_cols (int): Number of columns per patch (default: 4)
+
+    Returns:
+        torch.Tensor: Reconstructed channel matrix of shape (n_samples, 2, original_rows, original_cols)
+                     where channel 0 is real and channel 1 is imaginary
+    """
+    # Step 1: Extract dimensions
+    n_samples, n_patches, patch_size = patches.shape
+    assert patch_size == patch_rows * patch_cols * 2, "Patch size does not match patch_rows * patch_cols * 2"
+
+    # Step 2: Compute the number of patches along rows and columns
+    # Use integer division since no padding is needed
+    n_patches_rows = original_rows // patch_rows
+    n_patches_cols = original_cols // patch_cols
+    assert n_patches == n_patches_rows * n_patches_cols, "Number of patches does not match expected grid"
+
+    # Step 3: Reshape patches back into 2D blocks
+    patches_2d = patches.reshape(n_samples, n_patches_rows, n_patches_cols, patch_rows, patch_cols * 2)
+
+    # Step 4: Reconstruct the interleaved matrix
+    # No padding, so use original dimensions directly
+    interleaved = torch.zeros((n_samples, original_rows, original_cols * 2), dtype=torch.float32, device=patches.device)
+    for i in range(n_patches_rows):
+        for j in range(n_patches_cols):
+            interleaved[:, i * patch_rows:(i + 1) * patch_rows, j * patch_cols * 2:(j + 1) * patch_cols * 2] = \
+                patches_2d[:, i, j, :, :]
+
+    # Step 5: De-interleave real and imaginary parts
+    flat_real = interleaved[:, :, 0::2]
+    flat_imag = interleaved[:, :, 1::2]
+
+    # Step 6: Stack real and imaginary parts as separate channels along axis=1
+    reconstructed = torch.stack((flat_real, flat_imag), dim=1)  # Shape: (n_samples, 2, original_rows, original_cols)
+
+    return reconstructed
+
+def create_train_dataloader(grouped_data, batch_size, shuffle):
+    """
+    Creates a dictionary of DataLoaders from grouped input data based on sequence lengths.
+
+    This function processes pre-grouped training data where each key in the dictionary corresponds
+    to a specific sequence length, and each value is a list of training samples. It converts the 
+    data into PyTorch tensors and constructs a separate DataLoader for each sequence length.
+
+    Args:
+        grouped_data (dict): A dictionary where keys are sequence lengths (e.g., 1, 2, ..., T), and
+                             values are lists of tuples (input_ids, masked_tokens, masked_pos).
+        batch_size (int): Batch size to use for the DataLoaders.
+        shuffle (bool): Whether to shuffle the data during loading.
+
+    Returns:
+        dict: A dictionary mapping each sequence length to its corresponding DataLoader.
+    """
+    dataloaders = {}
+
+    for seq_length, group in grouped_data.items():
+        
+        print(f"dataloader in progress ...\nkey: {seq_length}")
+        
+        ## Uncomment the following line if you run out of memory during pre-training
+        # batch_size = batch_size // 8 if seq_length >= 5 else batch_size
+        
+        # Unpack samples for the current group
+        input_ids, masked_tokens, masked_pos = zip(*group)
+
+        # Convert to tensors
+        input_ids_tensor = torch.tensor(input_ids, dtype=torch.float32)
+        masked_tokens_tensor = torch.tensor(masked_tokens, dtype=torch.float32)
+        masked_pos_tensor = torch.tensor(masked_pos, dtype=torch.long)
+
+        # Create TensorDataset and DataLoader
+        dataset = TensorDataset(input_ids_tensor, masked_tokens_tensor, masked_pos_tensor)
+        dataloaders[seq_length] = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, pin_memory=True)
+
+    return dataloaders
+
+def count_parameters(model):
+    """
+    Counts the number of trainable parameters in a PyTorch model.
+
+    Args:
+        model (torch.nn.Module): The model to inspect.
+
+    Returns:
+        int: Total number of trainable parameters (i.e., those with requires_grad=True).
+    """
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+def nmse_loss(y_true, y_pred):
+    """Compute Normalized Mean Squared Error (NMSE) using PyTorch.
+    
+    Args:
+        y_true (torch.Tensor): Ground truth tensor.
+        y_pred (torch.Tensor): Predicted tensor.
+        
+    Returns:
+        torch.Tensor: NMSE value.
+    """
+    # Ensure inputs are torch tensors
+    y_true = torch.as_tensor(y_true)
+    y_pred = torch.as_tensor(y_pred)
+    
+    # Compute NMSE: mean((y_true - y_pred)^2) / mean(y_true^2)
+    squared_diff = torch.mean((y_true - y_pred) ** 2)
+    squared_true = torch.mean(y_true ** 2)
+    nmse = squared_diff / squared_true
+    
+    return nmse
+
+def train_lwm(model, train_loaders, val_loaders, optimizer, scheduler, epochs, device, save_dir="models", log_file="training_log.csv"):
+    """
+    Trains the Large Wireless Model (LWM) using masked channel modeling on grouped datasets of various sequence lengths.
+
+    The training alternates between training and evaluation every 2 epochs. For each sequence length, 
+    a separate DataLoader is used. MSE is used for the training objective, and both MSE and NMSE are computed
+    during validation for performance tracking.
+
+    Args:
+        model (torch.nn.Module): The LWM model to train.
+        train_loaders (dict): Dictionary mapping sequence length to DataLoader for training data.
+        val_loaders (dict): Dictionary mapping sequence length to DataLoader for validation data.
+        optimizer (torch.optim.Optimizer): Optimizer to use for training.
+        scheduler (torch.optim.lr_scheduler._LRScheduler): Learning rate scheduler.
+        epochs (int): Total number of epochs to train for.
+        device (torch.device): Device to train on ('cuda' or 'cpu').
+        save_dir (str, optional): Directory to save the best model checkpoints. Default is "models".
+        log_file (str, optional): CSV path to log training/validation metrics. Default is "training_log.csv".
+
+    Returns:
+        model (torch.nn.Module): The trained model with the best checkpoint (based on validation MSE).
+    """
+    # Create save directory if it doesn't exist
+    os.makedirs(save_dir, exist_ok=True)
+
+    # Initialize loss criterion
+    criterion = nn.MSELoss(reduction='sum')  # Sum reduction for manual averaging
+
+    # Initialize lists to store losses
+    train_mse_losses = []
+    val_mse_losses = []
+    val_nmse_losses = []
+    best_val_mse = float('inf')
+
+    for epoch in range(epochs):
+        # Training loop
+        model.train()
+        train_mse = 0.0
+        train_samples = 0
+
+        print(f"\nEpoch {epoch + 1}/{epochs} [Training]")
+        for length, train_loader in train_loaders.items():
+            print(f"Processing sequences of length {length}")
+            with tqdm(train_loader, desc=f"Length {length} [Training]", unit="batch") as t:
+                for batch in t:
+                    optimizer.zero_grad()
+
+                    # Move data to device
+                    input_ids, masked_tokens, masked_pos = [b.to(device) for b in batch]
+
+                    # Forward pass
+                    logits_lm = model(input_ids, masked_pos)[0]
+
+                    # Compute MSE loss
+                    loss = criterion(masked_tokens, logits_lm)
+                    loss.backward()
+                    optimizer.step()
+                    scheduler.step()
+
+                    train_mse += loss.item()
+                    train_samples += input_ids.shape[0]
+
+                    # Update progress bar with MSE
+                    t.set_postfix({"mse": train_mse / train_samples, "lr": scheduler.get_last_lr()[0]})
+
+        # Average MSE across training samples
+        train_mse = train_mse / max(train_samples, 1)
+        train_mse_losses.append(train_mse)
+
+        # Validation loop every 2 epochs
+        if epoch % 2 == 0:
+            model.eval()
+            val_mse = 0.0
+            val_nmse = 0.0
+            val_samples = 0
+
+            with torch.no_grad():
+                print(f"\nEpoch {epoch + 1}/{epochs} [Validation]")
+                for length, val_loader in val_loaders.items():
+                    print(f"Processing sequences of length {length}")
+                    with tqdm(val_loader, desc=f"Length {length} [Validation]", unit="batch") as t:
+                        for batch in t:
+                            # Move data to device
+                            input_ids, masked_tokens, masked_pos = [b.to(device) for b in batch]
+
+                            # Forward pass
+                            logits_lm = model(input_ids, masked_pos)[0]
+
+                            # Compute MSE loss 
+                            mse = criterion(masked_tokens, logits_lm)
+                            val_mse += mse.item()
+
+                            # Compute NMSE for reporting
+                            masked_tokens_np = masked_tokens.cpu().numpy()
+                            logits_lm_np = logits_lm.cpu().numpy()
+                            nmse = nmse_loss(masked_tokens_np, logits_lm_np)
+                            val_nmse += nmse * input_ids.shape[0]
+
+                            val_samples += input_ids.shape[0]
+
+                            # Update progress bar with both MSE and NMSE
+                            t.set_postfix({"mse": val_mse / val_samples, "nmse": val_nmse / val_samples})
+
+            # Average MSE and NMSE across validation samples
+            val_mse = val_mse / max(val_samples, 1)
+            val_nmse = val_nmse / max(val_samples, 1)
+            val_mse_losses.append(val_mse)
+            val_nmse_losses.append(val_nmse)
+
+            # Save model if validation MSE improves
+            if val_mse < best_val_mse:
+                best_val_mse = val_mse
+                model_path = os.path.join(save_dir, f"lwm_epoch{epoch+1}_train{train_mse:.4f}_val{val_mse:.4f}.pth")
+                torch.save(model.state_dict(), model_path)
+                print(f"Model saved: {model_path}")
+
+        # Log the results
+        print(f"  Train MSE: {train_mse:.4f}")
+        if epoch % 2 == 0:
+            print(f"  Validation MSE: {val_mse:.4f}")
+            print(f"  Validation NMSE: {val_nmse:.4f}")
+        print(f"  Learning Rate: {scheduler.get_last_lr()[0]:.6e}")
+
+        # Plot losses after each epoch
+        plt.figure(figsize=(10, 6))
+        plt.plot(range(1, len(train_mse_losses) + 1), train_mse_losses, label="Train MSE")
+        if val_mse_losses:  # Plot validation only if it exists
+            plt.plot(range(1, len(val_mse_losses) + 1), val_mse_losses, label="Validation MSE")
+            plt.plot(range(1, len(val_nmse_losses) + 1), val_nmse_losses, label="Validation NMSE")
+        plt.xlabel("Epochs")
+        plt.ylabel("Loss")
+        plt.title("Training and Validation Losses")
+        plt.legend()
+        plt.grid(True)
+        plt.show()
+
+    print("Training and validation complete.")
+    return model
+
+def inference(model, tokens, batch_size=128, device="cuda"):
+    """
+    Runs inference using a trained model to extract embeddings from input tokens.
+
+    This function processes the input tokens in batches (without shuffling), moves them to the 
+    specified device, passes them through the model, and aggregates the outputs.
+
+    Args:
+        model (torch.nn.Module): The trained model used for inference.
+        tokens (torch.Tensor): Input tensor of shape (N, ...) representing the tokenized data.
+        batch_size (int, optional): Batch size for inference. Default is 128.
+        device (str or torch.device, optional): Device to run inference on. Default is "cuda".
+
+    Returns:
+        torch.Tensor: A tensor of shape (N, D), where D is the embedding dimension produced by the model.
+    """
+    dataset = TensorDataset(tokens)
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+    
+    embeddings = []
+    model.eval()
+    with torch.no_grad():
+        with tqdm(dataloader, desc="Inference", unit="batch") as t:
+            for batch in t:
+                
+                input_ids = batch[0].to(device)
+                output = model(input_ids)
+                embeddings.append(output)
+                    
+    output_total = torch.cat(embeddings, dim=0).float()
+    return output_total
+
+def visualize_embeddings(embeddings, labels=None, method="tsne", label=None):
+    """
+    Visualizes high-dimensional embeddings in 2D using PCA, UMAP, or t-SNE.
+
+    This function reduces the dimensionality of embeddings to two components and visualizes them
+    with an optional color-coding based on provided labels. It supports three reduction methods:
+    PCA (linear), UMAP (nonlinear, preserves local/global structure), and t-SNE (nonlinear, local structure).
+
+    Args:
+        embeddings (torch.Tensor or np.ndarray): Embedding matrix of shape (n_samples, n_features).
+        labels (torch.Tensor or np.ndarray, optional): Class labels of shape (n_samples,). If provided,
+            each class will be visualized with a distinct color.
+        method (str): Dimensionality reduction method: one of {'pca', 'umap', 'tsne'}. Default is 'tsne'.
+        label (str, optional): Title for the plot. Defaults to "Embedding Visualization" if not provided.
+
+    Raises:
+        ValueError: If an unsupported dimensionality reduction method is specified.
+
+    Returns:
+        None. Displays a 2D scatter plot of the embeddings.
+    """
+    # to numpy
+    if isinstance(embeddings, torch.Tensor):
+        embeddings = embeddings.cpu().numpy()
+    if labels is not None and isinstance(labels, torch.Tensor):
+        labels = labels.cpu().numpy()
+
+    # choose reducer
+    m = method.lower()
+    if m == "pca":
+        reducer = PCA(n_components=2)
+    elif m == "umap":
+        reducer = umap.UMAP(n_components=2, n_neighbors=16, random_state=42)
+    elif m == "tsne":
+        reducer = TSNE(n_components=2, random_state=42, init="random")
+    else:
+        raise ValueError("Invalid method. Choose 'pca', 'umap', or 'tsne'.")
+
+    Z = reducer.fit_transform(embeddings)
+
+    plt.figure(figsize=(10, 8))
+    if labels is not None:
+        num_classes = len(np.unique(labels))
+        colors = plt.cm.get_cmap("tab10", num_classes)
+
+        for class_idx in range(num_classes):
+            class_points = Z[labels == class_idx]
+            plt.scatter(
+                class_points[:, 0], class_points[:, 1],
+                label=f"Class {class_idx}",
+                alpha=0.6,
+                cmap=colors
+            )
+    else:
+        plt.scatter(Z[:, 0], Z[:, 1], color="C0", alpha=0.6, label="Samples")
+
+    title = label or "Embedding Visualization"
+    plt.title(f"{title} ({method.upper()})")
+    plt.xlabel("Component 1")
+    plt.ylabel("Component 2")
+    plt.show()
+
+
+def embedding_space_visual(model, data, input_type="cls_emb", device="cpu", batch_size=64, task=None, visualization=False, labels=None, visualization_method="tsne", selected_tokens=0):
+    """
+    Extracts embeddings from a model and optionally visualizes the embedding space.
+
+    Supports different types of embeddings from the model (e.g., CLS token, mean-pooled embeddings),
+    and provides 2D visualization using t-SNE, PCA, or UMAP. Also supports angular-based clustering 
+    if the task is "ChannelCharting".
+
+    Args:
+        model (torch.nn.Module): The trained model for embedding extraction.
+        data (torch.Tensor): Input tensor of shape (N, ...). For non-'raw' types, it's tokenized input.
+        input_type (str): Type of embedding to extract. Options:
+            - 'cls_emb': Extract CLS token (first token) embeddings.
+            - 'channel_emb': Extract all non-CLS token embeddings.
+            - 'combined': Use the full output from the model.
+            - 'mean_pooled': Mean of all token embeddings.
+            - 'arbitrary_concat': Select specific token index (given by `selected_tokens`).
+            - 'arbitrary_meanPooled': Mean-pool over selected token indices.
+            - 'raw': Use the input `data` as is.
+        device (str or torch.device): The device for computation ('cuda' or 'cpu').
+        batch_size (int): Batch size for inference.
+        task (str, optional): If "ChannelCharting", performs position-based angular clustering.
+        visualization (bool): Whether to visualize the embeddings.
+        labels (torch.Tensor or np.ndarray, optional): Ground-truth labels or 2D positions used for coloring the plot.
+        visualization_method (str): One of {'tsne', 'pca', 'umap'} for 2D visualization.
+        selected_tokens (int, list, or tensor): Used for 'arbitrary_concat' or 'arbitrary_meanPooled' types.
+
+    Returns:
+        torch.Tensor: Output embedding tensor of shape (N, D) ready for downstream tasks or analysis.
+    """
+    print("\nPreparing for LWM inference and embedding space visualization ...")
+    if input_type == "raw":
+        output_total = data
+    else:
+        dataset = TensorDataset(data)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+        
+        embeddings = []
+        model.eval()
+        with torch.no_grad():
+            with tqdm(dataloader, desc="Inference", unit="batch") as t:
+                for batch in t:
+                    
+                    input_ids = batch[0].to(device)
+                    output = model(input_ids)
+                
+                    if input_type == "cls_emb":
+                        batch_embeddings = output[:, 0] 
+                    elif input_type == "channel_emb":
+                        batch_embeddings = output[:, 1:] 
+                    elif input_type == "combined":
+                        batch_embeddings = output
+                    elif input_type == "mean_pooled":
+                        batch_embeddings = torch.mean(output, dim=1).unsqueeze(1)
+                    elif input_type == "arbitrary_concat":
+                        batch_embeddings = output[:, selected_tokens]
+                    elif input_type == "arbitrary_meanPooled":
+                        batch_embeddings = torch.mean(output[:, selected_tokens], dim=1).unsqueeze(1)
+                    
+                    embeddings.append(batch_embeddings)
+                    
+                    
+        output_total = torch.cat(embeddings, dim=0).float()
+        
+        if visualization:
+            
+            if task in ["ChannelCharting"]:
+                
+                positions = labels.cpu().numpy()
+                x_coords = positions[:, 0] 
+                y_coords = positions[:, 1]  
+                center_x = np.mean(x_coords)
+                center_y = np.mean(y_coords)
+                x_shifted = x_coords - center_x
+                y_shifted = y_coords - center_y
+                angles = np.arctan2(y_shifted, x_shifted)
+                angles = (angles + 2 * np.pi) % (2 * np.pi)
+                n_clusters = 8
+                sector_size = (2 * np.pi) / n_clusters  # Size of each sector in radians
+                labels = np.floor(angles / sector_size).astype(int)
+                labels = np.clip(labels, 0, n_clusters - 1)
+
+                plt.figure(figsize=(10, 8))
+                if labels is not None:
+                    # Color-code by labels if provided
+                    num_classes = len(np.unique(labels))
+                    colors = plt.cm.get_cmap("tab10", num_classes)
+    
+                    for class_idx in range(num_classes):
+                        class_points = positions[labels == class_idx]
+                        plt.scatter(
+                            class_points[:, 0], class_points[:, 1],
+                            label=f"Class {class_idx}",
+                            alpha=0.6,
+                            cmap=colors
+                        )
+                else:
+                    # Plot all points in a single color if no labels
+                    plt.scatter(
+                        positions[:, 0], positions[:, 1],
+                        color="blue",  # Default color for unlabeled data
+                        alpha=0.6,
+                        label="Samples"
+                    )
+                plt.title("Original 2D Positions with 5 Clusters")
+                plt.xlabel("X Coordinate")
+                plt.ylabel("Y Coordinate")
+                plt.legend()
+                plt.grid(True, linestyle="--", alpha=0.3)
+                plt.show()
+
+            visualize_embeddings(output_total.view(output_total.size(0), -1), 
+                                 labels=labels, 
+                                 method=visualization_method, 
+                                 label="Embedding Space")
+        
+    return output_total
+
+def plot_radar_chart(task_names, optimized_scores, baseline_scores, title="Task Performance Comparison", figsize=(8, 8), save_path="submission/chart.png"):
+    """
+    Plot a dark-themed radar chart comparing optimized and baseline scores.
+
+    Args:
+        task_names (list): List of task names (e.g., ["LoS/NLoS Classification", ...]).
+        optimized_scores (list): List of optimized performance scores.
+        baseline_scores (list): List of baseline performance scores.
+        title (str): Title of the chart (default: "Task Performance Comparison").
+        figsize (tuple): Figure size (width, height) in inches (default: (8, 8)).
+        save_path (str): Path to save the figure (default: "submission/chart.png").
+
+    Raises:
+        ValueError: If input lists have mismatched lengths or are empty.
+    """
+    # Input validation
+    if not task_names or not optimized_scores or not baseline_scores:
+        raise ValueError("All input lists (task_names, optimized_scores, baseline_scores) must not be empty")
+    if not (len(task_names) == len(optimized_scores) == len(baseline_scores)):
+        raise ValueError("All input lists must have the same length")
+
+    # Number of variables (tasks)
+    num_tasks = len(task_names)
+
+    # Repeat the first score to close the radar chart
+    angles = [n / float(num_tasks) * 2 * pi for n in range(num_tasks)]
+    angles += angles[:1]
+    optimized = optimized_scores + optimized_scores[:1]
+    baseline = baseline_scores + baseline_scores[:1]
+
+    # Figure & Axes Setup
+    plt.style.use('default')  # Reset to default style to avoid global dark background
+    fig, ax = plt.subplots(
+        figsize=figsize, dpi=300,
+        subplot_kw=dict(polar=True)
+    )
+    fig.patch.set_facecolor('#1a1a1a')  # Dark gray background for the figure (outside the circle)
+    ax.set_facecolor('#1a1a1a')  # Dark background for the plot (inside the circle)
+    ax.patch.set_alpha(1)
+
+    ax.set_theta_offset(pi / 2)
+    ax.set_theta_direction(-1)
+
+    # Grid & Labels
+    ax.set_xticks(angles[:-1])
+    ax.set_xticklabels(task_names, fontsize=14, fontweight='bold', color='white')
+    ax.set_rlabel_position(0)
+    ax.set_yticks([0.2, 0.4, 0.6, 0.8, 1.0])
+    ax.set_yticklabels(['20%', '40%', '60%', '80%', '100%'], fontsize=11, color='#bbbbbb')
+    ax.set_ylim(0, 1.15)
+    ax.grid(color='#bbbbbb', linestyle=':', linewidth=0.8)
+
+    # Glow Effect Function
+    def plot_with_glow(x, y, color, label):
+        for lw, alpha in zip([15, 10, 6, 3], [0.02, 0.04, 0.06, 0.08]):
+            ax.plot(x, y, linewidth=lw, color=color, alpha=alpha, zorder=1)
+        ax.plot(x, y, linewidth=2.5, color=color, label=label, zorder=2)
+        ax.scatter(x, y, s=100, color=color, edgecolors='white', linewidth=1.5, zorder=3)
+
+    # Plot optimized & baseline
+    plot_with_glow(angles, optimized, color='#00ffff', label='Optimized')
+    plot_with_glow(angles, baseline, color='#ff5588', label='Baseline')
+
+    # Layered Fill
+    ax.fill_between(angles, optimized, color='#00ffff', alpha=0.12, zorder=1)
+    ax.fill_between(angles, baseline, color='#ff5588', alpha=0.12, zorder=1)
+
+    # Custom Legend
+    legend_elements = [
+        plt.Line2D([0], [0], color='#00ffff', lw=3, label='Optimized'),
+        plt.Line2D([0], [0], color='#ff5588', lw=3, label='Baseline')
+    ]
+    ax.legend(
+        handles=legend_elements,
+        loc='upper right', bbox_to_anchor=(1.15, 1.15),
+        frameon=True, facecolor='#bbbbbb', edgecolor='#555555', fontsize=18
+    )
+    
+    # Adjust layout
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight', transparent=False)
+    plt.show()
+