tensorflow_probability/python/distributions/von_mises.py
def random_von_mises(shape, concentration, dtype=tf.float32, seed=None):
  """Samples from the standardized von Mises distribution.

  The distribution is vonMises(loc=0, concentration=concentration), so the mean
  is zero.
  The location can then be changed by adding it to the samples.

  The sampling algorithm is rejection sampling with wrapped Cauchy proposal [1].
  The samples are pathwise differentiable using the approach of [2].

  Arguments:
    shape: The output sample shape.
    concentration: The concentration parameter of the von Mises distribution.
    dtype: The data type of concentration and the outputs.
    seed: (optional) The random seed.

  Returns:
    Differentiable samples of standardized von Mises.

  References:
    [1] Luc Devroye "Non-Uniform Random Variate Generation", Springer-Verlag,
    1986; Chapter 9, p. 473-476.
    http://www.nrbook.com/devroye/Devroye_files/chapter_nine.pdf
    + corrections http://www.nrbook.com/devroye/Devroye_files/errors.pdf
    [2] Michael Figurnov, Shakir Mohamed, Andriy Mnih. "Implicit
    Reparameterization Gradients", 2018.
  """
  seed = SeedStream(seed, salt="von_mises")
  concentration = tf.convert_to_tensor(
      value=concentration, dtype=dtype, name="concentration")

  @tf.custom_gradient
  def rejection_sample_with_gradient(concentration):
    """Performs rejection sampling for standardized von Mises.

    A nested function is required because @tf.custom_gradient does not handle
    non-tensor inputs such as dtype. Instead, they are captured by the outer
    scope.

    Arguments:
      concentration: The concentration parameter of the distribution.

    Returns:
      Differentiable samples of standardized von Mises.
    """
    r = 1. + tf.sqrt(1. + 4. * concentration ** 2)
    rho = (r - tf.sqrt(2. * r)) / (2. * concentration)

    s_exact = (1. + rho ** 2) / (2. * rho)

    # For low concentration, s becomes numerically unstable.
    # To fix that, we use an approximation. Here is the derivation.
    # First-order Taylor expansion at conc = 0 gives
    #   sqrt(1 + 4 concentration^2) ~= 1 + (2 concentration)^2 / 2.
    # Therefore, r ~= 2 + 2 concentration. By plugging this into rho, we have
    #   rho ~= conc + 1 / conc - sqrt(1 + 1 / concentration^2).
    # Let's expand the last term at concentration=0 up to the linear term:
    #   sqrt(1 + 1 / concentration^2) ~= 1 / concentration + concentration / 2
    # Thus, rho ~= concentration / 2. Finally,
    #   s = 1 / (2 rho) + rho / 2 ~= 1 / concentration + concentration / 4.
    # Since concentration is small, we drop the second term and simply use
    #   s ~= 1 / concentration.
    s_approximate = 1. / concentration

    # To compute the cutoff, we compute s_exact using mpmath with 30 decimal
    # digits precision and compare that to the s_exact and s_approximate
    # computed with dtype. Then, the cutoff is the largest concentration for
    # which abs(s_exact - s_exact_mpmath) > abs(s_approximate - s_exact_mpmath).
    s_concentration_cutoff_dict = {
        tf.float16: 1.8e-1,
        tf.float32: 2e-2,
        tf.float64: 1.2e-4,
    }
    s_concentration_cutoff = s_concentration_cutoff_dict[dtype]

    s = tf.where(concentration > s_concentration_cutoff, s_exact, s_approximate)

    def loop_body(done, u, w):
      """Resample the non-accepted points."""
      # We resample u each time completely. Only its sign is used outside the
      # loop, which is random.
      u = tf.random.uniform(
          shape, minval=-1., maxval=1., dtype=dtype, seed=seed())
      z = tf.cos(np.pi * u)
      # Update the non-accepted points.
      w = tf.where(done, w, (1. + s * z) / (s + z))
      y = concentration * (s - w)

      v = tf.random.uniform(
          shape, minval=0., maxval=1., dtype=dtype, seed=seed())
      accept = (y * (2. - y) >= v) | (tf.math.log(y / v) + 1. >= y)

      return done | accept, u, w

    _, u, w = tf.while_loop(
        cond=lambda done, *_: ~tf.reduce_all(input_tensor=done),
        body=loop_body,
        loop_vars=(
            tf.zeros(shape, dtype=tf.bool, name="done"),
            tf.zeros(shape, dtype=dtype, name="u"),
            tf.zeros(shape, dtype=dtype, name="w"),
        ),
        # The expected number of iterations depends on concentration.
        # It monotonically increases from one iteration for concentration = 0 to
        # sqrt(2 pi / e) ~= 1.52 iterations for concentration = +inf [1].
        # We use a limit of 100 iterations to avoid infinite loops
        # for very large / nan concentration.
        maximum_iterations=100,
        parallel_iterations=1 if seed.original_seed is None else 10,
    )

    x = tf.sign(u) * tf.math.acos(w)

    def grad(dy):
      """The gradient of the von Mises samples w.r.t. concentration."""
      broadcast_concentration = concentration + tf.zeros_like(x)
      _, dcdf_dconcentration = value_and_gradient(
          lambda conc: von_mises_cdf(x, conc), broadcast_concentration)
      inv_prob = tf.exp(-broadcast_concentration * (tf.cos(x) - 1.)) * (
          (2. * np.pi) * tf.math.bessel_i0e(broadcast_concentration))
      # Compute the implicit reparameterization gradient [2],
      # dz/dconc = -(dF(z; conc) / dconc) / p(z; conc)
      ret = dy * (-inv_prob * dcdf_dconcentration)
      # Sum over the sample dimensions. Assume that they are always the first
      # ones.
      num_sample_dimensions = (tf.rank(broadcast_concentration) -
                               tf.rank(concentration))
      return tf.reduce_sum(
          input_tensor=ret, axis=tf.range(num_sample_dimensions))

    return x, grad

  return rejection_sample_with_gradient(concentration)
