fitting

fit.go

@@ -0,0 +1,61 @@
+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])
+	}
+}