logsumexp with sign indicator - enable calculation with negative signs

[ingmarschuster]

This is a repost of numpy/numpy#5652 to where a logsumexp request actually belongs.

So my problem has been the following: I had to calculate the logsumexp of some numbers in logspace for accuracy reasons, but some of the elements of the array have negative sign. I implemented the following code, which you might want to integrate into the upstream logsumexp implementation of numpy.
The sign argument is a matrix assumed to be filled with +1 or -1 depending on the sign of the number before the log transformation.

from numpy import exp, log, sqrt
from scipy.misc import logsumexp
import numpy as np
import scipy.stats as stats


def log_sign(a):
    a = np.array(a)
    sign_indicator = ((a < 0 ) * -2 + 1)
    return (log(np.abs(a)), sign_indicator)

def exp_sign(a, sign_indicator):
    return exp(a) * sign_indicator    

def logsumexp_sign(a, axis = None, b = None, sign = None):
    if sign is None:
        abs_res = logsumexp(a, axis = axis, b = b)
        sign_res = np.ones_like(abs_res)
        return (abs_res, sign_res)
    else:
        if not (np.abs(sign) == 1).all():
            raise ValueError("sign arguments expected to contain only +1 or -1 elements")

        m = np.copy(a)
        m[sign == -1] = -np.infty # don't look at negative numbers
        m = logsumexp(m, axis = axis, b = b)
        m[np.isnan(m)] = -np.infty # replace NaNs resulting from summing infty

        s = np.copy(a)
        s[sign == +1] = -np.infty # don't look at positive numbers
        s = logsumexp(s, axis = axis, b = b)
        s[np.isnan(s)] = -np.infty # replace NaNs resulting from summing infty


        sign_res =  np.ones(np.broadcast(m, s).shape) #default to positive sign
        abs_res  = -np.infty * sign_res #default to log(0)

        idx = np.where(m > s)
        abs_res[idx] = log(1 - exp(s[idx] - m[idx])) + m[idx]

        idx = np.where(m < s)
        sign_res[idx] = -1
        abs_res[idx] = log(1 - exp(m[idx] - s[idx])) + s[idx]
        return (abs_res, sign_res)

I'm not sure I passed on the b parameter correctly (I never made use of it before). Here is a test method:

def test_logsumexp_sign():    
    for ax in (0, 1):
        for a in [np.array([[1,1,-1], [-1,-1,1]]),
                  np.array([[1,1,1], [-1,-1,-1]])]:
            (la, sa) = log_sign(a)
            log_res = logsumexp_sign(la, axis = ax, sign = sa)
            res = exp_sign(*log_res)
            if not np.all(np.abs(res - np.sum(a, axis = ax)).mean() <= 0.1):
                raise RuntimeError("problem when summing accros", ax,
                                   "for", repr(a), "getting", log_res)

[aarchiba]

The b array almost allows you to do this: you can just put your sign indicators in there. The only thing lacking is that the result will be NaN if the answer comes out negative; so a simple implementation could just detect that, flip the sign on all the b values, and rerun the function. You'd still want a wrapper that returned sign information.

[ingmarschuster]

Ok, that would be great. Would an implementetion detecting NaN results and flipping signs be faster than mine? Also, yes, you'd want a wrapper returning sign information and you'd want this use of the b parameter to be very transparent in the documentation (transparent in the code is nicer of course, which should be easy when implementing this as a wrapper of course).

[ev-br]

Has been fixed by gh-4859, so closing.

[ingmarschuster]

Wonderful to see this effect. Next time I'll not be bound to PhD work, so I might help beyond filing issues. Thanks to Anne!