samples-fast.py (1471B)
1 import numpy as np 2 rng = np.random.default_rng(123) 3 DEFAULT_N = 1000000 4 5 6 def normal(mean, std, n=DEFAULT_N): 7 return rng.normal(mean, std, n) 8 9 10 def lognormal(mean, std, n=DEFAULT_N): 11 return rng.lognormal(mean, std, n) 12 13 14 def to(low, high, n=DEFAULT_N): 15 normal95confidencePoint = 1.6448536269514722 16 logLow = np.log(low) 17 logHigh = np.log(high) 18 meanlog = (logLow + logHigh) / 2 19 sdlog = (logHigh - logLow) / (2 * normal95confidencePoint) 20 return lognormal(meanlog, sdlog, n) 21 22 23 def optimized_mixture(samples_funcs, weights_array, n=DEFAULT_N): 24 normalized_weights = weights_array / sum(weights_array) 25 cummulative_sums = np.cumsum(normalized_weights) 26 helper_probs = rng.random(n) 27 results = np.empty(n) 28 for i, (start, end) in enumerate(zip([0]+list(cummulative_sums[:-1]), cummulative_sums)): 29 idx = np.where((helper_probs >= start) & (helper_probs < end))[0] 30 # Generate only as many samples as needed for each distribution 31 samples_func = samples_funcs[i] 32 results[idx] = samples_func(n=len(idx)) 33 return results 34 35 36 p_a = 0.8 37 p_b = 0.5 38 p_c = p_a * p_b 39 dists = [ 40 lambda n=1: np.zeros(n), # Distribution returning 0 41 lambda n=1: np.ones(n), # Distribution returning 1 42 lambda n=1: to(1, 3, n), 43 lambda n=1: to(2, 10, n) 44 ] 45 weights = np.array([1 - p_c, p_c/2, p_c/4, p_c/4]) 46 result = optimized_mixture(dists, weights) 47 mean_result = np.mean(result) 48 print(f'Mean result: {mean_result}')