// simeq_newton3.java  solve nonlinear system of equations
//                     method: newton iteration using Jacobian
//                     use list for higher order terms
//
//  Given problem  A X = Y  where X may have terms x1, x2, x3, x4,
//                 and higher order such as: x1*x2, x1*x1*x3, x4*x4*x4, ...
//                 A sparse matrix may be used, coefficients given
//                 Y is vector of reals given
//                 independent unknowns are x1, x2, x3, x4
//
//  for testing, generate A using pseudo random numbers
//               choose x1=1.1 x2=1.2 x3=1.4 x4 =1.5, compute products
//               compute terms of Y using Y = A X
//
//  Solve by initial guess at values of x1, x2, x3, x4 computing products
//    X_next = X_initial - J_initial^-1 * (A *  X_initial - Y)
//    in general X_next = X - (J_prev^-1 * (A * X - Y))*b
//                                 where 0 < b < 1, often 0.5, for stability
//
//  solved when  abs sum each row A * X_next -Y < epsilon
//
//  It may stall, stop if abs(X_next-X)<epsilon
//  It may diverge, stop, indicate no solution
//                        (or try a different initial guess)
//  It may oscillate, stop, indicate no solution
//                        (or try a different initial guess)
//
//
//  The matrix equation for this case is:
//
//   x1     x2     x3     x4     x1*x2 ... x3*x4  (may be subset of nl terms) 
//
// | A1,1   A1,2   A1,3   A1,4   A1,5 ...  A1,10  | | x1    | |y1 |
// | A2,1   A2,2   A2,3   A2,4   A2,5 ...  A2,10  | | x2    | |y2 |
// | A3,1   A3,2   A3,3   A3,4   A3,5 ...  A3,10  | | x3    | |y3 |
// | A4,1   A4,2   A4,3   A4,4   A4,5 ...  A4,10  |*| x4    |=|y4 |
// | A5,1   A5,2   A5,3   A5,4   A5,5 ...  A5,10  | | x1*x2 | |y5 |
// | A6,1   A6,2   A6,3   A6,4   A6,5 ...  A6,10  | | x1*x3 | |y6 |
// | ...                                          | |       | |   |
// | A10,1  A10,2  A10,3  A10,4  A10,5 ... A10,10 | | x3*x4 | |y10|
//
//  The Jacobian, J is the numerically computed derivatives based on A, X
//  The derivative coefficients may be computed once based on A and the
//  unknown variables. The numeric value plugs in the X values.
//
//  J1,1 = A1,1 + A1,5*x2 + A1,6*x3 + A1,7*x4   deriv Row 1 wrt x1
//  J1,2 = A1,2 + A1,5*x1 + A1,8*x3 + A1,9*x4   deriv Row 1 wrt x2
//  J1,3 = A1,3 + A1,6*x1 + A1,8*x2 + A1,10*x4  deriv Row 1 wrt x3
//  J1,4 = A1,4 + A1,7*x1 + A1,9*x2 + A1,10*x3  deriv Row 1 wrt x4
//  J1,5 = A1,5                                 deriv Row 1 wrt x1, wrt x2
//  J1,6 = A1,6                                 deriv Row 1 wrt x1, wrt x3
//  J1,7 = A1,7                                 deriv Row 1 wrt x1, wrt x4
//  J1,8 = A1,8                                 deriv Row 1 wrt x2, wrt x3
//  J1,9 = A1,9                                 deriv Row 1 wrt x2, wrt x4
//  J1,10 = A1,10                               deriv Row 1 wrt x3, wrt x4
//
//  etc.
//
//  j4,1 ... j4,10  should be last ones needed? others not independent
//
//  Code or a data structure must be available to know the equation
//  of the entries in the Y vector. Numeric computation of
//  derivatives is used, based on var1, var2, var3.
//

