Add quantile function for Student-T distribution

stan/math/fwd/prob.hpp

@@ -5,5 +5,6 @@
 #include <stan/math/fwd/fun/Eigen_NumTraits.hpp>
 
 #include <stan/math/fwd/prob/std_normal_log_qf.hpp>
+#include <stan/math/fwd/prob/student_t_qf.hpp>
 
 #endif

stan/math/fwd/prob/student_t_qf.hpp

@@ -0,0 +1,88 @@
+#ifndef STAN_MATH_FWD_PROB_STUDENT_T_QF_HPP
+#define STAN_MATH_FWD_PROB_STUDENT_T_QF_HPP
+
+#include <stan/math/fwd/meta.hpp>
+#include <stan/math/fwd/fun/sqrt.hpp>
+#include <stan/math/fwd/fun/inv_inc_beta.hpp>
+#include <stan/math/prim/prob/student_t_lpdf.hpp>
+
+namespace stan {
+namespace math {
+
+template <typename T_p, typename T_nu, typename T_mu, typename T_sigma,
+          typename FvarT = return_type_t<T_p, T_mu, T_sigma, T_nu>,
+          require_all_stan_scalar_t<T_p, T_mu, T_sigma, T_nu>* = nullptr,
+          require_fvar_t<FvarT>* = nullptr>
+FvarT student_t_qf(const T_p& p, const T_nu& nu, const T_mu& mu,
+                   const T_sigma& sigma) {
+  static constexpr const char* function = "student_t_qf";
+  using T_partials = partials_type_t<FvarT>;
+
+  const auto& p_val = value_of(p);
+  const auto& nu_val = value_of(nu);
+  const auto& mu_val = value_of(mu);
+  const auto& sigma_val = value_of(sigma);
+
+  check_nonnegative(function, "Degrees of freedom parameter", nu_val);
+  check_positive(function, "Scale parameter", sigma_val);
+  check_bounded(function, "Probability parameter", p_val, 0.0, 1.0);
+
+  if (p_val == 0.0) {
+    return {NEGATIVE_INFTY, 0.0};
+  } else if (p_val == 1.0) {
+    return {INFTY, 0.0};
+  } else if (p_val == 0.5) {
+    return {mu_val, 0.0};
+  }
+
+  const T_partials p_val_flip = p_val < 0.5 ? p_val : 1.0 - p_val;
+  const double p_sign = value_of_rec(p_val) < 0.5 ? -1.0 : 1.0;
+
+  T_partials ibeta_arg = inv_inc_beta(0.5 * nu_val, 0.5, 2 * p_val_flip);
+  T_partials rtn_val = mu_val
+                       + p_sign * sigma_val * sqrt(nu_val)
+                             * math::sqrt(-1.0 + 1.0 / ibeta_arg);
+
+  FvarT rtn(rtn_val, 0.0);
+
+  if constexpr (!is_constant<T_p>::value) {
+    rtn.d_ += p.d_ * exp(-student_t_lpdf(rtn_val, nu_val, mu_val, sigma_val));
+  }
+
+  if constexpr (!is_constant<T_nu>::value) {
+    const T_partials half_nu = nu_val / 2.0;
+    Eigen::Matrix<T_partials, -1, 1> hyper_arg_a(3);
+    hyper_arg_a << 0.5, half_nu, half_nu;
+    Eigen::Matrix<T_partials, -1, 1> hyper_arg_b(2);
+    hyper_arg_b << 1 + half_nu, 1 + half_nu;
+    const T_partials hyper_arg
+        = hypergeometric_pFq(hyper_arg_a, hyper_arg_b, ibeta_arg);
+    const T_partials hyper2f1
+        = hypergeometric_2F1(1, (1 + nu_val) / 2, (2 + nu_val) / 2, ibeta_arg);
+    const T_partials digamma_a1 = digamma(half_nu);
+    const T_partials digamma_a2 = digamma((1 + nu_val) / 2);
+
+    const T_partials arg_1 = (4 * hyper_arg * sqrt(1 - ibeta_arg)) / nu_val;
+    const T_partials arg_2 = -2 * hyper2f1 * (-1 + ibeta_arg)
+                             * (log(ibeta_arg) - digamma_a1 + digamma_a2);
+
+    const T_partials num1 = sigma_val * (-2 + (2 - arg_1 + arg_2) / ibeta_arg);
+    const T_partials den1 = 4 * sqrt(nu_val) * sqrt(-1 + 1 / ibeta_arg);
+    rtn.d_ += nu.d_ * p_sign * num1 / den1;
+  }
+
+  if constexpr (!is_constant<T_mu>::value) {
+    rtn.d_ += mu.d_;
+  }
+
+  if constexpr (!is_constant<T_sigma>::value) {
+    rtn.d_ += sigma.d_ * p_sign * sqrt(nu_val)
+              * math::sqrt(-1.0 + 1.0 / ibeta_arg);
+  }
+
+  return rtn;
+}
+}  // namespace math
+}  // namespace stan
+
+#endif

