Solver.frink

Download or view Solver.frink in plain text format


use derivatives.frink
use integrals.frink
use solvingTransformations.frink
use powerTransformations.frink

// This class solves a system of equations.  It contains a graph of
// EquationNodes.  Each EquationNode represents an equation in the system.
// These nodes are connected by Edges which indicate that various
// equations are interrelated by containing the same variables.
class Solver
{
   // This is an array of EquationNodes that make up the graph.
   var equationNodes

   // This is a set of strings indicating the variables that we are not going
   // to solve for.
   var ignoreSet

   // Boolean flag indicating if we have done phase 1 simplifications
   var initialized

   // The cached final solutions for each variable, keyed by variable name
   var finalSolutions

   // A (currently-unused) set of aliases that will simplify equations.
   var aliases

   // Create a new Solver.
   new[eqs, ignoreList=[]] :=
   {
      ignoreSet = new set
      equationNodes = new array
      finalSolutions = new dict
      
      for i = ignoreList
         ignoreSet.put[i]

      for eq = eqs
         addEquation[eq]

      initialized = false
   }
   
   // Add an equation to the system.  This creates a new EquationNode and
   // automatically connects it to all of the other EquationNodes in the
   // system that share its equations.
   //
   // If index is undef. this will push the node onto the end of the list.
   // if index is a number, this will replace the specified node.
   //
   // (public)
   addEquation[eq, index=undef] :=
   {
      // THINK ABOUT:  Call transformExpression to canonicalize first?

      // Check for duplicate equations.
      for n = equationNodes
         if structureEquals[eq, n.getOriginalEquation[]]
         {
            println["Eliminating duplicate equation $eq"]
            return
         }
            
      reducedUnknowns = setDifference[getSymbols[eq], ignoreSet]
      if length[reducedUnknowns] == 0
         println["WARNING:  Equation $eq has no unknowns!"]
      node = new EquationNode[eq, reducedUnknowns]

      // If replacing, disconnect the other node
      if index != undef
         remove[index]

      for other = equationNodes
      {
         // THINK ABOUT:  Should we check to see if one equation is a proper
         // subset of the variables of the other and push the simpler into
         // the more complex right now?
         
         // This is a set of shared variables between the two equations.
         sharedVars = intersection[reducedUnknowns, other.getUnknowns[]]

         for varName = sharedVars
            connect[node, other, varName]
      }

      if index == undef
         equationNodes.push[node]
      else
         equationNodes.insert[index, node]

      // Discard any previously-found solutions
      finalSolutions = new dict
      initialized = false
   }


   // Method to initialize and simplify the system.
   initialize[] :=
   {
      if ! initialized
      {
         pushSimpler[]
//         draw[]
         changed = solveSimultaneous[]
//         draw[]
//         pushAliases[]
//         draw[]
         if changed
            pushSimpler[]
         initialized = true
      }
   }

   // Removes the specified node from the graph.  This removes all connections
   // to the specified node.
   // (public)
   remove[index] :=
   {
      node = equationNodes.remove[index]
      disconnect[node]
      initialized = false
      finalSolutions = new dict // Discard any previously-found solutions.
   }

   // Disconnect the specified node from the graph.
   // (private)
   disconnect[node] :=
   {
      node.disconnectAllEdges[]
   }

   // Connect the two specified equations by the specified variable.
   // (private)
   connect[n1 is EquationNode, n2 is EquationNode, varName is string] :=
   {
      e = new Edge[n1, n2, varName]
      n1.addEdge[e]
      n2.addEdge[e]
   }

   // Return a count of the EquationNodes in the system
   // (public)
   getEquationCount[] := length[equationNodes]

   // Returns the EquationNode with the specified index
   // (public)
   getEquationNode[index] := equationNodes@index

   // Returns an array with each element in the array being an array
   // [ unknowns, index ]
   // of the equations in the system, ordered with the simplest equations
   // (those with the fewest unknowns) first.  Unknowns is a set.
   getEquationsSortedByComplexity[] :=
   {
      list = new array
      i = 0
      last = length[equationNodes] - 1
      for i = 0 to last
         list.push[[equationNodes@i.getUnknowns[], i]]

      sort[list, {|a,b| length[a@0] <=> length[b@0]}]
      return list
   }


   // Returns the unknowns for the specified index.
   getUnknowns[index] := equationNodes@index.getUnknowns[]

