CODEWITHSUNDEEP
pythonmatplotlibpython-2.x
Jzl5325
Jzl5325

Reputation: 3974

matplotlib color map - predefine mappings to values?

I have an array that I am viewing using imshow(). (imsave() really, but the process should be identical).

I know that the values in the array will be between 0-9 and wonder if it is possible to use cmap to set each output to a specific 'color'. Perhaps by mapping these to a dict?

Upvotes: 2

Views: 5757

Answers (2)

sglbl
sglbl

Reputation: 617

You can use something like this:

def get_product_number_or_colormap(product=None, get_colormap=False):
    '''
    Get the number of the product or the colormap \\
    If get_colormap is True, return the colormap
    '''
    # Define product families and their corresponding numbers and colors
    product_families = {
        'Other': 'w',  # 'white
        'Semilla de Girasol': 'tab:pink',
        'MAIZ': 'tab:orange',
        'Trigo panificable': 'tab:green',
        'Cebada': 'tab:purple',
        'Urea': 'tab:grey',
        'Yara': 'tab:brown',
        'ACTH': 'tab:green',
        'AXAN': 'tab:red',
        'NAC27': 'tab:blue',
        'SULFAN': 'tab:cyan',
    }
    
    product_families_with_nums = {prod: {'number': num, 'color': product_families[prod]} for num, prod in enumerate(product_families)}
    # print(f"{product_families_with_nums=}")
    
    if get_colormap:
        # convert all colors to a list
        color_list = [product_families[prod] for prod in product_families]
        return colors.ListedColormap(color_list)
    
    # Else: Get the number and color for the product
    try:
        product_number = product_families_with_nums[product]['number']
    except KeyError:
        product_number = 1
        print(f"Product: {product} did not match any of the predefined products")
    
    return product_number

After you call this function: Surface you want to paint:

surface[:, surface_size] = get_product_number_or_colormap(product)

Use function you want: imshow/imsave/pcolormesh (This code is extended so it also would work with pcolormesh)
You can get the cmap using:

cmap = get_product_number_or_colormap(get_colormap=True)     
axs.flatten()[i].pcolormesh(x, y, surface, cmap=cmap, alpha=0.7, vmin=0, vmax=7)

Upvotes: 0

Joe Kington
Joe Kington

Reputation: 284562

Just use a ListedColormap.

As a quick (but ugly) example:

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

cmap = ListedColormap(['red', 'green', 'blue', 'black'], 'indexed')

fig, ax = plt.subplots()
im = ax.imshow([range(4)], interpolation='none', cmap=cmap)
fig.colorbar(im)
plt.show()

enter image description here

Upvotes: 5

Related Questions