/
negative_binomial.go
60 lines (54 loc) · 1.63 KB
/
negative_binomial.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
package inference
import (
"math"
"github.com/umbralcalc/stochadex/pkg/simulator"
"golang.org/x/exp/rand"
"gonum.org/v1/gonum/mat"
"gonum.org/v1/gonum/stat/distuv"
)
// NegativeBinomialLikelihoodDistribution assumes the real data are well
// described by a negative binomial distribution, given the input mean
// and covariance matrix.
type NegativeBinomialLikelihoodDistribution struct {
Src rand.Source
}
func (n *NegativeBinomialLikelihoodDistribution) Configure(
partitionIndex int,
settings *simulator.Settings,
) {
n.Src = rand.NewSource(settings.Seeds[partitionIndex])
}
func (n *NegativeBinomialLikelihoodDistribution) EvaluateLogLike(
mean *mat.VecDense,
covariance mat.Symmetric,
data []float64,
) float64 {
logLike := 0.0
for i := 0; i < mean.Len(); i++ {
r := mean.AtVec(i) * mean.AtVec(i) /
(covariance.At(i, i) - mean.AtVec(i))
p := mean.AtVec(i) / covariance.At(i, i)
lg1, _ := math.Lgamma(r + data[i])
lg2, _ := math.Lgamma(data[i] + 1.0)
lg3, _ := math.Lgamma(data[i])
logLike += lg1 + lg2 + lg3 + (r * math.Log(p)) +
(data[i] * math.Log(1.0-p))
}
return logLike
}
func (n *NegativeBinomialLikelihoodDistribution) GenerateNewSamples(
mean *mat.VecDense,
covariance mat.Symmetric,
) []float64 {
samples := make([]float64, 0)
distPoisson := &distuv.Poisson{Lambda: 1.0, Src: n.Src}
distGamma := &distuv.Gamma{Alpha: 1.0, Beta: 1.0, Src: n.Src}
for i := 0; i < mean.Len(); i++ {
distGamma.Beta = 1.0 /
((covariance.At(i, i) / mean.AtVec(i)) - 1.0)
distGamma.Alpha = mean.AtVec(i) * distGamma.Beta
distPoisson.Lambda = distGamma.Rand()
samples = append(samples, distPoisson.Rand())
}
return samples
}