fixes on src/models

config/linear.yaml

@@ -0,0 +1,2 @@
+model_type: 'linear'
+device: 'cpu' 

src/config/config_loader.py

@@ -13,31 +13,30 @@ class ConfigLoader:
     def __init__(self):
         self.logger = logging.getLogger(__name__)
 
-    def load_and_validate(self, system_config_path: Union[str, Path], model_config_path: Optional[Union[str, Path]] = None) -> Tuple[SystemConfig, Optional[ModelConfig]]:
+    def load_and_validate(self, system_config_path: Union[str, Path], model_config_path: Union[str, Path]) -> Tuple[SystemConfig, ModelConfig]:
         """
         Load and validate configuration files from YAML files.
 
         Args:
             system_config_path: Path to YAML configuration file for OFDM-related parameters
-            model_config_path: Optional path to YAML configuration file for model-related parameters
+            model_config_path: Path to YAML configuration file for model-related parameters
 
         Returns:
-            Tuple of (SystemConfig, Optional[ModelConfig]): Validated configuration objects
+            Tuple of (SystemConfig, ModelConfig): Validated configuration objects
 
         Raises:
-            FileNotFoundError: If the system config file doesn't exist
+            FileNotFoundError: If either config file doesn't exist
             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)
+        model_config_path = Path(model_config_path)
 
         if not system_config_path.exists():
             raise FileNotFoundError(f"System configuration file not found: {system_config_path}")
 
+        if not model_config_path.exists():
+            raise FileNotFoundError(f"Model configuration file not found: {model_config_path}")
+
         try:
             with open(system_config_path, 'r') as f:
                 system_raw_config = yaml.safe_load(f)
@@ -50,22 +49,22 @@ class ConfigLoader:
         except ValidationError as e:
             raise ValueError(f"System configuration validation for {system_config_path} failed:\n{e}")
 
-        if model_config_path is not None and model_config_path.exists():
-            try:
-                with open(model_config_path, 'r') as f:
-                    model_raw_config = yaml.safe_load(f)
-            except yaml.YAMLError as e:
-                raise ValueError(f"Failed to parse YAML file {model_config_path}: {e}")
-            try:
-                model_config = ModelConfig(**model_raw_config)
-                self.logger.info(f"Successfully loaded model config from {model_config_path}")
-            except ValidationError as e:
-                raise ValueError(f"Model configuration validation for {model_config_path} failed:\n{e}")
+        try:
+            with open(model_config_path, 'r') as f:
+                model_raw_config = yaml.safe_load(f)
+        except yaml.YAMLError as e:
+            raise ValueError(f"Failed to parse YAML file {model_config_path}: {e}")
+        
+        try:
+            model_config = ModelConfig(**model_raw_config)
+            self.logger.info(f"Successfully loaded model config from {model_config_path}")
+        except ValidationError as e:
+            raise ValueError(f"Model configuration validation for {model_config_path} failed:\n{e}")
 
         return system_config, model_config
 
 
-def load_config(system_config_path: Union[str, Path], model_config_path: Optional[Union[str, Path]] = None) -> Tuple[SystemConfig, Optional[ModelConfig]]:
+def load_config(system_config_path: Union[str, Path], model_config_path: Union[str, Path]) -> Tuple[SystemConfig, ModelConfig]:
     """Convenience function to load and validate config."""
     config_loader = ConfigLoader()
     return config_loader.load_and_validate(system_config_path, model_config_path)

src/config/schemas.py

@@ -1,5 +1,5 @@
 from pydantic import BaseModel, Field, model_validator
-from typing import Self, Tuple, List, Optional, Literal
+from typing import Self, Tuple, List, Optional, Literal, Union
 import torch
 
 
@@ -40,59 +40,11 @@ class SystemConfig(BaseModel):
     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: 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: 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 (required for AdaFortiTran)"
-    )
+class BaseConfig(BaseModel):
+    """Base configuration class with device validation."""
+    
     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."""
@@ -152,4 +104,59 @@ class ModelConfig(BaseModel):
             f"Available devices: {available_devices}"
         )
 
+
+class ModelConfig(BaseConfig):
+    model_type: Literal["linear", "fortitran", "adafortitran"] = Field(
+        default="fortitran",
+        description="Type of model (linear, fortitran, or adafortitran)"
+    )
+    patch_size: Tuple[int, int] = Field(..., description="Patch size as (subcarriers_per_patch, symbols_per_patch)")
+    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: 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: 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 (required for AdaFortiTran)"
+    )
+
+    @model_validator(mode='after')
+    def validate_model_specific_requirements(self) -> Self:
+        """Validate model-specific configuration requirements."""
+        if self.model_type == "linear":
+            # Linear model only needs device, no additional validation required
+            pass
+        elif 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"
+                )
+        elif 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_config = {"extra": "forbid"}

