CODEWITHSUNDEEP
pythonfunctionnumpygeneratorlarge-data
Clément Moissard
Clément Moissard

Reputation: 161

How to create an array based on a generator?

Say I have a large array:

A = 2*np.ones([100, 100, 100])

I want to do some calculations on it, for example:

def squared_elements(M):
    yield M**2

I choose to use a generator function because my array is very big and I don't need all the results. As a matter of fact, I only need, say, a cube of length 10 at the center of the matrix.

If it was a normal function, I could just write:

result = squared_elements(A)[45:55, 45:55, 45:55]

However, generators are not subscriptable, so this last expression doesn't work.

How can I get the same result using my generator function?

Upvotes: 0

Views: 227

Answers (1)

Thomas
Thomas

Reputation: 899

Solution

result = next(squared_elements(A))
result[45:55, 45:55, 45:55]

Solution to the updated problem

You can slice the subset before squaring

def squared_elements(A, slice_):
    return A[slice_] ** 2

result = squared_elements(A, [slice(45, 55), slice(45, 55), slice(45, 55)])

Notes

  • more on next() here.
  • more on generator functions here.

I hope this helps!

Upvotes: 1

Related Questions