tensorstore_demo.DimExpression.__getitem__

tensorstore_demo.DimExpression.__getitem__( self, indices: Any) -> DimExpression[source]

Applies a NumPy-style indexing operation with default index array semantics.

When using NumPy-style indexing with a dimension expression, all selected dimensions must be consumed by a term of the indexing spec; there is no implicit addition of an Ellipsis term to consume any remaining dimensions.

Returns:¶: Dimension expression with the indexing operation added.

Examples

Integer indexing¶

>>> transform = ts.IndexTransform(input_labels=['x', 'y', 'z'])
>>> transform[ts.d['x'][5]]
Rank 2 -> 3 index space transform:
  Input domain:
    0: (-inf*, +inf*) "y"
    1: (-inf*, +inf*) "z"
  Output index maps:
    out[0] = 5
    out[1] = 0 + 1 * in[0]
    out[2] = 0 + 1 * in[1]
>>> transform[ts.d['x', 'z'][5, 6]]
Rank 1 -> 3 index space transform:
  Input domain:
    0: (-inf*, +inf*) "y"
  Output index maps:
    out[0] = 5
    out[1] = 0 + 1 * in[0]
    out[2] = 6

A single scalar index term applies to all selected dimensions:

>>> transform[ts.d['x', 'y'][5]]
Rank 1 -> 3 index space transform:
  Input domain:
    0: (-inf*, +inf*) "z"
  Output index maps:
    out[0] = 5
    out[1] = 5
    out[2] = 0 + 1 * in[0]

Interval indexing¶

>>> transform = ts.IndexTransform(input_labels=['x', 'y', 'z'])
>>> transform[ts.d['x'][5:10]]
Rank 3 -> 3 index space transform:
  Input domain:
    0: [5, 10) "x"
    1: (-inf*, +inf*) "y"
    2: (-inf*, +inf*) "z"
  Output index maps:
    out[0] = 0 + 1 * in[0]
    out[1] = 0 + 1 * in[1]
    out[2] = 0 + 1 * in[2]
>>> transform[ts.d['x', 'z'][5:10, 20:30]]
Rank 3 -> 3 index space transform:
  Input domain:
    0: [5, 10) "x"
    1: (-inf*, +inf*) "y"
    2: [20, 30) "z"
  Output index maps:
    out[0] = 0 + 1 * in[0]
    out[1] = 0 + 1 * in[1]
    out[2] = 0 + 1 * in[2]

As an extension, TensorStore allows the start, stop, and step slice terms to be vectors rather than scalars:

>>> transform[ts.d['x', 'z'][[5, 20]:[10, 30]]]
Rank 3 -> 3 index space transform:
  Input domain:
    0: [5, 10) "x"
    1: (-inf*, +inf*) "y"
    2: [20, 30) "z"
  Output index maps:
    out[0] = 0 + 1 * in[0]
    out[1] = 0 + 1 * in[1]
    out[2] = 0 + 1 * in[2]
>>> transform[ts.d['x', 'z'][[5, 20]:30]]
Rank 3 -> 3 index space transform:
  Input domain:
    0: [5, 30) "x"
    1: (-inf*, +inf*) "y"
    2: [20, 30) "z"
  Output index maps:
    out[0] = 0 + 1 * in[0]
    out[1] = 0 + 1 * in[1]
    out[2] = 0 + 1 * in[2]

As with integer indexing, a single scalar slice applies to all selected dimensions:

>>> transform[ts.d['x', 'z'][5:30]]
Rank 3 -> 3 index space transform:
  Input domain:
    0: [5, 30) "x"
    1: (-inf*, +inf*) "y"
    2: [5, 30) "z"
  Output index maps:
    out[0] = 0 + 1 * in[0]
    out[1] = 0 + 1 * in[1]
    out[2] = 0 + 1 * in[2]

Adding singleton dimensions¶

Specifying a value of newaxis (equal to None) adds a new dummy/singleton dimension with implicit bounds \([0, 1)\):

>>> transform = ts.IndexTransform(input_labels=['x', 'y'])
>>> transform[ts.d[1][ts.newaxis]]
Rank 3 -> 2 index space transform:
  Input domain:
    0: (-inf*, +inf*) "x"
    1: [0*, 1*)
    2: (-inf*, +inf*) "y"
  Output index maps:
    out[0] = 0 + 1 * in[0]
    out[1] = 0 + 1 * in[2]
>>> transform[ts.d[0, -1][ts.newaxis, ts.newaxis]]
Rank 4 -> 2 index space transform:
  Input domain:
    0: [0*, 1*)
    1: (-inf*, +inf*) "x"
    2: (-inf*, +inf*) "y"
    3: [0*, 1*)
  Output index maps:
    out[0] = 0 + 1 * in[1]
    out[1] = 0 + 1 * in[2]

As with integer indexing, if only a single ts.newaxis term is specified, it applies to all selected dimensions:

>>> transform[ts.d[0, -1][ts.newaxis]]
Rank 4 -> 2 index space transform:
  Input domain:
    0: [0*, 1*)
    1: (-inf*, +inf*) "x"
    2: (-inf*, +inf*) "y"
    3: [0*, 1*)
  Output index maps:
    out[0] = 0 + 1 * in[1]
    out[1] = 0 + 1 * in[2]

newaxis terms are only permitted in the first operation of a dimension expression, since in subsequent operations all dimensions of the dimension selection necessarily refer to existing dimensions:

Error

>>> transform[ts.d[0, 1].translate_by[5][ts.newaxis]]
Traceback (most recent call last):
    ...
IndexError: tensorstore_demo.newaxis (`None`) not valid in chained indexing operations

It is also an error to use newaxis with dimensions specified by label:

Error

