diff --git a/alias.go b/alias.go index 5ec71c3..2de71d2 100644 --- a/alias.go +++ b/alias.go @@ -14,6 +14,10 @@ import ( type Alias struct { table []ipiece + maxRi uint32 + maxRj uint32 + avgP uint32 + dummy uint32 } type fpiece struct { @@ -26,6 +30,22 @@ 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) +} + +// 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}) @@ -83,6 +103,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] @@ -117,33 +141,167 @@ 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 } // Generates a random number according to the distribution using the rng passed. func (al *Alias) Gen(rng *rand.Rand) uint32 { - ri := uint32(rng.Int31()) +begin: + r := rng.Int63() + ri := uint32(r & (1<<31 - 1)) + rj := uint32((r >> 31) & (1<<31 - 1)) + if ri > al.maxRi || rj > al.maxRj { + goto begin + } w := ri % uint32(len(al.table)) - if ri > 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") } @@ -168,5 +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 fd0f686..6e380fe 100644 --- a/alias_test.go +++ b/alias_test.go @@ -5,10 +5,13 @@ package alias import ( + "fmt" "math" "math/rand" "reflect" "testing" + + stat "github.com/ematvey/gostat" ) const distributionCount = 1000000 @@ -49,31 +52,169 @@ 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}, - {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) +}