tensor.py

import torch

def merge_tensors_fft2(v0: torch.Tensor, v1: torch.Tensor, t: float) -> torch.Tensor:
 """
 Merges two tensors using 2D Fourier transform interpolation.
 
 Parameters:
 - v0 (torch.Tensor): The first input tensor.
 - v1 (torch.Tensor): The second input tensor.
 - t (float): Interpolation parameter (0 <= t <= 1).
 
 Returns:
 - torch.Tensor: The tensor resulting from the interpolated inverse FFT.
 """
 # Ensure the input tensors are on the same device and dtype
 v0 = v0.to("cuda:0")
 v1 = v1.to("cuda:0")

 if len(v0.shape) == 1:
 fft_v0 = torch.fft.fft(v0)
 fft_v1 = torch.fft.fft(v1)
 result_fft = torch.zeros_like(fft_v0)
 
 real_v0 = fft_v0.real
 real_v1 = fft_v1.real
 abs_real_v0 = real_v0.abs() 
 abs_real_v1 = real_v1.abs() 

 sign_mask = real_v0.sign() == real_v1.sign()
 larger_values_mask = abs_real_v0 > abs_real_v1

 result_fft.real[sign_mask] = torch.max(real_v0[sign_mask], real_v1[sign_mask])
 result_fft.real[~sign_mask] = torch.where(larger_values_mask[~sign_mask], real_v0[~sign_mask], real_v1[~sign_mask])

 imag_v0 = fft_v0.imag
 imag_v1 = fft_v1.imag
 abs_imag_v0 = imag_v0.abs()
 abs_imag_v1 = imag_v1.abs()
 larger_values_mask_imag = abs_imag_v0 > abs_imag_v1

 result_fft.imag[sign_mask] = (1 - t) * imag_v0[sign_mask] + t * imag_v1[sign_mask]
 result_fft.imag[~sign_mask] = torch.where(larger_values_mask_imag[~sign_mask], imag_v0[~sign_mask], imag_v1[~sign_mask])

 merged_tensor = torch.fft.ifft(result_fft).real # Taking the real part
 del v0, v1, fft_v0, fft_v1, result_fft
 return merged_tensor

 # Perform the 2D FFT on both tensors
 fft_v0 = torch.fft.fftn(v0, dim=(-2, -1))
 fft_v1 = torch.fft.fftn(v1, dim=(-2, -1))
 
 # Initialize the result FFT tensor
 result_fft = torch.zeros_like(fft_v0)
 
 # Compare real parts of the coefficients
 real_v0 = fft_v0.real
 real_v1 = fft_v1.real
 abs_real_v0 = real_v0.abs()
 abs_real_v1 = real_v1.abs()
 
 # Create masks for where signs match and where they do not
 sign_mask = real_v0.sign() == real_v1.sign()
 larger_values_mask = abs_real_v0 > abs_real_v1
 
 # Where signs match, interpolate; where signs do not match, take the larger by magnitude
 result_fft.real[sign_mask] = torch.max(real_v0[sign_mask], real_v1[sign_mask])
 result_fft.real[~sign_mask] = torch.where(larger_values_mask[~sign_mask], real_v0[~sign_mask], real_v1[~sign_mask])
 
 # Assuming the imaginary part should be treated similarly, adjust this if not
 imag_v0 = fft_v0.imag
 imag_v1 = fft_v1.imag
 abs_imag_v0 = imag_v0.abs()
 abs_imag_v1 = imag_v1.abs()
 larger_values_mask_imag = abs_imag_v0 > abs_imag_v1
 
 result_fft.imag[sign_mask] = (1 - t) * imag_v0[sign_mask] + t * imag_v1[sign_mask]
 result_fft.imag[~sign_mask] = torch.where(larger_values_mask_imag[~sign_mask], imag_v0[~sign_mask], imag_v1[~sign_mask])
 
 # Perform the inverse FFT to go back to the spatial domain
 merged_tensor = torch.fft.ifftn(result_fft, dim=(-2, -1)).real # Taking the real part
 
 return merged_tensor

def merge_tensors_fft_shell(v0: torch.Tensor, v1: torch.Tensor) -> torch.Tensor:
 """
 Merges two tensors using 2D Fourier transform interpolation.
 
 Parameters:
 - v0 (torch.Tensor): The first input tensor.
 - v1 (torch.Tensor): The second input tensor.
 
 Returns:
 - torch.Tensor: The tensor resulting from the maximal interpolated inverse FFT.
 """
 # Ensure the input tensors are on the same device and dtype
 v0 = v0.to("cuda:0")
 v1 = v1.to("cuda:0")

 if len(v0.shape) == 1:
 fft_v0 = torch.fft.fft(v0)
 fft_v1 = torch.fft.fft(v1)
 result_fft = torch.zeros_like(fft_v0)
 
 real_v0 = fft_v0.real
 real_v1 = fft_v1.real
 abs_real_v0 = real_v0.abs() 
 abs_real_v1 = real_v1.abs() 

 sign_mask = real_v0.sign() == real_v1.sign()
 larger_values_mask = abs_real_v0 > abs_real_v1

 result_fft.real[sign_mask] = torch.max(real_v0[sign_mask], real_v1[sign_mask])
 result_fft.real[~sign_mask] = torch.where(larger_values_mask[~sign_mask], real_v0[~sign_mask], real_v1[~sign_mask])

 imag_v0 = fft_v0.imag
 imag_v1 = fft_v1.imag
 abs_imag_v0 = imag_v0.abs()
 abs_imag_v1 = imag_v1.abs()
 larger_values_mask_imag = abs_imag_v0 > abs_imag_v1

 result_fft.imag[sign_mask] = torch.max(imag_v0[sign_mask], imag_v1[sign_mask])
 result_fft.imag[~sign_mask] = torch.where(larger_values_mask_imag[~sign_mask], imag_v0[~sign_mask], imag_v1[~sign_mask])

 merged_tensor = torch.fft.ifft(result_fft).real # Taking the real part
 del v0, v1, fft_v0, fft_v1, result_fft
 return merged_tensor

 # Perform the 2D FFT on both tensors
 fft_v0 = torch.fft.fftn(v0, dim=(-2, -1))
 fft_v1 = torch.fft.fftn(v1, dim=(-2, -1))
 
 # Initialize the result FFT tensor
 result_fft = torch.zeros_like(fft_v0)
 
 # Compare real parts of the coefficients
 real_v0 = fft_v0.real
 real_v1 = fft_v1.real
 abs_real_v0 = real_v0.abs()
 abs_real_v1 = real_v1.abs()
 
 # Create masks for where signs match and where they do not
 sign_mask = real_v0.sign() == real_v1.sign()
 larger_values_mask = abs_real_v0 > abs_real_v1
 
 # Where signs match, interpolate; where signs do not match, take the larger by magnitude
 result_fft.real[sign_mask] = torch.max(real_v0[sign_mask], real_v1[sign_mask])
 result_fft.real[~sign_mask] = torch.where(larger_values_mask[~sign_mask], real_v0[~sign_mask], real_v1[~sign_mask])
 
 # Assuming the imaginary part should be treated similarly, adjust this if not
 imag_v0 = fft_v0.imag
 imag_v1 = fft_v1.imag
 abs_imag_v0 = imag_v0.abs()
 abs_imag_v1 = imag_v1.abs()
 larger_values_mask_imag = abs_imag_v0 > abs_imag_v1
 
 result_fft.imag[sign_mask] = torch.max(imag_v0[sign_mask], imag_v1[sign_mask])
 result_fft.imag[~sign_mask] = torch.where(larger_values_mask_imag[~sign_mask], imag_v0[~sign_mask], imag_v1[~sign_mask])
 
 # Perform the inverse FFT to go back to the spatial domain
 merged_tensor = torch.fft.ifftn(result_fft, dim=(-2, -1)).real # Taking the real part
 
 return merged_tensor