diff --git a/equationcollection/equationcollection.py b/equationcollection/equationcollection.py index ede7d66fbbf6265fb92ba55349b3340786b1efce..8c6c6f32fef341296fd5b3163eb185f690bfd921 100644 --- a/equationcollection/equationcollection.py +++ b/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 diff --git a/equationcollection/simplifications.py b/equationcollection/simplifications.py index c8e93702b58207507341b6104d740c798f20d79f..412a889cca6273a1da429c2b65fca7fd1dcdcfc2 100644 --- a/equationcollection/simplifications.py +++ b/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) diff --git a/equationcollection/simplificationstrategy.py b/equationcollection/simplificationstrategy.py index 0caa0f269a60088609a69e73d2ade0639fb0b361..8f2078ba49ea302f08ad0bceb8ce9613634b2de1 100644 --- a/equationcollection/simplificationstrategy.py +++ b/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