Repository: dask/dask

Issue:
Wrapper always pass None to random.choice
From dask/array/random.py:

>     with ignoring(AttributeError):
>         @doc_wraps(np.random.RandomState.choice)
>         def choice(self, a, size=None, replace=True, p=None, chunks=None):
>             return self._wrap(np.random.RandomState.choice, a,
>                               size=size, replace=True, p=None, chunks=chunks)

Wrapper gets non-None p argument -> Wrapper passes None to p kwarg of np.random.RandomState.choice anyway.

Test:

> import dask
> import dask.array as da
> import numpy as np
> import dask.array.random as dar
> import matplotlib.pyplot as plt
> 
> p = np.arange(12)
> p = p*1.0 / np.sum(p)
> 
> p = da.from_array(p, chunks=6)
> s = dar.RandomState()
> 
> c = s.choice(len(p), p=p,size=10000,chunks=8)
> cc = c.compute()
> b = np.bincount(cc)
> print(b)
> 
> c = np.random.choice(len(p), p=np.array(p), size=10000)
> b = np.bincount(c)
> print(b)
> 

Result from numpy and dask are different.

> [806 815 833 873 877 809 850 818 837 792 892 798]
> [   0  117  332  467  637  731  901 1079 1251 1316 1483 1686]




Discussion:

Resolving patch:
diff --git a/dask/array/random.py b/dask/array/random.py
--- a/dask/array/random.py
+++ b/dask/array/random.py
@@ -1,11 +1,12 @@
 from __future__ import absolute_import, division, print_function
 
 from itertools import product
+from numbers import Integral
 from operator import getitem
 
 import numpy as np
 
-from .core import (normalize_chunks, Array, slices_from_chunks,
+from .core import (normalize_chunks, Array, slices_from_chunks, asarray,
                    broadcast_shapes, broadcast_to)
 from .. import sharedict
 from ..base import tokenize
@@ -170,8 +171,62 @@ def chisquare(self, df, size=None, chunks=None):
     with ignoring(AttributeError):
         @doc_wraps(np.random.RandomState.choice)
         def choice(self, a, size=None, replace=True, p=None, chunks=None):
-            return self._wrap(np.random.RandomState.choice, a,
-                              size=size, replace=True, p=None, chunks=chunks)
+            dsks = []
+            # Normalize and validate `a`
+            if isinstance(a, Integral):
+                # On windows the output dtype differs if p is provided or
+                # absent, see https://github.com/numpy/numpy/issues/9867
+                dummy_p = np.array([1]) if p is not None else p
+                dtype = np.random.choice(1, size=(), p=dummy_p).dtype
+                len_a = a
+                if a < 0:
+                    raise ValueError("a must be greater than 0")
+            else:
+                a = asarray(a).rechunk(a.shape)
+                dtype = a.dtype
+                if a.ndim != 1:
+                    raise ValueError("a must be one dimensional")
+                len_a = len(a)
+                dsks.append(a.dask)
+                a = a._keys()[0]
+
+            # Normalize and validate `p`
+            if p is not None:
+                if not isinstance(p, Array):
+                    # If p is not a dask array, first check the sum is close
+                    # to 1 before converting.
+                    p = np.asarray(p)
+                    if not np.isclose(p.sum(), 1, rtol=1e-7, atol=0):
+                        raise ValueError("probabilities do not sum to 1")
+                    p = asarray(p)
+                else:
+                    p = p.rechunk(p.shape)
+
+                if p.ndim != 1:
+                    raise ValueError("p must be one dimensional")
+                if len(p) != len_a:
+                    raise ValueError("a and p must have the same size")
+
+                dsks.append(p.dask)
+                p = p._keys()[0]
+
+            if size is None:
+                size = ()
+            elif not isinstance(size, (tuple, list)):
+                size = (size,)
+
+            chunks = normalize_chunks(chunks, size)
+            sizes = list(product(*chunks))
+            state_data = random_state_data(len(sizes), self._numpy_state)
+
+            name = 'da.random.choice-%s' % tokenize(state_data, size, chunks,
+                                                    a, replace, p)
+            keys = product([name], *(range(len(bd)) for bd in chunks))
+            dsk = {k: (_choice, state, a, size, replace, p) for
+                   k, state, size in zip(keys, state_data, sizes)}
+
+            return Array(sharedict.merge((name, dsk), *dsks),
+                         name, chunks, dtype=dtype)
 
     # @doc_wraps(np.random.RandomState.dirichlet)
     # def dirichlet(self, alpha, size=None, chunks=None):
@@ -352,6 +407,11 @@ def zipf(self, a, size=None, chunks=None):
                           size=size, chunks=chunks)
 
 
+def _choice(state_data, a, size, replace, p):
+    state = np.random.RandomState(state_data)
+    return state.choice(a, size=size, replace=replace, p=p)
+
+
 def _apply_random(func, state_data, size, args, kwargs):
     """Apply RandomState method with seed"""
     state = np.random.RandomState(state_data)
@@ -368,6 +428,8 @@ def _apply_random(func, state_data, size, args, kwargs):
 beta = _state.beta
 binomial = _state.binomial
 chisquare = _state.chisquare
+if hasattr(_state, 'choice'):
+    choice = _state.choice
 exponential = _state.exponential
 f = _state.f
 gamma = _state.gamma


Test patch:
diff --git a/dask/array/tests/test_random.py b/dask/array/tests/test_random.py
--- a/dask/array/tests/test_random.py
+++ b/dask/array/tests/test_random.py
@@ -216,3 +216,52 @@ def test_multinomial():
         y = np.random.multinomial(20, [1 / 6.] * 6, size=size)
 
         assert x.shape == y.shape == x.compute().shape
+
+
+def test_choice():
+    np_dtype = np.random.choice(1, size=()).dtype
+    size = (10, 3)
+    chunks = 4
+    x = da.random.choice(3, size=size, chunks=chunks)
+    assert x.dtype == np_dtype
+    assert x.shape == size
+    res = x.compute()
+    assert res.dtype == np_dtype
+    assert res.shape == size
+
+    np_a = np.array([1, 3, 5, 7, 9], dtype='f8')
+    da_a = da.from_array(np_a, chunks=2)
+
+    for a in [np_a, da_a]:
+        x = da.random.choice(a, size=size, chunks=chunks)
+        res = x.compute()
+        assert x.dtype == np_a.dtype
+        assert res.dtype == np_a.dtype
+        assert set(np.unique(res)).issubset(np_a)
+
+    np_p = np.array([0, 0.2, 0.2, 0.3, 0.3])
+    da_p = da.from_array(np_p, chunks=2)
+
+    for a, p in [(da_a, np_p), (np_a, da_p)]:
+        x = da.random.choice(a, size=size, chunks=chunks, p=p)
+        res = x.compute()
+        assert x.dtype == np_a.dtype
+        assert res.dtype == np_a.dtype
+        assert set(np.unique(res)).issubset(np_a[1:])
+
+    np_dtype = np.random.choice(1, size=(), p=np.array([1])).dtype
+    x = da.random.choice(5, size=size, chunks=chunks, p=np_p)
+    res = x.compute()
+    assert x.dtype == np_dtype
+    assert res.dtype == np_dtype
+
+    errs = [(-1, None),             # negative a
+            (np_a[:, None], None),  # a must be 1D
+            (np_a, np_p[:, None]),  # p must be 1D
+            (np_a, np_p[:-2]),      # a and p must match
+            (3, np_p),              # a and p must match
+            (4, [0.2, 0.2, 0.3])]   # p must sum to 1
+
+    for (a, p) in errs:
+        with pytest.raises(ValueError):
+            da.random.choice(a, size=size, chunks=chunks, p=p)