Add quantiles to Student-T, Beta and SigmoidBeta.

PiperOrigin-RevId: 473158097

Author: srvasude

tensorflow_probability/python/distributions/beta.py

@@ -332,6 +332,12 @@ def _cdf(self, x):
     return distribution_util.extend_cdf_outside_support(
         x, answer, low=0., high=1.)
 
+  @distribution_util.AppendDocstring(_beta_sample_note)
+  def _quantile(self, p):
+    concentration1 = tf.convert_to_tensor(self.concentration1)
+    concentration0 = tf.convert_to_tensor(self.concentration0)
+    return special.betaincinv(concentration1, concentration0, p)
+
   def _log_unnormalized_prob(self, x, concentration1, concentration0):
     return (tf.math.xlogy(concentration1 - 1., x) +
             tf.math.xlog1py(concentration0 - 1., -x))

tensorflow_probability/python/distributions/beta_test.py

@@ -291,7 +291,7 @@ def testBetaSampleMultidimensional(self):
 
   @parameterized.parameters((np.float32, 5e-3), (np.float64, 1e-4))
   def testBetaCdf(self, dt, rtol):
-    shape = (30, 40, 50)
+    shape = (30, 4, 5)
     a = 10. * np.random.random(shape).astype(dt)
     b = 10. * np.random.random(shape).astype(dt)
     x = np.random.random(shape).astype(dt)
@@ -304,9 +304,27 @@ def testBetaCdfBeyondSupport(self):
     cdf = beta.Beta(2., 3., validate_args=False).cdf([-3.7, 1.03])
     self.assertAllEqual([0., 1.], self.evaluate(cdf))
 
+  @parameterized.parameters((np.float32, 5e-3), (np.float64, 1e-4))
+  def testBetaQuantile(self, dt, rtol):
+    shape = (30, 4, 5)
+    a = 5. * np.random.random(shape).astype(dt)
+    b = 5. * np.random.random(shape).astype(dt)
+    p = np.random.uniform(low=0., high=1., size=shape).astype(dt)
+    quantile = tf.function(beta.Beta(a, b).quantile)
+    actual = self.evaluate(quantile(p))
+    # Pass f64 values to avoid errors in scipy.
+    self.assertAllClose(
+        sp_stats.beta.ppf(
+            p.astype(np.float64),
+            a.astype(np.float64),
+            b.astype(np.float64)),
+        actual,
+        rtol=rtol,
+        atol=1e-10)
+
   @parameterized.parameters((np.float32, 5e-3), (np.float64, 1e-4))
   def testBetaLogCdf(self, dt, rtol):
-    shape = (30, 40, 50)
+    shape = (30, 4, 5)
     a = 10. * np.random.random(shape).astype(dt)
     b = 10. * np.random.random(shape).astype(dt)
     x = np.random.random(shape).astype(dt)

tensorflow_probability/python/distributions/sigmoid_beta.py

@@ -244,5 +244,11 @@ def _cdf(self, x):
     sig_x = tf.math.sigmoid(x)
     return special.betainc(concentration1, concentration0, sig_x)
 
+  def _quantile(self, p):
+    concentration1 = tf.convert_to_tensor(self.concentration1)
+    concentration0 = tf.convert_to_tensor(self.concentration0)
+    y = special.betaincinv(concentration1, concentration0, p)
+    return tf.math.log(y) - tf.math.log1p(-y)
+
   def _mode(self):
     return tf.math.log(self.concentration1 / self.concentration0)

tensorflow_probability/python/distributions/sigmoid_beta_test.py

@@ -176,6 +176,20 @@ def testCDF(self):
         bijector=invert.Invert(sigmoid.Sigmoid())).cdf(x)
     self.assertAllClose(cdf, expected_cdf)
 
+  def testQuantile(self):
+    batch_shape = [6, 1]
+    a = 2. * np.ones(batch_shape, dtype=np.float32)
+    b = 3. * np.ones(batch_shape, dtype=np.float32)
+    p = np.logspace(-4., -0.01, 20).astype(np.float32)
+    dist = sigmoid_beta.SigmoidBeta(a, b, validate_args=True)
+    quantile = dist.quantile(p)
+    self.assertEqual(quantile.shape, (6, 20))
+
+    expected_quantile = transformed_distribution.TransformedDistribution(
+        distribution=beta.Beta(a, b),
+        bijector=invert.Invert(sigmoid.Sigmoid())).quantile(p)
+    self.assertAllClose(quantile, expected_quantile)
+
   def testMode(self):
     a = tf.constant(1.0)
     b = tf.constant(2.0)

