Created Exp in operator_file.py to represent Power

Implemented Expand and regroup

python_symb/TreeModification/basic_modif.py

@@ -0,0 +1,153 @@
+from python_symb.Expressions.expr import Expr
+from python_symb.MathTypes.symbols import Var
+from python_symb.MathTypes.operator_file import Operator, BinOperator, Add, Mul
+from typing import Union
+
+
+Number = Union[int, float]
+
+
+def expand(expr: Expr) -> Expr:
+    """
+    Expand an expression
+
+    :param expr: expression to expand
+    :return: expanded expression
+
+    example :
+    5*(a+b) -> 5*a + 5*b
+    Expr(Mul, [5, Expr(Add, [Expr(a), Expr(b)])]) -> Expr(Add, [Expr(Mul, [5, Expr(a)]), Expr(Mul, [5, Expr(b)])])
+    """
+
+    if expr.is_leaf:
+        return expr
+
+    match expr:
+        case Expr(BinOperator() as Op1, [Expr(Op2, op2_children), right]) if Op2.name in Op1.properties.left_distributive:
+            return expand(Expr(Op2, [Expr(Op1, [expand(op2_child), expand(right)]) for op2_child in op2_children]))
+
+        case Expr(BinOperator() as Op1, [left, Expr(Op2, op2_children)]) if Op2.name in Op1.properties.right_distributive:
+            return expand(Expr(Op2, [Expr(Op1, [expand(left), expand(op2_child)]) for op2_child in op2_children]))
+
+        case Expr(BinOperator() as Op, [left, right]):
+            return Expr(Op, [expand(left), expand(right)])
+
+    return expr
+
+
+def _regroup(expr: Expr, focus_op: BinOperator) -> Expr:
+    """
+    regroup an expression, with the contraint that the value of expr is focus_op
+    Will be used to regroup an expression
+
+    :param expr: expression to regroup
+    :param focus_op: operator to regroup
+    :return
+    x+x+x+x -> 4*x
+    Expr(Add, [Expr(x), Expr(Add, [Expr(x), Expr(Add, [Expr(x), Expr(x)])])]) -> Expr(Mul, [4, Expr(x)])
+
+    with Mul == Add.repeated_op
+    """
+    assert focus_op.repeated_op is not None, f'{focus_op} has no repeated_op'
+    assert expr.value == focus_op, f'{expr.value} is not a {focus_op}'
+
+    # Motifs : Key : (Expr) -> Value : int
+    # represent number of times the expression appears in the expression,
+    # custom hash make for instance x+y and y+x the same when counting
+    motifs = {}
+
+    def collect_motifs(expr: Expr):
+        match expr:
+            case Expr(BinOperator() as op, [left, right]) if op == focus_op:
+                collect_motifs(left)
+                collect_motifs(right)
+            case Expr(BinOperator() as op, [left, right]) if op == focus_op.repeated_op and isinstance(right.value, Number):
+                motifs[left] = motifs.get(expr, 0) + right.value
+            case Expr(BinOperator() as op, [left, right]) if op == focus_op.repeated_op and op.properties.commutative and isinstance(left.value, Number):
+
+                motifs[right] = 1 if right not in motifs else motifs[right] + left.value
+
+            case _:
+                motifs[expr] = 1 if expr not in motifs else motifs[expr] + 1
+
+    collect_motifs(expr)
+    tuple_motifs = list(motifs.items())
+
+    def reconstruct(tuple_motifs):
+        match tuple_motifs:
+            case [(expr, int(a))]:
+                if focus_op.repeated_op.properties.commutative:
+                    return Expr(focus_op.repeated_op, [Expr(a), expr])
+                return Expr(focus_op.repeated_op, [expr, Expr(a)])
+            case [(expr, int(a)), *rest]:
+                if focus_op.repeated_op.properties.commutative:
+                    return Expr(focus_op, [Expr(focus_op.repeated_op, [Expr(a), expr]), reconstruct(rest)])
+                return Expr(focus_op, [Expr(focus_op.repeated_op, [expr, Expr(a)]), reconstruct(rest)])
+
+    return reconstruct(tuple_motifs)
+
+def regroup(expr: Expr, focus_op: BinOperator) -> Expr:
+    """
+    Regroup an expression
+
+    :param expr: expression to regroup
+    :param focus_op: operator to regroup
+    :return: regrouped expression
+
+    example :
+    x+x+x+x -> 4*x
+    Expr(Add, [Expr(x), Expr(Add, [Expr(x), Expr(Add, [Expr(x), Expr(x)])])]) -> Expr(Mul, [4, Expr(x)])
+    """
+
+    if expr.is_leaf:
+        return expr
+
+    match expr:
+        case Expr(BinOperator() as op, [left, right]) if op == focus_op:
+            return _regroup(expr, focus_op)
+
+        case Expr(BinOperator() as op, [left, right]):
+            return Expr(op, [regroup(left, focus_op), regroup(right, focus_op)])
+
+    return expr
+
+
+def test():
+
+    x, y = Var('x'), Var('y')
+    a, b = Var('a'), Var('b')
+
+    def test_expand():
+        expr = (x+y)*(x+y)*(x+y)
+        expr = expand(expr)
+        expr = regroup(expr, Add)
+        print(f"(x+y)*(x+y)*(x+y) -> {expr.to_infix_str()}")
+        expr = (x+y+a)*b
+        print(f"(x+y+a)*b -> {expand(expr).to_infix_str()}")
+
+    def test_regroup():
+        expr = x+2*x+y+y+2*y
+        print(f"x+2*x+y+y+2*y -> {regroup(expr, Add).to_infix_str()}")
+
+    def test_power():
+        expr = x*x*x + y*y*y
+        print(f"x*x*x -> {regroup(expr, Mul).to_infix_str()}")
+
+    def test_all():
+        expr = (x+y)*(x+y)*(x+y)
+        expanded_expr = expand(expr)
+        regrouped_expr = regroup(expanded_expr, Add)
+        print(f"(x+y)*(x+y)*(x+y) -> {regrouped_expr.to_infix_str()}")
+
+
+
+
+    test_expand()
+    test_regroup()
+    test_power()
+    #test_all()
+
+if __name__ == "__main__":
+    test()
+
+