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.
CodePudding user response:
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