o  i @sdZddlmZddlZddlZddlZddlmZddlm Z ddl m Z m Z     d%d&ddZ d'ddZd(ddZd)ddZd*dd Zd'd!d"Zd(d#d$ZdS)+zJ This script contains utility functions for complex-value PyTorch tensor. ) annotationsN)Tensor)NdarrayOrTensor)convert_to_numpyconvert_to_tensorTFdata$NdarrayOrTensor | list | int | floatdtypetorch.dtype | Nonedevicetorch.device | None wrap_sequencebool track_metareturnrcCs t|trt|st|||||d}|Snt|s&t|||||d}|St|tjr8tj|j|j gdd}nMt|tj r^t d|j jdur]|jdkrRt|}tj|j|j fdd}n't|ttfrm|j|j gg}nt|trt|dd }tj|j|j fdd}t|||||d}|S) a# Convert complex-valued data to a 2-channel PyTorch tensor. The real and imaginary parts are stacked along the last dimension. This function relies on 'monai.utils.type_conversion.convert_to_tensor' Args: data: input data can be PyTorch Tensor, numpy array, list, int, and float. will convert Tensor, Numpy array, float, int, bool to Tensor, strings and objects keep the original. for list, convert every item to a Tensor if applicable. dtype: target data type to when converting to Tensor. device: target device to put the converted Tensor data. wrap_sequence: if `False`, then lists will recursively call this function. E.g., `[1, 2]` -> `[tensor(1), tensor(2)]`. If `True`, then `[1, 2]` -> `tensor([1, 2])`. track_meta: whether to track the meta information, if `True`, will convert to `MetaTensor`. default to `False`. Returns: PyTorch version of the data Example: .. code-block:: python import numpy as np data = np.array([ [1+1j, 1-1j], [2+2j, 2-2j] ]) # the following line prints (2,2) print(data.shape) # the following line prints torch.Size([2, 2, 2]) print(convert_to_tensor_complex(data).shape) )r r r rdimz[SaUO]NraxisT)r ) isinstancertorch is_complexrnp iscomplexobjstackrealimagndarrayresearchr strndimascontiguousarrayfloatintlistrtolist)rr r r rZconverted_datar(i/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/reconstruction/complex_utils.pyconvert_to_tensor_complexs: %           r*xcCs@|jddkrtd|jdd|dd|dddS)z Compute the absolute value of a complex tensor. Args: x: Input tensor with 2 channels in the last dimension representing real and imaginary parts. Returns: Absolute value along the last dimension rzx.shape[-1] is not 2 (z)..r.g?)shape ValueErrorr+r(r(r) complex_abs_tfs r3rcCst|S)a Compute the absolute value of a complex array. Args: x: Input array/tensor with 2 channels in the last dimension representing real and imaginary parts. Returns: Absolute value along the last dimension Example: .. code-block:: python import numpy as np x = np.array([3,4])[np.newaxis] # the following line prints 5 print(complex_abs(x)) )r3r2r(r(r) complex_absusr4ycCs|jddks|jddkrtd|jdd|jdd|d|d|d|d}|d|d|d|d}tj||fddS) a Compute complex-valued multiplication. Supports Ndim inputs with last dim equal to 2 (real/imaginary channels) Args: x: Input tensor with 2 channels in the last dimension representing real and imaginary parts. y: Input tensor with 2 channels in the last dimension representing real and imaginary parts. Returns: Complex multiplication of x and y rr,'last dim must be 2, but x.shape[-1] is  and y.shape[-1] is .r-r.rr0r1rr)r+r5 real_part imag_partr(r(r) complex_mul_ts  "  r<cCs|jddks|jddkrtd|jdd|jddt|tr)t||S|d|d|d|d}|d|d|d|d}tj||fdd}|S) a Compute complex-valued multiplication. Supports Ndim inputs with last dim equal to 2 (real/imaginary channels) Args: x: Input array/tensor with 2 channels in the last dimension representing real and imaginary parts. y: Input array/tensor with 2 channels in the last dimension representing real and imaginary parts. Returns: Complex multiplication of x and y Example: .. code-block:: python import numpy as np x = np.array([[1,2],[3,4]]) y = np.array([[1,1],[1,1]]) # the following line prints array([[-1, 3], [-1, 7]]) print(complex_mul(x,y)) rr,r6r7r8r-r.r)r0r1rrr<rr)r+r5r:r;multr(r(r) complex_muls"    r>cCs@|jddkrtd|jddtj|d|d fddS)a Compute complex conjugate of a tensor. Supports Ndim inputs with last dim equal to 2 (real/imaginary channels) Args: x: Input tensor with 2 channels in the last dimension representing real and imaginary parts. Returns: Complex conjugate of x rr,r6r8r-r.rr9r2r(r(r)complex_conj_ts r?cCsV|jddkrtd|jddt|trt|Stj|d|d fdd}|S)a Compute complex conjugate of an/a array/tensor. Supports Ndim inputs with last dim equal to 2 (real/imaginary channels) Args: x: Input array/tensor with 2 channels in the last dimension representing real and imaginary parts. Returns: Complex conjugate of x Example: .. code-block:: python import numpy as np x = np.array([[1,2],[3,4]]) # the following line prints array([[ 1, -2], [ 3, -4]]) print(complex_conj(x)) rr,r6r8r-r.r)r0r1rrr?rr)r+Znp_conjr(r(r) complex_conjs  r@)NNTF) rrr r r r r rrrrr)r+rrr)r+rrr)r+rr5rrr)r+rr5rrr)__doc__ __future__rrnumpyrrrmonai.config.type_definitionsrmonai.utils.type_conversionrrr*r3r4r<r>r?r@r(r(r(r)s&    K    "