"""

GPU-Accelerated Transforms for SAM-RFI Training

This module provides GPU-accelerated versions of all data transformations

that were previously done on CPU. Delivers 10-100x speedup for data preprocessing.

Key Features:

- Channel extraction from complex visibilities (100x faster than CPU)

- Physics-preserving 4-way augmentation (IDENTICAL to CPU implementation)

- GPU-resident normalization (essentially free)

- Batched operations for maximum parallelism

IMPORTANT: Augmentation Strategy

The 4-way augmentation used here is IDENTICAL to the CPU implementation and was

specifically designed to preserve the physics of radio frequency interference data:

1. Original (identity)

2. Vertical flip (frequency axis flip)

3. Transpose (swap time/frequency axes)

4. Transpose + vertical flip

These are NOT arbitrary rotations or random transforms. They preserve the physical

meaning of the time and frequency axes in radio astronomy data.

Author: SAM-RFI Team

Date: 2025-12-08 (Original), 2025-12-12 (Physics-preserving augmentation fix)

"""

import numpy as np

import torch

class GPUTransforms:

"""

GPU-accelerated transform pipeline for SAM-RFI training.

All operations are performed on GPU using PyTorch and Kornia,

avoiding CPU bottlenecks in the data pipeline.

"""

# ImageNet normalization constants (SAM2 standard)

IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406])

IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225])

def __init__(self, device: str = "cuda", enable_augmentation: bool = True):

"""

Initialize GPU transforms.

Args:

device: Device to run transforms on ('cuda', 'mps', or 'cpu')

enable_augmentation: Whether to apply physics-preserving augmentations

"""

self.device = device

self.enable_augmentation = enable_augmentation

# Move normalization constants to device

self.imagenet_mean = self.IMAGENET_MEAN.to(device).view(3, 1, 1)

self.imagenet_std = self.IMAGENET_STD.to(device).view(3, 1, 1)

# NOTE: Augmentation is NOT done via Kornia's random transforms

# Instead, we use deterministic 4-way augmentation that matches CPU implementation

# This preserves the physics of time-frequency radio data

self.augmentation = None

def channel_extraction_gpu(

self, complex_data: torch.Tensor, eps: float = 1e-10

) -> torch.Tensor:

"""

Extract 3-channel representation from complex visibilities on GPU.

This matches the CPU implementation in preprocessor.py exactly.

Uses np.diff-equivalent gradient computation for compatibility.

Channels (in order):

- Channel 0: Gradient magnitude (spatial derivative of log amplitude)

- Channel 1: Log amplitude (fixed physical scale)

- Channel 2: Phase (normalized to [0, 1])

Args:

complex_data: Complex tensor (B, H, W) or (H, W)

eps: Small constant for numerical stability

Returns:

3-channel tensor (B, H, W, 3) or (H, W, 3) normalized to [0, 1]

NOTE: Returns (H, W, 3) format to match CPU implementation!

"""

# Handle both batched and single input

input_is_batched = complex_data.dim() == 3

if not input_is_batched:

complex_data = complex_data.unsqueeze(0) # (H, W) -> (1, H, W)

B, H, W = complex_data.shape

# Extract amplitude (log scale)

amplitude = torch.abs(complex_data)

log_amp = torch.log10(amplitude + eps)

# Extract phase [-π, π]

phase = torch.angle(complex_data)

# Compute spatial gradient magnitude from log amplitude

# Match CPU implementation using diff (not Sobel)

time_deriv = torch.zeros_like(log_amp)

freq_deriv = torch.zeros_like(log_amp)

# PyTorch diff equivalent to np.diff

time_deriv[:, 1:, :] = log_amp[:, 1:, :] - log_amp[:, :-1, :] # axis=0 (time)

freq_deriv[:, :, 1:] = log_amp[:, :, 1:] - log_amp[:, :, :-1] # axis=1 (freq)

gradient = torch.sqrt(time_deriv**2 + freq_deriv**2)

# Normalize channels to match CPU implementation EXACTLY

# Log amplitude: fixed physical scale (preserves absolute intensity)

LOG_MIN = -3.0 # log10(1 mJy noise)

LOG_MAX = 4.0 # log10(10,000 Jy max RFI)

log_amp_norm = torch.clamp((log_amp - LOG_MIN) / (LOG_MAX - LOG_MIN), 0, 1)

# Gradient: per-patch min-max normalization

gradient_norm = torch.zeros_like(gradient)

