o ) iQ@sdddlmZddlZddlmZdd d ZdddZd ddZd ddZd!d"ddZd!d#ddZ dS)$) annotationsN)Tensorxrshiftint 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)rrrleftrightrc/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/blocks/fft_utils_t.pyroll_1ds r list[int] shift_dimscCsRt|t|krtdt|dt|dt||D] \}}t|||}q|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 rcCs>dgt|}t|D] \}}|j|d||<q t|||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_numrrrfftshiftEs r"cCsBdgt|}t|D]\}}|j|dd||<q t|||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 rrrr rrr ifftshift[s r$Tksp spatial_dims is_complexboolcCtt| d}|r*|jddkrtd|jddtt|dd|D}nt||}ttjj ||dd }t |d d|D}|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 ().cSg|]}|dqSr#r.0rrrr z$ifftn_centered_t..orthor normcSr,r-rr.rrrr0r1) listrangerrr view_as_complexr$ view_as_realfftifftnr")r%r&r'dimsroutrrrifftn_centered_tq r=imcCr)) 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+cSr,r-rr.rrrr0r1z#fftn_centered_t..r2r3cSr,r-rr.rrrr0r1) r5r6rrr r7r$r8r9fftnr")r?r&r'r;rr<rrrfftn_centered_tr>rA)rrrrrrrr)rrrrrrrr)rrrrrr)T)r%rr&rr'r(rr)r?rr&rr'r(rr) __future__rr rrrr"r$r=rArrrrs       ,