/* 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;
    }

    /**
     * Returns average relative deviation (relative to the uncertainty) of the
     * given fit from the given data.
     * <p>
     * chi = (<(diff/error)^2>/N)^.5
     */
    public static double getChi(int order, double[] x, double[] y, double[] w, double[] poly) {
        double chiSqSum = 0, nData = 0;
        for (int i = 0; i < x.length; i++) {
            double xp = 1;
            double yp = 0;
            for (int ipower = 0; ipower <= order; ipower++) {
                yp += poly[ipower] * xp;
            }
            double diff = y[i] - yp;
            chiSqSum += diff * diff * w[i];
            if (w[i] > 0) nData++;
        }
        return Math.sqrt(chiSqSum / nData);
    }
    
    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]);
    }
}

@@ -1702,32 +1702,69 @@
 .length; i++) %7B%0A
+            if (w%5Bi%5D == 0) continue;%0A
             doub
@@ -1773,24 +1773,24 @@
  xp = w%5Bi%5D;%0A
-
             
@@ -1926,21 +1926,20 @@
 t col = 
-2*ord
+ipow
 er - iro
@@ -1939,25 +1939,16 @@
 r - irow
- - ipower
 ;%0A      
@@ -2102,27 +2102,23 @@
 if (
-order - ipower %3E -1
+ipower %3C= order
 ) %7B%0A
@@ -2139,22 +2139,16 @@
       b%5B
-order-
 ipower%5D 
@@ -2152,17 +2152,19 @@
 r%5D += xp
-*
+ * 
 y%5Bi%5D;%0A  
@@ -2451,16 +2451,11 @@
 get(
-order-
 i,
+ 
 0);%0A
@@ -2707,27 +2707,16 @@
  getChi(
-int order, 
 double%5B%5D
@@ -2837,32 +2837,69 @@
 .length; i++) %7B%0A
+            if (w%5Bi%5D == 0) continue;%0A
             doub
@@ -2977,23 +2977,28 @@
 ipower %3C
-= order
+ poly.length
 ; ipower
@@ -3040,24 +3040,52 @@
 ower%5D * xp;%0A
+                xp *= x%5Bi%5D;%0A
             
@@ -3163,24 +3163,24 @@
 iff * w%5Bi%5D;%0A
+
             
@@ -3182,22 +3182,8 @@
     
- if (w%5Bi%5D %3E 0)
  nDa

fix getChi, list coefficients in order in M