/
fit.go
79 lines (63 loc) · 1.68 KB
/
fit.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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
// Copyright 2017 The go-hep 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 fit provides functions to fit data.
package fit // import "go-hep.org/x/hep/fit"
import (
"github.com/gonum/diff/fd"
)
// Func1D describes a 1D function to fit some data.
type Func1D struct {
// F is the function to minimize.
// ps is the slice of parameters to optimize during the fit.
F func(x float64, ps []float64) float64
// N is the number of parameters to optimize during the fit.
// If N is 0, Ps must not be nil.
N int
// Ps is the initial values for the parameters.
// If Ps is nil, the set of initial parameters values is a slice of
// length N filled with zeros.
Ps []float64
X []float64
Y []float64
Err []float64
sig2 []float64 // inverse of squares of measurement errors along Y.
fct func(ps []float64) float64 // cost function (objective function)
grad func(grad, ps []float64)
}
func (f *Func1D) init() {
f.sig2 = make([]float64, len(f.Y))
switch {
default:
for i := range f.Y {
f.sig2[i] = 1
}
case f.Err != nil:
for i, v := range f.Err {
f.sig2[i] = 1 / (v * v)
}
}
if f.Ps == nil {
f.Ps = make([]float64, f.N)
}
if len(f.Ps) == 0 {
panic("fit: invalid number of initial parameters")
}
if len(f.X) != len(f.Y) {
panic("fit: mismatch length")
}
if len(f.sig2) != len(f.Y) {
panic("fit: mismatch length")
}
f.fct = func(ps []float64) float64 {
var chi2 float64
for i := range f.X {
res := f.F(f.X[i], ps) - f.Y[i]
chi2 += res * res * f.sig2[i]
}
return 0.5 * chi2
}
f.grad = func(grad, ps []float64) {
fd.Gradient(grad, f.fct, ps, nil)
}
}