github.com/zuoyebang/bitalostable@v1.0.1-0.20240229032404-e3b99a834294/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  	"github.com/stretchr/testify/require"
    23  )
    24  
    25  func TestZeta(t *testing.T) {
    26  	var zetaTests = []struct {
    27  		n        uint64
    28  		theta    float64
    29  		expected float64
    30  	}{
    31  		{20, 0.99, 3.64309060779367},
    32  		{200, 0.99, 6.02031118558},
    33  		{1000, 0.99, 7.72895321728},
    34  		{2000, 0.99, 8.47398788329},
    35  		{10000, 0.99, 10.2243614596},
    36  		{100000, 0.99, 12.7783380626},
    37  		{1000000, 0.99, 15.391849746},
    38  		{10000000, 0.99, 18.066242575},
    39  		// TODO(peter): The last test case takes an excessively long time to run (7s).
    40  		// {100000000, 0.99, 20.80293049},
    41  	}
    42  
    43  	t.Run("FromScratch", func(t *testing.T) {
    44  		for _, test := range zetaTests {
    45  			computedZeta := computeZetaFromScratch(test.n, test.theta)
    46  			if math.Abs(computedZeta-test.expected) > 0.000000001 {
    47  				t.Fatalf("expected %6.4f, got %6.4f", test.expected, computedZeta)
    48  			}
    49  		}
    50  	})
    51  
    52  	t.Run("Incrementally", func(t *testing.T) {
    53  		// Theta cannot be 1 by definition, so this is a safe initial value.
    54  		oldTheta := 1.0
    55  		var oldZetaN float64
    56  		var oldN uint64
    57  		for _, test := range zetaTests {
    58  			// If theta has changed, recompute from scratch
    59  			if test.theta != oldTheta {
    60  				oldZetaN = computeZetaFromScratch(test.n, test.theta)
    61  				oldN = test.n
    62  				continue
    63  			}
    64  
    65  			computedZeta := computeZetaIncrementally(oldN, test.n, test.theta, oldZetaN)
    66  			if math.Abs(computedZeta-test.expected) > 0.000000001 {
    67  				t.Fatalf("expected %6.4f, got %6.4f", test.expected, computedZeta)
    68  			}
    69  
    70  			oldZetaN = computedZeta
    71  			oldN = test.n
    72  		}
    73  	})
    74  }
    75  
    76  func TestZetaIncMax(t *testing.T) {
    77  	// Construct a zipf generator covering the range [0,10] incrementally.
    78  	z0, err := NewZipf(0, 0, 0.99)
    79  	require.NoError(t, err)
    80  
    81  	for i := 0; i < 10; i++ {
    82  		z0.IncMax(1)
    83  	}
    84  
    85  	// Construct a zipf generator covering the range [0,10] via the constructor.
    86  	z10, err := NewZipf(0, 10, 0.99)
    87  	require.NoError(t, err)
    88  
    89  	z0.mu.Lock()
    90  	defer z0.mu.Unlock()
    91  	z10.mu.Lock()
    92  	defer z10.mu.Unlock()
    93  	if z0.mu.zetaN != z10.mu.zetaN {
    94  		t.Fatalf("expected zetaN %v, but found %v", z10.mu.zetaN, z0.mu.zetaN)
    95  	}
    96  	if z0.mu.eta != z10.mu.eta {
    97  		t.Fatalf("expected eta %v, but found %v", z10.mu.eta, z0.mu.eta)
    98  	}
    99  }
   100  
   101  func TestNewZipf(t *testing.T) {
   102  	var gens = []struct {
   103  		min, max uint64
   104  		theta    float64
   105  	}{
   106  		{0, 100, 0.99},
   107  		{0, 100, 1.01},
   108  	}
   109  
   110  	for _, gen := range gens {
   111  		_, err := NewZipf(gen.min, gen.max, gen.theta)
   112  		require.NoError(t, err)
   113  	}
   114  }
   115  
   116  func TestZipf(t *testing.T) {
   117  	rng := NewRand()
   118  	z, err := NewZipf(0, 99, 0.99)
   119  	require.NoError(t, err)
   120  
   121  	x := make([]int, 10000)
   122  	for i := range x {
   123  		x[i] = int(z.Uint64(rng))
   124  	}
   125  
   126  	if testing.Verbose() {
   127  		dumpSamples(x)
   128  	}
   129  }