matplotlib - 使用 pcolormesh 指定 colors 與 rgba 數組

Question

我在(n_y,n_x)網格上定義了數據，我已在(n_y,n_x,4) np.ndarray 。 我想使用pcolormesh 。

我嘗試將facecolors參數傳遞給pcolormesh ，它沒有做任何事情，並使用ListedColormap到 map 每個(y,x)單元格到一個顏色，這也不起作用。

下面的代碼重現了我遇到的問題。

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

'''
Define some arbitrary data
'''

dx = 0.01

x_range = [5,8]
y_range = [0,2]

x = np.arange(x_range[0],x_range[1],dx)
y = np.arange(y_range[0],y_range[1],dx)
X,Y = np.meshgrid(x,y)

data = X+Y**2

'''
Define colors based on the data
'''

def data_to_colors(data):
    colors = np.moveaxis(np.zeros_like([data]*4),0,-1) # shape (n_y,n_x,4)
    
    # make the data correspond to the blue channel, scaled between 0 and 1
    colors[...,2] = (data - data.min()) / (data.max()-data.min())

    # make red and green colors 0.5
    colors[...,0] = 0.5
    colors[...,1] = 0.5
    
    # make the alpha values all 1
    colors[...,-1] = 1
    
    return colors

'''
Show with imshow and pcolormesh
'''

fig,axs = plt.subplots(1,3,sharex=True,sharey=True,figsize=(12,4))

# show with imshow
extent = [x_range[0]-dx/2, x_range[-1]+dx/2, y_range[0]-dx/2, y_range[-1]+dx/2]
axs[0].imshow(data_to_colors(data),extent=extent,origin='lower')
axs[0].set_title('imshow (correct)')

# show with pcolormesh and facecolors
axs[1].pcolormesh(X,Y,np.ones_like(X),facecolors=data_to_colors(data.flatten()))
axs[1].set_title('pcolormesh, specifying facecolors')

# show using a ListedColorMap mapping each individual (row,column) to a color
ixs = np.arange(len(x)*len(y))
colors = data_to_colors(data.flatten())
axs[2].pcolormesh(X,Y,ixs.reshape(len(y),len(x)),cmap=ListedColormap(colors))
axs[2].set_title('pcolormesh, using a ListedColormap')

for ax in axs:
    ax.set_aspect('equal')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
fig.tight_layout()

有沒有一種方法可以讓我使用pcolormesh從imshow獲得相同的結果（或者當數據的行/列不一定對應於 y/x 的常量值時可以使用的任何方法）？

Answer 1

對於ListedColormap ，您需要（最多）256 個不同的 colors 的列表。 您可以通過向data_to_colors() function 提供 256 個值的列表來創建它們。 創建這些輸入值的最簡單方法是使用np.linspace(0, 1, 256) 。

對於facecolors方法，似乎 matplotlib 需要調用fig.canvas.draw()來創建 facecolors 數組。 此后，可以通過.set_facecolors設置它們。 同樣重要的是pcolormesh在網格頂點之間繪制面（與為每個頂點提供顏色的imshow相反）。 因此，網格中的頂點數比網格中的頂點數少一行一列。 您需要向網格中添加額外的行和額外的列，或者從 facecolors 中省略一行和一列。

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

dx = 0.01
x_range = [5, 8]
y_range = [0, 2]
x = np.arange(x_range[0], x_range[1], dx)
y = np.arange(y_range[0], y_range[1], dx)
X, Y = np.meshgrid(x, y)
data = X + Y ** 2

# Define colors based on the data
def data_to_colors(data):
    colors = np.moveaxis(np.zeros_like([data] * 4), 0, -1)  # shape (n_y,n_x,4)
    # make the data correspond to the blue channel, scaled between 0 and 1
    colors[..., 2] = (data - data.min()) / (data.max() - data.min())
    # make red and green colors 0.5
    colors[..., 0] = 0.5
    colors[..., 1] = 0.5
    # make the alpha values all 1
    colors[..., -1] = 1
    return colors

fig, axs = plt.subplots(1, 3, sharex=True, sharey=True, figsize=(12, 4))

# show with imshow
extent = [x_range[0] - dx / 2, x_range[-1] + dx / 2, y_range[0] - dx / 2, y_range[-1] + dx / 2]
axs[0].imshow(data_to_colors(data), extent=extent, origin='lower')
axs[0].set_title('imshow (correct)')

# show by updating the facecolors with set_facecolors
pcmesh = axs[1].pcolormesh(X, Y, data)
fig.canvas.draw()
pcmesh.set_facecolors(data_to_colors(data[:-1, :-1]).reshape(-1, 4))
axs[1].set_title('pcolormesh, using facecolors')

# show using a ListedColorMap mapping each individual (row,column) to a color
colors = data_to_colors(np.linspace(0, 1, 256))
axs[2].pcolormesh(X, Y, data, cmap=ListedColormap(colors))
axs[2].set_title('pcolormesh, using a ListedColormap')

for ax in axs:
    ax.set_aspect('equal')
fig.tight_layout()
plt.show()