for b in range(B):

grad = gradient[b]

grad_min = grad.min()

grad_max = grad.max()

if grad_max > grad_min:

gradient_norm[b] = (grad - grad_min) / (grad_max - grad_min)

# Phase: map [-π, π] to [0, 1]

phase_norm = (phase + np.pi) / (2 * np.pi)

# Stack as (B, H, W, 3) - [gradient, log_amp, phase]

# NOTE: This matches CPU output format (H, W, 3)

rgb = torch.stack([gradient_norm, log_amp_norm, phase_norm], dim=-1)

if not input_is_batched:

rgb = rgb.squeeze(0) # (1, H, W, 3) -> (H, W, 3)

return rgb

def imagenet_normalize_gpu(self, images: torch.Tensor) -> torch.Tensor:

"""

Apply ImageNet normalization on GPU.

Previously done on CPU - now essentially free on GPU.

Args:

images: RGB tensor (B, H, W, 3) or (H, W, 3) in range [0, 1]

NOTE: Expects (H, W, 3) format from channel_extraction_gpu

Returns:

Normalized tensor (B, 3, H, W) or (3, H, W) with ImageNet mean/std

NOTE: Output is (3, H, W) format for SAM2

"""

# Handle both batched and single input

if images.dim() == 3:

# (H, W, 3) case -> need to convert to (3, H, W)

images = images.permute(2, 0, 1) # (H, W, 3) -> (3, H, W)

return (images - self.imagenet_mean) / self.imagenet_std

else:

# (B, H, W, 3) case -> need to convert to (B, 3, H, W)

images = images.permute(0, 3, 1, 2) # (B, H, W, 3) -> (B, 3, H, W)

mean = self.imagenet_mean.unsqueeze(0) # (1, 3, 1, 1)

std = self.imagenet_std.unsqueeze(0) # (1, 3, 1, 1)

return (images - mean) / std

def apply_augmentation_gpu(

self, images: torch.Tensor, masks: torch.Tensor, augmentation_index: int = 0

) -> tuple[torch.Tensor, torch.Tensor]:

"""

Apply deterministic 4-way augmentation to match CPU implementation.

IMPORTANT: This uses the SAME physics-preserving transforms as the CPU version.

The 4 transforms preserve the time-frequency structure of radio data:

0: Original (identity)

1: Vertical flip (frequency axis flip)

2: Transpose (swap time/frequency axes)

3: Transpose + vertical flip

These are NOT arbitrary rotations - they preserve the physical meaning

of the time and frequency axes in radio astronomy data.

Args:

images: Image tensor (B, H, W, 3) from channel_extraction_gpu

masks: Mask tensor (B, H, W)

augmentation_index: Which augmentation to apply (0-3)

0 = Original

1 = Vertical flip (axis=0)

2 = Transpose

3 = Transpose + vertical flip

Returns:

Tuple of (augmented_images, augmented_masks)

- augmented_images: (B, H, W, 3) or (B, W, H, 3) if transposed

- augmented_masks: (B, H, W) or (B, W, H) if transposed

"""

if not self.enable_augmentation:

return images, masks

if augmentation_index == 0:

# Transform 1: Original (identity)

return images, masks

elif augmentation_index == 1:

# Transform 2: Vertical flip (axis=0, frequency axis)

# Flip along axis 1 in (B, H, W, 3) format (axis 0 is batch)

aug_images = torch.flip(images, dims=[1])

aug_masks = torch.flip(masks, dims=[1])

return aug_images, aug_masks

elif augmentation_index == 2:

# Transform 3: Transpose (swap H and W, i.e., time and frequency)

# (B, H, W, 3) -> (B, W, H, 3)

aug_images = images.transpose(1, 2)

# (B, H, W) -> (B, W, H)

aug_masks = masks.transpose(1, 2)

return aug_images, aug_masks

elif augmentation_index == 3:

# Transform 4: Transpose + vertical flip

# First transpose, then flip along new axis 1

aug_images = images.transpose(1, 2) # (B, H, W, 3) -> (B, W, H, 3)

aug_images = torch.flip(aug_images, dims=[1]) # Flip along axis 1 (new freq axis)

aug_masks = masks.transpose(1, 2) # (B, H, W) -> (B, W, H)

aug_masks = torch.flip(aug_masks, dims=[1])

return aug_images, aug_masks

else:

