Fast Fourier Transform

def fft(a, n=None, axis=-1, norm=None):
    """Compute 1D discrete Fourier Transform.
    
    Args:
        a: Input array
        n: Length of transformed axis
        axis: Axis over which to compute FFT
        norm: Normalization mode ('backward', 'ortho', 'forward')
        
    Returns:
        cupy.ndarray: Complex FFT result
    """

def ifft(a, n=None, axis=-1, norm=None):
    """Compute 1D inverse discrete Fourier Transform."""

def rfft(a, n=None, axis=-1, norm=None):
    """Compute 1D FFT of real input.
    
    Returns:
        cupy.ndarray: Complex FFT with conjugate symmetry
    """

def irfft(a, n=None, axis=-1, norm=None):
    """Compute inverse of rfft.
    
    Returns:
        cupy.ndarray: Real inverse FFT result
    """

def hfft(a, n=None, axis=-1, norm=None):
    """Compute FFT of signal with Hermitian symmetry."""

def ihfft(a, n=None, axis=-1, norm=None):
    """Compute inverse of hfft."""

def fft2(a, s=None, axes=(-2, -1), norm=None):
    """Compute 2D discrete Fourier Transform.
    
    Args:
        a: Input array
        s: Shape of transformed axes
        axes: Axes over which to compute FFT
        norm: Normalization mode
        
    Returns:
        cupy.ndarray: 2D FFT result
    """

def ifft2(a, s=None, axes=(-2, -1), norm=None):
    """Compute 2D inverse discrete Fourier Transform."""

def rfft2(a, s=None, axes=(-2, -1), norm=None):
    """Compute 2D FFT of real input."""

def irfft2(a, s=None, axes=(-2, -1), norm=None):
    """Compute 2D inverse of rfft2."""

def fftn(a, s=None, axes=None, norm=None):
    """Compute N-dimensional discrete Fourier Transform.
    
    Args:
        a: Input array
        s: Shape of transformed axes
        axes: Axes over which to compute FFT
        norm: Normalization mode
        
    Returns:
        cupy.ndarray: N-D FFT result
    """

def ifftn(a, s=None, axes=None, norm=None):
    """Compute N-dimensional inverse discrete Fourier Transform."""

def rfftn(a, s=None, axes=None, norm=None):
    """Compute N-dimensional FFT of real input."""

def irfftn(a, s=None, axes=None, norm=None):
    """Compute N-dimensional inverse of rfftn."""

def fftfreq(n, d=1.0):
    """Return discrete Fourier Transform sample frequencies.
    
    Args:
        n: Window length
        d: Sample spacing
        
    Returns:
        cupy.ndarray: Array of frequencies
    """

def rfftfreq(n, d=1.0):
    """Return sample frequencies for rfft."""

def fftshift(x, axes=None):
    """Shift zero-frequency component to center.
    
    Args:
        x: Input array
        axes: Axes to shift
        
    Returns:
        cupy.ndarray: Shifted array
    """

def ifftshift(x, axes=None):
    """Inverse of fftshift."""

import cupy as cp
import matplotlib.pyplot as plt

# Create a test signal
t = cp.linspace(0, 1, 1000, endpoint=False)
signal = cp.sin(2 * cp.pi * 50 * t) + cp.sin(2 * cp.pi * 120 * t)

# Add some noise
signal += 0.2 * cp.random.randn(len(t))

# Compute FFT
fft_result = cp.fft.fft(signal)
freqs = cp.fft.fftfreq(len(signal), d=t[1]-t[0])

# Get magnitude spectrum
magnitude = cp.abs(fft_result)

print(f"Signal length: {len(signal)}")
print(f"FFT result shape: {fft_result.shape}")
print(f"Peak frequencies found at: {freqs[cp.argsort(magnitude)[-3:-1]]}")

# For real-valued signals, use rfft for efficiency
real_signal = cp.cos(2 * cp.pi * 10 * t) + cp.cos(2 * cp.pi * 25 * t)

# Real FFT (more efficient for real inputs)
rfft_result = cp.fft.rfft(real_signal)
rfreqs = cp.fft.rfftfreq(len(real_signal), d=t[1]-t[0])

# Reconstruct original signal
reconstructed = cp.fft.irfft(rfft_result)

print(f"Original signal shape: {real_signal.shape}")
print(f"RFFT result shape: {rfft_result.shape}")
print(f"Reconstruction error: {cp.mean(cp.abs(real_signal - reconstructed)):.2e}")

# Create a 2D test pattern
x = cp.linspace(-5, 5, 256)
y = cp.linspace(-5, 5, 256)
X, Y = cp.meshgrid(x, y)

# Create a 2D sinusoidal pattern
image = cp.sin(2 * cp.pi * 0.5 * X) * cp.cos(2 * cp.pi * 0.3 * Y)

# Compute 2D FFT
fft2_result = cp.fft.fft2(image)

# Shift zero frequency to center for visualization
fft2_shifted = cp.fft.fftshift(fft2_result)

# Get magnitude spectrum
magnitude_2d = cp.abs(fft2_shifted)

print(f"Image shape: {image.shape}")
print(f"2D FFT shape: {fft2_result.shape}")
print(f"Max magnitude: {cp.max(magnitude_2d):.2f}")

