0

I have a 3-dimensional array A, and I need to access A[:,1,1] through the tuple x=[1,1]. Something like this:

x = [1,1]
A[:,*x]

However, doing that I get a syntax error. I would love to be able to access the elements of A[:,1,1] using the variable x, how can I do that?

Thank you!


Second question:

How to do the same but instead of slicing : do it with a boolean array. For example if t is an array of booleans, obtain A[t, *x]

Liubove
  • 57
  • 6

2 Answers2

1

You can do the following:

import numpy as np

A = np.arange(12).reshape((2, 3, 2))
print(A)

x = [1, 1]
print(A[(slice(None), *x)])

You can use slice(None) instead of : to build a tuple of slices. The tuple environment allows for value unpacking with the * operator.

Output:

[[[ 0  1]
  [ 2  3]
  [ 4  5]]

 [[ 6  7]
  [ 8  9]
  [10 11]]]

[3 9]

To verify it matches:

import numpy as np

A = np.arange(12).reshape((2, 3, 2))
x = [1, 1]
s = (slice(None), *x)
print(np.allclose(A[s], A[:, 1, 1]))  # True

*This is a modification of answers found here: Slicing a numpy array along a dynamically specified axis


Edit to reflect edit on question and comment:

To clarify, you can unpack any iterable you like in the tuple environment. The * operator functions normally in within the tuple. Order your elements however you like. Mix in different iterables, types, slice(None), how ever you want to build your slices, as long as you end up with a valid sequence of values, it will behave as expected.

import numpy as np

A = np.arange(12).reshape((2, 3, 2))
t = [True, False]
x = [1, 1]
print(np.allclose(A[(*t, *x)], A[True, False, 1, 1]))  # True

You can also add full lists as well in the tuple:

print(np.allclose(A[(t, *x)], A[[True, False], 1, 1]))  # True
Henry Ecker
  • 34,399
  • 18
  • 41
  • 57
  • Thank you! Follow up question: How would you do if instead of a slice `:` I have an array of booleans as first argument? For example if I have `t = [False, True]` and I want `A[t, *x]` – Liubove Apr 21 '21 at 03:23
  • Answer is updated to include this as well. – Henry Ecker Apr 21 '21 at 03:46
0

you can use slice(None) instead of :, So

y = tuple([slice[None]] + x)
A[y]

is what you need.

Terran
  • 649
  • 10
  • 24
  • because I want it more general. I am writing a flexible program that might well have `A` with more than three dimensions, and so `x` would have more dimensions too – Liubove Apr 21 '21 at 03:09