001/*
002 * Licensed to the Apache Software Foundation (ASF) under one or more
003 * contributor license agreements.  See the NOTICE file distributed with
004 * this work for additional information regarding copyright ownership.
005 * The ASF licenses this file to You under the Apache License, Version 2.0
006 * (the "License"); you may not use this file except in compliance with
007 * the License.  You may obtain a copy of the License at
008 *
009 *      http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017
018package org.apache.commons.rng.sampling.distribution;
019
020import org.apache.commons.rng.UniformRandomProvider;
021
022/**
023 * Discrete uniform distribution sampler.
024 *
025 * <p>Sampling uses {@link UniformRandomProvider#nextInt}.</p>
026 *
027 * <p>When the range is a power of two the number of calls is 1 per sample.
028 * Otherwise a rejection algorithm is used to ensure uniformity. In the worst
029 * case scenario where the range spans half the range of an {@code int}
030 * (2<sup>31</sup> + 1) the expected number of calls is 2 per sample.</p>
031 *
032 * <p>This sampler can be used as a replacement for {@link UniformRandomProvider#nextInt}
033 * with appropriate adjustment of the upper bound to be inclusive and will outperform that
034 * method when the range is not a power of two. The advantage is gained by pre-computation
035 * of the rejection threshold.</p>
036 *
037 * <p>The sampling algorithm is described in:</p>
038 *
039 * <blockquote>
040 *  Lemire, D (2019). <i>Fast Random Integer Generation in an Interval.</i>
041 *  <strong>ACM Transactions on Modeling and Computer Simulation</strong> 29 (1).
042 * </blockquote>
043 *
044 * <p>The number of {@code int} values required per sample follows a geometric distribution with
045 * a probability of success p of {@code 1 - ((2^32 % n) / 2^32)}. This requires on average 1/p random
046 * {@code int} values per sample.</p>
047 *
048 * @see <a href="https://arxiv.org/abs/1805.10941">Fast Random Integer Generation in an Interval</a>
049 *
050 * @since 1.0
051 */
052public class DiscreteUniformSampler
053    extends SamplerBase
054    implements SharedStateDiscreteSampler {
055
056    /** The appropriate uniform sampler for the parameters. */
057    private final SharedStateDiscreteSampler delegate;
058
059    /**
060     * Base class for a sampler from a discrete uniform distribution. This contains the
061     * source of randomness.
062     */
063    private abstract static class AbstractDiscreteUniformSampler
064            implements SharedStateDiscreteSampler {
065        /** Underlying source of randomness. */
066        protected final UniformRandomProvider rng;
067
068        /**
069         * @param rng Generator of uniformly distributed random numbers.
070         */
071        AbstractDiscreteUniformSampler(UniformRandomProvider rng) {
072            this.rng = rng;
073        }
074
075        @Override
076        public String toString() {
077            return "Uniform deviate [" + rng.toString() + "]";
078        }
079    }
080
081    /**
082     * Discrete uniform distribution sampler when the sample value is fixed.
083     */
084    private static final class FixedDiscreteUniformSampler
085            extends AbstractDiscreteUniformSampler {
086        /** The value. */
087        private final int value;
088
089        /**
090         * @param value The value.
091         */
092        FixedDiscreteUniformSampler(int value) {
093            // No requirement for the RNG
094            super(null);
095            this.value = value;
096        }
097
098        @Override
099        public int sample() {
100            return value;
101        }
102
103        @Override
104        public String toString() {
105            // No RNG to include in the string
106            return "Uniform deviate [X=" + value + "]";
107        }
108
109        @Override
110        public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
111            // No requirement for the RNG
112            return this;
113        }
114    }
115
116    /**
117     * Discrete uniform distribution sampler when the range is a power of 2 and greater than 1.
118     * This sampler assumes the lower bound of the range is 0.
119     *
120     * <p>Note: This cannot be used when the range is 1 (2^0) as the shift would be 32-bits
121     * which is ignored by the shift operator.</p>
122     */
123    private static final class PowerOf2RangeDiscreteUniformSampler
124            extends AbstractDiscreteUniformSampler {
125        /** Bit shift to apply to the integer sample. */
126        private final int shift;
127
128        /**
129         * @param rng Generator of uniformly distributed random numbers.
130         * @param range Maximum range of the sample (exclusive).
131         * Must be a power of 2 greater than 2^0.
132         */
133        PowerOf2RangeDiscreteUniformSampler(UniformRandomProvider rng,
134                                            int range) {
135            super(rng);
136            this.shift = Integer.numberOfLeadingZeros(range) + 1;
137        }
138
139        /**
140         * @param rng Generator of uniformly distributed random numbers.
141         * @param source Source to copy.
142         */
143        PowerOf2RangeDiscreteUniformSampler(UniformRandomProvider rng,
144                                            PowerOf2RangeDiscreteUniformSampler source) {
145            super(rng);
146            this.shift = source.shift;
147        }
148
149        @Override
150        public int sample() {
151            // Use a bit shift to favour the most significant bits.
152            // Note: The result would be the same as the rejection method used in the
153            // small range sampler when there is no rejection threshold.
154            return rng.nextInt() >>> shift;
155        }
156
157        @Override
158        public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
159            return new PowerOf2RangeDiscreteUniformSampler(rng, this);
160        }
161    }
162
163    /**
164     * Discrete uniform distribution sampler when the range is small
165     * enough to fit in a positive integer.
166     * This sampler assumes the lower bound of the range is 0.
167     *
168     * <p>Implements the algorithm of Lemire (2019).</p>
169     *
170     * @see <a href="https://arxiv.org/abs/1805.10941">Fast Random Integer Generation in an Interval</a>
171     */
172    private static final class SmallRangeDiscreteUniformSampler
173            extends AbstractDiscreteUniformSampler {
174        /** Maximum range of the sample (exclusive). */
175        private final long n;
176
177        /**
178         * The level below which samples are rejected based on the fraction remainder.
179         *
180         * <p>Any remainder below this denotes that there are still floor(2^32 / n) more
181         * observations of this sample from the interval [0, 2^32), where n is the range.</p>
182         */
183        private final long threshold;
184
185        /**
186         * @param rng Generator of uniformly distributed random numbers.
187         * @param range Maximum range of the sample (exclusive).
188         */
189        SmallRangeDiscreteUniformSampler(UniformRandomProvider rng,
190                                         int range) {
191            super(rng);
192            // Handle range as an unsigned 32-bit integer
193            this.n = range & 0xffffffffL;
194            // Compute 2^32 % n
195            threshold = (1L << 32) % n;
196        }
197
198        /**
199         * @param rng Generator of uniformly distributed random numbers.
200         * @param source Source to copy.
201         */
202        SmallRangeDiscreteUniformSampler(UniformRandomProvider rng,
203                                         SmallRangeDiscreteUniformSampler source) {
204            super(rng);
205            this.n = source.n;
206            this.threshold = source.threshold;
207        }
208
209        @Override
210        public int sample() {
211            // Rejection method using multiply by a fraction:
212            // n * [0, 2^32 - 1)
213            //     -------------
214            //         2^32
215            // The result is mapped back to an integer and will be in the range [0, n).
216            // Note this is comparable to range * rng.nextDouble() but with compensation for
217            // non-uniformity due to round-off.
218            long result;
219            do {
220                // Compute 64-bit unsigned product of n * [0, 2^32 - 1).
221                // The upper 32-bits contains the sample value in the range [0, n), i.e. result / 2^32.
222                // The lower 32-bits contains the remainder, i.e. result % 2^32.
223                result = n * (rng.nextInt() & 0xffffffffL);
224
225                // Test the sample uniformity.
226                // Samples are observed on average (2^32 / n) times at a frequency of either
227                // floor(2^32 / n) or ceil(2^32 / n).
228                // To ensure all samples have a frequency of floor(2^32 / n) reject any results with
229                // a remainder < (2^32 % n), i.e. the level below which denotes that there are still
230                // floor(2^32 / n) more observations of this sample.
231            } while ((result & 0xffffffffL) < threshold);
232            // Divide by 2^32 to get the sample.
233            return (int)(result >>> 32);
234        }
235
236        @Override
237        public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
238            return new SmallRangeDiscreteUniformSampler(rng, this);
239        }
240    }
241
242    /**
243     * Discrete uniform distribution sampler when the range between lower and upper is too large
244     * to fit in a positive integer.
245     */
246    private static final class LargeRangeDiscreteUniformSampler
247            extends AbstractDiscreteUniformSampler {
248        /** Lower bound. */
249        private final int lower;
250        /** Upper bound. */
251        private final int upper;
252
253        /**
254         * @param rng Generator of uniformly distributed random numbers.
255         * @param lower Lower bound (inclusive) of the distribution.
256         * @param upper Upper bound (inclusive) of the distribution.
257         */
258        LargeRangeDiscreteUniformSampler(UniformRandomProvider rng,
259                                         int lower,
260                                         int upper) {
261            super(rng);
262            this.lower = lower;
263            this.upper = upper;
264        }
265
266        @Override
267        public int sample() {
268            // Use a simple rejection method.
269            // This is used when (upper-lower) >= Integer.MAX_VALUE.
270            // This will loop on average 2 times in the worst case scenario
271            // when (upper-lower) == Integer.MAX_VALUE.
272            while (true) {
273                final int r = rng.nextInt();
274                if (r >= lower &&
275                    r <= upper) {
276                    return r;
277                }
278            }
279        }
280
281        @Override
282        public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
283            return new LargeRangeDiscreteUniformSampler(rng, lower, upper);
284        }
285    }
286
287    /**
288     * Adds an offset to an underlying discrete sampler.
289     */
290    private static final class OffsetDiscreteUniformSampler
291            extends AbstractDiscreteUniformSampler {
292        /** The offset. */
293        private final int offset;
294        /** The discrete sampler. */
295        private final SharedStateDiscreteSampler sampler;
296
297        /**
298         * @param offset The offset for the sample.
299         * @param sampler The discrete sampler.
300         */
301        OffsetDiscreteUniformSampler(int offset,
302                                     SharedStateDiscreteSampler sampler) {
303            super(null);
304            this.offset = offset;
305            this.sampler = sampler;
306        }
307
308        @Override
309        public int sample() {
310            return offset + sampler.sample();
311        }
312
313        @Override
314        public String toString() {
315            return sampler.toString();
316        }
317
318        @Override
319        public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
320            return new OffsetDiscreteUniformSampler(offset, sampler.withUniformRandomProvider(rng));
321        }
322    }
323
324    /**
325     * This instance delegates sampling. Use the factory method
326     * {@link #of(UniformRandomProvider, int, int)} to create an optimal sampler.
327     *
328     * @param rng Generator of uniformly distributed random numbers.
329     * @param lower Lower bound (inclusive) of the distribution.
330     * @param upper Upper bound (inclusive) of the distribution.
331     * @throws IllegalArgumentException if {@code lower > upper}.
332     */
333    public DiscreteUniformSampler(UniformRandomProvider rng,
334                                  int lower,
335                                  int upper) {
336        this(of(rng, lower, upper));
337    }
338
339    /**
340     * Private constructor used by to prevent partially initialized object if the construction
341     * of the delegate throws. In future versions the public constructor should be removed.
342     *
343     * @param delegate Delegate.
344     */
345    private DiscreteUniformSampler(SharedStateDiscreteSampler delegate) {
346        super(null);
347        this.delegate = delegate;
348    }
349
350    /** {@inheritDoc} */
351    @Override
352    public int sample() {
353        return delegate.sample();
354    }
355
356    /** {@inheritDoc} */
357    @Override
358    public String toString() {
359        return delegate.toString();
360    }
361
362    /**
363     * {@inheritDoc}
364     *
365     * @since 1.3
366     */
367    @Override
368    public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
369        // Direct return of the optimised sampler
370        return delegate.withUniformRandomProvider(rng);
371    }
372
373    /**
374     * Creates a new discrete uniform distribution sampler.
375     *
376     * @param rng Generator of uniformly distributed random numbers.
377     * @param lower Lower bound (inclusive) of the distribution.
378     * @param upper Upper bound (inclusive) of the distribution.
379     * @return the sampler
380     * @throws IllegalArgumentException if {@code lower > upper}.
381     * @since 1.3
382     */
383    public static SharedStateDiscreteSampler of(UniformRandomProvider rng,
384                                                int lower,
385                                                int upper) {
386        if (lower > upper) {
387            throw new IllegalArgumentException(lower  + " > " + upper);
388        }
389
390        // Choose the algorithm depending on the range
391
392        // Edge case for no range.
393        // This must be done first as the methods to handle lower == 0
394        // do not handle upper == 0.
395        if (upper == lower) {
396            return new FixedDiscreteUniformSampler(lower);
397        }
398
399        // Algorithms to ignore the lower bound if it is zero.
400        if (lower == 0) {
401            return createZeroBoundedSampler(rng, upper);
402        }
403
404        final int range = (upper - lower) + 1;
405        // Check power of 2 first to handle range == 2^31.
406        if (isPowerOf2(range)) {
407            return new OffsetDiscreteUniformSampler(lower,
408                                                    new PowerOf2RangeDiscreteUniformSampler(rng, range));
409        }
410        if (range <= 0) {
411            // The range is too wide to fit in a positive int (larger
412            // than 2^31); use a simple rejection method.
413            // Note: if range == 0 then the input is [Integer.MIN_VALUE, Integer.MAX_VALUE].
414            // No specialisation exists for this and it is handled as a large range.
415            return new LargeRangeDiscreteUniformSampler(rng, lower, upper);
416        }
417        // Use a sample from the range added to the lower bound.
418        return new OffsetDiscreteUniformSampler(lower,
419                                                new SmallRangeDiscreteUniformSampler(rng, range));
420    }
421
422    /**
423     * Create a new sampler for the range {@code 0} inclusive to {@code upper} inclusive.
424     *
425     * <p>This can handle any positive {@code upper}.
426     *
427     * @param rng Generator of uniformly distributed random numbers.
428     * @param upper Upper bound (inclusive) of the distribution. Must be positive.
429     * @return the sampler
430     */
431    private static AbstractDiscreteUniformSampler createZeroBoundedSampler(UniformRandomProvider rng,
432                                                                           int upper) {
433        // Note: Handle any range up to 2^31 (which is negative as a signed
434        // 32-bit integer but handled as a power of 2)
435        final int range = upper + 1;
436        return isPowerOf2(range) ?
437            new PowerOf2RangeDiscreteUniformSampler(rng, range) :
438            new SmallRangeDiscreteUniformSampler(rng, range);
439    }
440
441    /**
442     * Checks if the value is a power of 2.
443     *
444     * <p>This returns {@code true} for the value {@code Integer.MIN_VALUE} which can be
445     * handled as an unsigned integer of 2^31.</p>
446     *
447     * @param value Value.
448     * @return {@code true} if a power of 2
449     */
450    private static boolean isPowerOf2(final int value) {
451        return value != 0 && (value & (value - 1)) == 0;
452    }
453}