r2rs-nmath 0.1.1

// Translation of nmath's dnbeta
//
// Copyright (C) 2020 neek-sss
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use nonstdfloat::f128;
use num_traits::{Float, ToPrimitive, Zero};
use strafe_type::{FloatConstraint, Positive64, Rational64, Real64};

use crate::{
    distribution::{beta::dbeta, pois::dpois_raw},
    traits::DPQ,
};

/// Computes the density of the noncentral beta distribution with
/// noncentrality parameter ncp.  The noncentral beta distribution
/// has density:
///
/// $ f(x|a,b,ncp) = \sum^{\infty}_{i=0} \frac{p(i) x^{a+i-1} (1-x)^{b-1}}{B(a+i,b)} $
///
/// where:
///
/// $ p(k) = e^{-ncp/2} \frac{(ncp/2)^k}{k!} $
///
/// $ B(a,b) = \frac{Gamma(a) Gamma(b)}{Gamma(a+b)} $
///
///
/// This can be computed efficiently by using the recursions:
///
/// $ p(k+1) = \frac{ncp/2}{(k+1) p(k)} $
///
/// $ B(a+k+1,b)   = \frac{a+k}{a+b+k} B(a+k,b) $
///
/// The new algorithm first determines for which k the k-th term is maximal,
/// and then sums outwards to both sides from the 'mid'.
pub fn dnbeta<
    RE: Into<Real64>,
    RA1: Into<Rational64>,
    RA2: Into<Rational64>,
    P: Into<Positive64>,
>(
    x: RE,
    a: RA1,
    b: RA2,
    ncp: P,
    log: bool,
) -> Real64 {
    let x = x.into().unwrap();
    let a = a.into().unwrap();
    let b = b.into().unwrap();
    let ncp = ncp.into().unwrap();

    static eps: f64 = 1.0e-15;

    let mut kMax = 0.0;
    let mut k = 0.0;
    let mut ncp2 = 0.0;
    let mut dx2 = 0.0;
    let mut d = 0.0;
    let mut D = 0.0;
    let mut sum = f128::zero();
    let mut term = f128::zero();
    let mut p_k = f128::zero();
    let mut q = f128::zero();

    if !a.is_finite() || !b.is_finite() || !ncp.is_finite() {
        return f64::nan().into();
    }

    if !(0.0..=1.0).contains(&x) {
        return f64::d_0(log).into();
    }
    if ncp == 0.0 {
        return dbeta(x, a, b, log);
    }

    /* New algorithm, starting with *largest* term : */
    ncp2 = 0.5 * ncp;
    dx2 = ncp2 * x;
    d = (dx2 - a - 1.0) / 2.0;
    D = d * d + dx2 * (a + b) - a;
    if D <= 0.0 {
        kMax = 0.0
    } else {
        D = (d + D.sqrt()).ceil();
        kMax = if D > 0.0 { D.trunc() } else { 0.0 }
    }

    /* The starting "middle term" --- first look at it's log scale: */
    term = f128::new(dbeta(x, a + kMax, b, true));
    p_k = f128::new(dpois_raw(kMax, ncp2, true));
    if x == 0.0 || !term.is_finite() || !p_k.to_f64().unwrap().is_finite() {
        /* if term = +Inf */
        return (p_k + term).to_f64().unwrap().d_exp(log).into();
    }

    /* Now if s_k := p_k * t_k  {here = (p_k + term).exp()} would underflow,
     * we should rather scale everything and re-scale at the end:*/

    p_k += term; /* = log(p_k) + log(t_k) == log(s_k) -- used at end to rescale */
    /* mid = 1 = the rescaled value, instead of  mid = p_k.exp(); */

    /* Now sum from the inside out */
    term = f128::new(1.0); /* = mid term */
    sum = term;
    /* middle to the left */
    k = kMax;
    while k > 0.0 && term > sum * f128::new(eps) {
        k -= 1.0;
        q = f128::new((k + 1.0) * (k + a) / (k + a + b) / dx2);
        term *= q;
        sum += term
    }
    /* middle to the right */
    term = f128::new(1.0);
    k = kMax;
    loop {
        q = f128::new(dx2 * (k + a + b) / (k + a) / (k + 1.0));
        k += 1.0;
        term *= q;
        sum += term;
        if !(term > sum * f128::new(eps)) {
            break;
        }
    }

    (p_k + sum.ln()).to_f64().unwrap().d_exp(log).into()
}