Merge pull request #27815 from anjabeck/side-option-violinplot

ksunden · web-flow · commit d57fa9721d24 · 2024-03-01T14:57:23.000-06:00
Add side option to violinplot
diff --git a/doc/users/next_whats_new/sides_violinplot.rst b/doc/users/next_whats_new/sides_violinplot.rst
@@ -0,0 +1,4 @@
+Add option to plot only one half of violin plot
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Setting the parameter *side* to 'low' or 'high' allows to only plot one half of the violin plot.
diff --git a/galleries/examples/statistics/violinplot.py b/galleries/examples/statistics/violinplot.py
@@ -28,55 +28,73 @@
 pos = [1, 2, 4, 5, 7, 8]
 data = [np.random.normal(0, std, size=100) for std in pos]
 
-fig, axs = plt.subplots(nrows=2, ncols=5, figsize=(10, 6))
+fig, axs = plt.subplots(nrows=2, ncols=6, figsize=(10, 4))
 
 axs[0, 0].violinplot(data, pos, points=20, widths=0.3,
                      showmeans=True, showextrema=True, showmedians=True)
-axs[0, 0].set_title('Custom violinplot 1', fontsize=fs)
+axs[0, 0].set_title('Custom violin 1', fontsize=fs)
 
 axs[0, 1].violinplot(data, pos, points=40, widths=0.5,
                      showmeans=True, showextrema=True, showmedians=True,
                      bw_method='silverman')
-axs[0, 1].set_title('Custom violinplot 2', fontsize=fs)
+axs[0, 1].set_title('Custom violin 2', fontsize=fs)
 
 axs[0, 2].violinplot(data, pos, points=60, widths=0.7, showmeans=True,
                      showextrema=True, showmedians=True, bw_method=0.5)
-axs[0, 2].set_title('Custom violinplot 3', fontsize=fs)
+axs[0, 2].set_title('Custom violin 3', fontsize=fs)
 
 axs[0, 3].violinplot(data, pos, points=60, widths=0.7, showmeans=True,
                      showextrema=True, showmedians=True, bw_method=0.5,
                      quantiles=[[0.1], [], [], [0.175, 0.954], [0.75], [0.25]])
-axs[0, 3].set_title('Custom violinplot 4', fontsize=fs)
+axs[0, 3].set_title('Custom violin 4', fontsize=fs)
 
 axs[0, 4].violinplot(data[-1:], pos[-1:], points=60, widths=0.7,
                      showmeans=True, showextrema=True, showmedians=True,
                      quantiles=[0.05, 0.1, 0.8, 0.9], bw_method=0.5)
-axs[0, 4].set_title('Custom violinplot 5', fontsize=fs)
+axs[0, 4].set_title('Custom violin 5', fontsize=fs)
+
+axs[0, 5].violinplot(data[-1:], pos[-1:], points=60, widths=0.7,
+                     showmeans=True, showextrema=True, showmedians=True,
+                     quantiles=[0.05, 0.1, 0.8, 0.9], bw_method=0.5, side='low')
+
+axs[0, 5].violinplot(data[-1:], pos[-1:], points=60, widths=0.7,
+                     showmeans=True, showextrema=True, showmedians=True,
+                     quantiles=[0.05, 0.1, 0.8, 0.9], bw_method=0.5, side='high')
+axs[0, 5].set_title('Custom violin 6', fontsize=fs)
 
 axs[1, 0].violinplot(data, pos, points=80, vert=False, widths=0.7,
                      showmeans=True, showextrema=True, showmedians=True)
-axs[1, 0].set_title('Custom violinplot 6', fontsize=fs)
+axs[1, 0].set_title('Custom violin 7', fontsize=fs)
 
 axs[1, 1].violinplot(data, pos, points=100, vert=False, widths=0.9,
                      showmeans=True, showextrema=True, showmedians=True,
                      bw_method='silverman')
-axs[1, 1].set_title('Custom violinplot 7', fontsize=fs)
+axs[1, 1].set_title('Custom violin 8', fontsize=fs)
 
 axs[1, 2].violinplot(data, pos, points=200, vert=False, widths=1.1,
                      showmeans=True, showextrema=True, showmedians=True,
                      bw_method=0.5)
