New: truncate color map

setup.py

@@ -1,6 +1,6 @@
 from setuptools import setup, find_packages
 
-VERSION = '0.42.8'
+VERSION = '0.42.9'
 DESCRIPTION = 'Machine Learning project startup utilities'
 LONG_DESCRIPTION = 'My commonly used utilities for machine learning projects'
 
@@ -13,7 +13,7 @@
     description=DESCRIPTION,
     long_description=LONG_DESCRIPTION,
     url='https://github.com/StefanHeng/stef-util',
-    download_url='https://github.com/StefanHeng/stef-util/archive/refs/tags/v0.42.8.tar.gz',
+    download_url='https://github.com/StefanHeng/stef-util/archive/refs/tags/v0.42.9Y.tar.gz',
     packages=find_packages(),
     include_package_data=True,
     install_requires=[

stefutil/plot.py

@@ -28,11 +28,12 @@
     import seaborn as sns
     from matplotlib.patches import Ellipse
     from matplotlib.colors import to_rgba
+    import matplotlib.colors as colors
     import matplotlib.transforms as transforms
 
     __all__ += [
         'set_plot_style', 'LN_KWARGS',
-        'change_bar_width', 'vals2colors', 'set_color_bar', 'barplot'
+        'change_bar_width', 'vals2colors', 'truncate_colormap', 'set_color_bar', 'barplot'
     ]
 
 
@@ -48,10 +49,8 @@ def set_plot_style():
         sns.set_style('darkgrid')
         sns.set_context(rc={'grid.linewidth': 0.5})
 
-
     LN_KWARGS = dict(marker='o', ms=0.3, lw=0.25)  # matplotlib line plot default args
 
-
     def change_bar_width(ax, width: float = 0.5, orient: str = 'v'):
         """
         Modifies the bar width of a matplotlib bar plot
@@ -66,7 +65,6 @@ def change_bar_width(ax, width: float = 0.5, orient: str = 'v'):
             patch.set_width(width) if is_vert else patch.set_height(width)
             patch.set_x(patch.get_x() + diff * 0.5) if is_vert else patch.set_y(patch.get_y() + diff * 0.5)
 
-
     def vals2colors(
             vals: Iterable[float], color_palette: str = 'Spectral_r', gap: float = None,
     ) -> List:
@@ -88,6 +86,13 @@ def vals2colors(
         norm = (vals - mi) / (ma - mi)
         return cmap(norm)
 
+    def truncate_colormap(cmap, vmin=0.0, vmax=1.0, n=100):
+        if isinstance(cmap, str):
+            cmap = plt.get_cmap(cmap)
+        cmap: colors.Colormap
+        new_cmap = colors.LinearSegmentedColormap.from_list(
+            name='trunc({n},{a:.2f},{b:.2f})'.format(n=cmap.name, a=vmin, b=vmax), colors=cmap(np.linspace(vmin, vmax, n)))
+        return new_cmap
 
     def set_color_bar(vals, ax, color_palette: str = 'Spectral_r', orientation: str = 'vertical'):
         """
@@ -102,7 +107,6 @@ def set_color_bar(vals, ax, color_palette: str = 'Spectral_r', orientation: str
         plt.colorbar(sm, cax=ax, orientation=orientation)
         # plt.xlabel('colorbar')  # doesn't seem to work
 
-
     def barplot(
             data: pd.DataFrame = None,
             x: Union[Iterable, str] = None, y: Union[Iterable[float], str] = None,
@@ -147,7 +151,6 @@ def barplot(
             plt.show()
         return ax
 
-
     def confidence_ellipse(ax_, x, y, n_std=1., **kws):
         """
         Modified from https://matplotlib.org/stable/gallery/statistics/confidence_ellipse.html