Fit negative binomial with weighted observations in R

fitWeightedNegativeBinomial.R

# Fit using a simple EM algorithm
# observations are x
# weights are w (must be same length as x)
# returns (r, p)
# r - dispersion parameter
# p - probability of success
weightedNegBinomFit <- function(x, w, maxsteps=30)
{
    sum.wx = sum(x*w)
    sum.w = sum(w)
    r = 1
    search.vals = exp(seq(log(1e-12),log(1),length.out=20))
    for(step in 1:maxsteps){
        # E-step fits p
        p = sum.wx / (sum.w * r + sum.wx)
        # M-step fits r
        r.delta = 1
        while(r.delta > 1e-6){
            gr = -sum(w * digamma(x + r)) + sum.w * digamma(r) - sum.w * log(1-p)
            hess = -sum(w * trigamma(x + r)) + sum.w * trigamma(r)
            r.prev = r
            rvals = r - search.vals * gr / hess
            evals = rep(Inf, 20)
            for(i in 1:20){
                if (rvals[i] <= 0){ break }
                evals[i] = -sum(w * log(gamma(x + rvals[i]))) + sum.w * log(gamma(rvals[i])) - sum.wx*log(p) - sum.w*rvals[i]*log(1-p)
            }
            r = rvals[which.min(evals)]
            r.delta = abs(r - r.prev)
        }
    }
    return(c(r,p))
}