CODEWITHSUNDEEP
pythonnumpy
Ulysses
Ulysses

Reputation: 357

Elementwise division, disregarding zeros

It's a python question: let's say I have an m+1-dimensional numpy array a consisting of non-negative numbers, and I would like to obtain an array b of the same size where the last coordinates are normalized so that they sum up to 1, or zero in case all of them were zeros. For example, if m = 2, my code would be as follows

import numpy as np

a = np.array([[[ 0.34      ,  0.66],
               [ 0.75      ,  0.25]],
              [[ 0.        ,  0.  ],
               [ 1.        ,  0.  ]]])

for i1 in range(len(a)):
    for i2 in range(len(a)):
        s = a[i1][i2].sum()
        if s > 0:
            a[i1][i2] = a[i1][i2]/s

however I find this method sloppy. Also, it works only for fixed m.

Upvotes: 1

Views: 101

Answers (1)

eickenberg
eickenberg

Reputation: 14377

This can be done by broadcasting. There are several ways to take into account the zero-sum exception. Without taking it into account, you could write

import numpy as np
shape = (2, 3, 4)
X = np.random.randn(*shape) ** 2
sums = X.sum(-1)

Y = X / sums[..., np.newaxis]

Now, in order to take into account potential zero-sum-ness of some lines, we set one line of the data to 0:

X[0, 0, :] = 0

sums = X.sum(-1)

nnz = sums != 0
Y = np.zeros_like(X)
Y[nnz, :] = X[nnz, :] / sums[nnz, np.newaxis]

You will observe that Y.sum(axis=-1) has the entry 0 in coordinate (0,0) reflecting the zero-ness of the corresponding line.

EDIT: Application to the concrete example

X = np.array(array([[[ 0.34      ,  0.66],
                     [ 0.75      ,  0.25]],
                    [[ 0.        ,  0.  ],
                     [ 1.        ,  0.  ]]]))

sums = X.sum(-1)
nnz = sums != 0
Y = np.zeros_like(X)
Y[nnz, :] = X[nnz, :] / sums[nnz, np.newaxis]

yields Y == X (because along the last axis the sum is already one or zero.)

Upvotes: 3

Related Questions