diff --git a/libcudacxx/include/cuda/std/__random/negative_binomial_distribution.h b/libcudacxx/include/cuda/std/__random/negative_binomial_distribution.h new file mode 100644 index 00000000000..2beed244168 --- /dev/null +++ b/libcudacxx/include/cuda/std/__random/negative_binomial_distribution.h @@ -0,0 +1,212 @@ +//===----------------------------------------------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. +// +//===----------------------------------------------------------------------====// + +#ifndef _CUDA_STD___NEGATIVE_BINOMIAL_DISTRIBUTION_H +#define _CUDA_STD___NEGATIVE_BINOMIAL_DISTRIBUTION_H + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#if !_CCCL_COMPILER(NVRTC) +# include +#endif // !_CCCL_COMPILER(NVRTC) + +#include + +_CCCL_BEGIN_NAMESPACE_CUDA_STD + +template +class negative_binomial_distribution +{ + static_assert(__libcpp_random_is_valid_inttype<_IntType>, "IntType must be a supported integer type"); + +public: + // types + using result_type = _IntType; + + class param_type + { + result_type __k_ = result_type{1}; + double __p_ = 0.5; + + public: + using distribution_type = negative_binomial_distribution; + + _CCCL_API constexpr explicit param_type(result_type __k = 1, double __p = 0.5) noexcept + : __k_{__k} + , __p_{__p} + {} + + [[nodiscard]] _CCCL_API constexpr result_type k() const noexcept + { + return __k_; + } + [[nodiscard]] _CCCL_API constexpr double p() const noexcept + { + return __p_; + } + + [[nodiscard]] _CCCL_API friend constexpr bool operator==(const param_type& __x, const param_type& __y) noexcept + { + return __x.__k_ == __y.__k_ && __x.__p_ == __y.__p_; + } +#if _CCCL_STD_VER <= 2017 + [[nodiscard]] _CCCL_API friend constexpr bool operator!=(const param_type& __x, const param_type& __y) noexcept + { + return !(__x == __y); + } +#endif // _CCCL_STD_VER <= 2017 + }; + +private: + param_type __p_{}; + +public: + // constructor and reset functions + constexpr negative_binomial_distribution() noexcept = default; + + _CCCL_API constexpr explicit negative_binomial_distribution(result_type __k, double __p = 0.5) noexcept + : __p_{__k, __p} + {} + _CCCL_API constexpr explicit negative_binomial_distribution(const param_type& __p) noexcept + : __p_{__p} + {} + _CCCL_API constexpr void reset() noexcept {} + + // generating functions + template + [[nodiscard]] _CCCL_API result_type operator()(_URng& __g) + { + return (*this)(__g, __p_); + } + template + [[nodiscard]] _CCCL_API result_type operator()(_URng& __urng, const param_type& __pr) + { + static_assert(__cccl_random_is_valid_urng<_URng>, "URng must meet the UniformRandomBitGenerator requirements"); + const result_type __k = __pr.k(); + const double __p = __pr.p(); + // When the number of bits in _IntType is small, we are too likely to + // overflow __f below to use this technique. + if (__k <= 21 * __p && sizeof(_IntType) > 1) + { + bernoulli_distribution __gen(__p); + result_type __f = 0; + result_type __s = 0; + while (__s < __k) + { + if (__gen(__urng)) + { + ++__s; + } + else + { + ++__f; + } + } + return __f; + } + return poisson_distribution(gamma_distribution(__k, (1 - __p) / __p)(__urng))(__urng); + } + + // property functions + [[nodiscard]] _CCCL_API constexpr result_type k() const noexcept + { + return __p_.k(); + } + [[nodiscard]] _CCCL_API constexpr double p() const noexcept + { + return __p_.p(); + } + + [[nodiscard]] _CCCL_API constexpr param_type param() const noexcept + { + return __p_; + } + _CCCL_API constexpr void param(const param_type& __p) noexcept + { + __p_ = __p; + } + + [[nodiscard]] _CCCL_API static constexpr result_type min() noexcept + { + return 0; + } + [[nodiscard]] _CCCL_API static constexpr result_type max() noexcept + { + return numeric_limits::max(); + } + + [[nodiscard]] _CCCL_API friend constexpr bool + operator==(const negative_binomial_distribution& __x, const negative_binomial_distribution& __y) noexcept + { + return __x.__p_ == __y.__p_; + } +#if _CCCL_STD_VER <= 2017 + [[nodiscard]] _CCCL_API friend constexpr bool + operator!=(const negative_binomial_distribution& __x, const negative_binomial_distribution& __y) noexcept + { + return !(__x == __y); + } +#endif // _CCCL_STD_VER <= 2017 + +#if !_CCCL_COMPILER(NVRTC) + template + friend ::std::basic_ostream<_CharT, _Traits>& + operator<<(::std::basic_ostream<_CharT, _Traits>& __os, const negative_binomial_distribution& __x) + { + using _Ostream = ::std::basic_ostream<_CharT, _Traits>; + auto __flags = __os.flags(); + __os.flags(_Ostream::dec | _Ostream::left | _Ostream::scientific); + _CharT __sp = __os.widen(' '); + _CharT __fill = __os.fill(__sp); + auto __precision = __os.precision(numeric_limits::max_digits10); + __os << __x.k() << __sp << __x.p(); + __os.precision(__precision); + __os.fill(__fill); + __os.flags(__flags); + return __os; + } + + template + friend ::std::basic_istream<_CharT, _Traits>& + operator>>(::std::basic_istream<_CharT, _Traits>& __is, negative_binomial_distribution& __x) + { + using _Istream = ::std::basic_istream<_CharT, _Traits>; + auto __flags = __is.flags(); + __is.flags(_Istream::skipws); + result_type __k; + double __p; + __is >> __k >> __p; + if (!__is.fail()) + { + __x.param(param_type(__k, __p)); + } + __is.flags(__flags); + return __is; + } +#endif // !_CCCL_COMPILER(NVRTC) +}; + +_CCCL_END_NAMESPACE_CUDA_STD + +#include + +#endif // _CUDA_STD___NEGATIVE_BINOMIAL_DISTRIBUTION_H diff --git a/libcudacxx/include/cuda/std/__random_ b/libcudacxx/include/cuda/std/__random_ index 78507fce790..cf0d111fe4a 100644 --- a/libcudacxx/include/cuda/std/__random_ +++ b/libcudacxx/include/cuda/std/__random_ @@ -31,6 +31,7 @@ #include #include #include +#include #include #include #include diff --git a/libcudacxx/test/libcudacxx/std/random/distribution/negative_binomial.pass.cpp b/libcudacxx/test/libcudacxx/std/random/distribution/negative_binomial.pass.cpp new file mode 100644 index 00000000000..65b81137710 --- /dev/null +++ b/libcudacxx/test/libcudacxx/std/random/distribution/negative_binomial.pass.cpp @@ -0,0 +1,60 @@ +//===----------------------------------------------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. +// +//===----------------------------------------------------------------------===// +// +// REQUIRES: long_tests + +// + +// template +// class negative_binomial_distribution + +#include +#include + +#include "random_utilities/stats_functions.h" +#include "random_utilities/test_distribution.h" +#include "test_macros.h" + +template +struct negative_binomial_cdf +{ + using P = typename cuda::std::negative_binomial_distribution::param_type; + + __host__ __device__ double operator()(double x, const P& p) const + { + // CDF: F(x; k, p) = I_p(k, x+1) where I_p is the regularized incomplete beta function + // This represents P(X <= x) where X is the number of failures before k successes + if (x < 0) + { + return 0.0; + } + double k = static_cast(p.k()); + double prob = p.p(); + return incomplete_beta(k, cuda::std::floor(x) + 1.0, prob); + } +}; + +template +__host__ __device__ void test() +{ + // Cannot be constexpr due to gamma_distribution and log/exp + [[maybe_unused]] const bool test_constexpr = false; + using D = cuda::std::negative_binomial_distribution; + using P = typename D::param_type; + using G = cuda::std::philox4x64; + cuda::std::array params = {P(5, 0.5), P(10, 0.3), P(3, 0.7), P(15, 0.4), P(20, 0.6)}; + test_distribution(params, negative_binomial_cdf{}); +} + +int main(int, char**) +{ + test(); + test(); + return 0; +}