   // Prints out the state of the solver for debugging.
   // (public)
   dump[] :=
   {
      last = getEquationCount[]-1
      for i = 0 to last
      {
         node = getEquationNode[i]
         print["$i\t" + node.getOriginalEquation[] + "\t"]
         for e = node.getEdges[]
         {
            other = e.getOtherNode[node]
            print["[" + e.getVariableName[] +  "," + getNodeIndex[other] +"] "]
         }
         println[]
      }

      println[]
   }

   // Draw a representation of the system.
   draw[g is graphics, left=0, top=0, right=1, bottom=2] :=
   {
      last = getEquationCount[]-1
      width = right-left
      height = bottom-top
      g.font["Serif", "italic", height/30]
      g.color[0,0,0]
      cy = top + height/2
      w = 0.7 width/20;

      for i = 0 to last
      {
         node = getEquationNode[i]
         [x,y] = getPosition[i, left, top, right, bottom]
         for e = node.getEdges[]
         {
            oi = getNodeIndex[e.getOtherNode[node]]
            if (i < oi)
            {
               [ox, oy] = getPosition[oi, left, top,right,bottom]
               g.color[0,0,0]
               g.line[x+randomFloat[-w,w],y+randomFloat[-w,w],ox+randomFloat[-w, w],oy+randomFloat[-w,w]]
            }
         }
         
         g.color[.9,.9,.9]
         g.fillEllipseCenter[x,y,width/10, height/2/10]
         g.color[0,0,0]
         g.drawEllipseCenter[x,y,width/10, height/2/10]
         g.text[i,x,y,"center","center"]
         g.text["$i: " + node.getOriginalEquation[], left, i * height/20 + cy, "left", "top"]
      }
   }

   // Calculates the center of the specified node in the graph.
   getPosition[index, left, top, right, bottom] :=
   {
      phase = circle / getEquationCount[]
      width = right-left
      height = bottom-top
      cx = left + width/2
      cy = top + height/4
      radius = 1/2 (.9) width
      x = radius sin[index phase] + cx
      y = radius * -cos[index phase] + cy
      return [x,y]
   }

   draw[] :=
   {
      g = new graphics
      draw[g]
      g.show[]
   }

   // Gets the node index given the node.
   getNodeIndex[node is EquationNode] :=
   {
      last = length[equationNodes]
      for i=0 to last
         if equationNodes@i == node
            return i
         
      return "Node not found!"
   }

   // Keep replacing simpler variables until we're done.
   pushSimpler[] :=
   {
//      dump[]
      while pushSimplerOnce[]
      {
//         dump[]
      }
   }

   // Pushes the simpler equations into the more complex
   // equations.  This returns true if changes were made to the system,
   // false otherwise.
   pushSimplerOnce[] :=
   {
      sortedEqs = getEquationsSortedByComplexity[]
      last = length[sortedEqs]-1
      for i=0 to last-1
      {
         [unknownsI, nodeI] = sortedEqs@i
         JLOOP:
         for j=i+1 to last
         {
            [unknownsJ, nodeJ] = sortedEqs@j
            if isProperSubset[unknownsI, unknownsJ]
            {
//               println["Node $nodeI is a proper subset of $nodeJ"]
               simpleVarList = pickSimplestVariable[nodeI, nodeJ]
               for [simpleVar, count] = simpleVarList
               {
                  sols = equationNodes@nodeI.getSolutions[simpleVar]
                  if length[sols] < 1
                  {
                     println["Ugh.  Has " + length[sols] + " solutions in pushSimplerOnce.  Maybe improve pickSimplestVariable?  Solutions are: $sols, equations are " + equationNodes@nodeI.getOriginalEquation[] + ", " + equationNodes@nodeJ.getOriginalEquation[] ]
                     next
                  } else
                  {
                     origEq = equationNodes@nodeJ.getOriginalEquation[]
//                     println["Removing node $nodeJ"]
                     remove[nodeJ]

                     if (length[sols] > 1)
                        println["Warning:  pushSimplerOnce split into " + length[sols] + " solutions: $sols"]
                     for sol = sols
                     {
//                        println["Replacing $simpleVar in $origEq with " + child[sol,1]]
                        newEq = replaceVar[origEq, simpleVar, child[sol,1]]
                        [newEq, fullyReduced] = prettify[newEq]
                        if !fullyReduced
                           newEq = transformExpression[newEq]
                        addEquation[newEq]
                     }
                     return true
                  }
               }
            }
         }
      }

      return false              // We made no changes.
   }

