CODEWITHSUNDEEP
pythonnumpyiteratorgenerator
user2162855
user2162855

Reputation: 13

python generator returning ndarray

I'm having the following odd behavior with a generator that returns a numpy array. In the real application, I am using scipy.sparse to implement an iterative linear algebra solver: The below is a stripped down version. I have the following:

import numpy as np

def iterable(n):

  nc = np.array([0,0,0])
  yield nc

  while nc[0] < n:
    nc += 1
    yield nc

  raise StopIteration

When i try and populate a list with the output

iter = iterable(4)

print list(iter)

I get

$ python itertest.py 
[array([4, 4, 4]), array([4, 4, 4]), array([4, 4, 4]), array([4, 4, 4]), array([4, 4, 4])]

i.e., the result of the final 'yield' repeated n times. If I change the yield nc to yield nc.tolist() then I get the expected result. If possible I would like to avoid that as I will end up having to re-make the array. What's going on?

Thanks in advance!

Upvotes: 1

Views: 133

Answers (1)

BrenBarn
BrenBarn

Reputation: 251538

You only create one ndarray. When you do nc += 1, you modify it in place. So the later iterations modify the same array already yielded, and then yield it again.

You can see this by printing out the generated arrays one at a time:

>>> it = iterable(3)
>>> x = next(it)
>>> print x
[0 0 0]
>>> y = next(it)
>>> print y
[1 1 1]
>>> print x   # x changed!
[1 1 1]
>>> x is y
True

You can see that the second iteration changes the object that was yielded on the first iteration, because they are the same object.

To fix it, change your nc += 1 line to nc = nc + 1. This will create a new array on each iteration.

Upvotes: 1

Related Questions