diff --git a/stan/math/prim/fun/quantile.hpp b/stan/math/prim/fun/quantile.hpp index a427cb6e9e9..f343e495419 100644 --- a/stan/math/prim/fun/quantile.hpp +++ b/stan/math/prim/fun/quantile.hpp @@ -6,7 +6,6 @@ #include #include #include -#include namespace stan { namespace math { @@ -41,21 +40,25 @@ inline double quantile(const T& samples_vec, const double p) { Eigen::VectorXd x = as_array_or_scalar(samples_vec); check_not_nan("quantile", "samples_vec", x); - if (n_sample == 1) + if (n_sample == 1) { return x.coeff(0); - else if (p == 0.) - return *std::min_element(x.data(), x.data() + n_sample); - else if (p == 1.) - return *std::max_element(x.data(), x.data() + n_sample); + } else if (p == 0.) { + return x.minCoeff(); + } else if (p == 1.) { + return x.maxCoeff(); + } - double index = (n_sample - 1) * p; - size_t lo = std::floor(index); - size_t hi = std::ceil(index); + const size_t nm1 = (n_sample - 1); + const double index = nm1 * p; + const size_t lo = static_cast(index); - std::sort(x.data(), x.data() + n_sample, std::less()); + std::nth_element(x.data(), x.data() + lo, x.data() + n_sample); - double h = index - lo; - return (1 - h) * x.coeff(lo) + h * x.coeff(hi); + const double h = index - lo; + if (h == 0) { + return x.coeff(lo); + } + return (1 - h) * x.coeff(lo) + h * x.tail(nm1 - lo).minCoeff(); } /** @@ -75,10 +78,12 @@ inline double quantile(const T& samples_vec, const double p) { * @throw std::invalid_argument If any of the values are NaN or size 0. * @throw std::domain_error If `p<0` or `p>1` for any p in ps. */ -template * = nullptr, +template , + require_all_vector_t* = nullptr, require_vector_vt* = nullptr, - require_std_vector_vt* = nullptr> -inline std::vector quantile(const T& samples_vec, const Tp& ps) { + require_vector_vt* = nullptr> +inline ReturnT quantile(const T& samples_vec, const Tp& ps) { check_not_nan("quantile", "ps", ps); check_bounded("quantile", "ps", ps, 0, 1); @@ -88,29 +93,24 @@ inline std::vector quantile(const T& samples_vec, const Tp& ps) { return {}; } - Eigen::VectorXd x = as_array_or_scalar(samples_vec); + plain_type_t x = samples_vec; check_not_nan("quantile", "samples_vec", x); - const auto& p = as_array_or_scalar(ps); - std::vector ret(n_ps, 0.0); + std::sort(x.begin(), x.end()); + ReturnT ret(n_ps); - std::sort(x.data(), x.data() + n_sample, std::less()); - Eigen::ArrayXd index = (n_sample - 1) * p; + const size_t nm1 = (n_sample - 1); - for (size_t i = 0; i < n_ps; ++i) { - if (p[i] == 0.) { - ret[i] = x.coeff(0); - } else if (p[i] == 1.) { - ret[i] = x.coeff(n_sample - 1); - } else { - size_t lo = std::floor(index[i]); - size_t hi = std::ceil(index[i]); - - double h = index[i] - lo; - - ret[i] = (1 - h) * x.coeff(lo) + h * x.coeff(hi); + for (size_t i = 0; i < n_ps; i++) { + const double sample_xi = nm1 * ps[i]; + const size_t smpl_xi_int = static_cast(sample_xi); + const double smpl_xi_frc = sample_xi - smpl_xi_int; + ret[i] = (1 - smpl_xi_frc) * x[smpl_xi_int]; + if (smpl_xi_frc != 0.0) { + ret[i] += smpl_xi_frc * x[smpl_xi_int + 1]; } } + return ret; } diff --git a/test/unit/math/prim/fun/quantile_test.cpp b/test/unit/math/prim/fun/quantile_test.cpp index a939a1098ac..a2cb4dbc74c 100644 --- a/test/unit/math/prim/fun/quantile_test.cpp +++ b/test/unit/math/prim/fun/quantile_test.cpp @@ -105,7 +105,7 @@ inline void test_quantile_double() { p[2] = 0.2; p[3] = 1; - std::vector ret = quantile(v, p); + Tp ret = quantile(v, p); EXPECT_FLOAT_EQ(ret[0], -0.28); EXPECT_FLOAT_EQ(ret[1], -0.196); EXPECT_FLOAT_EQ(ret[2], -0.112); @@ -145,7 +145,7 @@ inline void test_quantile_double() { // check size 1 first argument works T v1(1); v1[0] = -0.07; - std::vector ret1 = quantile(v1, p); + Tp ret1 = quantile(v1, p); EXPECT_FLOAT_EQ(ret1[0], -0.07); EXPECT_FLOAT_EQ(ret1[1], -0.07); EXPECT_FLOAT_EQ(ret1[2], -0.07); @@ -160,6 +160,14 @@ TEST(MathFunctions, quantileEigenVectorXdStdVecDouble) { test_quantile_double>(); } +TEST(MathFunctions, quantileEigenVectorXdEigenVectorXd) { + test_quantile_double(); +} + TEST(MathFunctions, quantileEigenRowVectorXdStdVecDouble) { test_quantile_double>(); } + +TEST(MathFunctions, quantileEigenRowVectorXdEigenRowVectorXd) { + test_quantile_double(); +}