Mapping certain value to a color using contourf in matplotlib-CodePudding

I've been trying to solve this problem for days, but none of the solutions I found online seem to work very well for me, despite how simple the problem seems.

import numpy as np
from matplotlib import pyplot as plt

grid_x = np.linspace(-5, 5, 500)
grid_y = np.linspace(-5, 5, 500)
xx, yy = np.meshgrid(grid_x, grid_y)

fig, ax = plt.subplots(2)
rand1 = np.ones((500,500))
rand1[:, 250:] = 2
ax[0].contourf(xx, yy, rand1, alpha=.3, cmap = 'jet')

rand2 = np.zeros((500,500))
rand2[:, 200:250] = 1
rand2[:,250:] = 2
ax[1].contourf(xx, yy, rand2, alpha=.3, cmap = 'jet')

This is the resulting plot by running the above:

On the first axis, the 1s are mapped to blue, and 2s are mapped to red. On the second axis, the 0s are mapped to blue, 1s are mapped to green, and 2s are mapped to red.

I would like it to be consistent, so that for example, 1s are always mapped to blue, 2s always mapped to red, and 0s always mapped to green.

As far as I am aware, this issue would be fixed by ensuring that both plots have values in (0,1,2) all show up at some point, however I cannot ensure that for my project, so that is not a valid solution.

Another thought I had was to map each region separately. However, I am not sure how to do this, and though I imagine it would be relatively easy for the simple example above, I need it to generalise to any arbitrarily shaped subregions.

Does anyone have suggestions on how to fix this issue? Thank you very much!

CodePudding user response：

To better see what's happening, it helps to add a colorbar. It shows the levels used and the corresponding colors.

To have the same colors for corresponding levels, the levels should be the same between the plots:

import numpy as np
from matplotlib import pyplot as plt

grid_x = np.linspace(-5, 5, 500)
grid_y = np.linspace(-5, 5, 500)
xx, yy = np.meshgrid(grid_x, grid_y)

rand1 = np.ones((500, 500))
rand1[:, 250:] = 2

rand2 = np.zeros((500, 500))
rand2[:, 200:250] = 1
rand2[:, 250:] = 2

# common_levels = np.linspace(0, 2, 11)
common_levels = np.linspace(min(rand1.min(), rand2.min()), max(rand1.max(), rand2.max()), 11)

fig, axs = plt.subplots(ncols=2, nrows=2, figsize=(12, 6))
contour00 = axs[0, 0].contourf(xx, yy, rand1, alpha=.3, cmap='jet')
axs[0, 0].set_title('values 1 and 2, default levels')
plt.colorbar(contour00, ax=axs[0, 0])

contour01 = axs[0, 1].contourf(xx, yy, rand1, levels=common_levels, alpha=.3, cmap='jet')
plt.colorbar(contour01, ax=axs[0, 1])
axs[0, 1].set_title('values 1 and 2, common levels')

contour10 = axs[1, 0].contourf(xx, yy, rand2, alpha=.3, cmap='jet')
axs[1, 0].set_title('values 0, 1 and 2, default levels')
plt.colorbar(contour10, ax=axs[1, 0])

contour11 = axs[1, 1].contourf(xx, yy, rand2, levels=common_levels, alpha=.3, cmap='jet')
axs[1, 1].set_title('values 0, 1 and 2, common levels')
plt.colorbar(contour11, ax=axs[1, 1])

plt.tight_layout()
plt.show()