src/main.py

@@ -44,6 +44,7 @@ from pathlib import Path
 from src.main.parser import parse_arguments
 from src.main.trainer import train
 from src.config import load_config
+from src.config.schemas import ModelConfig
 
 
 def setup_logging(log_level: str) -> None:
@@ -88,10 +89,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")
-        if model_config is not None:
-            logger.info(f"Model architecture: {model_config.num_layers} layers, {model_config.model_dim} dimensions")
+        if model_config.model_type == "linear":
+            logger.info(f"Linear model with device: {model_config.device}")
         else:
-            logger.info("Using Linear model (no model config required)")
+            logger.info(f"Model architecture: {model_config.num_layers} layers, {model_config.model_dim} dimensions")
         
         # Start training
         logger.info("Initializing training...")

src/main/trainer.py

@@ -69,7 +69,7 @@ class ModelTrainer:
 
     EXP_LR_GAMMA = 0.995
 
-    def __init__(self, system_config: SystemConfig, model_config: ModelConfig | None, args: TrainingArguments):
+    def __init__(self, system_config: SystemConfig, model_config: ModelConfig, args: TrainingArguments):
         """
         Initialize the ModelTrainer.
 
@@ -121,14 +121,10 @@ class ModelTrainer:
             Initialized model instance of the specified type
         """
         if self.args.model_name == "linear":
-            model = LinearEstimator(self.system_config, device=str(self.device))
+            model = LinearEstimator(self.system_config, self.model_config)
         elif self.args.model_name == "adafortitran":
-            if self.model_config is None:
-                raise ValueError("model_config must be provided for AdaFortiTranEstimator.")
             model = AdaFortiTranEstimator(self.system_config, self.model_config)
         elif self.args.model_name == "fortitran":
-            if self.model_config is None:
-                raise ValueError("model_config must be provided for FortiTranEstimator.")
             model = FortiTranEstimator(self.system_config, self.model_config)
         else:
             raise ValueError(f"Unknown model name: {self.args.model_name}")
@@ -406,7 +402,7 @@ class ModelTrainer:
         self.writer.close()
 
 
-def train(system_config: SystemConfig, model_config: ModelConfig | None, args: TrainingArguments) -> None:
+def train(system_config: SystemConfig, model_config: ModelConfig, args: TrainingArguments) -> None:
     """
     Train an OFDM channel estimation model.
 

src/models/blocks/enhancers.py

@@ -23,9 +23,9 @@ class ConvEnhancer(nn.Module):
         """Forward pass through the convolutional enhancement network.
 
         Args:
-            x (torch.Tensor): Input tensor of shape (batch_size, 1, height, width)
+            x (torch.Tensor): Input tensor of shape (batch_size, 1, num_subcarriers, num_symbols)
 
         Returns:
-            torch.Tensor: Enhanced tensor of shape (batch_size, 1, height, width)
+            torch.Tensor: Enhanced tensor of shape (batch_size, 1, num_subcarriers, num_symbols)
         """
         return self.conv_block(x)

src/models/blocks/patch_processors.py

@@ -15,7 +15,7 @@ class PatchEmbedding(nn.Module):
         """Initialize the PatchEmbedding layer.
 
         Args:
-            patch_size: Size of patches to extract (height, width)
+            patch_size: Size of patches to extract (subcarriers_per_patch, symbols_per_patch)
         """
         super().__init__()
         self.patch_size = patch_size
@@ -25,11 +25,11 @@ class PatchEmbedding(nn.Module):
         """Transform input tensor into patch embeddings.
 
         Args:
-            x: Input tensor of shape (batch_size, height, width)
+            x: Input tensor of shape (batch_size, num_subcarriers, num_symbols)
 
         Returns:
             Tensor of shape (batch_size, num_patches, patch_size[0]*patch_size[1])
