org.matheclipse.core.reflection.system.NSolve Maven / Gradle / Ivy

package org.matheclipse.core.reflection.system;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;

import org.matheclipse.core.eval.exception.Validate;
import org.matheclipse.core.eval.exception.WrongArgumentType;
import org.matheclipse.core.eval.interfaces.AbstractFunctionEvaluator;
import org.matheclipse.core.expression.F;
import org.matheclipse.core.generic.Predicates;
import org.matheclipse.core.interfaces.IAST;
import org.matheclipse.core.interfaces.IExpr;
import org.matheclipse.core.interfaces.ISymbol;

/**
 * Try to solve a set of equations (i.e. Equal[...] expressions).
 */
public class NSolve extends AbstractFunctionEvaluator {
	/**
	 * Analyze an expression, if it has linear, polynomial or other form.
	 * 
	 */
	private static class ExprAnalyzer implements Comparable {

		final static public int LINEAR = 0;
		final static public int OTHERS = 2;
		final static public int POLYNOMIAL = 1;

		private int equationType;
		private IExpr expr;
		private IExpr numer;
		private IExpr denom;
		private int leafCount;

		IAST row;
		HashSet symbolSet;
		IAST value;

		final IAST vars;

		public ExprAnalyzer(IExpr expr, IAST vars) {
			super();
			this.expr = expr;
			this.numer = expr;
			this.denom = F.C1;
			if (this.expr.isAST()) {
				this.expr = Together.together((IAST) this.expr);
				// split expr into numerator and denominator
				this.denom = F.eval(F.Denominator(this.expr));
				if (!this.denom.isOne()) {
					// search roots for the numerator expression
					this.numer = F.eval(F.Numerator(this.expr));
				}
			}
			this.vars = vars;
			this.symbolSet = new HashSet();
			this.leafCount = 0;
			reset();
		}

		public void analyze() {
			analyze(getNumerator());
		}

		/**
		 * Analyze an expression, if it has linear, polynomial or other form.
		 * 
		 */
		private void analyze(IExpr eqExpr) {
			if (eqExpr.isFree(Predicates.in(vars), true)) {
				leafCount++;
				value.add(eqExpr);
			} else if (eqExpr.isPlus()) {
				leafCount++;
				IAST arg = (IAST) eqExpr;
				IExpr expr;
				for (int i = 1; i < arg.size(); i++) {
					expr = arg.get(i);
					if (expr.isFree(Predicates.in(vars), true)) {
						leafCount++;
						value.add(expr);
					} else {
						getPlusEquationType(expr);
					}
				}
			} else {
				getPlusEquationType(eqExpr);
			}
		}

		@Override
		public int compareTo(ExprAnalyzer o) {
			if (symbolSet.size() != o.symbolSet.size()) {
				if (symbolSet.size() < o.symbolSet.size()) {
					return -1;
				}
				return 1;
			}
			if (equationType != o.equationType) {
				if (equationType < o.equationType) {
					return -1;
				}
				return 1;
			}
			if (leafCount != o.leafCount) {
				if (leafCount < o.leafCount) {
					return -1;
				}
				return 1;
			}

			return 0;
		}

		/**
		 * @return the expr
		 */
		public IExpr getExpr() {
			return expr;
		}

		public IExpr getNumerator() {
			return numer;
		}

		public IExpr getDenominator() {
			return denom;
		}

		public int getNumberOfVars() {
			return symbolSet.size();
		}

		private void getPlusEquationType(IExpr eqExpr) {
			if (eqExpr.isTimes()) {
				ISymbol sym = null;
				leafCount++;
				IAST arg = (IAST) eqExpr;
				IExpr expr;
				for (int i = 1; i < arg.size(); i++) {
					expr = arg.get(i);
					if (expr.isFree(Predicates.in(vars), true)) {
						leafCount++;
					} else if (expr.isSymbol()) {
						leafCount++;
						for (int j = 1; j < vars.size(); j++) {
							if (vars.get(j).equals(expr)) {
								symbolSet.add((ISymbol) expr);
								if (sym != null) {
									if (equationType == LINEAR) {
										equationType = POLYNOMIAL;
									}
								} else {
									sym = (ISymbol) expr;
									if (equationType == LINEAR) {
										IAST cloned = arg.clone();
										cloned.remove(i);
										row.set(j, F.Plus(row.get(j), cloned));
									}
								}
							}
						}
					} else if (expr.isPower() && (expr.getAt(2).isInteger() || expr.getAt(2).isNumIntValue())) {
						// (JASConvert.getExponent((IAST) expr) > 0)) {
						if (equationType == LINEAR) {
							equationType = POLYNOMIAL;
						}
						getTimesEquationType(((IAST) expr).get(1));
					} else {
						leafCount += LeafCount.leafCount(eqExpr);
						if (equationType <= POLYNOMIAL) {
							equationType = OTHERS;
						}
					}
				}
				if (equationType == LINEAR) {
					if (sym == null) {
						// should never happen??
						System.out.println("sym == null???");
					}
				}
			} else {
				getTimesEquationType(eqExpr);
			}
		}

