View Javadoc
1   /*
2    * Licensed to the Apache Software Foundation (ASF) under one or more
3    * contributor license agreements.  See the NOTICE file distributed with
4    * this work for additional information regarding copyright ownership.
5    * The ASF licenses this file to You under the Apache License, Version 2.0
6    * (the "License"); you may not use this file except in compliance with
7    * the License.  You may obtain a copy of the License at
8    *
9    *      http://www.apache.org/licenses/LICENSE-2.0
10   *
11   * Unless required by applicable law or agreed to in writing, software
12   * distributed under the License is distributed on an "AS IS" BASIS,
13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14   * See the License for the specific language governing permissions and
15   * limitations under the License.
16   */
17  
18  package org.apache.commons.rng.sampling.distribution;
19  
20  import org.apache.commons.rng.UniformRandomProvider;
21  
22  /**
23   * Discrete uniform distribution sampler.
24   *
25   * <p>Sampling uses {@link UniformRandomProvider#nextInt}.</p>
26   *
27   * <p>When the range is a power of two the number of calls is 1 per sample.
28   * Otherwise a rejection algorithm is used to ensure uniformity. In the worst
29   * case scenario where the range spans half the range of an {@code int}
30   * (2<sup>31</sup> + 1) the expected number of calls is 2 per sample.</p>
31   *
32   * <p>This sampler can be used as a replacement for {@link UniformRandomProvider#nextInt}
33   * with appropriate adjustment of the upper bound to be inclusive and will outperform that
34   * method when the range is not a power of two. The advantage is gained by pre-computation
35   * of the rejection threshold.</p>
36   *
37   * <p>The sampling algorithm is described in:</p>
38   *
39   * <blockquote>
40   *  Lemire, D (2019). <i>Fast Random Integer Generation in an Interval.</i>
41   *  <strong>ACM Transactions on Modeling and Computer Simulation</strong> 29 (1).
42   * </blockquote>
43   *
44   * <p>The number of {@code int} values required per sample follows a geometric distribution with
45   * a probability of success p of {@code 1 - ((2^32 % n) / 2^32)}. This requires on average 1/p random
46   * {@code int} values per sample.</p>
47   *
48   * @see <a href="https://arxiv.org/abs/1805.10941">Fast Random Integer Generation in an Interval</a>
49   *
50   * @since 1.0
51   */
52  public class DiscreteUniformSampler
53      extends SamplerBase
54      implements SharedStateDiscreteSampler {
55  
56      /** The appropriate uniform sampler for the parameters. */
57      private final SharedStateDiscreteSampler delegate;
58  
59      /**
60       * Base class for a sampler from a discrete uniform distribution. This contains the
61       * source of randomness.
62       */
63      private abstract static class AbstractDiscreteUniformSampler
64              implements SharedStateDiscreteSampler {
65          /** Underlying source of randomness. */
66          protected final UniformRandomProvider rng;
67  
68          /**
69           * @param rng Generator of uniformly distributed random numbers.
70           */
71          AbstractDiscreteUniformSampler(UniformRandomProvider rng) {
72              this.rng = rng;
73          }
74  
75          @Override
76          public String toString() {
77              return "Uniform deviate [" + rng.toString() + "]";
78          }
79      }
80  
81      /**
82       * Discrete uniform distribution sampler when the sample value is fixed.
83       */
84      private static class FixedDiscreteUniformSampler
85              extends AbstractDiscreteUniformSampler {
86          /** The value. */
87          private final int value;
88  
89          /**
90           * @param rng Generator of uniformly distributed random numbers.
91           * @param value The value.
92           */
93          FixedDiscreteUniformSampler(UniformRandomProvider rng,
94                                      int value) {
95              super(rng);
96              this.value = value;
97          }
98  
99          @Override
100         public int sample() {
101             return value;
102         }
103 
104         @Override
105         public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
106             // No requirement for the RNG
107             return this;
108         }
109     }
110 
111     /**
112      * Discrete uniform distribution sampler when the range is a power of 2 and greater than 1.
113      * This sampler assumes the lower bound of the range is 0.
114      *
115      * <p>Note: This cannot be used when the range is 1 (2^0) as the shift would be 32-bits
116      * which is ignored by the shift operator.</p>
117      */
118     private static class PowerOf2RangeDiscreteUniformSampler
119             extends AbstractDiscreteUniformSampler {
120         /** Bit shift to apply to the integer sample. */
121         private final int shift;
122 
123         /**
124          * @param rng Generator of uniformly distributed random numbers.
125          * @param range Maximum range of the sample (exclusive).
126          * Must be a power of 2 greater than 2^0.
127          */
128         PowerOf2RangeDiscreteUniformSampler(UniformRandomProvider rng,
129                                             int range) {
130             super(rng);
131             this.shift = Integer.numberOfLeadingZeros(range) + 1;
132         }
133 
134         /**
135          * @param rng Generator of uniformly distributed random numbers.
136          * @param source Source to copy.
137          */
138         PowerOf2RangeDiscreteUniformSampler(UniformRandomProvider rng,
139                                             PowerOf2RangeDiscreteUniformSampler source) {
140             super(rng);
141             this.shift = source.shift;
142         }
143 
144         @Override
145         public int sample() {
146             // Use a bit shift to favour the most significant bits.
147             // Note: The result would be the same as the rejection method used in the
148             // small range sampler when there is no rejection threshold.
149             return rng.nextInt() >>> shift;
150         }
151 
152         @Override
153         public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
154             return new PowerOf2RangeDiscreteUniformSampler(rng, this);
155         }
156     }
157 
158     /**
159      * Discrete uniform distribution sampler when the range is small
160      * enough to fit in a positive integer.
161      * This sampler assumes the lower bound of the range is 0.
162      *
163      * <p>Implements the algorithm of Lemire (2019).</p>
164      *
165      * @see <a href="https://arxiv.org/abs/1805.10941">Fast Random Integer Generation in an Interval</a>
166      */
167     private static class SmallRangeDiscreteUniformSampler
168             extends AbstractDiscreteUniformSampler {
169         /** Maximum range of the sample (exclusive). */
170         private final long n;
171 
172         /**
173          * The level below which samples are rejected based on the fraction remainder.
174          *
175          * <p>Any remainder below this denotes that there are still floor(2^32 / n) more
176          * observations of this sample from the interval [0, 2^32), where n is the range.</p>
177          */
178         private final long threshold;
179 
180         /**
181          * @param rng Generator of uniformly distributed random numbers.
182          * @param range Maximum range of the sample (exclusive).
183          */
184         SmallRangeDiscreteUniformSampler(UniformRandomProvider rng,
185                                          int range) {
186             super(rng);
187             // Handle range as an unsigned 32-bit integer
188             this.n = range & 0xffffffffL;
189             // Compute 2^32 % n
190             threshold = (1L << 32) % n;
191         }
192 
193         /**
194          * @param rng Generator of uniformly distributed random numbers.
195          * @param source Source to copy.
196          */
197         SmallRangeDiscreteUniformSampler(UniformRandomProvider rng,
198                                          SmallRangeDiscreteUniformSampler source) {
199             super(rng);
200             this.n = source.n;
201             this.threshold = source.threshold;
202         }
203 
204         @Override
205         public int sample() {
206             // Rejection method using multiply by a fraction:
207             // n * [0, 2^32 - 1)
208             //     -------------
209             //         2^32
210             // The result is mapped back to an integer and will be in the range [0, n).
211             // Note this is comparable to range * rng.nextDouble() but with compensation for
212             // non-uniformity due to round-off.
213             long result;
214             do {
215                 // Compute 64-bit unsigned product of n * [0, 2^32 - 1).
216                 // The upper 32-bits contains the sample value in the range [0, n), i.e. result / 2^32.
217                 // The lower 32-bits contains the remainder, i.e. result % 2^32.
218                 result = n * (rng.nextInt() & 0xffffffffL);
219 
220                 // Test the sample uniformity.
221                 // Samples are observed on average (2^32 / n) times at a frequency of either
222                 // floor(2^32 / n) or ceil(2^32 / n).
223                 // To ensure all samples have a frequency of floor(2^32 / n) reject any results with
224                 // a remainder < (2^32 % n), i.e. the level below which denotes that there are still
225                 // floor(2^32 / n) more observations of this sample.
226             } while ((result & 0xffffffffL) < threshold);
227             // Divide by 2^32 to get the sample.
228             return (int)(result >>> 32);
229         }
230 
231         @Override
232         public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
233             return new SmallRangeDiscreteUniformSampler(rng, this);
234         }
235     }
236 
237     /**
238      * Discrete uniform distribution sampler when the range between lower and upper is too large
239      * to fit in a positive integer.
240      */
241     private static class LargeRangeDiscreteUniformSampler
242             extends AbstractDiscreteUniformSampler {
243         /** Lower bound. */
244         private final int lower;
245         /** Upper bound. */
246         private final int upper;
247 
248         /**
249          * @param rng Generator of uniformly distributed random numbers.
250          * @param lower Lower bound (inclusive) of the distribution.
251          * @param upper Upper bound (inclusive) of the distribution.
252          */
253         LargeRangeDiscreteUniformSampler(UniformRandomProvider rng,
254                                          int lower,
255                                          int upper) {
256             super(rng);
257             this.lower = lower;
258             this.upper = upper;
259         }
260 
261         @Override
262         public int sample() {
263             // Use a simple rejection method.
264             // This is used when (upper-lower) >= Integer.MAX_VALUE.
265             // This will loop on average 2 times in the worst case scenario
266             // when (upper-lower) == Integer.MAX_VALUE.
267             while (true) {
268                 final int r = rng.nextInt();
269                 if (r >= lower &&
270                     r <= upper) {
271                     return r;
272                 }
273             }
274         }
275 
276         @Override
277         public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
278             return new LargeRangeDiscreteUniformSampler(rng, lower, upper);
279         }
280     }
281 
282     /**
283      * Adds an offset to an underlying discrete sampler.
284      */
285     private static class OffsetDiscreteUniformSampler
286             extends AbstractDiscreteUniformSampler {
287         /** The offset. */
288         private final int offset;
289         /** The discrete sampler. */
290         private final SharedStateDiscreteSampler sampler;
291 
292         /**
293          * @param offset The offset for the sample.
294          * @param sampler The discrete sampler.
295          */
296         OffsetDiscreteUniformSampler(int offset,
297                                      SharedStateDiscreteSampler sampler) {
298             super(null);
299             this.offset = offset;
300             this.sampler = sampler;
301         }
302 
303         @Override
304         public int sample() {
305             return offset + sampler.sample();
306         }
307 
308         /** {@inheritDoc} */
309         @Override
310         public String toString() {
311             return sampler.toString();
312         }
313 
314         @Override
315         public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
316             return new OffsetDiscreteUniformSampler(offset, sampler.withUniformRandomProvider(rng));
317         }
318     }
319 
320     /**
321      * This instance delegates sampling. Use the factory method
322      * {@link #of(UniformRandomProvider, int, int)} to create an optimal sampler.
323      *
324      * @param rng Generator of uniformly distributed random numbers.
325      * @param lower Lower bound (inclusive) of the distribution.
326      * @param upper Upper bound (inclusive) of the distribution.
327      * @throws IllegalArgumentException if {@code lower > upper}.
328      */
329     public DiscreteUniformSampler(UniformRandomProvider rng,
330                                   int lower,
331                                   int upper) {
332         super(null);
333         delegate = of(rng, lower, upper);
334     }
335 
336     /** {@inheritDoc} */
337     @Override
338     public int sample() {
339         return delegate.sample();
340     }
341 
342     /** {@inheritDoc} */
343     @Override
344     public String toString() {
345         return delegate.toString();
346     }
347 
348     /**
349      * {@inheritDoc}
350      *
351      * @since 1.3
352      */
353     @Override
354     public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
355         // Direct return of the optimised sampler
356         return delegate.withUniformRandomProvider(rng);
357     }
358 
359     /**
360      * Creates a new discrete uniform distribution sampler.
361      *
362      * @param rng Generator of uniformly distributed random numbers.
363      * @param lower Lower bound (inclusive) of the distribution.
364      * @param upper Upper bound (inclusive) of the distribution.
365      * @return the sampler
366      * @throws IllegalArgumentException if {@code lower > upper}.
367      * @since 1.3
368      */
369     public static SharedStateDiscreteSampler of(UniformRandomProvider rng,
370                                                 int lower,
371                                                 int upper) {
372         if (lower > upper) {
373             throw new IllegalArgumentException(lower  + " > " + upper);
374         }
375 
376         // Choose the algorithm depending on the range
377 
378         // Edge case for no range.
379         // This must be done first as the methods to handle lower == 0
380         // do not handle upper == 0.
381         if (upper == lower) {
382             return new FixedDiscreteUniformSampler(rng, lower);
383         }
384 
385         // Algorithms to ignore the lower bound if it is zero.
386         if (lower == 0) {
387             return createZeroBoundedSampler(rng, upper);
388         }
389 
390         final int range = (upper - lower) + 1;
391         // Check power of 2 first to handle range == 2^31.
392         if (isPowerOf2(range)) {
393             return new OffsetDiscreteUniformSampler(lower,
394                                                     new PowerOf2RangeDiscreteUniformSampler(rng, range));
395         }
396         if (range <= 0) {
397             // The range is too wide to fit in a positive int (larger
398             // than 2^31); use a simple rejection method.
399             // Note: if range == 0 then the input is [Integer.MIN_VALUE, Integer.MAX_VALUE].
400             // No specialisation exists for this and it is handled as a large range.
401             return new LargeRangeDiscreteUniformSampler(rng, lower, upper);
402         }
403         // Use a sample from the range added to the lower bound.
404         return new OffsetDiscreteUniformSampler(lower,
405                                                 new SmallRangeDiscreteUniformSampler(rng, range));
406     }
407 
408     /**
409      * Create a new sampler for the range {@code 0} inclusive to {@code upper} inclusive.
410      *
411      * <p>This can handle any positive {@code upper}.
412      *
413      * @param rng Generator of uniformly distributed random numbers.
414      * @param upper Upper bound (inclusive) of the distribution. Must be positive.
415      * @return the sampler
416      */
417     private static AbstractDiscreteUniformSampler createZeroBoundedSampler(UniformRandomProvider rng,
418                                                                            int upper) {
419         // Note: Handle any range up to 2^31 (which is negative as a signed
420         // 32-bit integer but handled as a power of 2)
421         final int range = upper + 1;
422         return isPowerOf2(range) ?
423             new PowerOf2RangeDiscreteUniformSampler(rng, range) :
424             new SmallRangeDiscreteUniformSampler(rng, range);
425     }
426 
427     /**
428      * Checks if the value is a power of 2.
429      *
430      * <p>This returns {@code true} for the value {@code Integer.MIN_VALUE} which can be
431      * handled as an unsigned integer of 2^31.</p>
432      *
433      * @param value Value.
434      * @return {@code true} if a power of 2
435      */
436     private static boolean isPowerOf2(final int value) {
437         return value != 0 && (value & (value - 1)) == 0;
438     }
439 }