2

I am working on a program which will create contour data out of numpy arrays, and trying to avoid calls to matplotlib.

I have an array of length L which contains NxN arrays of booleans. I want to convert this into an LxNxN array where, for example, the "True"s in the first inner array get replaced by "red", in the second, by "blue" and so forth.

The following code works as expected:

import numpy as np
import pdb

def new_layer(N,p):
    return np.random.choice(a=[False,True],size=(N,N),p=[p,1-p])

a = np.array([new_layer(3,0.5),new_layer(3,0.5),new_layer(3,0.5)]).astype('object')

colors = np.array(["red","green","blue"])

for i in range(np.shape(a)[0]):
    b = a[i]
    b[np.where(b==True)] = colors[i]
    a[i] = b
print(a)

But I am wondering if there is a way to accomplish the same using Numpy's built-in tools, e.g., indexing. I am a newcomer to Numpy and I suspect there is a better way to do this but I can't think what it would be. Thank you.

D. Lef
  • 239
  • 3
  • 13

3 Answers3

2

You could use np.copyto:

np.copyto(a, colors[:, None, None], where=a.astype(bool))
Paul Panzer
  • 51,835
  • 3
  • 54
  • 99
0

Here's one way -

a_bool = a.astype(bool)
a[a_bool] = np.repeat(colors,a_bool.sum((1,2)))

Another with extending colors to 3D -

a_bool = a.astype(bool)
colors3D = np.broadcast_to(colors[:,None,None],a.shape)
a[a_bool] = colors3D[a_bool]
Divakar
  • 218,885
  • 19
  • 262
  • 358
0

You can use a combination of boolean indexes and np.indices. Also you can use a as index to itself. Then you could do what you did in the for loop with this line (although I don't think it necessarily is a good idea):

a[a.astype(bool)] = colors[np.indices(a.shape)[0][a.astype(bool)]]

Also, for the new_layer function you could just use np.random.rand(N,N) > p (not sure if the actual distribution will be exactly the same as what you had).

fvizeus
  • 88
  • 6