pythonpandasmatplotlibplot

Conditional marker for scatterplot Matplotlib


I am trying to plot x1 and x2 against each other, and if class == 1, then marker should be +, if class == 0, then the marker should just be a dot.

The file.csv is structured like this, the class column which is the conditional will only be either 1 or 0:

x1 x2 mark
1 2 0
9 4 1
0 5 1
2 6 0

Here's the code I have:

df = pd.read_csv('file.csv')
print(df)
x1 = df['x1'].to_numpy()
x2 = df['x2'].to_numpy()
mark = df['class'].to_numpy()
        
figure, ax = plt.subplots()

for i in range(0,80,1):
    if mark[i] > 0:
        ax.plot(x1[i], color="red", marker='o')
    else:
        ax.plot(x1[i], color="blue", marker='x')

This is what the result SHOULD look like: result


Solution

  • I would plot it this way for all the data, if you're going to select a subset, consider using a new dataframe to select the subset:

    import matplotlib.pyplot as plt
    
    # the rows where class is 1 and plot them with '+' marker
    plt.scatter(df[df['class']==1]['x1'], df[df['class']==1]['x2'], marker='+')
    
    # the rows where class is 0 and plot them with 'o' marker 
    plt.scatter(df[df['class']==0]['x1'], df[df['class']==0]['x2'], marker='o')
    
    plt.xlabel('x1') 
    plt.ylabel('x2') 
    plt.title('Scatter plot x1 vs x2') 
    plt.show()