-axs[1, 2].set_title('Custom violinplot 8', fontsize=fs)
+axs[1, 2].set_title('Custom violin 9', fontsize=fs)
 
 axs[1, 3].violinplot(data, pos, points=200, vert=False, widths=1.1,
                      showmeans=True, showextrema=True, showmedians=True,
                      quantiles=[[0.1], [], [], [0.175, 0.954], [0.75], [0.25]],
                      bw_method=0.5)
-axs[1, 3].set_title('Custom violinplot 9', fontsize=fs)
+axs[1, 3].set_title('Custom violin 10', fontsize=fs)
 
 axs[1, 4].violinplot(data[-1:], pos[-1:], points=200, vert=False, widths=1.1,
                      showmeans=True, showextrema=True, showmedians=True,
                      quantiles=[0.05, 0.1, 0.8, 0.9], bw_method=0.5)
-axs[1, 4].set_title('Custom violinplot 10', fontsize=fs)
+axs[1, 4].set_title('Custom violin 11', fontsize=fs)
+
+axs[1, 5].violinplot(data[-1:], pos[-1:], points=200, vert=False, widths=1.1,
+                     showmeans=True, showextrema=True, showmedians=True,
+                     quantiles=[0.05, 0.1, 0.8, 0.9], bw_method=0.5, side='low')
+
+axs[1, 5].violinplot(data[-1:], pos[-1:], points=200, vert=False, widths=1.1,
+                     showmeans=True, showextrema=True, showmedians=True,
+                     quantiles=[0.05, 0.1, 0.8, 0.9], bw_method=0.5, side='high')
+axs[1, 5].set_title('Custom violin 12', fontsize=fs)
 
 
 for ax in axs.flat:
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -8205,7 +8205,7 @@ def matshow(self, Z, **kwargs):
     @_preprocess_data(replace_names=["dataset"])
     def violinplot(self, dataset, positions=None, vert=True, widths=0.5,
                    showmeans=False, showextrema=True, showmedians=False,
-                   quantiles=None, points=100, bw_method=None):
+                   quantiles=None, points=100, bw_method=None, side='both'):
         """
         Make a violin plot.
 
@@ -8256,6 +8256,10 @@ def violinplot(self, dataset, positions=None, vert=True, widths=0.5,
           callable, it should take a `matplotlib.mlab.GaussianKDE` instance as
           its only parameter and return a float.
 
+        side : {'both', 'low', 'high'}, default: 'both'
+            'both' plots standard violins. 'low'/'high' only
+            plots the side below/above the positions value.
+
         data : indexable object, optional
             DATA_PARAMETER_PLACEHOLDER
 
@@ -8307,10 +8311,10 @@ def _kde_method(X, coords):
                                      quantiles=quantiles)
         return self.violin(vpstats, positions=positions, vert=vert,
                            widths=widths, showmeans=showmeans,
-                           showextrema=showextrema, showmedians=showmedians)
+                           showextrema=showextrema, showmedians=showmedians, side=side)
 
     def violin(self, vpstats, positions=None, vert=True, widths=0.5,
-               showmeans=False, showextrema=True, showmedians=False):
+               showmeans=False, showextrema=True, showmedians=False, side='both'):
         """
         Draw a violin plot from pre-computed statistics.
 
