Implement merge shuffle algorithm for Numba

PySDM/backends/impl_common/index.py

@@ -37,11 +37,19 @@ def from_ndarray(array):
         def sort_by_key(self, keys):
             backend.sort_by_key(self, keys)
 
-        def shuffle(self, temporary, parts=None):
+        def shuffle(self, temporary, parts=None, merge_shuffle=False):
             if parts is None:
-                backend.shuffle_global(
-                    idx=self.data, length=self.length, u01=temporary.data
-                )
+                if merge_shuffle:
+                    backend.shuffle_global(
+                        idx=self.data,
+                        length=self.length,
+                        u01=temporary.data,
+                        cutoff=0x100000,
+                    )
+                else:
+                    backend.shuffle_global(
+                        idx=self.data, length=self.length, u01=temporary.data
+                    )
             else:
                 backend.shuffle_local(
                     idx=self.data, u01=temporary.data, cell_start=parts.data

PySDM/backends/impl_numba/methods/index_methods.py

@@ -2,11 +2,45 @@
 CPU implementation of shuffling and sorting backend methods
 """
 import numba
+import numpy as np
 
 from PySDM.backends.impl_common.backend_methods import BackendMethods
 from PySDM.backends.impl_numba import conf
 
 
+@numba.njit(**{**conf.JIT_FLAGS, **{"parallel": False}})
+def draw_random_int(start: int, end: int, u01: float):
+    return min(int(start + u01 * (end - start + 1)), end)
+
+
+@numba.njit(**{**conf.JIT_FLAGS, **{"parallel": False}})
+def fisher_yates_shuffle(idx, u01, start, end, random_offset=0):
+    for i in range(end - 1, start, -1):
+        j = draw_random_int(start=start, end=i, u01=u01[random_offset + i])
+        idx[i], idx[j] = idx[j], idx[i]
+
+
+@numba.njit(**{**conf.JIT_FLAGS, **{"parallel": False}})
+def merge(
+    idx, u01, start, middle, end, random_offset
+):  # pylint: disable=too-many-arguments
+    i = start
+    j = middle
+
+    while True:
+        if u01[random_offset + i] > 0.5:
+            if j == end:
+                break
+            idx[i], idx[j] = idx[j], idx[i]
+            j += 1
+        else:
+            if i == j:
+                break
+        i += 1
+
+    fisher_yates_shuffle(idx, u01, i, end, random_offset)
+
+
 class IndexMethods(BackendMethods):
     @staticmethod
     @numba.njit(**conf.JIT_FLAGS)
@@ -15,19 +49,34 @@ def identity_index(idx):
             idx[i] = i
 
     @staticmethod
-    @numba.njit(**{**conf.JIT_FLAGS, **{"parallel": False}})
-    def shuffle_global(idx, length, u01):
-        for i in range(length - 1, 0, -1):
-            j = int(u01[i] * (i + 1))
-            idx[i], idx[j] = idx[j], idx[i]
+    @numba.njit(**{**conf.JIT_FLAGS})
+    def shuffle_global(idx, length, u01, cutoff=None):
+        depth = 0
+        while cutoff is not None and (length >> depth) > cutoff:
+            depth += 1
+
+        split_start = np.linspace(0, length, (1 << depth) + 1).astype(np.uint32)
+
+        for c in numba.prange(len(split_start) - 1):  # pylint: disable=not-an-iterable
+            fisher_yates_shuffle(idx, u01, split_start[c], split_start[c + 1])
+
+        for i in range(1, depth + 1):
+            for c in numba.prange(1 << (depth - i)):  # pylint: disable=not-an-iterable
+                start = c * i * 2
+                merge(
+                    idx,
+                    u01,
+                    split_start[start],
+                    split_start[start + i],
+                    split_start[start + i + i],
+                    length * i,
+                )
 
     @staticmethod
     @numba.njit(**conf.JIT_FLAGS)
     def shuffle_local(idx, u01, cell_start):
         for c in numba.prange(len(cell_start) - 1):  # pylint: disable=not-an-iterable
-            for i in range(cell_start[c + 1] - 1, cell_start[c], -1):
-                j = int(cell_start[c] + u01[i] * (cell_start[c + 1] - cell_start[c]))
-                idx[i], idx[j] = idx[j], idx[i]
+            fisher_yates_shuffle(idx, u01, cell_start[c], cell_start[c + 1])
 
     @staticmethod
     def sort_by_key(idx, attr):

PySDM/backends/impl_thrust_rtc/methods/index_methods.py

@@ -22,7 +22,10 @@ def identity_index(idx):
 
     @staticmethod
     @nice_thrust(**NICE_THRUST_FLAGS)
-    def shuffle_global(idx, length, u01):
+    def shuffle_global(idx, length, u01, cutoff=None):
+        if cutoff is not None:
+            raise NotImplementedError()
+
         # WARNING: ineffective implementation
 
         # TODO #328 : Thrust modifies key array, conflicts with rand_reuse logic

PySDM/dynamics/collisions/collision.py

@@ -233,7 +233,9 @@ def step(self):
     def toss_candidate_pairs_and_sort_within_pair_by_multiplicity(
         self, is_first_in_pair, u01
     ):
-        self.particulator.attributes.permutation(u01, self.croupier == "local")
+        self.particulator.attributes.permutation(
+            u01, self.croupier == "local", merge_shuffle=False
+        )
         is_first_in_pair.update(
             self.particulator.attributes.cell_start,
             self.particulator.attributes.cell_idx,

PySDM/impl/particle_attributes.py

@@ -88,13 +88,13 @@ def __getitem__(self, item):
     def __contains__(self, key):
         return key in self.__attributes
 
-    def permutation(self, u01, local):
-        """apply Fisher-Yates algorithm to all super-droplets (local=False) or
+    def permutation(self, u01, local, merge_shuffle=False):
+        """apply Fisher-Yates or MergeShuffle algorithm to all super-droplets (local=False) or
         otherwise on a per-cell basis"""
         if local:
             self.__idx.shuffle(u01, parts=self.cell_start)
         else:
-            self.__idx.shuffle(u01)
+            self.__idx.shuffle(u01, merge_shuffle=merge_shuffle)
             self.__sorted = False
 
     def __sort_by_cell_id(self):

tests/unit_tests/backends/test_index_methods.py

@@ -0,0 +1,101 @@
+# pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring
+import matplotlib.pyplot as plt
+import numpy as np
+import pytest
+
+from PySDM.backends.impl_numba.methods.index_methods import (
+    IndexMethods,
+    draw_random_int,
+    fisher_yates_shuffle,
+)
+
+
+@pytest.mark.parametrize(
+    "a, b, u01, expected",
+    (
+        (0, 100, 0.0, 0),
+        (0, 100, 1.0, 100),
+        (0, 1, 0.5, 1),
+        (0, 1, 0.49, 0),
+        (0, 3, 0.49, 1),
+        (0, 3, 0.245, 0),
+        (0, 2, 0.332, 0),
+        (0, 2, 0.333, 0),
+        (0, 2, 0.334, 1),
+        (0, 2, 0.665, 1),
+        (0, 2, 0.666, 1),
+        (0, 2, 0.667, 2),
+        (0, 2, 0.999, 2),
+    ),
+)
+def test_draw_random_int(a, b, u01, expected):
+    # act
+    actual = draw_random_int(a, b, u01)
+
+    # assert
+    assert actual == expected
+
+
+def test_fisher_yates_shuffle():
+    # arrange
+    n = 10
+
+    idx = np.arange(n)
+    random_nums = np.linspace(1, 0, n + 2)[1:-1]
+
+    # act
+    fisher_yates_shuffle(idx, random_nums, 0, len(idx))
+
+    # assert
+    expected = np.array([9, 8, 3, 4, 5, 6, 7, 2, 1, 0])
+    assert np.all(expected == idx)
+
+
+@pytest.mark.parametrize("seed", (1, 2, 3, 1231, 42))
+# pylint: disable=redefined-outer-name
+def test_shuffle_global_generates_uniform_distribution(seed, plot=False):
+    # arrange
+    np.random.seed(seed)
+
+    n = 4
+    possible_permutations_num = np.math.factorial(n)
+    coverage = 1000
+
+    n_runs = coverage * possible_permutations_num
+
+    # act
+    all_permutations = {}
+
+    for _ in range(n_runs):
+        idx = np.arange(0, n)
+        u01 = np.random.uniform(0, 1, n)
+
+        IndexMethods.shuffle_global(idx, len(idx), u01)
+        # np.random.shuffle(idx)
+
+        perm_tuple = (*idx,)
+        if perm_tuple in all_permutations:
+            all_permutations[perm_tuple] += 1
+        else:
+            all_permutations[perm_tuple] = 1
+
+    all_permutations_vals = np.array(list(all_permutations.values()))
+    deviations_from_mean = all_permutations_vals - coverage
+
+    plt.bar(np.arange(possible_permutations_num), deviations_from_mean)
+    plt.axhline(0, color="green")
+
+    if plot:
+        plt.show()
+
+    # assert
+    assert len(all_permutations.keys()) == possible_permutations_num
+    assert np.amax(np.abs(deviations_from_mean)) < coverage * 0.15
+
+    tmp = np.abs(all_permutations_vals - coverage) ** 2.0
+    std = np.sqrt(np.mean(tmp))
+    assert std < coverage * 0.05
+
+
+def test_merge_shuffle_equals_fisher_yates_when_depth_0():
+    pass  # TODO:

tests/unit_tests/impl/test_particle_attributes.py

@@ -1,7 +1,10 @@
 # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring
 import numpy as np
 import pytest
+from matplotlib import pyplot
+from scipy import stats
 
+from PySDM import Formulae
 from PySDM.backends import CPU, GPU, ThrustRTC
 from PySDM.backends.impl_common.index import make_Index
 from PySDM.backends.impl_common.indexed_storage import make_IndexedStorage
@@ -20,6 +23,16 @@ def make_indexed_storage(backend, iterable, idx=None):
     return result
 
 
+def calculate_permutation_index(permutation, length):
+    result = 0
+    for i in range(length):
+        result = result * (length - i)
+        for j in range(i + 1, length):
+            if permutation[i] > permutation[j]:
+                result += 1
+    return result
+
+
 # pylint: disable=protected-access
 class TestParticleAttributes:
     @staticmethod
@@ -166,7 +179,7 @@ def test_permutation_global_as_implemented_in_numba():
         sut.permutation(u01, local=False)
 
         # Assert
-        expected = np.array([1, 3, 5, 7, 6, 0, 4, 2])
+        expected = np.array([5, 3, 0, 7, 6, 1, 4, 2])
         np.testing.assert_array_equal(sut._ParticleAttributes__idx, expected)
         assert not sut._ParticleAttributes__sorted
 
@@ -280,3 +293,65 @@ def test_permutation_local_repeatable(backend_class):
         np.testing.assert_array_equal(
             sut._ParticleAttributes__idx.to_ndarray(), expected
         )
+
+    @staticmethod
+    @pytest.mark.parametrize("seed", (1, 2, 3))
+    # pylint: disable=redefined-outer-name
+    def test_permutation_global_uniform_distribution(
+        seed, backend_class=CPU, plot=False
+    ):
+        n_sd = 4
+        possible_permutations_num = np.math.factorial(n_sd)
+        coverage = 1000
+
+        random_numbers = np.linspace(
+            0.0, 1.0, n_sd * possible_permutations_num * coverage
+        )
+        np.random.seed(seed)
+        np.random.shuffle(random_numbers)
+
+        # Arrange
+        particulator = DummyParticulator(CPU, n_sd=n_sd, formulae=Formulae(seed=seed))
+
+        sut = ParticleAttributesFactory.empty_particles(particulator, n_sd)
+        idx_length = len(sut._ParticleAttributes__idx)
+        sut._ParticleAttributes__tmp_idx = make_indexed_storage(
+            particulator.backend, [0] * idx_length
+        )
+        sut._ParticleAttributes__sorted = True
+        sut._ParticleAttributes__n_sd = particulator.n_sd
+
+        # Act
+        permutation_ids = np.zeros((possible_permutations_num * coverage,))
+
+        for i in range(possible_permutations_num * coverage):
+            u01 = make_indexed_storage(
+                particulator.backend, random_numbers[i * n_sd : (i + 1) * n_sd]
+            )
+            sut.permutation(u01, local=False)
+            permutation_ids[i] = calculate_permutation_index(
+                sut._ParticleAttributes__idx, idx_length
+            )
+
+        # Plot
+        counts, _ = np.histogram(permutation_ids, bins=possible_permutations_num)
+        _, uniformity = stats.chisquare(counts)
+
+        avg = np.mean(counts)
+        std = np.std(counts)
+
+        pyplot.plot(counts, marker=".")
+        pyplot.xlabel("permutation id")
+        pyplot.ylabel("occurrence count")
+        pyplot.xlim(0, possible_permutations_num)
+        pyplot.axhline(coverage, color="black", label="coverage")
+        pyplot.axhline(avg, color="green", label="mean +/- std")
+        for offset in (-std, +std):
+            pyplot.axhline(avg + offset, color="green", linestyle="--")
+        pyplot.legend()
+        if plot:
+            pyplot.show()
+
+        # Assert
+        assert abs(avg - coverage) / coverage < 1e-6
+        assert uniformity > 0.9