augmentrum 0.0.1__py3-none-any.whl

augmentrum/sampling/coil_average_sampler.py

@@ -0,0 +1,274 @@
+####################################################################################################
+#                                   coil_average_sampler.py                                        #
+####################################################################################################
+#                                                                                                  #
+# Authors: J. P. Merkofer (j.p.merkofer@tue.nl)                                                    #
+#                                                                                                  #
+# Created: 2025-10-07                                                                              #
+#                                                                                                  #
+# Purpose: Implements sampling of coils and averages from raw MRS data.                            #
+#                                                                                                  #
+####################################################################################################
+
+
+#*************#
+#   imports   #
+#*************#
+import torch
+
+from fsl_mrs.core.nifti_mrs import split
+
+# own
+from augmentrum.processing.utils import update_processing_prov
+from augmentrum.core.base_module import BaseModule
+from augmentrum.core import Backend
+
+
+#**************************************************************************************************#
+#                                        Class CoilAverageSampler                                  #
+#**************************************************************************************************#
+#                                                                                                  #
+# Samples coils and averages from raw MRS data. Supports random/deterministic selection            #
+# and reweighting.                                                                                 #
+#                                                                                                  #
+#**************************************************************************************************#
+class CoilAverageSampler(BaseModule):
+    """
+    Samples coils and averages from raw MRS data.
+
+    Uses FSL-MRS `split` on NIfTI-MRS objects, so this module operates
+    natively on the NIFTI_LIST backend.  For any other backend the base
+    class automatically routes the call through process_nifti_list and
+    returns the result in the original backend format.
+    """
+
+    SUPPORTED_BACKENDS = [Backend.NIFTI_LIST]
+
+    def __init__(self, mode='random', n_coils=(1, None), n_averages=(1, None), reweight=False):
+        """
+        Initializes the sampler.
+
+        Args:
+            mode (str): 'random' or 'deterministic' sampling mode.
+            n_coils (tuple or int): Min/max number of coils to sample, or exact value.
+                                   Examples: (1, 8) = range, 4 = exactly 4
+            n_averages (tuple or int): Min/max number of averages to sample, or exact value.
+            reweight (bool): Whether to reweight signals after sampling.
+        """
+        super().__init__(mode=mode, n_coils=n_coils, n_averages=n_averages, reweight=reweight)
+
+        self.mode = mode
+
+        # Normalize to tuples: int → (value, value), tuple → tuple
+        self.n_coils = self._normalize_param(n_coils)
+        self.n_averages = self._normalize_param(n_averages)
+        self.reweight = reweight
+
+    def _normalize_param(self, param):
+        """
+        Normalize parameter to tuple format.
+
+        Args:
+            param: Either tuple (min, max) or int/float (exact value)
+
+        Returns:
+            Tuple (min, max)
+
+        Examples:
+            4 → (4, 4)  # Exactly 4
+            (1, 8) → (1, 8)  # Range 1-8
+            None → (1, None)  # Default
+        """
+        if param is None:
+            return (1, None)
+        elif isinstance(param, (int, float)):
+            # Single value → exact value
+            return (int(param), int(param))
+        elif isinstance(param, tuple):
+            return param
+        else:
+            raise TypeError(f"Parameter must be int, float, or tuple, got {type(param)}")
+
+
+    def process_nifti_list(self, data_list, water_list=None, coil_indices=None, average_indices=None, **kwargs):
+        """
+        Samples coils and averages from each NIFTI_MRS in the list.
+
+        Args:
+            data_list: List of metabolite MRS data (NIFTI_MRS objects).
+            water_list: List of water reference MRS data (NIFTI_MRS objects), optional.
+            coil_indices: List of coil indices to select (for deterministic mode).
+            average_indices: List of average indices to select (for deterministic mode).
+            **kwargs: Additional arguments.
+
+        Returns:
+            Tuple of (processed_data_list, processed_water_list)
+        """
+        processed_data = []
+        processed_water = []
+
+        for i, data_met in enumerate(data_list):
+            data_wat = water_list[i] if water_list is not None else None
+
+            # Sample coils and averages for this subject
+            data_met, data_wat = self._process_single(data_met, data_wat, coil_indices, average_indices)
+
+            processed_data.append(data_met)
+            if water_list is not None:
+                processed_water.append(data_wat if data_wat is not None else water_list[i])
+
+        return processed_data, (processed_water if water_list is not None else None)
+
+    def _process_single(self, data_met, data_wat=None, coil_indices=None, average_indices=None):
+        """
+        Samples coils and averages from a single subject's data.
+
+        Args:
+            data_met: Metabolite MRS data (NiftiMRS object).
+            data_wat: Water reference MRS data (NiftiMRS object), optional.
+            coil_indices: List of coil indices to select (for deterministic mode).
+            average_indices: List of average indices to select (for deterministic mode).
+        """
+        data_met, data_wat = self.sample_coils(data_met, data_wat, coil_indices)
+        data_met, data_wat = self.sample_averages(data_met, data_wat, average_indices)
+
+        if self.reweight:
+            data_met, data_wat = self._reweight_signals(data_met, data_wat)
+        return data_met, data_wat
+
+    def sample_coils(self, data_met, data_wat=None, coil_indices=None):
+        """
+        Samples specified coils from the data.
+
+        Args:
+            data_met: Metabolite MRS data (NiftiMRS object).
+            data_wat: Water reference MRS data (NiftiMRS object), optional
+            coil_indices: List of coil indices to select (if provided, use deterministic mode).
+        """
+        dim_tags = getattr(data_met, 'dim_tags', [])
+        has_coil = 'DIM_COIL' in dim_tags
+
+        # Determine mode at runtime: if coil_indices provided, use them (deterministic)
+        # Otherwise, sample randomly (if mode='random') or use all (if mode='deterministic' but no indices)
+        if coil_indices is not None:
+            # Deterministic mode: use provided indices
+            pass  # Will apply indices below
+        elif self.mode == 'random' and has_coil:
+            # Random mode: sample random coils
+            coil_dim = data_met.dim_position('DIM_COIL')
+            n_total  = data_met.shape[coil_dim]          # actual coil count
+            min_c, max_c = self._get_limits(self.n_coils, n_total)
+            if min_c < max_c:
+                num_coils   = torch.randint(min_c, max_c, (1,)).item()
+                coil_indices = torch.randperm(n_total)[:num_coils].tolist()
+        elif self.mode == 'deterministic' and has_coil:
+            # Deterministic mode but no indices provided: use all coils
+            coil_dim = data_met.dim_position('DIM_COIL')
+            coil_indices = list(range(data_met.shape[coil_dim]))
+
+        # apply sampling
+        if has_coil and coil_indices is not None:
+            if isinstance(coil_indices, list) or (isinstance(coil_indices, int) and coil_indices > 0):
+                _, data_met = split(data_met, 'DIM_COIL', coil_indices)
+                if data_wat is not None:
+                    _, data_wat = split(data_wat, 'DIM_COIL', coil_indices)
+
+            # update processing provenance
+            processing_info = f'{__name__}.sample_coils, '
+            processing_info += f'coil_indices={coil_indices}.'
+            update_processing_prov(data_met, 'Random Coil Sampling', processing_info)
+            if data_wat is not None:
+                update_processing_prov(data_wat, 'Random Coil Sampling', processing_info)
+
+        return data_met, data_wat
+
+    def sample_averages(self, data_met, data_wat=None, average_indices=None):
+        """
+        Samples specified averages from the data.
+
+        Args:
+            data_met: Metabolite MRS data (NiftiMRS object).
+            data_wat: Water reference MRS data (NiftiMRS object), optional.
+            average_indices: List of average indices to select.
+        """
+        dim_tags = getattr(data_met, 'dim_tags', [])
+        has_dyn = 'DIM_DYN' in dim_tags
+
+        # Determine mode at runtime: if average_indices provided, use them (deterministic)
+        # Otherwise, sample randomly (if mode='random') or use all (if mode='deterministic' but no indices)
+        if average_indices is not None:
+            # Deterministic mode: use provided indices
+            pass  # Will apply indices below
+        elif self.mode == 'random' and has_dyn:
+            # Random mode: sample random averages
+            dyn_dim = data_met.dim_position('DIM_DYN')
+            n_total = data_met.shape[dyn_dim]            # actual average count
+            min_a, max_a = self._get_limits(self.n_averages, n_total)
+            if min_a < max_a:
+                num_averages    = torch.randint(min_a, max_a, (1,)).item()
+                average_indices = torch.randperm(n_total)[:num_averages].tolist()
+        elif self.mode == 'deterministic' and has_dyn:
+            # Deterministic mode but no indices provided: use all averages
+            dyn_dim = data_met.dim_position('DIM_DYN')
+            average_indices = list(range(data_met.shape[dyn_dim]))
+
+        # apply sampling
+        if has_dyn and average_indices is not None:
+            if isinstance(average_indices, list) or (isinstance(average_indices, int) and average_indices > 0):
+                _, data_met = split(data_met, 'DIM_DYN', average_indices)
+
+            # update processing provenance
+            processing_info = f'{__name__}.sample_averages, '
+            processing_info += f'coil_indices={average_indices}.'
+            update_processing_prov(data_met, 'Random Average Sampling', processing_info)
+
+        return data_met, data_wat
+
+    def _reweight_signals(self, data, water):
+        """
+        Reweights coils and averages based on random weights.
+        TODO: implement reweighting.
+
+        Args:
+            data: MRS data (NiftiMRS object).
+            water: Water reference MRS data (NiftiMRS object), optional.
+        """
+        raise NotImplementedError("Reweighting not implemented yet.")
+
+    def _sample_snr_improvement(self, data, water):
+        """
+        Samples coils and averages to improve SNR.
+        TODO: implement SNR improvement sampling.
+
+        Args:
+            data: MRS data (NiftiMRS object).
+            water: Water reference MRS data (NiftiMRS object), optional.
+        """
+        raise NotImplementedError("SNR improvement sampling not implemented yet.")
+
+    def _get_limits(self, n, n_total):
+        """
+        Return ``(min_n, max_exclusive)`` for ``torch.randint(min_n, max_exclusive)``.
+
+        Args:
+            n (tuple): ``(min, max)`` — both *inclusive*, either can be ``None``.
+                       ``None`` on min → 1.  ``None`` on max → ``n_total`` (use all).
+            n_total (int): Total number of items available (coils / averages).
+
+        Examples::
+
+            _get_limits((1, None), 4)  →  (1, 5)   # randint → 1, 2, 3 or 4  ✓
+            _get_limits((2, 3),    4)  →  (2, 4)   # randint → 2 or 3         ✓
+            _get_limits((4, 4),    4)  →  (4, 5)   # randint → always 4       ✓
+            _get_limits((6, None), 4)  →  (4, 5)   # clamps to available      ✓
+        """
+        if n is None:
+            return 1, n_total + 1
+
+        min_val, max_val = n
+
+        min_n = 1 if min_val is None else max(1, min(int(min_val), n_total))
+        max_n = n_total if max_val is None else min(int(max_val), n_total)
+
+        # +1 because torch.randint upper bound is exclusive
+        return min_n, max_n + 1

augmentrum/sampling/kspace_reconstructor.py

@@ -0,0 +1,275 @@
+####################################################################################################
+#                                 kspace_reconstructor.py                                          #
+####################################################################################################
+#                                                                                                  #
+# Authors: J. T. LaMaster (john.t.lamaster@gmail.com)                                              #
+#                                                                                                  #
+# Created: 2025-10-22                                                                              #
+#                                                                                                  #
+# Purpose: Implements reconstruct_with_masking which accepts subsampled k-space data and its       #
+#          trajectory coordinates to perform density compesnation and regridding.                  #
+#                                                                                                  #
+####################################################################################################
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchkbnufft as tkbn
+from typing import Union, List, Tuple
+
+
+__all__ = ['reconstruct_with_masking']
+
+
+# --- Constants for normalization ---
+# TorchKbNufft expects coordinates in the range [-pi, pi]
+PI = torch.pi
+
+
+def prepare_data_and_coords(
+    coords: torch.Tensor, 
+    kdata: Optional[torch.Tensor] = None
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Reshapes coordinates and k-data to the format required by TorchKbNufft:
+    Coords: [B, D, N_total], Kdata: [B, C, N_total]
+    N_total = S * L (shots x length).
+    """
+    # Assuming input coords shape: [B, S, D, L] (e.g., [4, 10, 3, 100])
+    B, S, D, L = coords.shape
+    N_total = S * L
+
+    # 1. Normalize Coordinates (if not already done)
+    # The NUFFT forward/adjoint operators typically assume coordinates in [-pi, pi].
+    # We assume your input coords are normalized to [-1, 1] relative to k-space extent.
+    coords_normalized = coords * PI 
+
+    # 2. Flatten Trajectory: [B, S, D, L] -> [B, D, S, L] -> [B, D, N_total]
+    k_traj_flat = coords_normalized.permute(0, 2, 1, 3).reshape(B, D, N_total)
+
+    kdata_flat = None
+    if kdata is not None:
+        # Assuming input kdata shape: [B, S, C, L] (e.g., [4, 10, 1, 100])
+        B_k, S_k, C, L_k = kdata.shape
+        
+        # Flatten K-data: [B, S, C, L] -> [B, C, S, L] -> [B, C, N_total]
+        kdata_flat = kdata.permute(0, 2, 1, 3).reshape(B, C, N_total)
+
+    return k_traj_flat, kdata_flat
+
+def calculate_nufft_dcf_weights(
+    coords: torch.Tensor,
+    image_size: Union[List[int], Tuple[int]],
+    oversampling_factor: Union[List[float], Tuple[float]]
+) -> torch.Tensor:
+    """
+    Calculates Density Compensation Function (DCF) weights using the
+    iterative Pipe-style method from torchkbnufft.
+    """
+    device = coords.device
+    B, S, D, L = coords.shape
+    
+    # 1. Prepare Trajectory Coordinates
+    k_traj_flat, _ = prepare_data_and_coords(coords) # [B, D, N_total]
+
+    # 2. Determine Grid Size based on Image Size and Oversampling
+    # Ensure all sizes are determined on the CPU and converted to a list/tuple of native types
+    
+    # Get sizes as native types (no longer need to keep them as tensors)
+    image_size_list = list(image_size)
+    oversampling_factor_tensor = torch.tensor(oversampling_factor, dtype=torch.float32)
+    
+    # Grid size = floor(Image Size * Oversampling Factor)
+    grid_size_tensor = torch.floor(torch.tensor(image_size) * oversampling_factor_tensor).to(torch.long)
+    grid_size_list = grid_size_tensor.tolist() # <--- Pass native list/tuple
+
+    # 3. Calculate DCF (Iterative Pipe-style)
+    # Pass 'im_size' and 'grid_size' as separate arguments, not the KbNufft object.
+    with torch.no_grad():
+        dcf_flat = tkbn.calc_density_compensation_function(
+            k_traj_flat, 
+            im_size=image_size_list,     # <--- FIX 1: Pass im_size
+            grid_size=grid_size_list     # <--- FIX 2: Pass grid_size
+        ).detach() # [B, 1, N_total]
+
+    # 4. Reshape DCF for easy multiplication with k-data: [B, 1, S*L] -> [B, S, L]
+    dcf_weights = dcf_flat.view(B, S, L)
+
+    return dcf_weights
+
+
+def regrid_kdata_to_image(
+    weighted_kdata: torch.Tensor,
+    coords: torch.Tensor,
+    image_size: Union[List[int], Tuple[int]],
+    oversampling_factor: Union[List[float], Tuple[float]]
+) -> torch.Tensor:
+    """
+    Performs Adjoint NUFFT (regridding) to map weighted k-space data to the image domain.
+
+    Args:
+        weighted_kdata (torch.Tensor): Density-compensated k-space data, shape [B, S, C, L].
+        coords (torch.Tensor): K-space trajectory coordinates, shape [B, S, D, L].
+        image_size (Union[List[int], Tuple[int]]): Target image dimensions (Nx, Ny, Nz).
+        oversampling_factor (Union[List[float], Tuple[float]]): Oversampling factors 
+                                                                for the regridding grid (rho_x, rho_y, rho_z).
+
+    Returns:
+        torch.Tensor: Reconstructed complex image volume, shape [B, C, Nx, Ny, Nz].
+    """
+    device = weighted_kdata.device
+    
+    # 1. Prepare Coords and K-data for NUFFT
+    # k_traj_flat: [B, D, N_total], kdata_flat: [B, C, N_total]
+    k_traj_flat, kdata_flat = prepare_data_and_coords(coords, weighted_kdata)
+    
+    # 2. Determine Grid Size
+    image_size = torch.tensor(image_size, dtype=torch.long, device=device)
+    oversampling_factor = torch.tensor(oversampling_factor, dtype=torch.float32, device=device)
+    
+    # Grid size = floor(Image Size * Oversampling Factor)
+    grid_size = torch.floor(image_size * oversampling_factor).to(torch.long)
+
+    # 3. Instantiate Adjoint NUFFT Operator (Regridding)
+    adj_op = tkbn.KbNufftAdjoint(
+        im_size=image_size.tolist(),  # [Nx, Ny, Nz]
+        grid_size=grid_size.tolist()  # [Grid_Nx, Grid_Ny, Grid_Nz]
+    ).to(device)
+
+    # 4. Perform Regridding
+    # The operation is: Adjoint(k_data) -> IFFT -> Image
+    image_domain_output = adj_op(kdata_flat, k_traj_flat)
+
+    # Output shape: [B, C, Nx, Ny, Nz]
+    return image_domain_output
+
+
+def reconstruct_with_masking(
+    coords: torch.Tensor,
+    kdata: torch.Tensor,
+    undersampling_mask: torch.Tensor,
+    image_size: Union[List[int], Tuple[int]],
+    oversampling_factor: Union[List[float], Tuple[float]]
+) -> torch.Tensor:
+    """
+    High-level wrapper to apply undersampling mask, calculate DCF, and perform regridding.
+
+    Args:
+        coords (torch.Tensor): K-space trajectory coordinates, shape [B, S, D, L].
+        kdata (torch.Tensor): Raw k-space data, shape [B, S, C, L].
+        undersampling_mask (torch.Tensor): Boolean mask indicating which shots to KEEP. 
+                                           Shape: [B, S].
+        image_size (Union[List[int], Tuple[int]]): Target image dimensions (Nx, Ny, Nz).
+        oversampling_factor (Union[List[float], Tuple[float]]): Oversampling factors 
+                                                                for the regridding grid (rho_x, rho_y, rho_z).
+
+    Returns:
+        torch.Tensor: Reconstructed complex image volume, shape [B, C, Nx, Ny, Nz].
+    """    
+    # --- 1. Apply Undersampling Mask (Zero-Filling) ---
+    # The mask is [B, S]. We need to expand it to [B, S, C, L] to multiply with kdata.
+    B, S, C, L = kdata.shape
+    
+    # Expand mask to match kdata shape: [B, S] -> [B, S, 1, 1] -> [B, S, C, L]
+    # Invert the boolean mask to create a float mask (True -> 0, False -> 1)
+    # Since the mask is for samples to KEEP (conventionally True), 
+    # we zero the kdata samples where the mask is False (discard).
+    
+    # Convert mask to float tensor (0 or 1) and expand
+    mask_float = undersampling_mask.to(dtype=kdata.real.dtype).unsqueeze(2).unsqueeze(3) 
+    
+    # Expand across Coil (C) and Length (L) dimensions
+    mask_bcast = mask_float.expand(B, S, C, L)
+    
+    # Apply mask (zero out discarded k-space samples/shots)
+    masked_kdata = kdata * mask_bcast
+
+    # --- 2. Calculate DCF Weights ---
+    # DCF calculation should use ALL coordinates to reflect the desired density.
+    dcf_weights = calculate_nufft_dcf_weights(
+        coords=coords, 
+        image_size=image_size, 
+        oversampling_factor=oversampling_factor
+    )
+    
+    # Reshape DCF for multiplication: [B, S, L] -> [B, S, 1, L]
+    dcf_bcast = dcf_weights.unsqueeze(2) 
+    
+    # --- 3. Apply DCF to Masked K-space Data ---
+    weighted_kdata = masked_kdata * dcf_bcast 
+    
+    # --- 4. Perform Regridding (Adjoint NUFFT) ---
+    reconstructed_image = regrid_kdata_to_image(
+        weighted_kdata=weighted_kdata,
+        coords=coords,
+        image_size=image_size,
+        oversampling_factor=oversampling_factor
+    )
+    
+    return reconstructed_image
+
+
+if __name__ == "__ main __":
+    # Define parameters for a 3D example
+    B, S, D, L, C = 4, 10, 3, 100, 1 # Batch, Shots, Dimensions, Length, Channels
+    Target_Image_Size = (64, 64, 64)
+    Oversampling_Factor = (1.5, 1.5, 1.5)
+    
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    # Dummy Data: Coords must be float, Kdata must be complex
+    coords_in = (torch.rand(B, S, D, L) * 2 - 1).to(dtype=torch.float32, device=device) # [-1, 1] range
+    kdata_in = torch.randn(B, S, C, L, dtype=torch.complex64, device=device) 
+    
+    # --- NEW: Create an Undersampling Mask ---
+    # Mask [B, S] (e.g., keeping only 50% of the shots randomly per batch sample)
+    undersampling_mask_in = torch.rand(B, S) < 0.5
+    undersampling_mask_in = undersampling_mask_in.to(device)
+
+    # --- 1. Full Reconstruction Pipeline with Masking ---
+    print("--- Starting Full Reconstruction Pipeline (Masked) ---")
+    reconstructed_image_masked = reconstruct_with_masking(
+        coords=coords_in,
+        kdata=kdata_in,
+        undersampling_mask=undersampling_mask_in,
+        image_size=Target_Image_Size,
+        oversampling_factor=Oversampling_Factor
+    )
+    
+    print(f"Masked Reconstruction successful. Final Image Shape: {reconstructed_image_masked.shape}")
+    print(f"Shots kept per batch: {undersampling_mask_in.sum(dim=1).tolist()}")
+    print(f"Expected Final Image Shape: [{B}, {C}, {Target_Image_Size[0]}, {Target_Image_Size[1]}, {Target_Image_Size[2]}]")
+
+    # The reconstructed image is now a complex tensor [B, C, Nx, Ny, Nz]
+
+    # Define parameters for a 2D example
+    B, S, D, L, C = 4, 10, 2, 100, 1 # Batch, Shots, Dimensions, Length, Channels
+    Target_Image_Size = (64, 64)
+    Oversampling_Factor = (1.5, 1.5)
+    
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    # Dummy Data: Coords must be float, Kdata must be complex
+    coords_in = (torch.rand(B, S, D, L) * 2 - 1).to(dtype=torch.float32, device=device) # [-1, 1] range
+    kdata_in = torch.randn(B, S, C, L, dtype=torch.complex64, device=device) 
+    
+    # --- NEW: Create an Undersampling Mask ---
+    # Mask [B, S] (e.g., keeping only 50% of the shots randomly per batch sample)
+    undersampling_mask_in = torch.rand(B, S) < 0.5
+    undersampling_mask_in = undersampling_mask_in.to(device)
+    
+    # --- 1. Full Reconstruction Pipeline with Masking ---
+    print("--- Starting Full Reconstruction Pipeline (Masked) ---")
+    reconstructed_image_masked = reconstruct_with_masking(
+        coords=coords_in,
+        kdata=kdata_in,
+        undersampling_mask=undersampling_mask_in,
+        image_size=Target_Image_Size,
+        oversampling_factor=Oversampling_Factor
+    )
+    
+    print(f"Masked Reconstruction successful. Final Image Shape: {reconstructed_image_masked.shape}")
+    print(f"Shots kept per batch: {undersampling_mask_in.sum(dim=1).tolist()}")
+    print("Expected Final Image Shape: [{}, {}, {}]".format(B, C, [x for x in Target_Image_Size]))
+    
+    # The reconstructed image is now a complex tensor [B, C, Nx, Ny]