stan/math/prim/meta.hpp

@@ -71,6 +71,7 @@
 #include <stan/math/prim/meta/ad_promotable.hpp>
 #include <stan/math/prim/meta/append_return_type.hpp>
 #include <stan/math/prim/meta/base_type.hpp>
+#include <stan/math/prim/meta/common_container_type.hpp>
 #include <stan/math/prim/meta/contains_std_vector.hpp>
 #include <stan/math/prim/meta/contains_tuple.hpp>
 #include <stan/math/prim/meta/error_index.hpp>

stan/math/prim/meta/common_container_type.hpp

@@ -0,0 +1,77 @@
+#ifndef STAN_MATH_PRIM_META_COMMON_CONTAINER_TYPE_HPP
+#define STAN_MATH_PRIM_META_COMMON_CONTAINER_TYPE_HPP
+
+#include <stan/math/prim/meta/is_container.hpp>
+#include <stan/math/prim/meta/is_tuple.hpp>
+#include <stan/math/prim/meta/is_detected.hpp>
+#include <stan/math/prim/meta/is_stan_scalar.hpp>
+#include <stan/math/prim/meta/is_var_matrix.hpp>
+#include <stan/math/prim/meta/plain_type.hpp>
+#include <stan/math/prim/meta/return_type.hpp>
+#include <stan/math/prim/meta/promote_scalar_type.hpp>
+#include <type_traits>
+
+namespace stan {
+namespace internal {
+template <typename T1, typename T2, typename = void, typename = void>
+struct common_container_type_impl;
+
+template <typename T1, typename T2>
+struct common_container_type_impl<T1, T2, require_stan_scalar_t<T1>,
+                                  require_stan_scalar_t<T2>> {
+  using type = return_type_t<T1, T2>;
+};
+
+template <typename T1, typename T2>
+struct common_container_type_impl<T1, T2, require_container_t<T1>,
+                                  require_container_t<T2>> {
+  using return_t = return_type_t<T1, T2>;
+  using container_type_1 = math::promote_scalar_t<return_t, plain_type_t<T1>>;
+  using container_type_2 = math::promote_scalar_t<return_t, plain_type_t<T2>>;
+  using type = std::conditional_t<
+      std::is_same<container_type_1, container_type_2>::value, container_type_1,
+      void>;
+};
+
+template <typename T1, typename T2>
+struct common_container_type_impl<T1, T2, require_stan_scalar_t<T1>,
+                                  require_container_t<T2>> {
+  using type = math::promote_scalar_t<return_type_t<T1, T2>, plain_type_t<T2>>;
+};
+
+template <typename T1, typename T2>
+struct common_container_type_impl<T1, T2, require_container_t<T1>,
+                                  require_stan_scalar_t<T2>> {
+  using type = math::promote_scalar_t<return_type_t<T1, T2>, plain_type_t<T1>>;
+};
+}  // namespace internal
+
+template <typename... Ts>
+struct common_container_type;
+
+template <typename T>
+struct common_container_type<T> {
+  using type = typename internal::common_container_type_impl<
+      T, double>::type;  // Use double for base case
+};
+
+/**
+ * Determine the common container type for a variadic list of types.
+ * If all types are scalars, then the common scalar type is returned.
+ * If all container types the same, but not necessarily the same scalar type,
+ * the common container type with the common scalar type is returned.
+ *
+ * If different container types are present, the result is `void`.
+ */
+template <typename T1, typename... Ts>
+struct common_container_type<T1, Ts...> {
+  using type = typename internal::common_container_type_impl<
+      T1, typename common_container_type<Ts...>::type>::type;
+};
+
+template <typename... Ts>
+using common_container_t = typename common_container_type<Ts...>::type;
+
+}  // namespace stan
+
+#endif  // STAN_MATH_PRIM_META_PLAIN_TYPE_HPP

stan/math/prim/prob.hpp

@@ -287,6 +287,7 @@
 #include <stan/math/prim/prob/student_t_lccdf.hpp>
 #include <stan/math/prim/prob/student_t_lcdf.hpp>
 #include <stan/math/prim/prob/student_t_lpdf.hpp>
+#include <stan/math/prim/prob/student_t_qf.hpp>
 #include <stan/math/prim/prob/student_t_rng.hpp>
 #include <stan/math/prim/prob/uniform_ccdf_log.hpp>
 #include <stan/math/prim/prob/uniform_cdf.hpp>

stan/math/prim/prob/student_t_qf.hpp

