dhokas dhokas - 20 days ago 8
Python Question

efficiently fill a tensor in numpy

I have a numpy array

X
of size
(n, m)
and of type
np.uint8
(so it only contains values in
[0, 255]
). I also have a mapping
f
from
[0, 255]
to
[0, 3]
.

I want to create an array
Y
of shape
(4, n, m)
such that
y_{k, i, j} = 1 if k == f(x_{i, j})
and 0 otherwise. For now, I do it like this:

Y = np.zeros((4, n, m))
for i in range(256):
Y[f(i), X == i] = 1


But this is super slow, and I'm not able to find a more efficient way to do it. Any ideas?

Answer

Assuming f could operate on all iterating values in one-go, you can use broadcasting -

Yout = (f(X) == np.arange(4)[:,None,None]).astype(int)

Runtime test and verification -

In [35]: def original_app(X,n,m):
    ...:     Y = np.zeros((4, n, m))
    ...:     for i in range(256):
    ...:         Y[f(i), X == i] = 1
    ...:     return Y
    ...: 

In [36]: # Setup Inputs
    ...: n,m = 2000,2000
    ...: X = np.random.randint(0,255,(n,m)).astype('uint8')
    ...: v = np.random.randint(4, size=(256,))
    ...: def f(x): 
    ...:     return v[x]
    ...: 

In [37]: Y = original_app(X,n,m)
    ...: Yout = (f(X) == np.arange(4)[:,None,None]).astype(int)
    ...: 

In [38]: np.allclose(Yout,Y)  # Verify
Out[38]: True

In [39]: %timeit original_app(X,n,m)
1 loops, best of 3: 3.77 s per loop

In [40]: %timeit (f(X) == np.arange(4)[:,None,None]).astype(int)
10 loops, best of 3: 74.5 ms per loop
Comments