tensorflow_probability/python/distributions/student_t.py

@@ -120,6 +120,28 @@ def cdf(x, df, loc, scale):
   return tf.where(y < 0., neg_cdf, 1. - neg_cdf)
 
 
+def quantile(p, df, loc, scale):
+  """Compute quantile function of Student T distribution.
+
+  Note that scale can be negative.
+
+  Args:
+    p: Floating-point `Tensor`. Probabilities from 0 to 1.
+    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.
+  """
+  df = tf.convert_to_tensor(df)
+  p_adjusted = tf.where(p < 0.5, p, 1. - p)
+  y = special.betaincinv(0.5 * df, 0.5, 2 * p_adjusted)
+  return loc + tf.math.abs(
+      scale) * tf.math.sign(p - 0.5) * tf.math.sqrt(df * (1 - y) / y)
+
+
 def entropy(df, scale, batch_shape, dtype):
   """Compute entropy of the StudentT distribution.
 
@@ -350,6 +372,10 @@ def _cdf(self, x):
     df = tf.convert_to_tensor(self.df)
     return cdf(x, df, self.loc, self.scale)
 
+  def _quantile(self, x):
+    df = tf.convert_to_tensor(self.df)
+    return quantile(x, df, self.loc, self.scale)
+
   def _entropy(self):
     df = tf.convert_to_tensor(self.df)
     scale = tf.convert_to_tensor(self.scale)

tensorflow_probability/python/distributions/student_t_test.py

@@ -42,10 +42,10 @@ def testStudentPDFAndLogPDF(self):
     student = student_t.StudentT(df, loc=mu, scale=-sigma, validate_args=True)
 
     log_pdf = student.log_prob(t)
-    self.assertEquals(log_pdf.shape, (6,))
+    self.assertAllEqual(log_pdf.shape, (6,))
     log_pdf_values = self.evaluate(log_pdf)
     pdf = student.prob(t)
-    self.assertEquals(pdf.shape, (6,))
+    self.assertAllEqual(pdf.shape, (6,))
     pdf_values = self.evaluate(pdf)
 
     expected_log_pdf = sp_stats.t.logpdf(t, df_v, loc=mu_v, scale=sigma_v)
@@ -92,10 +92,10 @@ def testStudentCDFAndLogCDF(self):
     student = student_t.StudentT(df, loc=mu, scale=sigma, validate_args=True)
 
     log_cdf = student.log_cdf(t)
-    self.assertEquals(log_cdf.shape, (6,))
+    self.assertAllEqual(log_cdf.shape, (6,))
     log_cdf_values = self.evaluate(log_cdf)
     cdf = student.cdf(t)
-    self.assertEquals(cdf.shape, (6,))
+    self.assertAllEqual(cdf.shape, (6,))
     cdf_values = self.evaluate(cdf)
 
     expected_log_cdf = sp_stats.t.logcdf(t, df_v, loc=mu_v, scale=sigma_v)
@@ -107,6 +107,23 @@ def testStudentCDFAndLogCDF(self):
     self.assertAllClose(
         np.exp(expected_log_cdf), cdf_values, atol=0., rtol=1e-5)
 
+  def testStudentQuantile(self):
+    batch_shape = (40, 1)
+    df = np.random.uniform(
+        low=0.5, high=10., size=batch_shape).astype(np.float32)
+    mu = 7. * np.ones(batch_shape, dtype=np.float32)
+    sigma = -8.
+    p = np.logspace(-4., -0.01, 20).astype(np.float32)
+    student = student_t.StudentT(
+        df, loc=mu, scale=sigma, validate_args=True)
+
+    quantile = student.quantile(p)
+    self.assertAllEqual(quantile.shape, (40, 20))
+    quantile_values = self.evaluate(quantile)
+
+    expected_quantile = sp_stats.t.ppf(p, df, loc=mu, scale=np.abs(sigma))
+    self.assertAllClose(expected_quantile, quantile_values, rtol=5e-4)
+
   def testStudentEntropy(self):
     df_v = np.array([[2., 3., 7.]])  # 1x3
     mu_v = np.array([[1., -1, 0]])  # 1x3