Added named_arrays.plt.scatter() function.

named_arrays/_functions/tests/test_functions.py

@@ -632,6 +632,12 @@ class TestPltPlotLikeFunctions(
         ):
             pass
 
+        @pytest.mark.skip
+        class TestPltScatter(
+            named_arrays.tests.test_core.AbstractTestAbstractArray.TestNamedArrayFunctions.TestPltScatter
+        ):
+            pass
+
         @pytest.mark.xfail
         class TestJacobian(
             named_arrays.tests.test_core.AbstractTestAbstractArray.TestNamedArrayFunctions.TestJacobian,

named_arrays/_scalars/scalar_named_array_functions.py

@@ -341,6 +341,100 @@
     return result
 
 
+@_implements(na.plt.scatter)
+def plt_scatter(
+        *args: na.AbstractScalarArray,
+        s: None | na.AbstractScalarArray = None,
+        c: None | na.AbstractScalarArray = None,
+        ax: None | matplotlib.axes.Axes | na.ScalarArray = None,
+        where: bool | na.AbstractScalarArray = True,
+        components: None | tuple[str, ...] = None,
+        **kwargs,
+) -> na.ScalarArray:
+
+    if components is not None:
+        raise ValueError(f"`components` should be `None` for scalars, got {components}")
+
+    try:
+        args = tuple(scalars._normalize(arg) for arg in args)
+        s = scalars._normalize(s)
+        c = scalars._normalize(c)
+        where = scalars._normalize(where)
+        kwargs = {k: scalars._normalize(kwargs[k]) for k in kwargs}
+    except na.ScalarTypeError:
+        return NotImplemented
+
+    if ax is None:
+        ax = plt.gca()
+    ax = na.as_named_array(ax)
+
+    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
+
+    args = tuple(arg.broadcast_to(shape) for arg in args)
+
+    if np.issubdtype(na.get_dtype(s), np.number):
+        s = na.broadcast_to(s, shape)
+    else:
+        s = na.broadcast_to(s, shape_orthogonal)
+
+    if np.issubdtype(na.get_dtype(c), np.number):
+        if "rgba" in c.shape:
+            c = na.broadcast_to(c, shape | dict(rgba=c.shape["rgba"]))
+        else:
+            c = na.broadcast_to(c, shape)
+    else:
+        c = na.broadcast_to(c, shape_orthogonal)
+
+    where = where.broadcast_to(shape)
+
+    args = tuple(np.where(where, arg, np.nan) for arg in args)
+
+    kwargs_broadcasted = dict()
+    for k in kwargs:
+        kwarg = kwargs[k]
+        if not set(na.shape(kwarg)).issubset(shape_orthogonal):
+            raise ValueError(
+                f"the shape of `{k}`, {na.shape(kwarg)}, "
+                f"should be a subset of the shape of `ax`, {shape_orthogonal}"
+            )
+        kwargs_broadcasted[k] = na.broadcast_to(kwarg, shape_orthogonal)
+    kwargs = kwargs_broadcasted
+
+    result = na.ScalarArray.empty(shape=shape_orthogonal, dtype=object)
+
+    for index in na.ndindex(shape_orthogonal):
+        func_matplotlib = getattr(ax[index].ndarray, "scatter")
+        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,
+            s=s_index,
+            c=c_index,
+            **kwargs_index,
+        )
+
+    return result
+
+
 @_implements(na.jacobian)
 def jacobian(
         function: Callable[[na.AbstractScalar], na.AbstractScalar],

named_arrays/_scalars/tests/test_scalars.py

@@ -888,6 +888,40 @@ class TestPltPlotLikeFunctions(
         ):
             pass
 
+        @pytest.mark.parametrize(
+            argnames="array_2",
+            argvalues=_scalar_arrays_2()[:8],
+        )
+        @pytest.mark.parametrize(
+            argnames="s",
+            argvalues=[
+                None,
+                5,
+            ]
+        )
+        @pytest.mark.parametrize(
+            argnames="c",
+            argvalues=[
+                None,
+                5,
+                na.ScalarArray(
+                    ndarray=np.array([.5, .5, .5, 1]),
+                    axes="rgba",
+                )
+            ]
+        )
+        @pytest.mark.parametrize(
+            argnames="where",
+            argvalues=[
+                True,
+                na.linspace(0, 1, axis="x", num=_num_x) > 0.5,
+            ],
+        )
+        class TestPltScatter(
+            tests.test_core.AbstractTestAbstractArray.TestNamedArrayFunctions.TestPltScatter,
+        ):
+            pass
+
         @pytest.mark.parametrize(
             argnames="function",
             argvalues=[

named_arrays/_scalars/uncertainties/tests/test_uncertainties.py

@@ -685,6 +685,34 @@ class TestPltPlotLikeFunctions(
         ):
             pass
 
+        @pytest.mark.parametrize(
+            argnames="array_2",
+            argvalues=_uncertain_scalar_arrays_2(),
+        )
+        @pytest.mark.parametrize(
+            argnames="s",
+            argvalues=[
+                None,
+            ]
+        )
+        @pytest.mark.parametrize(
+            argnames="c",
+            argvalues=[
+                None,
+            ]
+        )
+        @pytest.mark.parametrize(
+            argnames="where",
+            argvalues=[
+                True,
+                na.linspace(0, 1, axis="x", num=_num_x) > 0.5,
+            ]
+        )
+        class TestPltScatter(
+            named_arrays.tests.test_core.AbstractTestAbstractArray.TestNamedArrayFunctions.TestPltScatter,
+        ):
+            pass
+
         @pytest.mark.parametrize(
             argnames="function",
             argvalues=[

named_arrays/_scalars/uncertainties/uncertainties.py

@@ -289,7 +289,7 @@ def __bool__(self):
         result = super().__bool__()
         nominal = bool(self.nominal)
         distribution = self.distribution
-        if self.axis_distribution in distribution.shape:
+        if self.axis_distribution in na.shape(distribution):
             distribution = np.all(distribution, axis=self.axis_distribution)
         distribution = bool(distribution)
         return result and nominal and distribution

named_arrays/_scalars/uncertainties/uncertainties_named_array_functions.py

@@ -248,6 +248,91 @@
     return result
 
 
+@_implements(na.plt.scatter)
+def plt_scatter(
+        *args: na.AbstractScalar,
+        s: None | na.AbstractScalar = None,
+        c: None | na.AbstractScalar = None,
+        ax: None | matplotlib.axes.Axes | na.ScalarArray = None,
+        where: bool | na.AbstractScalar = True,
+        components: None | tuple[str, ...] = None,
+        **kwargs,
+) -> na.UncertainScalarArray:
+
+    if components is not None:
+        raise ValueError(
+            f"`components` should be `None` for scalars, got {components}"
+        )
+
+    try:
+        args = tuple(uncertainties._normalize(arg) for arg in args)
+        s = uncertainties._normalize(s)
+        c = uncertainties._normalize(c)
+        where = uncertainties._normalize(where)
+        kwargs = {k: uncertainties._normalize(kwargs[k]) for k in kwargs}
+    except na.UncertainScalarTypeError:
+        return NotImplemented
+
+    if ax is None:
+        ax = plt.gca()
+    ax = na.as_named_array(ax)
+
+    axis_distribution = args[0].axis_distribution
+    shape_distribution = na.broadcast_shapes(*[arg.shape_distribution for arg in args])
+
+    if axis_distribution in shape_distribution:
+        num_distribution = shape_distribution[axis_distribution]
+    else:
+        num_distribution = 1
+
+    if num_distribution == 0:
+        alpha = 1
+    else:
+        alpha = max(1 / num_distribution, 1 / 255)
+
+    if "alpha" in kwargs:
+        kwargs["alpha"] = kwargs["alpha"] * alpha
+    else:
+        kwargs["alpha"] = na.UncertainScalarArray(1, alpha)
+
+    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():
+            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=result_nominal,
+        distribution=result_distribution,
+    )
+
+    return result
+
+
 @_implements(na.jacobian)
 def jacobian(
         function: Callable[[na.AbstractScalar], na.AbstractScalar],

named_arrays/_vectors/tests/test_vectors.py

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