# Create noisy signal
t = cp.linspace(0, 2, 2000, endpoint=False)
clean_signal = cp.sin(2 * cp.pi * 5 * t)  # 5 Hz signal
noisy_signal = clean_signal + 0.5 * cp.sin(2 * cp.pi * 50 * t)  # Add 50 Hz noise

# FFT of noisy signal
fft_noisy = cp.fft.fft(noisy_signal)
freqs = cp.fft.fftfreq(len(noisy_signal), d=t[1]-t[0])

# Create low-pass filter (remove frequencies above 10 Hz)
filter_mask = cp.abs(freqs) < 10
fft_filtered = fft_noisy * filter_mask

# Inverse FFT to get filtered signal
filtered_signal = cp.fft.ifft(fft_filtered).real

# Compare results
original_power = cp.mean(clean_signal**2)
noisy_power = cp.mean(noisy_signal**2)
filtered_power = cp.mean(filtered_signal**2)

print(f"Original signal power: {original_power:.3f}")
print(f"Noisy signal power: {noisy_power:.3f}")
print(f"Filtered signal power: {filtered_power:.3f}")
print(f"Filter effectiveness: {cp.mean(cp.abs(clean_signal - filtered_signal)):.4f}")

# FFT-based convolution is efficient for large arrays
signal_length = 10000
kernel_length = 100

# Create test signal and kernel
signal = cp.random.randn(signal_length)
kernel = cp.exp(-cp.linspace(-2, 2, kernel_length)**2)  # Gaussian kernel

# Pad for proper convolution
total_length = signal_length + kernel_length - 1
signal_padded = cp.zeros(total_length)
signal_padded[:signal_length] = signal

kernel_padded = cp.zeros(total_length)
kernel_padded[:kernel_length] = kernel

# FFT-based convolution
signal_fft = cp.fft.fft(signal_padded)
kernel_fft = cp.fft.fft(kernel_padded)
convolved_fft = signal_fft * kernel_fft
convolved = cp.fft.ifft(convolved_fft).real

# Compare with direct convolution (for verification on smaller arrays)
if signal_length < 1000:
    direct_conv = cp.convolve(signal, kernel, mode='full')
    fft_conv_trimmed = convolved[:len(direct_conv)]
    error = cp.mean(cp.abs(direct_conv - fft_conv_trimmed))
    print(f"FFT convolution error: {error:.2e}")

print(f"Original signal length: {signal_length}")
print(f"Convolved signal length: {len(convolved)}")

# Power spectral density estimation
fs = 1000  # Sampling frequency
t = cp.arange(0, 2, 1/fs)  # 2 seconds of data

# Create signal with multiple frequency components
signal = (cp.sin(2*cp.pi*50*t) + 
         0.5*cp.sin(2*cp.pi*120*t) + 
         0.2*cp.sin(2*cp.pi*200*t) +
         0.1*cp.random.randn(len(t)))

# Compute power spectral density
fft_signal = cp.fft.fft(signal)
freqs = cp.fft.fftfreq(len(signal), 1/fs)
power_spectrum = cp.abs(fft_signal)**2

# Find peaks in spectrum
positive_freqs = freqs[:len(freqs)//2]
positive_power = power_spectrum[:len(power_spectrum)//2]

# Get indices of largest peaks
peak_indices = cp.argsort(positive_power)[-5:]
peak_freqs = positive_freqs[peak_indices]
peak_powers = positive_power[peak_indices]

print("Top 5 frequency components:")
for freq, power in zip(peak_freqs, peak_powers):
    print(f"  {freq:.1f} Hz: {power:.2e}")

# 3D FFT for volumetric data
volume_shape = (64, 64, 64)
volume_data = cp.random.randn(*volume_shape)

# Add some structured pattern
x, y, z = cp.meshgrid(cp.arange(64), cp.arange(64), cp.arange(64), indexing='ij')
pattern = cp.sin(2*cp.pi*x/16) * cp.cos(2*cp.pi*y/16) * cp.sin(2*cp.pi*z/16)
volume_data += pattern

# 3D FFT
fft3d = cp.fft.fftn(volume_data)

# Get magnitude
magnitude_3d = cp.abs(fft3d)

print(f"3D volume shape: {volume_data.shape}")
print(f"3D FFT shape: {fft3d.shape}")
print(f"Peak magnitude: {cp.max(magnitude_3d):.2f}")
print(f"Memory usage: {fft3d.nbytes / 1024**2:.1f} MB")

# FFT performance depends on array size
# Powers of 2 are typically fastest
sizes = [1000, 1024, 2000, 2048, 4000, 4096]

for size in sizes:
    data = cp.random.randn(size)
    
    # Time the FFT
    start = cp.cuda.Event()
    end = cp.cuda.Event()
    
    start.record()
    result = cp.fft.fft(data)
    end.record()
    end.synchronize()
    
    elapsed = start.elapsed_time(end)
    print(f"Size {size:4d}: {elapsed:.2f} ms")

# For repeated FFTs of the same size, CuPy caches plans
data = cp.random.randn(1024)
for i in range(5):
    start = cp.cuda.Event()
    end = cp.cuda.Event()
    
    start.record()
    result = cp.fft.fft(data)
    end.record()
    end.synchronize()
    
    elapsed = start.elapsed_time(end)
    print(f"Iteration {i+1}: {elapsed:.2f} ms")