combined adafortitran and fortitran with basefortitran to minimize code repetition. Added dataset.py

requirements.txt

@@ -1,3 +1,4 @@
 torch
 pydantic
-yaml
+yaml
+scipy

src/config/__init__.py

@@ -0,0 +1 @@
+from src.config.schemas import ModelConfig, SystemConfig

src/data/dataset.py

@@ -0,0 +1,238 @@
+"""Module for loading and processing .mat files containing channel estimates for PyTorch."""
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, List, Optional, Tuple, Union
+
+import scipy.io as sio
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+from src.utils import extract_values
+
+__all__ = ['MatDataset', 'get_test_dataloaders']
+
+
+@dataclass
+class PilotDimensions:
+    """Container for pilot signal dimensions.
+
+    Stores and validates the dimensions of pilot signals used in channel estimation.
+
+    Attributes:
+        num_subcarriers: Number of subcarriers in the pilot signal
+        num_ofdm_symbols: Number of OFDM symbols in the pilot signal
+    """
+    num_subcarriers: int
+    num_ofdm_symbols: int
+
+    def __post_init__(self):
+        """Validate dimensions after initialization.
+
+        Raises:
+            ValueError: If either dimension is not a positive integer
+        """
+        if self.num_subcarriers <= 0 or self.num_ofdm_symbols <= 0:
+            raise ValueError("Pilot dimensions must be positive integers")
+
+    def as_tuple(self) -> Tuple[int, int]:
+        """Return dimensions as a tuple.
+
+        Returns:
+            Tuple of (num_subcarriers, num_ofdm_symbols)
+        """
+        return self.num_subcarriers, self.num_ofdm_symbols
+
+
+class MatDataset(Dataset):
+    """Dataset for loading and formatting .mat files containing channel estimates.
+
+    Processes .mat files containing channel estimation data and converts them into
+    PyTorch complex tensors for channel estimation tasks.
+    """
+
+    def __init__(
+            self,
+            data_dir: Union[str, Path],
+            pilot_dims: List[int],
+            transform: Optional[Callable] = None
+    ) -> None:
+        """Initialize the MatDataset.
+
+        Args:
+            data_dir: Path to the directory containing the dataset.
+            pilot_dims: Dimensions of pilot data as [num_subcarriers, num_ofdm_symbols].
+            transform: Optional transformation to apply to samples.
+
+        Raises:
+            ValueError: If pilot dimensions are invalid.
+            FileNotFoundError: If data_dir doesn't exist.
+        """
+        self.data_dir = Path(data_dir)
+        self.pilot_dims = PilotDimensions(*pilot_dims)
+        self.transform = transform
+
+        if not self.data_dir.exists():
+            raise FileNotFoundError(f"Data directory not found: {self.data_dir}")
+
+        self.file_list = list(self.data_dir.glob("*.mat"))
+        if not self.file_list:
+            raise ValueError(f"No .mat files found in {self.data_dir}")
+
+    def __len__(self) -> int:
+        """Return the total number of files in the dataset.
+
+        Returns:
+            Integer count of .mat files in the dataset directory
+        """
+        return len(self.file_list)
+
+    def _process_channel_data(
+            self,
+            h_ideal: torch.Tensor,
+            mat_data: dict
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Process channel data and extract pilot values from LS estimates.
+
+        Extracts pilot values from LS channel estimates with zero entries removed,
+        returning complex-valued tensors for both estimate and ground truth.
+
+        Args:
+            h_ideal: Ground truth channel tensor
+            mat_data: Loaded .mat file data
+
+        Returns:
+            Tuple of (pilot LS estimate, ground truth channel)
+
+        Raises:
+            ValueError: If the data format is unexpected or processing fails
+        """
+        try:
+            # Extract LS channel estimate with zero entries
+            hzero_ls = torch.tensor(mat_data['H'][:, :, 1], dtype=torch.cfloat)
+
+            # Remove zero entries, keep only pilot values
+            zero_complex = torch.complex(torch.tensor(0.0), torch.tensor(0.0))
+            hp_ls = hzero_ls[hzero_ls != zero_complex]
+
+            # Validate expected number of pilot values
+            expected_pilots = self.pilot_dims.num_subcarriers * self.pilot_dims.num_ofdm_symbols
+            if hp_ls.numel() != expected_pilots:
+                raise ValueError(
+                    f"Expected {expected_pilots} pilot values, got {hp_ls.numel()}"
+                )
+
+            # Reshape to pilot grid dimensions [subcarriers, symbols]
+            hp_ls = hp_ls.unsqueeze(dim=1).view(
+                self.pilot_dims.num_ofdm_symbols,
+                self.pilot_dims.num_subcarriers
+            ).t()
+
+            return hp_ls, h_ideal
+
+        except Exception as e:
+            raise ValueError(f"Error processing channel data: {str(e)}")
+
+    def __getitem__(
+            self,
+            idx: int
+    ) -> Tuple[torch.Tensor, torch.Tensor, Tuple]:
+        """Load and process a .mat file at the given index.
+
+        Args:
+            idx: Index of the file to load.
+
+        Returns:
+            Tuple containing:
+                - Pilot LS channel estimate (complex tensor)
+                - Ground truth channel estimate (complex tensor)
+                - Metadata extracted from filename
+
+        Raises:
+            ValueError: If file format is invalid or processing fails.
+            IndexError: If idx is out of range.
+        """
+        if not 0 <= idx < len(self):
+            raise IndexError(f"Index {idx} out of range for dataset of size {len(self)}")
+
+        try:
+            # Load .mat file
+            mat_data = sio.loadmat(self.file_list[idx])
+            if 'H' not in mat_data or mat_data['H'].shape[-1] < 3:
+                raise ValueError("Invalid .mat file format: missing required data")
+
+            # Extract ground truth channel
+            h_ideal = torch.tensor(mat_data['H'][:, :, 0], dtype=torch.cfloat)
+
+            # Process channel data to extract pilot estimates
+            h_est, h_ideal = self._process_channel_data(h_ideal, mat_data)
+
+            # Extract metadata from filename
+            meta_data = extract_values(self.file_list[idx].name)
+            if meta_data is None:
+                raise ValueError(f"Unrecognized filename format: {self.file_list[idx].name}")
+
+            # Apply optional transforms
+            if self.transform:
+                h_est = self.transform(h_est)
+                h_ideal = self.transform(h_ideal)
+
+            return h_est, h_ideal, meta_data
+
+        except Exception as e:
+            raise ValueError(f"Error processing file {self.file_list[idx]}: {str(e)}")
+
+
+def get_test_dataloaders(
+        dataset_dir: Union[str, Path],
+        params: dict
+) -> List[Tuple[str, DataLoader]]:
+    """Create DataLoaders for each subdirectory in the dataset directory.
+
+    Automatically discovers and creates appropriate DataLoader instances for
+    all subdirectories in the specified dataset directory, useful for testing
+    across multiple test conditions or scenarios.
+
+    Args:
+        dataset_dir: Path to main directory containing dataset subdirectories
+        params: Configuration parameters including:
+            - pilot_dims: List of [num_subcarriers, num_ofdm_symbols]
+            - batch_size: Number of samples per batch
+
+    Returns:
+        List of tuples containing (subdirectory_name, corresponding_dataloader)
+
+    Raises:
+        FileNotFoundError: If dataset_dir doesn't exist
+        ValueError: If params are invalid or no valid subdirectories are found
+    """
+    dataset_dir = Path(dataset_dir)
+    if not dataset_dir.exists():
+        raise FileNotFoundError(f"Dataset directory not found: {dataset_dir}")
+
+    if not isinstance(params, dict) or "pilot_dims" not in params or "batch_size" not in params:
+        raise ValueError("params must be a dict containing 'pilot_dims' and 'batch_size'")
+
+    subdirs = [d for d in dataset_dir.iterdir() if d.is_dir()]
+    if not subdirs:
+        raise ValueError(f"No subdirectories found in {dataset_dir}")
+
+    test_datasets = [
+        (
+            subdir.name,
+            MatDataset(
+                subdir,
+                params["pilot_dims"]
+            )
+        )
+        for subdir in subdirs
+    ]
+
+    return [
+        (name, DataLoader(
+            dataset,
+            batch_size=params["batch_size"],
+            shuffle=False,
+            num_workers=0
+        ))
+        for name, dataset in test_datasets
+    ]

src/models/adafortitran.py

@@ -1,25 +1,14 @@
-import torch
-from torch import nn
-import logging
-from typing import Tuple, List
+from .fortitran import BaseFortiTranEstimator
+from src.config import SystemConfig, ModelConfig
 
-from src.config.schemas import SystemConfig, ModelConfig
-from src.models.blocks import ConvEnhancer, PatchEmbedding, InversePatchEmbedding, TransformerEncoderForChannels, ChannelAdapter
-
-
-class AdaFortiTranEstimator(nn.Module):
 
+class AdaFortiTranEstimator(BaseFortiTranEstimator):
     """
-    Hybrid CNN-Transformer Channel Estimator for OFDM Systems with channel adaptation.
+    Adaptive Hybrid CNN-Transformer Channel Estimator for OFDM Systems with channel adaptation.
 
-    This model performs channel estimation by:
-    1. Upsampling pilot symbols to full OFDM grid size
-    2. Applying convolutional enhancement for spatial features
-    3. Converting to patch embeddings for transformer processing
-    4. Concatenating channel statistics priors to channel patches
-    5. Using transformer encoder to capture long-range dependencies
-    6. Reconstructing spatial representation and applying residual connections
-    7. Final convolutional refinement for high-quality channel estimates
+    This model extends the base estimator with channel adaptation capabilities,
+    incorporating channel conditions (SNR, delay spread, Doppler shift) into
+    the estimation process through conditional attention mechanisms.
     """
 
     def __init__(self, system_config: SystemConfig, model_config: ModelConfig) -> None:
@@ -30,187 +19,4 @@ class AdaFortiTranEstimator(nn.Module):
             system_config: OFDM system configuration (subcarriers, symbols, pilot arrangement)
             model_config: Model architecture configuration (patch size, layers, etc.)
         """
-        super().__init__()
-
-        self.system_config = system_config
-        self.model_config = model_config
-        self.device = torch.device(model_config.device)
-        self.logger = logging.getLogger(self.__class__.__name__)
-
-        # Cache key dimensions for efficiency
-        self._setup_dimensions()
-
-        # Initialize model components
-        self._build_architecture()
-
-        # Move model to specified device
-        self.to(self.device)
-
-        self._log_initialization_info()
-
-    def _setup_dimensions(self) -> None:
-        """Calculate and cache key dimensions from configuration."""
-        # OFDM grid dimensions
-        self.ofdm_size = (
-            self.system_config.ofdm.num_scs,
-            self.system_config.ofdm.num_symbols
-        )
-
-        # Pilot arrangement dimensions
-        self.pilot_size = (
-            self.system_config.pilot.num_scs,
-            self.system_config.pilot.num_symbols
-        )
-
-        # Feature dimensions for linear layers
-        self.pilot_features = self.pilot_size[0] * self.pilot_size[1]
-        self.ofdm_features = self.ofdm_size[0] * self.ofdm_size[1]
-
-        # Patch processing dimensions
-        self.patch_length = (
-                self.model_config.patch_size[0] * self.model_config.patch_size[1]
-        )
-
-        self.adaptive_patch_length = self.patch_length + self.model_config.adaptive_token_length
-
-    def _build_architecture(self) -> None:
-        """Construct the model architecture components."""
-        # 1. Pilot-to-OFDM upsampling
-        self.pilot_upsampler = nn.Linear(self.pilot_features, self.ofdm_features)
-        # 2. Initial convolutional enhancement
-        self.initial_enhancer = ConvEnhancer()
-
-        # 3. Patch embedding for transformer processing
-        self.patch_embedder = PatchEmbedding(self.model_config.patch_size)
-
-        # 4. Channel adapter for conditional attention
-        self.channel_adapter = ChannelAdapter(self.model_config.channel_adaptivity_hidden_sizes)
-
-        # 5. Transformer encoder for sequence modeling
-        self.transformer_encoder = TransformerEncoderForChannels(
-            input_dim=self.adaptive_patch_length,
-            output_dim=self.patch_length,
-            model_dim=self.model_config.model_dim,
-            num_head=self.model_config.num_head,
-            activation=self.model_config.activation,
-            dropout=self.model_config.dropout,
-            num_layers=self.model_config.num_layers,
-            max_len=self.model_config.max_seq_len,
-            pos_encoding_type=self.model_config.pos_encoding_type
-        )
-
-        # 6. Patch reconstruction
-        self.patch_reconstructor = InversePatchEmbedding(
-            self.ofdm_size,
-            self.model_config.patch_size
-        )
-
-        # 7. Final convolutional refinement
-        self.final_refiner = ConvEnhancer()
-
-    def _log_initialization_info(self) -> None:
-        """Log model initialization details."""
-        self.logger.info("AdaFortiTranEstimator initialized successfully:")
-        self.logger.info(f"  OFDM grid: {self.ofdm_size[0]}×{self.ofdm_size[1]} = {self.ofdm_features} elements")
-        self.logger.info(f"  Pilot grid: {self.pilot_size[0]}×{self.pilot_size[1]} = {self.pilot_features} elements")
-        self.logger.info(f"  Patch size: {self.model_config.patch_size}")
-        self.logger.info(f"  Model dimension: {self.model_config.model_dim}")
-        self.logger.info(f"  Transformer layers: {self.model_config.num_layers}")
-        self.logger.info(f"  Device: {self.device}")
-
-        total_params = sum(p.numel() for p in self.parameters())
-        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
-        self.logger.info(f"  Total parameters: {total_params:,}")
-        self.logger.info(f"  Trainable parameters: {trainable_params:,}")
-
-    def forward(self, pilot_symbols: torch.Tensor, meta_data: Tuple) -> torch.Tensor:
-        """
-        Forward pass for channel estimation.
-
-        Args:
-            pilot_symbols: Complex pilot symbols of shape [batch, pilot_scs, pilot_symbols]
-            meta_data: TODO: Add complete type annotation.
-
-        Returns:
-            Estimated channel matrix of shape [batch, ofdm_scs, ofdm_symbols]
-        """
-
-        # Extract and move channel conditions to device
-        _, snr, delay_spread, max_dop_shift, _, _ = meta_data
-        channel_conditions = [
-            tensor.to(self.device)
-            for tensor in (snr, delay_spread, max_dop_shift)
-        ]
-
-        # Ensure input is on correct device
-        pilot_symbols = pilot_symbols.to(self.device)
-
-        # Process real and imaginary parts separately
-        real_estimate = self._forward_real_valued(pilot_symbols.real, channel_conditions)
-        imag_estimate = self._forward_real_valued(pilot_symbols.imag, channel_conditions)
-
-        # Combine into complex tensor
-        channel_estimate = torch.complex(real_estimate, imag_estimate)
-
-        return channel_estimate
-
-    def _forward_real_valued(self, x: torch.Tensor, channel_conditions: List[torch.Tensor]) -> torch.Tensor:
-        """
-        Process real-valued input through the estimation pipeline.
-
-        Args:
-            x: Real-valued input tensor [batch, pilot_features] or [batch, pilot_scs, pilot_symbols]
-
-        Returns:
-            Real-valued channel estimate [batch, ofdm_scs, ofdm_symbols]
-        """
-        batch_size = x.shape[0]
-
-        # Flatten spatial dimensions for linear upsampling
-        if x.dim() > 2:
-            x = x.view(batch_size, -1)
-
-        # Stage 1: Upsample from pilot grid to OFDM grid
-        upsampled = self.pilot_upsampler(x)
-
-        # Reshape for convolutional processing
-        upsampled_2d = upsampled.view(batch_size, 1, *self.ofdm_size)
-
-        # Stage 2: Initial convolutional enhancement
-        conv_enhanced = torch.squeeze(self.initial_enhancer(upsampled_2d), dim=1)
-
-        # Stage 3: Convert to patch embeddings
-        patch_embeddings = self.patch_embedder(conv_enhanced)
-
-        # Stage 4: Get conditioned channel encodings
-        encoded_channel_condition = self.channel_adapter(*channel_conditions)
-        conditioned_channel_encodings = torch.cat((patch_embeddings, encoded_channel_condition), dim=2)
-
-        # Stage 5: Transformer processing for long-range dependencies
-        transformer_output = self.transformer_encoder(conditioned_channel_encodings)
-
-        # Stage 6: Reconstruct spatial representation
-        reconstructed = self.patch_reconstructor(transformer_output)
-
-        # Stage 7: Apply residual connection
-        residual_combined = conv_enhanced + reconstructed
-
-        # Stage 8: Final convolutional refinement
-        refined_output = torch.squeeze(self.final_refiner(torch.unsqueeze(residual_combined, dim=1)), dim=1)
-
-        return refined_output
-
-    def get_model_info(self) -> dict:
-        """Return model configuration and statistics."""
-        return {
-            'model_name': self.__class__.__name__,
-            'ofdm_size': self.ofdm_size,
-            'pilot_size': self.pilot_size,
-            'patch_size': self.model_config.patch_size,
-            'patch_length': self.patch_length,
-            'model_dim': self.model_config.model_dim,
-            'num_layers': self.model_config.num_layers,
-            'device': str(self.device),
-            'total_parameters': sum(p.numel() for p in self.parameters()),
-            'trainable_parameters': sum(p.numel() for p in self.parameters() if p.requires_grad)
-        }
+        super().__init__(system_config, model_config, use_channel_adaptation=True)

src/models/fortitran.py

@@ -1,14 +1,16 @@
 import torch
 from torch import nn
 import logging
+from typing import Tuple, List, Optional
 
-from src.config.schemas import SystemConfig, ModelConfig
-from src.models.blocks import ConvEnhancer, PatchEmbedding, InversePatchEmbedding, TransformerEncoderForChannels
+from src.config import SystemConfig, ModelConfig
+from src.models.blocks import ConvEnhancer, PatchEmbedding, InversePatchEmbedding, TransformerEncoderForChannels, \
+    ChannelAdapter
 
 
-class FortiTranEstimator(nn.Module):
+class BaseFortiTranEstimator(nn.Module):
     """
-    Hybrid CNN-Transformer Channel Estimator for OFDM Systems.
+    Base Hybrid CNN-Transformer Channel Estimator for OFDM Systems.
 
     This model performs channel estimation by:
     1. Upsampling pilot symbols to full OFDM grid size
@@ -19,18 +21,21 @@ class FortiTranEstimator(nn.Module):
     6. Final convolutional refinement for high-quality channel estimates
     """
 
-    def __init__(self, system_config: SystemConfig, model_config: ModelConfig) -> None:
+    def __init__(self, system_config: SystemConfig, model_config: ModelConfig,
+                 use_channel_adaptation: bool = False) -> None:
         """
-        Initialize the FortiTranEstimator.
+        Initialize the BaseFortiTranEstimator.
 
         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
         """
         super().__init__()
 
         self.system_config = system_config
         self.model_config = model_config
+        self.use_channel_adaptation = use_channel_adaptation
         self.device = torch.device(model_config.device)
         self.logger = logging.getLogger(self.__class__.__name__)
 
@@ -68,41 +73,57 @@ class FortiTranEstimator(nn.Module):
                 self.model_config.patch_size[0] * self.model_config.patch_size[1]
         )
 