		/**
		 * @return the row
		 */
		public IAST getRow() {
			return row;
		}

		/**
		 * @return the symbolSet
		 */
		public HashSet getSymbolSet() {
			return symbolSet;
		}

		private void getTimesEquationType(IExpr expr) {
			if (expr.isSymbol()) {
				leafCount++;
				for (int i = 1; i < vars.size(); i++) {
					if (vars.get(i).equals(expr)) {
						symbolSet.add((ISymbol) expr);
						if (equationType == LINEAR) {
							row.set(i, F.Plus(row.get(i), F.C1));
						}
					}
				}
				return;
			}
			if (expr.isFree(Predicates.in(vars), true)) {
				leafCount++;
				value.add(expr);
				return;
			}
			if (expr.isPower()) {
				if (((IAST) expr).get(2).isInteger()) {
					if (equationType == LINEAR) {
						equationType = POLYNOMIAL;
					}
					getTimesEquationType(((IAST) expr).get(1));
					return;
				}
				if (((IAST) expr).get(2).isNumIntValue()) {
					if (equationType == LINEAR) {
						equationType = POLYNOMIAL;
					}
					getTimesEquationType(((IAST) expr).get(1));
					return;
				}
			}
			leafCount += LeafCount.leafCount(expr);
			if (equationType <= POLYNOMIAL) {
				equationType = OTHERS;
			}

		}

		/**
		 * @return the value
		 */
		public IAST getValue() {
			return value;
		}

		/**
		 * Return true if the expression is linear.
		 * 
		 * @return true if the expression is linear
		 */
		public boolean isLinear() {
			return equationType == LINEAR;
		}

		public boolean isLinearOrPolynomial() {
			return equationType == LINEAR || equationType == POLYNOMIAL;
		}

		public void reset() {
			this.row = F.List();
			for (int i = 1; i < vars.size(); i++) {
				row.add(F.C0);
			}
			this.value = F.Plus();
			this.equationType = LINEAR;
		}

	}

	@SuppressWarnings("serial")
	private static class NoSolution extends Exception {
		/**
		 * Solution couldn't be found.
		 */
		final static public int NO_SOLUTION_FOUND = 1;

		/**
		 * Definitely wrong solution.
		 */
		final static public int WRONG_SOLUTION = 0;

		final int solType;

		public NoSolution(int solType) {
			super();
			this.solType = solType;
		}

		public int getType() {
			return solType;
		}
	}

	/**
	 * 
	 * @param analyzerList
	 * @param vars
	 * @param resultList
	 * @param matrix
	 * @param vector
	 * @return null if the solution couldn't be found
	 */
	private static IAST analyzeSublist(ArrayList analyzerList, IAST vars, IAST resultList, IAST matrix, IAST vector)
			throws NoSolution {
		ExprAnalyzer exprAnalyzer;
		Collections.sort(analyzerList);
		int currEquation = 0;
		while (currEquation < analyzerList.size()) {
			exprAnalyzer = analyzerList.get(currEquation);
			if (exprAnalyzer.getNumberOfVars() == 0) {
				// check if the equation equals zero.
				IExpr expr = exprAnalyzer.getNumerator();
				if (!expr.isZero()) {
					if (expr.isNumber()) {
						throw new NoSolution(NoSolution.WRONG_SOLUTION);
					}
					if (!PossibleZeroQ.possibleZeroQ(expr)) {
						throw new NoSolution(NoSolution.NO_SOLUTION_FOUND);
					}
				}
			} else if (exprAnalyzer.getNumberOfVars() == 1 && exprAnalyzer.isLinearOrPolynomial()) {
				IAST listOfRules = rootsOfUnivariatePolynomial(exprAnalyzer);
				if (listOfRules != null) {
					boolean evaled = false;
					++currEquation;
					for (int k = 1; k < listOfRules.size(); k++) {
						if (currEquation >= analyzerList.size()) {
							resultList.add(F.List(listOfRules.getAST(k)));
							evaled = true;
						} else {

							ArrayList subAnalyzerList = new ArrayList();
							// collect linear and univariate polynomial
							// equations:
							for (int i = currEquation; i < analyzerList.size(); i++) {
								IExpr expr = analyzerList.get(i).getExpr();
								IExpr temp = expr.replaceAll(listOfRules.getAST(k));
								if (temp != null) {
									expr = F.eval(temp);
									exprAnalyzer = new ExprAnalyzer(expr, vars);
									exprAnalyzer.analyze();
								} else {
									// reuse old analyzer; expression hasn't
									// changed
									exprAnalyzer = analyzerList.get(i);
								}
								subAnalyzerList.add(exprAnalyzer);
							}
							try {
								IAST subResultList = analyzeSublist(subAnalyzerList, vars, F.List(), matrix, vector);
								if (subResultList != null) {
									evaled = true;
									for (IExpr expr : subResultList) {
										if (expr.isList()) {
											IAST list = (IAST) expr;
											list.add(1, listOfRules.getAST(k));
											resultList.add(list);
										} else {
											resultList.add(expr);
										}
									}
								}
							} catch (NoSolution e) {
								if (e.getType() == NoSolution.WRONG_SOLUTION) {
									evaled = true;
								}
							}
						}
					}
					if (evaled) {
						return resultList;
					}
				}
				throw new NoSolution(NoSolution.NO_SOLUTION_FOUND);
			} else if (exprAnalyzer.isLinear()) {
				matrix.add(F.eval(exprAnalyzer.getRow()));
				vector.add(F.eval(F.Negate(exprAnalyzer.getValue())));
			} else {
				throw new NoSolution(NoSolution.NO_SOLUTION_FOUND);
			}
			currEquation++;
		}
		return resultList;
	}

