I am trying to create a custom distribution in scipy (COM-Poisson).

import numpy as np
from mpmath import nsum, fac, inf
from scipy import stats


def pmf(k, mu, nu):
    return float((mu ** k) / ((fac(k) ** nu) * nsum(lambda j: (mu ** j) / (fac(j) ** nu), [0, inf])))


vectorized_pmf = np.vectorize(pmf)


class com_poisson_gen(stats.rv_discrete):
    def _argcheck(self, mu, nu):
        return mu >= 0 & nu >= 0

    def _pmf(self, k, mu, nu):
        return vectorized_pmf(k, mu, nu)


com_poisson = com_poisson_gen(name="com-poisson", longname='Conway-Maxwell-Poisson')
com_poisson.rvs(2, 4, size=100000)

I would like to cache the normalising constant so it is calculated only once, how can I achieve this?

CodePudding user response：

A suggestion of an lru_cahce decorator from the guys at scipy.stats helped solve this.

from functools import lru_cache

import numpy as np
from mpmath import nsum, fac
from numpy import inf
from scipy import stats
from scipy.special import factorial


class com_poisson_gen(stats.rv_discrete):
    @lru_cache(maxsize=None)
    def _normalization_factor(self, mu, nu):
        return float(nsum(lambda j: (mu ** j) / (fac(j) ** nu), [0, inf]))

    def _argcheck(self, mu, nu):
        return (np.asarray(mu) > 0) & (np.asarray(nu) > 0)

    def _pmf(self, k, mu, nu):
        broadcast = np.broadcast(mu, nu)
        normalization_factor = np.empty(broadcast.shape)
        normalization_factor.flat = [self._normalization_factor(x, y) for x, y in broadcast]
        return (mu ** k) / ((factorial(k) ** nu) * normalization_factor)

com_poisson = com_poisson_gen()