Different color for each populations in cmap

Question

Say I have four arrays and also an x and y array

arr1
arr2
arr3
arr4
xarr
yarr

Each of these 4 arrays has the same length as my x and y array and they each have a value of how "strong" that population is at that particular point in x and y. My plan is to plot this as a scatter plot, coloring each point in the scatter plot by which population has the highest value (so say I'd color it red if arr1 was strongest at that particular point, blue if arr2 was strongest, green if arr3 was strongest, and black if arr4 was strongest.)

I already can make this plot (attached below). Don't worry about labels or subplots here, I'm just attaching this so you can have somewhat of a visual. My problem is that I also want to take into account the value of strength. I'm not sure how to explain this quickly so bear with me. In the plot I attached below my colorbar is just 4 solid colors, each for a separate population. I still want the colorbar to have 4 distinct colors, but I want each of those four colors to have gradients, so that I could still have a sense of which regions the population is stronger and weaker - even if its still the strongest out of all four. Matplotlib has already a colormap that is like this (tab20b or tab20c) but I am just not sure how to implement this in how I'm doing things.

Basically how I make this plot is a loop over x and y where I compare each population and then assign my color array that I feed into the scatter plot (color) with either a 0, 1, 2, or 3 depending on who is strongest:

       for i in range(0, len(xarr)):
          for j in range(0, len(yarr)):
              if arr1[i,j] > max(arr2[i,j], arr3[i,j], arr4[i,j]):
                  color = np.append(color, 0)
              elif arr2[i,j] > max(arr1[i,j], arr3[i,j], arr4[i,j]):
                  color = np.append(color, 1)
              elif arr3[i,j] > max(arr1[i,j], arr2[i,j], arr4[i,j]):
                  color = np.append(color, 2)
              else:
                  color = np.append(color, 3)

I want now instead to find a way to store the value of arr[i,j] in my color while also still being able to have each different population colored differently. Any ideas on how to do this? I have never created my own colorbar but I know it's a thing you can do.

Answer 1

If you put all your arrays together as in arr = np.array([arr1, arr2, arr3, arr4]), you can calculate the maximums as maxs = arr.max(axis=0) and create filters to plot each part of the dots. Each of the arrays then can be plotted with a different colormap.

Each colormap can be visualized in a separate colorbar, which can be positioned via subplots.

import matplotlib.pyplot as plt
import numpy as np

M, N = 50, 30
arr = [np.random.randn(M, N).cumsum(axis=1).cumsum(axis=0).ravel() for _ in range(4)]
arr = np.array([arr_i - arr_i.mean() for arr_i in arr])
maxs = arr.max(axis=0)
x = np.tile(np.arange(M), N)
y = np.repeat(np.arange(N), M)

cmaps = ['Reds', 'Greens', 'Blues', 'Greys']
fig, axes = plt.subplots(ncols=5, figsize=(12, 4), gridspec_kw={'width_ratios': [30, 1, 1, 1, 1], 'wspace': 0.4})
for arr_i, cmap, cbar_ax in zip(arr, cmaps, axes[1:]):
    filter = arr_i == maxs
    scat = axes[0].scatter(x[filter], y[filter], c=arr_i[filter], cmap=cmap)
    plt.colorbar(scat, cax=cbar_ax)
plt.show()