new_lbm: SRT simplifications working

equationcollection/equationcollection.py

@@ -53,13 +53,16 @@ class EquationCollection:
         res.subexpressionSymbolNameGenerator = self.subexpressionSymbolNameGenerator
         return res
 
-    def newWithSubstitutionsApplied(self, substitutionDict):
+    def newWithSubstitutionsApplied(self, substitutionDict, addSubstitutionsAsSubexpresions=False):
         """
         Returns a new equation collection, where terms are substituted according to the passed `substitutionDict`.
         Substitutions are made in the subexpression terms and the main equations
         """
         newSubexpressions = [fastSubs(eq, substitutionDict) for eq in self.subexpressions]
         newEquations = [fastSubs(eq, substitutionDict) for eq in self.mainEquations]
+        if addSubstitutionsAsSubexpresions:
+            newSubexpressions = [sp.Eq(b, a) for a, b in substitutionDict.items()] + newSubexpressions
+
         res = EquationCollection(newEquations, newSubexpressions, self.simplificationHints)
         res.subexpressionSymbolNameGenerator = self.subexpressionSymbolNameGenerator
         return res

equationcollection/simplifications.py

@@ -53,7 +53,7 @@ def subexpressionSubstitutionInExistingSubexpressions(equationCollection):
                               equationCollection.subexpressionSymbolNameGenerator)
 
 
-def subexpressionSubstitutionInUpdateEquations(equationCollection):
+def subexpressionSubstitutionInMainEquations(equationCollection):
     """Replaces already existing subexpressions in the equations of the equationCollection"""
     result = []
     for s in equationCollection.mainEquations:
@@ -62,3 +62,22 @@ def subexpressionSubstitutionInUpdateEquations(equationCollection):
             newRhs = replaceAdditive(newRhs, subExpr.lhs, subExpr.rhs, requiredMatchReplacement=1.0)
         result.append(sp.Eq(s.lhs, newRhs))
     return equationCollection.newWithAdditionalSubexpressions(result, [])
+
+
+def addSubexpressionsForDivisions(equationCollection):
+    divisors = set()
+
+    def searchDivisors(term):
+        if term.func == sp.Pow:
+            if term.exp.is_integer and term.exp.is_number and term.exp < 0:
+                divisors.add(term)
+        else:
+            for a in term.args:
+                searchDivisors(a)
+
+    for eq in equationCollection.allEquations:
+        searchDivisors(eq.rhs)
+
+    newSymbolGen = equationCollection.subexpressionSymbolNameGenerator
+    substitutions = {divisor: newSymbol for newSymbol, divisor in zip(newSymbolGen, divisors)}
+    return equationCollection.newWithSubstitutionsApplied(substitutions, True)

equationcollection/simplificationstrategy.py

@@ -14,7 +14,7 @@ class SimplificationStrategy:
     def __init__(self):
         self._rules = []
 
-    def addSimplificationRule(self, rule):
+    def add(self, rule):
         """
         Adds the given simplification rule to the end of the collection.
         :param rule: function that taking one equation collection and returning a (simplified) equation collection