Photrek / Nonlinear-Statistical-Coupling

Apache License 2.0
5 stars 1 forks source link

CoupledNormal - log_prob function #6

Closed Kevin-Chen0 closed 3 years ago

Kevin-Chen0 commented 3 years ago

Implement the log PDF for CoupledNormal. Use the commented out TFP's Student log_prob code as a foundation.

See our latest nsc code here.

See original StudentT code here:

def log_prob(x, df, loc, scale):
  """Compute log probability of Student T distribution.
  Note that scale can be negative.
  Args:
    x: Floating-point `Tensor`. Where to compute the log probabilities.
    df: Floating-point `Tensor`. The degrees of freedom of the
      distribution(s). `df` must contain only positive values.
    loc: Floating-point `Tensor`; the location(s) of the distribution(s).
    scale: Floating-point `Tensor`; the scale(s) of the distribution(s).
  Returns:
    A `Tensor` with shape broadcast according to the arguments.
  """
  # Writing `y` this way reduces XLA mem copies.
  y = (x - loc) * (tf.math.rsqrt(df) / scale)
  log_unnormalized_prob = -0.5 * (df + 1.) * log1psquare(y)
  log_normalization = (
      tf.math.log(tf.abs(scale)) + 0.5 * tf.math.log(df) +
      0.5 * np.log(np.pi) + tfp_math.log_gamma_difference(0.5, 0.5 * df))
  return log_unnormalized_prob - log_normalization