/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package etomica.math.numerical;

import Jama.Matrix;

/**
 * Class that performs a polynomial fit to x,y data, optionally taking weights
 * associated with each data point.
 * 
 * @author Andrew Schultz
 */
public class PolynomialFit {

    /**
     * Perform polynomial fit of the given order.  order=2 would fit
     *     y = a*x^2 + b*x + c
     * The polynomial coefficients are returned as a double array in order of
     * increasing power of x.  So,
     *   double[] r = doFit();
     * r[i] is the coefficient for x^i
     */
    public static double[] doFit(int order, double[] x, double[] y) {
        double[] w = new double[x.length];
        for (int i=0; i<x.length; i++) {
            w[i] = 1;
        }
        return doFit(order, x, y, w);
    }

    /**
     * Perform polynomial fit of the given order with the given weights (w).
     * order=2 would fit
     *     y = a*x^2 + b*x + c
     * The polynomial coefficients are returned as a double array in order of
     * increasing power of x.  So,
     *   double[] r = doFit();
     * r[i] is the coefficient for x^i
     */
    public static double[] doFit(int order, double[] x, double[] y, double[] w) {
        if (x.length != y.length || x.length != w.length || x.length < order+1) {
            // We need at least order+1 data points to do a meaningful fit.
            return null;
        }

        double[][] M = new double[order+1][order+1];
        double[] b = new double[order+1];
        for (int i=0; i<x.length; i++) {
            double xp = w[i];
            for (int ipower = 0; ipower<2*order+1; ipower++) {
                for (int irow=order; irow>-1; irow--) {
                    int col = 2*order - irow - ipower;
                    if (col > -1 && col < order+1) {
                        M[irow][col] += xp;
                    }
                }
                if (order - ipower > -1) {
                    b[order-ipower] += xp*y[i];
                }
                xp *= x[i];
            }
        }
        Matrix mat = new Matrix(M);
        Matrix sol = mat.solve(new Matrix(b,order+1));
        double[] result = new double[order+1];
        for (int i=0; i<result.length; i++) {
            result[i] = sol.get(order-i,0);
        }
        return result;
    }
    
    public static void main(String[] args) {
        double[] c = doFit(2, new double[]{1,2,3}, new double[]{9,6,1});
        System.out.println("a="+c[2]+" b="+c[1]+" c="+c[0]);
    }
}

@@ -2474,24 +2474,734 @@
 esult;%0A    %7D
+%0A%0A    /**%0A     * Returns average relative deviation (relative to the uncertainty) of the%0A     * given fit from the given data.%0A     * %3Cp%3E%0A     * chi = (%3C(diff/error)%5E2%3E/N)%5E.5%0A     */%0A    public static double getChi(int order, double%5B%5D x, double%5B%5D y, double%5B%5D w, double%5B%5D poly) %7B%0A        double chiSqSum = 0, nData = 0;%0A        for (int i = 0; i %3C x.length; i++) %7B%0A            double xp = 1;%0A            double yp = 0;%0A            for (int ipower = 0; ipower %3C= order; ipower++) %7B%0A                yp += poly%5Bipower%5D * xp;%0A            %7D%0A            double diff = y%5Bi%5D - yp;%0A            chiSqSum += diff * diff * w%5Bi%5D;%0A            if (w%5Bi%5D %3E 0) nData++;%0A        %7D%0A        return Math.sqrt(chiSqSum / nData);%0A    %7D
 %0A    %0A    pu

add method to compute chi = (<(diff/error)^2>/N)^.5