ceremonyclient/pebble/internal/randvar/zipf_test.go

// Copyright 2017 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License. See the AUTHORS file
// for names of contributors.

package randvar

import (
	"math"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestZeta(t *testing.T) {
	var zetaTests = []struct {
		n        uint64
		theta    float64
		expected float64
	}{
		{20, 0.99, 3.64309060779367},
		{200, 0.99, 6.02031118558},
		{1000, 0.99, 7.72895321728},
		{2000, 0.99, 8.47398788329},
		{10000, 0.99, 10.2243614596},
		{100000, 0.99, 12.7783380626},
		{1000000, 0.99, 15.391849746},
		{10000000, 0.99, 18.066242575},
		// TODO(peter): The last test case takes an excessively long time to run (7s).
		// {100000000, 0.99, 20.80293049},
	}

	t.Run("FromScratch", func(t *testing.T) {
		for _, test := range zetaTests {
			computedZeta := computeZetaFromScratch(test.n, test.theta)
			if math.Abs(computedZeta-test.expected) > 0.000000001 {
				t.Fatalf("expected %6.4f, got %6.4f", test.expected, computedZeta)
			}
		}
	})

	t.Run("Incrementally", func(t *testing.T) {
		// Theta cannot be 1 by definition, so this is a safe initial value.
		oldTheta := 1.0
		var oldZetaN float64
		var oldN uint64
		for _, test := range zetaTests {
			// If theta has changed, recompute from scratch
			if test.theta != oldTheta {
				oldZetaN = computeZetaFromScratch(test.n, test.theta)
				oldN = test.n
				continue
			}

			computedZeta := computeZetaIncrementally(oldN, test.n, test.theta, oldZetaN)
			if math.Abs(computedZeta-test.expected) > 0.000000001 {
				t.Fatalf("expected %6.4f, got %6.4f", test.expected, computedZeta)
			}

			oldZetaN = computedZeta
			oldN = test.n
		}
	})
}

func TestZetaIncMax(t *testing.T) {
	// Construct a zipf generator covering the range [0,10] incrementally.
	z0, err := NewZipf(0, 0, 0.99)
	require.NoError(t, err)

	for i := 0; i < 10; i++ {
		z0.IncMax(1)
	}

	// Construct a zipf generator covering the range [0,10] via the constructor.
	z10, err := NewZipf(0, 10, 0.99)
	require.NoError(t, err)

	z0.mu.Lock()
	defer z0.mu.Unlock()
	z10.mu.Lock()
	defer z10.mu.Unlock()
	if z0.mu.zetaN != z10.mu.zetaN {
		t.Fatalf("expected zetaN %v, but found %v", z10.mu.zetaN, z0.mu.zetaN)
	}
	if z0.mu.eta != z10.mu.eta {
		t.Fatalf("expected eta %v, but found %v", z10.mu.eta, z0.mu.eta)
	}
}

func TestNewZipf(t *testing.T) {
	var gens = []struct {
		min, max uint64
		theta    float64
	}{
		{0, 100, 0.99},
		{0, 100, 1.01},
	}

	for _, gen := range gens {
		_, err := NewZipf(gen.min, gen.max, gen.theta)
		require.NoError(t, err)
	}
}

func TestZipf(t *testing.T) {
	rng := NewRand()
	z, err := NewZipf(0, 99, 0.99)
	require.NoError(t, err)

	x := make([]int, 10000)
	for i := range x {
		x[i] = int(z.Uint64(rng))
	}

	if testing.Verbose() {
		dumpSamples(x)
	}
}