U PhG @sddlmZddlZddlmZdddddddZdd d dd d d Zdd dd ddZdd dd ddZddddddddZddddddddZ dS)) annotationsN)Tensorrint)xshift shift_dimreturncCs\|||}|dkr|S||d|||}||||||}tj||f|dS)a{ Similar to roll but for only one dim. Args: x: input data (k-space or image) that can be 1) real-valued: the shape is (C,H,W) for 2D spatial inputs and (C,H,W,D) for 3D, or 2) complex-valued: the shape is (C,H,W,2) for 2D spatial data and (C,H,W,D,2) for 3D. C is the number of channels. shift: the amount of shift along each of shift_dims dimension shift_dim: the dimension over which the shift is applied Returns: 1d-shifted version of x Note: This function is called when fftshift and ifftshift are not available in the running pytorch version r)dim)sizenarrowtorchcat)rrrleftrightrV/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/blocks/fft_utils_t.pyroll_1ds rz list[int])rr shift_dimsrcCsRt|t|kr.tdt|dt|dt||D]\}}t|||}q8|S)a Similar to np.roll but applies to PyTorch Tensors Args: x: input data (k-space or image) that can be 1) real-valued: the shape is (C,H,W) for 2D spatial inputs and (C,H,W,D) for 3D, or 2) complex-valued: the shape is (C,H,W,2) for 2D spatial data and (C,H,W,D,2) for 3D. C is the number of channels. shift: the amount of shift along each of shift_dims dimensions shift_dims: dimensions over which the shift is applied Returns: shifted version of x Note: This function is called when fftshift and ifftshift are not available in the running pytorch version z$len(shift) != len(shift_dims), got fz and f.)len ValueErrorzipr)rrrsdrrrroll-s r)rrrcCs>dgt|}t|D]\}}|j|d||<qt|||S)a9 Similar to np.fft.fftshift but applies to PyTorch Tensors Args: x: input data (k-space or image) that can be 1) real-valued: the shape is (C,H,W) for 2D spatial inputs and (C,H,W,D) for 3D, or 2) complex-valued: the shape is (C,H,W,2) for 2D spatial data and (C,H,W,D,2) for 3D. C is the number of channels. shift_dims: dimensions over which the shift is applied Returns: fft-shifted version of x Note: This function is called when fftshift is not available in the running pytorch version rr enumerateshaperrrriZdim_numrrrfftshiftEsr!cCsBdgt|}t|D]\}}|j|dd||<qt|||S)a< Similar to np.fft.ifftshift but applies to PyTorch Tensors Args: x: input data (k-space or image) that can be 1) real-valued: the shape is (C,H,W) for 2D spatial inputs and (C,H,W,D) for 3D, or 2) complex-valued: the shape is (C,H,W,2) for 2D spatial data and (C,H,W,D,2) for 3D. C is the number of channels. shift_dims: dimensions over which the shift is applied Returns: ifft-shifted version of x Note: This function is called when ifftshift is not available in the running pytorch version rrrrrrr ifftshift[sr#Tbool)ksp spatial_dims is_complexrcCstt| d}|rL|jddkr8td|jddtt| dd}tt| d}t||}|rttjjt ||dd}nttjj||dd}t ||}|S) aC Pytorch-based ifft for spatial_dims-dim signals. "centered" means this function automatically takes care of the required ifft and fft shifts. This is equivalent to do fft in numpy based on numpy.fft.ifftn, numpy.fft.fftshift, and numpy.fft.ifftshift Args: ksp: k-space data that can be 1) real-valued: the shape is (C,H,W) for 2D spatial inputs and (C,H,W,D) for 3D, or 2) complex-valued: the shape is (C,H,W,2) for 2D spatial data and (C,H,W,D,2) for 3D. C is the number of channels. spatial_dims: number of spatial dimensions (e.g., is 2 for an image, and is 3 for a volume) is_complex: if True, then the last dimension of the input ksp is expected to be 2 (representing real and imaginary channels) Returns: "out" which is the output image (inverse fourier of ksp) Example: .. code-block:: python import torch ksp = torch.ones(1,3,3,2) # the last dim belongs to real/imaginary parts # output1 and output2 will be identical output1 = torch.fft.ifftn(torch.view_as_complex(torch.fft.ifftshift(ksp,dim=(-3,-2))), dim=(-2,-1), norm="ortho") output1 = torch.fft.fftshift( torch.view_as_real(output1), dim=(-3,-2) ) output2 = ifftn_centered(ksp, spatial_dims=2, is_complex=True) rrzksp.shape[-1] is not 2 ().r"orthor norm) listrangerrr#r view_as_realfftifftnview_as_complexr!)r%r&r'rdimsroutrrrifftn_centered_tqs   r5)imr&r'rcCstt| d}|rL|jddkr8td|jddtt| dd}tt| d}t||}|rttjjt ||dd}nttjj||dd}t ||}|S) a, Pytorch-based fft for spatial_dims-dim signals. "centered" means this function automatically takes care of the required ifft and fft shifts. This is equivalent to do ifft in numpy based on numpy.fft.fftn, numpy.fft.fftshift, and numpy.fft.ifftshift Args: im: image that can be 1) real-valued: the shape is (C,H,W) for 2D spatial inputs and (C,H,W,D) for 3D, or 2) complex-valued: the shape is (C,H,W,2) for 2D spatial data and (C,H,W,D,2) for 3D. C is the number of channels. spatial_dims: number of spatial dimensions (e.g., is 2 for an image, and is 3 for a volume) is_complex: if True, then the last dimension of the input im is expected to be 2 (representing real and imaginary channels) Returns: "out" which is the output kspace (fourier of im) Example: .. code-block:: python import torch im = torch.ones(1,3,3,2) # the last dim belongs to real/imaginary parts # output1 and output2 will be identical output1 = torch.fft.fftn(torch.view_as_complex(torch.fft.ifftshift(im,dim=(-3,-2))), dim=(-2,-1), norm="ortho") output1 = torch.fft.fftshift( torch.view_as_real(output1), dim=(-3,-2) ) output2 = fftn_centered(im, spatial_dims=2, is_complex=True) rr(rzimg.shape[-1] is not 2 (r)r"r*r+) r-r.rrr#r r/r0fftnr2r!)r6r&r'rr3rr4rrrfftn_centered_ts   r8)T)T) __future__rr rrrr!r#r5r8rrrr s  0