   // Pushes aliases like a == b or a == 2 b around the system so that
   // the system is as simple and disconnected as possible.
   pushAliases[] :=
   {
      size = length[equationNodes]
      i = 0
      while i<size
      {
         node = equationNodes@i
         eq = equationNodes@i.getOriginalEquation[]
         left =  child[eq, 0]
         lu = getSymbols[left]
         lu = setDifference[lu, ignoreSet]
         if length[lu] == 1
         {
            right = child[eq, 1]
            ru = getSymbols[right]
            ru = setDifference[ru, ignoreSet]
            if length[ru] == 1
            {
               ls = array[lu]@0
               rs = array[ru]@0

               first = ls > rs ? ls : rs
               sols = node.getSolutions[first]
               if length[sols] != 1
               {
                  println["Ugh.  Has " + length[sols] + " solutions in pushAliases.  Solutions are:\n" + sols]
               } else
               {
                  println["Replacing " + sols@0]
                  replaceAll[sols@0, i]
                  remove[i]
                  i = i - 1     // Don't increment i
                  size = size - 1
               }
            }
         }
         i = i + 1
      }
   }
   
   // Simplifies an equation to the form a === 10 if only one variable
   // remains.
   // returns [equation, remainingSymbol, reduced]
   //   where reduced is a boolean flag indicating if the equation has been
   //   simplified to the form above.
   prettify[eq] :=
   {
      reduced = false
      solvedUnknowns = getSymbols[eq]
      solvedUnknowns = setDifference[solvedUnknowns, ignoreSet]
      remainingSymbol = undef
      if length[solvedUnknowns] == 1
      {
         remainingSymbol = array[solvedUnknowns]@0
         prettified = array[solveSingle[eq, remainingSymbol]]

         if length[prettified] == 1
         {
            reduced = true
            eq = prettified@0
         } else
            eq = transformExpression[eq]

      }
      return [eq, reduced]
   }

   // Picks the simplest variable shared by nodes 1 and 2.  Node 1 should be
   // the simpler node.  Returns a string.
   pickSimplestVariable[index1, index2] :=
   {
      node1 = equationNodes@index1
      node2 = equationNodes@index2
      u1 = node1.getUnknowns[]
      u2 = node2.getUnknowns[]
      intersection = intersection[u1, u2]

      results = new array
      sortedUnknowns = getSymbolsByComplexity[node1.getOriginalEquation[]]
      //println["Sorted unknowns is $sortedUnknowns"]
      for [unknown, count] = sortedUnknowns
         if intersection.contains[unknown]
            results.push[[unknown, count]]

      sortedUnknowns = getSymbolsByComplexity[node2.getOriginalEquation[]]
      for [unknown, count] = sortedUnknowns
         if intersection.contains[unknown]
            for i = 0 to length[results]-1
               if results@i@0 == unknown
                  results@i@1 = results@i@1 + count

      sort[results, {|a,b| a@1 <=> b@1}]
      return results       
   }

   // Returns true if the node at index 1 contains unknowns which are a
   // proper subset of the unknowns in index 2.  This means that node 1
   // is a "simpler" version of equation 2, and its values should be
   // substituted into equation 2.
   isSimpler[index1, index2] :=
   {
      return isProperSubset[equationNodes@index1.getUnknowns[],
                            equationNodes@index2.getUnknowns[]]
   }

   // Eliminate simultaneous equations in the system.
   // returns true if the system has been changed.
   solveSimultaneous[] :=
   {
      changed = false
      // TODO:  Sort by simplest equations?
      size = length[equationNodes]
      for i=0 to size-2
      {
         JLOOP:
         for j = i+1 to size-1
         {
            nodeI = equationNodes@i
            nodeJ = equationNodes@j
            ui = nodeI.getUnknowns[]
            uj = nodeJ.getUnknowns[]
            sharedUnknowns = intersection[ui, uj]

//            println[nodeI.getOriginalEquation[]]
//            println[nodeJ.getOriginalEquation[]]
//            println["$i: $ui\t$j: $uj"]
//            println["$i $j Shared unknowns are $sharedUnknowns"]

            if length[sharedUnknowns] >= 2
            {
               varsToReplace = pickSimplestVariable[i, j]
//               println["varsToReplace is $varsToReplace"]
               for [varToReplace, count] = varsToReplace
               {
                  skipNode = i
                  solution = nodeI.getSolutions[varToReplace]
                  if length[solution] != 1
                  {
                     // Didn't find single solution, try solving
                     // and replacing from the other node.
                     solution = nodeJ.getSolutions[varToReplace]
                     skipNode = j
                  }
                  
                  if length[solution] == 1
                  {
                     replaceAll[solution@0, skipNode]
                     //                  dump[]
                     changed = true
                     break JLOOP
                  }
               }

               println["Ugh.  SolveSimultaneous fell through without replacing.  Equations were " + nodeI.getOriginalEquation[] + " and " + nodeJ.getOriginalEquation[]]
            }
         }
      }
      return changed
   }

