Python: Coloring 1-D swarmplots in seaborn

Question

I am trying to make a 1-D swarm plot using the below dataframe(only first few rows are shown):

      value    color
0   0.446928    1
1   0.258523    0
2   0.716512    2
3   0.288698    0
4   0.132203    0
5   0.871158    3
6   0.613292    2
7   0.697033    2
8   0.333995    1
9   0.549433    2

Using the below code I made a 1D swarm plot:

import matplotlib.pyplot as plt
import seaborn as sns

plt.style.use('ggplot')
sns.swarmplot(x=df['value'], hue=df['color'])

Here is the plot generated:

The problem here is I want to define colours to each group like value from 0 to 0.3 should have red colour(say), 0.3 to 0.5 should have the second colour(say green) and so on. For that, I have defined the colour column in the dataframe but on passing hue=df['color'], the plot is not generating colours. I want my plot to look something like this:

What changes should I make in my code so that my program works as expected?

user_12046409 · Answer

The generated scatter points are stored in ax.collections[-1]. You can extract their x-coordinates, use np.digitize() to divide them into groups and use those values for coloring:

import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
import seaborn as sns
import pandas as pd
import numpy as np

plt.style.use('ggplot')
df = pd.DataFrame({'value': np.random.uniform(0, 1, 1000)})
ax = sns.swarmplot(x=df['value'])
scatter_dots = ax.collections[-1]
xpos = scatter_dots.get_offsets()[:, 0]
boundaries = [0, .3, .5, .7, 1]  # one more than the number of colors
colors = np.digitize(xpos, boundaries)
scatter_dots.set_array(colors)
scatter_dots.set_cmap(ListedColormap(['crimson', 'limegreen', 'gold', 'skyblue']))
plt.show()

Python: Coloring 1-D swarmplots in seaborn

1 Answers