# 🔬 LWM Spectrogram Research Framework

**Core Research Focus**: Comprehensive research framework to validate Large Wireless Model (LWM) representation learning effectiveness on spectrogram data.

This project systematically compares and analyzes **performance differences between using and not using LWM**.

## 🏗️ Research Framework Structure

### 📂 **3-Stage Research Pipeline**
```
pretraining/ → downstream/ → baselines/
     ↓             ↓             ↓
  LWM Training    Transfer Learning    Direct Learning
  (Representation)  (Embedding Usage)  (Raw Data)
```

### 🎯 **Core Research Questions**
- **LWM Representation Learning Effectiveness**: How helpful is pre-trained LWM for spectrogram understanding?
- **Transfer Learning Efficiency**: How much does LWM embedding improve downstream task performance?
- **Cost-Benefit Trade-off**: Is the performance improvement from LWM addition justified by the computational cost?

### 🚀 **Key Features**
- **Comprehensive Comparative Analysis**: LWM usage vs non-usage performance comparison
- **Multiple Model Support**: ResNet, MobileNet, EfficientNet, SqueezeNet, etc.
- **Real-time Monitoring**: tqdm-based progress tracking and performance metrics
- **Automatic Visualization**: Training curves and performance metric graphs
- **Organized Result Storage**: Timestamp-based folder structure

## 📋 Requirements

```bash
Python 3.10+
PyTorch 2.0.0+
torchvision 0.15.0+
torchaudio 2.0.0+
CUDA 12.1+ (recommended for GPU training)
```

**Core Dependencies:**
- torch>=2.0.0, torchvision>=0.15.0, torchaudio>=2.0.0 (PyTorch ecosystem)
- numpy>=1.21.0,<2.0, scipy>=1.15.3, scikit-learn>=1.6.1 (scientific computing)
- matplotlib>=3.10.3, seaborn>=0.11.0 (visualization)
- tqdm>=4.67.1 (progress bars)
- DeepMIMO>=4.0.0b9 (wireless channel modeling)
- umap-learn>=0.5.7 (dimensionality reduction)

## 🛠️ Installation

### Option 1: Recommended (pyproject.toml) - Standard Installation
```bash
git clone https://github.com/yourusername/lwm-spectro.git
cd lwm-spectro

# Install in editable mode (recommended for development)
pip install -e .

# Install DeepMIMO separately (important for compatibility)
pip install --no-deps "DeepMIMO>=4.0.0b9"

# Or install with development dependencies
pip install -e ".[dev]"

# For GPU support on CUDA systems
pip install -e ".[gpu]"
```

### Option 1.5: Lambda/AWS Optimized Installation
```bash
git clone https://github.com/yourusername/lwm-spectro.git
cd lwm-spectro

# Install with Lambda-compatible versions
pip install -e ".[lambda]"

# Install DeepMIMO separately (important for Lambda compatibility)
pip install --no-deps "DeepMIMO>=4.0.0b9"

# Verify installation
python -c "
import torch
import numpy as np
import matplotlib
import deepmimo as dm
print('✅ Lambda installation successful!')
print(f'PyTorch: {torch.__version__}')
print(f'NumPy: {np.__version__}')
print(f'matplotlib: {matplotlib.__version__}')
print(f'DeepMIMO: {dm.__version__}')
"
```

### Option 2: Legacy (requirements.txt)
```bash
git clone https://github.com/yourusername/lwm-spectro.git
cd lwm-spectro

# Install from requirements.txt
pip install -r requirements.txt
```

### Option 3: Conda Environment
```bash
# If you have conda environment 'lwm' set up:
conda activate lwm
pip install -e .
```

