o  iI@sddlmZddlmZddlmZddlZddlZddl m Z m Z ddl m Z mZgdZdkd d ZdldmddZdnddZddZdoddZ dpdqd#d$Zdrdsd&d'Zdtd(d)Zdudvd,d-Zdkd.d/Zdod0d1Zdwd2d3Zdwd4d5Zdkd6d7Zdxd8d9Zdyd:d;Zdzd{d?d@Z d|dodAdBZ!dkdCdDZ"d}d~dFdGZ#d|ddIdJZ$dkdKdLZ%edMZ&ddPdQZ'ddSdTZ(dddWdXZ)ddYdZZ*dd[d\Z+d|dd^d_Z,d|dd`daZ-d|ddbdcZ.d|ddddeZ/dpddgdhZ0d|ddidjZ1dS)) annotations)Sequence)TypeVarN)NdarrayOrTensor NdarrayTensor)convert_data_typeconvert_to_dst_type)allclosemoveaxisin1dclip percentilewhereargwhereargsortnonzero floor_divide unravel_indexunravel_indicesravel any_np_ptmaximum concatenatecumsumisfinite searchsortedrepeatisnanascontiguousarraystackmodeuniquemaxminmedianmeanstdsoftplusxrreturncCs0t|tjrtt||Stt||S)zstable softplus through `np.logaddexp` with equivalent implementation for torch. Args: x: array/tensor. Returns: Softplus of the input. ) isinstancenpndarray logaddexp zeros_liketorchr(r1r/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/transforms/utils_pytorch_numpy_unification.pyr':s r'h㈵>:0yE>FarbboolcCsFt||dd^}}t|tjrtj|||||dStj|||||dS)z7`np.allclose` with equivalent implementation for torch.T) wrap_sequence)rtolatol equal_nan)rr*r+r,r r/)r5r6r9r:r;_r1r1r2r Hs r srcint | Sequence[int]dstcCs(t|tjr t|||St|||S)z `moveaxis` for pytorch and numpy)r*r/Tensormovedimr+r )r(r=r?r1r1r2r Ps r cCs<t|tjr t||S|dtj||jdkddS)z3`np.in1d` with equivalent implementation for torch.).Ndevice) r*r+r,isinr/tensorrCanyview)r(yr1r1r2r Ws  $r cCs0t|tjrt|||}|St|||}|S)z3`np.clip` with equivalent implementation for torch.)r*r+r,r r/clamp)r5a_mina_maxresultr1r1r2r ^s r dim int | NonekeepdimNdarrayOrTensor | float | intcKst|tjddd}|dk|dkBrtd|dt|tjs/t|tjrNt|dkrNt|tjdd}tj ||f||d |}t ||d}|St |d |d}tj ||||d }|S) a(`np.percentile` with equivalent implementation for torch. Pytorch uses `quantile`. For more details please refer to: https://pytorch.org/docs/stable/generated/torch.quantile.html. https://numpy.org/doc/stable/reference/generated/numpy.percentile.html. Args: x: input data. q: percentile to compute (should in range 0 <= q <= 100). dim: the dim along which the percentiles are computed. default is to compute the percentile along a flattened version of the array. keepdim: whether the output data has dim retained or not. kwargs: if `x` is numpy array, additional args for `np.percentile`, more details: https://numpy.org/doc/stable/reference/generated/numpy.percentile.html. Returns: Resulting value (scalar) T) output_typer8rdz*q values must be in [0, 100], got values: .i@B)rR)axiskeepdimsgY@)rNrP) rr+r,rG ValueErrorr*r/r@numelr rquantile)r(qrNrPkwargsZq_np_xrMr1r1r2r hs&r conditioncCst|tjr|durt|||}|St|}|S|dur9tj||jd}tj||j|jd}t|||}|St|}|S)zA Note that `torch.where` may convert y.dtype to x.dtype. NrB)rCdtype)r*r+r,rr/ as_tensorrCr^)r]r(rIrMr1r1r2rs   rcC t|tjr t|St|S)a`np.argwhere` with equivalent implementation for torch. Args: a: input data. Returns: Indices of elements that are non-zero. Indices are grouped by element. This array will have shape (N, a.ndim) where N is the number of non-zero items. )r*r+r,rr/)r5r1r1r2rs  rrDrUcCs(t|tjr tj||dStj||dS)z`np.argsort` with equivalent implementation for torch. Args: a: the array/tensor to sort. axis: axis along which to sort. Returns: Array/Tensor of indices that sort a along the specified axis. rU)rN)r*r+r,rr/)r5rUr1r1r2rs rcCs(t|tjr t|dSt|S)z`np.nonzero` with equivalent implementation for torch. Args: x: array/tensor. Returns: Index unravelled for given shape r)r*r+r,rr/flattenr0r1r1r2rs rcCs*t|tjr t||Stt||S)aD`np.floor_divide` with equivalent implementation for torch. As of pt1.8, use `torch.div(..., rounding_mode="floor")`, and before that, use `torch.floor_divide`. Args: a: first array/tensor b: scalar to divide by Returns: Element-wise floor division between two arrays/tensors. )r*r/r@rr+asarrayr5r6r1r1r2rs rcCs\t|tjr%g}t|D]}|||t||}q t|dddStt ||S)z`np.unravel_index` with equivalent implementation for torch. Args: idx: index to unravel. shape: shape of array/tensor. Returns: Index unravelled for given shape NrD) r*r/r@reversedappendrrr+rcr)idxshapecoordrNr1r1r2rs   rcs2t|dtjr tjntj}|fdd|DS)zComputing unravel coordinates from indices. Args: idx: a sequence of indices to unravel. shape: shape of array/tensor. Returns: Stacked indices unravelled for given shape rcsg|]}t|qSr1)r).0irhr1r2 sz#unravel_indices..)r*r/r@rr+)rgrhZ lib_stackr1rlr2rs rcCs4t|tjrttdr|S|St|S)z`np.ravel` with equivalent implementation for torch. Args: x: array/tensor to ravel. Returns: Return a contiguous flattened array/tensor. r)r*r/r@hasattrrrb contiguousr+r0r1r1r2rs   rc Cslt|tjr t||St|ts|gn|}|D]}zt||}Wqty3t||}Yqw|S)a`np.any` with equivalent implementation for torch. For pytorch, convert to boolean for compatibility with older versions. Args: x: input array/tensor. axis: axis to perform `any` over. Returns: Return a contiguous flattened array/tensor. )r*r+r,rGrr/ RuntimeErrorr7)r(rUaxr1r1r2rs  rcCs0t|tjrt|tjrt||St||S)z`np.maximum` with equivalent implementation for torch. Args: a: first array/tensor. b: second array/tensor. Returns: Element-wise maximum between two arrays/tensors. )r*r/r@rr+rdr1r1r2r)s  rto_catSequence[NdarrayOrTensor]intcCs.t|dtjrt|||Stj|||dS)zH`np.concatenate` with equivalent implementation for torch (`torch.cat`).r)rNout)r*r+r,rr/cat)rrrUrur1r1r2r8srcKsRt|tjr t||S|durtj|dddfi|Stj|fd|i|S)aH `np.cumsum` with equivalent implementation for torch. Args: a: input data to compute cumsum. axis: expected axis to compute cumsum. kwargs: if `a` is PyTorch Tensor, additional args for `torch.cumsum`, more details: https://pytorch.org/docs/stable/generated/torch.cumsum.html. NrrN)r*r+r,rr/)r5rUr[r1r1r2r?s rcCs t|tjs t|St|S)z7`np.isfinite` with equivalent implementation for torch.)r*r/r@r+rr0r1r1r2rRs   rvcKs@|rdnd}t|tjrt||||Stj||fd|i|S)ay `np.searchsorted` with equivalent implementation for torch. Args: a: numpy array or tensor, containing monotonically increasing sequence on the innermost dimension. v: containing the search values. right: if False, return the first suitable location that is found, if True, return the last such index. sorter: if `a` is numpy array, optional array of integer indices that sort array `a` into ascending order. kwargs: if `a` is PyTorch Tensor, additional args for `torch.searchsorted`, more details: https://pytorch.org/docs/stable/generated/torch.searchsorted.html. rightleft)r*r+r,rr/)r5rwrxsorterr[sider1r1r2rYs rrepeatscKs2t|tjr t|||Stj||fd|i|S)a `np.repeat` with equivalent implementation for torch (`repeat_interleave`). Args: a: input data to repeat. repeats: number of repetitions for each element, repeats is broadcast to fit the shape of the given axis. axis: axis along which to repeat values. kwargs: if `a` is PyTorch Tensor, additional args for `torch.repeat_interleave`, more details: https://pytorch.org/docs/stable/generated/torch.repeat_interleave.html. rN)r*r+r,rr/repeat_interleave)r5r|rUr[r1r1r2rls rcCr`)z^`np.isnan` with equivalent implementation for torch. Args: x: array/tensor. )r*r+r,rr/r0r1r1r2r}s   rTNdarrayTensor | TNdarrayOrTensor | TcKsDt|tjr|jdkr |St|St|tjr |jdi|S|S)a-`np.ascontiguousarray` with equivalent implementation for torch (`contiguous`). Args: x: array/tensor. kwargs: if `x` is PyTorch Tensor, additional args for `torch.contiguous`, more details: https://pytorch.org/docs/stable/generated/torch.Tensor.contiguous.html. rNr1)r*r+r,ndimrr/r@ror(r[r1r1r2rs   rSequence[NdarrayTensor]cCs(t|dtjrt||St||S)z`np.stack` with equivalent implementation for torch. Args: x: array/tensor. dim: dimension along which to perform the stack (referred to as `axis` by numpy). r)r*r+r,rr/)r(rNr1r1r2rs  rTto_longcCsB|rtjnd}t|tj|d^}}t||j}t||^}}|S)z`torch.mode` with equivalent implementation for numpy. Args: x: array/tensor. dim: dimension along which to perform `mode` (referred to as `axis` by numpy). to_long: convert input to long before performing mode. N)r^)r/int64rr@r valuesr)r(rNrr^Zx_tr<Zo_tor1r1r2r s r cKs4t|tjtfrtj|fi|Stj|fi|S)za`torch.unique` with equivalent implementation for numpy. Args: x: array/tensor. )r*r+r,listr!r/rr1r1r2r!s4r!cCsDt|tjrttdrt|St|tjrtj|Stj|S)ze`torch.linalg.inv` with equivalent implementation for numpy. Args: x: array/tensor. inverse)r*r/r@rnrlinalginvr+r0r1r1r2 linalg_invs $rint | tuple | NonecK|durt|tjtfrtj|fi|ntj|fi|}n t|tjtfr3tj|fd|i|}n tj|t|fi|}t|trH|dS|S)z`torch.max` with equivalent implementation for numpy Args: x: array/tensor. Returns: the maximum of x. NrUr)r*r+r,rr"r/rttupler(rNr[retr1r1r2r"s  6r"cK|dur"t|tjtfrtj|fi|}|Stj|fi|}|St|tjtfr7tj|fd|i|}|Stj|t|fi|}|S)z`torch.mean` with equivalent implementation for numpy Args: x: array/tensor. Returns: the mean of x NrU)r*r+r,rr%r/rtrr1r1r2r% "r%cKr)z`torch.median` with equivalent implementation for numpy Args: x: array/tensor. Returns the median of x. NrU)r*r+r,rr$r/rtrr1r1r2r$rr$cKr)z`torch.min` with equivalent implementation for numpy Args: x: array/tensor. Returns: the minimum of x. NrUr)r*r+r,rr#r/rtrrr1r1r2r#s  6r#unbiasedcCsn|durt|tjtfrt|}|St||}|St|tjtfr,tj||d}|St|t||}|S)z`torch.std` with equivalent implementation for numpy Args: x: array/tensor. Returns: the standard deviation of x. Nra)r*r+r,rr&r/rt)r(rNrrr1r1r2r&&s  r&cKr)z`torch.sum` with equivalent implementation for numpy Args: x: array/tensor. Returns: the sum of x. NrU)r*r+r,rsumr/rtrr1r1r2r<rr)r(rr)r)r3r4F)r5rr6rr)r7)r(rr=r>r?r>r)r)r5rr)r)NF)r(rrNrOrPr7r)rQ)NN)r]rr)r)r5rr)r)rD)r5rrUrOr)r)r)r)r(rrUr>r)r)r5rr6rr)r)rN)rrrsrUrtr)r)N)FN)r5rrwrr)r)r5rr|rtrUrOr)r)r(rr)r)r(rrNrtr)r)rDT)r(rrNrtrr7r)r)r(rr)r)r(rrNrr)r)r(rrNrrr7r)r)2 __future__rcollections.abcrtypingrnumpyr+r/monai.config.type_definitionsrrZmonai.utils.type_conversionrr__all__r'r r r r r rrrrrrrrrrrrrrrrr~rrr r!rr"r%r$r#r&rr1r1r1r2sV   #    #