numpy: Efficiently avoid 0s when taking log(matrix)

Question

from numpy import *

m = array([[1,0],
           [2,3]])

I would like to compute the element-wise log2(m), but only in the places where m is not 0. In those places, I would like to have 0 as a result.

I am now fighting against:

RuntimeWarning: divide by zero encountered in log2

---

Try 1: using where

res = where(m != 0, log2(m), 0)

which computes me the correct result, but I still get logged a RuntimeWarning: divide by zero encountered in log2. It looks like (and syntactically it is quite obvious) numpy still computes log2(m) on the full matrix and only afterwards where picks the values to keep.

I would like to avoid this warning.

---

Try 2: using masks

from numpy import ma

res = ma.filled(log2(ma.masked_equal(m, 0)), 0)

Sure masking away the zeros will prevent log2 to get applied to them, won't it? Unfortunately not: We still get RuntimeWarning: divide by zero encountered in log2.

Even though the matrix is masked, log2 still seems to be applied to every element.

---

How can I efficiently compute the element-wise log of a numpy array without getting division-by-zero warnings?

• Of course I could temporarily disable the logging of these warnings using seterr, but that doesn't look like a clean solution.
• And sure a double for loop would help with treating 0s specially, but defeats the efficiency of numpy.

Any ideas?

Answer 1

We can use masked arrays for this:

>>> from numpy import *
>>> m = array([[1,0], [2,3]])
>>> x = ma.log(m)
>>> print x.filled(0)
[[ 0.          0.        ]
 [ 0.69314718  1.09861229]]