minor fixes on src/config

README.md

@@ -1,4 +1,4 @@
-# Official implementation of ICC 2025 paper [AdaFortiTran: An Adaptive Transformer Model for Robust OFDM Channel Estimation](https://arxiv.org/abs/2505.09076)
+# Official implementation of [AdaFortiTran: An Adaptive Transformer Model for Robust OFDM Channel Estimation](https://arxiv.org/abs/2505.09076) accepted at ICC 2025, Montreal, Canada.
 
 
 ## License

config/adafortitran.yaml

@@ -1,3 +1,4 @@
+model_type: 'adafortitran'
 patch_size: [3, 2]
 num_layers: 6
 model_dim: 128

config/fortitran.yaml

@@ -1,3 +1,4 @@
+model_type: 'fortitran'
 patch_size: [3, 2]
 num_layers: 6
 model_dim: 128
@@ -5,5 +6,4 @@ num_head: 4
 activation: 'gelu'
 dropout: 0.1
 max_seq_len: 512
-pos_encoding_type: 'learnable'
-adaptive_token_length: 6
+pos_encoding_type: 'learnable'

scripts/add_gitkeep.py

@@ -44,7 +44,6 @@ def add_gitkeep_to_directories(root_path: str | Path):
     print(f"\nTotal .gitkeep files added: {gitkeep_count}")
 
 if __name__ == "__main__":
-    # Add .gitkeep to all subdirectories in the data folder
     data_path = Path("data")
     
     print(f"Adding .gitkeep files to subdirectories in {data_path.absolute()}")

src/config/__init__.py

@@ -1 +1,5 @@
-from src.config.schemas import ModelConfig, SystemConfig
+"""This module provides a clean interface for loading and validating configuration files."""
+
+from .config_loader import load_config
+
+__all__ = ["load_config"]

src/config/config_loader.py

@@ -29,6 +29,8 @@ class ConfigLoader:
             ValueError: If configuration validation fails
         """
         system_config_path = Path(system_config_path)
+
+        # certain models may not have a model config
         model_config = None
         if model_config_path is not None:
             model_config_path = Path(model_config_path)
@@ -48,7 +50,6 @@ class ConfigLoader:
         except ValidationError as e:
             raise ValueError(f"System configuration validation for {system_config_path} failed:\n{e}")
 
-        # Only load model config if path is provided and file exists
         if model_config_path is not None and model_config_path.exists():
             try:
                 with open(model_config_path, 'r') as f:

src/config/schemas.py

@@ -1,14 +1,18 @@
 from pydantic import BaseModel, Field, model_validator
-from typing import Self, Tuple, List, Optional
+from typing import Self, Tuple, List, Optional, Literal
 import torch
 
 
 class OFDMParams(BaseModel):
+    # ... means required (i.e. no default value)
+    # gt=0 means greater than 0
     num_scs: int = Field(..., gt=0, description="Number of sub-carriers")
     num_symbols: int = Field(..., gt=0, description="Number of OFDM symbols")
 
 
 class PilotParams(BaseModel):
+    # ... means required (i.e. no default value)
+    # gt=0 means greater than 0
     num_scs: int = Field(..., gt=0, description="Number of pilots across sub-carriers")
     num_symbols: int = Field(..., gt=0, description="Number of pilots across OFDM symbols")
 
@@ -17,7 +21,7 @@ class SystemConfig(BaseModel):
     ofdm: OFDMParams
     pilot: PilotParams
 
-    @model_validator(mode='after')
+    @model_validator(mode='after')  # validate after all fields are initialized
     def validate_pilot_constraints(self) -> Self:
         """Ensure pilot parameters don't exceed OFDM parameters."""
         if self.pilot.num_scs > self.ofdm.num_scs:
@@ -33,25 +37,62 @@ class SystemConfig(BaseModel):
             )
         return self
 
-    model_config = {"extra": "forbid"}
+    model_config = {"extra": "forbid"}  # forbid extra fields
 
 
 class ModelConfig(BaseModel):
