FastLoadedDiceRollerDiscreteSampler.java
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.commons.rng.sampling.distribution;
import java.math.BigInteger;
import java.util.Arrays;
import org.apache.commons.rng.UniformRandomProvider;
/**
* Distribution sampler that uses the Fast Loaded Dice Roller (FLDR). It can be used to
* sample from {@code n} values each with an associated relative weight. If all unique items
* are assigned the same weight it is more efficient to use the {@link DiscreteUniformSampler}.
*
* <p>Given a list {@code L} of {@code n} positive numbers,
* where {@code L[i]} represents the relative weight of the {@code i}th side, FLDR returns
* integer {@code i} with relative probability {@code L[i]}.
*
* <p>FLDR produces <em>exact</em> samples from the specified probability distribution.
* <ul>
* <li>For integer weights, the probability of returning {@code i} is precisely equal to the
* rational number {@code L[i] / m}, where {@code m} is the sum of {@code L}.
* <li>For floating-points weights, each weight {@code L[i]} is converted to the
* corresponding rational number {@code p[i] / q[i]} where {@code p[i]} is a positive integer and
* {@code q[i]} is a power of 2. The rational weights are then normalized (exactly) to sum to unity.
* </ul>
*
* <p>Note that if <em>exact</em> samples are not required then an alternative sampler that
* ignores very small relative weights may have improved sampling performance.
*
* <p>This implementation is based on the algorithm in:
*
* <blockquote>
* Feras A. Saad, Cameron E. Freer, Martin C. Rinard, and Vikash K. Mansinghka.
* The Fast Loaded Dice Roller: A Near-Optimal Exact Sampler for Discrete Probability
* Distributions. In AISTATS 2020: Proceedings of the 23rd International Conference on
* Artificial Intelligence and Statistics, Proceedings of Machine Learning Research 108,
* Palermo, Sicily, Italy, 2020.
* </blockquote>
*
* <p>Sampling uses {@link UniformRandomProvider#nextInt()} as the source of random bits.
*
* @see <a href="https://arxiv.org/abs/2003.03830">Saad et al (2020)
* Proceedings of the 23rd International Conference on Artificial Intelligence and Statistics,
* PMLR 108:1036-1046.</a>
* @since 1.5
*/
public abstract class FastLoadedDiceRollerDiscreteSampler
implements SharedStateDiscreteSampler {
/**
* The maximum size of an array.
*
* <p>This value is taken from the limit in Open JDK 8 {@code java.util.ArrayList}.
* It allows VMs to reserve some header words in an array.
*/
private static final int MAX_ARRAY_SIZE = Integer.MAX_VALUE - 8;
/** The maximum biased exponent for a finite double.
* This is offset by 1023 from {@code Math.getExponent(Double.MAX_VALUE)}. */
private static final int MAX_BIASED_EXPONENT = 2046;
/** Size of the mantissa of a double. Equal to 52 bits. */
private static final int MANTISSA_SIZE = 52;
/** Mask to extract the 52-bit mantissa from a long representation of a double. */
private static final long MANTISSA_MASK = 0x000f_ffff_ffff_ffffL;
/** BigInteger representation of {@link Long#MAX_VALUE}. */
private static final BigInteger MAX_LONG = BigInteger.valueOf(Long.MAX_VALUE);
/** The maximum offset that will avoid loss of bits for a left shift of a 53-bit value.
* The value will remain positive for any shift {@code <=} this value. */
private static final int MAX_OFFSET = 10;
/** Initial value for no leaf node label. */
private static final int NO_LABEL = Integer.MAX_VALUE;
/** Name of the sampler. */
private static final String SAMPLER_NAME = "Fast Loaded Dice Roller";
/**
* Class to handle the edge case of observations in only one category.
*/
private static final class FixedValueDiscreteSampler extends FastLoadedDiceRollerDiscreteSampler {
/** The sample value. */
private final int sampleValue;
/**
* @param sampleValue Sample value.
*/
FixedValueDiscreteSampler(int sampleValue) {
this.sampleValue = sampleValue;
}
@Override
public int sample() {
return sampleValue;
}
@Override
public FastLoadedDiceRollerDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
return this;
}
@Override
public String toString() {
return SAMPLER_NAME;
}
}
/**
* Class to implement the FLDR sample algorithm.
*/
private static final class FLDRSampler extends FastLoadedDiceRollerDiscreteSampler {
/** Empty boolean source. This is the location of the sign-bit after 31 right shifts on
* the boolean source. */
private static final int EMPTY_BOOL_SOURCE = 1;
/** Underlying source of randomness. */
private final UniformRandomProvider rng;
/** Number of categories. */
private final int n;
/** Number of levels in the discrete distribution generating (DDG) tree.
* Equal to {@code ceil(log2(m))} where {@code m} is the sum of observations. */
private final int k;
/** Number of leaf nodes at each level. */
private final int[] h;
/** Stores the leaf node labels in increasing order. Named {@code H} in the FLDR paper. */
private final int[] lH;
/**
* Provides a bit source for booleans.
*
* <p>A cached value from a call to {@link UniformRandomProvider#nextInt()}.
*
* <p>Only stores 31-bits when full as 1 bit has already been consumed.
* The sign bit is a flag that shifts down so the source eventually equals 1
* when all bits are consumed and will trigger a refill.
*/
private int booleanSource = EMPTY_BOOL_SOURCE;
/**
* Creates a sampler.
*
* <p>The input parameters are not validated and must be correctly computed tables.
*
* @param rng Generator of uniformly distributed random numbers.
* @param n Number of categories
* @param k Number of levels in the discrete distribution generating (DDG) tree.
* Equal to {@code ceil(log2(m))} where {@code m} is the sum of observations.
* @param h Number of leaf nodes at each level.
* @param lH Stores the leaf node labels in increasing order.
*/
FLDRSampler(UniformRandomProvider rng,
int n,
int k,
int[] h,
int[] lH) {
this.rng = rng;
this.n = n;
this.k = k;
// Deliberate direct storage of input arrays
this.h = h;
this.lH = lH;
}
/**
* Creates a copy with a new source of randomness.
*
* @param rng Generator of uniformly distributed random numbers.
* @param source Source to copy.
*/
private FLDRSampler(UniformRandomProvider rng,
FLDRSampler source) {
this.rng = rng;
this.n = source.n;
this.k = source.k;
this.h = source.h;
this.lH = source.lH;
}
/** {@inheritDoc} */
@Override
public int sample() {
// ALGORITHM 5: SAMPLE
int c = 0;
int d = 0;
for (;;) {
// b = flip()
// d = 2 * d + (1 - b)
d = (d << 1) + flip();
if (d < h[c]) {
// z = H[d][c]
final int z = lH[d * k + c];
// assert z != NO_LABEL
if (z < n) {
return z;
}
d = 0;
c = 0;
} else {
d = d - h[c];
c++;
}
}
}
/**
* Provides a source of boolean bits.
*
* <p>Note: This replicates the boolean cache functionality of
* {@code o.a.c.rng.core.source32.IntProvider}. The method has been simplified to return
* an {@code int} value rather than a {@code boolean}.
*
* @return the bit (0 or 1)
*/
private int flip() {
int bits = booleanSource;
if (bits == 1) {
// Refill
bits = rng.nextInt();
// Store a refill flag in the sign bit and the unused 31 bits, return lowest bit
booleanSource = Integer.MIN_VALUE | (bits >>> 1);
return bits & 0x1;
}
// Shift down eventually triggering refill, return current lowest bit
booleanSource = bits >>> 1;
return bits & 0x1;
}
/** {@inheritDoc} */
@Override
public String toString() {
return SAMPLER_NAME + " [" + rng.toString() + "]";
}
/** {@inheritDoc} */
@Override
public FastLoadedDiceRollerDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
return new FLDRSampler(rng, this);
}
}
/** Package-private constructor. */
FastLoadedDiceRollerDiscreteSampler() {
// Intentionally empty
}
/** {@inheritDoc} */
// Redeclare the signature to return a FastLoadedDiceRollerSampler not a SharedStateLongSampler
@Override
public abstract FastLoadedDiceRollerDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng);
/**
* Creates a sampler.
*
* <p>Note: The discrete distribution generating (DDG) tree requires {@code (n + 1) * k} entries
* where {@code n} is the number of categories, {@code k == ceil(log2(m))} and {@code m}
* is the sum of the observed frequencies. An exception is raised if this cannot be allocated
* as a single array.
*
* <p>For reference the sum is limited to {@link Long#MAX_VALUE} and the value {@code k} to 63.
* The number of categories is limited to approximately {@code ((2^31 - 1) / k) = 34,087,042}
* when the sum of frequencies is large enough to create k=63.
*
* @param rng Generator of uniformly distributed random numbers.
* @param frequencies Observed frequencies of the discrete distribution.
* @return the sampler
* @throws IllegalArgumentException if {@code frequencies} is null or empty, a
* frequency is negative, the sum of all frequencies is either zero or
* above {@link Long#MAX_VALUE}, or the size of the discrete distribution generating tree
* is too large.
*/
public static FastLoadedDiceRollerDiscreteSampler of(UniformRandomProvider rng,
long[] frequencies) {
final long m = sum(frequencies);
// Obtain indices of non-zero frequencies
final int[] indices = indicesOfNonZero(frequencies);
// Edge case for 1 non-zero weight. This also handles edge case for 1 observation
// (as log2(m) == 0 will break the computation of the DDG tree).
if (indices.length == 1) {
return new FixedValueDiscreteSampler(indexOfNonZero(frequencies));
}
return createSampler(rng, frequencies, indices, m);
}
/**
* Creates a sampler.
*
* <p>Weights are converted to rational numbers {@code p / q} where {@code q} is a power of 2.
* The numerators {@code p} are scaled to use a common denominator before summing.
*
* <p>All weights are used to create the sampler. Weights with a small magnitude relative
* to the largest weight can be excluded using the constructor method with the
* relative magnitude parameter {@code alpha} (see {@link #of(UniformRandomProvider, double[], int)}).
*
* @param rng Generator of uniformly distributed random numbers.
* @param weights Weights of the discrete distribution.
* @return the sampler
* @throws IllegalArgumentException if {@code weights} is null or empty, a
* weight is negative, infinite or {@code NaN}, the sum of all weights is zero, or the size
* of the discrete distribution generating tree is too large.
* @see #of(UniformRandomProvider, double[], int)
*/
public static FastLoadedDiceRollerDiscreteSampler of(UniformRandomProvider rng,
double[] weights) {
return of(rng, weights, 0);
}
/**
* Creates a sampler.
*
* <p>Weights are converted to rational numbers {@code p / q} where {@code q} is
* a power of 2. The numerators {@code p} are scaled to use a common
* denominator before summing.
*
* <p>Note: The discrete distribution generating (DDG) tree requires
* {@code (n + 1) * k} entries where {@code n} is the number of categories,
* {@code k == ceil(log2(m))} and {@code m} is the sum of the weight numerators
* {@code q}. An exception is raised if this cannot be allocated as a single
* array.
*
* <p>For reference the value {@code k} is equal to or greater than the ratio of
* the largest to the smallest weight expressed as a power of 2. For
* {@code Double.MAX_VALUE / Double.MIN_VALUE} this is ~2098. The value
* {@code k} increases with the sum of the weight numerators. A number of
* weights in excess of 1,000,000 with values equal to {@link Double#MAX_VALUE}
* would be required to raise an exception when the minimum weight is
* {@link Double#MIN_VALUE}.
*
* <p>Weights with a small magnitude relative to the largest weight can be
* excluded using the relative magnitude parameter {@code alpha}. This will set
* any weight to zero if the magnitude is approximately 2<sup>alpha</sup>
* <em>smaller</em> than the largest weight. This comparison is made using only
* the exponent of the input weights. The {@code alpha} parameter is ignored if
* not above zero. Note that a small {@code alpha} parameter will exclude more
* weights than a large {@code alpha} parameter.
*
* <p>The alpha parameter can be used to exclude categories that
* have a very low probability of occurrence and will improve the construction
* performance of the sampler. The effect on sampling performance depends on
* the relative weights of the excluded categories; typically a high {@code alpha}
* is used to exclude categories that would be visited with a very low probability
* and the sampling performance is unchanged.
*
* <p><b>Implementation Note</b>
*
* <p>This method creates a sampler with <em>exact</em> samples from the
* specified probability distribution. It is recommended to use this method:
* <ul>
* <li>if the weights are computed, for example from a probability mass function; or
* <li>if the weights sum to an infinite value.
* </ul>
*
* <p>If the weights are computed from empirical observations then it is
* recommended to use the factory method
* {@link #of(UniformRandomProvider, long[]) accepting frequencies}. This
* requires the total number of observations to be representable as a long
* integer.
*
* <p>Note that if all weights are scaled by a power of 2 to be integers, and
* each integer can be represented as a positive 64-bit long value, then the
* sampler created using this method will match the output from a sampler
* created with the scaled weights converted to long values for the factory
* method {@link #of(UniformRandomProvider, long[]) accepting frequencies}. This
* assumes the sum of the integer values does not overflow.
*
* <p>It should be noted that the conversion of weights to rational numbers has
* a performance overhead during construction (sampling performance is not
* affected). This may be avoided by first converting them to integer values
* that can be summed without overflow. For example by scaling values by
* {@code 2^62 / sum} and converting to long by casting or rounding.
*
* <p>This approach may increase the efficiency of construction. The resulting
* sampler may no longer produce <em>exact</em> samples from the distribution.
* In particular any weights with a converted frequency of zero cannot be
* sampled.
*
* @param rng Generator of uniformly distributed random numbers.
* @param weights Weights of the discrete distribution.
* @param alpha Alpha parameter.
* @return the sampler
* @throws IllegalArgumentException if {@code weights} is null or empty, a
* weight is negative, infinite or {@code NaN}, the sum of all weights is zero,
* or the size of the discrete distribution generating tree is too large.
* @see #of(UniformRandomProvider, long[])
*/
public static FastLoadedDiceRollerDiscreteSampler of(UniformRandomProvider rng,
double[] weights,
int alpha) {
final int n = checkWeightsNonZeroLength(weights);
// Convert floating-point double to a relative weight
// using a shifted integer representation
final long[] frequencies = new long[n];
final int[] offsets = new int[n];
convertToIntegers(weights, frequencies, offsets, alpha);
// Obtain indices of non-zero weights
final int[] indices = indicesOfNonZero(frequencies);
// Edge case for 1 non-zero weight.
if (indices.length == 1) {
return new FixedValueDiscreteSampler(indexOfNonZero(frequencies));
}
final BigInteger m = sum(frequencies, offsets, indices);
// Use long arithmetic if possible. This occurs when the weights are similar in magnitude.
if (m.compareTo(MAX_LONG) <= 0) {
// Apply the offset
for (int i = 0; i < n; i++) {
frequencies[i] <<= offsets[i];
}
return createSampler(rng, frequencies, indices, m.longValue());
}
return createSampler(rng, frequencies, offsets, indices, m);
}
/**
* Sum the frequencies.
*
* @param frequencies Frequencies.
* @return the sum
* @throws IllegalArgumentException if {@code frequencies} is null or empty, a
* frequency is negative, or the sum of all frequencies is either zero or above
* {@link Long#MAX_VALUE}
*/
private static long sum(long[] frequencies) {
// Validate
if (frequencies == null || frequencies.length == 0) {
throw new IllegalArgumentException("frequencies must contain at least 1 value");
}
// Sum the values.
// Combine all the sign bits in the observations and the intermediate sum in a flag.
long m = 0;
long signFlag = 0;
for (final long o : frequencies) {
m += o;
signFlag |= o | m;
}
// Check for a sign-bit.
if (signFlag < 0) {
// One or more observations were negative, or the sum overflowed.
for (final long o : frequencies) {
if (o < 0) {
throw new IllegalArgumentException("frequencies must contain positive values: " + o);
}
}
throw new IllegalArgumentException("Overflow when summing frequencies");
}
if (m == 0) {
throw new IllegalArgumentException("Sum of frequencies is zero");
}
return m;
}
/**
* Convert the floating-point weights to relative weights represented as
* integers {@code value * 2^exponent}. The relative weight as an integer is:
*
* <pre>
* BigInteger.valueOf(value).shiftLeft(exponent)
* </pre>
*
* <p>Note that the weights are created using a common power-of-2 scaling
* operation so the minimum exponent is zero.
*
* <p>A positive {@code alpha} parameter is used to set any weight to zero if
* the magnitude is approximately 2<sup>alpha</sup> <em>smaller</em> than the
* largest weight. This comparison is made using only the exponent of the input
* weights.
*
* @param weights Weights of the discrete distribution.
* @param values Output floating-point mantissas converted to 53-bit integers.
* @param exponents Output power of 2 exponent.
* @param alpha Alpha parameter.
* @throws IllegalArgumentException if a weight is negative, infinite or
* {@code NaN}, or the sum of all weights is zero.
*/
private static void convertToIntegers(double[] weights, long[] values, int[] exponents, int alpha) {
int maxExponent = Integer.MIN_VALUE;
for (int i = 0; i < weights.length; i++) {
final double weight = weights[i];
// Ignore zero.
// When creating the integer value later using bit shifts the result will remain zero.
if (weight == 0) {
continue;
}
final long bits = Double.doubleToRawLongBits(weight);
// For the IEEE 754 format see Double.longBitsToDouble(long).
// Extract the exponent (with the sign bit)
int exp = (int) (bits >>> MANTISSA_SIZE);
// Detect negative, infinite or NaN.
// Note: Negative values sign bit will cause the exponent to be too high.
if (exp > MAX_BIASED_EXPONENT) {
throw new IllegalArgumentException("Invalid weight: " + weight);
}
long mantissa;
if (exp == 0) {
// Sub-normal number:
mantissa = (bits & MANTISSA_MASK) << 1;
// Here we convert to a normalised number by counting the leading zeros
// to obtain the number of shifts of the most significant bit in
// the mantissa that is required to get a 1 at position 53 (i.e. as
// if it were a normal number with assumed leading bit).
final int shift = Long.numberOfLeadingZeros(mantissa << 11);
mantissa <<= shift;
exp -= shift;
} else {
// Normal number. Add the implicit leading 1-bit.
mantissa = (bits & MANTISSA_MASK) | (1L << MANTISSA_SIZE);
}
// Here the floating-point number is equal to:
// mantissa * 2^(exp-1075)
values[i] = mantissa;
exponents[i] = exp;
maxExponent = Math.max(maxExponent, exp);
}
// No exponent indicates that all weights are zero
if (maxExponent == Integer.MIN_VALUE) {
throw new IllegalArgumentException("Sum of weights is zero");
}
filterWeights(values, exponents, alpha, maxExponent);
scaleWeights(values, exponents);
}
/**
* Filters small weights using the {@code alpha} parameter.
* A positive {@code alpha} parameter is used to set any weight to zero if
* the magnitude is approximately 2<sup>alpha</sup> <em>smaller</em> than the
* largest weight. This comparison is made using only the exponent of the input
* weights.
*
* @param values 53-bit values.
* @param exponents Power of 2 exponent.
* @param alpha Alpha parameter.
* @param maxExponent Maximum exponent.
*/
private static void filterWeights(long[] values, int[] exponents, int alpha, int maxExponent) {
if (alpha > 0) {
// Filter weights. This must be done before the values are shifted so
// the exponent represents the approximate magnitude of the value.
for (int i = 0; i < exponents.length; i++) {
if (maxExponent - exponents[i] > alpha) {
values[i] = 0;
}
}
}
}
/**
* Scale the weights represented as integers {@code value * 2^exponent} to use a
* minimum exponent of zero. The values are scaled to remove any common trailing zeros
* in their representation. This ultimately reduces the size of the discrete distribution
* generating (DGG) tree.
*
* @param values 53-bit values.
* @param exponents Power of 2 exponent.
*/
private static void scaleWeights(long[] values, int[] exponents) {
// Find the minimum exponent and common trailing zeros.
int minExponent = Integer.MAX_VALUE;
for (int i = 0; i < exponents.length; i++) {
if (values[i] != 0) {
minExponent = Math.min(minExponent, exponents[i]);
}
}
// Trailing zeros occur when the original weights have a representation with
// less than 52 binary digits, e.g. {1.5, 0.5, 0.25}.
int trailingZeros = Long.SIZE;
for (int i = 0; i < values.length && trailingZeros != 0; i++) {
trailingZeros = Math.min(trailingZeros, Long.numberOfTrailingZeros(values[i]));
}
// Scale by a power of 2 so the minimum exponent is zero.
for (int i = 0; i < exponents.length; i++) {
exponents[i] -= minExponent;
}
// Remove common trailing zeros.
if (trailingZeros != 0) {
for (int i = 0; i < values.length; i++) {
values[i] >>>= trailingZeros;
}
}
}
/**
* Sum the integers at the specified indices.
* Integers are represented as {@code value * 2^exponent}.
*
* @param values 53-bit values.
* @param exponents Power of 2 exponent.
* @param indices Indices to sum.
* @return the sum
*/
private static BigInteger sum(long[] values, int[] exponents, int[] indices) {
BigInteger m = BigInteger.ZERO;
for (final int i : indices) {
m = m.add(toBigInteger(values[i], exponents[i]));
}
return m;
}
/**
* Convert the value and left shift offset to a BigInteger.
* It is assumed the value is at most 53-bits. This allows optimising the left
* shift if it is below 11 bits.
*
* @param value 53-bit value.
* @param offset Left shift offset (must be positive).
* @return the BigInteger
*/
private static BigInteger toBigInteger(long value, int offset) {
// Ignore zeros. The sum method uses indices of non-zero values.
if (offset <= MAX_OFFSET) {
// Assume (value << offset) <= Long.MAX_VALUE
return BigInteger.valueOf(value << offset);
}
return BigInteger.valueOf(value).shiftLeft(offset);
}
/**
* Creates the sampler.
*
* <p>It is assumed the frequencies are all positive and the sum does not
* overflow.
*
* @param rng Generator of uniformly distributed random numbers.
* @param frequencies Observed frequencies of the discrete distribution.
* @param indices Indices of non-zero frequencies.
* @param m Sum of the frequencies.
* @return the sampler
*/
private static FastLoadedDiceRollerDiscreteSampler createSampler(UniformRandomProvider rng,
long[] frequencies,
int[] indices,
long m) {
// ALGORITHM 5: PREPROCESS
// a == frequencies
// m = sum(a)
// h = leaf node count
// H = leaf node label (lH)
final int n = frequencies.length;
// k = ceil(log2(m))
final int k = 64 - Long.numberOfLeadingZeros(m - 1);
// r = a(n+1) = 2^k - m
final long r = (1L << k) - m;
// Note:
// A sparse matrix can often be used for H, as most of its entries are empty.
// This implementation uses a 1D array for efficiency at the cost of memory.
// This is limited to approximately ((2^31 - 1) / k), e.g. 34087042 when the sum of
// observations is large enough to create k=63.
// This could be handled using a 2D array. In practice a number of categories this
// large is not expected and is currently not supported.
final int[] h = new int[k];
final int[] lH = new int[checkArraySize((n + 1L) * k)];
Arrays.fill(lH, NO_LABEL);
int d;
for (int j = 0; j < k; j++) {
final int shift = (k - 1) - j;
final long bitMask = 1L << shift;
d = 0;
for (final int i : indices) {
// bool w ← (a[i] >> (k − 1) − j)) & 1
// h[j] = h[j] + w
// if w then:
if ((frequencies[i] & bitMask) != 0) {
h[j]++;
// H[d][j] = i
lH[d * k + j] = i;
d++;
}
}
// process a(n+1) without extending the input frequencies array by 1
if ((r & bitMask) != 0) {
h[j]++;
lH[d * k + j] = n;
}
}
return new FLDRSampler(rng, n, k, h, lH);
}
/**
* Creates the sampler. Frequencies are represented as a 53-bit value with a
* left-shift offset.
* <pre>
* BigInteger.valueOf(value).shiftLeft(offset)
* </pre>
*
* <p>It is assumed the frequencies are all positive.
*
* @param rng Generator of uniformly distributed random numbers.
* @param frequencies Observed frequencies of the discrete distribution.
* @param offsets Left shift offsets (must be positive).
* @param indices Indices of non-zero frequencies.
* @param m Sum of the frequencies.
* @return the sampler
*/
private static FastLoadedDiceRollerDiscreteSampler createSampler(UniformRandomProvider rng,
long[] frequencies,
int[] offsets,
int[] indices,
BigInteger m) {
// Repeat the logic from createSampler(...) using extended arithmetic to test the bits
// ALGORITHM 5: PREPROCESS
// a == frequencies
// m = sum(a)
// h = leaf node count
// H = leaf node label (lH)
final int n = frequencies.length;
// k = ceil(log2(m))
final int k = m.subtract(BigInteger.ONE).bitLength();
// r = a(n+1) = 2^k - m
final BigInteger r = BigInteger.ONE.shiftLeft(k).subtract(m);
final int[] h = new int[k];
final int[] lH = new int[checkArraySize((n + 1L) * k)];
Arrays.fill(lH, NO_LABEL);
int d;
for (int j = 0; j < k; j++) {
final int shift = (k - 1) - j;
d = 0;
for (final int i : indices) {
// bool w ← (a[i] >> (k − 1) − j)) & 1
// h[j] = h[j] + w
// if w then:
if (testBit(frequencies[i], offsets[i], shift)) {
h[j]++;
// H[d][j] = i
lH[d * k + j] = i;
d++;
}
}
// process a(n+1) without extending the input frequencies array by 1
if (r.testBit(shift)) {
h[j]++;
lH[d * k + j] = n;
}
}
return new FLDRSampler(rng, n, k, h, lH);
}
/**
* Test the logical bit of the shifted integer representation.
* The value is assumed to have at most 53-bits of information. The offset
* is assumed to be positive. This is functionally equivalent to:
* <pre>
* BigInteger.valueOf(value).shiftLeft(offset).testBit(n)
* </pre>
*
* @param value 53-bit value.
* @param offset Left shift offset.
* @param n Index of bit to test.
* @return true if the bit is 1
*/
private static boolean testBit(long value, int offset, int n) {
if (n < offset) {
// All logical trailing bits are zero
return false;
}
// Test if outside the 53-bit value (note that the implicit 1 bit
// has been added to the 52-bit mantissas for 'normal' floating-point numbers).
final int bit = n - offset;
return bit <= MANTISSA_SIZE && (value & (1L << bit)) != 0;
}
/**
* Check the weights have a non-zero length.
*
* @param weights Weights.
* @return the length
*/
private static int checkWeightsNonZeroLength(double[] weights) {
if (weights == null || weights.length == 0) {
throw new IllegalArgumentException("weights must contain at least 1 value");
}
return weights.length;
}
/**
* Create the indices of non-zero values.
*
* @param values Values.
* @return the indices
*/
private static int[] indicesOfNonZero(long[] values) {
int n = 0;
final int[] indices = new int[values.length];
for (int i = 0; i < values.length; i++) {
if (values[i] != 0) {
indices[n++] = i;
}
}
return Arrays.copyOf(indices, n);
}
/**
* Find the index of the first non-zero frequency.
*
* @param frequencies Frequencies.
* @return the index
* @throws IllegalStateException if all frequencies are zero.
*/
static int indexOfNonZero(long[] frequencies) {
for (int i = 0; i < frequencies.length; i++) {
if (frequencies[i] != 0) {
return i;
}
}
throw new IllegalStateException("All frequencies are zero");
}
/**
* Check the size is valid for a 1D array.
*
* @param size Size
* @return the size as an {@code int}
* @throws IllegalArgumentException if the size is too large for a 1D array.
*/
static int checkArraySize(long size) {
if (size > MAX_ARRAY_SIZE) {
throw new IllegalArgumentException("Unable to allocate array of size: " + size);
}
return (int) size;
}
}