go-hep.org/x/hep@v0.38.1/fit/poly.go (about)

     1  // Copyright ©2025 The go-hep Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package fit // import "go-hep.org/x/hep/fit"
     6  
     7  import (
     8  	"fmt"
     9  	"slices"
    10  
    11  	"gonum.org/v1/gonum/mat"
    12  )
    13  
    14  // Poly fits a polynomial p of degree `degree` to points (x, y):
    15  //
    16  //	p(x) = p[0] * x^deg + ... + p[deg]
    17  func Poly(xs, ys []float64, degree int) ([]float64, error) {
    18  	var (
    19  		a = vandermonde(xs, degree+1)
    20  		b = mat.NewDense(len(ys), 1, ys)
    21  		o = make([]float64, degree+1)
    22  		c = mat.NewDense(degree+1, 1, o)
    23  	)
    24  
    25  	var qr mat.QR
    26  	qr.Factorize(a)
    27  
    28  	const trans = false
    29  	err := qr.SolveTo(c, trans, b)
    30  	if err != nil {
    31  		return nil, fmt.Errorf("could not solve QR: %w", err)
    32  	}
    33  
    34  	slices.Reverse(o)
    35  	return o, nil
    36  }
    37  
    38  func vandermonde(a []float64, d int) *mat.Dense {
    39  	x := mat.NewDense(len(a), d, nil)
    40  	for i := range a {
    41  		for j, p := 0, 1.0; j < d; j, p = j+1, p*a[i] {
    42  			x.Set(i, j, p)
    43  		}
    44  	}
    45  	return x
    46  }