fix: Use Fisher-Yates for shuffledIndices

dataframe-operations/src/DataFrame/Operations/Permutation.hs

@@ -191,14 +191,12 @@ shuffledIndices pureGen k
     shuffleVec :: (RandomGen g) => g -> VU.Vector Int
     shuffleVec g = runST $ do
         vm <- VUM.generate k id
-        let (n, nGen) = randomR (1, k - 1) g
-        go vm n nGen
+        go vm (k - 1) g
         VU.unsafeFreeze vm
 
-    go _v (-1) _ = pure ()
-    go _v 0 _ = pure ()
-    go v maxInd gen =
+    go _v i _ | i <= 0 = pure ()
+    go v i gen =
         let
-            (n, nextGen) = randomR (1, maxInd) gen
+            (j, nextGen) = randomR (0, i) gen
          in
-            VUM.swap v 0 n *> go (VUM.tail v) (maxInd - 1) nextGen
+            VUM.swap v i j *> go v (i - 1) nextGen

tests/Operations/Shuffle.hs

@@ -5,11 +5,13 @@ module Operations.Shuffle where
 
 import qualified DataFrame as D
 
+import Data.List (permutations)
+import qualified Data.Map.Strict as M
 import qualified Data.Set as Set
 import qualified Data.Vector.Unboxed as VU
 import DataFrame.Operations.Permutation (shuffle, shuffledIndices)
 import System.Random (mkStdGen)
-import Test.HUnit (Test (..), assertEqual)
+import Test.HUnit (Test (..), assertBool, assertEqual)
 
 testDataFrame :: D.DataFrame
 testDataFrame =
@@ -91,9 +93,81 @@ shuffleDoesNotAddOrDropIndices =
             , TestCase (assertEqual "There are no repeated indecis" computed actual)
             ]
 
+-- A one-row frame has exactly one permutation.
+shuffleSingleRow :: Test
+shuffleSingleRow =
+    TestCase
+        ( assertEqual
+            "shuffling one index yields that index"
+            (VU.fromList [0 :: Int])
+            (shuffledIndices (mkStdGen 7) 1)
+        )
+
+{- | Chi-squared statistic of observed counts against a flat expectation:
+sum over cells of (observed - expected)^2 / expected.
+-}
+chiSquared :: [Int] -> Double
+chiSquared counts =
+    let expected = fromIntegral (sum counts) / fromIntegral (length counts)
+     in sum [(fromIntegral o - expected) ^ (2 :: Int) / expected | o <- counts]
+
+{- | Every permutation of n items is equally likely under a uniform shuffle,
+so the counts over all n! outcomes are chi-squared with n! - 1 degrees of
+freedom. Testing the whole permutation, rather than one position at a time,
+also catches a shuffle whose positions are individually uniform but
+correlated. Seeds are fixed, so the sample -- and the verdict -- is
+deterministic.
+
+n = 5 gives 120 outcomes; 12000 draws puts 100 in each on average. The bound
+is the 0.999 quantile of chi-squared with 119 degrees of freedom.
+-}
+shufflePermutationsAreUniform :: Test
+shufflePermutationsAreUniform =
+    let n = 5
+        trials = 12000
+        observed =
+            M.fromListWith
+                (+)
+                [(VU.toList (shuffledIndices (mkStdGen s) n), 1 :: Int) | s <- [1 .. trials]]
+        counts = [M.findWithDefault 0 p observed | p <- permutations [0 .. n - 1]]
+        stat = chiSquared counts
+     in TestCase
+            ( assertBool
+                ("chi-squared over all permutations is " ++ show stat ++ ", above 172.4")
+                (stat < 172.4)
+            )
+
+{- | The frequency test from Knuth 3.3.2: each item lands in each position with
+probability 1/n, so the n x n position-by-item table is chi-squared with
+(n - 1)^2 degrees of freedom. A larger n than the permutation test can afford,
+to catch bias that only shows at scale, such as a shuffle that leaves a
+suffix untouched or never leaves an item in place.
+
+n = 10 and 5000 draws put 500 in each cell. The bound is the 0.999 quantile of
+chi-squared with 81 degrees of freedom.
+-}
+shufflePositionsAreUniform :: Test
+shufflePositionsAreUniform =
+    let n = 10
+        trials = 5000
+        samples = [VU.toList (shuffledIndices (mkStdGen s) n) | s <- [1 .. trials]]
+        cell p i = length [() | xs <- samples, xs !! p == i]
+        stat = chiSquared [cell p i | p <- [0 .. n - 1], i <- [0 .. n - 1]]
+     in TestCase
+            ( assertBool
+                ( "chi-squared over the position-by-item table is "
+                    ++ show stat
+                    ++ ", above 126.1"
+                )
+                (stat < 126.1)
+            )
+
 tests :: [Test]
 tests =
-    [ TestLabel "shuffleShuffles" shuffleShuffles
+    [ TestLabel "shuffleSingleRow" shuffleSingleRow
+    , TestLabel "shufflePermutationsAreUniform" shufflePermutationsAreUniform
+    , TestLabel "shufflePositionsAreUniform" shufflePositionsAreUniform
+    , TestLabel "shuffleShuffles" shuffleShuffles
     , TestLabel "shufflePreservesData" shufflePreservesData
     , TestLabel "shufflePreservesColumnNames" shufflePreservesColumnNames
     , TestLabel "shuffleSameSeedIsSameShuffle" shuffleSameSeedIsSameShuffle