Skip to content

Commit

Permalink
Merge pull request #8163 from steppi/complex-digamma
Browse files Browse the repository at this point in the history 
Add complex support for the digamma function
  • Loading branch information
@asi1024
asi1024 committed Feb 9, 2024
2 parents f643379 + 6006b59 commit 97ff6d3
Show file tree
Hide file tree
Showing 19 changed files with 2,198 additions and 167 deletions.
85 changes: 85 additions & 0 deletions cupy/_core/include/cupy/special/binom.h
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
/* Translated from Cython into C++ by SciPy developers in 2024.
*
* Original authors: Pauli Virtanen, Eric Moore
*/

// Binomial coefficient

#pragma once

#include "config.h"

#include "cephes/beta.h"
#include "cephes/gamma.h"

namespace special {

SPECFUN_HOST_DEVICE inline double binom(double n, double k) {
double kx, nx, num, den, dk, sgn;

if (n < 0) {
nx = std::floor(n);
if (n == nx) {
// Undefined
return std::numeric_limits<double>::quiet_NaN();
}
}

kx = std::floor(k);
if (k == kx && (std::abs(n) > 1E-8 || n == 0)) {
/* Integer case: use multiplication formula for less rounding
* error for cases where the result is an integer.
*
* This cannot be used for small nonzero n due to loss of
* precision. */
nx = std::floor(n);
if (nx == n && kx > nx / 2 && nx > 0) {
// Reduce kx by symmetry
kx = nx - kx;
}

if (kx >= 0 && kx < 20) {
num = 1.0;
den = 1.0;
for (int i = 1; i < 1 + static_cast<int>(kx); i++) {
num *= i + n - kx;
den *= i;
if (std::abs(num) > 1E50) {
num /= den;
den = 1.0;
}
}
return num / den;
}
}

// general case
if (n >= 1E10 * k and k > 0) {
// avoid under/overflows intermediate results
return std::exp(-cephes::lbeta(1 + n - k, 1 + k) - std::log(n + 1));
}
if (k > 1E8 * std::abs(n)) {
// avoid loss of precision
num = cephes::Gamma(1 + n) / std::abs(k) + cephes::Gamma(1 + n) * n / (2 * k * k); // + ...
num /= M_PI * std::pow(std::abs(k), n);
if (k > 0) {
kx = std::floor(k);
if (static_cast<int>(kx) == kx) {
dk = k - kx;
sgn = (static_cast<int>(kx) % 2 == 0) ? 1 : -1;
} else {
dk = k;
sgn = 1;
}
return num * std::sin((dk - n) * M_PI) * sgn;
}
kx = std::floor(k);
if (static_cast<int>(kx) == kx) {
return 0;
}
return num * std::sin(k * M_PI);
}
return 1 / (n + 1) / cephes::beta(1 + n - k, 1 + k);
}

} // namespace special
255 changes: 255 additions & 0 deletions cupy/_core/include/cupy/special/cephes/beta.h
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,255 @@
/* Translated into C++ by SciPy developers in 2024.
* Original header with Copyright information appears below.
*/

/* beta.c
*
* Beta function
*
*
*
* SYNOPSIS:
*
* double a, b, y, beta();
*
* y = beta( a, b );
*
*
*
* DESCRIPTION:
*
* - -
* | (a) | (b)
* beta( a, b ) = -----------.
* -
* | (a+b)
*
* For large arguments the logarithm of the function is
* evaluated using lgam(), then exponentiated.
*
*
*
* ACCURACY:
*
* Relative error:
* arithmetic domain # trials peak rms
* IEEE 0,30 30000 8.1e-14 1.1e-14
*
* ERROR MESSAGES:
*
* message condition value returned
* beta overflow log(beta) > MAXLOG 0.0
* a or b <0 integer 0.0
*
*/

/*
* Cephes Math Library Release 2.0: April, 1987
* Copyright 1984, 1987 by Stephen L. Moshier
* Direct inquiries to 30 Frost Street, Cambridge, MA 02140
*/
#pragma once

#include "../config.h"
#include "const.h"
#include "gamma.h"