+    model_type: Literal["fortitran", "adafortitran"] = Field(
+        default="fortitran",
+        description="Type of model (fortitran or adafortitran)"
+    )
     patch_size: Tuple[int, int] = Field(..., description="Patch size as (height, width)")
     num_layers: int = Field(..., gt=0, description="Number of transformer layers")
     model_dim: int = Field(..., gt=0, description="Model dimension")
     num_head: int = Field(..., gt=0, description="Number of attention heads")
-    activation: str = Field(default="gelu", description="Activation function")
-    dropout: float = Field(default=0.1, ge=0.0, le=1.0, description="Dropout rate")
+    activation: Literal["relu", "gelu"] = Field(
+        default="gelu", 
+        description="Activation function used within the transformer's FFN"
+    )
+    dropout: float = Field(default=0.1, ge=0.0, le=1.0, description="Dropout rate used within the transformer's FFN")
     max_seq_len: int = Field(default=512, gt=0, description="Maximum sequence length")
-    pos_encoding_type: str = Field(default="learnable", description="Position encoding type")
-    adaptive_token_length: int = Field(default=6, gt=0, description="Adaptive token length")
+    pos_encoding_type: Literal["learnable", "sinusoidal"] = Field(
+        default="learnable", 
+        description="Positional encoding type"
+    )
+    adaptive_token_length: Optional[int] = Field(
+        default=None, 
+        gt=0, 
+        description="Adaptive token length (required for AdaFortiTran)"
+    )
     channel_adaptivity_hidden_sizes: Optional[List[int]] = Field(
         default=None, 
-        description="Hidden sizes for channel adaptation layers"
+        description="Hidden sizes for channel adaptation layers (required for AdaFortiTran)"
     )
     device: str = Field(default="cpu", description="Device to use")
 
+    @model_validator(mode='after')
+    def validate_model_specific_requirements(self) -> Self:
+        """Validate model-specific configuration requirements."""
+        if self.model_type == "adafortitran":
+            if self.channel_adaptivity_hidden_sizes is None:
+                raise ValueError(
+                    "channel_adaptivity_hidden_sizes is required for AdaFortiTran model"
+                )
+            if self.adaptive_token_length is None:
+                raise ValueError(
+                    "adaptive_token_length is required for AdaFortiTran model"
+                )
+        
+        if self.model_type == "fortitran":
+            if self.channel_adaptivity_hidden_sizes is not None:
+                # Note: channel_adaptivity_hidden_sizes will be ignored for FortiTran
+                pass
+            if self.adaptive_token_length is not None:
+                # Note: adaptive_token_length will be ignored for FortiTran
+                pass
+        
+        return self
+
     @model_validator(mode='after')
     def validate_device(self) -> Self:
         """Validate that the specified device is available."""
@@ -67,7 +108,6 @@ class ModelConfig(BaseModel):
                 self.device = 'cpu'
             return self
 
-        # Validate CPU
         if device_str == 'cpu':
             return self
 

src/main.py

@@ -13,7 +13,7 @@ from pathlib import Path
 
 from src.main.parser import parse_arguments
 from src.main.trainer import train
-from src.config.config_loader import load_config
+from src.config import load_config
 
 
 def setup_logging(log_level: str) -> None:
@@ -58,7 +58,10 @@ def main() -> None:
         logger.info("Configuration loaded successfully")
         logger.info(f"OFDM dimensions: {system_config.ofdm.num_scs} subcarriers x {system_config.ofdm.num_symbols} symbols")
         logger.info(f"Pilot dimensions: {system_config.pilot.num_scs} subcarriers x {system_config.pilot.num_symbols} symbols")
-        logger.info(f"Model architecture: {model_config.num_layers} layers, {model_config.model_dim} dimensions")
+        if model_config is not None:
+            logger.info(f"Model architecture: {model_config.num_layers} layers, {model_config.model_dim} dimensions")
+        else:
+            logger.info("Using Linear model (no model config required)")
         
         # Start training
         logger.info("Initializing training...")