<p>In general, there are <a href="https://stackoverflow.com/questions/5027757/data-structure-for-loaded-dice/63166311#63166311">many ways to choose an integer</a> with a custom distribution, but most of them take weights that are <em>proportional</em> to the given probabilities. If the weights are log probabilities instead, then a slightly different approach is needed. Perhaps the simplest algorithm for this is rejection sampling, described below and implemented in Python. In the following algorithm, the maximum log-probability is <code>max</code>, and there are <code>k</code> integers to choose from.</p>
<ol>
<li>Choose a uniform random integer <code>i</code> in [0, <code>k</code>).</li>
<li>Get the log-weight corresponding to <code>i</code>, then generate an exponential(1) random number, call it <code>ex</code>.</li>
<li>If <code>max</code> minus <code>ex</code> is less than the log-weight, return <code>i</code>. Otherwise, go to step 1.</li>
</ol>
<p>The time complexity for rejection sampling is constant on average, especially if <code>max</code> is set to equal the true maximum weight. On the other hand, the expected number of iterations per sample depends greatly on the shape of the distribution. See also <a href="https://www.keithschwarz.com/darts-dice-coins/" rel="nofollow noreferrer">Keith Schwarz's discussion</a> on the &quot;Fair Die/Biased Coin Loaded Die&quot; algorithm.</p>
<p>Now, Python code for this algorithm follows.</p>
<pre><code>import random
import math

def categ(c):
 # Do a weighted choice of an item with the
 # given log-probabilities.
 cm=max(c) # Find max log probability
 while True:
      # Choose an item at random
      x=random.randint(0,len(c)-1)
      # Choose it with probability proportional
      # to exp(c[x])
      y=cm-random.expovariate(1)
      # Alternatively: y=math.log(random.random())+cm
      if y&lt;c[x]:
          return x
</code></pre>
<p>The code above generates one variate at a time and uses only Python's base modules, rather than NumPy. <a href="https://stackoverflow.com/a/64881410/815724">Another answer</a> shows how rejection sampling can be implemented in NumPy by blocks of random variates at a time (demonstrated on a different random sampling task, though).</p>
<hr />
<p>The so-called &quot;<a href="https://arxiv.org/abs/2110.01515" rel="nofollow noreferrer">Gumbel max trick</a>&quot;, used above all in machine learning, can be used to sample from a distribution with unnormalized log probabilities. This involves—</p>
<ol>
<li>(&quot;Gumbel&quot;) adding a separate <em>Gumbel</em> random variate to each log probability, namely −ln(−ln(<em>U</em>)) where <em>U</em> is a random variate greater than 0 and less than 1, then</li>
<li>(&quot;max&quot;) choosing the item corresponding to the highest log probability.</li>
</ol>
<p>However, the time complexity for this algorithm is linear in the number of items.</p>
<p>The following code illustrates the Gumbel max trick:</p>
<pre><code>import random
import math

def categ(c):
 # Do a weighted choice of an item with the
 # given log-probabilities, using the Gumbel max trick
 return max([[c[i]-math.log(-math.log(random.random())),i] \
      for i in range(len(c))])[1]
 # Or:
 # return max([[c[i]-math.log(random.expovariate(1)),i] \
 #   for i in range(len(c))])[1]
</code></pre>
Sampling from a set according to unnormalized log-probabilities in NumPy
<p>I have a 1-D <code>np.ndarray</code> filled with unnormalized log-probabilities that define a categorical distribution. I would like to sample an integer index from this distribution. Since many of the probabilities are small, normalizing and exponentiating the log-probabilities introduces significant numerical error, therefore I cannot use <a href="https://numpy.org/doc/stable/reference/random/generated/numpy.random.choice.html" rel="nofollow noreferrer"><code>np.random.choice</code></a>. Effectively, I am looking for a NumPy equivalent to TensorFlow's <a href="https://www.tensorflow.org/api_docs/python/tf/random/categorical" rel="nofollow noreferrer"><code>tf.random.categorical</code></a>, which works on unnormalized log-probabilities.</p>
<p>If there is not a function in NumPy that achieves this directly, what is an efficient manner to implement such sampling?</p>