62 lines
1.3 KiB
Go
62 lines
1.3 KiB
Go
package imageproc
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"gonum.org/v1/gonum/mat"
|
|
)
|
|
|
|
// PolynomialFit fits a polynomial of given degree to the data points (x, y).
|
|
func PolynomialFit(x, y []float64, degree int) []float64 {
|
|
n := len(x)
|
|
// Create the Vandermonde matrix
|
|
vander := mat.NewDense(n, degree+1, nil)
|
|
for i := 0; i < n; i++ {
|
|
for j := 0; j <= degree; j++ {
|
|
vander.Set(i, j, pow(x[i], j))
|
|
}
|
|
}
|
|
|
|
// Create the right-hand side vector
|
|
yVec := mat.NewVecDense(n, y)
|
|
|
|
// Solve the least squares problem
|
|
var qr mat.QR
|
|
qr.Factorize(vander)
|
|
coeffs := mat.NewDense(degree+1, 1, nil)
|
|
err := qr.SolveTo(coeffs, false, yVec)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
return coeffs.RawMatrix().Data
|
|
}
|
|
|
|
// pow is a helper function to calculate power of a number.
|
|
func pow(x float64, n int) float64 {
|
|
if n == 0 {
|
|
return 1
|
|
}
|
|
res := x
|
|
for i := 1; i < n; i++ {
|
|
res *= x
|
|
}
|
|
return res
|
|
}
|
|
|
|
func main() {
|
|
// Input data
|
|
x := []float64{1.0, 2.0, 3.0, 4.0, 5.0}
|
|
y := []float64{1.0, 1.5, 2.5, 3.5, 5.0}
|
|
|
|
// Fit a second degree polynomial (quadratic)
|
|
degree := 2
|
|
coeffs := PolynomialFit(x, y, degree)
|
|
|
|
// Print the coefficients
|
|
fmt.Printf("Fitted polynomial coefficients: %v\n", coeffs)
|
|
if degree == 2 {
|
|
fmt.Printf("Fitted quadratic: y = %.4fx^2 + %.4fx + %.4f\n", coeffs[2], coeffs[1], coeffs[0])
|
|
}
|
|
}
|