Skip to content

Instantly share code, notes, and snippets.

@tansey
Created July 2, 2017 16:18
Show Gist options
  • Save tansey/c55c1fad155d6ececba24af026358460 to your computer and use it in GitHub Desktop.
Save tansey/c55c1fad155d6ececba24af026358460 to your computer and use it in GitHub Desktop.
Fit negative binomial with weighted observations in 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))
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment