Repository: scipy/scipy
Pull request: stats/distributions: make rv_sample public, allow subclassing

#### Reference issue
<!--Example: Closes gh-WXYZ.-->

closes gh-8057

#### What does this implement/fix?
<!--Please explain your changes.-->

Expose `rv_sample` to the main stats namespace. Make it play ball with subclassing: previously, subclassing `rv_discrete(values=(..., ...))` was broken. With this PR, subclassing `rv_sample` just works.

#### Additional information
<!--Any additional information you think is important.-->

After several failed attempts to fix the `__new__` dispatch in `rv_discrete` I'm throwing a towel on that. The dispatch to a subclass is a hack which apparently breaks all sorts of invariants. Which is, after all, expected --- this was originally devised as a hack around the `rv_discrete` quirk which was doing very different things in its `__init__`, depending on whether `values` are provided or not.

So, this PR instead introduces a new `rv_sample` subclass of `rv_discrete`. It is constructed with `xk, pk` arguments directly --- thus sidestepping the dispatch on `values=...`. I think it actually makes the signature lighter and more reasonable, too :-).

Deprecating `rv_discrete(values='...)` is out of scope of this PR.
Whether to actually deprecate it or not in a follow-up PR depends a bit on the timeline for the plans for revamping the distributions infrastructure. If it's planned in a finite amount of time from now, maybe there is no point causing downstream churn.


Problem statement:
Subclass of scipy.stats.rv_discrete initiated with values throws an AttributeError on pmf method
Subclass of `scipy.stats.rv_discrete` initiated with values (i.e `rv_sample`) throws an error if `pmf `method is called. The same without subclassing works fine.

### Reproducing code example:
```
from scipy import stats 

class subc_rv_discrete(stats.rv_discrete):
    pass

xk = (0,1,2,3)
pk = (0.1, 0.2, 0.3, 0.4)

inst = subc_rv_discrete(values=(xk, pk))
inst.pmf(xk)
```
### Error message:
```
Traceback (most recent call last):

  File "<ipython-input-48-129655c38e6a>", line 10, in <module>
    inst.pmf(xk)

  File "C:\Anaconda3\lib\site-packages\scipy\stats\_distn_infrastructure.py", line 2832, in pmf
    args, loc, _ = self._parse_args(*args, **kwds)

AttributeError: 'rv_sample' object has no attribute '_parse_args'
```
### Scipy/Numpy/Python version information:
```
0.19.0 1.12.1 sys.version_info(major=3, minor=6, micro=1, releaselevel='final', serial=0)
```

Hints from issue discussion:
Can you share a use case where you need such a subclass? 
In general, I wanted to adjust the `rv_discrete` to do Bayesian statistics. In particular, I wanted to add a  likelihood function, mixture etc. in a similar way that is described in the docs. That was a simple tinkering experiment and I understand that there are dedicated modules for Bayesian statistics. 

Anyways, briefly looking at the source code I could not figure out if this behavior is intended or not and decided to point at this issue. Especially as the [docs](https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.rv_discrete.html) suggest the `rv_discrete` subclassing as a productive feature. 
trying this out a bit
going through the regular subclassing creates a rv_sample class and doesn't __init__ the correct class

the following works for me with 0.18.1 (which I have open right now)

```
from scipy.stats._distn_infrastructure import rv_sample


class subc_rv_discrete(rv_sample):


    def __new__(cls, *args, **kwds):
        return super(subc_rv_discrete, cls).__new__(cls)


xk = [0,1,2,3]
pk = [0.1, 0.2, 0.3, 0.4]

inst = subc_rv_discrete(values=(xk, pk))
print(inst.pmf(xk))
print(inst.__class__)
```
Ok, indeed, in general, subclassing is the way to go. Unless you hit this ugly corner :-). The dispatch in `__new__` is ugly and I'm not sure it can be made to play with subclassing. Like, try giving your class both __init__ and __new__? Maybe the correct thing to do is to subclassing rv_sample direcyly
rv_sample will need to be added to a public import path.

subclassing rv_sample alone is not enough, it needs also the `__new__` override (currently or with my scipy).
AFAICS, rv_discrete.__new__ is called by default which creates a rv_sample instance instead of the subclass.

If rv_discrete is subclassed directly, then it has the wrong `__init__` which doesn't handle `values` anymore.
```
        if values is not None:
            raise ValueError("rv_discrete.__init__(..., values != None, ...)")
```

