########################################################################################### # Copyright ironArray SL 2021. # # All rights reserved. # # This software is the confidential and proprietary information of ironArray SL # ("Confidential Information"). You shall not disclose such Confidential Information # and shall use it only in accordance with the terms of the license agreement. ########################################################################################### import iarray as ia from iarray.expr_udf import expr_udf def fuse_operands(operands1, operands2): new_operands = {} dup_operands = {} new_pos = len(operands1) for k2, v2 in operands2.items(): try: k1 = list(operands1.keys())[list(operands1.values()).index(v2)] # The operand is duplicated; keep track of it dup_operands[k2] = k1 except ValueError: # The value is not among operands1, so rebase it new_op = f"o{new_pos}" new_pos += 1 new_operands[new_op] = operands2[k2] return new_operands, dup_operands def fuse_expressions(expr, new_base, dup_op): new_expr = "" skip_to_char = 0 old_base = 0 prev_pos = {} for i in range(len(expr)): if i < skip_to_char: continue if expr[i] == "o": if i > 0 and (expr[i - 1] != " " and expr[i - 1] != "("): # Not a variable new_expr += expr[i] continue # This is a variable. Find the end of it. j = i + 1 for k in range(len(expr[j:])): if expr[j + k] in " )[": j = k break if expr[i + j] == ")": j -= 1 old_pos = int(expr[i + 1 : i + j + 1]) old_op = f"o{old_pos}" if old_op not in dup_op: if old_pos in prev_pos: # Keep track of duplicated old positions inside expr new_pos = prev_pos[old_pos] else: new_pos = old_base + new_base old_base += 1 new_expr += f"o{new_pos}" prev_pos[old_pos] = new_pos else: new_expr += dup_op[old_op] skip_to_char = i + j + 1 else: new_expr += expr[i] return new_expr class LazyExpr: """Class for hosting lazy expressions. This is not meant to be called directly from user space. Once the lazy expression is created, it can be evaluated via :func:`LazyExpr.eval`. """ def __init__(self, new_op): value1, op, value2 = new_op if op is not None and op.startswith(f"{ia.dflt_ulib}"): # A scalar UDF call args = tuple(value2) nops = 0 self.operands = {} new_args = [] for arg in args: if isinstance(arg, ia.IArray): self.operands[f"o{nops}"] = arg new_args.append(f"o{nops}") nops += 1 else: new_args.append(f"{arg}") new_args = ", ".join(new_args) self.expression = f"{op}({new_args})" return if value2 is None: # ufunc if isinstance(value1, LazyExpr): self.expression = f"{op}({self.expression})" else: self.operands = {"o0": value1} self.expression = "o0" if op is None else f"{op}(o0)" return elif op in ("atan2", "pow"): self.operands = {"o0": value1, "o1": value2} self.expression = f"{op}(o0, o1)" return if isinstance(value1, (int, float)) and isinstance(value2, (int, float)): self.expression = f"({value1} {op} {value2})" elif isinstance(value2, (int, float)): self.operands = {"o0": value1} self.expression = f"(o0 {op} {value2})" elif isinstance(value1, (int, float)): self.operands = {"o0": value2} self.expression = f"({value1} {op} o0)" else: if value1 is value2: self.operands = {"o0": value1} self.expression = f"(o0 {op} o0)" elif isinstance(value1, LazyExpr) or isinstance(value2, LazyExpr): if isinstance(value1, LazyExpr): self.expression = value1.expression self.operands = {"o0": value2} else: self.expression = value2.expression self.operands = {"o0": value1} self.update_expr(new_op) else: # This is the very first time that a LazyExpr is formed from two operands # that are not LazyExpr themselves self.operands = {"o0": value1, "o1": value2} self.expression = f"(o0 {op} o1)" def update_expr(self, new_op): # We use a lot the original IArray.__eq__ as 'is', so deactivate the overloaded one ia._disable_overloaded_equal = True # One of the two operands are LazyExpr instances value1, op, value2 = new_op if isinstance(value1, LazyExpr) and isinstance(value2, LazyExpr): # Expression fusion # Fuse operands in expressions and detect duplicates new_op, dup_op = fuse_operands(value1.operands, value2.operands) # Take expression 2 and rebase the operands while removing duplicates new_expr = fuse_expressions(value2.expression, len(value1.operands), dup_op) self.expression = f"({self.expression} {op} {new_expr})" self.operands.update(new_op) elif isinstance(value1, LazyExpr): if op == "not": self.expression = f"({op}{self.expression})" elif isinstance(value2, (int, float)): self.expression = f"({self.expression} {op} {value2})" else: try: op_name = list(value1.operands.keys())[ list(value1.operands.values()).index(value2) ] except ValueError: op_name = f"o{len(self.operands)}" self.operands[op_name] = value2 self.expression = f"({self.expression} {op} {op_name})" else: if isinstance(value1, (int, float)): self.expression = f"({value1} {op} {self.expression})" else: try: op_name = list(value2.operands.keys())[ list(value2.operands.values()).index(value1) ] except ValueError: op_name = f"o{len(self.operands)}" self.operands[op_name] = value1 if op == "[]": # syntactic sugar for slicing self.expression = f"({op_name}[{self.expression}])" else: self.expression = f"({op_name} {op} {self.expression})" ia._disable_overloaded_equal = False return self def __add__(self, value): return self.update_expr(new_op=(self, "+", value)) def __radd__(self, value): return self.update_expr(new_op=(value, "+", self)) def __sub__(self, value): return self.update_expr(new_op=(self, "-", value)) def __rsub__(self, value): return self.update_expr(new_op=(value, "-", self)) def __mul__(self, value): return self.update_expr(new_op=(self, "*", value)) def __rmul__(self, value): return self.update_expr(new_op=(value, "*", self)) def __truediv__(self, value): return self.update_expr(new_op=(self, "/", value)) def __rtruediv__(self, value): return self.update_expr(new_op=(value, "/", self)) def __and__(self, value): return self.update_expr(new_op=(self, "and", value)) def __rand__(self, value): return self.update_expr(new_op=(value, "and", self)) def __or__(self, value): return self.update_expr(new_op=(self, "or", value)) def __ror__(self, value): return self.update_expr(new_op=(value, "or", self)) def __invert__(self): return self.update_expr(new_op=(self, "not", None)) def eval(self, debug=0, cfg: ia.Config = None, **kwargs) -> ia.IArray: """Evaluate the lazy expression in self. Parameters ---------- cfg : :class:`Config` The configuration for this operation. If None (default), the current configuration will be used. kwargs : dict A dictionary for setting some or all of the fields in the :class:`Config` dataclass that should override the current configuration. Returns ------- :ref:`IArray` The output array. """ if cfg is None: cfg = ia.get_config_defaults() with ia.config(cfg=cfg, **kwargs) as cfg: expr = ia.expr_from_string(self.expression, self.operands, debug=debug, cfg=cfg) out = expr.eval() return out def __str__(self): expression = f"{self.expression}" return expression if __name__ == "__main__": # Check representations of default config import numpy as np print(ia.get_config_defaults()) print() # Create initial containers dtshape_ = [40, 20] a1 = ia.linspace(dtshape_, 0, 10) a2 = ia.linspace(dtshape_, 0, 10) a3 = ia.linspace(dtshape_, 0, 10) a4 = ia.linspace(dtshape_, 0, 10) # Evaluate with different methods # a3 = ia.tan(a1) * (ia.sin(a2) * ia.sin(a2) + ia.cos(a3)) + (ia.sqrt(a4) * 2) ia_expr = ia.sin(a1) + 2 * a1 + 1 ia_expr += 2 print(ia_expr) ia_res = ia_expr.eval().data np_res = np.sin(ia.iarray2numpy(a1)) + 2 * ia.iarray2numpy(a1) + 1 + 2 # print(np_res) np.testing.assert_allclose(ia_res, np_res) print("Everything is working fine")