Removing Loop from Numpy XOR

Question

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?

Cheers

Answer 1

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)