	/**
	 * Evaluate the roots of a univariate polynomial with the Roots[] function.
	 * 
	 * @param exprAnalyzer
	 * @param vars
	 * @return
	 */
	private static IAST rootsOfUnivariatePolynomial(ExprAnalyzer exprAnalyzer) {
		IExpr expr = exprAnalyzer.getNumerator();
		IExpr denom = exprAnalyzer.getDenominator();
		// try to solve the expr for a symbol in the symbol set
		for (ISymbol sym : exprAnalyzer.getSymbolSet()) {
			IExpr temp = Roots.rootsOfVariable(expr, denom, F.List(sym), true);
			if (temp != null) {
				IAST resultList = F.List();
				if (temp.isASTSizeGE(F.List, 2)) {
					IAST rootsList = (IAST) temp;
					for (IExpr root : rootsList) {
						IAST rule = F.Rule(sym, root);
						resultList.add(rule);
					}
					return resultList;
				}
				return null;
			}
		}
		return null;
	}

	public NSolve() {
	}

	/**
	 * Check if the argument at the given position is an equation (i.e.
	 * Equal[a,b]) or a list of equations and return a list of expressions, which
	 * should be equal to 0.
	 * 
	 * @param ast
	 * @param position
	 * @return
	 */
	private IAST checkEquations(final IAST ast, int position) {
		IAST termsEqualZeroList = F.List();
		IAST eqns = null;
		IAST eq;
		if (ast.get(position).isList()) {
			eqns = (IAST) ast.get(position);
			for (int i = 1; i < eqns.size(); i++) {
				if (eqns.get(i).isAST(F.Equal, 3)) {
					eq = (IAST) eqns.get(i);
					termsEqualZeroList.add(F.evalExpandAll(F.Subtract(eq.get(1), eq.get(2))));
				} else {
					// not an equation
					throw new WrongArgumentType(eqns, eqns.get(i), i, "Equal[] expression (a==b) expected");
				}
			}
		} else {
			if (ast.get(position).isAST(F.Equal, 3)) {
				eq = (IAST) ast.get(position);
				termsEqualZeroList.add(F.evalExpandAll(F.Subtract(eq.get(1), eq.get(2))));
			} else {
				// not an equation
				throw new WrongArgumentType(ast, ast.get(1), 1, "Equal[] expression (a==b) expected");
			}
		}
		return termsEqualZeroList;
	}

	@Override
	public IExpr evaluate(final IAST ast) {
		Validate.checkSize(ast, 3);
		IAST vars = Validate.checkSymbolOrSymbolList(ast, 2);
		IAST termsEqualZeroList = checkEquations(ast, 1);

		ExprAnalyzer exprAnalyzer;
		ArrayList analyzerList = new ArrayList();
		// collect linear and univariate polynomial equations:
		for (IExpr expr : termsEqualZeroList) {
			exprAnalyzer = new ExprAnalyzer(expr, vars);
			exprAnalyzer.analyze();
			analyzerList.add(exprAnalyzer);
		}
		IAST matrix = F.List();
		IAST vector = F.List();
		try {
			IAST resultList = F.List();
			resultList = analyzeSublist(analyzerList, vars, resultList, matrix, vector);

			if (vector.size() > 1) {
				// solve a linear equation matrix.x == vector
				IExpr temp = F.eval(F.LinearSolve(matrix, vector));
				if (temp.isASTSizeGE(F.List, 2)) {
					IAST rootsList = (IAST) temp;
					IAST list = F.List();
					for (int j = 1; j < vars.size(); j++) {
						IAST rule = F.Rule(vars.get(j), rootsList.get(j));
						list.add(rule);
					}
					resultList.add(list);
				} else {
					return null;
				}
			}

			return resultList;
		} catch (NoSolution e) {
			if (e.getType() == NoSolution.WRONG_SOLUTION) {
				return F.List();
			}
			return null;
		}
	}
}