Repository: dask/dask

Issue:
Multinomial random generator
`dask/array/random.py` is missing a multinomial random generator (there is aplaceholder `# multinomial`).

Will dask have a multinomial random generator at some point? Does it require a significantly different approach than the other generators?

Discussion:

Patch:
diff --git a/dask/array/random.py b/dask/array/random.py
index 71050b304..9a1b0b364 100644
--- a/dask/array/random.py
+++ b/dask/array/random.py
@@ -44,8 +44,13 @@ class RandomState(object):
         self._numpy_state.seed(seed)
 
     def _wrap(self, func, *args, **kwargs):
+        """ Wrap numpy random function to produce dask.array random function
+
+        extra_chunks should be a chunks tuple to append to the end of chunks
+        """
         size = kwargs.pop('size')
         chunks = kwargs.pop('chunks')
+        extra_chunks = kwargs.pop('extra_chunks', ())
 
         if not isinstance(size, (tuple, list)):
             size = (size,)
@@ -62,12 +67,13 @@ class RandomState(object):
         seeds = different_seeds(len(sizes), self._numpy_state)
         token = tokenize(seeds, size, chunks, args, kwargs)
         name = 'da.random.{0}-{1}'.format(func.__name__, token)
-        keys = product([name], *[range(len(bd)) for bd in chunks])
+        keys = product([name], *([range(len(bd)) for bd in chunks]
+                               + [[0]] * len(extra_chunks)))
         vals = ((_apply_random, func.__name__, seed, size, args, kwargs)
                 for seed, size in zip(seeds, sizes))
         dsk = dict(zip(keys, vals))
 
-        return Array(dsk, name, chunks, dtype=dtype)
+        return Array(dsk, name, chunks + extra_chunks, dtype=dtype)
 
     @doc_wraps(np.random.RandomState.beta)
     def beta(self, a, b, size=None, chunks=None):
@@ -144,7 +150,11 @@ class RandomState(object):
         return self._wrap(np.random.RandomState.logseries, p,
                          size=size, chunks=chunks)
 
-    # multinomial
+    @doc_wraps(np.random.RandomState.multinomial)
+    def multinomial(self, n, pvals, size=None, chunks=None):
+        return self._wrap(np.random.RandomState.multinomial, n, pvals,
+                          size=size, chunks=chunks,
+                          extra_chunks=((len(pvals),),))
 
     @doc_wraps(np.random.RandomState.negative_binomial)
     def negative_binomial(self, n, p, size=None, chunks=None):
@@ -295,6 +305,7 @@ laplace = _state.laplace
 logistic = _state.logistic
 lognormal = _state.lognormal
 logseries = _state.logseries
+multinomial = _state.multinomial
 negative_binomial = _state.negative_binomial
 noncentral_chisquare = _state.noncentral_chisquare
 noncentral_f = _state.noncentral_f


Test patch:
diff --git a/dask/array/tests/test_random.py b/dask/array/tests/test_random.py
index 855b200fd..1112a13ae 100644
--- a/dask/array/tests/test_random.py
+++ b/dask/array/tests/test_random.py
@@ -134,6 +134,7 @@ def test_random_all():
     da.random.logistic(size=5, chunks=3).compute()
     da.random.lognormal(size=5, chunks=3).compute()
     da.random.logseries(0.5, size=5, chunks=3).compute()
+    da.random.multinomial(20, [1/6.]*6, size=5, chunks=3).compute()
     da.random.negative_binomial(5, 0.5, size=5, chunks=3).compute()
     da.random.noncentral_chisquare(2, 2, size=5, chunks=3).compute()
 
@@ -159,3 +160,11 @@ def test_random_all():
     da.random.standard_gamma(2, size=5, chunks=3).compute()
     da.random.standard_normal(size=5, chunks=3).compute()
     da.random.standard_t(2, size=5, chunks=3).compute()
+
+
+def test_multinomial():
+    for size, chunks in [(5, 3), ((5, 4), (2, 3))]:
+        x = da.random.multinomial(20, [1/6.]*6, size=size, chunks=chunks)
+        y = np.random.multinomial(20, [1/6.]*6, size=size)
+
+        assert x.shape == y.shape == x.compute().shape