Merge pull request #125 from kafitzgerald/colormap

fix and update truncate_colormap
NCAR · May 13, 2023 · 5a0aa0f · 5a0aa0f
2 parents 4c0642c + f666791
commit 5a0aa0f
Showing 1 changed file with 18 additions and 10 deletions.
diff --git a/src/geocat/viz/util.py b/src/geocat/viz/util.py
@@ -14,7 +14,6 @@
 import matplotlib.path as mpath
 import matplotlib.pyplot as plt
 import matplotlib.ticker as tic
-import matplotlib.cm as cm
 
 import cartopy.crs as ccrs
 import cartopy.util as cutil
@@ -639,10 +638,11 @@ def set_axes_limits_and_ticks(
 def truncate_colormap(cmap: matplotlib.colors.Colormap,
                       minval: typing.Union[int, float] = 0.0,
                       maxval: typing.Union[int, float] = 1.0,
-                      n: int = 100,
-                      name: str = None):
+                      n: int = 256,
+                      name: str = None,
+                      force: bool = True):
     """Utility function that truncates a colormap. Registers the new colormap
-    by name in plt.cm, and also returns the updated map.
+    by name and returns the corresponding colormap object.
 
     `Copied from Stack Overflow <https://stackoverflow.com/questions/18926031/how-to-extract-a-subset-of-a-colormap-as-a-new-colormap-in-matplotlib>`_
 
@@ -652,10 +652,10 @@ def truncate_colormap(cmap: matplotlib.colors.Colormap,
         Colormap to be truncated.
 
     minval : int, float
-        Minimum value to be used for truncation of the color map.
+        Minimum normalized value to be used for truncation of the color map.
 
     maxval : int, float
-        Maximum value to be used for truncation of the color map.
+        Maximum normalized value to be used for truncation of the color map.
 
     n : int
         Number of color values in the new color map.
@@ -664,6 +664,11 @@ def truncate_colormap(cmap: matplotlib.colors.Colormap,
         Optional name of the new color map. If not set, a new name is generated by using the name of the input
         colormap as well as min and max values.
 
+    force : bool, default = True
+        If False, a ValueError is raised if trying to overwrite an already registered name. True supports
+        overwriting registered colormaps other than the builtin colormaps.
+
+
     Examples
     --------
     All usage examples are within the GeoCAT-Examples Gallery. To see more usage cases, search the function on the
@@ -681,10 +686,13 @@ def truncate_colormap(cmap: matplotlib.colors.Colormap,
                                                    a=minval,
                                                    b=maxval)
     new_cmap = mpl.colors.LinearSegmentedColormap.from_list(
-        name=name,
-        colors=cmap(np.linspace(minval, maxval, n)),
-    )
-    cm.register_cmap(name, new_cmap)
+        name=name, colors=cmap(np.linspace(minval, maxval)), N=n)
+
+    try:
+        mpl.colormaps.register(new_cmap, force=force)
+    except AttributeError:
+        mpl.cm.register_cmap(name=name, cmap=new_cmap)
+
     return new_cmap