From 9e38674a443aa1df7ea76f85614f581b5aa50266 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?H=C3=A5kon=20Wiik=20=C3=85nes?= <hwaanes@gmail.com>
Date: Wed, 28 Sep 2022 21:19:07 +0200
Subject: [PATCH] Make stereo. scatter() allow c/color and s/sizes to set
 color/sizes
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Håkon Wiik Ånes <hwaanes@gmail.com>
---
 CHANGELOG.rst                              |  9 ++++++++
 orix/plot/stereographic_plot.py            | 22 ++++++++++++++------
 orix/tests/plot/test_stereographic_plot.py | 24 ++++++++++++++++++++++
 3 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 2dd3ba58..f378aebb 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -7,6 +7,15 @@ All notable changes to the ``orix`` project are documented in this file.
 The format is based on `Keep a Changelog <https://keepachangelog.com/en/1.0.0/>`_, and
 this project adheres to `Semantic Versioning <https://semver.org/spec/v2.0.0.html>`_.
 
+2022-09-30 - version 0.10.1
+===========================
+
+Fixed
+-----
+- ``StereographicPlot.scatter()`` now accepts both ``c``/``color`` and ``s``/``sizes``
+  to set the color and sizes of scatter points, in line with
+  ``matplotlib.axes.Axes.scatter()``.
+
 2022-09-22 - version 0.10.0
 ===========================
 
diff --git a/orix/plot/stereographic_plot.py b/orix/plot/stereographic_plot.py
index 4656bb2e..f818463b 100644
--- a/orix/plot/stereographic_plot.py
+++ b/orix/plot/stereographic_plot.py
@@ -202,10 +202,20 @@ def scatter(
         if x.size == 0:
             return
 
-        # Color(s) and size(s)
-        c = updated_kwargs.pop("c", "C0")
+        # Color(s)
+        if "color" in updated_kwargs.keys():
+            key_color = "color"
+        else:
+            key_color = "c"
+        c = updated_kwargs.pop(key_color, "C0")
         c = _get_array_of_values(value=c, visible=visible)
-        s = updated_kwargs.pop("s", None)
+
+        # Size(s)
+        if "sizes" in updated_kwargs.keys():
+            key_size = "sizes"
+        else:
+            key_size = "s"
+        s = updated_kwargs.pop(key_size, None)
         if s is not None:
             s = _get_array_of_values(value=s, visible=visible)
 
@@ -333,8 +343,6 @@ def pole_density_function(
             sigma=sigma,
             log=log,
             hemisphere=self.hemisphere,
-            symmetry=None,
-            mrd=True,
             weights=weights,
         )
 
@@ -758,7 +766,9 @@ def _prepare_to_call_inherited_method(
         return x, y, visible, updated_kwargs
 
     def _pretransform_input(
-        self, values: Union[Vector3d, Tuple[np.ndarray, np.ndarray]], sort: bool = False
+        self,
+        values: Union[Vector3d, Tuple[Vector3d], Tuple[np.ndarray, np.ndarray]],
+        sort: bool = False,
     ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
         """Return arrays of (x, y) from input data.
 
diff --git a/orix/tests/plot/test_stereographic_plot.py b/orix/tests/plot/test_stereographic_plot.py
index 7f8b9989..a9f33937 100644
--- a/orix/tests/plot/test_stereographic_plot.py
+++ b/orix/tests/plot/test_stereographic_plot.py
@@ -245,6 +245,30 @@ def test_order_in_hemisphere(self):
 
         plt.close("all")
 
+    def test_color_parameter(self):
+        """Pass either ``color`` or ``c`` to color scatter points."""
+        v = Vector3d([[1, 0, 0], [1, 1, 0], [1, 1, 1]])
+
+        colors = [f"C{i}" for i in range(v.size)]
+        colors_rgba = np.array([mcolors.to_rgba(c) for c in colors])
+
+        fig = v.scatter(color=colors, return_figure=True)
+        assert np.allclose(fig.axes[0].collections[0].get_facecolors(), colors_rgba)
+
+        fig2 = v.scatter(c=colors, return_figure=True)
+        assert np.allclose(fig2.axes[0].collections[0].get_facecolors(), colors_rgba)
+
+    def test_size_parameter(self):
+        """Pass either ``sizes`` or ``s`` to set scatter points sizes."""
+        v = Vector3d([[1, 0, 0], [1, 1, 0], [1, 1, 1]])
+        sizes = np.arange(v.size)
+
+        fig = v.scatter(sizes=sizes, return_figure=True)
+        assert np.allclose(fig.axes[0].collections[0].get_sizes(), sizes)
+
+        fig2 = v.scatter(s=sizes, return_figure=True)
+        assert np.allclose(fig2.axes[0].collections[0].get_sizes(), sizes)
+
 
 class TestSymmetryMarker:
     def test_properties(self):