@@ -8366,6 +8370,10 @@ def violin(self, vpstats, positions=None, vert=True, widths=0.5,
         showmedians : bool, default: False
           Whether to show the median with a line.
 
+        side : {'both', 'low', 'high'}, default: 'both'
+            'both' plots standard violins. 'low'/'high' only
+            plots the side below/above the positions value.
+
         Returns
         -------
         dict
@@ -8428,8 +8436,13 @@ def violin(self, vpstats, positions=None, vert=True, widths=0.5,
         elif len(widths) != N:
             raise ValueError(datashape_message.format("widths"))
 
+        # Validate side
+        _api.check_in_list(["both", "low", "high"], side=side)
+
         # Calculate ranges for statistics lines (shape (2, N)).
-        line_ends = [[-0.25], [0.25]] * np.array(widths) + positions
+        line_ends = [[-0.25 if side in ['both', 'low'] else 0],
+                     [0.25 if side in ['both', 'high'] else 0]] \
+                          * np.array(widths) + positions
 
         # Colors.
         if mpl.rcParams['_internal.classic_mode']:
@@ -8441,20 +8454,34 @@ def violin(self, vpstats, positions=None, vert=True, widths=0.5,
         # Check whether we are rendering vertically or horizontally
         if vert:
             fill = self.fill_betweenx
-            perp_lines = functools.partial(self.hlines, colors=linecolor)
-            par_lines = functools.partial(self.vlines, colors=linecolor)
+            if side in ['low', 'high']:
+                perp_lines = functools.partial(self.hlines, colors=linecolor,
+                                                capstyle='projecting')
+                par_lines = functools.partial(self.vlines, colors=linecolor,
+                                                capstyle='projecting')
+            else:
+                perp_lines = functools.partial(self.hlines, colors=linecolor)
+                par_lines = functools.partial(self.vlines, colors=linecolor)
         else:
             fill = self.fill_between
-            perp_lines = functools.partial(self.vlines, colors=linecolor)
-            par_lines = functools.partial(self.hlines, colors=linecolor)
+            if side in ['low', 'high']:
+                perp_lines = functools.partial(self.vlines, colors=linecolor,
+                                                capstyle='projecting')
+                par_lines = functools.partial(self.hlines, colors=linecolor,
+                                                capstyle='projecting')
+            else:
+                perp_lines = functools.partial(self.vlines, colors=linecolor)
+                par_lines = functools.partial(self.hlines, colors=linecolor)
 
         # Render violins
         bodies = []
         for stats, pos, width in zip(vpstats, positions, widths):
             # The 0.5 factor reflects the fact that we plot from v-p to v+p.
             vals = np.array(stats['vals'])
             vals = 0.5 * width * vals / vals.max()
-            bodies += [fill(stats['coords'], -vals + pos, vals + pos,
+            bodies += [fill(stats['coords'],
+                            -vals + pos if side in ['both', 'low'] else pos,
+                            vals + pos if side in ['both', 'high'] else pos,
                             facecolor=fillcolor, alpha=0.3)]
             means.append(stats['mean'])
             mins.append(stats['min'])
diff --git a/lib/matplotlib/axes/_axes.pyi b/lib/matplotlib/axes/_axes.pyi
@@ -743,6 +743,7 @@ class Axes(_AxesBase):
         | float
         | Callable[[GaussianKDE], float]
         | None = ...,
+        side: Literal["both", "low", "high"] = ...,
         *,
         data=...,
     ) -> dict[str, Collection]: ...
@@ -755,6 +756,7 @@ class Axes(_AxesBase):
         showmeans: bool = ...,
         showextrema: bool = ...,
         showmedians: bool = ...,
+        side: Literal["both", "low", "high"] = ...,
     ) -> dict[str, Collection]: ...
 
     table = mtable.table
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -4075,6 +4075,7 @@ def violinplot(
     | float
     | Callable[[GaussianKDE], float]
     | None = None,
+    side: Literal["both", "low", "high"] = "both",
     *,
     data=None,
 ) -> dict[str, Collection]:
@@ -4089,6 +4090,7 @@ def violinplot(
         quantiles=quantiles,
         points=points,
         bw_method=bw_method,
+        side=side,
         **({"data": data} if data is not None else {}),
     )
 
diff --git a/lib/matplotlib/tests/baseline_images/test_axes/violinplot_sides.png b/lib/matplotlib/tests/baseline_images/test_axes/violinplot_sides.png
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -3794,6 +3794,21 @@ def test_horiz_violinplot_custompoints_200():
                   showextrema=False, showmedians=False, points=200)
 
 
+@image_comparison(['violinplot_sides.png'], remove_text=True, style='mpl20')
+def test_violinplot_sides():
+    ax = plt.axes()
+    np.random.seed(19680801)
+    data = [np.random.normal(size=100)]
+    # Check horizontal violinplot
+    for pos, side in zip([0, -0.5, 0.5], ['both', 'low', 'high']):
+        ax.violinplot(data, positions=[pos], vert=False, showmeans=False,
+                      showextrema=True, showmedians=True, side=side)
+    # Check vertical violinplot
+    for pos, side in zip([4, 3.5, 4.5], ['both', 'low', 'high']):
+        ax.violinplot(data, positions=[pos], vert=True, showmeans=False,
+                      showextrema=True, showmedians=True, side=side)
+
+
 def test_violinplot_bad_positions():
     ax = plt.axes()
     # First 9 digits of frac(sqrt(47))
