Reputation: 17170
I have a Python code containing the following lines:
# Poisson model
lambda_param = 3.0
x_model = numpy.linspace(0, 10, num=11, dtype=int)
y_model = numpy.zeros(10)
index = 0
for x in x_model:
y_model[index] = numpy.power(lambda_param, x) / math.factorial(x) * numpy.exp(-lambda_param)
index += 1
I want to write this in a more Pythonic way - without using a for loop.
I tried this:
y_model = numpy.power(lambda_param, x_model) / math.factorial(x_model) * numpy.exp(-lambda_param)
But this didn't seem to work, I would guess because math.factorial
doesn't know how to operate on an array like object.
Is there a way to do this?
Upvotes: -1
Views: 89
Reputation: 79288
This is the pmf
for a poisson distribution. Use scipy.stats.poisson
:
Note that for poisson, x
must be integers:
lambda_param = 3.0
x_model = np.arange(10)
y_model = np.zeros(10)
from scipy import stats
stats.poisson(lambda_param).pmf(x_model)
array([0.04978707, 0.14936121, 0.22404181, 0.22404181, 0.16803136,
0.10081881, 0.05040941, 0.02160403, 0.00810151, 0.0027005 ])
Upvotes: 2
Reputation: 14434
You can use scipy.special.factorial
as a factorial function that supports broadcasting over numpy arrays.
import numpy as np
import scipy.special
>>> scipy.special.factorial(np.arange(7))
array([ 1., 1., 2., 6., 24., 120., 720.])
Upvotes: 1