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

Question

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.

score 1 · Accepted Answer · answered Nov 18 '21 at 11:45

1

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 Nov 18 '21 at 11:45

Jody Klymak

4,979
2
15
31

ah yes thank you. this works in this code. Unfortunately the real use has a more complex decision function that pulls from different databases and compares different values so the function needs to be invoked. my bad for the confusion. – Kurt Rojas Nov 18 '21 at 12:28
if the function _needs_ to be invoked https://stackoverflow.com/questions/35215161/most-efficient-way-to-map-function-over-numpy-array – Jody Klymak Nov 18 '21 at 13:08

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

1 Answers1