ENH: DataFrame.plot.scatter argument c now accepts a column of strings, where rows with the same string are colored identically

[michaelmannino]

• Closes Scatter plot with colour_by and size_by variables #16827
• Closes Meaning of color argument in DataFrame.plot.scatter() #16485
• Tests added and passed if fixing a bug or adding a new feature
• All code checks passed.
• Added type annotations to new arguments/methods/functions.
• Added an entry in the latest doc/source/whatsnew/vX.X.X.rst file if fixing a bug or adding a new feature.

[michaelmannino]

import pandas as pd
import random as rand

df = pd.DataFrame({
    'dataX': [rand.randint(0, 100) for _ in range(100)],
    'dataY': [rand.randint(0, 100) for _ in range(100)],
    'fav_fruit': [rand.choice(['Apples', 'Bananas', 'Grapes', 'Peaches']) for _ in range(100)],
})
df.plot.scatter('dataX', 'dataY', c='fav_fruit')

[WillAyd]

@Dr-Irv in case you have any feedback as well

[WillAyd]

In what instances is c_values a list? Might be misreading but would be better if we only worked with a pd.Series and could call .unique on that, instead of checking every single value in a loop

[michaelmannino]

Updated to take a pd.Series, not np.ndarray | list

[WillAyd]

I am no matplotlib expert but I think we need to defer to that somehow to get the desired colors, instead of trying to write this out ourselves

[michaelmannino]

I did some research looking into how other libs which support this functionality, including seaborn, handle this workflow. Each utilize the current mpl.colormap and draw colors from a linear space across this map. This allows users to change the choice colors of the same way they would with all other graphs

Additionally, I've added an automatic legend to this type of plot since the chosen colors are not exposed to the user, similar to how a colorbar is drawn in some cases of the same function

[Dr-Irv]

I'm not 100% sure that the test here should be in the tm.assert_produces_warning(None) context. My intuition is that should be removed.

There should also be tests where the column state contains values that are colors, as well as a mix of colors and non-colors.

[michaelmannino]

I removed the tm.assert_produces_warning(None). I added 2 test cases: 1 containing only strings and 1 containing both valid mpl strings and invalid mpl strings.

Additionally, I updated some other functionality tests to confirm they are not executing the new logic added and coloring properly

[michaelmannino]

df = pd.DataFrame({
    'dataX': [rand.randint(0, 100) for _ in range(100)],
    'dataY': [rand.randint(0, 100) for _ in range(100)],
    'fav_fruit': [rand.choice(['Grapes', 'Bananas', 'Apple', 'Peaches']) for _ in range(100)],
})

df.plot.scatter('dataX', 'dataY', c='fav_fruit')

df = pd.DataFrame({
    'dataX': [rand.randint(0, 100) for _ in range(100)],
    'dataY': [rand.randint(0, 100) for _ in range(100)],
    'fav_fruit': [rand.choice(['Grapes', 'Bananas', 'Apple', 'Peaches']) for _ in range(100)],
})

df.plot.scatter('dataX', 'dataY', c='fav_fruit', cmap='rainbow')

[michaelmannino]

Not sure if I need to add this as part of this PR, but there is some documentation (and maybe more) about using the c arg in df.plot.scatter found in: https://github.com/pandas-dev/pandas/blob/main/doc/source/user_guide/visualization.rst?plain=1#L636-L637

Should I update this to include this new feature?

[Dr-Irv]

I'm OK with this. Will leave it to @WillAyd to decide/merge

[WillAyd]

Thanks @michaelmannino - very impressive work! Thanks for all the efforts on the refactors / feedback