diff --git a/dataframe-operations/src/DataFrame/Operations/Permutation.hs b/dataframe-operations/src/DataFrame/Operations/Permutation.hs index 155e5c37..531bb024 100644 --- a/dataframe-operations/src/DataFrame/Operations/Permutation.hs +++ b/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 diff --git a/tests/Operations/Shuffle.hs b/tests/Operations/Shuffle.hs index 4ccf3e0c..b7f71c96 100644 --- a/tests/Operations/Shuffle.hs +++ b/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