blob: 60869c73dd0280654b6cc69fcd07947c413e08d8 [file] [log] [blame]
 // Copyright ©2016 The Gonum Authors. All rights reserved. // Use of this source code is governed by a BSD-style // license that can be found in the LICENSE file. package distmv import ( "math" "golang.org/x/exp/rand" "gonum.org/v1/gonum/floats" "gonum.org/v1/gonum/mat" "gonum.org/v1/gonum/stat/distuv" ) // Dirichlet implements the Dirichlet probability distribution. // // The Dirichlet distribution is a continuous probability distribution that // generates elements over the probability simplex, i.e. ||x||_1 = 1. The Dirichlet // distribution is the conjugate prior to the categorical distribution and the // multivariate version of the beta distribution. The probability of a point x is // 1/Beta(α) \prod_i x_i^(α_i - 1) // where Beta(α) is the multivariate Beta function (see the mathext package). // // For more information see https://en.wikipedia.org/wiki/Dirichlet_distribution type Dirichlet struct { alpha []float64 dim int src rand.Source lbeta float64 sumAlpha float64 } // NewDirichlet creates a new dirichlet distribution with the given parameters alpha. // NewDirichlet will panic if len(alpha) == 0, or if any alpha is <= 0. func NewDirichlet(alpha []float64, src rand.Source) *Dirichlet { dim := len(alpha) if dim == 0 { panic(badZeroDimension) } for _, v := range alpha { if v <= 0 { panic("dirichlet: non-positive alpha") } } a := make([]float64, len(alpha)) copy(a, alpha) d := &Dirichlet{ alpha: a, dim: dim, src: src, } d.lbeta, d.sumAlpha = d.genLBeta(a) return d } // CovarianceMatrix returns the covariance matrix of the distribution. Upon // return, the value at element {i, j} of the covariance matrix is equal to // the covariance of the i^th and j^th variables. // covariance(i, j) = E[(x_i - E[x_i])(x_j - E[x_j])] // If the input matrix is nil a new matrix is allocated, otherwise the result // is stored in-place into the input. func (d *Dirichlet) CovarianceMatrix(cov *mat.SymDense) *mat.SymDense { if cov == nil { cov = mat.NewSymDense(d.Dim(), nil) } else if cov.Symmetric() == 0 { *cov = *(cov.GrowSquare(d.dim).(*mat.SymDense)) } else if cov.Symmetric() != d.dim { panic("normal: input matrix size mismatch") } scale := 1 / (d.sumAlpha * d.sumAlpha * (d.sumAlpha + 1)) for i := 0; i < d.dim; i++ { ai := d.alpha[i] v := ai * (d.sumAlpha - ai) * scale cov.SetSym(i, i, v) for j := i + 1; j < d.dim; j++ { aj := d.alpha[j] v := -ai * aj * scale cov.SetSym(i, j, v) } } return cov } // genLBeta computes the generalized LBeta function. func (d *Dirichlet) genLBeta(alpha []float64) (lbeta, sumAlpha float64) { for _, alpha := range d.alpha { lg, _ := math.Lgamma(alpha) lbeta += lg sumAlpha += alpha } lg, _ := math.Lgamma(sumAlpha) return lbeta - lg, sumAlpha } // Dim returns the dimension of the distribution. func (d *Dirichlet) Dim() int { return d.dim } // LogProb computes the log of the pdf of the point x. // // It does not check that ||x||_1 = 1. func (d *Dirichlet) LogProb(x []float64) float64 { dim := d.dim if len(x) != dim { panic(badSizeMismatch) } var lprob float64 for i, x := range x { lprob += (d.alpha[i] - 1) * math.Log(x) } lprob -= d.lbeta return lprob } // Mean returns the mean of the probability distribution at x. If the // input argument is nil, a new slice will be allocated, otherwise the result // will be put in-place into the receiver. func (d *Dirichlet) Mean(x []float64) []float64 { x = reuseAs(x, d.dim) copy(x, d.alpha) floats.Scale(1/d.sumAlpha, x) return x } // Prob computes the value of the probability density function at x. func (d *Dirichlet) Prob(x []float64) float64 { return math.Exp(d.LogProb(x)) } // Rand generates a random number according to the distributon. // If the input slice is nil, new memory is allocated, otherwise the result is stored // in place. func (d *Dirichlet) Rand(x []float64) []float64 { x = reuseAs(x, d.dim) for i := range x { x[i] = distuv.Gamma{Alpha: d.alpha[i], Beta: 1, Src: d.src}.Rand() } sum := floats.Sum(x) floats.Scale(1/sum, x) return x }