Added named_arrays.plt.pcolormesh() function.

named_arrays/_functions/function_named_array_functions.py

@@ -1,5 +1,6 @@
 from typing import Callable
 import numpy as np
+import matplotlib
 import astropy.units as u
 import named_arrays as na
 import named_arrays._scalars.scalar_named_array_functions
@@ -88,3 +89,37 @@ def unit_normalized(
         squeeze: bool = True,
 ) -> u.UnitBase | na.AbstractArray:
     return na.unit_normalized(a.outputs, squeeze=squeeze)
+
+
+@_implements(na.plt.pcolormesh)
+def pcolormesh(
+    *XY: na.AbstractArray,
+    C: na.AbstractFunctionArray,
+    components: None | tuple[str, str] = None,
+    axis_rgb: None | str = None,
+    ax: None | matplotlib.axes.Axes | na.AbstractArray = None,
+    cmap: None | str | matplotlib.colors.Colormap = None,
+    norm: None | str | matplotlib.colors.Normalize = None,
+    vmin: None | na.ArrayLike = None,
+    vmax: None | na.ArrayLike = None,
+    **kwargs,
+) -> na.ScalarArray:
+
+    if len(XY) != 0:
+        raise ValueError(
+            "if `C` is an instance of `na.AbstractFunctionArray`, "
+            "`XY` must not be specified."
+        )
+
+    return na.plt.pcolormesh(
+        C.inputs,
+        C=C.outputs,
+        components=components,
+        axis_rgb=axis_rgb,
+        ax=ax,
+        cmap=cmap,
+        norm=norm,
+        vmin=vmin,
+        vmax=vmax,
+        **kwargs,
+    )

named_arrays/_functions/tests/test_functions.py

@@ -638,6 +638,53 @@ class TestPltScatter(
         ):
             pass
 
