CODEWITHSUNDEEP
pythonmatplotlibplot
Thomas Kimber
Thomas Kimber

Reputation: 11067

Display sections of a matplotlib violinplot in different colours (colors) either side of a threshold value

I have some code that expresses the balance of an election result in the form of a violin plot:

plt.plot((0, 0), (0.85, 1.15), 'k-')
p = plt.violinplot(np.array(results['Balance']),vert=False, bw_method='scott', widths=0.2,showextrema=False)

This returns a plot that looks like the following:

Violin Plot

I would like to colour this plot either side of the "Decision Line" to reflect voting intention. Something like the mock-up shown below:

enter image description here

I've tried plotting the two distinct sets independently of one another, e.g.

p = plt.violinplot(np.array(results[results['Balance']>=0]['Balance']),vert=False, bw_method='scott', widths=0.2,showextrema=False)
n = plt.violinplot(np.array(results[results['Balance']<0]['Balance']),vert=False, bw_method='scott', widths=0.2,showextrema=False)

I might then use the method discussed here to colour each PolyCollection differently.

But the shape returned no longer reflects the original distribution, and so is too far from what I'm looking for to be helpful in this instance.

skewed distribution

Does anyone have any ideas or techniques for achieving something closer to my coloured mockup?

Upvotes: 1

Views: 1029

Answers (1)

ImportanceOfBeingErnest
ImportanceOfBeingErnest

Reputation: 339340

The problem is that the violinplot will compute a different violin if the data is cut. So One would need to work with the same violin for both sides of the separating line.

An idea could be to use the path of the violin, which can be obtained via path = p['bodies'][0].get_paths()[0], to cut out part of two differently-colored, plot-filling rectangles on either side of the separating line.

import matplotlib.pyplot as plt
import matplotlib.patches
import numpy as np
import pandas as pd

#generate some data
a = np.log(1+np.random.poisson(size=1500))
b = 0.2+np.random.rand(1500)*0.3
c = 0.1+np.random.rand(1500)*0.6
results=pd.DataFrame({'Balance':np.r_[a*b,-a*c]})

#create figure and axes
fig, ax=plt.subplots()
# sep is the point where the separation should occur
sep = -0.05
plt.plot((sep, sep), (0.85, 1.15), 'k-')
# plot the violin
p = plt.violinplot(np.array(results['Balance']),vert=False,
                   bw_method='scott', widths=0.2,showextrema=False)
# obtain path of violin surrounding
path = p['bodies'][0].get_paths()[0]

#create two rectangles left and right of the separation line
r =  matplotlib.patches.Rectangle((results['Balance'].min(),0.85), 
                  width=sep-results['Balance'].min(), height=0.3, facecolor="r")
r2 =  matplotlib.patches.Rectangle((sep,0.85), 
                  width=results['Balance'].max()-sep, height=0.3, facecolor="b")
# clip the rectangles with the violin path
r.set_clip_path(path, transform=ax.transData)
r2.set_clip_path(path, transform=ax.transData)
ax.add_patch(r)
ax.add_patch(r2)

#optionally add edge around violin.
s = matplotlib.patches.PathPatch(path, linewidth=1, edgecolor="k", fill=False)
ax.add_patch(s)
plt.show()

enter image description here

Upvotes: 3

Related Questions