@@ -0,0 +1,99 @@
+#ifndef STAN_MATH_PRIM_PROB_STUDENT_T_QF_HPP
+#define STAN_MATH_PRIM_PROB_STUDENT_T_QF_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/sqrt.hpp>
+#include <stan/math/prim/fun/inv_inc_beta.hpp>
+#include <stan/math/prim/fun/max_size.hpp>
+
+namespace stan {
+namespace math {
+
+/**
+ * The quantile function of the Student's t-distribution.
+ *
+ * @tparam T_p type of the probability parameter
+ * @tparam T_nu type of the degrees of freedom parameter
+ * @tparam T_mu type of the location parameter
+ * @tparam T_sigma type of the scale parameter
+ * @param p Probability in the range [0, 1].
+ * @param nu Degrees of freedom, must be non-negative.
+ * @param mu Location parameter.
+ * @param sigma Scale parameter, must be positive.
+ * @return Quantile function value.
+ * @throw std::domain_error if `nu` is negative or `sigma` is not positive,
+ * or if `p` is not in [0, 1].
+ */
+template <typename T_p, typename T_nu, typename T_mu, typename T_sigma,
+          require_all_stan_scalar_t<T_p, T_nu, T_mu, T_sigma>* = nullptr,
+          require_all_arithmetic_t<T_p, T_nu, T_mu, T_sigma>* = nullptr>
+double student_t_qf(const T_p& p, const T_nu& nu, const T_mu& mu,
+                    const T_sigma& sigma) {
+  static constexpr const char* function = "student_t_qf";
+  check_nonnegative(function, "Degrees of freedom parameter", nu);
+  check_positive(function, "Scale parameter", sigma);
+  check_bounded(function, "Probability parameter", p, 0.0, 1.0);
+
+  if (p == 0.0) {
+    return NEGATIVE_INFTY;
+  } else if (p == 1.0) {
+    return INFTY;
+  } else if (p == 0.5) {
+    return mu;
+  }
+
+  const double p_val_flip = p < 0.5 ? p : 1.0 - p;
+  const double p_sign = p < 0.5 ? -1.0 : 1.0;
+  const auto& ibeta_arg = inv_inc_beta(0.5 * nu, 0.5, 2 * p_val_flip);
+
+  return mu + p_sign * sigma * sqrt(nu) * sqrt(-1.0 + 1.0 / ibeta_arg);
+}
+
+/**
+ * A vectorized version of the Student's t quantile function that accepts
+ * std::vectors, Eigen Matrix/Array objects, or expressions, and containers of
+ * these.
+ *
+ * @tparam T_p type of the probability parameter
+ * @tparam T_nu type of the degrees of freedom parameter
+ * @tparam T_mu type of the location parameter
+ * @tparam T_sigma type of the scale parameter
+ * @tparam T_container type of the container to hold results
+ * @param p Probability in the range [0, 1].
+ * @param nu Degrees of freedom, must be non-negative.
+ * @param mu Location parameter.
+ * @param sigma Scale parameter, must be positive.
+ * @return Container with quantile function values for each input.
+ */
+template <typename T_p, typename T_nu, typename T_mu, typename T_sigma,
+          typename T_container = common_container_t<T_p, T_nu, T_mu, T_sigma>,
+          require_any_vector_t<T_p, T_nu, T_mu, T_sigma>* = nullptr,
+          require_not_t<std::is_void<T_container>>* = nullptr>
+auto student_t_qf(const T_p& p, const T_nu& nu, const T_mu& mu,
+                  const T_sigma& sigma) {
+  static constexpr const char* function = "student_t_qf";
+  const size_t max_size_all = max_size(p, nu, mu, sigma);
+  T_container result(max_size_all);
+
+  ref_type_t<T_p> p_ref = p;
+  ref_type_t<T_nu> nu_ref = nu;
+  ref_type_t<T_mu> mu_ref = mu;
+  ref_type_t<T_sigma> sigma_ref = sigma;
+
+  scalar_seq_view<ref_type_t<T_p>> p_vec(p_ref);
+  scalar_seq_view<ref_type_t<T_nu>> nu_vec(nu_ref);
+  scalar_seq_view<ref_type_t<T_mu>> mu_vec(mu_ref);
+  scalar_seq_view<ref_type_t<T_sigma>> sigma_vec(sigma_ref);
+
+  for (size_t i = 0; i < max_size_all; ++i) {
+    result[i] = student_t_qf(p_vec[i], nu_vec[i], mu_vec[i], sigma_vec[i]);
+  }
+
+  return result;
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

stan/math/rev/prob.hpp

@@ -4,5 +4,6 @@
 #include <stan/math/prim/fun/Eigen.hpp>
 
 #include <stan/math/rev/prob/std_normal_log_qf.hpp>
+#include <stan/math/rev/prob/student_t_qf.hpp>
 
 #endif