def freeze_expression(expr): """ Reconstruct ``expr`` turning all :class:`sympy.Mul` and :class:`sympy.Add` into, respectively, :class:`devito.Mul` and :class:`devito.Add`. """ if expr.is_Atom or expr.is_Indexed: return expr elif expr.is_Add: rebuilt_args = [freeze_expression(e) for e in expr.args] return Add(*rebuilt_args, evaluate=False) elif expr.is_Mul: rebuilt_args = [freeze_expression(e) for e in expr.args] return Mul(*rebuilt_args, evaluate=False) elif expr.is_Pow: rebuilt_args = [freeze_expression(e) for e in expr.args] return Pow(*rebuilt_args, evaluate=False) elif expr.is_Equality: rebuilt_args = [freeze_expression(e) for e in expr.args] if isinstance(expr, FrozenExpr): # Avoid dropping metadata associated with /expr/ return expr.func(*rebuilt_args) else: return Eq(*rebuilt_args, evaluate=False) else: return expr.func(*[freeze_expression(e) for e in expr.args])
def freeze(expr): """ Reconstruct ``expr`` turning all sympy.Mul and sympy.Add into FrozenExpr equivalents. """ if expr.is_Atom or expr.is_Indexed: return expr elif expr.is_Add: rebuilt_args = [freeze(e) for e in expr.args] return Add(*rebuilt_args, evaluate=False) elif expr.is_Mul: rebuilt_args = [freeze(e) for e in expr.args] return Mul(*rebuilt_args, evaluate=False) elif expr.is_Pow: rebuilt_args = [freeze(e) for e in expr.args] return Pow(*rebuilt_args, evaluate=False) elif expr.is_Equality: rebuilt_args = [freeze(e) for e in expr.args] if isinstance(expr, FrozenExpr): # Avoid dropping metadata associated with /expr/ return expr.func(*rebuilt_args) else: return Eq(*rebuilt_args, evaluate=False) else: return expr.func(*[freeze(e) for e in expr.args])