tfp.substrates.jax.mcmc.sample_halton_sequence

View source on GitHub

Returns a sample from the dim dimensional Halton sequence.

tfp.substrates.jax.mcmc.sample_halton_sequence(
    dim,
    num_results=None,
    sequence_indices=None,
    dtype=tf.float32,
    randomized=True,
    seed=None,
    name=None
)

Computes the members of the low discrepancy Halton sequence in dimension dim. The dim-dimensional sequence takes values in the unit hypercube in dim dimensions. Currently, only dimensions up to 10000 are supported. The prime base for the k-th axes is the k-th prime starting from 2. For example, if dim = 3, then the bases will be [2, 3, 5] respectively and the first element of the non-randomized sequence will be: [0.5, 0.333, 0.2]. For a more complete description of the Halton sequences see here. For low discrepancy sequences and their applications see here.

If randomized is true, this function produces a scrambled version of the Halton sequence introduced by [Owen (2017)][1]. For the advantages of randomization of low discrepancy sequences see here.

The number of samples produced is controlled by the num_results and sequence_indices parameters. The user must supply either num_results or sequence_indices but not both. The former is the number of samples to produce starting from the first element. If sequence_indices is given instead, the specified elements of the sequence are generated. For example, sequence_indices=tf.range(10) is equivalent to specifying n=10.

Examples

from tensorflow_probability.python.internal.backend import jax as tf
import tensorflow_probability as tfp; tfp = tfp.substrates.jax

# Produce the first 1000 members of the Halton sequence in 3 dimensions.
num_results = 1000
dim = 3
sample = tfp.mcmc.sample_halton_sequence(
  dim,
  num_results=num_results,
  seed=127)

# Evaluate the integral of x_1 * x_2^2 * x_3^3  over the three dimensional
# hypercube.
powers = tf.range(1.0, limit=dim + 1)
integral = tf.reduce_mean(tf.reduce_prod(sample ** powers, axis=-1))
true_value = 1.0 / tf.reduce_prod(powers + 1.0)
with tf.Session() as session:
  values = session.run((integral, true_value))

# Produces a relative absolute error of 1.7%.
print ("Estimated: %f, True Value: %f" % values)

# Now skip the first 1000 samples and recompute the integral with the next
# thousand samples. The sequence_indices argument can be used to do this.


sequence_indices = tf.range(start=1000, limit=1000 + num_results,
                            dtype=tf.int32)
sample_leaped = tfp.mcmc.sample_halton_sequence(
    dim,
    sequence_indices=sequence_indices,
    seed=111217)

integral_leaped = tf.reduce_mean(tf.reduce_prod(sample_leaped ** powers,
                                                axis=-1))
with tf.Session() as session:
  values = session.run((integral_leaped, true_value))
# Now produces a relative absolute error of 0.05%.
print ("Leaped Estimated: %f, True Value: %f" % values)

References

[1]: Art B. Owen. A randomized Halton algorithm in R. arXiv preprint arXiv:1706.02808, 2017. https://arxiv.org/abs/1706.02808