Updated named_arrays.plt.scatter() to handle matplotlib >= 3.8

named_arrays/_scalars/scalar_named_array_functions.py

@@ -368,7 +368,12 @@ def plt_scatter(
         ax = plt.gca()
     ax = na.as_named_array(ax)
 
-    shape = na.shape_broadcasted(*args, s, c, ax, where)
+    shape_c = c.shape
+    if "rgba" in c.shape:
+        shape_c.pop("rgba")
+
+    shape = na.shape_broadcasted(*args, s, ax, where)
+    shape = na.broadcast_shapes(shape, shape_c)
 
     shape_orthogonal = ax.shape
 
@@ -382,7 +387,10 @@ def plt_scatter(
     if np.issubdtype(na.get_dtype(c), np.number):
         c = na.broadcast_to(c, shape)
     else:
-        c = na.broadcast_to(c, shape_orthogonal)
+        if "rgba" in c.shape:
+            c = na.broadcast_to(c, shape_orthogonal | dict(rgba=c.shape["rgba"]))
+        else:
+            c = na.broadcast_to(c, shape_orthogonal)
 
     where = where.broadcast_to(shape)
 
@@ -403,9 +411,19 @@ def plt_scatter(
 
     for index in na.ndindex(shape_orthogonal):
         func_matplotlib = getattr(ax[index].ndarray, "scatter")
-        args_index = tuple(arg[index].ndarray for arg in args)
+        args_index = tuple(arg[index].ndarray.reshape(-1) for arg in args)
+
         s_index = s[index].ndarray
+        if s_index is not None:
+            s_index = s_index.reshape(-1)
+
         c_index = c[index].ndarray
+        if c_index is not None:
+            if "rgba" in c.shape:
+                c_index = c[index].ndarray.reshape(-1, c.shape["rgba"])
+            else:
+                c_index = c[index].ndarray.reshape(-1)
+
         kwargs_index = {k: kwargs[k][index].ndarray for k in kwargs}
         result[index] = func_matplotlib(
             *args_index,

named_arrays/_scalars/uncertainties/uncertainties_named_array_functions.py

@@ -267,7 +267,7 @@ def plt_scatter(
     try:
         args = tuple(uncertainties._normalize(arg) for arg in args)
         s = uncertainties._normalize(s)
-        c = uncertainties._normalize(c) if c is not None else c
+        c = uncertainties._normalize(c)
         where = uncertainties._normalize(where)
         kwargs = {k: uncertainties._normalize(kwargs[k]) for k in kwargs}
     except na.UncertainScalarTypeError:
@@ -295,33 +295,42 @@ def plt_scatter(
     else:
         kwargs["alpha"] = na.UncertainScalarArray(1, alpha)
 
-    if c is None:
-        c = na.ScalarArray.empty(shape=ax.shape, dtype=object)
+    result_nominal = na.plt.scatter(
+        *[na.as_named_array(arg.nominal) for arg in args],
+        s=s.nominal,
+        c=c.nominal,
+        ax=ax,
+        where=where.nominal,
+        components=components,
+        **{k: kwargs[k].nominal for k in kwargs},
+    )
+
+    if c.distribution is None:
+        c_distribution = na.ScalarArray.zeros(shape=ax.shape | dict(rgba=4))
         for index in ax.ndindex():
-            c[index] = next(ax[index].ndarray._get_lines.prop_cycler)['color']
-        c = na.UncertainScalarArray(c, c)
+            facecolor = result_nominal[index].ndarray.get_facecolor()[0]
+            c_distribution[index] = na.ScalarArray(
+                ndarray=facecolor,
+                axes="rgba",
+            )
+        c.distribution = c_distribution
+
+    result_distribution = na.plt.scatter(
+        *[na.as_named_array(arg.distribution) for arg in args],
+        s=s.distribution,
+        c=c.distribution,
+        ax=ax,
+        where=where.distribution,
+        components=components,
+        **{k: kwargs[k].distribution for k in kwargs},
+    )
 
     result = na.UncertainScalarArray(
-        nominal=na.plt.scatter(
-            *[na.as_named_array(arg.nominal) for arg in args],
-            s=s.nominal,
-            c=c.nominal,
-            ax=ax,
-            where=where.nominal,
-            components=components,
-            **{k: kwargs[k].nominal for k in kwargs},
-        ),
-        distribution=na.plt.scatter(
-            *[na.as_named_array(arg.distribution) for arg in args],
-            s=s.distribution,
-            c=c.distribution,
-            ax=ax,
-            where=where.distribution,
-            components=components,
-            **{k: kwargs[k].distribution for k in kwargs},
-        )
+        nominal=result_nominal,
+        distribution=result_distribution,
     )
 
+    return result
 
 
 @_implements(na.jacobian)

named_arrays/plt.py

@@ -282,7 +282,7 @@ def scatter(
         The component names of ``*args`` to plot, helpful if ``*args`` are an instance of
         :class:`named_arrays.AbstractVectorArray`.
     kwargs
-        Additional keyword arguments passed to :meth:`matplotlib.axes.Axes.plot`.
+        Additional keyword arguments passed to :meth:`matplotlib.axes.Axes.scatter`.
         These can be instances of :class:`named_arrays.AbstractArray`.
 
     Examples
@@ -296,7 +296,7 @@ def scatter(
         import matplotlib.pyplot as plt
         import named_arrays as na
 
-        x = na.linspace(0, 2 * np.pi, axis="x",  num=101)
+        x = na.linspace(0, 2 * np.pi, axis="x",  num=51)
         y = np.sin(x)
 
         plt.figure();
@@ -329,7 +329,7 @@ def scatter(
 
         fig, ax = na.plt.subplots(axis_rows="z", nrows=z.shape["z"], sharex=True)
 
-        na.plt.scatter(x, y, ax=ax, axis="x");
+        na.plt.scatter(x, y, ax=ax);
     """
     if transformation is not None:
         args = tuple(transformation(arg) for arg in args)