def random_von_mises(shape, concentration, dtype=tf.float32, seed=None):
  """Samples from the standardized von Mises distribution.

  The distribution is vonMises(loc=0, concentration=concentration), so the mean
  is zero.
  The location can then be changed by adding it to the samples.

  The sampling algorithm is rejection sampling with wrapped Cauchy proposal [1].
  The samples are pathwise differentiable using the approach of [2].

  Arguments:
    shape: The output sample shape.
    concentration: The concentration parameter of the von Mises distribution.
    dtype: The data type of concentration and the outputs.
    seed: (optional) The random seed.

  Returns:
    Differentiable samples of standardized von Mises.

  References:
    [1] Luc Devroye "Non-Uniform Random Variate Generation", Springer-Verlag,
    1986; Chapter 9, p. 473-476.
    http://www.nrbook.com/devroye/Devroye_files/chapter_nine.pdf
    + corrections http://www.nrbook.com/devroye/Devroye_files/errors.pdf
    [2] Michael Figurnov, Shakir Mohamed, Andriy Mnih. "Implicit
    Reparameterization Gradients", 2018.
  """
  seed = SeedStream(seed, salt="von_mises")
  concentration = tf.convert_to_tensor(
      value=concentration, dtype=dtype, name="concentration")

  @tf.custom_gradient
  def rejection_sample_with_gradient(concentration):
    """Performs rejection sampling for standardized von Mises.

    A nested function is required because @tf.custom_gradient does not handle
    non-tensor inputs such as dtype. Instead, they are captured by the outer
    scope.

    Arguments:
      concentration: The concentration parameter of the distribution.

    Returns:
      Differentiable samples of standardized von Mises.
    """
    r = 1. + tf.sqrt(1. + 4. * concentration ** 2)
    rho = (r - tf.sqrt(2. * r)) / (2. * concentration)

    s_exact = (1. + rho ** 2) / (2. * rho)

    # For low concentration, s becomes numerically unstable.
    # To fix that, we use an approximation. Here is the derivation.
    # First-order Taylor expansion at conc = 0 gives
    #   sqrt(1 + 4 concentration^2) ~= 1 + (2 concentration)^2 / 2.
    # Therefore, r ~= 2 + 2 concentration. By plugging this into rho, we have
    #   rho ~= conc + 1 / conc - sqrt(1 + 1 / concentration^2).
    # Let's expand the last term at concentration=0 up to the linear term:
    #   sqrt(1 + 1 / concentration^2) ~= 1 / concentration + concentration / 2
    # Thus, rho ~= concentration / 2. Finally,
    #   s = 1 / (2 rho) + rho / 2 ~= 1 / concentration + concentration / 4.
    # Since concentration is small, we drop the second term and simply use
    #   s ~= 1 / concentration.
    s_approximate = 1. / concentration

    # To compute the cutoff, we compute s_exact using mpmath with 30 decimal
    # digits precision and compare that to the s_exact and s_approximate
    # computed with dtype. Then, the cutoff is the largest concentration for
    # which abs(s_exact - s_exact_mpmath) > abs(s_approximate - s_exact_mpmath).
    s_concentration_cutoff_dict = {
        tf.float16: 1.8e-1,
        tf.float32: 2e-2,
        tf.float64: 1.2e-4,
    }
    s_concentration_cutoff = s_concentration_cutoff_dict[dtype]

    s = tf.where(concentration > s_concentration_cutoff, s_exact, s_approximate)

    def loop_body(done, u, w):
      """Resample the non-accepted points."""
      # We resample u each time completely. Only its sign is used outside the
      # loop, which is random.
      u = tf.random.uniform(
          shape, minval=-1., maxval=1., dtype=dtype, seed=seed())
      z = tf.cos(np.pi * u)
      # Update the non-accepted points.
      w = tf.where(done, w, (1. + s * z) / (s + z))
      y = concentration * (s - w)

      v = tf.random.uniform(
          shape, minval=0., maxval=1., dtype=dtype, seed=seed())
      accept = (y * (2. - y) >= v) | (tf.math.log(y / v) + 1. >= y)

      return done | accept, u, w

    _, u, w = tf.while_loop(
        cond=lambda done, *_: ~tf.reduce_all(input_tensor=done),
        body=loop_body,
        loop_vars=(
            tf.zeros(shape, dtype=tf.bool, name="done"),
            tf.zeros(shape, dtype=dtype, name="u"),
            tf.zeros(shape, dtype=dtype, name="w"),
        ),
        # The expected number of iterations depends on concentration.
        # It monotonically increases from one iteration for concentration = 0 to
        # sqrt(2 pi / e) ~= 1.52 iterations for concentration = +inf [1].
        # We use a limit of 100 iterations to avoid infinite loops
        # for very large / nan concentration.
        maximum_iterations=100,
        parallel_iterations=1 if seed.original_seed is None else 10,
    )

    x = tf.sign(u) * tf.math.acos(w)

    def grad(dy):
      """The gradient of the von Mises samples w.r.t. concentration."""
      broadcast_concentration = concentration + tf.zeros_like(x)
      _, dcdf_dconcentration = value_and_gradient(
          lambda conc: von_mises_cdf(x, conc), broadcast_concentration)
      inv_prob = tf.exp(-broadcast_concentration * (tf.cos(x) - 1.)) * (
          (2. * np.pi) * tf.math.bessel_i0e(broadcast_concentration))
      # Compute the implicit reparameterization gradient [2],
      # dz/dconc = -(dF(z; conc) / dconc) / p(z; conc)
      ret = dy * (-inv_prob * dcdf_dconcentration)
      # Sum over the sample dimensions. Assume that they are always the first
      # ones.
      num_sample_dimensions = (tf.rank(broadcast_concentration) -
                               tf.rank(concentration))
      return tf.reduce_sum(
          input_tensor=ret, axis=tf.range(num_sample_dimensions))

    return x, grad

  return rejection_sample_with_gradient(concentration)
Samples from the standardized von Mises distribution.

  The distribution is vonMises(loc=0, concentration=concentration), so the mean
  is zero.
  The location can then be changed by adding it to the samples.

  The sampling algorithm is rejection sampling with wrapped Cauchy proposal [1].
  The samples are pathwise differentiable using the approach of [2].

  Arguments:
    shape: The output sample shape.
    concentration: The concentration parameter of the von Mises distribution.
    dtype: The data type of concentration and the outputs.
    seed: (optional) The random seed.

  Returns:
    Differentiable samples of standardized von Mises.

  References:
    [1] Luc Devroye "Non-Uniform Random Variate Generation", Springer-Verlag,
    1986; Chapter 9, p. 473-476.
    http://www.nrbook.com/devroye/Devroye_files/chapter_nine.pdf
    + corrections http://www.nrbook.com/devroye/Devroye_files/errors.pdf
    [2] Michael Figurnov, Shakir Mohamed, Andriy Mnih. "Implicit
    Reparameterization Gradients", 2018.