squiggle.go (3522B)
1 package main 2 3 import "fmt" 4 import "math" 5 import "sync" 6 import rand "math/rand/v2" 7 8 // https://pkg.go.dev/math/rand/v2 9 10 type src = *rand.Rand 11 type func64 = func(src) float64 12 13 func sample_unit_uniform(r src) float64 { 14 return r.Float64() 15 } 16 17 func sample_unit_normal(r src) float64 { 18 return r.NormFloat64() 19 } 20 21 func sample_uniform(start float64, end float64, r src) float64 { 22 return sample_unit_uniform(r)*(end-start) + start 23 } 24 25 func sample_normal(mean float64, sigma float64, r src) float64 { 26 return mean + sample_unit_normal(r)*sigma 27 } 28 29 func sample_lognormal(logmean float64, logstd float64, r src) float64 { 30 return (math.Exp(sample_normal(logmean, logstd, r))) 31 } 32 33 func sample_normal_from_90_ci(low float64, high float64, r src) float64 { 34 var normal90 float64 = 1.6448536269514727 35 var mean float64 = (high + low) / 2.0 36 var std float64 = (high - low) / (2.0 * normal90) 37 return sample_normal(mean, std, r) 38 39 } 40 41 func sample_to(low float64, high float64, r src) float64 { 42 // Given a (positive) 90% confidence interval, 43 // returns a sample from a lognorma with a matching 90% c.i. 44 // Key idea: If we want a lognormal with 90% confidence interval [a, b] 45 // we need but get a normal with 90% confidence interval [log(a), log(b)]. 46 // Then see code for sample_normal_from_90_ci 47 var loglow float64 = math.Log(low) 48 var loghigh float64 = math.Log(high) 49 return math.Exp(sample_normal_from_90_ci(loglow, loghigh, r)) 50 } 51 52 func sample_mixture(fs []func64, weights []float64, r src) float64 { 53 54 // fmt.Println("weights initially: ", weights) 55 var sum_weights float64 = 0 56 for _, weight := range weights { 57 sum_weights += weight 58 } 59 60 var total float64 = 0 61 var cumsummed_normalized_weights = append([]float64(nil), weights...) 62 for i, weight := range weights { 63 total += weight / sum_weights 64 cumsummed_normalized_weights[i] = total 65 } 66 67 var result float64 68 var flag int = 0 69 var p float64 = r.Float64() 70 71 for i, cnw := range cumsummed_normalized_weights { 72 if p < cnw { 73 result = fs[i](r) 74 flag = 1 75 break 76 } 77 } 78 79 if flag == 0 { 80 result = fs[len(fs)-1](r) 81 } 82 return result 83 84 } 85 86 func sample_parallel(f func64, n_samples int) []float64 { 87 var num_threads = 16 88 var xs = make([]float64, n_samples) 89 var wg sync.WaitGroup 90 var h = n_samples / num_threads 91 wg.Add(num_threads) 92 for i := range num_threads { 93 var xs_i = xs[i*h : (i+1)*h] 94 go func(f func64) { 95 defer wg.Done() 96 var r = rand.New(rand.NewPCG(uint64(i), uint64(i+1))) 97 for i := range xs_i { 98 xs_i[i] = f(r) 99 } 100 }(f) 101 } 102 103 wg.Wait() 104 return xs 105 106 } 107 108 func main() { 109 110 var p_a float64 = 0.8 111 var p_b float64 = 0.5 112 var p_c float64 = p_a * p_b 113 ws := [4](float64){1 - p_c, p_c / 2, p_c / 4, p_c / 4} 114 115 sample_0 := func(r src) float64 { return 0 } 116 sample_1 := func(r src) float64 { return 1 } 117 sample_few := func(r src) float64 { return sample_to(1, 3, r) } 118 sample_many := func(r src) float64 { return sample_to(2, 10, r) } 119 fs := [4](func64){sample_0, sample_1, sample_few, sample_many} 120 121 model := func(r src) float64 { return sample_mixture(fs[0:], ws[0:], r) } 122 n_samples := 1_000_000 123 xs := sample_parallel(model, n_samples) 124 var avg float64 = 0 125 for _, x := range xs { 126 avg += x 127 } 128 avg = avg / float64(n_samples) 129 fmt.Printf("Average: %v\n", avg) 130 /* 131 // Without concurrency: 132 n_samples := 1_000_000 133 var r = rand.New(rand.NewPCG(uint64(1), uint64(2))) 134 var avg float64 = 0 135 for i := 0; i < n_samples; i++ { 136 avg += sample_mixture(fs[0:], ws[0:], r) 137 } 138 avg = avg / float64(n_samples) 139 fmt.Printf("Average: %v\n", avg) 140 */ 141 }