-            where num_patches = (height // patch_size[0]) * (width // patch_size[1])
+            where num_patches = (num_subcarriers // patch_size[0]) * (num_symbols // patch_size[1])
         """
         x = self.unfold(torch.unsqueeze(x, dim=1))
         return torch.permute(x, dims=(0, 2, 1))
@@ -46,8 +46,8 @@ class InversePatchEmbedding(nn.Module):
         """Initialize the InversePatchEmbedding layer.
 
         Args:
-            output_size: Size of output matrix (height, width)
-            patch_size: Size of input patches (height, width)
+            output_size: Size of output matrix (num_subcarriers, num_symbols)
+            patch_size: Size of input patches (subcarriers_per_patch, symbols_per_patch)
         """
         super().__init__()
         self.fold = nn.Fold(
@@ -64,7 +64,7 @@ class InversePatchEmbedding(nn.Module):
               where num_patches = (output_size[0] // patch_size[0]) * (output_size[1] // patch_size[1])
 
         Returns:
-            Tensor of shape (batch_size, output_size[0], output_size[1])
+            Tensor of shape (batch_size, num_subcarriers, num_symbols)
         """
         x = torch.permute(x, dims=(0, 2, 1))
         x = self.fold(x)

src/models/fortitran.py

@@ -4,8 +4,7 @@ import logging
 from typing import Tuple, List, Optional
 
 from src.config.schemas import SystemConfig, ModelConfig
-from src.models.blocks import ConvEnhancer, PatchEmbedding, InversePatchEmbedding, TransformerEncoderForChannels, \
-    ChannelAdapter
+from src.models.blocks import ConvEnhancer, PatchEmbedding, InversePatchEmbedding, TransformerEncoderForChannels, ChannelAdapter
 
 
 class BaseFortiTranEstimator(nn.Module):
@@ -13,11 +12,11 @@ class BaseFortiTranEstimator(nn.Module):
     Base Hybrid CNN-Transformer Channel Estimator for OFDM Systems.
 
     This model performs channel estimation by:
-    1. Upsampling pilot symbols to full OFDM grid size
-    2. Applying convolutional enhancement for spatial features
+    1. Upsampling pilot symbols to full OFDM grid size (with linear layer)
+    2. Applying convolutional enhancement for subcarrier-symbol features
     3. Converting to patch embeddings for transformer processing
     4. Using transformer encoder to capture long-range dependencies
-    5. Reconstructing spatial representation and applying residual connections
+    5. Reconstructing subcarrier-symbol representation and applying residual connections
     6. Final convolutional refinement for high-quality channel estimates
     """
 
@@ -29,7 +28,7 @@ class BaseFortiTranEstimator(nn.Module):
         Args:
             system_config: OFDM system configuration (subcarriers, symbols, pilot arrangement)
             model_config: Model architecture configuration (patch size, layers, etc.)
-            use_channel_adaptation: Whether to enable channel adaptation features
+            use_channel_adaptation: Whether to enable channel adaptation features (disabled for FortiTran)
         """
         super().__init__()
 
@@ -73,11 +72,13 @@ class BaseFortiTranEstimator(nn.Module):
                 self.model_config.patch_size[0] * self.model_config.patch_size[1]
         )
 
-        # Adaptive patch length (only used if channel adaptation is enabled)
+        # Transformer input dimension (includes channel tokens if adaptation is enabled)
         if self.use_channel_adaptation:
-            self.adaptive_patch_length = self.patch_length + self.model_config.adaptive_token_length
+            if self.model_config.adaptive_token_length is None:
+                raise ValueError("adaptive_token_length must be set when channel adaptation is enabled")
+            self.transformer_input_dim = self.patch_length + self.model_config.adaptive_token_length
         else:
-            self.adaptive_patch_length = self.patch_length
+            self.transformer_input_dim = self.patch_length
 
     def _build_architecture(self) -> None:
         """Construct the model architecture components."""
@@ -92,14 +93,19 @@ class BaseFortiTranEstimator(nn.Module):
 
         # 4. Channel adapter (conditional on use_channel_adaptation)
         if self.use_channel_adaptation:
-            self.channel_adapter = ChannelAdapter(self.model_config.channel_adaptivity_hidden_sizes)
+            if self.model_config.channel_adaptivity_hidden_sizes is None:
+                raise ValueError("channel_adaptivity_hidden_sizes must be set when channel adaptation is enabled")
+            # Convert list to tuple as expected by ChannelAdapter (exactly 3 values)
+            hidden_sizes = tuple(self.model_config.channel_adaptivity_hidden_sizes)
+            if len(hidden_sizes) != 3:
+                raise ValueError("channel_adaptivity_hidden_sizes must have exactly 3 values")
+            self.channel_adapter = ChannelAdapter(hidden_sizes)
 
         # 5. Transformer encoder for sequence modeling
-        transformer_input_dim = self.adaptive_patch_length if self.use_channel_adaptation else self.patch_length
         transformer_output_dim = self.patch_length  # Always output standard patch length
 
         self.transformer_encoder = TransformerEncoderForChannels(
-            input_dim=transformer_input_dim,
+            input_dim=self.transformer_input_dim,
             output_dim=transformer_output_dim,
             model_dim=self.model_config.model_dim,
             num_head=self.model_config.num_head,
@@ -189,7 +195,7 @@ class BaseFortiTranEstimator(nn.Module):
         """
         batch_size = x.shape[0]
 
-        # Flatten spatial dimensions for linear upsampling
+        # Flatten subcarrier and symbol dimensions for linear upsampling
         if x.dim() > 2:
             x = x.view(batch_size, -1)
 
@@ -215,7 +221,7 @@ class BaseFortiTranEstimator(nn.Module):
         # Stage 5: Transformer processing for long-range dependencies
         transformer_output = self.transformer_encoder(transformer_input)
 
-        # Stage 6: Reconstruct spatial representation
+        # Stage 6: Reconstruct subcarrier-symbol representation
         reconstructed = self.patch_reconstructor(transformer_output)
 
         # Stage 7: Apply residual connection
@@ -235,7 +241,7 @@ class BaseFortiTranEstimator(nn.Module):
             'pilot_size': self.pilot_size,
             'patch_size': self.model_config.patch_size,
             'patch_length': self.patch_length,
-            'adaptive_patch_length': self.adaptive_patch_length,
+            'transformer_input_dim': self.transformer_input_dim,
             'model_dim': self.model_config.model_dim,
             'num_layers': self.model_config.num_layers,
             'device': str(self.device),

src/models/linear.py

@@ -10,43 +10,47 @@ import logging
 import torch
 import torch.nn as nn
 
-from src.config.schemas import SystemConfig
+from src.config.schemas import SystemConfig, ModelConfig
 
 
 class LinearEstimator(nn.Module):
     """Learned MMSE estimator.
 
+    Find W such that W*h_pilot = h_hat, where h_hat is the estimated channel by stochastic gradient descent on |h_hat - h_ideal|^2
+
     Attributes:
         device (torch.device): Target device for computation
-        config (SystemConfig): Validated configuration object
-        ofdm_size (Tuple[int, int]): Dimensions of OFDM frame as (height, width)
-            height (int): number of sub-carriers
-            width (int): number of OFDM symbols
-        pilot_size (Tuple[int, int]): Dimensions of pilot signal as (height, width)
-            height (int): number of pilots across sub-carriers
-            width (int): number of pilots across OFDM symbols
+        system_config (SystemConfig): Validated configuration object for OFDM system parameters
+        model_config (ModelConfig): Validated configuration object for model parameters
+        ofdm_size (Tuple[int, int]): Dimensions of OFDM frame as (num_subcarriers, num_symbols)
+            num_subcarriers (int): number of sub-carriers
+            num_symbols (int): number of OFDM symbols
+        pilot_size (Tuple[int, int]): Dimensions of pilot signal as (num_subcarriers, num_symbols)
+            num_subcarriers (int): number of pilots across sub-carriers
+            num_symbols (int): number of pilots across OFDM symbols
     """
 
-    def __init__(self, config: SystemConfig, device: str = "cpu") -> None:
+    def __init__(self, system_config: SystemConfig, model_config: ModelConfig) -> None:
         """Initialize the MMSE estimator.
 
         Args:
-            config: Validated SystemConfig object containing OFDM system parameters
-            device: Device to use for computation (cpu, cuda, etc.)
+            system_config: Validated SystemConfig object containing OFDM system parameters
+            model_config: Validated ModelConfig object containing model parameters
         """
         super().__init__()
 
-        self.config = config
-        self.device = torch.device(device)
+        self.system_config = system_config
+        self.model_config = model_config
+        self.device = torch.device(model_config.device)
         self.logger = logging.getLogger(__name__)
 
         # Extract dimensions from validated config
-        self.ofdm_size = (config.ofdm.num_scs, config.ofdm.num_symbols)
-        self.pilot_size = (config.pilot.num_scs, config.pilot.num_symbols)
+        self.ofdm_size = (system_config.ofdm.num_scs, system_config.ofdm.num_symbols)
+        self.pilot_size = (system_config.pilot.num_scs, system_config.pilot.num_symbols)
 
         # Calculate feature dimensions
-        in_feature_dim = config.pilot.num_scs * config.pilot.num_symbols
-        out_feature_dim = config.ofdm.num_scs * config.ofdm.num_symbols
+        in_feature_dim = system_config.pilot.num_scs * system_config.pilot.num_symbols
+        out_feature_dim = system_config.ofdm.num_scs * system_config.ofdm.num_symbols
 
         self.logger.info(f"Initializing LinearEstimator:")
         self.logger.info(f"  OFDM size: {self.ofdm_size}")
@@ -70,7 +74,7 @@ class LinearEstimator(nn.Module):
             Estimated OFDM signal tensor with shape
             (batch_size, ofdm_size[0], ofdm_size[1])
         """
-        # pytorch does nothin if input is already on correct device
+        # pytorch does nothing if input is already on correct device
         x = x.to(self.device)
         self.logger.debug(f"Input shape: {x.size()}")
 
@@ -95,14 +99,6 @@ class LinearEstimator(nn.Module):
 
         return x
 
-    def get_config(self) -> SystemConfig:
-        """Get the configuration used by this estimator.
-
-        Returns:
-            SystemConfig: The configuration object
-        """
-        return self.config
-
     def __repr__(self) -> str:
         """String representation of the estimator."""
         return (