From f443a612132b5b9605e139245dce54af29ed2e98 Mon Sep 17 00:00:00 2001
From: Stefanie Lumnitz <stefanie.lumnitz@gmail.com>
Date: Mon, 25 Jun 2018 15:31:37 -0700
Subject: [PATCH] add utility functions `shift_colormap` and
 `truncate_colormap`

---
 splot/_viz_utils.py | 94 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 94 insertions(+)

diff --git a/splot/_viz_utils.py b/splot/_viz_utils.py
index 993c8cdc..809e9a4a 100644
--- a/splot/_viz_utils.py
+++ b/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
\ No newline at end of file