Custom matplotlib colormap using an existing one

gistfile1.py

import copy

import matplotlib
from matplotlib import cm
import matplotlib.pyplot as plt
import numpy as np


def plot_colormap(cmapname='jet'):
    """
    Plot our colormap to see what it looks like
    """
    a = np.linspace(0, 1, 256).reshape(1, -1)
    a = np.vstack((a, a))
    fig = plt.figure(figsize=(5, 5))
    plt.axis('off')
    plt.imshow(a, aspect='auto', cmap=plt.get_cmap(cmapname))
    fig.text(0.5, 0.95, cmapname, fontsize=10, verticalalignment='center',
             horizontalalignment='center')


def find_location(colour_list, location):
    # Find where to add our new entry to the colour list
    return [n for n, i in enumerate([item[0] for item in
                                     colour_list]) if i > location][0]


def gen_custom_map(cmapname='jet', custom_map_name='jet_w_white'):
    cdict = copy.deepcopy(cm.datad[cmapname])
    print 'cm jet: {}'.format(cdict)

    # Turn them from tuples to lists so that they are mutable
    blue = [list(item) for item in cdict['blue']]
    green = [list(item) for item in cdict['green']]
    red = [list(item) for item in cdict['red']]

    # Define the values and their location to add
    # 0.5 corresponds to half way
    location = 0.5
    colour_spec = [location, 1, 1]

    # Insert new entry at a suitable location
    center = find_location(blue, location)
    blue.insert(center, colour_spec)

    center = find_location(green, location)
    green.insert(center, colour_spec)

    center = find_location(red, location)
    red.insert(center, colour_spec)

    # Turn them back into tuples to create a new colormap
    cdict['blue'] = tuple([tuple(item) for item in blue])
    cdict['green'] = tuple([tuple(item) for item in green])
    cdict['red'] = tuple([tuple(item) for item in red])

    print 'cm jet custom: {}'.format(cdict)

    # Register our new colormap
    custom_cm = matplotlib.colors.LinearSegmentedColormap(custom_map_name,
                                                          cdict)
    cm.register_cmap(name=custom_map_name, cmap=custom_cm)


if __name__ == '__main__':
    gen_custom_map(cmapname='jet', custom_map_name='jet_w_white')
    plot_colormap('jet')
    plot_colormap('jet_w_white')
    plt.show()