github.com/MetalBlockchain/metalgo@v1.11.9/utils/sampler/weighted_benchmark_test.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package sampler
     5  
     6  import (
     7  	"fmt"
     8  	"math"
     9  	"testing"
    10  
    11  	safemath "github.com/MetalBlockchain/metalgo/utils/math"
    12  )
    13  
    14  // BenchmarkAllWeightedSampling
    15  func BenchmarkAllWeightedSampling(b *testing.B) {
    16  	pows := []float64{
    17  		0,
    18  		1,
    19  		2,
    20  		3,
    21  	}
    22  	sizes := []int{
    23  		10,
    24  		500,
    25  		1000,
    26  		50000,
    27  		100000,
    28  	}
    29  	for _, s := range weightedSamplers {
    30  		for _, pow := range pows {
    31  			for _, size := range sizes {
    32  				if WeightedPowBenchmarkSampler(b, s.sampler, pow, size) {
    33  					b.Run(fmt.Sprintf("sampler %s with %d elements at x^%.1f", s.name, size, pow), func(b *testing.B) {
    34  						WeightedPowBenchmarkSampler(b, s.sampler, pow, size)
    35  					})
    36  				}
    37  			}
    38  		}
    39  		for _, size := range sizes {
    40  			if WeightedSingletonBenchmarkSampler(b, s.sampler, size) {
    41  				b.Run(fmt.Sprintf("sampler %s with %d singleton elements", s.name, size), func(b *testing.B) {
    42  					WeightedSingletonBenchmarkSampler(b, s.sampler, size)
    43  				})
    44  			}
    45  		}
    46  	}
    47  }
    48  
    49  // BenchmarkAllWeightedInitializer
    50  func BenchmarkAllWeightedInitializer(b *testing.B) {
    51  	pows := []float64{
    52  		0,
    53  		1,
    54  		2,
    55  		3,
    56  	}
    57  	sizes := []int{
    58  		10,
    59  		500,
    60  		1000,
    61  		50000,
    62  		100000,
    63  	}
    64  	for _, s := range weightedSamplers {
    65  		for _, pow := range pows {
    66  			for _, size := range sizes {
    67  				if WeightedPowBenchmarkSampler(b, s.sampler, pow, size) {
    68  					b.Run(fmt.Sprintf("initializer %s with %d elements at x^%.1f", s.name, size, pow), func(b *testing.B) {
    69  						WeightedPowBenchmarkInitializer(b, s.sampler, pow, size)
    70  					})
    71  				}
    72  			}
    73  		}
    74  		for _, size := range sizes {
    75  			if WeightedSingletonBenchmarkSampler(b, s.sampler, size) {
    76  				b.Run(fmt.Sprintf("initializer %s with %d singleton elements", s.name, size), func(b *testing.B) {
    77  					WeightedSingletonBenchmarkInitializer(b, s.sampler, size)
    78  				})
    79  			}
    80  		}
    81  	}
    82  }
    83  
    84  func CalcWeightedPoW(exponent float64, size int) (uint64, []uint64, error) {
    85  	weights := make([]uint64, size)
    86  	totalWeight := uint64(0)
    87  	for i := range weights {
    88  		weight := uint64(math.Pow(float64(i+1), exponent))
    89  		weights[i] = weight
    90  
    91  		newWeight, err := safemath.Add64(totalWeight, weight)
    92  		if err != nil {
    93  			return 0, nil, err
    94  		}
    95  		totalWeight = newWeight
    96  	}
    97  	return totalWeight, weights, nil
    98  }
    99  
   100  func WeightedPowBenchmarkSampler(
   101  	b *testing.B,
   102  	s Weighted,
   103  	exponent float64,
   104  	size int,
   105  ) bool {
   106  	totalWeight, weights, err := CalcWeightedPoW(exponent, size)
   107  	if err != nil {
   108  		return false
   109  	}
   110  	if err := s.Initialize(weights); err != nil {
   111  		return false
   112  	}
   113  
   114  	b.ResetTimer()
   115  	for i := 0; i < b.N; i++ {
   116  		_, _ = s.Sample(globalRNG.Uint64Inclusive(totalWeight - 1))
   117  	}
   118  	return true
   119  }
   120  
   121  func WeightedSingletonBenchmarkSampler(b *testing.B, s Weighted, size int) bool {
   122  	weights := make([]uint64, size)
   123  	weights[0] = math.MaxUint64 - uint64(size-1)
   124  	for i := 1; i < len(weights); i++ {
   125  		weights[i] = 1
   126  	}
   127  
   128  	err := s.Initialize(weights)
   129  	if err != nil {
   130  		return false
   131  	}
   132  
   133  	b.ResetTimer()
   134  	for i := 0; i < b.N; i++ {
   135  		_, _ = s.Sample(globalRNG.Uint64Inclusive(math.MaxUint64 - 1))
   136  	}
   137  	return true
   138  }
   139  
   140  func WeightedPowBenchmarkInitializer(
   141  	b *testing.B,
   142  	s Weighted,
   143  	exponent float64,
   144  	size int,
   145  ) {
   146  	_, weights, _ := CalcWeightedPoW(exponent, size)
   147  
   148  	b.ResetTimer()
   149  	for i := 0; i < b.N; i++ {
   150  		_ = s.Initialize(weights)
   151  	}
   152  }
   153  
   154  func WeightedSingletonBenchmarkInitializer(b *testing.B, s Weighted, size int) {
   155  	weights := make([]uint64, size)
   156  	weights[0] = math.MaxUint64 - uint64(size-1)
   157  	for i := 1; i < len(weights); i++ {
   158  		weights[i] = 1
   159  	}
   160  
   161  	b.ResetTimer()
   162  	for i := 0; i < b.N; i++ {
   163  		_ = s.Initialize(weights)
   164  	}
   165  }