   // Replace the specified symbol, recursively, in all equations except
   // the index specified.
   // (private)
   replaceAll[solution, skipIndex] :=
   {
      size = length[equationNodes]
      sym = child[solution,0]
      rep = child[solution,1]   // Right-hand-side of solution
      // Substitute result into other equations.
      for k = 0 to size-1
         if k != skipIndex
         {
            orig = equationNodes@k.getOriginalEquation[]
            subst = substituteExpression[orig, sym, rep]
//            println["orig is $orig, sym is $sym, solution is $solution, rep is $rep, subst is $subst"]
            if orig != subst // and length[getSymbols[subst]] <= length[getSymbols[orig]]
            {
               [subst, eqSolved]  = prettify[subst]
               subst2 = transformExpression[subst]   // THINK ABOUT:  Do this?
               if structureEquals[_a === _b, subst2] //  and structureEquals[child[subst2,0],sym] and ! expressionContains[child[subst2,1], sym]
                  subst = subst2
               else
                  println["Warning:  In replaceAll, did not get solution.  Input was $solution, output was $subst2"]

//               println["Substituted $sym to " + rep + " in $orig, result is $subst"]
               addEquation[subst, k]  // Replace equation.
               if eqSolved
               {
//                  println["Going to recursively replace $sym"]
                  replaceAll[subst, k]  // Recursively replace others
               }
            }
         }
   }

   // Return a set of all unknowns in the system.
   getAllUnknowns[] :=
   {
      allUnknowns = new set
      for node = equationNodes
         allUnknowns = union[allUnknowns, node.getUnknowns[]]

      return allUnknowns
   }
   
   // Solves for all variables in the system.
   solveAll[] :=
   {
      allUnknowns = getAllUnknowns[]

      results = new array
      for u = allUnknowns
      {
         res = solveFor[u]
         for eq = res
            results.push[eq]
      }
      return results
   }

   // Solves the system for the specified variable name.
   // (public)
   solveFor[varName] :=
   {
      if !initialized
         initialize[]

      cached = finalSolutions@varName
      if cached
         return cached

      results = new array
      size = getEquationCount[]
      for i=0 to size-1
      {
         if getUnknowns[i].contains[varName]
         {
            partialResults = solveNodeForVariable[i, varName]
            for r = partialResults
               results.push[r]
         }
      }

      // Cache results.
      finalSolutions@varName = results
      return results
   }


   // Solve for the specified variable name, substituting the list of
   // arguments.  Args is an array of ["varname", value] pairs.
   // The answer will be returned symbolically as an equation in the form
   // varName === solution
   // with constants and units still intact.
   solveForSymbolic[varName, args] :=
   {
      results = new array
      sols = solveFor[varName]
      for sol = sols
      {
         for [arg, val] = args
         {
            sym = constructExpression["Symbol", arg]
            sol = substituteExpression[sol, sym, val]
         }

         // THINK ABOUT:  Transform expression here to simplify?
         // res = transformExpression[res]

         results.push[eval[sol]]
      }

      return eliminateOverconstrained[results, false, false]
   }

   // Solve for the specified variable name, substituting the list of
   // arguments.  The result is a list of evaluated solutions.
   solveFor[varName, args] :=
   {
      sols = solveForSymbolic[varName, args]
      results = new array
      for sol = sols
      {
         right = child[sol,1]
         final = eval[right]
         exists = false
         CHECKDUPLICATES:
         for r = results
            if (final conforms r) and (final == r)
            {
               exists = true
               break CHECKDUPLICATES
            }
         if ! exists        
            results.push[final]
      }

      return results
   }