>>> transform[ts.d['x'][ts.newaxis]]
Traceback (most recent call last):
    ...
IndexError: New dimensions cannot be specified by label

Ellipsis¶

Specifying the special Ellipsis value (...) is equivalent to specifying as many full slices : as needed to consume the remaining selected dimensions not consumed by other indexing terms:

>>> transform = ts.IndexTransform(input_rank=4)
>>> transform[ts.d[:][1, ..., 5].translate_by[3]]
Rank 2 -> 4 index space transform:
  Input domain:
    0: (-inf*, +inf*)
    1: (-inf*, +inf*)
  Output index maps:
    out[0] = 1
    out[1] = -3 + 1 * in[0]
    out[2] = -3 + 1 * in[1]
    out[3] = 5

An indexing spec consisting solely of an Ellipsis term has no effect:

>>> transform[ts.d[:][...]]
Rank 4 -> 4 index space transform:
  Input domain:
    0: (-inf*, +inf*)
    1: (-inf*, +inf*)
    2: (-inf*, +inf*)
    3: (-inf*, +inf*)
  Output index maps:
    out[0] = 0 + 1 * in[0]
    out[1] = 0 + 1 * in[1]
    out[2] = 0 + 1 * in[2]
    out[3] = 0 + 1 * in[3]

Integer array indexing¶

Specifying an array_like index array of integer values selects the coordinates given by the elements of the array of the selected dimension:

>>> x = ts.array([[1, 2, 3], [4, 5, 6]], dtype=ts.int32)
>>> x = x[ts.d[:].label['x', 'y']]
>>> x[ts.d['y'][[1, 1, 0]]]
TensorStore({
  'array': [[2, 2, 1], [5, 5, 4]],
  'context': {'data_copy_concurrency': {}},
  'driver': 'array',
  'dtype': 'int32',
  'transform': {
    'input_exclusive_max': [2, 3],
    'input_inclusive_min': [0, 0],
    'input_labels': ['x', ''],
  },
})

As in the example above, if only a single index array term is specified, the dimensions of the index array are added to the result domain in place of the selected dimension, consistent with direct NumPy-style indexing in the default index array mode.

However, when using NumPy-style indexing with a dimension expression, if more than one index array term is specified, the broadcast dimensions of the index arrays are always added to the beginning of the result domain, i.e. exactly the behavior of DimExpression.vindex. Unlike with direct NumPy-style indexing (not with a dimension expression), the behavior does not depend on whether the index array terms apply to consecutive dimensions, since consecutive dimensions are not well-defined for dimension expressions:

>>> x = ts.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], dtype=ts.int32)
>>> x = x[ts.d[:].label['x', 'y', 'z']]
>>> x[ts.d['z', 'y'][[1, 0], [1, 1]]]
TensorStore({
  'array': [[4, 3], [8, 7]],
  'context': {'data_copy_concurrency': {}},
  'driver': 'array',
  'dtype': 'int32',
  'transform': {
    'input_exclusive_max': [2, 2],
    'input_inclusive_min': [0, 0],
    'input_labels': ['x', ''],
  },
})

Boolean array indexing¶

Specifying an array_like of bool values is equivalent to specifying a sequence of integer index arrays containing the coordinates of True values (in C order), e.g. as obtained from numpy.nonzero:

Specifying a 1-d bool array is equivalent to a single index array of the non-zero coordinates:

>>> x = ts.array([[1, 2, 3], [4, 5, 6]], dtype=ts.int32)
>>> x = x[ts.d[:].label['x', 'y']]
>>> x[ts.d['y'][[False, True, True]]]
TensorStore({
  'array': [[2, 3], [5, 6]],
  'context': {'data_copy_concurrency': {}},
  'driver': 'array',
  'dtype': 'int32',
  'transform': {
    'input_exclusive_max': [2, 2],
    'input_inclusive_min': [0, 0],
    'input_labels': ['x', ''],
  },
})

Equivalently, using an index array:

>>> x[ts.d['y'][[1, 2]]]
TensorStore({
  'array': [[2, 3], [5, 6]],
  'context': {'data_copy_concurrency': {}},
  'driver': 'array',
  'dtype': 'int32',
  'transform': {
    'input_exclusive_max': [2, 2],
    'input_inclusive_min': [0, 0],
    'input_labels': ['x', ''],
  },
})

More generally, specifying an n-dimensional bool array is equivalent to specifying n 1-dimensional index arrays, where the ith index array specifies the ith coordinate of the True values:

>>> x = ts.array([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]],
...              dtype=ts.int32)
>>> x = x[ts.d[:].label['x', 'y', 'z']]
>>> x[ts.d['x', 'z'][[[True, False, False], [True, True, False]]]]
TensorStore({
  'array': [[1, 4], [7, 10], [8, 11]],
  'context': {'data_copy_concurrency': {}},
  'driver': 'array',
  'dtype': 'int32',
  'transform': {
    'input_exclusive_max': [3, 2],
    'input_inclusive_min': [0, 0],
    'input_labels': ['', 'y'],
  },
})

Equivalently, using an index array:

>>> x[ts.d['x', 'z'][[0, 1, 1], [0, 0, 1]]]
TensorStore({
  'array': [[1, 4], [7, 10], [8, 11]],
  'context': {'data_copy_concurrency': {}},
  'driver': 'array',
  'dtype': 'int32',
  'transform': {
    'input_exclusive_max': [3, 2],
    'input_inclusive_min': [0, 0],
    'input_labels': ['', 'y'],
  },
})

Note that as with integer array indexing, when using NumPy-styling indexing with a dimension expression, if boolean arrays are applied to more than one selected dimension, the added dimension corresponding to the True values is always added to the beginning of the result domain, i.e. exactly the behavior of DimExpression.vindex.

Last update: Nov 16, 2024