| 85 | |
| 86 | |
| 87 | def alias_setup(probs): |
| 88 | import numpy as np |
| 89 | |
| 90 | """ |
| 91 | Compute utility lists for non-uniform sampling from discrete distributions. |
| 92 | Refer to https://hips.seas.harvard.edu/blog/2013/03/03/the-alias-method-efficient-sampling-with-many-discrete-outcomes/ |
| 93 | for details |
| 94 | """ |
| 95 | K = len(probs) |
| 96 | q = np.zeros(K) |
| 97 | J = np.zeros(K, dtype=int) |
| 98 | |
| 99 | smaller = [] |
| 100 | larger = [] |
| 101 | for kk, prob in enumerate(probs): |
| 102 | q[kk] = K * prob |
| 103 | if q[kk] < 1.0: |
| 104 | smaller.append(kk) |
| 105 | else: |
| 106 | larger.append(kk) |
| 107 | |
| 108 | while len(smaller) > 0 and len(larger) > 0: |
| 109 | small = smaller.pop() |
| 110 | large = larger.pop() |
| 111 | |
| 112 | J[small] = large |
| 113 | q[large] = q[large] + q[small] - 1.0 |
| 114 | if q[large] < 1.0: |
| 115 | smaller.append(large) |
| 116 | else: |
| 117 | larger.append(large) |
| 118 | |
| 119 | return J, q |