Cleanup internals and fix 2-level bug (fixes #407)

lukelbd · lukelbd · commit d283935fe43e · 2023-03-11T15:41:43.000-07:00
diff --git a/proplot/colors.py b/proplot/colors.py
@@ -339,7 +339,7 @@
 docstring._snippet_manager['colors.from_list'] = _from_list_docstring
 
 
-def _clip_colors(colors, clip=True, gray=0.2, warn=False):
+def _clip_colors(colors, clip=True, gray=0.2, warn_invalid=False):
     """
     Clip impossible colors rendered in an HSL-to-RGB colorspace
     conversion. Used by `PerceptualColormap`.
@@ -349,13 +349,11 @@ def _clip_colors(colors, clip=True, gray=0.2, warn=False):
     colors : sequence of 3-tuple
         The RGB colors.
     clip : bool, optional
-        If `clip` is ``True`` (the default), RGB channel values >1 are
-        clipped to 1. Otherwise, the color is masked out as gray.
+        If `clip` is ``True`` (the default), RGB channel values >1
+        are clipped to 1. Otherwise, the color is masked out as gray.
     gray : float, optional
-        The identical RGB channel values (gray color) to be used if
-        `clip` is ``True``.
-    warn : bool, optional
-        Whether to issue warning when colors are clipped.
+        The identical RGB channel values (gray color) to be used
+        if `clip` is ``True``.
     """
     colors = np.asarray(colors)
     under = colors < 0
@@ -364,15 +362,15 @@ def _clip_colors(colors, clip=True, gray=0.2, warn=False):
         colors[under], colors[over] = 0, 1
     else:
         colors[under | over] = gray
-    if warn:
+    if warn_invalid:
         msg = 'Clipped' if clip else 'Invalid'
         for i, name in enumerate('rgb'):
             if np.any(under[:, i]) or np.any(over[:, i]):
                 warnings._warn_proplot(f'{msg} {name!r} channel.')
     return colors
 
 
-def _get_channel(color, channel, space='hcl'):
+def _color_channel(color, channel, space='hcl'):
     """
     Get the hue, saturation, or luminance channel value from the input color. The
     color name `color` can optionally be a string with the format ``'color+x'``
@@ -586,7 +584,7 @@ def _load_colors(path, warn_on_failure=True):
         If ``True``, issue a warning when loading fails instead of raising an error.
     """
     # Warn or raise error (matches Colormap._from_file behavior)
-    if not os.path.exists(path):
+    if not os.path.isfile(path):
         message = f'Failed to load color data file {path!r}. File not found.'
         if warn_on_failure:
             warnings._warn_proplot(message)
@@ -782,7 +780,7 @@ def _warn_or_raise(descrip, error=RuntimeError):
                 warnings._warn_proplot(prefix + ' ' + descrip)
             else:
                 raise error(prefix + ' ' + descrip)
-        if not os.path.exists(path):
+        if not os.path.isfile(path):
             return _warn_or_raise('File not found.', FileNotFoundError)
 
         # Directly read segmentdata json file
@@ -1944,7 +1942,7 @@ def __init__(
             for i, xyy in enumerate(array):
                 xyy = list(xyy)  # make copy!
                 for j, y in enumerate(xyy[1:]):  # modify the y values
-                    xyy[j + 1] = _get_channel(y, key, space)
+                    xyy[j + 1] = _color_channel(y, key, space)
                 data[key][i] = xyy
         # Initialize
         super().__init__(name, data, gamma=1.0, N=N, **kwargs)
@@ -2425,10 +2423,10 @@ def __init__(
         # Instead user-reversed levels will always get passed here just as
         # they are passed to SegmentedNorm inside plot.py
         levels, descending = _sanitize_levels(levels)
+        vcenter = getattr(norm, 'vcenter', None)
         vmin = norm.vmin = np.min(levels)
         vmax = norm.vmax = np.max(levels)
         bins, _ = _sanitize_levels(norm(levels))
-        vcenter = getattr(norm, 'vcenter', None)
         mids = np.zeros((levels.size + 1,))
         mids[1:-1] = 0.5 * (levels[1:] + levels[:-1])
         mids[0], mids[-1] = mids[1], mids[-2]
@@ -2442,12 +2440,20 @@ def __init__(
         # minimum 0 maximum 1, would mess up color distribution. However this is still
         # not perfect... get asymmetric color intensity either side of central point.
         # So we add special handling for diverging norms below to improve symmetry.
-        if unique in ('min', 'both'):
-            mids[0] += step * (mids[1] - mids[2])
-        if unique in ('max', 'both'):
-            mids[-1] += step * (mids[-2] - mids[-3])
-        mmin, mmax = np.min(mids), np.max(mids)
-        if vcenter is None:
+        if len(levels) == 2:
+            step = 0.5  # dummy step
+            mids[0] += step * (levels[0] - levels[1])
+            mids[-1] += step * (levels[-1] - levels[-2])
+        else:
+            if unique in ('min', 'both'):
+                offset = mids[1] - mids[2]
+                mids[0] += step * offset
+            if unique in ('max', 'both'):
+                offset = mids[-2] - mids[-3]
+                mids[-1] += step * offset
+        mmin = np.min(mids)
+        mmax = np.max(mids)
+        if vcenter is None:  # not diverging norm or centered segmented norm
             mids = _interpolate_scalar(mids, mmin, mmax, vmin, vmax)
         else:
             mask1, mask2 = mids < vcenter, mids >= vcenter
