Lewis Morris
Lewis Morris

Reputation: 2124

Filling a 2d numpy array with object that takes index as a parameter

Imagine a class like this.

class Foo:
   def __init__(self,y,x):
      self.y = y
      self.x = x

Then I have an array

obj_arr = np.zeros((200,200), dtype=object)

I want to fill it with full of the class but with the input parameters as the index of the array.

I've toyed with the idea of np.where but cant see how I can add the parameters.

np.where(obj_arr == None, obj_arr, Foo()) 

Currently I'm doing this disgusting nested loop. 1) i know its not the right way 2) it takes forever.

for y in range(obj_arr.shape[0]):
    for x in range(obj_arr.shape[1]):
        obj_arr[y,x] == Foo(y,x)

Someone please push me in the right direction before I go crazy.

Upvotes: 0

Views: 68

Answers (2)

hpaulj
hpaulj

Reputation: 231395

class Foo:
    def __init__(self, y, x):
        self.y = y
        self.x = x
    def __repr__(self):
        return f'Foo({self.y}, {self.x})'

In [69]: Foo(1,2)
Out[69]: Foo(1, 2)

The proposed nditer solution:

def foo1(n):
    a = np.empty((n,n), dtype=object)
    with np.nditer(a, flags=['multi_index', 'refs_ok'], op_flags=['writeonly']) as it:
        for x in it:
            x[...] = Foo(it.multi_index[1], it.multi_index[0])
    return a

In [70]: foo1(2)
Out[70]: 
array([[Foo(0, 0), Foo(1, 0)],
       [Foo(0, 1), Foo(1, 1)]], dtype=object)

The nested loop:

def foo2(n):
    a = np.empty((n,n), dtype=object)
    for i in range(n):
        for j in range(n):
            a[i,j] = Foo(i,j)
    return a

In [71]: foo2(2)
Out[71]: 
array([[Foo(0, 0), Foo(0, 1)],
       [Foo(1, 0), Foo(1, 1)]], dtype=object)

my favorite, frompyfunc:

def foo3(n):
    f = np.frompyfunc(Foo, 2, 1)
    I,J = np.meshgrid(np.arange(n),np.arange(n), sparse=True)
    return f(I,J)

In [72]: foo3(2)
Out[72]: 
array([[Foo(0, 0), Foo(1, 0)],
       [Foo(0, 1), Foo(1, 1)]], dtype=object)

timings:

In [73]: timeit foo1(200)
144 ms ± 305 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [74]: timeit foo2(200)
25.7 ms ± 958 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [75]: timeit foo3(200)
17.7 ms ± 40.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

And for comparison, a nested list:

def foo4(n):
    alist = []
    for i in range(n):
        blist = []
        for j in range(n):
            blist.append(Foo(i,j))
        alist.append(blist)
    return alist

In [77]: foo4(2)
Out[77]: [[Foo(0, 0), Foo(0, 1)], [Foo(1, 0), Foo(1, 1)]]
In [78]: timeit foo4(200)
18.6 ms ± 149 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Upvotes: 2

sai
sai

Reputation: 1784

Came across this well documented Iterating Over Arrays site-

This should be able to solve your problem, I've tried it for 200x200 as well and it hardly takes any time.

Note: Make sure the numpy version is >=1.15

import numpy as np

class Foo:
    def __init__(self, y, x):
        self.y = y
        self.x = x


a = np.empty((2, 2), dtype=object)

with np.nditer(a, flags=['multi_index', 'refs_ok'], op_flags=['writeonly']) as it:
    for x in it:
        x[...] = Foo(it.multi_index[1], it.multi_index[0])

print(a)
[[<__main__.Foo object at 0x7f91e2613370>
  <__main__.Foo object at 0x7f91e25b2070>]
 [<__main__.Foo object at 0x7f91e11ef220>
  <__main__.Foo object at 0x7f91e11ef160>]]

Upvotes: 0

Related Questions