cunumeric.random - sort and shuffle-based operations

cunumeric/eager.py

@@ -711,9 +711,9 @@ def bitgenerator_integers(
             )
         else:
             if self.array.size == 1:
-                self.array.fill(np.random.random_integers(low, high))
+                self.array.fill(np.random.randint(low, high))
             else:
-                a = np.random.random_integers(low, high, size=self.array.shape)
+                a = np.random.randint(low, high, size=self.array.shape)
                 self.array[:] = a
 
     def bitgenerator_lognormal(

cunumeric/random/random.py

@@ -20,6 +20,7 @@
 from cunumeric.array import ndarray
 from cunumeric.runtime import runtime
 
+import cunumeric.sort
 from cunumeric.random import generator
 
 if TYPE_CHECKING:
@@ -216,6 +217,96 @@ def chisquare(
     return generator.get_static_generator().chisquare(df, size, dtype)
 
 
+def choice(
+    a: Union[int, ndarray],
+    size: Union[NdShapeLike, None] = None,
+    replace: bool = True,
+    p: Union[None, ndarray] = None,
+) -> Union[int, ndarray, npt.NDArray[Any]]:
+    """
+    Returns an array of random values from a given 1-D array.
+    Each element of the returned array is independently sampled
+    from ``a`` according to ``p`` or uniformly.
+
+    Note
+    ----
+       Currently ``p`` is not supported when ``replace=False``.
+
+    Parameters
+    ----------
+    a : ndarray or int : If an array-like,
+            a random sample is generated from its elements.
+            If an int, the random sample is generated as if ``a`` was
+            ``np.arange(n)``
+    size : int or tuple of ints : The shape of the array.
+    replace (boolean): Whether the sample is with or without replacement.
+    p : 1-D array-like :
+            The probabilities associated with each entry in ``a``.
+            If not given the sample assumes a uniform distribution over all
+            entries in ``a``.
+
+    Returns
+    -------
+    out : ndarray : The generated random samples.
+
+
+    Raises
+    ------
+    ValueError
+        If a is an int and less than zero, if a or p are not 1-dimensional,
+        if a is an array-like of size 0, if p is not a vector of
+        probabilities, if a and p have different lengths, or if
+        replace=False and the sample size is greater than the population
+        size
+
+    See Also
+    --------
+    numpy.random.choice
+
+    Availability
+    --------
+    Multiple GPUs, Multiple CPUs
+    """
+
+    if isinstance(a, int) and a <= 0:
+        raise ValueError
+    if not isinstance(size, int):
+        raise ValueError
+    if p is None:
+        if isinstance(a, int):
+            if replace:
+                return cunumeric.random.randint(0, a, size)
+            else:
+                if a < size:
+                    raise ValueError
+                return cunumeric.random.permutation(a)[:size]
+        else:
+            if replace:
+                indices = cunumeric.random.randint(0, len(a), size)
+                return a[indices]
+            else:
+                if len(a) < size:
+                    raise ValueError
+                return cunumeric.random.permutation(a)[:size]
+    else:
+        if not replace:
+            raise ValueError
+        cump = p.cumsum()
+        # check if p is a probability distribution
+        if abs(cump[-1] - 1.0) > len(p) * np.finfo(p.dtype).eps:
+            # does not sum up to 1
+            raise ValueError
+        cump[-1] = 1.0  # fix rounding issues
+        # draw uniforms
+        uni = cunumeric.random.uniform(0.0, 1.0, size)
+        # inverse distribution
+        idx = cunumeric.searchsorted(cump, uni, "right")
+        # return samples
+        if isinstance(a, int):
+            return cunumeric.arange(a)[idx]
+        return a[idx]
+
+
 def exponential(
     scale: float = 1.0,
     size: Union[NdShapeLike, None] = None,
@@ -924,6 +1015,40 @@ def pareto(
     return generator.get_static_generator().pareto(a, size, dtype)
 
 
+def permutation(x: Union[int, ndarray]) -> ndarray:
+    """
+    Returns a permuted range or a permutation of an array.
+
+    Parameters
+    ----------
+    a : (int or ndarray): The range or the array to be shuffled.
+
+    Returns
+    -------
+    out: If `a` is an integer, it is permutation range between 0
+        and `a` - 1.
+        Otherwise, it is a permutation of `a`.
+
+    See Also
+    --------
+    numpy.random.permutation
+
+    Availability
+    --------
+    Multiple GPUs, Multiple CPUs
+    """
+    if isinstance(x, int):
+        count = x
+    else:
+        count = len(x)
+    key = uniform(0.0, 1.0, count)
+    indices = cunumeric.argsort(key)
+    if isinstance(x, int):
+        return indices
+    else:
+        return x[indices]
+
+
 def poisson(
     lam: float = 1.0, size: Union[NdShapeLike, None] = None
 ) -> ndarray:
@@ -1277,6 +1402,26 @@ def rayleigh(
 sample = random_sample
 
 
+def shuffle(a: ndarray) -> None:
+    """
+    Shuffles an array.
+
+    Parameters
+    ----------
+    a : (ndarray): The array to be shuffled.
+
+    See Also
+    --------
+    numpy.random.shuffle
+
+    Availability
+    --------
+    Multiple GPUs, Multiple CPUs
+    """
+    b = cunumeric.random.permutation(a)
+    a[:] = b
+
+
 def standard_cauchy(
     size: Union[NdShapeLike, None] = None,
     dtype: npt.DTypeLike = np.float64,

tests/integration/test_random_permsort.py

@@ -0,0 +1,116 @@
+# Copyright 2022 NVIDIA Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import numpy as np
+import pytest
+
+import cunumeric as num
+
+
+def test_permutation_int():
+    count = 1024
+    p = num.random.permutation(count)
+    p.sort()
+    assert num.linalg.norm(p - np.arange(count)) == 0.0
+
+
+def test_permutation_array():
+    count = 1024
+    x = num.arange(count)
+    p = num.random.permutation(x)
+    assert num.linalg.norm(x - p) != 0.0
+    p.sort()
+    assert num.linalg.norm(x - p) == 0.0
+
+
+def test_shuffle():
+    count = 16
+    p = num.arange(count)
+    x = num.arange(count)
+    num.random.shuffle(x)
+    assert num.linalg.norm(x - p) != 0.0
+    x.sort()
+    assert num.linalg.norm(x - p) == 0.0
+
+
+def test_choice_1(maxvalue=1024, count=42):
+    a = num.random.choice(maxvalue, count)
+    assert len(a) == count
+    assert num.amax(a) <= maxvalue
+    assert num.amin(a) >= 0
+
+
+def test_choice_2(maxvalue=1024, count=42):
+    a = num.random.choice(maxvalue, count, False)
+    assert len(a) == count
+    assert num.amax(a) <= maxvalue
+    assert num.amin(a) >= 0
+    for i in range(count):
+        for j in range(count):
+            if i == j:
+                continue
+            assert a[i] != a[j]
+
+
+def test_choice_3(maxvalue=1024, count=42):
+    values = num.random.random_integers(0, maxvalue, maxvalue)
+
+    a = num.random.choice(values, count)
+    assert len(a) == count
+    assert num.amax(a) <= num.amax(values)
+    assert num.amin(a) >= num.amin(values)
+
+
+def test_choice_4(maxvalue=1024, count=42):
+    values = num.arange(maxvalue)
+
+    a = num.random.choice(values, count, False)
+    assert len(a) == count
+    assert num.amax(a) <= num.amax(values)
+    assert num.amin(a) >= num.amin(values)
+    for i in range(count):
+        for j in range(count):
+            if i == j:
+                continue
+            assert a[i] != a[j]
+
+
+def test_choice_5(maxvalue=1024, count=42):
+    values = num.arange(maxvalue)
+
+    p = num.random.uniform(0, 1, maxvalue)
+    p /= p.sum()
+
+    a = num.random.choice(values, count, True, p)
+    assert len(a) == count
+    assert num.amax(a) <= num.amax(values)
+    assert num.amin(a) >= num.amin(values)
+
+
+def test_choice_6(maxvalue=1024, count=42):
+    values = num.random.random_integers(0, maxvalue, maxvalue)
+
+    p = num.random.uniform(0, 1, maxvalue)
+    p /= p.sum()
+
+    a = num.random.choice(values, count, True, p)
+    assert len(a) == count
+    assert num.amax(a) <= num.amax(values)
+    assert num.amin(a) >= num.amin(values)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(sys.argv))