namespace special {
namespace cephes {

SPECFUN_HOST_DEVICE double beta(double, double);
SPECFUN_HOST_DEVICE double lbeta(double, double);

namespace detail {
constexpr double beta_ASYMP_FACTOR = 1e6;

/*
* Asymptotic expansion for ln(|B(a, b)|) for a > ASYMP_FACTOR*max(|b|, 1).
*/
SPECFUN_HOST_DEVICE inline double lbeta_asymp(double a, double b, int *sgn) {
double r = lgam_sgn(b, sgn);
r -= b * std::log(a);

r += b * (1 - b) / (2 * a);
r += b * (1 - b) * (1 - 2 * b) / (12 * a * a);
r += -b * b * (1 - b) * (1 - b) / (12 * a * a * a);

return r;
}

/*
* Special case for a negative integer argument
*/

SPECFUN_HOST_DEVICE inline double beta_negint(int a, double b) {
int sgn;
if (b == static_cast<int>(b) && 1 - a - b > 0) {
sgn = (static_cast<int>(b) % 2 == 0) ? 1 : -1;
return sgn * special::cephes::beta(1 - a - b, b);
} else {
set_error("lbeta", SF_ERROR_OVERFLOW, NULL);
return std::numeric_limits<double>::infinity();
}
}

SPECFUN_HOST_DEVICE inline double lbeta_negint(int a, double b) {
double r;
if (b == static_cast<int>(b) && 1 - a - b > 0) {
r = special::cephes::lbeta(1 - a - b, b);
return r;
} else {
set_error("lbeta", SF_ERROR_OVERFLOW, NULL);
return std::numeric_limits<double>::infinity();
}
}
} // namespace detail

SPECFUN_HOST_DEVICE inline double beta(double a, double b) {
double y;
int sign = 1;

if (a <= 0.0) {
if (a == std::floor(a)) {
if (a == static_cast<int>(a)) {
return detail::beta_negint(static_cast<int>(a), b);
} else {
goto overflow;
}
}
}

if (b <= 0.0) {
if (b == std::floor(b)) {
if (b == static_cast<int>(b)) {
return detail::beta_negint(static_cast<int>(b), a);
} else {
goto overflow;
}
}
}

if (std::abs(a) < std::abs(b)) {
y = a;
a = b;
b = y;
}

if (std::abs(a) > detail::beta_ASYMP_FACTOR * std::abs(b) && a > detail::beta_ASYMP_FACTOR) {
/* Avoid loss of precision in lgam(a + b) - lgam(a) */
y = detail::lbeta_asymp(a, b, &sign);
return sign * std::exp(y);
}

y = a + b;
if (std::abs(y) > detail::MAXGAM || std::abs(a) > detail::MAXGAM || std::abs(b) > detail::MAXGAM) {
int sgngam;
y = detail::lgam_sgn(y, &sgngam);
sign *= sgngam; /* keep track of the sign */
y = detail::lgam_sgn(b, &sgngam) - y;
sign *= sgngam;
y = detail::lgam_sgn(a, &sgngam) + y;
sign *= sgngam;
if (y > detail::MAXLOG) {
goto overflow;
}
return (sign * std::exp(y));
}

y = Gamma(y);
a = Gamma(a);
b = Gamma(b);
if (y == 0.0)
goto overflow;

if (std::abs(std::abs(a) - std::abs(y)) > std::abs(std::abs(b) - std::abs(y))) {
y = b / y;
y *= a;
} else {
y = a / y;
y *= b;
}

return (y);

overflow:
set_error("beta", SF_ERROR_OVERFLOW, NULL);
return (sign * std::numeric_limits<double>::infinity());
}

/* Natural log of |beta|. */

SPECFUN_HOST_DEVICE inline double lbeta(double a, double b) {
double y;
int sign;

sign = 1;

if (a <= 0.0) {
if (a == std::floor(a)) {
if (a == static_cast<int>(a)) {
return detail::lbeta_negint(static_cast<int>(a), b);
} else {
goto over;
}
}
}

if (b <= 0.0) {
if (b == std::floor(b)) {
if (b == static_cast<int>(b)) {
return detail::lbeta_negint(static_cast<int>(b), a);
} else {
goto over;
}
}
}

if (std::abs(a) < std::abs(b)) {
y = a;
a = b;
b = y;
}

if (std::abs(a) > detail::beta_ASYMP_FACTOR * std::abs(b) && a > detail::beta_ASYMP_FACTOR) {
/* Avoid loss of precision in lgam(a + b) - lgam(a) */
y = detail::lbeta_asymp(a, b, &sign);
return y;
}

y = a + b;
if (std::abs(y) > detail::MAXGAM || std::abs(a) > detail::MAXGAM || std::abs(b) > detail::MAXGAM) {
int sgngam;
y = detail::lgam_sgn(y, &sgngam);
sign *= sgngam; /* keep track of the sign */
y = detail::lgam_sgn(b, &sgngam) - y;
sign *= sgngam;
y = detail::lgam_sgn(a, &sgngam) + y;
sign *= sgngam;
return (y);
}

y = Gamma(y);
a = Gamma(a);
b = Gamma(b);
if (y == 0.0) {
over:
set_error("lbeta", SF_ERROR_OVERFLOW, NULL);
return (sign * std::numeric_limits<double>::infinity());
}

if (std::abs(std::abs(a) - std::abs(y)) > std::abs(std::abs(b) - std::abs(y))) {
y = b / y;
y *= a;
} else {
y = a / y;
y *= b;
}

if (y < 0) {
y = -y;
}

return (std::log(y));
}
} // namespace cephes
} // namespace special

0 comments on commit 97ff6d3

Please sign in to comment.