I have been playing around with numpy and matplotlib.

My little project was to create a scatter plot ranging from -1 to 1 on both X and Y, but where the shading is done with the XOR scheme.

The following is the code that I have implemented (it is in a cell of a jupyter notebook, hence the trailing ; after plt.scatter):

arr_size = 2000

X = np.random.uniform(low=-1.0, high=1.0, size=(arr_size, arr_size))

colour = np.zeros(arr_size)

for i in range(arr_size):
    if X[0, i] > 0 and X[1, i] < 0:
        colour[i] = 1
    elif X[0, i] < 0 and X[1, i] > 0:
        colour[i] = 1

plt.scatter(X[0], X[1], c=colour);

The output that I generated was:

enter image description here

Which is the desired output.

However, I am on a bit of a campaign to make my numpy code run faster (I am on an ML course and we have been taught to remove for loops wherever possible). Can anyone show me how to make this code more efficient?



Since (a < 0 and b > 0) or (a > 0 and b < 0) can be summarized as a*b < 0, moreover * and < work on vectors, we get the one-liner:

plt.scatter(X[0], X[1], c=X[0]*X[1] < 0)

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.