   // Recursive method to find the solutions for the specified variable
   // starting from the specified node.  This recursively enumerates all
   // of the permutations of substitutions in the system.
   // This method just sets up parameters for the recursive call.
   // (Private method)
   solveNodeForVariable[index, varName, cachedSolutions = undef] :=
   {
      if cachedSolutions == undef
         cachedSolutions = new dict
      
      node = getEquationNode[index]
      sols = node.getSolutions[varName]
//      println["Solutions for $varName are $sols"]
      results = solveNodeForVariable[node, varName, sols, new set, cachedSolutions]
      results = transformExpression[results]
//      return results
      return eliminateOverconstrained[results, true, false]
   }

   
   // The recursive (private) call to solve for the particular variable.
   solveNodeForVariable[node, varName, inEqs, usedEdges, cachedSolutions] :=
   {
//      print["Solving for $varName in " + getNodeIndex[node] + ", {"]
//      for e = usedEdges
//         print[e.getVariableName[] + " "]
//      println["}"]
      
      // Return partial solution from cache if possible.
      if cachedSolutions.containsKey[node]
      {
         varDict = cachedSolutions@node
         if varDict.containsKey[varName]
         {
            edgeDict = varDict@varName
            if edgeDict.containsKey[usedEdges]
               return edgeDict@usedEdges
         }
      }

      results = inEqs.shallowCopy[]
      edges = setDifference[node.getEdges[], usedEdges]
      for e = edges
         if e.getVariableName[] == varName
            edges.remove[e]
            
      len = length[edges]
      if (len == 0)             // No more replacements to do.
      {
         putCache[cachedSolutions, node, varName, usedEdges, results]
         return results
      }

      // Set up states array to enumerate through permutations.
      states = new array 
      for i=0 to len-2
         states@i = false
      
      states@(len-1) = true     // Skip all-false state (no replacements)

      i = len-1
      edgeArray = array[edges]
      
      //newUsedEdges = union[node.getEdges[], usedEdges]
      
      while i >= 0
      {
         newUsedEdges = usedEdges.shallowCopy[]
         for j = 0 to len-1
            if states@j
               newUsedEdges.put[edgeArray@j]
            
         // Perform replacements on each edge
         EDGELOOP:
         for j = 0 to len-1
         {
            edge = edgeArray@j
//            newUsedEdges.put[edge]  // Mark this edge as used.
            replacingVar = edge.getVariableName[]
            if states@j
            {
               replacingSymbol = edge.getSymbol[]
               otherNode = edge.getOtherNode[node]
               //                  newGlobalUsedNodes.put[newUsedNodesHere]
               // Recursively solve the other node for the variable
               // represented by this edge.
               repList = solveNodeForVariable[otherNode,
                                              replacingVar,
                                              otherNode.getSolutions[replacingVar],
                                              newUsedEdges,
                                              cachedSolutions ]
               //               println["repList is $repList"]
               for repWithFull = repList
               {
                  repWith = child[repWithFull, 1]  // Get right-hand-side
                  for eq = inEqs
                  {
                     res = substituteExpression[eq, replacingSymbol, repWith]
                     //                     println["Replacing $replacingVar with $repWith in $eq, result is $res"]

                     // Check to see if the variable we're solving for occurs on the right
                     rightSyms = getSymbols[child[res,1]]
                     if rightSyms.contains[varName]
                     {
                        //println["WARNING:  Right side contains $varName in $res"]
                        res2 = solveSingle[res, varName]
                        //println["Re-solving:  $res2"]
                        
                        // TODO:  This may return a whole lot of solutions.
                        // We need to evaluate each one and push them all
                        // onto the solutions list.
                        varSymbol = constructExpression["Symbol", varName]
                        for subR = array[res2]
                           if structureEquals[_a === _b, subR] and structureEquals[child[subR,0],varSymbol] and ! expressionContains[child[subR,1], varSymbol]
                           {
                              //println["Re-solving successful."]
                              results.push[subR]
                           } else
                           {
//                              println["WARNING:  Right side contains $varName in $res"]
                              println["Re-solving FAILED:  $res2"]
//                              println["Re-solving FAILED."]
                           }
                        } else
                        {
                           //                      res = transformExpression[res]
                           results.push[res]
                        }
                     }
                  }
               }
         }
         
         // Advance to next binary state
         flipped = false
         i = len-1
         while i>=0 and !flipped
         {
            // Enter next state
            if states@i == false
            {
               states@i = true
               flipped = true
            } else
            {  // Carry
               states@i = false
               i = i - 1
            }
         }

         // i now contains the last index flipped.  If i < 0, we're done
      }

      results = eliminateOverconstrained[results, true, false]
      putCache[cachedSolutions, node, varName, usedEdges, results]
      return results
   }


