Internal cleanup and fix locator _vcenter bug

lukelbd · lukelbd · commit ab8de2d9f48c · 2023-03-11T19:34:23.000-07:00
diff --git a/proplot/axes/base.py b/proplot/axes/base.py
@@ -1083,7 +1083,8 @@ def _add_colorbar(
         # DiscreteLocator or else get issues (see mpl #22233).
         norm = mappable.norm
         source = getattr(norm, '_norm', None)
-        vcenter = {'vcenter': getattr(source, 'vcenter', 0.0)}
+        vcenter = getattr(source, 'vcenter', None)
+        vcenter = {} if vcenter is None else {'vcenter': vcenter}
         formatter = _not_none(formatter, getattr(norm, '_labels', None), 'auto')
         formatter = constructor.Formatter(formatter, **formatter_kw)
         categorical = isinstance(formatter, mticker.FixedFormatter)
diff --git a/proplot/axes/plot.py b/proplot/axes/plot.py
@@ -2500,7 +2500,7 @@ def _parse_level_num(
             Unused arguments.
         """
         # Input args
-        # NOTE: Some of this is adapted from the hidden contour.ContourSet._autolev
+        # NOTE: Some of this is adapted from contour.ContourSet._autolev
         # NOTE: We use 'symmetric' with MaxNLocator to ensure boundaries
         # include a zero level but may trim many of these levels below.
         norm_kw = norm_kw or {}
@@ -2551,7 +2551,7 @@ def _parse_level_num(
             vmin, vmax, levels = vmin + vcenter, vmax + vcenter, levels + vcenter
 
         # Possibly trim levels far outside of 'vmin' and 'vmax'
-        # NOTE: This part is mostly copied from contour.ContourSet._autolev
+        # NOTE: This part is mostly copied from matplotlib _autolev
         if not symmetric:
             i0, i1 = 0, len(levels)  # defaults
             under, = np.where(levels < vmin)
@@ -2832,7 +2832,11 @@ def _parse_level_norm(
 
         # Generate DiscreteNorm and update "child" norm with vmin and vmax from
         # levels. This lets the colorbar set tick locations properly!
-        if not isinstance(norm, mcolors.BoundaryNorm) and len(levels) > 1:
+        if len(levels) == 1:
+            pass  # e.g. contours
+        elif isinstance(norm, mcolors.BoundaryNorm):
+            pass  # override with native matplotlib normalizer
+        else:
             norm = pcolors.DiscreteNorm(
                 levels, norm=norm, unique=unique, step=step,
                 ticks=discrete_ticks, labels=discrete_labels,
diff --git a/proplot/colors.py b/proplot/colors.py
@@ -2440,19 +2440,17 @@ 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 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
+        if unique in ('min', 'both'):
+            scale = levels[0] - levels[1] if len(levels) == 2 else mids[1] - mids[2]
+            mids[0] += step * scale
+        if unique in ('max', 'both'):
+            scale = levels[-1] - levels[-2] if len(levels) == 2 else mids[-2] - mids[-3]
+            mids[-1] += step * scale
         mmin = np.min(mids)
         mmax = np.max(mids)
+        if np.isclose(mmin, mmax):
+            mmin = mmin - (mmin or 1) * 1e-10
+            mmax = mmax + (mmax or 1) * 1e-10
         if vcenter is None:  # not diverging norm or centered segmented norm
             mids = _interpolate_scalar(mids, mmin, mmax, vmin, vmax)
         else:
