tensorflow_probability/python/distributions/von_mises.py
def _kl_von_mises_von_mises(d1, d2, name=None):
  """Batchwise KL divergence KL(d1 || d2) with d1 and d2 von Mises.

  Args:
    d1: instance of a von Mises distribution object.
    d2: instance of a a von Mises distribution object.
    name: (optional) Name to use for created operations.
      default is "kl_von_mises_von_mises".

  Returns:
    Batchwise KL(d1 || d2)
  """
  with tf.name_scope(name or "kl_von_mises_von_mises"):
    # The density of von Mises is (abbreviating the concentration for conc):
    #   vonMises(x; loc, conc) = exp(conc cos(x - loc)) / (2 pi I_0 (conc) )
    # We need two properties:
    # 1. Standardization: if z ~ vonMises(0, conc), then
    #    z + loc ~ vonMises(loc, conc).
    # 2. Expectation of cosine:
    #    E_q(z | 0, conc) cos z = I_1 (conc) / I_0 (conc)
    # Now,
    # KL(d1 || d2)
    #   = E_vonMises(x; loc1, conc1) log vonMises(x; loc1, conc1)
    #                                     / vonMises(x; loc2, conc2)
    # Plugging the densities and rearranging, we have
    #   log I_0(conc2) / I_0(conc1)
    #     + E_vonMises(x; loc1, conc1) [ conc1 cos (z - loc1)
    #                                     - conc2 cos (z - loc2) ]
    # Let's transform the second term using the standardization property:
    #   E_vonMises(x; 0, conc1) [conc1 cos z - conc2 cos (z - (loc2 - loc1))]
    # Applying the cos (x - y) = cos x cos y + sin x sin y expansion, we get
    #   E_vonMises(x; 0, conc1) [conc1 cos z - conc2 cos (loc2 - loc1) cos z
    #     - conc2 sin(loc2 - loc1) sin z]
    # Because the distribution is symmetric around zero, the last term vanishes
    # in expectation. The remaining two terms are computed using the
    # "expectation of cosine" property:
    #   (conc1 - conc2 cos (loc2 - loc1) E_vonMises(x; 0, conc1) cos z
    #     = (conc1 - conc2 cos (loc2 - loc1)) I_1(conc1) / I_0(conc1)
    # In total, we have
    #   KL(d1 || d2) = log I_0(conc2) / I_0(conc1)
    #     + (conc1 - conc2 cos (loc2 - loc1)) I_1(conc1) / I_0(conc1)
    # To improve the numerical stability, we can replace I_j(k) functions with
    # the exponentially scaled versions using the equality
    # I_j(k) = I_j^E(k) exp(k) (which holds for k >= 0):
    #   KL(d1 || d2) = (conc2 - conc1) + log I_0^E(conc2) / I_0^E(conc1)
    #     + (conc1 - conc2 cos (loc2 - loc1)) I_1^E(conc1) / I_0^E(conc1)
    # Note that this formula is numerically stable for conc1 = 0 and/or
    # conc2 = 0 because I_0 (0) = I_0^E (0) = 1.
    i0e_concentration1 = tf.math.bessel_i0e(d1.concentration)
    i1e_concentration1 = tf.math.bessel_i1e(d1.concentration)
    i0e_concentration2 = tf.math.bessel_i0e(d2.concentration)
    return ((d2.concentration - d1.concentration) +
            tf.math.log(i0e_concentration2 / i0e_concentration1) +
            (d1.concentration - d2.concentration * tf.cos(d1.loc - d2.loc)) *
            (i1e_concentration1 / i0e_concentration1))
def _kl_von_mises_von_mises(d1, d2, name=None):
  """Batchwise KL divergence KL(d1 || d2) with d1 and d2 von Mises.

  Args:
    d1: instance of a von Mises distribution object.
    d2: instance of a a von Mises distribution object.
    name: (optional) Name to use for created operations.
      default is "kl_von_mises_von_mises".

  Returns:
    Batchwise KL(d1 || d2)
  """
  with tf.name_scope(name or "kl_von_mises_von_mises"):
    # The density of von Mises is (abbreviating the concentration for conc):
    #   vonMises(x; loc, conc) = exp(conc cos(x - loc)) / (2 pi I_0 (conc) )
    # We need two properties:
    # 1. Standardization: if z ~ vonMises(0, conc), then
    #    z + loc ~ vonMises(loc, conc).
    # 2. Expectation of cosine:
    #    E_q(z | 0, conc) cos z = I_1 (conc) / I_0 (conc)
    # Now,
    # KL(d1 || d2)
    #   = E_vonMises(x; loc1, conc1) log vonMises(x; loc1, conc1)
    #                                     / vonMises(x; loc2, conc2)
    # Plugging the densities and rearranging, we have
    #   log I_0(conc2) / I_0(conc1)
    #     + E_vonMises(x; loc1, conc1) [ conc1 cos (z - loc1)
    #                                     - conc2 cos (z - loc2) ]
    # Let's transform the second term using the standardization property:
    #   E_vonMises(x; 0, conc1) [conc1 cos z - conc2 cos (z - (loc2 - loc1))]
    # Applying the cos (x - y) = cos x cos y + sin x sin y expansion, we get
    #   E_vonMises(x; 0, conc1) [conc1 cos z - conc2 cos (loc2 - loc1) cos z
    #     - conc2 sin(loc2 - loc1) sin z]
    # Because the distribution is symmetric around zero, the last term vanishes
    # in expectation. The remaining two terms are computed using the
    # "expectation of cosine" property:
    #   (conc1 - conc2 cos (loc2 - loc1) E_vonMises(x; 0, conc1) cos z
    #     = (conc1 - conc2 cos (loc2 - loc1)) I_1(conc1) / I_0(conc1)
    # In total, we have
    #   KL(d1 || d2) = log I_0(conc2) / I_0(conc1)
    #     + (conc1 - conc2 cos (loc2 - loc1)) I_1(conc1) / I_0(conc1)
    # To improve the numerical stability, we can replace I_j(k) functions with
    # the exponentially scaled versions using the equality
    # I_j(k) = I_j^E(k) exp(k) (which holds for k >= 0):
    #   KL(d1 || d2) = (conc2 - conc1) + log I_0^E(conc2) / I_0^E(conc1)
    #     + (conc1 - conc2 cos (loc2 - loc1)) I_1^E(conc1) / I_0^E(conc1)
    # Note that this formula is numerically stable for conc1 = 0 and/or
    # conc2 = 0 because I_0 (0) = I_0^E (0) = 1.
    i0e_concentration1 = tf.math.bessel_i0e(d1.concentration)
    i1e_concentration1 = tf.math.bessel_i1e(d1.concentration)
    i0e_concentration2 = tf.math.bessel_i0e(d2.concentration)
    return ((d2.concentration - d1.concentration) +
            tf.math.log(i0e_concentration2 / i0e_concentration1) +
            (d1.concentration - d2.concentration * tf.cos(d1.loc - d2.loc)) *
            (i1e_concentration1 / i0e_concentration1))
Batchwise KL divergence KL(d1 || d2) with d1 and d2 von Mises.

  Args:
    d1: instance of a von Mises distribution object.
    d2: instance of a a von Mises distribution object.
    name: (optional) Name to use for created operations.
      default is "kl_von_mises_von_mises".

  Returns:
    Batchwise KL(d1 || d2)