+        # Adaptive patch length (only used if channel adaptation is enabled)
+        if self.use_channel_adaptation:
+            self.adaptive_patch_length = self.patch_length + self.model_config.adaptive_token_length
+        else:
+            self.adaptive_patch_length = self.patch_length
+
     def _build_architecture(self) -> None:
         """Construct the model architecture components."""
         # 1. Pilot-to-OFDM upsampling
         self.pilot_upsampler = nn.Linear(self.pilot_features, self.ofdm_features)
+
         # 2. Initial convolutional enhancement
         self.initial_enhancer = ConvEnhancer()
 
         # 3. Patch embedding for transformer processing
         self.patch_embedder = PatchEmbedding(self.model_config.patch_size)
 
-        # 4. Transformer encoder for sequence modeling
+        # 4. Channel adapter (conditional on use_channel_adaptation)
+        if self.use_channel_adaptation:
+            self.channel_adapter = ChannelAdapter(self.model_config.channel_adaptivity_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=self.patch_length,
-            output_dim=self.patch_length,
+            input_dim=transformer_input_dim,
+            output_dim=transformer_output_dim,
             model_dim=self.model_config.model_dim,
             num_head=self.model_config.num_head,
             activation=self.model_config.activation,
             dropout=self.model_config.dropout,
             num_layers=self.model_config.num_layers,
             max_len=self.model_config.max_seq_len,
-            pos_encoding_type=self.model_config.pos_encoding_type,
+            pos_encoding_type=self.model_config.pos_encoding_type
         )
 
