:mod:`heat.indexing`
=========================
.. py:module:: heat.core.indexing

.. autoapi-nested-parse::

   Functions relating to indices of items within DNDarrays, i.e. `where()`



Module Contents
---------------


.. function:: nonzero(x: heat.core.dndarray.DNDarray, as_tuple: bool = True) -> tuple[heat.core.dndarray.DNDarray, ...] | heat.core.dndarray.DNDarray

   Return a tuple of :class:`~heat.core.dndarray.DNDarray`s, one for each dimension of ``x``,
   containing the indices of the non-zero elements in that dimension. If ``x`` is split then
   the result is split in the first dimension. However, this :class:`~heat.core.dndarray.DNDarray`
   can be UNBALANCED as it contains the indices of the non-zero elements on each node.
   The values in ``x`` are always tested and returned in row-major, C-style order.
   The corresponding non-zero values can be obtained with: ``x[nonzero(x)]``.

   :param x: Input array
   :type x: DNDarray
   :param as_tuple: Default is True for numpy-style nonzero output. If False, the output is a torch-style single 2D ``DNDarray`` of shape `(num_nonzero, ndim)` containing the indices of the non-zero elements.
   :type as_tuple: bool, optional

   .. rubric:: Examples

   >>> import heat as ht
   >>> x = ht.array([[3, 0, 0], [0, 4, 1], [0, 6, 0]], split=0)
   >>> ht.nonzero(x)
   (DNDarray([0, 1, 1, 2], dtype=ht.int64, device=cpu:0, split=None),
       DNDarray([0, 1, 2, 1], dtype=ht.int64, device=cpu:0, split=None))
   >>> y = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=0)
   >>> y > 3
   DNDarray([[False, False, False],
             [ True,  True,  True],
             [ True,  True,  True]], dtype=ht.bool, device=cpu:0, split=0)
   >>> ht.nonzero(y > 3)
   DNDarray([[1, 0],
             [1, 1],
             [1, 2],
             [2, 0],
             [2, 1],
             [2, 2]], dtype=ht.int64, device=cpu:0, split=0)
   (DNDarray([1, 1, 1, 2, 2, 2], dtype=ht.int64, device=cpu:0, split=None),
       DNDarray([0, 1, 2, 0, 1, 2], dtype=ht.int64, device=cpu:0, split=None))
   >>> y[ht.nonzero(y > 3)]
   DNDarray([4, 5, 6, 7, 8, 9], dtype=ht.int64, device=cpu:0, split=0)


.. function:: where(cond: heat.core.dndarray.DNDarray, x: None | int | float | heat.core.dndarray.DNDarray = None, y: None | int | float | heat.core.dndarray.DNDarray = None) -> heat.core.dndarray.DNDarray

   Return a :class:`~heat.core.dndarray.DNDarray` containing elements chosen from ``x`` or ``y`` depending on condition.
   Result is a :class:`~heat.core.dndarray.DNDarray` with elements from ``x`` where ``cond`` is True, and from ``y`` elsewhere.

   If only ``cond`` is provided, this function acts as a shorthand for :func:`nonzero`.

   :param cond: Condition of interest, where true yield ``x`` otherwise yield ``y``
   :type cond: DNDarray
   :param x: Values from which to choose. ``x``, ``y`` and condition need to be broadcastable to some shape.
   :type x: DNDarray or int or float, optional
   :param y: Values from which to choose. ``x``, ``y`` and condition need to be broadcastable to some shape.
   :type y: DNDarray or int or float, optional

   :raises NotImplementedError: if splits of the two input :class:`~heat.core.dndarray.DNDarray` differ
   :raises TypeError: if only x or y is given or both are not DNDarrays or numerical scalars

   .. rubric:: Notes

   When only condition is provided, this function is a shorthand for :func:`nonzero` and the function returns a tuple
   of :class:`~heat.core.dndarray.DNDarray`, analogously to ``numpy.where``.

   .. rubric:: Examples

   >>> import heat as ht
   >>> x = ht.arange(10, split=0)
   >>> ht.where(x < 5, x, 10 * x)
   DNDarray(MPI-rank: 0, Shape: (10,), Split: 0, Local Shape: (10,), Device: cpu:0, Dtype: int32, Data:
        [ 0,  1,  2,  3,  4, 50, 60, 70, 80, 90])
   >>> y = ht.array([[0, 1, 2], [0, 2, 4], [0, 3, 6]])
   >>> ht.where(y < 4, y, -1)
   DNDarray(MPI-rank: 0, Shape: (3, 3), Split: None, Local Shape: (3, 3), Device: cpu:0, Dtype: int64, Data:
        [[ 0,  1,  2],
         [ 0,  2, -1],
         [ 0,  3, -1]])