From 6d873e2741280c37bb3ab385c8eb4405790caab4 Mon Sep 17 00:00:00 2001 From: Boris Nagaev Date: Fri, 10 Mar 2017 17:11:36 +0100 Subject: [PATCH 1/4] fix bias in returned values There were two issues. * Bias of distribution of remainders of Int31. Demonstrated in TestTail. Items from the first half of the table get one addition remainder. Can be fixed without overhead by rethrowing RNG in that rare case when its value goes belongs to that tail. * The random number used to choose bucket is not independent from the random number used to decide whether to switch to alias. Demonstrated in TestBalanceInsideBucket. I don't know how to fix without overhead of second call to RNG. Both issues were resolved in this commit. Statistical tests added. Impact on performance: Before: BenchmarkGen5-8 100000000 24.1 ns/op BenchmarkGen50-8 50000000 28.7 ns/op BenchmarkGen500-8 50000000 27.2 ns/op BenchmarkGen5000-8 50000000 27.9 ns/op BenchmarkGen50000-8 50000000 31.3 ns/op After: BenchmarkGen5-8 50000000 38.8 ns/op BenchmarkGen50-8 30000000 53.0 ns/op BenchmarkGen500-8 30000000 40.9 ns/op BenchmarkGen5000-8 30000000 41.5 ns/op BenchmarkGen50000-8 30000000 43.9 ns/op --- alias.go | 5 ++--- alias_test.go | 61 +++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 3 deletions(-) diff --git a/alias.go b/alias.go index 5ec71c3..1c98a29 100644 --- a/alias.go +++ b/alias.go @@ -122,9 +122,8 @@ func New(prob []float64) (*Alias, error) { // Generates a random number according to the distribution using the rng passed. func (al *Alias) Gen(rng *rand.Rand) uint32 { - ri := uint32(rng.Int31()) - w := ri % uint32(len(al.table)) - if ri > al.table[w].prob { + w := uint32(rng.Int31n(int32(len(al.table)))) + if uint32(rng.Int31()) > al.table[w].prob { return al.table[w].alias } return w diff --git a/alias_test.go b/alias_test.go index fd0f686..b2a1daf 100644 --- a/alias_test.go +++ b/alias_test.go @@ -9,6 +9,8 @@ import ( "math/rand" "reflect" "testing" + + stat "github.com/ematvey/gostat" ) const distributionCount = 1000000 @@ -49,6 +51,65 @@ func TestDistribution(t *testing.T) { testDistribution(t, []float64{1000, 1, 3, 10}, 61) } +func TestTail(t *testing.T) { + const size = 33294320 + const half = size / 2 + const tries = 1000000 + const alpha = 0.05 + dist := make([]float64, size) + for i := range dist { + dist[i] = 1 + } + a, err := New(dist) + if err != nil { + t.Fatalf("Got an error during creation:", err) + } + rng := rand.New(rand.NewSource(42)) + var k int64 // [0,half) (one half) + for i := 0; i < tries; i++ { + if a.Gen(rng) < half { + k++ + } + } + // Expected probability of getting k <= k_observed if p == 0.5. + p := stat.Binomial_CDF_At(0.5, tries, k) + if p < alpha/2 || p > (1-alpha/2) { + t.Errorf("The distribution is biased. %d of %d results were in the first half. Binomial_CDF = %f", k, tries, p) + } +} + +func TestBalanceInsideBucket(t *testing.T) { + const size = 33294320 + //const size = 8 + const half = size / 2 + const tries = 1000000 + const alpha = 0.05 + dist := make([]float64, size) + for i := range dist { + if i < half { + dist[i] = 1 + } else { + dist[i] = 3 + } + } + a, err := New(dist) + if err != nil { + t.Fatalf("Got an error during creation:", err) + } + rng := rand.New(rand.NewSource(421)) + var k int64 // [0,half) (one half) + for i := 0; i < tries; i++ { + if a.Gen(rng) < half { + k++ + } + } + // Expected probability of getting k <= k_observed if p == 0.5. + p := stat.Binomial_CDF_At(0.25, tries, k) + if p < alpha/2 || p > (1-alpha/2) { + t.Errorf("The distribution is biased. %d of %d results were in the first half. Binomial_CDF = %f", k, tries, p) + } +} + func TestMarshalBinary(t *testing.T) { distributions := [][]float64{ {1}, From 6d01a3e3fb1b9b231caaddde7aad08bc7feab7f2 Mon Sep 17 00:00:00 2001 From: Boris Nagaev Date: Fri, 10 Mar 2017 19:18:26 +0100 Subject: [PATCH 2/4] add sanity check --- alias.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/alias.go b/alias.go index 1c98a29..8c2ca04 100644 --- a/alias.go +++ b/alias.go @@ -83,6 +83,10 @@ func New(prob []float64) (*Alias, error) { } } + if lgBot != smTop + 1 { + panic("alias.New: internal error") + } + for smTop >= 0 && lgBot < n { // pair off a small and large block, taking the chunk from the large block that's wanted l := twins[smTop] From db3cba382e79e928763a0408f582bb519636f36e Mon Sep 17 00:00:00 2001 From: Boris Nagaev Date: Fri, 10 Mar 2017 19:47:10 +0100 Subject: [PATCH 3/4] use rand.Int63 instead of two rand.Int31 Now performance is good again. --- alias.go | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/alias.go b/alias.go index 8c2ca04..a456fd9 100644 --- a/alias.go +++ b/alias.go @@ -14,6 +14,7 @@ import ( type Alias struct { table []ipiece + maxRi uint32 } type fpiece struct { @@ -26,6 +27,11 @@ type ipiece struct { alias uint32 } +func calcMax(n uint32) uint32 { + // Taken from math/rand.Rand.Int31n source. + return (1 << 31) - 1 - (1<<31)%uint32(n) +} + // Create a new alias object. // For example, // var v = alias.New([]float64{8,10,2}) @@ -83,7 +89,7 @@ func New(prob []float64) (*Alias, error) { } } - if lgBot != smTop + 1 { + if lgBot != smTop+1 { panic("alias.New: internal error") } @@ -121,13 +127,22 @@ func New(prob []float64) (*Alias, error) { al.table[twins[i].alias].prob = 1<<31 - 1 } + al.maxRi = calcMax(uint32(n)) + return &al, nil } // Generates a random number according to the distribution using the rng passed. func (al *Alias) Gen(rng *rand.Rand) uint32 { - w := uint32(rng.Int31n(int32(len(al.table)))) - if uint32(rng.Int31()) > al.table[w].prob { +begin: + r := rng.Int63() + ri := uint32(r & (1<<31 - 1)) + rj := uint32((r >> 31) & (1<<31 - 1)) + if ri > al.maxRi { + goto begin + } + w := ri % uint32(len(al.table)) + if rj > al.table[w].prob { return al.table[w].alias } return w @@ -171,5 +186,7 @@ func (al *Alias) UnmarshalBinary(p []byte) error { al.table[i].alias = alias } + al.maxRi = calcMax(uint32(len(al.table))) + return nil } From 9bc51d94857827ac3ff17d0f28cd0e3a61cd972f Mon Sep 17 00:00:00 2001 From: Boris Nagaev Date: Sat, 11 Mar 2017 09:04:17 +0100 Subject: [PATCH 4/4] add construction of Alias from integer weights The original Alias is based on floats, which can result in biased distributions for some inputs because of float errors. The new Alias variant returns results with probabilities that are exactly proportional to input weights. In int based algo number 1.0 was replaced with an average weight (avg). If the average weight of input array is not an an integer, we add a dummy element with such a weight that the average weight becomes an integer. We do not return the dummy element as a result of Gen(): if it is about to be returned, we jump to the beginning of Gen() method using goto. For float based Alias average weight is always (1<<31)-1. The random value rj (used inside a piece to decide whether to switch to an alias) is now in range [0, avg). Member field divRj stores avg. Member maxRj stores maximum allowed value of the second RNG. (The second RNG provides input for the choice inside piece.) New kind of Alias is generated with NewInt() function. Both of alias kinds use the same struct Alias. New fields were added to this structto support integer based alias: dummy, maxRj and divRj. They have fixed values in float based Alias. Operation of method Gen() was slowed down by this change, but for float based Alias this effect was small: divRj is a power a 2, so modulus operation is fast; maxRj is the maximum possible int32, so it never causes rethrowing; dummy is always set to a value that can not be returned, so it also never causes rethrowing. Int based Alias operates slower, mostly due to rethrowing. Probability of rethrowing can be decreased by using small input values (this effect was demonstrated in the benchmark). Benchmark results on the same machine. Before: BenchmarkGen5-4 100000000 17.5 ns/op BenchmarkGen50-4 100000000 21.9 ns/op BenchmarkGen500-4 100000000 20.7 ns/op BenchmarkGen5000-4 100000000 21.6 ns/op BenchmarkGen50000-4 50000000 24.6 ns/op After: BenchmarkGen5-4 100000000 19.7 ns/op BenchmarkGen50-4 50000000 23.8 ns/op BenchmarkGen500-4 100000000 22.6 ns/op BenchmarkGen5000-4 100000000 23.3 ns/op BenchmarkGen50000-4 50000000 25.5 ns/op BenchmarkGenIntLargeP5-4 50000000 42.4 ns/op BenchmarkGenIntLargeP50-4 50000000 23.7 ns/op BenchmarkGenIntLargeP500-4 30000000 41.6 ns/op BenchmarkGenIntLargeP5000-4 30000000 42.9 ns/op BenchmarkGenIntLargeP50000-4 50000000 44.4 ns/op BenchmarkGenIntSmallP5-4 50000000 31.0 ns/op BenchmarkGenIntSmallP50-4 50000000 25.7 ns/op BenchmarkGenIntSmallP500-4 50000000 25.6 ns/op BenchmarkGenIntSmallP5000-4 50000000 24.1 ns/op BenchmarkGenIntSmallP50000-4 50000000 29.2 ns/op --- alias.go | 169 ++++++++++++++++++++++++++++++++++++++++++++++++-- alias_test.go | 100 ++++++++++++++++++++++++++--- bench_test.go | 89 ++++++++++++++++++++++++++ 3 files changed, 343 insertions(+), 15 deletions(-) diff --git a/alias.go b/alias.go index a456fd9..2de71d2 100644 --- a/alias.go +++ b/alias.go @@ -15,6 +15,9 @@ import ( type Alias struct { table []ipiece maxRi uint32 + maxRj uint32 + avgP uint32 + dummy uint32 } type fpiece struct { @@ -32,6 +35,17 @@ func calcMax(n uint32) uint32 { return (1 << 31) - 1 - (1<<31)%uint32(n) } +// checkAvgP checks assumption that piece with prob = al.avgP - 1 exists. +// This assumption is used in UnmarshalBinary. +func checkAvgP(al *Alias) { + for _, p := range al.table { + if p.prob == al.avgP-1 { + return + } + } + panic("Internal error: no piece with prob = al.avgP-1 found.") +} + // Create a new alias object. // For example, // var v = alias.New([]float64{8,10,2}) @@ -127,7 +141,124 @@ func New(prob []float64) (*Alias, error) { al.table[twins[i].alias].prob = 1<<31 - 1 } + al.avgP = (1 << 31) + al.maxRi = calcMax(uint32(len(al.table))) + al.maxRj = calcMax(al.avgP) + al.dummy = uint32(len(al.table)) + + checkAvgP(&al) + + return &al, nil +} + +// Create a new alias object with integer weights. +// For example, +// var v = alias.NewInt([]int32{8,10,2}) +// creates an alias that returns 0 40% of the time, 1 50% of the time, and +// 2 10% of the time. +func NewInt(prob []int32) (*Alias, error) { + + // This implementation is based on + // http://www.keithschwarz.com/darts-dice-coins/ + + n := len(prob) + + if n < 1 { + return nil, errors.New("too few probabilities") + } + + if int(uint32(n)) != n { + return nil, errors.New("too many probabilities") + } + + total := uint64(0) + for _, v := range prob { + if v <= 0 { + return nil, errors.New("a probability is non-positive") + } + total += uint64(v) + } + + var al Alias + al.dummy = uint32(n) + + avg := uint32(total / uint64(n)) + if uint64(avg)*uint64(n) < total { + // Really add dummy. Its index is already set to al.dummy. + avg += 1 + n += 1 + dummy := uint64(avg)*uint64(n) - total + total += dummy + prob = append(prob, int32(dummy)) + } + + // Michael Vose's algorithm + + // "small" stack grows from the bottom of this array + // "large" stack from the top + twins := make([]ipiece, n) + + smTop := -1 + lgBot := n + + // invariant: smTop < lgBot, that is, the twin stacks don't collide + + for i, p := range prob { + + // push large items (>=1 probability) into the large stack + // others in the small stack + if uint32(p) >= avg { + lgBot-- + twins[lgBot] = ipiece{uint32(p), uint32(i)} + } else { + smTop++ + twins[smTop] = ipiece{uint32(p), uint32(i)} + } + } + + if lgBot != smTop+1 { + panic("alias.New: internal error") + } + + al.table = make([]ipiece, n) + + for smTop >= 0 && lgBot < n { + // pair off a small and large block, taking the chunk from the large block that's wanted + l := twins[smTop] + smTop-- + + g := twins[lgBot] + lgBot++ + + al.table[l.alias].prob = l.prob - 1 + al.table[l.alias].alias = g.alias + + g.prob = (g.prob + l.prob) - avg + + // put the rest of the large block back in a list + if g.prob < avg { + smTop++ + twins[smTop] = g + } else { + lgBot-- + twins[lgBot] = g + } + } + + // clear out any remaining blocks + for i := n - 1; i >= lgBot; i-- { + al.table[twins[i].alias].prob = avg - 1 + } + + if smTop != -1 { + panic("alias.NewInt: internal error") + } + + al.avgP = avg al.maxRi = calcMax(uint32(n)) + al.maxRj = calcMax(al.avgP) + + checkAvgP(&al) return &al, nil } @@ -138,30 +269,39 @@ begin: r := rng.Int63() ri := uint32(r & (1<<31 - 1)) rj := uint32((r >> 31) & (1<<31 - 1)) - if ri > al.maxRi { + if ri > al.maxRi || rj > al.maxRj { goto begin } w := ri % uint32(len(al.table)) - if rj > al.table[w].prob { - return al.table[w].alias + x := rj % al.avgP + if x > al.table[w].prob { + w = al.table[w].alias + } + if w == al.dummy { + goto begin } return w } // MarshalBinary implements encoding.BinaryMarshaller. func (al *Alias) MarshalBinary() ([]byte, error) { - out := make([]byte, len(al.table)*8) + out := make([]byte, len(al.table)*8, len(al.table)*8+4) for i, piece := range al.table { bin := out[i*8 : 8+i*8] binary.LittleEndian.PutUint32(bin[0:4], piece.prob) binary.LittleEndian.PutUint32(bin[4:8], piece.alias) } + if al.dummy != uint32(len(al.table)) { + dummy := make([]byte, 4) + binary.LittleEndian.PutUint32(dummy, al.dummy) + out = append(out, dummy...) + } return out, nil } // UnmarshalBinary implements encoding.BinaryUnmarshaller. func (al *Alias) UnmarshalBinary(p []byte) error { - if len(p)%8 != 0 { + if len(p)%4 != 0 { return errors.New("bad data length") } @@ -186,7 +326,26 @@ func (al *Alias) UnmarshalBinary(p []byte) error { al.table[i].alias = alias } + // How we calculate avgP... There must be at least one piece with + // probability of 1.0. TODO: prove it. In our model such a probability + // is represented as avgP - 1. (-1 comes from the way we compare prob + // with random.) So avgP = max(prob) + 1. + maxProb := uint32(0) + for _, piece := range al.table { + if piece.prob > maxProb { + maxProb = piece.prob + } + } + al.avgP = maxProb + 1 + al.maxRi = calcMax(uint32(len(al.table))) + al.maxRj = calcMax(al.avgP) + + al.dummy = uint32(len(al.table)) + if len(p)%8 != 0 { + dummy := p[len(p)-4:] + al.dummy = binary.LittleEndian.Uint32(dummy) + } return nil } diff --git a/alias_test.go b/alias_test.go index b2a1daf..6e380fe 100644 --- a/alias_test.go +++ b/alias_test.go @@ -5,6 +5,7 @@ package alias import ( + "fmt" "math" "math/rand" "reflect" @@ -111,30 +112,109 @@ func TestBalanceInsideBucket(t *testing.T) { } func TestMarshalBinary(t *testing.T) { - distributions := [][]float64{ - {1}, - {1, 1}, - {1, 2, 3, 4, 5, 6, 7, 8, 9, 1000}, + makeFloat := func(p []float64) *Alias { + a, err := New(p) + if err != nil { + t.Fatalf("Couldn't create alias: %v", err) + } + return a } - for _, distribution := range distributions { - a, err := New(distribution) + makeInt := func(p []int32) *Alias { + a, err := NewInt(p) if err != nil { t.Fatalf("Couldn't create alias: %v", err) } - + return a + } + aliases := []*Alias{ + makeFloat([]float64{1}), + makeFloat([]float64{1, 1}), + makeFloat([]float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 1000}), + makeInt([]int32{1}), + makeInt([]int32{1, 1}), + makeInt([]int32{1, 2, 3, 4, 5, 6, 7, 8, 9, 1000}), + } + for i, a := range aliases { data, err := a.MarshalBinary() if err != nil { - t.Fatalf("Couldn't MarshalBinary: %v", err) + t.Errorf("Couldn't MarshalBinary: %v", err) } a2 := &Alias{} err = a2.UnmarshalBinary(data) if err != nil { - t.Fatalf("Couldn't UnmarshalBinary: %v", err) + t.Errorf("Couldn't UnmarshalBinary: %v", err) } if !reflect.DeepEqual(a, a2) { - t.Fatalf("Unmarshalled version was not the same as original") + fmt.Println(a, a2) + t.Errorf("case %d: Unmarshalled version %v was not the same as original %v", i, a2, a) } } } + +func testIntDistribution(t *testing.T, dist []int32, seed int64) { + sum := uint64(0) + for i := 0; i < len(dist); i++ { + sum += uint64(dist[i]) + } + + a, err := NewInt(dist) + if err != nil { + t.Error("Got an error during creation:", err) + return + } + + rng := rand.New(rand.NewSource(seed)) + + counts := make([]int64, len(dist)) + for i := 0; i < distributionCount; i++ { + counts[a.Gen(rng)]++ + } + + for i := 0; i < len(dist); i++ { + p := float64(counts[i]) / distributionCount + if math.Abs(p-float64(dist[i])/float64(sum)) > errorBound { + t.Error("Distribution did not match, seed", seed, "- got ", p, "expected", float64(dist[i])/float64(sum)) + } + } +} + +func TestIntDistribution(t *testing.T) { + testIntDistribution(t, []int32{1, 1}, 1) + testIntDistribution(t, []int32{1, 2, 3}, 2) + testIntDistribution(t, []int32{9, 8, 1, 4, 2}, 5) + testIntDistribution(t, []int32{1000, 1, 3, 10}, 39) + testIntDistribution(t, []int32{1000, 1, 3, 10}, 61) +} + +func TestIntBalanceInsideBucket(t *testing.T) { + const size = 33294320 + const half = size / 2 + const tries = 1000000 + const alpha = 0.05 + dist := make([]int32, size) + for i := range dist { + if i < half { + dist[i] = 1 + } else { + dist[i] = 3 + } + } + a, err := NewInt(dist) + if err != nil { + t.Fatalf("Got an error during creation:", err) + } + rng := rand.New(rand.NewSource(421)) + var k int64 // [0,half) (one half) + for i := 0; i < tries; i++ { + if a.Gen(rng) < half { + k++ + } + } + // Expected probability of getting k <= k_observed if p == 0.5. + p := stat.Binomial_CDF_At(0.25, tries, k) + if p < alpha/2 || p > (1-alpha/2) { + t.Errorf("The distribution is biased. %d of %d results were in the first half. Binomial_CDF = %f", k, tries, p) + } +} diff --git a/bench_test.go b/bench_test.go index 0360748..ac6747c 100644 --- a/bench_test.go +++ b/bench_test.go @@ -85,3 +85,92 @@ func BenchmarkCreate5000(b *testing.B) { func BenchmarkCreate50000(b *testing.B) { benchCreationSize(b, 50000) } + +func benchGenInt(b *testing.B, size int, maxP int32) { + b.StopTimer() + arr := make([]int32, size) + for i := 0; i < size; i++ { + arr[i] = rand.Int31n(maxP) + } + a, err := NewInt(arr) + if err != nil { + b.Error("Got an error during creation:", err) + } + rng := rand.New(rand.NewSource(99)) + b.StartTimer() + for i := 0; i < b.N; i++ { + a.Gen(rng) + } +} + +func BenchmarkGenIntLargeP5(b *testing.B) { + benchGenInt(b, 5, (1<<31)-1) +} + +func BenchmarkGenIntLargeP50(b *testing.B) { + benchGenInt(b, 50, (1<<31)-1) +} + +func BenchmarkGenIntLargeP500(b *testing.B) { + benchGenInt(b, 500, (1<<31)-1) +} + +func BenchmarkGenIntLargeP5000(b *testing.B) { + benchGenInt(b, 5000, (1<<31)-1) +} + +func BenchmarkGenIntLargeP50000(b *testing.B) { + benchGenInt(b, 50000, (1<<31)-1) +} + +func BenchmarkGenIntSmallP5(b *testing.B) { + benchGenInt(b, 5, 1<<29) +} + +func BenchmarkGenIntSmallP50(b *testing.B) { + benchGenInt(b, 50, 1<<29) +} + +func BenchmarkGenIntSmallP500(b *testing.B) { + benchGenInt(b, 500, 1<<29) +} + +func BenchmarkGenIntSmallP5000(b *testing.B) { + benchGenInt(b, 5000, 1<<29) +} + +func BenchmarkGenIntSmallP50000(b *testing.B) { + benchGenInt(b, 50000, 1<<29) +} + +func benchCreationSizeInt(b *testing.B, size int) { + b.StopTimer() + arr := make([]int32, size) + for i := 0; i < size; i++ { + arr[i] = rand.Int31() + } + b.StartTimer() + for i := 0; i < b.N; i++ { + NewInt(arr) + } +} + +func BenchmarkCreateInt5(b *testing.B) { + benchCreationSizeInt(b, 5) +} + +func BenchmarkCreateInt50(b *testing.B) { + benchCreationSizeInt(b, 50) +} + +func BenchmarkCreateInt500(b *testing.B) { + benchCreationSizeInt(b, 500) +} + +func BenchmarkCreateInt5000(b *testing.B) { + benchCreationSizeInt(b, 5000) +} + +func BenchmarkCreateInt50000(b *testing.B) { + benchCreationSizeInt(b, 50000) +}