add utility functions shift_colormap and truncate_colormap

splot/_viz_utils.py

@@ -195,3 +195,97 @@ def calc_data_aspect(plot_height, plot_width, bounds):
     ymin -= (ymax - ymin) / 2 * margin
     ymax += (ymax - ymin) / 2 * margin
     return xmin, xmax, ymin, ymax
+
+
+#Utility functions for viz of continuous variables
+
+#Utility function #1 - forces diverging colormap to be centered at zero
+def shift_colormap(cmap, start=0, midpoint=0.5, stop=1.0, name='shiftedcmap'):
+    '''
+    Function to offset the "center" of a colormap. Useful for
+    data with a negative min and positive max and you want the
+    middle of the colormap's dynamic range to be at zero
+
+    Parameters
+    ----------
+    cmap : matplotlib colormap
+        colormap to be altered
+    start : float, optional
+        Offset from lowest point in the colormap's range.
+        Should be between 0.0 and `midpoint`.
+        Default =0.0 (no lower ofset).
+    midpoint : float, optional
+        The new center of the colormap.Should be between 0.0 and
+        1.0. In general, this should be 1 - vmax/(vmax + abs(vmin)).
+        For example if your data range from -15.0 to +5.0 and
+        you want the center of the colormap at 0.0, `midpoint`
+        should be set to  1 - 5/(5 + 15)) or 0.75.
+        Default =0.5 (no shift).
+    stop : float, optional
+        Offset from highets point in the colormap's range.
+        Should be between `midpoint` and 1.0.
+        Default =1.0 (no upper ofset).
+    name : str, optional
+        Name of the new colormap.
+    
+    Returns
+    -------
+    new_cmap : A new colormap that has been shifted. 
+    '''
+
+    cdict = {
+        'red': [],
+        'green': [],
+        'blue': [],
+        'alpha': []
+    }
+
+    # regular index to compute the colors
+    reg_index = np.linspace(start, stop, 257)
+
+    # shifted index to match the data
+    shift_index = np.hstack([
+        np.linspace(0.0, midpoint, 128, endpoint=False), 
+        np.linspace(midpoint, 1.0, 129, endpoint=True)
+    ])
+
+    for ri, si in zip(reg_index, shift_index):
+        r, g, b, a = cmap(ri)
+
+        cdict['red'].append((si, r, r))
+        cdict['green'].append((si, g, g))
+        cdict['blue'].append((si, b, b))
+        cdict['alpha'].append((si, a, a))
+
+    new_cmap = mpl.colors.LinearSegmentedColormap(name, cdict)
+    plt.register_cmap(cmap=new_cmap)
+    return new_cmap
+
+
+#Utility #2 - truncate colorcap in order to grab only positive or negative portion
+def truncate_colormap(cmap, minval=0.0, maxval=1.0, n=100):
+    '''
+    Function to truncate a colormap by selecting a subset of
+    the original colormap's values
+
+    Parameters
+    ----------
+    cmap : Mmatplotlib colormap
+        Colormap to be altered
+    minval : float, optional
+        Minimum value of the original colormap to include
+        in the truncated colormap. Default =0.0.
+    maxval : Maximum value of the original colormap to
+        include in the truncated colormap. Default =1.0.
+    n : int, optional
+        Number of intervals between the min and max values
+        for the gradient of the truncated colormap. Default =100.
+          
+    Returns
+    -------
+    new_cmap : A new colormap that has been shifted. 
+    '''
+    new_cmap = mpl.colors.LinearSegmentedColormap.from_list(
+        'trunc({n},{a:.2f},{b:.2f})'.format(n=cmap.name, a=minval, b=maxval),
+        cmap(np.linspace(minval, maxval, n)))
+    return new_cmap