ENH: allow image to interpolate post RGBA

jklymak · jklymak · commit f50ab9acba6d · 2021-03-21T09:01:20.000-07:00
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5392,9 +5392,10 @@ def fill_betweenx(self, y, x1, x2=0, where=None,
     #### plotting z(x, y): imshow, pcolor and relatives, contour
     @_preprocess_data()
     def imshow(self, X, cmap=None, norm=None, aspect=None,
-               interpolation=None, alpha=None, vmin=None, vmax=None,
-               origin=None, extent=None, *, filternorm=True, filterrad=4.0,
-               resample=None, url=None, **kwargs):
+               interpolation=None, interp_postrgba=False, alpha=None,
+               vmin=None, vmax=None, origin=None, extent=None, *,
+               filternorm=True, filterrad=4.0, resample=None, url=None,
+               **kwargs):
         """
         Display data as an image, i.e., on a 2D regular raster.
 
@@ -5485,6 +5486,14 @@ def imshow(self, X, cmap=None, norm=None, aspect=None,
             which can be set by *filterrad*. Additionally, the antigrain image
             resize filter is controlled by the parameter *filternorm*.
 
+        interp_postrgba : bool, default: False
+            Whether the interpolation above is carried out before or after
+            the image data has been passed through the *norm* and colormapping.
+            This will cause the interpolation to take place in colorspace
+            rather than dataspace, which can be better for visual
+            anti-aliasing.  Note that this is ignored if *X* is passed in as
+            RGB(A).  
+
         alpha : float or array-like, optional
             The alpha blending value, between 0 (transparent) and 1 (opaque).
             If *alpha* is an array, the alpha blending values are applied pixel
@@ -5582,9 +5591,9 @@ def imshow(self, X, cmap=None, norm=None, aspect=None,
         if aspect is None:
             aspect = rcParams['image.aspect']
         self.set_aspect(aspect)
-        im = mimage.AxesImage(self, cmap, norm, interpolation, origin, extent,
-                              filternorm=filternorm, filterrad=filterrad,
-                              resample=resample, **kwargs)
+        im = mimage.AxesImage(self, cmap, norm, interpolation, interp_postrgba,
+                              origin, extent, filternorm=filternorm,
+                              filterrad=filterrad, resample=resample, **kwargs)
 
         im.set_data(X)
         im.set_alpha(alpha)
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -234,6 +234,7 @@ def __init__(self, ax,
                  cmap=None,
                  norm=None,
                  interpolation=None,
+                 interp_postrgba=False,
                  origin=None,
                  filternorm=True,
                  filterrad=4.0,
@@ -249,6 +250,7 @@ def __init__(self, ax,
         self.set_filternorm(filternorm)
         self.set_filterrad(filterrad)
         self.set_interpolation(interpolation)
+        self.set_interp_postrgba(interp_postrgba)
         self.set_resample(resample)
         self.axes = ax
 
@@ -393,7 +395,7 @@ def _make_image(self, A, in_bbox, out_bbox, clip_bbox, magnification=1.0,
             if not (A.ndim == 2 or A.ndim == 3 and A.shape[-1] in (3, 4)):
                 raise ValueError(f"Invalid shape {A.shape} for image data")
 
-            if A.ndim == 2:
+            if A.ndim == 2 and not self._interp_postrgba:
                 # if we are a 2D array, then we are running through the
                 # norm + colormap transformation.  However, in general the
                 # input data is not going to match the size on the screen so we
@@ -543,6 +545,9 @@ def _make_image(self, A, in_bbox, out_bbox, clip_bbox, magnification=1.0,
                                        ):
                     output = self.norm(resampled_masked)
             else:
+                if A.ndim == 2:  # _interp_postrgba is True
+                    self.norm.autoscale_None(A)
+                    A = self.to_rgba(A)
                 if A.shape[2] == 3:
                     A = _rgb_to_rgba(A)
                 alpha = self._get_scalar_alpha()
@@ -775,6 +780,20 @@ def set_interpolation(self, s):
         self._interpolation = s
         self.stale = True
 
+    def set_interp_postrgba(self, s):
+        """
+        Set whether interpolation happens after the image has been
+        transformed to RGBA.
+
+        Parameters
+        ----------
+        s : bool or None
+        """
+        if s is None:
+            s = False  #  placeholder for maybe having rcParam
+        self._interp_postrgba = s
+        self.stale = True
+
     def can_composite(self):
         """Return whether the image can be composited with its neighbors."""
         trans = self.get_transform()
@@ -887,6 +906,7 @@ def __init__(self, ax,
                  cmap=None,
                  norm=None,
                  interpolation=None,
+                 interp_postrgba=False,
                  origin=None,
                  extent=None,
                  filternorm=True,
@@ -902,6 +922,7 @@ def __init__(self, ax,
             cmap=cmap,
             norm=norm,
             interpolation=interpolation,
+            interp_postrgba=interp_postrgba,
             origin=origin,
             filternorm=filternorm,
             filterrad=filterrad,
