Dealing with zeros in numpy array normalization

Question

I have a numpy array of 2D vectors, which I am trying to normalize as below. The array can have vectors with magnitude zero.

x = np.array([[0.0, 0.0], [1.0, 0.0]])
norms = np.array([np.linalg.norm(a) for a in x])

>>> x/norms
array([[ nan,   0.],
       [ inf,   0.]])

>>> nonzero = norms > 0.0
>>> nonzero
array([False,  True], dtype=bool)

Can I somehow use nonzero to apply the division only to x[i] such that nonzero[i] is True ? (I can write a loop for this - just wondering if there's a numpy way of doing this)

Or is there a better way of normalizing the array of vectors, skipping all zero vectors in the process?

Answer 1

If you can do the normalization in place, you can use your boolean indexing array like this:

nonzero = norms > 0
x[nonzero] /= norms[nonzero]

Answer 2

Here's one possible way of doing this

norms = np.sqrt((x**2).sum(axis=1,keepdims=True))
x[:] = np.where(norms!=0,x/norms,0.)

This uses np.where to do the substitution you need.

Note: in this case x is modified in place.

Answer 3

It's probably easiest just to do the calculation and then modify the results to be what you'd like:

y = x/norms
y[np.isnan(y) | np.isinf(y)]=0

#y = array([[ 0.,  0.],
#       [ 0.,  0.]])