I've created a scipy.stats.rv_continuous
subclass, and it seems to be doing what I want, but it's extremely slow. Code and test results below.
The distribution functions I'm using (broken power-law) are easy to integrate and calculate properties, so is there another internal method which I should subclass with analytic values to make it faster? The documentation is unclear on how the rvs
are actually drawn, but presumably it's finding the inverse of the cdf
.
class Broken_Power_Law(sp.stats.rv_continuous):
def __init__(self, slopes, breaks, name='Broken_Power_Law'):
super().__init__(a=np.min(breaks), b=np.max(breaks), name=name)
nums = len(slopes)
pdf_norms = np.array([np.power(breaks[ii], slopes[ii-1] - slopes[ii]) if ii > 0 else 1.0
for ii in range(nums)])
pdf_norms = np.cumprod(pdf_norms)
cdf_offsets = np.array([(an/(alp+1))*(np.power(breaks[ii+1], alp+1) -
np.power(breaks[ii], alp+1))
for ii, (alp, an) in enumerate(zip(slopes, pdf_norms))])
off_sum = cdf_offsets.sum()
cdf_offsets = np.cumsum(cdf_offsets)
pdf_norms /= off_sum
cdf_offsets /= off_sum
self.breaks = breaks
self.slopes = slopes
self.pdf_norms = pdf_norms
self.cdf_offsets = cdf_offsets
self.num_segments = nums
return
def _pdf(self, xx):
mm = np.atleast_1d(xx)
yy = np.zeros_like(mm)
for ii in range(self.num_segments):
idx = (self.breaks[ii] < mm) & (mm <= self.breaks[ii+1])
aa = self.slopes[ii]
an = self.pdf_norms[ii]
yy[idx] = an * np.power(mm[idx], aa)
return yy
def _cdf(self, xx):
mm = np.atleast_1d(xx)
yy = np.zeros_like(mm)
off = 0.0
for ii in range(self.num_segments):
off = self.cdf_offsets[ii-1] if ii > 0 else 0.0
idx = (self.breaks[ii] < mm) & (mm <= self.breaks[ii+1])
aa = self.slopes[ii]
an = self.pdf_norms[ii]
ap1 = aa + 1
yy[idx] = (an/(ap1)) * (np.power(mm[idx], ap1) - np.power(self.breaks[ii], ap1)) + off
return yy
When testing it:
> test1 = sp.stats.norm()
> %timeit rvs = test1.rvs(size=100)
46.3 µs ± 1.87 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
> test2 = Broken_Power_Law([-1.3, -2.2, -2.7], [0.08, 0.5, 1.0, 150.0])
> %timeit rvs = test2.rvs(size=100)
200 ms ± 8.57 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
i.e. 5000x slower!!!
Aucun commentaire:
Enregistrer un commentaire