github.com/petermattis/pebble@v0.0.0-20190905164901-ab51a2166067/internal/randvar/zipf_test.go (about)

     1  // Copyright 2017 The Cockroach Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
    12  // implied. See the License for the specific language governing
    13  // permissions and limitations under the License. See the AUTHORS file
    14  // for names of contributors.
    15  
    16  package randvar
    17  
    18  import (
    19  	"math"
    20  	"testing"
    21  )
    22  
    23  func TestZeta(t *testing.T) {
    24  	var zetaTests = []struct {
    25  		n        uint64
    26  		theta    float64
    27  		expected float64
    28  	}{
    29  		{20, 0.99, 3.64309060779367},
    30  		{200, 0.99, 6.02031118558},
    31  		{1000, 0.99, 7.72895321728},
    32  		{2000, 0.99, 8.47398788329},
    33  		{10000, 0.99, 10.2243614596},
    34  		{100000, 0.99, 12.7783380626},
    35  		{1000000, 0.99, 15.391849746},
    36  		{10000000, 0.99, 18.066242575},
    37  		// TODO(peter): The last test case takes an excessively long time to run (7s).
    38  		// {100000000, 0.99, 20.80293049},
    39  	}
    40  
    41  	t.Run("FromScratch", func(t *testing.T) {
    42  		for _, test := range zetaTests {
    43  			computedZeta := computeZetaFromScratch(test.n, test.theta)
    44  			if math.Abs(computedZeta-test.expected) > 0.000000001 {
    45  				t.Fatalf("expected %6.4f, got %6.4f", test.expected, computedZeta)
    46  			}
    47  		}
    48  	})
    49  
    50  	t.Run("Incrementally", func(t *testing.T) {
    51  		// Theta cannot be 1 by definition, so this is a safe initial value.
    52  		oldTheta := 1.0
    53  		var oldZetaN float64
    54  		var oldN uint64
    55  		for _, test := range zetaTests {
    56  			// If theta has changed, recompute from scratch
    57  			if test.theta != oldTheta {
    58  				oldZetaN = computeZetaFromScratch(test.n, test.theta)
    59  				oldN = test.n
    60  				continue
    61  			}
    62  
    63  			computedZeta := computeZetaIncrementally(oldN, test.n, test.theta, oldZetaN)
    64  			if math.Abs(computedZeta-test.expected) > 0.000000001 {
    65  				t.Fatalf("expected %6.4f, got %6.4f", test.expected, computedZeta)
    66  			}
    67  
    68  			oldZetaN = computedZeta
    69  			oldN = test.n
    70  		}
    71  	})
    72  }
    73  
    74  func TestZetaIncMax(t *testing.T) {
    75  	// Construct a zipf generator covering the range [0,10] incrementally.
    76  	z0, err := NewZipf(nil, 0, 0, 0.99)
    77  	if err != nil {
    78  		t.Fatal(err)
    79  	}
    80  	for i := 0; i < 10; i++ {
    81  		z0.IncMax(1)
    82  	}
    83  
    84  	// Contruct a zipf generator covering the range [0,10] via the constructor.
    85  	z10, err := NewZipf(nil, 0, 10, 0.99)
    86  	if err != nil {
    87  		t.Fatal(err)
    88  	}
    89  
    90  	z0.mu.Lock()
    91  	defer z0.mu.Unlock()
    92  	z10.mu.Lock()
    93  	defer z10.mu.Unlock()
    94  	if z0.mu.zetaN != z10.mu.zetaN {
    95  		t.Fatalf("expected zetaN %v, but found %v", z10.mu.zetaN, z0.mu.zetaN)
    96  	}
    97  	if z0.mu.eta != z10.mu.eta {
    98  		t.Fatalf("expected eta %v, but found %v", z10.mu.eta, z0.mu.eta)
    99  	}
   100  }
   101  
   102  func TestNewZipf(t *testing.T) {
   103  	var gens = []struct {
   104  		min, max uint64
   105  		theta    float64
   106  	}{
   107  		{0, 100, 0.99},
   108  		{0, 100, 1.01},
   109  	}
   110  
   111  	for _, gen := range gens {
   112  		_, err := NewZipf(nil, gen.min, gen.max, gen.theta)
   113  		if err != nil {
   114  			t.Fatal(err)
   115  		}
   116  	}
   117  }
   118  
   119  func TestZipf(t *testing.T) {
   120  	z, err := NewZipf(nil, 0, 99, 0.99)
   121  	if err != nil {
   122  		t.Fatal(err)
   123  	}
   124  
   125  	x := make([]int, 10000)
   126  	for i := range x {
   127  		x[i] = int(z.Uint64())
   128  	}
   129  
   130  	if testing.Verbose() {
   131  		dumpSamples(x)
   132  	}
   133  }