   // This function eliminates overconstrained equations.  For example, a
   // system containing the solutions a===1/2 c r  and  a===c d^-1 r^2  is
   // overconstrained because a value can always be obtained with the first
   // equation.  The second is not necessary, and could lead to
   // inconsistent results.   This method ignores any symbols listed in the
   // ignoreSymbols list, (these are probably units,) eliminating them from
   // the equations.
   eliminateOverconstrained[eqArray, dupsOnly, debug=false] :=
   {
      size = length[eqArray]
      unknowns = new array
      lefts = new array
      for i = 0 to size-1
      {
         lefts@i = child[eqArray@i, 0]
         unknowns@i = setDifference[getSymbols[child[eqArray@i,1]], ignoreSet]
      }

      res = new array

      // Check for duplicates.
      for i=0 to size-1
      {
         remove = false
         j = 0
         do
         {
            if i != j and structureEquals[lefts@i, lefts@j]
            {
               remove = (i<j and structureEquals[eqArray@i, eqArray@j]) or ((! dupsOnly) and isProperSubset[unknowns@j, unknowns@i])
               if remove
                  if debug
                     println[eqArray@j + " is a proper subset or match of " + eqArray@i]
            }
            j=j+1
         } while (j < size) and ! remove

         if (! remove)
            res.push[eqArray@i]  // If we got here, no j is a proper subset of i.
      }

      return res
   }

   // Puts the specified values into the cache.
   class putCache[cachedSolutions, node, varName, usedEdges, vals] :=
   {
      if ! cachedSolutions.containsKey[node]
      {
         nodeDict = new dict
         cachedSolutions@node = nodeDict
      } else
         nodeDict = cachedSolutions@node
         
      if ! nodeDict.containsKey[varName]
      {
         varDict = new dict
         nodeDict@varName = varDict
      } else
        varDict = nodeDict@varName 

      varDict@usedEdges = vals
   }

   // This is an experimental function that uses Frink's multi-input capability
   // to interactively plug numbers into a solution.
   interact[] :=
   {
      allUnknowns = sort[array[getAllUnknowns[]]]
      opts = new array
      for u = allUnknowns
      {
         sols = solveFor[u]
         println["$u:\t$sols"]
         if length[sols] != 1
            opts.push[u]
      }

      vals = input["Enter values: ", opts]
   }
}


// This is a node in the graph that represents an equation.  It stores various
// information about the variables stored in the equation and its solutions
// for those variables.  Users will not create these directly, but rather
// call methods on the System class to create these nodes and connect them
// properly.
class EquationNode
{
   // The original equation
   var origEq

   // A set of unknowns in the equation.
   var unknowns

   // An set of edges that connect this node to other nodes.
   var edges

   // This is a dictionary whose key is the variable name (as a string)
   // and the value is an object of type SolutionPart.
   // If this is undef, it means that the equation has not 
   var solvedDict

   // Create a new equation.  The equation should contain a === expression.
   // The set reducedUnknowns is the unknowns with the ignoreSet removed.
   new[eq, reducedUnknowns is set] :=
   {
      origEq = eq
      unknowns = reducedUnknowns
      edges = new set
      solvedDict = new dict
   }

   // Add a new edge that connects to another node
   addEdge[e is Edge] :=
   {
      edges.put[e]
   }

   // Returns the primary equation for this node.
   getOriginalEquation[] := origEq

   // Return the set of unknowns in this equation.
   getUnknowns[] := unknowns

   // Disconnect any edges that connect this node to other nodes.
   disconnectAllEdges[] :=
   {
      for e = edges
      {
         other = e.getOtherNode[this]
         other.removeEdgesTo[this]
      }

      edges = new set
   }

   // Remove any edges that connect to the specified EquationNode.
   // This NOT recursive and should only be called from disconnectAllEdges
   // (private method)
   removeEdgesTo[node is EquationNode] :=
   {
      for e = edges
         if e.connectsTo[node]
            edges.remove[e]
   }

