Fix contour level color mapping with explicit limits (#599)

* Fix contour level color mapping with explicit limits

* Update ultraplot/axes/plot.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Fix `_parse_level_norm` docstring to reflect conditional return type (#600)

* Initial plan

* Update _parse_level_norm docstring Returns section to reflect possible return types

Co-authored-by: cvanelteren <19485143+cvanelteren@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: Casper van Elteren <caspervanelteren@gmail.com>
Co-authored-by: cvanelteren <19485143+cvanelteren@users.noreply.github.com>

* Scope explicit contour limits to line contours

* Pass vmin/vmax through automatic level generation

* Restore default tricontour discrete mapping

* Format contour norm routing changes

* Clarify contour norm routing docs

* Refactor contour norm routing flow

* Refactor contour norm routing flags

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com>
Co-authored-by: cvanelteren <19485143+cvanelteren@users.noreply.github.com>

Author: 3 people

ultraplot/axes/plot.py

@@ -4141,6 +4141,7 @@ def _parse_cmap(
         # NOTE: Unlike xarray, but like matplotlib, vmin and vmax only approximately
         # determine level range. Levels are selected with Locator.tick_values().
         levels = None  # unused
+        explicit_limits = False
         isdiverging = False
         if not discrete and not skip_autolev:
             vmin, vmax, kwargs = self._parse_level_lim(
@@ -4150,7 +4151,15 @@ def _parse_cmap(
                 if abs(np.sign(vmax) - np.sign(vmin)) == 2:
                     isdiverging = True
         if discrete:
-            levels, vmin, vmax, norm, norm_kw, kwargs = self._parse_level_vals(
+            (
+                levels,
+                vmin,
+                vmax,
+                norm,
+                norm_kw,
+                explicit_limits,
+                kwargs,
+            ) = self._parse_level_vals(
                 *args,
                 vmin=vmin,
                 vmax=vmax,
@@ -4199,6 +4208,7 @@ def _parse_cmap(
                 center_levels=center_levels,
                 extend=extend,
                 min_levels=min_levels,
+                explicit_limits=explicit_limits,
                 **kwargs,
             )
         params = _pop_params(kwargs, *self._level_parsers, ignore_internal=True)
@@ -4462,7 +4472,6 @@ def _parse_level_num(
         center_levels = _not_none(center_levels, rc["colorbar.center_levels"])
         vmin = _not_none(vmin=vmin, norm_kw_vmin=norm_kw.pop("vmin", None))
         vmax = _not_none(vmax=vmax, norm_kw_vmax=norm_kw.pop("vmax", None))
-        norm = constructor.Norm(norm or "linear", **norm_kw)
         symmetric = _not_none(
             symmetric=symmetric,
             locator_kw_symmetric=locator_kw.pop("symmetric", None),
@@ -4555,6 +4564,8 @@ def _parse_level_vals(
         nozero=False,
         norm=None,
         norm_kw=None,
+        vmin=None,
+        vmax=None,
         skip_autolev=False,
         min_levels=None,
         center_levels=None,
@@ -4577,7 +4588,9 @@ def _parse_level_vals(
             Whether to remove out non-positive, non-negative, and zero-valued
             levels. The latter is useful for single-color contour plots.
         norm, norm_kw : optional
-            Passed to `Norm`. Used to possbily infer levels or to convert values.
+            Passed to `Norm`. Used to possibly infer levels or to convert values.
+        vmin, vmax : float, optional
+            The user input normalization range.
         skip_autolev : bool, optional
             Whether to skip automatic level generation.
         min_levels : int, optional
@@ -4587,6 +4600,8 @@ def _parse_level_vals(
         -------
         levels : list of float
             The level edges.
+        explicit_limits : bool
+            Whether the user explicitly provided `vmin` and/or `vmax`.
         **kwargs
             Unused arguments.
         """
@@ -4625,7 +4640,9 @@ def _sanitize_levels(key, array, minsize):
             return array
 
         # Parse input arguments and resolve incompatibilities
-        vmin = vmax = None
+        explicit_limits = vmin is not None or vmax is not None
+        line_contours = min_levels == 1
+        keep_explicit_line_limits = line_contours and explicit_limits
         levels = _not_none(N=N, levels=levels, norm_kw_levs=norm_kw.pop("levels", None))
         if positive and negative:
             warnings._warn_ultraplot(
@@ -4684,6 +4701,8 @@ def _sanitize_levels(key, array, minsize):
             levels, kwargs = self._parse_level_num(
                 *args,
                 levels=levels,
+                vmin=vmin,
+                vmax=vmax,
                 norm=norm,
                 norm_kw=norm_kw,
                 extend=extend,
@@ -4696,24 +4715,30 @@ def _sanitize_levels(key, array, minsize):
             levels = values = None
 
         # Determine default colorbar locator and norm and apply filters
-        # NOTE: DiscreteNorm does not currently support vmin and
-        # vmax different from level list minimum and maximum.
+        # NOTE: Preserve explicit vmin/vmax only for line contours, where levels
+        # represent contour values rather than filled bins.
         # NOTE: The level restriction should have no effect if levels were generated
         # automatically. However want to apply these to manual-input levels as well.
         if levels is not None:
             levels = _restrict_levels(levels)
             if len(levels) == 0:  # skip
                 pass
             elif len(levels) == 1:  # use central colormap color
-                vmin, vmax = levels[0] - 1, levels[0] + 1
+                if not keep_explicit_line_limits or vmin is None:
+                    vmin = levels[0] - 1
+                if not keep_explicit_line_limits or vmax is None:
+                    vmax = levels[0] + 1
             else:  # use minimum and maximum
-                vmin, vmax = np.min(levels), np.max(levels)
+                if not keep_explicit_line_limits or vmin is None:
+                    vmin = np.min(levels)
+                if not keep_explicit_line_limits or vmax is None:
+                    vmax = np.max(levels)
                 if not np.allclose(levels[1] - levels[0], np.diff(levels)):
                     norm = _not_none(norm, "segmented")
             if norm in ("segments", "segmented"):
                 norm_kw["levels"] = levels
 
-        return levels, vmin, vmax, norm, norm_kw, kwargs
+        return levels, vmin, vmax, norm, norm_kw, explicit_limits, kwargs
 
     @staticmethod
     def _parse_level_norm(
@@ -4726,6 +4751,7 @@ def _parse_level_norm(
         discrete_ticks=None,
         discrete_labels=None,
         center_levels=None,
+        explicit_limits=False,
         **kwargs,
     ):
         """
@@ -4748,11 +4774,14 @@ def _parse_level_norm(
             The colorbar locations to tick.
         discrete_labels : array-like, optional
             The colorbar tick labels.
+        explicit_limits : bool, optional
+            Whether `vmin`/`vmax` were explicitly provided by the user.
 
         Returns
         -------
-        norm : `~ultraplot.colors.DiscreteNorm`
-            The discrete normalizer.
+        norm : `~ultraplot.colors.DiscreteNorm` or `~matplotlib.colors.Normalize`
+            The discrete normalizer, or the original continuous normalizer when
+            line contours have explicit limits or use qualitative color lists.
         cmap : `~matplotlib.colors.Colormap`
             The possibly-modified colormap.
         kwargs
@@ -4814,10 +4843,16 @@ def _parse_level_norm(
                 elif extend == "max":
                     unique = "neither"
 
-        # Generate DiscreteNorm and update "child" norm with vmin and vmax from
-        # levels. This lets the colorbar set tick locations properly!
+        # Generate DiscreteNorm for filled-contour style bins. For line contours
+        # with explicit limits or qualitative color lists, keep the continuous
+        # normalizer to preserve one-to-one value->color mapping.
         center_levels = _not_none(center_levels, rc["colorbar.center_levels"])
-        if not isinstance(norm, mcolors.BoundaryNorm) and len(levels) > 1:
+        preserve_line_mapping = min_levels == 1 and (explicit_limits or qualitative)
+        if (
+            not preserve_line_mapping
+            and not isinstance(norm, mcolors.BoundaryNorm)
+            and len(levels) > 1
+        ):
             norm = pcolors.DiscreteNorm(
                 levels,
                 norm=norm,

ultraplot/tests/test_2dplots.py

@@ -6,6 +6,7 @@
 import numpy as np
 import pytest
 import xarray as xr
+from matplotlib.colors import Normalize
 
 import ultraplot as uplt, warnings
 
@@ -291,6 +292,126 @@ def test_levels_with_vmin_vmax(rng):
     return fig
 
 
+def test_contour_levels_respect_explicit_vmin_vmax():
+    """
+    Explicit `vmin` and `vmax` should be preserved for line contours.
+    """
+    data = np.linspace(0, 10, 25).reshape((5, 5))
+    levels = [2, 4, 6]
+    _, ax = uplt.subplots()
+    m = ax.contour(data, levels=levels, cmap="viridis", vmin=0, vmax=10)
+    assert m.norm.vmin == pytest.approx(0)
+    assert m.norm.vmax == pytest.approx(10)
+    assert m.norm(3) == pytest.approx(0.3)
+    assert m.norm(5) == pytest.approx(0.5)
+
+
+def test_contour_levels_default_stretch():
+    """
+    Without explicit limits, level bins should continue to span full cmap range.
+    """
+    data = np.linspace(0, 10, 25).reshape((5, 5))
+    levels = [2, 4, 6]
+    _, ax = uplt.subplots()
+    m = ax.contourf(data, levels=levels, cmap="viridis")
+    assert m.norm(3) == pytest.approx(0.0)
+    assert m.norm(5) == pytest.approx(1.0)
+
+
+def test_contour_levels_default_use_discrete_norm():
+    """
+    Line contours should retain DiscreteNorm behavior unless limits are explicit.
+    """
+    data = np.linspace(0, 10, 25).reshape((5, 5))
+    levels = [2, 4, 6]
+    _, ax = uplt.subplots()
+    m = ax.contour(data, levels=levels, cmap="viridis")
+    assert hasattr(m.norm, "_norm")
+    assert m.norm(3) == pytest.approx(0.0)
+    assert m.norm(5) == pytest.approx(1.0)
+
+
+def test_contourf_levels_keep_level_range_with_explicit_vmin_vmax():
+    """
+    Filled contour bins keep level-based discrete scaling.
+    """
+    data = np.linspace(0, 10, 25).reshape((5, 5))
+    levels = [2, 4, 6]
+    _, ax = uplt.subplots()
+    m = ax.contourf(data, levels=levels, cmap="viridis", vmin=0, vmax=10)
+    assert m.norm.vmin == pytest.approx(2)
+    assert m.norm.vmax == pytest.approx(6)
+    assert m.norm._norm.vmin == pytest.approx(2)
+    assert m.norm._norm.vmax == pytest.approx(6)
+    assert m.norm(3) == pytest.approx(0.0)
+    assert m.norm(5) == pytest.approx(1.0)
+
+
+def test_contour_explicit_colors_match_levels():
+    """
+    Explicit contour line colors should map one-to-one with contour levels.
+    """
+    x = np.linspace(-1, 1, 100)
+    y = np.linspace(-1, 1, 100)
+    X, Y = np.meshgrid(x, y)
+    Z = np.exp(-(X**2 + Y**2))
+    levels = [0.3, 0.6, 0.9]
+    turbo = uplt.Colormap("turbo")
+    colors = turbo(Normalize(vmin=0, vmax=1)(levels))
+    _, ax = uplt.subplots()
+    m = ax.contour(X, Y, Z, levels=levels, colors=colors, linewidths=1)
+    assert np.allclose(np.asarray(m.get_edgecolor()), colors)
+
+
+def test_tricontour_default_use_discrete_norm():
+    """
+    Triangular line contours should default to DiscreteNorm bin mapping.
+    """
+    rng = np.random.default_rng(51423)
+    x = rng.random(40)
+    y = rng.random(40)
+    z = np.sin(3 * x) + np.cos(3 * y)
+    levels = [-1.0, 0.0, 1.0]
+    _, ax = uplt.subplots()
+    m = ax.tricontour(x, y, z, levels=levels, cmap="viridis")
+    assert hasattr(m.norm, "_norm")
+    assert m.norm(-0.5) == pytest.approx(0.0)
+    assert m.norm(0.5) == pytest.approx(1.0)
+
+
+def test_tricontour_levels_respect_explicit_vmin_vmax():
+    """
+    Triangular line contours preserve explicit normalization limits.
+    """
+    rng = np.random.default_rng(51423)
+    x = rng.random(40)
+    y = rng.random(40)
+    z = np.sin(3 * x) + np.cos(3 * y)
+    levels = [-1.0, 0.0, 1.0]
+    _, ax = uplt.subplots()
+    m = ax.tricontour(x, y, z, levels=levels, cmap="viridis", vmin=-2, vmax=2)
+    assert m.norm.vmin == pytest.approx(-2)
+    assert m.norm.vmax == pytest.approx(2)
+    assert m.norm(-0.5) == pytest.approx(0.375)
+    assert m.norm(0.5) == pytest.approx(0.625)
+
+
+def test_tricontour_explicit_colors_match_levels():
+    """
+    Explicit triangular contour colors should map one-to-one with levels.
+    """
+    rng = np.random.default_rng(51423)
+    x = rng.random(40)
+    y = rng.random(40)
+    z = np.sin(3 * x) + np.cos(3 * y)
+    levels = [-1.0, 0.0, 1.0]
+    turbo = uplt.Colormap("turbo")
+    colors = turbo(Normalize(vmin=-2, vmax=2)(levels))
+    _, ax = uplt.subplots()
+    m = ax.tricontour(x, y, z, levels=levels, colors=colors, linewidths=1)
+    assert np.allclose(np.asarray(m.get_edgecolor()), colors)
+
+
 @pytest.mark.mpl_image_compare
 def test_level_restriction(rng):
     """