Package: torch-frft
Trainable Fractional Fourier TransformA differentiable fractional Fourier transform (FRFT) implementation with layers that can be trained end-to-end with the rest of the network. This package provides implementations of both fast computations of continuous FRFT and discrete FRFT (DFRFT) and pre-configured layers that are eligible for use in neural networks.The fast transform approximates the continuous FRFT and is based onDigital computation of the fractional Fourier transformpaper. The DFRFT is based onThediscrete fractional Fourier transform_](https://ieeexplore.ieee.org/document/839980) paper. MATLAB implementations of both approaches are provided onHaldun M. Özaktaş's pageasfracF.manddFRT.m, respectively.This package implements these approaches in PyTorch with specific optimizations and, most notably, adds the ability to apply the transform along a particular tensor dimension.We provide primer layers that extendtorch.nn.Modulefor continuous and discrete transforms, an example of the custom layer implementation, is also provided in theREADME.mdfile.Table of ContentsTrainable Fractional Fourier TransformTable of ContentsInstallationFor UsageFor DevelopmentUsageTransformsPre-configured LayersCustom LayersFRFT ShiftInstallationFor UsageYou can install the package directly fromPYPIusingpiporpoetryas follows:pipinstalltorch-frftorpoetryaddtorch-frftFor DevelopmentThis codebase utilizesPoetryfor package management. To install the dependencies:poetryinstallor one can install the dependencies provided inrequirements.txtusingpiporconda, e,g.,pipinstall-rrequirements.txtUsageTransforms:warning: Transforms applied in the same device as the input tensor. If the input tensor is on GPU, the transform will also be applied on GPU.The package provides transform functions that operate on the $n^{th}$ dimension of an input tensor,frft()anddfrft(), which correspond to the fast computation of continuous fractional Fourier transform (FRFT) and discrete fractional Fourier transform (DFRFT), respectively. It also provides a functiondfrftmtx(), which computes the DFRFT matrix for a given length and order, similar to MATLAB'sdftmtx()function for the ordinary DFT matrix. Note that thefrft()only operates on even-sized lengths as in the original MATLAB implementationfracF.m.importtorchfromtorch_frft.frft_moduleimportfrftfromtorch_frft.dfrft_moduleimportdfrft,dfrftmtxdevice=torch.device("cuda"iftorch.cuda.is_available()else"cpu")N=128a=0.5X=torch.rand(N,N,device=device)Y1=frft(X,a)# equivalent to dim=-1Y2=frft(X,a,dim=0)# 2D FRFTa0,a1=1.25,0.75Y3=frft(frft(X,a0,dim=0),a1,dim=1)Pre-configured LayersThe package also provides two differentiable FRFT layers,FrFTLayerandDFrFTLayer, which can be used as follows:importtorchimporttorch.nnasnnfromtorch_frft.dfrft_moduleimportdfrftfromtorch_frft.layerimportDFrFTLayer,FrFTLayer# FRFT with initial order 1.25, operating on the last dimensionmodel=nn.Sequential(FrFTLayer(order=1.25,dim=-1))# DFRFT with initial order 0.75, operating on the first dimensionmodel=nn.Sequential(DFrFTLayer(order=0.75,dim=0))Then, the simplest toy example to train the layer is as follows:device=torch.device("cuda"iftorch.cuda.is_available()else"cpu")num_samples,seq_length=100,16a_original=1.1a_initial=1.25X=torch.randn(num_samples,seq_length,dtype=torch.float32,device=device)Y=dfrft(X,a_original)model=DFrFTLayer(order=a_initial).to(device)optim=torch.optim.Adam(model.parameters(),lr=1e-3)epochs=1000forepochinrange(1+epochs):optim.zero_grad()loss=torch.norm(Y-model(X))loss.backward()optim.step()print("Original  a:",a_original)print("Estimated a:",model.order.item())One can also place these layers directly into thetorch.nn.Sequential. Remark that these transforms generate complex-valued outputs, so one may need to convert them to real-valued outputs, e.g., taking the real part, absolute value, etc. For example, the following code snippet implements a simple fully connected network with real parts ofFrFTLayerandDFrFTLayerin between:classReal(nn.Module):def__init__(self)->None:super().__init__()defforward(self,x:torch.Tensor)->torch.Tensor:returnx.realmodel=nn.Sequential(nn.Linear(16,6),nn.ReLU(),DFrFTLayer(1.35,dim=-1),Real(),nn.ReLU(),nn.Linear(6,1),FrFTLayer(0.65,dim=0),Real(),nn.ReLU(),)Custom LayersCreating custom layers with the providedfrft()anddfrft()functions is also possible. The below example contains in-betweenLinearandReLUlayers and the same fractional order for forward and backward DFRFT transforms is as follows:importtorchimporttorch.nnasnnfromtorch_frft.dfrft_moduleimportdfrftclassCustomLayer(nn.Module):def__init__(self,in_feat:int,out_feat:int,*,order:float=1.0,dim:int=-1)->None:super().__init__()self.in_features=in_featself.out_features=out_featself.order=nn.Parameter(torch.tensor(order,dtype=torch.float32),requires_grad=True)self.dim=dimdefforward(self,x:torch.Tensor)->torch.Tensor:x1=dfrft(x,self.order,dim=self.dim)a1=nn.ReLU()(x1.real)+1j*nn.ReLU()(x1.imag)x2=nn.Linear(self.in_features,self.in_features,dtype=a1.dtype,device=x.device)(a1)a2=nn.ReLU()(x2.real)+1j*nn.ReLU()(x2.imag)x3=dfrft(a2,-self.order,dim=self.dim)a3=nn.ReLU()(x3.real)+1j*nn.ReLU()(x3.imag)x4=nn.Linear(self.in_features,self.out_features,dtype=a3.dtype,device=x.device)(a3)a4=nn.ReLU()(x4.real)returna4Then, a simple training example for the givenCustomLayercan be given as follows:device=torch.device("cuda"iftorch.cuda.is_available()else"cpu")num_samples,seq_length,out_length=100,32,5X=torch.rand(num_samples,seq_length,device=device)y=torch.rand(num_samples,out_length,device=device)model=CustomLayer(seq_length,out_length,order=1.25)optim=torch.optim.Adam(model.parameters(),lr=1e-3)epochs=1000forepochinrange(epochs):optim.zero_grad()loss=torch.nn.MSELoss()output=loss(model(X),y)output.backward()optim.step()ifepoch%100==0:print(f"Epoch{epoch:4d}| Loss{output.item():.4f}")print("Final a:",model.order)FRFT ShiftNote that the fast computation of continuous FRFT is defined for the central grid of $\left[-\lfloor\frac{N}{2}\rfloor, \lfloor\frac{N-1}{2}\rfloor\right]$. Therefore, we needfftshift()to create equivalence with the original FFT when the transform order is precisely $1$. In this package, we also provide a shifted version of the fast FRFT computation,frft_shifted(), which operates with the assumption that the grid is $[0, N-1]$. The latter interval is not the default behavior since we want consistency with the original MATLAB implementation. The all there lines below are equivalent:importtorchfromtorch.fftimportfft,fftshiftfromtorch_frft.frft_moduleimportfrft,frft_shiftedtorch.manual_seed(0)x=torch.rand(100)y1=fft(x,norm="ortho")y2=fftshift(frft(fftshift(x),1.0))y3=frft_shifted(x,1.0)asserttorch.allclose(y1,y2,atol=1e-5)asserttorch.allclose(y1,y3,atol=1e-5)