-        # 5. Patch reconstruction
+        # 6. Patch reconstruction
         self.patch_reconstructor = InversePatchEmbedding(
             self.ofdm_size,
             self.model_config.patch_size
         )
 
-        # 6. Final convolutional refinement
+        # 7. Final convolutional refinement
         self.final_refiner = ConvEnhancer()
 
     def _log_initialization_info(self) -> None:
         """Log model initialization details."""
-        self.logger.info("FortiTranEstimator initialized successfully:")
+        adaptation_status = "enabled" if self.use_channel_adaptation else "disabled"
+        self.logger.info(f"{self.__class__.__name__} initialized successfully:")
+        self.logger.info(f"  Channel adaptation: {adaptation_status}")
         self.logger.info(f"  OFDM grid: {self.ofdm_size[0]}×{self.ofdm_size[1]} = {self.ofdm_features} elements")
         self.logger.info(f"  Pilot grid: {self.pilot_size[0]}×{self.pilot_size[1]} = {self.pilot_features} elements")
         self.logger.info(f"  Patch size: {self.model_config.patch_size}")
@@ -115,34 +136,53 @@ class FortiTranEstimator(nn.Module):
         self.logger.info(f"  Total parameters: {total_params:,}")
         self.logger.info(f"  Trainable parameters: {trainable_params:,}")
 