### Server Environment Setup (NumPy 2.3 Compatible)
For remote servers with NumPy 2.3+ environment:
```bash
# Update system and install basic dependencies
sudo apt update && sudo apt install -y git python3.12 python3.12-venv

# Clone and install
git clone https://github.com/yourusername/lwm-spectro.git
cd lwm-spectro

# Create virtual environment with Python 3.12
python3.12 -m venv venv
source venv/bin/activate

# Install with server-compatible versions (NumPy 2.3+ support)
pip install -e .

# Or install from requirements.txt (updated for server compatibility)
pip install -r requirements.txt

# For CUDA support (if GPU available)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# Verify NumPy version compatibility
python -c "
import numpy as np
import numba
print(f'✅ NumPy version: {np.__version__}')
print(f'✅ Numba version: {numba.__version__}')
print('✅ Server environment ready!')
"
```

### Conda Environment Setup (Recommended for Server)
```bash
# Create conda environment from provided environment.yml
conda env create -f environment.yml
conda activate lwm-spectro-server

# Verify installation
python -c "
import torch
import numpy as np
import numba
import matplotlib
print('✅ All dependencies installed successfully!')
print(f'PyTorch: {torch.__version__}')
print(f'NumPy: {np.__version__}')
print(f'Numba: {numba.__version__}')
print(f'Matplotlib: {matplotlib.__version__}')
"
```

### Troubleshooting NumPy Compatibility Issues
If you encounter NumPy/numba compatibility issues:
```bash
# Option 1: Update numba to latest version (supports NumPy 2.3)
pip install --upgrade numba

# Option 2: Downgrade NumPy if needed (not recommended for server)
pip install 'numpy>=1.21.0,<2.3.0'

# Option 3: Use conda environment (recommended)
conda env create -f environment.yml
conda activate lwm-spectro-server

# Option 4: Force reinstall with compatible versions
pip uninstall numpy numba -y
pip install 'numpy>=1.21.0,<2.4.0' 'numba>=0.59.0'
```

## 🎯 Usage

### Basic Training

Train with all available spectrogram sizes:
```bash
python mcs_classification.py
```

### Specify Spectrogram Size

Train with 32x32 spectrograms only:
```bash
python mcs_classification.py 32x32
```

Train with 128x128 spectrograms only:
```bash
python mcs_classification.py 128x128
```

### Choose Model Architecture

Use different model architectures:
```bash
# Custom lightweight CNN (default)
python mcs_classification.py 32x32 --model custom

# ResNet18 (balanced performance)
python mcs_classification.py 32x32 --model resnet18

# MobileNetV2 (lightweight mobile-optimized)
python mcs_classification.py 32x32 --model mobilenet_v2

# EfficientNet-B0 (efficient architecture)
python mcs_classification.py 32x32 --model efficientnet_b0

# SqueezeNet (minimal memory usage)
python mcs_classification.py 32x32 --model squeezenet
```

## 📊 Data Structure

The project uses spectrogram data stored in pickle files:

- **32x32 Spectrograms**: Generated from 128FFT (smaller size, faster training)
- **128x128 Spectrograms**: Generated from 512FFT (larger size, more features)

Each pickle file contains:
```python
{
    'spectrograms': numpy.ndarray,  # Shape: (1000, H, W)
    'configuration': dict          # Metadata
}
```

## 🏷️ MCS Classes (7 Classes)

| MCS Level | Modulation | Code Rate | Description |
|-----------|------------|-----------|-------------|
| 0 | QPSK | 1/2 | Basic QPSK with 1/2 code rate |
| 1 | QPSK | 3/4 | QPSK with 3/4 code rate |
| 2 | QAM16 | 1/2 | 16-QAM with 1/2 code rate |
| 3 | QAM16 | 3/4 | 16-QAM with 3/4 code rate |
| 4 | QAM64 | 1/2 | 64-QAM with 1/2 code rate |
| 5 | QAM64 | 2/3 | 64-QAM with 2/3 code rate |
| 6 | QAM64 | 3/4 | 64-QAM with 3/4 code rate |

## 📁 Output Files

Training automatically creates organized output in `models/[model_type]_[size]/`:

### Model Files
- `best.pth` - Best performing model (highest F1 score)
- `latest.pth` - Most recent model checkpoint

