Search code examples
pythonnumpyscientific-computing

Collapse Numpy arrays to scalar when, e.g., multiplying by zero


I'm a bit surprised that I couldn't find an easy when to collapse dimensions of a Numpy array containing identical values. Let me explain.

I might have to multiply to time series implemented as arrays, say a * b. In most cases, this is fine, but sometimes a or b are scalars representing constant signals. This works also smoothly because Numpy knows how to broadcast them, and computes this very quickly (twice as fast as if b were an array full of identical values).

Now sometimes a = 0; and as a result, I get an array full of zeros. I want to collapse it to the scalar 0 because it is a constant signal, but I can't seem to find any easy way to do it -- and I mean, without adding a condition to check whether a or b is 0 and treat it as a special case every singly time I operate on my arrays.

Do you know of any simple ways to achieve this?


Solution

  • The test case is your best bet. Numpy does not generally presume to know the contents of your data, so it only does what you tell it. Specifically, as you noted, the result of a broadcasted operation will be whatever the broadcasting comes out to. Anything short of that would voilate Python's principle of least surprise.

    Here is what I would argue is the simplest way to do it:

    c = a * b
    if not c.any():
        c = c[0].item()
    

    There are other alternatives. The first that comes to mind is using (c == 0).all(), but that involves a second temporary array, wastes memory and effectively undoes the benefit of short-circuiting.

    Another alternative is to check if either a or b are zeros up front. But to do that thoroughly, you would need to check for things like this:

    a = [1, 0, 2, 0, 3, 0]
    b = [0, 1, 0, 2, 0, 3]
    

    So simply multiplying the two arrays and inspecting afterwards is likely to be your best bet.