-    def forward(self, pilot_symbols: torch.Tensor) -> torch.Tensor:
+    def forward(self, pilot_symbols: torch.Tensor, meta_data: Optional[Tuple] = None) -> torch.Tensor:
         """
         Forward pass for channel estimation.
 
         Args:
             pilot_symbols: Complex pilot symbols of shape [batch, pilot_scs, pilot_symbols]
+            meta_data: Channel conditions (only used if channel adaptation is enabled)
 
         Returns:
             Estimated channel matrix of shape [batch, ofdm_scs, ofdm_symbols]
         """
+        # Validate inputs based on adaptation mode
+        if self.use_channel_adaptation and meta_data is None:
+            raise ValueError("meta_data is required when channel adaptation is enabled")
+
+        if not self.use_channel_adaptation and meta_data is not None:
+            self.logger.warning("meta_data provided but channel adaptation is disabled - ignoring meta_data")
+
+        # Extract channel conditions if adaptation is enabled
+        channel_conditions = None
+        if self.use_channel_adaptation and meta_data is not None:
+            _, snr, delay_spread, max_dop_shift, _, _ = meta_data
+            channel_conditions = [
+                tensor.to(self.device)
+                for tensor in (snr, delay_spread, max_dop_shift)
+            ]
+
         # Ensure input is on correct device
         pilot_symbols = pilot_symbols.to(self.device)
 
         # Process real and imaginary parts separately
