stat/distmv: Add ScoreInput function for the Normal distribution

stat/distmv/normal.go

@@ -277,6 +277,35 @@ func (n *Normal) Rand(x []float64) []float64 {
 	return x
 }
 
+// ScoreInput returns the gradient of the log-probability with respect to the
+// input x. That is, ScoreInput computes
+//  ∇_x log(p(x))
+// If score is nil, a new slice will be allocated and returned. If score is of
+// length the dimension of Normal, then the result will be put in-place into score.
+// If neither of these is true, ScoreInput will panic.
+func (n *Normal) ScoreInput(score, x []float64) []float64 {
+	// Normal log probability is
+	//  c - 0.5*(x-μ)' Σ^-1 (x-μ).
+	// So the derivative is just
+	//  -Σ^-1 (x-μ).
+	if len(x) != n.Dim() {
+		panic(badInputLength)
+	}
+	if score == nil {
+		score = make([]float64, len(x))
+	}
+	if len(score) != len(x) {
+		panic(badSizeMismatch)
+	}
+	tmp := make([]float64, len(x))
+	copy(tmp, x)
+	floats.Sub(tmp, n.mu)
+
+	n.chol.SolveVec(mat.NewVector(len(score), score), mat.NewVector(len(tmp), tmp))
+	floats.Scale(-1, score)
+	return score
+}
+
 // SetMean changes the mean of the normal distribution. SetMean panics if len(mu)
 // does not equal the dimension of the normal distribution.
 func (n *Normal) SetMean(mu []float64) {

stat/distmv/normal_test.go

@@ -9,6 +9,7 @@ import (
 	"math/rand"
 	"testing"
 
+	"gonum.org/v1/gonum/diff/fd"
 	"gonum.org/v1/gonum/floats"
 	"gonum.org/v1/gonum/mat"
 	"gonum.org/v1/gonum/stat"
@@ -536,3 +537,37 @@ func TestMarginalSingle(t *testing.T) {
 		}
 	}
 }
+
+func TestNormalScoreInput(t *testing.T) {
+	for cas, test := range []struct {
+		mu    []float64
+		sigma *mat.SymDense
+		x     []float64
+	}{
+		{
+			mu:    []float64{2, 3, 4},
+			sigma: mat.NewSymDense(3, []float64{2, 0.5, 3, 0.5, 1, 0.6, 3, 0.6, 10}),
+			x:     []float64{1, 3.1, -2},
+		},
+		{
+			mu:    []float64{2, 3, 4, 5},
+			sigma: mat.NewSymDense(4, []float64{2, 0.5, 3, 0.1, 0.5, 1, 0.6, 0.2, 3, 0.6, 10, 0.3, 0.1, 0.2, 0.3, 3}),
+			x:     []float64{1, 3.1, -2, 5},
+		},
+	} {
+		normal, ok := NewNormal(test.mu, test.sigma, nil)
+		if !ok {
+			t.Fatalf("Bad test, covariance matrix not positive definite")
+		}
+		x := make([]float64, len(test.x))
+		copy(x, test.x)
+		score := normal.ScoreInput(nil, x)
+		if !floats.Equal(x, test.x) {
+			t.Errorf("x modified during call to ScoreInput")
+		}
+		scoreFD := fd.Gradient(nil, normal.LogProb, x, nil)
+		if !floats.EqualApprox(score, scoreFD, 1e-4) {
+			t.Errorf("Case %d: derivative mismatch. Got %v, want %v", cas, score, scoreFD)
+		}
+	}
+}