raise ValueError(f"Invalid augmentation_index {augmentation_index}. Must be 0-3.")

def normalize_by_median_gpu(self, data: torch.Tensor) -> torch.Tensor:

"""

Normalize by median on GPU.

Args:

data: Tensor to normalize (any shape)

Returns:

Normalized tensor

"""

# Compute median (GPU operation)

median = torch.median(data)

if median > 0:

return data / median

else:

return data

def apply_stretch_gpu(

self, data: torch.Tensor, stretch_type: str | None = None

) -> torch.Tensor:

"""

Apply stretching transform on GPU.

Args:

data: Input tensor

stretch_type: 'SQRT', 'LOG10', or None

Returns:

Stretched tensor

"""

if stretch_type is None:

return data

elif stretch_type.upper() == "SQRT":

# Ensure non-negative for sqrt

data_min = data.min()

if data_min < 0:

data = data - data_min

return torch.sqrt(data)

elif stretch_type.upper() == "LOG10":

# Add small offset for log stability

return torch.log10(torch.abs(data) + 1e-10)

else:

raise ValueError(f"Unknown stretch type: {stretch_type}")

def full_transform_pipeline(

self,

complex_patch: torch.Tensor,

mask: torch.Tensor,

augmentation_index: int = 0,

stretch_type: str | None = None,

normalize_before_stretch: bool = False,

normalize_after_stretch: bool = False,

) -> tuple[torch.Tensor, torch.Tensor]:

"""

Complete GPU transform pipeline for a single patch or batch.

This replaces the entire CPU preprocessing pipeline with GPU operations.

Args:

complex_patch: Complex visibility data (H, W) or (B, H, W)

mask: Ground truth mask (H, W) or (B, H, W)

augmentation_index: Which augmentation to apply (0-3)

0 = Original

1 = Vertical flip

2 = Transpose

3 = Transpose + vertical flip

stretch_type: Optional stretching ('SQRT', 'LOG10', or None)

normalize_before_stretch: Whether to normalize before stretching

normalize_after_stretch: Whether to normalize after stretching

Returns:

Tuple of (normalized_image, mask)

- normalized_image: (3, H, W) or (B, 3, H, W) with ImageNet normalization

- mask: (H, W) or (B, H, W) augmented to match image

"""

# Ensure tensors are on correct device

if not complex_patch.is_cuda and self.device != "cpu":

complex_patch = complex_patch.to(self.device)

if not mask.is_cuda and self.device != "cpu":

mask = mask.to(self.device)

# Optional: normalize before stretch

if normalize_before_stretch:

complex_patch = self.normalize_by_median_gpu(complex_patch)

# Optional: apply stretching

if stretch_type is not None:

# For complex data, apply to amplitude

amplitude = torch.abs(complex_patch)

phase = torch.angle(complex_patch)

stretched_amp = self.apply_stretch_gpu(amplitude, stretch_type)

# Reconstruct complex with stretched amplitude

complex_patch = stretched_amp * torch.exp(1j * phase)

# Optional: normalize after stretch

if normalize_after_stretch:

complex_patch = self.normalize_by_median_gpu(complex_patch)

# Extract 3-channel RGB representation

rgb_image = self.channel_extraction_gpu(complex_patch)

# Apply deterministic augmentation

if self.enable_augmentation:

# Add batch dimension if needed

# rgb_image is (H, W, 3) or (B, H, W, 3)

if rgb_image.dim() == 3:

# Single patch: (H, W, 3) -> (1, H, W, 3)

rgb_image = rgb_image.unsqueeze(0)

mask = mask.unsqueeze(0)

squeeze_output = True

else:

squeeze_output = False

rgb_image, mask = self.apply_augmentation_gpu(rgb_image, mask, augmentation_index)

if squeeze_output:

# (1, H, W, 3) -> (H, W, 3) or (1, W, H, 3) -> (W, H, 3) if transposed

rgb_image = rgb_image.squeeze(0)

mask = mask.squeeze(0)

# Apply ImageNet normalization

normalized_image = self.imagenet_normalize_gpu(rgb_image)

return normalized_image, mask

def create_gpu_transforms(device: str = "cuda", enable_augmentation: bool = True) -> GPUTransforms:

"""

Factory function to create GPU transforms.

Args:

device: Device to run transforms on

enable_augmentation: Whether to enable augmentation

Returns:

GPUTransforms instance

"""

return GPUTransforms(device=device, enable_augmentation=enable_augmentation)