-        real_estimate = self._forward_real_valued(pilot_symbols.real)
-        imag_estimate = self._forward_real_valued(pilot_symbols.imag)
+        real_estimate = self._forward_real_valued(pilot_symbols.real, channel_conditions)
+        imag_estimate = self._forward_real_valued(pilot_symbols.imag, channel_conditions)
 
         # Combine into complex tensor
         channel_estimate = torch.complex(real_estimate, imag_estimate)
 
         return channel_estimate
 
-    def _forward_real_valued(self, x: torch.Tensor) -> torch.Tensor:
+    def _forward_real_valued(self, x: torch.Tensor,
+                             channel_conditions: Optional[List[torch.Tensor]] = None) -> torch.Tensor:
         """
         Process real-valued input through the estimation pipeline.
 
         Args:
             x: Real-valued input tensor [batch, pilot_features] or [batch, pilot_scs, pilot_symbols]
+            channel_conditions: Channel conditions for adaptation (optional)
 
         Returns:
             Real-valued channel estimate [batch, ofdm_scs, ofdm_symbols]
@@ -165,16 +205,23 @@ class FortiTranEstimator(nn.Module):
         # Stage 3: Convert to patch embeddings
         patch_embeddings = self.patch_embedder(conv_enhanced)
 
-        # Stage 4: Transformer processing for long-range dependencies
-        transformer_output = self.transformer_encoder(patch_embeddings)
+        # Stage 4: Apply channel adaptation if enabled
+        if self.use_channel_adaptation and channel_conditions is not None:
+            encoded_channel_condition = self.channel_adapter(*channel_conditions)
+            transformer_input = torch.cat((patch_embeddings, encoded_channel_condition), dim=2)
+        else:
+            transformer_input = patch_embeddings
+
+        # Stage 5: Transformer processing for long-range dependencies
+        transformer_output = self.transformer_encoder(transformer_input)
 
-        # Stage 5: Reconstruct spatial representation
+        # Stage 6: Reconstruct spatial representation
         reconstructed = self.patch_reconstructor(transformer_output)
 
-        # Stage 6: Apply residual connection
+        # Stage 7: Apply residual connection
         residual_combined = conv_enhanced + reconstructed
 
-        # Stage 7: Final convolutional refinement
+        # Stage 8: Final convolutional refinement
         refined_output = torch.squeeze(self.final_refiner(torch.unsqueeze(residual_combined, dim=1)), dim=1)
 
         return refined_output
@@ -183,13 +230,33 @@ class FortiTranEstimator(nn.Module):
         """Return model configuration and statistics."""
         return {
             'model_name': self.__class__.__name__,
+            'channel_adaptation': self.use_channel_adaptation,
             'ofdm_size': self.ofdm_size,
             'pilot_size': self.pilot_size,
             'patch_size': self.model_config.patch_size,
             'patch_length': self.patch_length,
+            'adaptive_patch_length': self.adaptive_patch_length,
             'model_dim': self.model_config.model_dim,
             'num_layers': self.model_config.num_layers,
             'device': str(self.device),
             'total_parameters': sum(p.numel() for p in self.parameters()),
             'trainable_parameters': sum(p.numel() for p in self.parameters() if p.requires_grad)
         }
+
+
+class FortiTranEstimator(BaseFortiTranEstimator):
+    """
+    Standard Hybrid CNN-Transformer Channel Estimator for OFDM Systems.
+
+    This is the base version without channel adaptation features.
+    """
+
+    def __init__(self, system_config: SystemConfig, model_config: ModelConfig) -> None:
+        """
+        Initialize the FortiTranEstimator.
+
+        Args:
+            system_config: OFDM system configuration (subcarriers, symbols, pilot arrangement)
+            model_config: Model architecture configuration (patch size, layers, etc.)
+        """
+        super().__init__(system_config, model_config, use_channel_adaptation=False)

src/utils.py

@@ -0,0 +1,56 @@
+"""Utility functions for OFDM channel estimation."""
+import re
+import torch
+
+def extract_values(file_name):
+    """
+    Extract channel information from a file name.
+
+    Parses file names with format:
+    '{number}_SNR-{snr}_DS-{delay_spread}_DOP-{doppler}_N-{pilot_freq}_{channel_type}.mat'
+
+    Args:
+        file_name: The file name from which to extract channel information
+
+    Returns:
+        tuple: A tuple containing:
+            - file_number (torch.Tensor): The file number
+            - snr (torch.Tensor): Signal-to-noise ratio value
+            - delay_spread (torch.Tensor): Delay spread value
+            - max_doppler_shift (torch.Tensor): Maximum Doppler shift value
+            - pilot_placement_frequency (torch.Tensor): Pilot placement frequency
+            - channel_type (list): The channel type
+
+    Raises:
+        ValueError: If the file name does not match the expected pattern
+    """
+    pattern = r'(\d+)_SNR-(\d+)_DS-(\d+)_DOP-(\d+)_N-(\d+)_([A-Z\-]+)\.mat'
+    match = re.match(pattern, file_name)
+    if match:
+        file_no = torch.tensor([int(match.group(1))], dtype=torch.float)
+        snr_value = torch.tensor([int(match.group(2))], dtype=torch.float)
+        ds_value = torch.tensor([int(match.group(3))], dtype=torch.float)
+        dop_value = torch.tensor([int(match.group(4))], dtype=torch.float)
+        n = torch.tensor([int(match.group(5))], dtype=torch.float)
+        channel_type = [match.group(6)]
+        return file_no, snr_value, ds_value, dop_value, n, channel_type
+    else:
+        raise ValueError("Cannot extract file information.")
+
+def concat_complex_channel(channel_matrix):
+    """
+    Convert a complex channel matrix into a real matrix by concatenating real and imaginary parts.
+
+    Transforms a complex tensor into a real-valued tensor by concatenating
+    the real and imaginary components along the specified dimension.
+
+    Args:
+        channel_matrix: Complex channel matrix
+
+    Returns:
+        Real-valued channel matrix with concatenated real and imaginary parts
+    """
+    real_channel_m = torch.real(channel_matrix)
+    imag_channel_m = torch.imag(channel_matrix)
+    cat_channel_m = torch.cat((real_channel_m, imag_channel_m), dim=1)
+    return cat_channel_m