CODEWITHSUNDEEP
pythonnumpymatplotlibplot
GiorgosPap31
GiorgosPap31

Reputation: 43

Scatter plot with specifc conditions

Suppose that I have this data set

x1 = np.array([0.1,0.3,0.1,0.6,0.4,0.6,0.5,0.9,0.4,0.7])
x2 = np.array([0.1,0.4,0.5,0.9,0.2,0.3,0.6,0.2,0.4,0.6])
c  = np.array([  1,  1,  1,  1,  1,  0,  0,  0,  0,  0])

I want a scatter plot of x1,x2 under the following condition: if the corresponding index for (x1[i],x2[i]) in c is c[i]==1 then plot this a marker as a red X , but if the corresponding index for (x1[i], x2[i]) in c is c[i]==0 then plot the marker at (x1[i], x2[i]) as a blue O.

Any idea on how to do it?

Upvotes: 0

Views: 1305

Answers (2)

sivic
sivic

Reputation: 529

Since you're using numpy, it is rather easy to get elements (or indices) where certain rule (comparison) is satisfied.

You can get an array of boolean values which corresponds to some condition acted on elements of existing array:

print(c==0)
[False False False False False  True  True  True  True  True]

Also, you can access only some elements of array using the array of booleans (same as the size of initial array):

print(x1[c==1])
[0.1 0.3 0.1 0.6 0.4]

or more complex operations and/or set values:

x1[x2 < 0.5] = 0
print(x1)
x2[x2 < 0.5] = 10
print(x2)

[0.  0.  0.1 0.6 0.  0.  0.5 0.  0.  0.7]
[10.  10.   0.5  0.9 10.  10.   0.6 10.  10.   0.6]

For even more complicated conditions, there are logical functions for NumPy arrays.

This approach (NumPy) is significantly faster than using loops and you should utilize it whenever it is possible!

Applying above, it is easy to solve your problem:

import numpy as np
import matplotlib.pyplot as plt

x1 = np.array([0.1,0.3,0.1,0.6,0.4,0.6,0.5,0.9,0.4,0.7])
x2 = np.array([0.1,0.4,0.5,0.9,0.2,0.3,0.6,0.2,0.4,0.6])
c=np.array([ 1,1,1,1,1,0,0,0,0,0 ])

plt.plot(x1[c==0], x2[c==0], 'bo')
plt.plot(x1[c==1], x2[c==1], 'rx')

enter image description here

Upvotes: 1

Rajan
Rajan

Reputation: 757

I don't know it is efficient way or not but you can do this,

i = 0
for (p1,p2) in zip(x1,x2):
    if(c[i]==1): 
        plt.scatter(p1, p2, marker='x', color='red')
    else: 
        plt.scatter(p1, p2, marker='o', color='blue')
    i+=1

check other marker shapes here,

https://matplotlib.org/3.1.1/api/markers_api.html

Upvotes: 0

Related Questions