public class simeq_newton3
{
  // derivatives not hand coded, computed from var1 and var2 and var3 to
  // load Ja, initial guess in X[] n+nlin size
  public simeq_newton3(int n, int nlin, double A[][], double Y[],
                       int var1[], int var2[], int var3[], double X[],
                       int uiter, double ub, double ueps, int monitor)
  {
    double eps = 1.0e-6; // possible default
    double b = 1.00; // stability convergence  factor
    double resid;  // residual from last iteration
    double presid; // residual from prior to last iteration
    int maxiter = 10; // maximum number of iterations
    double Ja[][] = new double[n][n]; // Jacobian inverted in place
    double X_resid[] = new double[n+nlin];
    double X_tmp2[] = new double[n+nlin];
    double X_next[] = new double[n+nlin];

    if(monitor>0) 
      System.out.println("simeq_newton3 running n="+n+", nlin="+nlin);
    eps = ueps;
    maxiter = uiter;
    b = ub;
    // setup nonlinear terms
    for(int i=n; i<n+nlin; i++) // in nonlinear range
    {
      X[i] = X[var1[i]]; // must be valid
      if(var2[i]>=0) X[i] *= X[var2[i]]; // -1 flag should not be used
      if(var3[i]>=0) X[i] *= X[var3[i]];
    }
    for(int i=0; i<n+nlin; i++)
    {
      X_next[i] = X[i];
    }

    presid = 1.0e12;
    // iterate to find solution
    for(int itr=0; itr<maxiter; itr++)
    {
      // compute residual
      resid = 0.0; 
      for(int i=0; i<n; i++)
      {
        X_resid[i] = 0.0;
        for(int j=0; j<n+nlin; j++)
        {
	  X_resid[i] += A[i][j]*X[j];
        }
        X_resid[i] -= Y[i]; // X_resid used later if not solved
        // resid += Math.abs(X_resid[i]);
        resid = Math.max(resid,Math.abs(X_resid[i]));
      }
      // debug print
      if(monitor>2)
      {
        for(int i=0; i<n; i++)
        {
	  System.out.println("residual X_resid["+i+"]="+X_resid[i]);
        }
        System.out.println(" ");
      }

      // check for convergence
      if(monitor>0)
      System.out.println("simeq_newton3 itr "+itr+", prev="+presid+
                         " residual="+resid);
      if(resid<eps) break;

      // compute Jacobian
      for(int i=0; i<n; i++) // equation
      {
	for(int j=0; j<n; j++) // variable
        {
	  Ja[i][j] = A[i][j]; // each linear term
	  for(int k=n; k<n+nlin; k++) // nonlinear
	  {
            /* add derivative of nonlinear term */
	    if(var1[k]==j && var2[k]==j && var3[k]==j)       /* X^3 */
	      Ja[i][j] += A[i][k]*3.0*X[var2[k]]*X[var3[k]]; /* 3X^2 */

	    else if(var1[k]==j && var2[k]==j && var3[k]>=0)  /* X^2 xa */
	      Ja[i][j] += A[i][k]*2.0*X[var2[k]]*X[var3[k]]; /* 2X xa */
	    else if(var1[k]==j && var2[k]==j && var3[k]<0)   /* X^2 */
	      Ja[i][j] += A[i][k]*2.0*X[var2[k]];            /* 2X */
	    else if(var2[k]==j && var3[k]==j && var1[k]>=0)  /* xa X^2 */
	      Ja[i][j] += A[i][k]*2.0*X[var3[k]]*X[var1[k]]; /* 2X xa */
	    else if(var2[k]==j && var3[k]==j && var1[k]<0)   /*  X^2 */
	      Ja[i][j] += A[i][k]*2.0*X[var3[k]]           ; /* 2X  */
	    else if(var3[k]==j && var1[k]==j && var2[k]>=0)  /* xa X^2 */
	      Ja[i][j] += A[i][k]*2.0*X[var1[k]]*X[var2[k]]; /* 2X xa */
	    else if(var3[k]==j && var1[k]==j && var2[k]<0)   /*  X^2 */
	      Ja[i][j] += A[i][k]*2.0*X[var1[k]];            /* 2X  */
            else if(var1[k]==j)
	    {
	      if(var2[k]>=0 && var3[k]>=0) 
                Ja[i][j] += A[i][k]*X[var2[k]]*X[var3[k]];
              else if(var2[k]>=0) Ja[i][j] += A[i][k]*X[var2[k]];
              else if(var3[k]>=0) Ja[i][j] += A[i][k]*X[var3[k]];
	    }
            else if(var2[k]==j)
	    {
              if(var1[k]>=0 && var3[k]>=0) 
                Ja[i][j] += A[i][k]*X[var1[k]]*X[var3[k]];
              else if(var1[k]>=0) Ja[i][j] += A[i][k]*X[var1[k]];
              else if(var3[k]>=0) Ja[i][j] += A[i][k]*X[var3[k]];
	    }
            else if(var3[k]==j)
	    {
              if(var2[k]>=0 && var1[k]>=0) 
                Ja[i][j] += A[i][k]*X[var2[k]]*X[var1[k]];
              else if(var2[k]>=0) Ja[i][j] += A[i][k]*X[var2[k]];
              else if(var1[k]>=0) Ja[i][j] += A[i][k]*X[var1[k]];
	    }
	  } // end k
	} // end j
      } // end i
      if(monitor>4)
      {
        System.out.println("Ja computed ");
        for(int i=0; i<n; i++)
        {
          for(int j=0; j<n; j++)
          {
	    System.out.println("Ja["+i+"]["+j+"]="+Ja[i][j]);
	  }
        }
        System.out.println(" ");
      }

      // invert Ja
      new invert(Ja);
      if(monitor>5)
      {
        System.out.println("Ja inverted ");
        for(int i=0; i<n; i++)
        {
          for(int j=0; j<n; j++)
          {
	    System.out.println("Ja["+i+"]["+j+"]="+Ja[i][j]);
	  }
        }
        System.out.println(" ");
      }

      // X_resid = A X - Y        error vector
      // X_tmp2 = J^-1 * X_resid  raw correction vector
      // X_next = X - (X_tmp2)*b weighted correction vector
      // X      = X_next         next X with computed higher terms

      for(int i=0; i<n; i++)
      {
	X_tmp2[i] = 0.0;
        for(int j=0; j<n; j++)
        {
	  X_tmp2[i] += Ja[i][j]*X_resid[j]; 
        }
      }
      // debug print
      if(monitor>3)
      {
        for(int i=0; i<n; i++)
        {
	  System.out.println("change X_tmp2["+i+"]="+X_tmp2[i]);
        }
        System.out.println(" ");
      }
      for(int i=0; i<n; i++)
      {
        X_next[i] = X[i] - b*X_tmp2[i];
      }
      for(int i=n; i<n+nlin; i++)
      {
	X_next[i] = X_next[var1[i]]; // must be valid
        if(var2[i]>=0) X_next[i] *= X_next[var2[i]];
        if(var3[i]>=0) X_next[i] *= X_next[var3[i]];
      }

      if(monitor>2)
      {
        for(int i=0; i<n+nlin; i++)
        {
          System.out.println("X_next["+i+"]="+X_next[i]);
        }
        System.out.println(" ");
      }

      // iterate
      presid = resid; 
      for(int i=0; i<n+nlin; i++)
      {
	  X[i] = X_next[i];
      }
      if(monitor>2) System.out.println(" ");

    } // end iteration    

    if(monitor>0) System.out.println("simeq_newton3.java finished");
  }
} // end class simeq_newton3