U Ph[I@sddlmZddlmZddlmZddlZddlZddl m Z m Z ddl m Z ddlmZmZdd d d d d ddddddddddddddddddd d!d"d#d$d%d&gZd'd'd(d)d&Zdkd-d'd.d/d0dZd'd1d1d'd2d3d Zd4d Zd'd'd5d6d Zdld'd7d.d8d9d:d Zdmd'd'd;d:0yE>Frbool)abr*cCsFt||dd^}}t|tjr2tj|||||dStj|||||dS)z7`np.allclose` with equivalent implementation for torch.T) wrap_sequence)rtolatol equal_nan)r r+r,r-r r0)r7r8r:r;r<_r2r2r3r Is zint | Sequence[int])r)srcdstr*cCs(t|tjrt|||St|||S)z `moveaxis` for pytorch and numpy)r+r0Tensormovedimr,r )r)r>r?r2r2r3r Qs cCs<t|tjrt||S|dtj||jdkddS)z3`np.in1d` with equivalent implementation for torch.).Ndevice) r+r,r-r r0tensorrCanyview)r)yr2r2r3r Xs  )r7r*cCs.t|tjrt|||}nt|||}|S)z3`np.clip` with equivalent implementation for torch.)r+r,r-r r0clamp)r7a_mina_maxresultr2r2r3r _s z int | NonezNdarrayOrTensor | float | int)r)dimkeepdimr*cKst|tjddd}|dk|dkBr8td|dt|tjs^t|tjrt|dkrt|tjdd}tj ||f||d |}t ||d}n$t |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_typer9rdz*q values must be in [0, 100], got values: .i@B)rO)axiskeepdimsgY@)rMrN) rr,r-rF ValueErrorr+r0r@numelrr quantile)r)qrMrNkwargsZq_np_xrLr2r2r3ris&) conditionr*cCszt|tjr0|dk r$t|||}qvt|}nF|dk rltj||jd}tj||j|jd}t|||}n t|}|S)zA Note that `torch.where` may convert y.dtype to x.dtype. NrB)rCdtype)r+r,r-rr0 as_tensorrCr[)rZr)rHrLr2r2r3rs   cCs t|tjrt|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-rr0)r7r2r2r3rs  rD)r7rRr*cCs(t|tjrtj||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. rR)rM)r+r,r-rr0)r7rRr2r2r3rs cCs(t|tjrt|dSt|S)z`np.nonzero` with equivalent implementation for torch. Args: x: array/tensor. Returns: Index unravelled for given shape r)r+r,r-rr0flattenr1r2r2r3rs cCs>t|tjr2ttdr&tj||ddSt||St||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. )rfloor) rounding_mode)r+r0r@rdivrr,r7r8r2r2r3rs   )r*cCs\t|tjrJg}t|D]}|||t||}qt|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+r0r@reversedappendrr r,asarrayr)idxshapecoordrMr2r2r3rs   cs2t|dtjrtjntj}|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|qSr2)r).0irir2r3 sz#unravel_indices..)r+r0r@r r,)rhriZ lib_stackr2rmr3rs cCs4t|tjr*ttdr|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+r0r@hasattrrr^ contiguousr,r1r2r2r3rs   )r)rRr*c Cspt|tjrt||St|ts(|gn|}|D]:}zt||}Wq0tk rht||}Yq0Xq0|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-rFrr0 RuntimeErrorr6)r)rRaxr2r2r3rs  cCs0t|tjr$t|tjr$t||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+r0r@rr,rdr2r2r3r+s  zSequence[NdarrayOrTensor]int)to_catrRr*cCs.t|dtjrt|||Stj|||dS)zH`np.concatenate` with equivalent implementation for torch (`torch.cat`).r)rMout)r+r,r-rr0cat)rtrRrur2r2r3r:scKsNt|tjrt||S|dkr8tj|dddf|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. NrrM)r+r,r-rr0)r7rRrXr2r2r3rAs  cCs t|tjst|St|S)z7`np.isfinite` with equivalent implementation for torch.)r+r0r@r,rr1r2r2r3rTs  )r7vr*cKs@|rdnd}t|tjr(t||||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-rr0)r7rwrxsorterrXsider2r2r3r[s  )r7repeatsrRr*cKs2t|tjrt|||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. rM)r+r,r-rr0repeat_interleave)r7r|rRrXr2r2r3rns cCs t|tjrt|St|S)z^`np.isnan` with equivalent implementation for torch. Args: x: array/tensor. )r+r,r-rr0r1r2r2r3rs  TzNdarrayTensor | TzNdarrayOrTensor | TcKs@t|tjr$|jdkr|St|St|tjr<|jf|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. r)r+r,r-ndimrr0r@rpr)rXr2r2r3rs     zSequence[NdarrayTensor])r)rMr*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-r r0)r)rMr2r2r3r s T)r)rMto_longr*cCsB|r tjnd}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[)r0int64rr@r!valuesr )r)rMrr[Zx_tr=o_tor2r2r3r!s cKs,t|tjtfrtj|f|Stj|f|S)za`torch.unique` with equivalent implementation for numpy. Args: x: array/tensor. )r+r,r-listr"r0rr2r2r3r"scCsDt|tjr ttdr t|St|tjr8tj|Stj|S)ze`torch.linalg.inv` with equivalent implementation for numpy. Args: x: array/tensor. inverse)r+r0r@rorlinalginvr,r1r2r2r3 linalg_invs rzint | tuple | NonecKsv|dkr6t|tjtfr&tj|f|n tj|f|}ns .r)r4r5F)NF)NN)rD)rN)N)FN)N)rDT)N)N)N)N)NF)N)4 __future__rcollections.abcrtypingrnumpyr,r0monai.config.type_definitionsrrmonai.utils.miscrZmonai.utils.type_conversionrr __all__r(r r r r rrrrrrrrrrrrrrrrrr~rr r!r"rr#r&r%r$r'rr2r2r2r3 s    # #