mardi 21 novembre 2017

speedup scipy custom continuous random variable

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