ENH: Add Multivariate Hypergeometric Distribution

[mdhaber]

Part of "A Solid Foundation for Statistics in Python with SciPy".

Add Multivariate Hypergeometric Distribution to scipy.stats.

Be sure to follow the guide Adding a New Statistics Distribution carefully and see gh-6375 for an example of adding a related distribution. See methods inherited from multi_rv_generic in _multivariate.py, although I guess there's not too much and most of it needs to be implemented from scratch.

[mdhaber]

@tirthasheshpatel if you're up for another challenge after finishing gh-12170, I'd be happy to help with this, too!

[tirthasheshpatel]

Sure! I will try to work on this, this week!

[tirthasheshpatel]

@mdhaber @WarrenWeckesser I am trying to implement this distribution and almost done. I am facing a problem with sampling from the distribution. @WarrenWeckesser implemented sampling methods for numpy in np.random.Generator but scipy still uses the old np.random.RandomState. Is scipy going to continue with that or is there any plan to upgrade to np.random.Generator? Just in case scipy continues to use np.random.RandomState, I tried to implement a count method in python for sampling from this distribution (but is not very efficient)! I ended up with this:

    def rvs(self, m, n, size=1, random_state=None):
        """
        m : array_like
            Number of each type of objects in the population.
        n : array_like
            Number of samples in each draw from the population.
        """
        m, n = self._process_parameters(m, n)

        random_state = self._get_random_state(random_state)
        try:
            # Generator is only available in numpy >= 1.17
            if isinstance(random_state, np.random.Generator):
                out = random_state.multivariate_hypergeometric(m, n, size)
                return _squeeze_output(out)
        except AttributeError:
            pass

        # NumPy's count method for sampling from the MHG.
        # We first broadcast the arrays `m` and `n`.
        # Broadcasted array has some shape `(..., dims)`
        # where `dims` is the dimension of the MHG.
        # We reshape that array to get a well-shaped array
        # of shape `(num_examples, dims)` and iterate over
        # each example and sample `size` number of variates.
        # Finally, Reshape the  array `rv` to retain the
        # original shape.
        def _mhg_rvs_count(m, n):
            rv = np.empty((size, len(m)), dtype=np.int_)
            c = np.repeat(np.arange(len(m)), m)
            for i in range(size):
                sel = random_state.choice(c, size=n,
                                          replace=False)
                rv[i] = np.bincount(sel, minlength=len(m))
            return rv

        m, n = np.broadcast_arrays(m, n)
        oshape = m.shape
        m, n = m.reshape(-1, m.shape[-1]), n[...,0].ravel()

        rv = np.empty((m.shape[0], size, m.shape[1]),
                      dtype=np.int_)
        for i in range(m.shape[0]):
            rv[i] = _mhg_rvs_count(m[i], n[i])

        return rv.reshape((size,) + oshape)

@WarrenWeckesser Any thought on how to make it more efficient with arbitrarily shaped arrays? If this is not the appropriate place for discussion, I will open a draft pr where we can continue our discussion!

[andyfaff]

You need to write rvs as if random_state is either a Generator or RandomState. All of the rv_continuous classes can now use either of them.

[tirthasheshpatel]

Oh, if I understand correctly, I can assume the random_state is a Generator object and implement the rvs method, right @andyfaff ?

[andyfaff]

No, it'll either be RandomState or Generator, so your code needs to work with both.

[tirthasheshpatel]

Then, because, RandomState doesn't have a sampler for the multivariate hypergeometric distribution, it needs to implemented manually using the count or marginals method... I tried implementing a count method for sampling which works with arbitrarily shaped arrays. I was asking if there is a more efficient method than the one I have implemented...

[mdhaber]

@tirthasheshpatel I would go ahead and open the PR. I think we can work on the rest of the code while you wait for feedback on this.

[WarrenWeckesser]

It would probably be safer to implement the marginals method instead of the count method. With large input values, the count method might allocate a huge temporary array.

[mdhaber]

BTW @tirthasheshpatel while I'd suggest following @WarrenWeckesser's suggestion, don't worry about speed too much here. If it is not fast enough, we have permission to add these random number generators under the BSD license, and we'll already be laying the groundwork for the noncentral hypergeometric distributions. It would be easy to also add the multinomial random number generator.

[tirthasheshpatel]

Sure @mdhaber. Didn't get some time from a couple of weeks to work on this but I am almost done with the code. I will open a PR as soon as the code is properly working! About random generation, yes, I still have just implemented the count method but I will look into the ones you have linked to. BTW, how do I add rngs from the c++ library you have linked? Do I just put the necessary files (with license) under some directory in stats module (and extend it for python) or do I just make a python port of the c++ code?