### Configuration & Logs
- `config.json` - Complete training configuration and dataset info
- `training_history.json` - Epoch-by-epoch metrics (loss, F1 scores, accuracy)
- `model_summary.txt` - Model architecture and parameter count
- `performance_metrics.json` - Final evaluation results

### Visualizations
- `training_curves.png` - Combined loss, F1, and accuracy curves
- `loss_curves.png` - Loss curves only
- `f1_progression.png` - F1 score progression over epochs

## 🔧 Model Architectures

### Custom CNN
- Lightweight custom convolutional neural network
- Optimized for speed and minimal resource usage
- Perfect for quick experiments and Apple Silicon

### ResNet (18/34)
- Residual Network architectures
- Excellent balance of performance and computational efficiency
- Pre-trained weights available

### MobileNet (V2/V3)
- Mobile-optimized architectures
- Extremely efficient for resource-constrained environments
- V3 includes neural architecture search optimizations

### EfficientNet-B0
- Compound scaling for optimal efficiency
- Balances network depth, width, and resolution
- State-of-the-art efficiency

### SqueezeNet
- Extreme compression with minimal parameter count
- Maintains accuracy with dramatically reduced model size
- Ideal for deployment on edge devices

## 🎛️ Configuration Options

### Training Parameters
- `batch_size`: Batch size for training (default: 32)
- `learning_rate`: Initial learning rate (default: 0.001)
- `num_epochs`: Maximum training epochs (default: 50)
- `patience_f1`: Epochs to wait for F1 improvement (default: 10)
- `patience_loss`: Epochs to wait for loss decrease (default: 5)

### Early Stopping Criteria
- F1 score improvement threshold (default: 0.001)
- Validation loss stagnation detection
- Configurable patience parameters

## 📈 Real-time Monitoring

The training process provides live updates:

```
🚀 MCS Classification Training (Size: 32x32, Model: custom)
📊 Dataset Info:
   Total samples: 147,000
   Training samples: 117,600
   Validation samples: 29,400
   Classes: 7 (QPSK 1/2, QPSK 3/4, QAM16 1/2, QAM16 3/4, QAM64 1/2, QAM64 2/3, QAM64 3/4)

🏗️ Model Info:
   Type: custom
   Total parameters: 619,015
   Trainable parameters: 619,015

🏃 Training with Smart Early Stopping...
Epoch 1/50 - Train Loss: 1.234, Val Loss: 0.987, F1: 0.654, F1 Improvement: 0.654
Epoch 2/50 - Train Loss: 0.876, Val Loss: 0.765, F1: 0.712, F1 Improvement: 0.058
...
```

## 🎯 Evaluation Metrics

The system evaluates using comprehensive F1 metrics:

- **F1 Macro**: Unweighted mean of per-class F1 scores
- **F1 Micro**: Global F1 score (accounts for class imbalance)
- **F1 Weighted**: Weighted mean by class support
- **Per-class F1**: Individual F1 scores for each MCS class
- **Accuracy**: Overall classification accuracy

## 🔍 Project Structure

```
lwm-spectro/
├── __init__.py              # Package initialization
├── pyproject.toml           # Modern Python packaging (recommended)
├── requirements.txt         # Legacy dependency management
├── README.md                # This documentation
├── .gitignore              # Git ignore rules
├── lwm_spectrogram_adapter.py  # Main LWM adapter script
├── mcs_classification.py   # MCS classification training
├── mcs_classifier_lwm.py   # LWM-based MCS classifier
├── train_32x32.py         # 32x32 model training
├── train_128x128.py       # 128x128 model training
├── utils.py               # Utility functions
├── pretrained_model.py    # Pretrained model definitions
├── train_heads_config.py  # Training configuration
├── check_cuda.py         # CUDA availability checker
├── models/               # Auto-generated model outputs
│   ├── custom_32x32/    # Model-specific folders
│   ├── resnet18_32x32/
│   └── ...
├── spectrograms/         # Spectrogram data directory
│   └── city_0_newyork/  # Dataset (not in repo)
│       ├── LTE/
│       ├── QPSK/
│       ├── QAM16/
│       └── QAM64/
```

