dhokas - 1 year ago 100
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?

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
``````
Recommended from our users: Dynamic Network Monitoring from WhatsUp Gold from IPSwitch. Free Download