@josef-pkt indeed this does the job. To spread the word, can you please answer the corresponding [stackoverflow](https://stackoverflow.com/questions/46635803/instance-of-scipy-stats-rv-discrete-subclass-throws-error-on-pmf-method) question and I will accept it.

Should the issue be closed or maybe a pull request be made? 
We need a PR with at least putting `rv_sample` in the `stats` and `stats.distributions` namespace.

@ev-br AFAICS, we could deprecate the `values` keyword in `rv_discrete` and then we could remove the `__new__` delegation to the rv_sample subclass in two versions or so. Users can use the `rv_sample` class directly.
That will need some adjustments to the unit tests, I guess.
Commenting to link in [another stackoverflow question](https://stackoverflow.com/questions/60981879/creating-scipy-stats-random-variable-subclass-does-not-result-in-expected-object) where this issue has popped up. The approach in josef-pkt's first answer doesn't seem to do it there.
When `rv_sample` is made public, it would be good to tweak the logic that detects whether an argument is within the support (i.e. equal to one of the discrete supported values) or not. As discussed in gh-15932, 
floating point errors make exact equality comparison too unreliable when, say, `loc != 0`.
@ev-br In-between `integrate` PRs sometime, would exposing `rv_sample` and deprecating that use of `rv_discrete` be something you're interested in doing? I'd be happy to review.

Commit message: MAINT: stats: rename rv_sample -> rv_count, add docstring

Patch to review:
diff --git a/scipy/stats/__init__.py b/scipy/stats/__init__.py
index ea8275fda544..153674b3c4a4 100644
--- a/scipy/stats/__init__.py
+++ b/scipy/stats/__init__.py
@@ -40,6 +40,7 @@
 
    rv_continuous
    rv_discrete
+   rv_count
    rv_histogram
 
 Continuous distributions
@@ -481,6 +482,7 @@
 from ._stats_py import *
 from ._variation import variation
 from .distributions import *
+from ._distn_infrastructure import rv_count
 from ._morestats import *
 from ._binomtest import binomtest
 from ._binned_statistic import *
diff --git a/scipy/stats/_distn_infrastructure.py b/scipy/stats/_distn_infrastructure.py
index f41c3c9c6d43..e92339ff0976 100644
--- a/scipy/stats/_distn_infrastructure.py
+++ b/scipy/stats/_distn_infrastructure.py
@@ -1054,7 +1054,7 @@ def rvs(self, *args, **kwds):
             self._random_state = random_state_saved
 
         # Cast to int if discrete
-        if discrete and not isinstance(self, rv_sample):
+        if discrete and not isinstance(self, _rv_sample):
             if size == ():
                 vals = int(vals)
             else:
@@ -3126,13 +3126,11 @@ class rv_discrete(rv_generic):
     >>> R = custm.rvs(size=100)
 
     """
-    def __new__(cls, a=0, b=inf, name=None, badvalue=None,
-                moment_tol=1e-8, values=None, inc=1, longname=None,
-                shapes=None, extradoc=None, seed=None):
-
+    def __new__(cls, *args, **kwds):
+        values = kwds.get('values', None)
         if values is not None:
             # dispatch to a subclass
-            return super(rv_discrete, cls).__new__(rv_sample)
+            return super(rv_discrete, cls).__new__(_rv_sample)
         else:
             # business as usual
             return super(rv_discrete, cls).__new__(cls)
@@ -3342,7 +3340,7 @@ def pmf(self, k, *args, **kwds):
         k = asarray((k-loc))
         cond0 = self._argcheck(*args)
         cond1 = (k >= _a) & (k <= _b)
-        if not isinstance(self, rv_sample):
+        if not isinstance(self, _rv_sample):
             cond1 = cond1 & self._nonzero(k, *args)
         cond = cond0 & cond1
         output = zeros(shape(cond), 'd')
@@ -3380,7 +3378,7 @@ def logpmf(self, k, *args, **kwds):
         k = asarray((k-loc))
         cond0 = self._argcheck(*args)
         cond1 = (k >= _a) & (k <= _b)
-        if not isinstance(self, rv_sample):
+        if not isinstance(self, _rv_sample):
             cond1 = cond1 & self._nonzero(k, *args)
         cond = cond0 & cond1
         output = empty(shape(cond), 'd')
@@ -3731,7 +3729,7 @@ def fun(x):
         else:
             invfac = 1.0
 
-        if isinstance(self, rv_sample):
+        if isinstance(self, _rv_sample):
             res = self._expect(fun, lb, ub)
             return res / invfac
 
@@ -3823,7 +3821,7 @@ def _iter_chunked(x0, x1, chunksize=4, inc=1):
         yield supp
 
 
-class rv_sample(rv_discrete):
+class _rv_sample(rv_discrete):
     """A 'sample' discrete distribution defined by the support and values.
 
     The ctor ignores most of the arguments, only needs the `values` argument.
@@ -3839,7 +3837,7 @@ def __init__(self, a=0, b=inf, name=None, badvalue=None,
                           "SciPy 1.11.0", DeprecationWarning)
 
         if values is None:
-            raise ValueError("rv_sample.__init__(..., values=None,...)")
+            raise ValueError("rv_count.__init__(..., values=None,...)")
 
         # cf generic freeze
         self._ctor_param = dict(
@@ -3913,6 +3911,175 @@ def _get_support(self, *args):
         """
         return self.a, self.b
 
+    def _pmf(self, x):
+        return rv_count._pmf(self, x)
+
+    def _cdf(self, x):
+        return rv_count._cdf(self, x)
+
+    def _ppf(self, q):
+        return rv_count._ppf(self, q)
+
+    def _rvs(self, size=None, random_state=None):
+        return rv_count._rvs(self, size, random_state)
+
+    def _entropy(self):
+        return rv_count._entropy(self)
+
+    def generic_moment(self, n):
+        return rv_count.generic_moment(self, n)
+
+    def _expect(self, fun, lb, ub, *args, **kwds):
+        return rv_count._expect(self, fun, lb, ub, *args, **kwds)
+
+
+class rv_count(rv_discrete):
+    """A discrete random variable defined by support points and probabilities.
+
+    Parameters
+    ----------
+    xk : array_like
+        Support points.
+    pk : array_like
+        The probabilities, corresponding to `xk`
+    name : str, optional
+        The name of the instance. This string is used to construct the default
+        example for distributions.
+    longname : str, optional
+        This string is used as part of the first line of the docstring returned
+        when a subclass has no docstring of its own. Note: `longname` exists
+        for backwards compatibility, do not use for new subclasses.
+    seed : {None, int, `numpy.random.Generator`, `numpy.random.RandomState`}, optional
+        If `seed` is None (or `np.random`), the `numpy.random.RandomState`
+        singleton is used.
+        If `seed` is an int, a new ``RandomState`` instance is used,
+        seeded with `seed`.
+        If `seed` is already a ``Generator`` or ``RandomState`` instance then
+        that instance is used.
+
+    Methods
+    -------
+    rvs
+    pmf
+    logpmf
+    cdf
+    logcdf
+    sf
+    logsf
+    ppf
+    isf
+    moment
+    stats
+    entropy
+    expect
+    median
+    mean
+    std
+    var
+    interval
+    __call__
+    support
+
+    Notes
+    -----
+    This class is a discrete analog of `rv_histogram`.
+
+    Examples
+    --------
+
+    >>> import numpy as np
+    >>> from scipy import stats
+    >>> xk = np.arange(7)
+    >>> pk = (0.1, 0.2, 0.3, 0.1, 0.1, 0.0, 0.2)
+    >>> custm = stats.rv_count(xk, pk, name='custm')
+    >>>
+    >>> import matplotlib.pyplot as plt
+    >>> fig, ax = plt.subplots(1, 1)
+    >>> ax.plot(xk, custm.pmf(xk), 'ro', ms=12, mec='r')
+    >>> ax.vlines(xk, 0, custm.pmf(xk), colors='r', lw=4)
+    >>> plt.show()
+
+    Random number generation:
+
+    >>> R = custm.rvs(size=100)
+
+    """
+    def __init__(self, xk, pk, name=None, longname=None, seed=None):
+        super(rv_discrete, self).__init__(seed)
+
+        # cf generic freeze
+        self._ctor_param = dict(xk=xk, pk=pk, seed=seed,
+                                name=name, longname=longname)
+
+        self.badvalue = np.nan
+        self.vecentropy = self._entropy
+        self.inc = 1
+
+        if np.shape(xk) != np.shape(pk):
+            raise ValueError("xk and pk must have the same shape.")
+        if np.less(pk, 0.0).any():
+            raise ValueError("All elements of pk must be non-negative.")
+        if not np.allclose(np.sum(pk), 1):
+            raise ValueError("The sum of provided pk is not 1.")
+
+        indx = np.argsort(np.ravel(xk))
+        self.xk = np.take(np.ravel(xk), indx, 0)
+        self.pk = np.take(np.ravel(pk), indx, 0)
+        self.a = self.xk[0]
+        self.b = self.xk[-1]
+
+        self.qvals = np.cumsum(self.pk, axis=0)
+
+        self.shapes = ' '   # bypass inspection
+
+        self._construct_argparser(meths_to_inspect=[self._pmf],
+                                  locscale_in='loc=0',
+                                  # scale=1 for discrete RVs
+                                  locscale_out='loc, 1')
+
+        self._attach_methods()
+
+        self._construct_docstrings(name, longname, extradoc=None)
+
+    def __getstate__(self):
+        dct = self.__dict__.copy()
+
+        # these methods will be remade in rv_generic.__setstate__,
+        # which calls rv_generic._attach_methods
+        attrs = ["_parse_args", "_parse_args_stats", "_parse_args_rvs"]
+        [dct.pop(attr, None) for attr in attrs]
+
+        return dct
+
+    def _updated_ctor_param(self):
+        """Return the current version of _ctor_param, possibly updated by user.
+
+        Used by freezing. Keep this in sync with the signature of __init__.
+        Just copy the _ctor_param for rv_count.
+        """
+        dct = self._ctor_param.copy()
+        return dct
+
+    def _attach_methods(self):
+        """Attaches dynamically created argparser methods."""
+        self._attach_argparser_methods()
+
+    def _get_support(self, *args):
+        """Return the support of the (unscaled, unshifted) distribution.
+
+        Parameters
+        ----------
+        arg1, arg2, ... : array_like
+            The shape parameter(s) for the distribution (see docstring of the
+            instance object for more information).
+
+        Returns
+        -------
+        a, b : numeric (float, or int or +/-np.inf)
+            end-points of the distribution's support.
+        """
+        return self.a, self.b
+
     def _pmf(self, x):
         return np.select([x == k for k in self.xk],
                          [np.broadcast_arrays(p, x)[0] for p in self.pk], 0)
diff --git a/scipy/stats/tests/test_discrete_basic.py b/scipy/stats/tests/test_discrete_basic.py
index d380bff05096..fef06bc666e7 100644
--- a/scipy/stats/tests/test_discrete_basic.py
+++ b/scipy/stats/tests/test_discrete_basic.py
@@ -17,7 +17,8 @@
 from scipy.stats._distn_infrastructure import rv_discrete_frozen
 
 vals = ([1, 2, 3, 4], [0.1, 0.2, 0.3, 0.4])
-distdiscrete += [[stats.rv_discrete(values=vals), ()]]
+distdiscrete += [[stats.rv_discrete(values=vals), ()],
+                 [stats.rv_count(xk=vals[0], pk=vals[1]), ()]]
 
 # For these distributions, test_discrete_basic only runs with test mode full
 distslow = {'zipfian', 'nhypergeom'}
diff --git a/scipy/stats/tests/test_distributions.py b/scipy/stats/tests/test_distributions.py
index c9a6e312e9c5..8dbcfde421b3 100755
--- a/scipy/stats/tests/test_distributions.py
+++ b/scipy/stats/tests/test_distributions.py
@@ -2812,6 +2812,20 @@ def test_expect2(self):
         assert_allclose(rv.expect(lambda x: x**2),
                         sum(v**2 * w for v, w in zip(y, py)), atol=1e-14)
 
+    def test_rv_count_subclassing(self):
+        # gh-8057: rv_discrete(values=(xk, pk)) cannot be subclassed easily
+        class S(stats.rv_count):
+            def extra(self):
+                return 42
+
+        s = S(xk=[1, 2, 3], pk=[0.2, 0.7, 0.1])
+        assert_allclose(s.pmf([2, 3, 1]), [0.7, 0.1, 0.2], atol=1e-15)
+        assert s.extra() == 42
+
+        # make sure subclass freezes correctly
+        frozen_s = s()
+        assert_allclose(frozen_s.pmf([2, 3, 1]), [0.7, 0.1, 0.2], atol=1e-15)
+
 
 class TestSkewCauchy:
     def test_cauchy(self):


Code review comments:
diff_hunk: @@ -1054,7 +1054,7 @@ def rvs(self, *args, **kwds):
             self._random_state = random_state_saved
 
         # Cast to int if discrete
-        if discrete and not isinstance(self, rv_sample):
+        if discrete and not isinstance(self, _rv_sample):
path: scipy/stats/_distn_infrastructure.py
text: @user1:
I suppose we could leave the name of `rv_sample` alone now that the public class is `rv_count`. Not necessary to change back; just a thought.

@author:
Good point, `rv_sample` is back in the last commit (while we are at it).