Efficient Sampling of A (Countable) Discrete Distribution

Question

I have a sequence of positive terms $(a_n)$, for which $\sum_{n=1}^\infty a_n = A <\infty$, and wish to take a psuedo-random sample from the discrete probability distribution

\begin{equation} \mathbf{P}[ X = n] = \frac{a_n}{A}. \end{equation}

The standard approach to sampling $X$ is to first sample a (continuous) uniform variable $U \sim \text{Unif}[0,1]$, and then set $X = x_U$, with \begin{equation} x_U = \min \Bigg\{n \, \colon \, \sum_{k=1}^n a_k > U \Bigg\}. \end{equation}

In general if the $a_n$ are arranged in decreasing order of magnitude, then for typical samples of $U$, $x_U$ can be calculated rather quickly. However, given that I want to sample the distribution many times, inevitably I will encounter cases where computing $x_U$ is very time consuming.

I am looking for an efficient implementation of this algorithm in Mathematica. My guess is that there is perhaps a way to do this with recursive functions, saving the totals of the 'large sums' (i.e. those over many $k$) to be used for future calculations.

Does anybody have any suggestions? Thanks!

EDIT To give an example of a distirbution, consider $a_n = \log(n)/n^{(3/2)}$. This sequence is decreasing, and the associated series converges to -Zeta'[3/2], i.e. $A \approx 3.93224$.

On my (admittedly not particularly swift) computer, if I sample $U =0.9$, it takes approximately 1.25 seconds to return $X =2496$. My simple (naive) implementation is as follows:

X[u_] := Module[{A = -Zeta'[3/2]},

k=0;
sumk=0;

While[(sumk/Z < u),
k+=1;
sumk += N[Log[k]/k^(3/2)];
];

k];

Answer 1

If you are willing to precompute some things, it can be pretty quick. Here we precalculate 100000 terms of the $a_n$ sequence. Then calculate the CDF (cumulative distribution function) by using Accumulate. To find the closest term to the u, use a NearestFunction.

capA = -Zeta'[3/2] // N;
aAll = (Log[#]/#^1.5 & /@ Range[100000])/capA;
accAll = Accumulate[aAll];
nf = Nearest[accAll -> Range[Length[accAll]]];

To find the first time the sum is near 0.9, calculate nf[0.9]. Successive values can be done without recomputing the nearest function, so they should be very fast.

Answer 2

This is pretty tricky due to the very heavy tail of the distribution. You can certainly get big speed ups as suggested by bill s by pre-computing some quantiles. However, there will always be a good chunk of the tail left to compute. I'm going to try to address the latter and borrow from Bill's solution for pre-computation.

s[n_] := Log[n]*n^(-1.5);    
A = Evaluate@Sum[s[k], {k, 1, Infinity}];

Since Mathematica can compute this sum in closed form we can use it in computing tail probabilities.

F[n_] := Evaluate[Sum[s[k], {k, 1, n}]]

This code effectively uses the derivative of the CDF to adaptively step to the next term. When the cdf is very flat and we are still a good distance from the desired quantile it steps quite far which avoids unnecessary sums. The algorithm converges when it can no longer take an integer step.

f[k_, q_] := If[#2 == #1, 0, Ceiling[(A*q - #1)/(#2 - #1)]] &[F[k], F[k + 1]]

tail[q_] :=
 Block[{k = 1., step},
  While[(step = f[k,q]) > 0, 
   k += step];
  k - 1
  ]

Now to wrap it up with the precomputed body. Note that I am using Interpolation rather than Nearest. This is because Nearest will be looking at closest point by Euclidean distance which isn't what we want. A little trickery with signs is necessary to get the right continuity on the zero-order interpolation.

maxTerms = 10^6;
x = Range[maxTerms];
y = Accumulate[s[x]/A];

ifun = Interpolation[Transpose[{-y, x}], InterpolationOrder -> 0];

c = F[maxTerms]/A;

Clear[invF];
SetAttributes[invF, Listable];
invF[q_ /; q <= c] := ifun[-q]
invF[q_/; q > c] := tail[q]

Now we generate some random uniforms and plot the result against the density for your distribution.

u = RandomReal[{0, 1}, 10000];   
rvs = invF[u];

Show[Histogram[rvs, {Range[0, 100, 1]~Join~{Max[rvs]}}, "PDF"], 
 DiscretePlot[s[k]/A, {k, 1, 200}], PlotRange -> {{0, 100}, All}]