How to add a geom_point layer based on a condition in python

Question

I am using python package plotnine based in R's ggplot and I am trying to create a plot where I want a line to be plotted, and then, I want to add a geom_point layer that plots only some points if a given criteria is verified.

So I have a dataframe with a variable that contains real numbers (that is the one that I want to plot as a line) and then a binary variable. I would like to plot points with the value of the first variable if the second variable has a value of 1, and nothing if the second variable has a value of 0.

I provide a small dataframe as an example

import pandas as pd
from plotnine import ggplot, aes, geom_line, geom_point
d = dict(
    a = np.arange(10),
    b = [0, 0, 0, 1, 0, 0, 1, 0, 1, 0]
)
d = pd.DataFrame(d)

(ggplot(df, aes(x='a')) + 
 geom_line(aes(y='a'), color='black')
 # Some code to plot points with value a[i] if b[i]==1
 )

The only way of doing this that I could think of was to create a secondary dataframe with the values of a if b=1 and None if b=0, but this feels a bit clumsy

Answer 1

You have to use a dataframe with the selected data that you want in the point layer.

(ggplot(df, aes(x='a'))
 + geom_line(aes(y='a'), color='black')
 + geom_point(df[df['b'] == 1], aes(y='a')) # or geom_point(df.query('b == 1'), aes(y='a'))
 )