+        class TestPltPcolormesh(
+            named_arrays.tests.test_core.AbstractTestAbstractArray.TestNamedArrayFunctions.TestPltPcolormesh,
+        ):
+            @pytest.mark.parametrize("axis_rgb", [None, "rgb"])
+            def test_pcolormesh(
+                self,
+                array: na.AbstractScalarArray,
+                axis_rgb: None | str
+            ):
+                if not isinstance(array.inputs, na.AbstractVectorArray):
+                    return
+
+                components = list(array.inputs.components.keys())[:2]
+
+                kwargs = dict(
+                    C=array,
+                    axis_rgb=axis_rgb,
+                    components=components,
+                )
+
+                if isinstance(array.outputs, na.AbstractVectorArray):
+                    with pytest.raises(TypeError):
+                        na.plt.pcolormesh(**kwargs)
+                    return
+                elif isinstance(array.outputs, na.AbstractUncertainScalarArray):
+                    with pytest.raises(TypeError):
+                        na.plt.pcolormesh(**kwargs)
+                    return
+
+                if any(isinstance(array.inputs.components[c], na.AbstractUncertainScalarArray) for c in components):
+                    with pytest.raises(TypeError):
+                        na.plt.pcolormesh(**kwargs)
+                    return
+
+                if axis_rgb is not None:
+                    with pytest.raises(ValueError):
+                        na.plt.pcolormesh(**kwargs)
+                    return
+
+                if array.ndim != 2:
+                    with pytest.raises(ValueError):
+                        na.plt.pcolormesh(**kwargs)
+                    return
+
+                result = na.plt.pcolormesh(**kwargs)
+                assert isinstance(result, na.ScalarArray)
+
         @pytest.mark.xfail
         class TestJacobian(
             named_arrays.tests.test_core.AbstractTestAbstractArray.TestNamedArrayFunctions.TestJacobian,

named_arrays/_scalars/scalar_named_array_functions.py

@@ -512,6 +512,64 @@ def plt_imshow(
     return result
 
 
+@_implements(na.plt.pcolormesh)
+def pcolormesh(
+    *XY: na.AbstractScalarArray,
+    C: na.AbstractScalarArray,
+    components: None | tuple[str, str] = None,
+    axis_rgb: None | str = None,
+    ax: None | matplotlib.axes.Axes | na.AbstractScalarArray = None,
+    cmap: None | str | matplotlib.colors.Colormap = None,
+    norm: None | str | matplotlib.colors.Normalize = None,
+    vmin: None | float | u.Quantity | na.AbstractScalarArray = None,
+    vmax: None | float | u.Quantity | na.AbstractScalarArray = None,
+    **kwargs,
+) -> na.ScalarArray:
+
+    if components is not None:
+        raise ValueError(f"`components` should be `None` for scalars, got {components}")
+
+    try:
+        XY = tuple(scalars._normalize(arg) for arg in XY)
+        C = scalars._normalize(C)
+        vmin = scalars._normalize(vmin) if vmin is not None else vmin
+        vmax = scalars._normalize(vmax) if vmax is not None else vmax
+    except na.ScalarTypeError:  # pragma: nocover
+        pass
+
+    if ax is None:
+        ax = plt.gca()
+    ax = na.as_named_array(ax)
+
+    if axis_rgb is not None:    # pragma: nocover
+        if axis_rgb not in C.shape:
+            raise ValueError(f"`{axis_rgb=}` must be a member of `{C.shape=}`")
+
+    shape_C = na.shape_broadcasted(*XY, C, ax)
+    shape = {a: shape_C[a] for a in shape_C if a != axis_rgb}
+    shape_orthogonal = ax.shape
+
+    XY = tuple(arg.broadcast_to(shape) for arg in XY)
+    C = C.broadcast_to(shape_C)
+    vmin = vmin.broadcast_to(shape_orthogonal) if vmin is not None else vmin
+    vmax = vmax.broadcast_to(shape_orthogonal) if vmax is not None else vmax
+
+    result = na.ScalarArray.empty(shape_orthogonal, dtype=object)
+
+    for index in na.ndindex(shape_orthogonal):
+        result[index] = ax[index].ndarray.pcolormesh(
+            *[arg[index].ndarray for arg in XY],
+            C[index].ndarray,
+            cmap=cmap,
+            norm=norm,
+            vmin=vmin[index].ndarray if vmin is not None else vmin,
+            vmax=vmax[index].ndarray if vmax is not None else vmax,
+            **kwargs,
+        )
+
+    return result
+
+
 @_implements(na.jacobian)
 def jacobian(
         function: Callable[[na.AbstractScalar], na.AbstractScalar],

named_arrays/_scalars/uncertainties/tests/test_uncertainties.py

@@ -713,6 +713,12 @@ class TestPltScatter(
         ):
             pass
 
+        @pytest.mark.skip
+        class TestPltPcolormesh(
+            named_arrays._scalars.tests.test_scalars.AbstractTestAbstractScalar.TestNamedArrayFunctions.TestPltPcolormesh,
+        ):
+            pass
+
         @pytest.mark.parametrize(
             argnames="function",
             argvalues=[

named_arrays/_vectors/tests/test_vectors.py

@@ -539,6 +539,12 @@ class TestPltScatter(
         ):
             pass
 
+        @pytest.mark.skip
+        class TestPltPcolormesh(
+            named_arrays.tests.test_core.AbstractTestAbstractArray.TestNamedArrayFunctions.TestPltPcolormesh,
+        ):
+            pass
+
         @pytest.mark.parametrize(
             argnames="function",
             argvalues=[

named_arrays/_vectors/vector_named_array_functions.py

@@ -4,6 +4,7 @@
 import matplotlib.axes
 import astropy.units as u
 import named_arrays as na
+from named_arrays._scalars import scalars
 import named_arrays._scalars.scalar_named_array_functions
 from . import vectors
 
@@ -264,6 +265,52 @@ def plt_plot_like(
     )
 
 
+@_implements(na.plt.pcolormesh)
+def pcolormesh(
+    *XY: na.AbstractVectorArray,
+    C: na.AbstractScalarArray,
+    components: None | tuple[str, str] = None,
+    axis_rgb: None | str = None,
+    ax: None | matplotlib.axes.Axes | na.AbstractScalarArray = None,
+    cmap: None | str | matplotlib.colors.Colormap = None,
+    norm: None | str | matplotlib.colors.Normalize = None,
+    vmin: None | float | u.Quantity | na.AbstractScalarArray = None,
+    vmax: None | float | u.Quantity | na.AbstractScalarArray = None,
+    **kwargs,
+) -> na.ScalarArray:
+    try:
+        C = scalars._normalize(C)
+        vmin = scalars._normalize(vmin) if vmin is not None else vmin
+        vmax = scalars._normalize(vmax) if vmax is not None else vmax
+    except na.ScalarTypeError:
+        return NotImplemented
+
+    try:
+        prototype = vectors._prototype(*XY)
+        XY = tuple(vectors._normalize(arg, prototype) for arg in XY)
+    except na.VectorTypeError:
+        return NotImplemented
+
+    if len(XY) != 1:
+        raise ValueError("if any element of `XY` is a vector, `XY` must have a length of 1")
+    XY = XY[0]
+
+    components_XY = XY.components
+    components = [components_XY[c] for c in components]
+
+    return na.plt.pcolormesh(
+        *components,
+        C=C,
+        axis_rgb=axis_rgb,
+        ax=ax,
+        cmap=cmap,
+        norm=norm,
+        vmin=vmin,
+        vmax=vmax,
+        **kwargs,
+    )
+
+
 @_implements(na.jacobian)
 def jacobian(
         function: Callable[[InputT], OutputT],

named_arrays/plt.py

@@ -11,6 +11,7 @@
     "fill",
     "scatter",
     "imshow",
+    "pcolormesh",
 ]
 
 
@@ -508,3 +509,121 @@ def imshow(
         extent=extent,
         **kwargs,
     )
+
+
+def pcolormesh(
+    *XY: na.AbstractArray,
+    C: na.AbstractArray,
+    components: None | tuple[str, str] = None,
+    axis_rgb: None | str = None,
+    ax: None | matplotlib.axes.Axes | na.AbstractArray = None,
+    cmap: None | str | matplotlib.colors.Colormap = None,
+    norm: None | str | matplotlib.colors.Normalize = None,
+    vmin: None | na.ArrayLike = None,
+    vmax: None | na.ArrayLike = None,
+    **kwargs,
+) -> na.AbstractScalar:
+    """
+    A thin wrapper around :func:`matplotlib.pyplot.pcolormesh` for named arrays.
+
+    Parameters
+    ----------
+    XY
+        The coordinates of the mesh.
+        If `C` is a scalar, `XY` can either be two scalars or one vector .
+        If `C` is a function, `XY` is not specified.
+        If `XY` is not specified as two scalars, the `components` must be given,
+        see below.
+    C
+        The mesh data.
+    components
+        If `XY` is not specified as two scalars, this parameter should
+        be a tuple of two strings, specifying the vector components of `XY`
+        to use as the horizontal and vertical components of the mesh.
+    axis_rgb
+        The optional logical axis along which the RGB color channels are
+        distributed.
+    ax
+        The instances of :class:`matplotlib.axes.Axes` to use.
+        If :obj:`None`, calls :func:`matplotlib.pyplot.gca` to get the current axes.
+        If an instance of :class:`named_arrays.ScalarArray`, ``ax.shape`` should be a subset of the broadcasted shape of
+        ``*args``.
+    cmap
+        The colormap used to map scalar data to colors.
+    norm
+        The normalization method used to scale data into the range [0, 1] before
+        mapping to colors.
+    vmin
+        The minimum value of the data range.
+    vmax
+        The maximum value of the data range.
+    kwargs
+        Additional keyword arguments accepted by `matplotlib.pyplot.pcolormesh`
+
+    Examples
+    --------
+
+    Plot a random 2D mesh
+
+    .. jupyter-execute::
+
+        import matplotlib.pyplot as plt
+        import named_arrays as na
+
+        # Define the size of the grid
+        shape = dict(x=16, y=16)
+
+        # Define a simple coordinate grid
+        x = na.linspace(-2, 2, axis="x", num=shape["x"])
+        y = na.linspace(-1, 1, axis="y", num=shape["y"])
+
+        # Define a random 2D array of values to plot
+        a = na.random.uniform(-1, 1, shape_random=shape)
+
+        # Plot the coordinates and values using pcolormesh
+        fig, ax = plt.subplots(constrained_layout=True)
+        na.plt.pcolormesh(x, y, C=a, ax=ax);
+
+    |
+
+    Plot a grid of random 2D meshes
+
+    .. jupyter-execute::
+
+        import named_arrays as na
+
+        # Define the size of the grid
+        shape = dict(row=2, col=3, x=16, y=16)
+
+        # Define a simple coordinate grid
+        x = na.linspace(-2, 2, axis="x", num=shape["x"])
+        y = na.linspace(-1, 1, axis="y", num=shape["y"])
+
+        # Define a random 2D array of values to plot
+        a = na.random.uniform(-1, 1, shape_random=shape)
+
+        # Plot the coordinates and values using pcolormesh
+        fig, ax = na.plt.subplots(
+            axis_rows="row",
+            nrows=shape["row"],
+            axis_cols="col",
+            ncols=shape["col"],
+            sharex=True,
+            sharey=True,
+            constrained_layout=True,
+        )
+        na.plt.pcolormesh(x, y, C=a, ax=ax);
+    """
+    return na._named_array_function(
+        pcolormesh,
+        *XY,
+        C=C,
+        axis_rgb=axis_rgb,
+        ax=ax,
+        cmap=cmap,
+        norm=norm,
+        vmin=vmin,
+        vmax=vmax,
+        components=components,
+        **kwargs,
+    )