   // Returns a set of all Edge objects.
   getEdges[] := edges

   // Gets the solutions to this equation for the specified variable.
   // This will fetch the value from the cache if it exists, otherwise it
   // will solve for it.
   getSolutions[varName] :=
   {
      if ! solvedDict.containsKey[varName]
      {
         solution = solveSingle[origEq, varName]
         addSolutions[varName, solution]
      }
      
      return solvedDict@varName.solutions
   }

   // Add one or more solutions for the specified variable to a node.
   // The equation may be a single equation or a list of equations.
   // This checks to ensure that the equations are properly solved for the
   // variable.
   addSolutions[varName, equation] :=
   {
      if solvedDict.containsKey[varName]
         sp = solvedDict@varName
      else
      {
         sp = new SolutionPart[varName]
         solvedDict@varName = sp
      }

      sym = sp.symbol
      for eq = flatten[array[equation]]
      {
         // Make sure that the solution was properly solved for the variable.
         // this ensures that the equation is of the form var == solution
         // where solution does not contain the variable.
         if structureEquals[_a === _b, eq] and structureEquals[child[eq,0],sym] and ! expressionContains[child[eq,1], sym]
            sp.addSolution[eq]
         else
            println["Could not solve $origEq for $varName!  Solution was $eq"]
      }
   }
}


// This represents an edge between two EquationNodes in a graph.  It
// contains a variable name which defines the variable that connects the
// two nodes.  Users will not create these directly, but rather
// call methods on the System class to create these nodes and connect them
// properly.
class Edge
{
   // A string containing the name of the variable that connects the two
   // EquationNodes
   var varName

   // A symbol representing the symbol.
   var varSymbol

   // One of the EquationNodes that this connects.
   var node1

   // The other one of the EquationNodes that this connects.
   var node2

   // Create a new Edge that connects the specified EquationNodes.
   new[n1 is EquationNode, n2 is EquationNode, name is string] :=
   {
      node1 = n1
      node2 = n2
      varName = name
      varSymbol = constructExpression["Symbol", name]
   }

   // Returns a string containing the variable name which indicates the
   // variable name that connects the two Nodes by this edge
   getVariableName[] := varName

   // Returns a symbol representing the variable.
   getSymbol[] := varSymbol
   
   // Returns the other node of this edge.
   getOtherNode[oneNode is EquationNode] :=
   {
      if oneNode == node1
         return node2
      else
         if oneNode == node2
            return node1
         else
         {
            println["getOtherNode:  Matching node not found!"]
            return undef
         }
   }

   // Returns true if this Edge connects to the specified node.
   connectsTo[node is EquationNode] :=
   {
      return node == node1 or node == node2
   }
}


// This is a helper class that contains the solutions for a particular
// variable.
class SolutionPart
{
   var symbol                   // The variable stored as a symbol.

   var solutions                // An array of solutions in the form
                                //  a === b + c

   // Construct a new SolutionPart given the string that represents the
   // variable.
   new[symbolString] :=
   {
      symbol = constructExpression["Symbol", symbolString]
      solutions = new array
   }

   // Adds a solution to the list.  This does not do any checking; that is
   // performed by EquationNode.addSolution.
   addSolution[eq] :=
   {
      solutions.push[eq]
   }
}


// Solve a single equation for the specified symbol.
// TODO:  Verify that the equation was solved appropriately?
solveSingle[eq, symbol] :=
{
   xSymbol = constructExpression["Symbol", symbol]
   // We could use this in symbolic mode, otherwise it warns.
   // solveEq = solve[eq, xSymbol]
   solveEq = constructExpression["FunctionCall", ["solve", eq, xSymbol]]
   
   return transformExpression[solveEq]
}


// Replace a variable in the specified equation with the specified value.
// You should probably wrap the "eq" and "value" in a noEval[] block if
// you're not passing them in from a variable.
replaceVar[eq, varName, value] :=
{
   sym = constructExpression["Symbol", varName]
   res = substituteExpression[eq, sym, value]
   // THINK ABOUT:  Transform expression here to simplify?
   // res = transformExpression[res]
   return res
}


Download or view Solver.frink in plain text format


This is a program written in the programming language Frink.
For more information, view the Frink Documentation or see More Sample Frink Programs.

Alan Eliasen was born 19935 days, 11 hours, 36 minutes ago.