Pythonic way to calculate cumulative sum with complex math over numpy array

Viewed 182

I'm performing Data Science and am calculating the Log Likelihood of a Poisson Distribution of arrival times.

def LogLikelihood(arrival_times, _lambda):
  """Calculate the likelihood that _lambda predicts the arrival_times well."""
  ll = 0
  for t in arrival_times:
    ll += -(_lambda) + t*np.log(_lambda) - np.log(factorial(t))
  return ll

Mathematically, the expression is on the last line:

Text

Is there a more pythonic way to perform this sum? Perhaps in one line?

3 Answers

Seems perfectly Pythonic to me; but since numpy is already here, why not to vectorize the whole thing?

return (
    -_lambda
    + arrival_times * np.log(_lambda)
    - np.log(np.vectorize(np.math.factorial)(arrival_times))
).sum()

If you have scipy available, use loggamma which is more robust than chaining log and factorial:

from scipy import special

def loglikeli(at,l):
    return (np.log(l)*at-l-special.loggamma(at+1)).sum()

### example 
rng = np.random.default_rng()
at = rng.integers(1,3,10).cumsum()
l = rng.uniform(0,1)

### check against OP's impementation
np.isclose(loglikeli(at,l),LogLikelihood(at,l))
# True

This looks ugly to me but it fits a single line:

def LogLikelihood(arrival_times, _lambda):
    return np.cumsum(list(map(lambda t: -(_lambda) + t*np.log(_lambda) - np.log(factorial(t)),arrival_times)))[-1]
Related