## 🚀 Research Workflow Execution Guide

### 1️⃣ **Step 1: LWM Pretraining** (Representation Learning)
```bash
# Pre-train LWM model with spectrogram data
cd pretraining
python lwm_spectrogram_adapter.py --max-samples-per-size 1000

# Results: LWM checkpoints saved in models/ folder
```

### 2️⃣ **Step 2: Downstream Tasks** (Transfer Learning)
```bash
# Classification using pre-trained LWM embeddings
cd downstream
python mcs_classification.py --model-type resnet18 --target-size 32x32

# Compare performance with different models
python mcs_classification.py --model-type efficientnet_b0 --target-size 128x128
```

### 3️⃣ **Step 3: Baselines** (Direct Learning Comparison)
```bash
# Raw spectrogram classification without LWM
cd baselines
python baseline_models_training.py --target-size 32x32
```

### 🔄 **Complete Comparative Workflow**
```bash
# 1. Measure baseline performance
cd baselines && python baseline_models_training.py

# 2. LWM pre-training
cd ../pretraining && python lwm_spectrogram_adapter.py

# 3. Measure LWM-enhanced performance
cd ../downstream && python mcs_classification.py --model-type resnet18

# 4. Compare and analyze results
```

## 📊 Research Methodology

### **Evaluation Metrics**
- **F1 Score**: Macro, Micro, Weighted averages
- **Accuracy**: Overall classification accuracy
- **Training Efficiency**: Convergence speed, memory usage comparison

### **Experimental Design**
- **Cross-validation**: Compare LWM usage vs non-usage with identical data
- **Model Diversity**: Validate LWM effectiveness across different architectures
- **Data Size Comparison**: Compare 32x32 vs 128x128 spectrograms

### **Key Features**
- **Apple Silicon Optimization**: MPS support for efficient training on Mac
- **Single-channel Input Optimization**: All models specialized for spectrograms
- **Automatic Result Storage**: Timestamp-based organized result management
- **Real-time Monitoring**: tqdm-based progress tracking
- **Automated Visualization**: Training curves and performance metric graphs

## 🌐 Hugging Face Model Hub

This project is available on Hugging Face for easy access and deployment!

### 📦 Model Card & Documentation

We've prepared comprehensive documentation for Hugging Face Hub:

- **MODEL_CARD.md**: Complete model description, architecture, usage examples, and results
- **example_inference.py**: Ready-to-use inference scripts with visualization
- **HUGGINGFACE_GUIDE.md**: Step-by-step upload guide

### 🚀 Quick Upload to Hugging Face

```bash
# Method 1: Using Bash script (recommended)
./upload_to_huggingface.sh

# Method 2: Using Python script
python upload_to_huggingface.py --username YOUR_USERNAME --upload-type minimal
```

Upload types:
- `minimal`: README + MoE checkpoint only
- `models`: All trained models
- `full`: Entire codebase

### 📚 Documentation Files

| File | Description |
|------|-------------|
| `MODEL_CARD.md` | Main Hugging Face model card (README) |
| `HUGGINGFACE_GUIDE.md` | Detailed upload instructions |
| `HUGGINGFACE_SUMMARY.md` | Quick reference guide |
| `example_inference.py` | Inference code examples |
| `upload_to_huggingface.sh` | Automated upload script (Bash) |
| `upload_to_huggingface.py` | Automated upload script (Python) |
| `.gitattributes` | Git LFS configuration |

For more details, see **[HUGGINGFACE_SUMMARY.md](HUGGINGFACE_SUMMARY.md)**.

## 🤝 Contributing

Contributions are welcome! Please feel free to submit issues and pull requests.

## 📄 License

This project is open source. Please check the license file for details.