Move palette plot to da.xrs.plot() accessor

The .xrs accessor (registered on all DataArrays by xrspatial) now has
a plot() method that checks for an embedded TIFF colormap in attrs.
If present, it applies BoundaryNorm with the ListedColormap so that
integer class indices map to the correct palette colors.

    da = read_geotiff('landcover.tif')
    da.xrs.plot()  # palette colors applied automatically

For non-palette DataArrays, falls through to the standard da.plot().
The old plot_geotiff() function is kept as a thin wrapper.

Author: brendancol

xrspatial/accessor.py

@@ -21,6 +21,33 @@ class XrsSpatialDataArrayAccessor:
     def __init__(self, obj):
         self._obj = obj
 
+    # ---- Plot ----
+
+    def plot(self, **kwargs):
+        """Plot the DataArray, using an embedded TIFF colormap if present.
+
+        For palette/indexed-color GeoTIFFs (read via ``read_geotiff``),
+        the TIFF's color table is applied automatically with correct
+        normalization.  For all other DataArrays, falls through to the
+        standard ``da.plot()``.
+
+        Usage::
+
+            da = read_geotiff('landcover.tif')
+            da.xrs.plot()  # palette colors used automatically
+        """
+        import numpy as np
+        cmap = self._obj.attrs.get('cmap')
+        if cmap is not None and 'cmap' not in kwargs:
+            from matplotlib.colors import BoundaryNorm
+            n_colors = len(cmap.colors)
+            boundaries = np.arange(n_colors + 1) - 0.5
+            norm = BoundaryNorm(boundaries, n_colors)
+            kwargs.setdefault('cmap', cmap)
+            kwargs.setdefault('norm', norm)
+            kwargs.setdefault('add_colorbar', True)
+        return self._obj.plot(**kwargs)
+
     # ---- Surface ----
 
     def slope(self, **kwargs):

xrspatial/geotiff/__init__.py

@@ -20,8 +20,7 @@
 from ._reader import read_to_array
 from ._writer import write
 
-__all__ = ['read_geotiff', 'write_geotiff', 'open_cog', 'read_geotiff_dask',
-           'plot_geotiff']
+__all__ = ['read_geotiff', 'write_geotiff', 'open_cog', 'read_geotiff_dask']
 
 
 def _geo_to_coords(geo_info, height: int, width: int) -> dict:
@@ -381,30 +380,8 @@ def _read():
 
 
 def plot_geotiff(da: xr.DataArray, **kwargs):
-    """Plot a DataArray read from a GeoTIFF, using its embedded colormap if present.
+    """Plot a DataArray using its embedded colormap if present.
 
-    For palette/indexed-color TIFFs, the TIFF's color table is used
-    automatically. For other TIFFs, falls through to xarray's default plot.
-
-    Parameters
-    ----------
-    da : xr.DataArray
-        DataArray from read_geotiff.
-    **kwargs
-        Additional keyword arguments passed to da.plot().
-
-    Returns
-    -------
-    matplotlib artist (from da.plot())
+    Deprecated: use ``da.xrs.plot()`` instead.
     """
-    cmap = da.attrs.get('cmap')
-    if cmap is not None and 'cmap' not in kwargs:
-        from matplotlib.colors import BoundaryNorm
-        n_colors = len(cmap.colors)
-        # Build a BoundaryNorm that maps integer index i to palette[i]
-        boundaries = np.arange(n_colors + 1) - 0.5
-        norm = BoundaryNorm(boundaries, n_colors)
-        kwargs.setdefault('cmap', cmap)
-        kwargs.setdefault('norm', norm)
-        kwargs.setdefault('add_colorbar', True)
-    return da.plot(**kwargs)
+    return da.xrs.plot(**kwargs)

xrspatial/geotiff/tests/test_features.py

@@ -920,11 +920,11 @@ def test_palette_cmap_works_with_plot(self, tmp_path):
         assert cmap(0)[:3] == pytest.approx((1.0, 0.0, 0.0), abs=0.01)
         assert cmap(1 / 255)[:3] == pytest.approx((0.0, 1.0, 0.0), abs=0.01)
 
-    def test_plot_geotiff_with_palette(self, tmp_path):
-        """plot_geotiff() uses the embedded colormap."""
+    def test_xrs_plot_with_palette(self, tmp_path):
+        """da.xrs.plot() uses the embedded colormap."""
         import matplotlib
-        matplotlib.use('Agg')  # non-interactive backend for tests
-        from xrspatial.geotiff import plot_geotiff
+        matplotlib.use('Agg')
+        import xrspatial.accessor  # register .xrs accessor
 
         palette = [
             (65535, 0, 0),
@@ -941,7 +941,42 @@ def test_plot_geotiff_with_palette(self, tmp_path):
             f.write(tiff_data)
 
         da = read_geotiff(path)
-        # Should not raise
+        artist = da.xrs.plot()
+        assert artist is not None
+        import matplotlib.pyplot as plt
+        plt.close('all')
+
+    def test_xrs_plot_no_palette(self, tmp_path):
+        """da.xrs.plot() falls through to normal plot for non-palette data."""
+        import matplotlib
+        matplotlib.use('Agg')
+        import xrspatial.accessor
+
+        arr = np.random.RandomState(42).rand(4, 4).astype(np.float32)
+        path = str(tmp_path / 'no_palette.tif')
+        write(arr, path, compression='none', tiled=False)
+
+        da = read_geotiff(path)
+        artist = da.xrs.plot()
+        assert artist is not None
+        import matplotlib.pyplot as plt
+        plt.close('all')
+
+    def test_plot_geotiff_deprecated(self, tmp_path):
+        """plot_geotiff still works as deprecated wrapper."""
+        import matplotlib
+        matplotlib.use('Agg')
+        import xrspatial.accessor
+        from xrspatial.geotiff import plot_geotiff
+
+        palette = [(65535, 0, 0), (0, 65535, 0)] + [(0, 0, 0)] * 254
+        pixels = np.array([[0, 1], [1, 0]], dtype=np.uint8)
+        tiff_data = _make_palette_tiff(2, 2, 8, pixels, palette)
+        path = str(tmp_path / 'deprecated.tif')
+        with open(path, 'wb') as f:
+            f.write(tiff_data)
+
+        da = read_geotiff(path)
         artist = plot_geotiff(da)
         assert artist is not None
         import matplotlib.pyplot as plt