[FIXED] Python & numpy, better way of creating a Z from x and y meshgrid using a a decision function

Issue

I am having a problem with regards preparing an array for use in a pcolormesh plot. Currently I am implementing a decision function to a combination of x and y arrays to produce z, which is kind of a isosurface value (in a way). the code goes like,

def region(nx, ny):
    if nx >= 0 and ny >= 0: 
        r = 1
    elif nx < 0 and ny >= 0:
        r = 2
    elif nx < 0 and ny < 0:
        r = 3
    elif nx >= 0 and ny < 0:
        r = 4  
    return r

x = np.linspace(-10, 10, 41)   
y = np.linspace(-20, 20, 41) 
z = np.array([region(i, j) for j in y for i in x])
X, Y = np.meshgrid(x, y)
Z = z.reshape(41, 41)

discrete_colors = [sns.color_palette("Paired")[0],
                   sns.color_palette("Paired")[2],
                   sns.color_palette("Paired")[6],
                   sns.color_palette("Paired")[8]]
my_colormap = colors.ListedColormap(discrete_colors)

fig, ax = plt.subplots()
cmap = plt.get_cmap('RdBu', 4)
# c = ax.pcolormesh(X, Y, Z,cmap=cmap, vmin = np.min(Z)-.5, vmax = np.max(Z)+.5)
c = ax.pcolormesh(X, Y, Z,cmap=my_colormap, vmin = np.min(Z)-.5, vmax = np.max(Z)+.5)
cbar = fig.colorbar(c, ax=ax, ticks=np.arange(np.min(Z),np.max(Z)+1))
cbar.set_ticklabels(['Region 1','Region 2','Region 3','Region 4'])
ax.set_xlim(min(x)-0.5, max(x)+0.5)
ax.set_ylim(min(y)-0.5, max(y)+0.5)
plt.show()

Is there a better way of creating the z array from x and y ? Maybe perhaps editing the decision function? I am using a very large data set so the current implementation is rather slow.

Solution

You can just use boolean indexing:

Z = np.zeros_like(X)
Z[(X>=0) & (Y>=0)] = 1
Z[(X<0) & (Y>=0)] = 2
Z[(X<0) & (Y<0)] = 3
Z[(X>=0) & (Y<0)] = 4

Answered By - Jody Klymak