Package ai.djl.ndarray.index

Class NDIndex

  • public class NDIndex
extends java.lang.Object
    The NDIndex allows you to specify a subset of an NDArray that can be used for fetching or updating.

    It accepts a different index option for each dimension, given in the order of the dimensions. Each dimension has options corresponding to:

    • Return all dimensions - Pass null to addIndices
    • A single value in the dimension - Pass the value to addIndices with a negative index -i corresponding to [dimensionLength - i]
    • A range of values - Use addSliceDim

    We recommend creating the NDIndex using NDIndex(String, Object...).

    See Also:
    NDIndex(String, Object...)

    • Constructor Summary

      Constructors 
      Constructor Description
      NDIndex()
      Creates an empty NDIndex to append values to.
      NDIndex​(long... indices)
      Creates an NDIndex with the given indices as specified values on the NDArray.
      NDIndex​(java.lang.String indices, java.lang.Object... args)
      Creates a NDIndex given the index values.

    • Method Summary

      All Methods Static Methods Instance Methods Concrete Methods 
      Modifier and Type Method Description
      NDIndex addAllDim()
      Appends a new index to get all values in the dimension.
      NDIndex addAllDim​(int count)
      Appends multiple new index to get all values in the dimension.
      NDIndex addBooleanIndex​(NDArray index)
      Updates the NDIndex by appending a boolean NDArray.
      NDIndex addIndices​(long... indices)
      Updates the NDIndex by appending indices as specified values on the NDArray.
      NDIndex addIndices​(java.lang.String indices, java.lang.Object... args)
      Updates the NDIndex by appending indices to the array.
      NDIndex addPickDim​(NDArray index)
      Appends a picking index that gets values by index in the axis.
      NDIndex addSliceDim​(long min, long max)
      Appends a new index to slice the dimension and returns a range of values.
      NDIndex addSliceDim​(long min, long max, long step)
      Appends a new index to slice the dimension and returns a range of values.
      NDIndexElement get​(int dimension)
      Returns the index affecting the given dimension.
      int getEllipsisIndex()
      Returns the index of the ellipsis.
      java.util.List<NDIndexElement> getIndices()
      Returns the indices.
      int getRank()
      Returns the number of dimensions specified in the Index.
      static NDIndex sliceAxis​(int axis, long min, long max)
      Creates an NDIndex that just has one slice in the given axis.
      java.util.stream.Stream<NDIndexElement> stream()
      Returns a stream of the NDIndexElements.

      • Methods inherited from class java.lang.Object

        clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

    • Constructor Detail

      • NDIndex

        public NDIndex()
        Creates an empty NDIndex to append values to.

      • NDIndex

        public NDIndex​(java.lang.String indices,
               java.lang.Object... args)
        Creates a NDIndex given the index values.

        Here are some examples of the indices format. 

             NDArray a = manager.ones(new Shape(5, 4, 3));

     // Gets a subsection of the NDArray in the first axis.
     assertEquals(a.get(new NDIndex("2")).getShape(), new Shape(4, 3));

     // Gets a subsection of the NDArray indexing from the end (-i == length - i).
     assertEquals(a.get(new NDIndex("-1")).getShape(), new Shape(4, 3));

     // Gets everything in the first axis and a subsection in the second axis.
     // You can use either : or * to represent everything
     assertEquals(a.get(new NDIndex(":, 2")).getShape(), new Shape(5, 3));
     assertEquals(a.get(new NDIndex("*, 2")).getShape(), new Shape(5, 3));

     // Gets a range of values along the second axis that is inclusive on the bottom and exclusive on the top.
     assertEquals(a.get(new NDIndex(":, 1:3")).getShape(), new Shape(5, 2, 3));

     // Excludes either the min or the max of the range to go all the way to the beginning or end.
     assertEquals(a.get(new NDIndex(":, :3")).getShape(), new Shape(5, 3, 3));
     assertEquals(a.get(new NDIndex(":, 1:")).getShape(), new Shape(5, 4, 3));

     // Uses the value after the second colon in a slicing range, the step, to get every other result.
     assertEquals(a.get(new NDIndex(":, 1::2")).getShape(), new Shape(5, 2, 3));

     // Uses a negative step to reverse along the dimension.
     assertEquals(a.get(new NDIndex("-1")).getShape(), new Shape(5, 4, 3));

     // Uses a variable argument to the index
     // It can replace any number in any of these formats with {} and then the value of {}
     // is specified in an argument following the indices string.
     assertEquals(a.get(new NDIndex("{}, {}:{}", 0, 1, 3)).getShape(), new Shape(2, 3));

     // Uses ellipsis to insert many full slices
     assertEquals(a.get(new NDIndex("...")).getShape(), new Shape(5, 4, 3));

     // Uses ellipsis to select all the dimensions except for last axis where we only get a subsection.
     assertEquals(a.get(new NDIndex("..., 2")).getShape(), new Shape(5, 4));
        Parameters:
        indices - a comma separated list of indices corresponding to either subsections, everything, or slices on a particular dimension
        args - arguments to replace the variable "{}" in the indices string. Can be an integer, long, boolean NDArray, or integer NDArray.
        See Also:
        Numpy Indexing

      • NDIndex

        public NDIndex​(long... indices)
        Creates an NDIndex with the given indices as specified values on the NDArray.
        Parameters:
        indices - the indices with each index corresponding to the dimensions and negative indices starting from the end

    • Method Detail

      • sliceAxis

        public static NDIndex sliceAxis​(int axis,
                                long min,
                                long max)
        Creates an NDIndex that just has one slice in the given axis.
        Parameters:
        axis - the axis to slice
        min - the min of the slice
        max - the max of the slice
        Returns:
        a new NDIndex with the given slice.

      • getRank

        public int getRank()
        Returns the number of dimensions specified in the Index.
        Returns:
        the number of dimensions specified in the Index

      • getEllipsisIndex

        public int getEllipsisIndex()
        Returns the index of the ellipsis.
        Returns:
        the index of the ellipsis within this index or -1 for none.

      • get

        public NDIndexElement get​(int dimension)
        Returns the index affecting the given dimension.
        Parameters:
        dimension - the affected dimension
        Returns:
        the index affecting the given dimension

      • getIndices

        public java.util.List<NDIndexElement> getIndices()
        Returns the indices.
        Returns:
        the indices

      • addIndices

        public final NDIndex addIndices​(java.lang.String indices,
                                java.lang.Object... args)
        Updates the NDIndex by appending indices to the array.
        Parameters:
        indices - the indices to add similar to NDIndex(String, Object...)
        args - arguments to replace the variable "{}" in the indices string. Can be an integer, long, boolean NDArray, or integer NDArray.
        Returns:
        the updated NDIndex
        See Also:
        NDIndex(String, Object...)

      • addIndices

        public final NDIndex addIndices​(long... indices)
        Updates the NDIndex by appending indices as specified values on the NDArray.
        Parameters:
        indices - with each index corresponding to the dimensions and negative indices starting from the end
        Returns:
        the updated NDIndex

      • addBooleanIndex

        public NDIndex addBooleanIndex​(NDArray index)
        Updates the NDIndex by appending a boolean NDArray.

        The NDArray should have a matching shape to the dimensions being fetched and will return where the values in NDIndex do not equal zero.

        Parameters:
        index - a boolean NDArray where all nonzero elements correspond to elements to return
        Returns:
        the updated NDIndex

      • addAllDim

        public NDIndex addAllDim()
        Appends a new index to get all values in the dimension.
        Returns:
        the updated NDIndex

      • addAllDim

        public NDIndex addAllDim​(int count)
        Appends multiple new index to get all values in the dimension.
        Parameters:
        count - how many axes of NDIndexAll to add.
        Returns:
        the updated NDIndex
        Throws:
        java.lang.IllegalArgumentException - if count is negative

      • addSliceDim

        public NDIndex addSliceDim​(long min,
                           long max)
        Appends a new index to slice the dimension and returns a range of values.
        Parameters:
        min - the minimum of the range
        max - the maximum of the range
        Returns:
        the updated NDIndex

      • addSliceDim

        public NDIndex addSliceDim​(long min,
                           long max,
                           long step)
        Appends a new index to slice the dimension and returns a range of values.
        Parameters:
        min - the minimum of the range
        max - the maximum of the range
        step - the step of the slice
        Returns:
        the updated NDIndex

      • addPickDim

        public NDIndex addPickDim​(NDArray index)
        Appends a picking index that gets values by index in the axis.
        Parameters:
        index - the indices should be NDArray. For each element in the indices array, it acts like a fixed index returning an element of that shape. So, the final shape would be indices.getShape().addAll(target.getShape().slice(1)) (assuming it is the first index element).
        Returns:
        the updated NDIndex

      • stream

        public java.util.stream.Stream<NDIndexElement> stream()
        Returns a stream of the NDIndexElements.
        Returns:
        a stream of the NDIndexElements