expr_udf.py

import ast
import math

import numpy as np

import iarray as ia
from iarray import udf
from iarray.py2llvm.py2llvm import MATH_FUNCS

try:
    # Python 3.9
    ast.unparse
except AttributeError:
    # Python 3.8
    from ast_decompiler import decompile
else:

    def decompile(tree):
        return ast.unparse(tree)


def name(id, ctx=ast.Load()):
    return ast.Name(id, ctx=ctx)


def constant(value):
    return ast.Constant(value, kind=None)


def For(dim, ndim, body):
    body = [For(dim + 1, ndim, body)] if dim < ndim - 1 else body
    return ast.For(
        target=name(f"i{dim}", ctx=ast.Store()),
        iter=ast.Call(
            func=name("range"),
            args=[
                ast.Subscript(
                    value=ast.Attribute(
                        value=name("out"),
                        attr="shape",
                        ctx=ast.Load(),
                    ),
                    slice=ast.Index(value=constant(dim)),
                    ctx=ast.Load(),
                )
            ],
            keywords=[],
        ),
        body=body,
        orelse=[],
    )


class Transformer(ast.NodeTransformer):
    def __init__(self, args):
        self.args = args
        # The function arguments
        dtype_map = {
            np.float32: "udf.float32",
            np.float64: "udf.float64",
            np.int8: "udf.int8",
            np.int16: "udf.int16",
            np.int32: "udf.int32",
            np.int64: "udf.int64",
            float: "udf.float64",
            int: "udf.float64",  # FIXME Should be int64
        }
        self.func_args = []
        for key, value in args.items():
            if isinstance(value, ia.IArray):
                ndim = value.ndim
                dtype = dtype_map[value.dtype]
                annotation = ast.parse(f"udf.Array({dtype}, {ndim})")
            else:
                dtype = dtype_map[type(value)]
                annotation = ast.parse(dtype)
            self.func_args.append(ast.arg(key, annotation=annotation))

        # The output is the first argument
        # FIXME output name is hardcoded, may conflict with expression names
        # FIXME Cast args types to find out type
        annotation = ast.parse(f"udf.Array({dtype}, {ndim})")
        self.func_args.insert(0, ast.arg("out", annotation))

        # Keep the ndim, and the index used to access the arrays
        self.ndim = ndim
        self.index = ast.Tuple(
            elts=[name(f"i{i}") for i in range(ndim)],
            ctx=ast.Load(),
        )

    def visit_Module(self, node):
        self.generic_visit(node)
        return ast.Module(
            body=[
                ast.FunctionDef(
                    name="f",
                    args=ast.arguments(
                        posonlyargs=[],
                        args=self.func_args,
                        vararg=None,
                        kwonlyargs=[],
                        kw_defaults=[],
                        kwarg=None,
                        defaults=[],
                    ),
                    body=[
                        For(0, self.ndim, node.body),
                        ast.Return(value=constant(0)),
                    ],
                    decorator_list=[],
                    returns=None,
                )
            ],
            type_ignores=[],
        )

    def visit_Call(self, node):
        self.generic_visit(node)

        # Translate negative(x) to math.copysign(x, -1.0)
        # https://github.com/inaos/iron-array/issues/559
        if isinstance(node.func, ast.Name):
            if node.func.id in {"negative", "negate"}:
                node.func = ast.Attribute(
                    value=name("math"),
                    attr="copysign",
                    ctx=node.func.ctx,
                )
                node.args.append(constant(-1.0))

        return node

    def visit_Expr(self, node):
        self.generic_visit(node)
        return ast.Assign(
            targets=[
                ast.Subscript(
                    value=name("out"),
                    slice=ast.Index(value=self.index),
                    ctx=ast.Store(),
                )
            ],
            value=node.value,
        )

    def visit_Name(self, node):
        # Translate math function names from those used in mingjugg to those
        # used in Python's math library
        translate_map = {
            "abs": "fabs",
            "absolute": "fabs",
            "arccos": "acos",
            "arcsin": "asin",
            "arctan": "atan",
            "arctan2": "atan2",
            "power": "pow",
        }
        node_id = translate_map.get(node.id, node.id)

        # Math functions
        if node_id in MATH_FUNCS:
            return ast.Attribute(
                value=name("math"),
                attr=node_id,
                ctx=node.ctx,
            )

        # Access to arrays
        arg = self.args.get(node_id)
        if isinstance(arg, ia.IArray):
            return ast.Subscript(
                value=node,
                slice=ast.Index(value=self.index),
                ctx=node.ctx,
            )

        return node

    def visit_Subscript(self, node):
        self.generic_visit(node)
        slice = node.slice

        # Python 3.8
        if isinstance(slice, ast.Index):
            slice = slice.value

        if isinstance(slice, (ast.UnaryOp, ast.BoolOp, ast.Compare)):
            return ast.IfExp(
                test=slice,
                body=node.value,
                orelse=constant(math.nan),
            )

        return node


def expr_udf(expr, args, cfg=None, debug=0, **kwargs):
    # There must be at least 1 argument
    assert len(args) > 0

    # Split input arrays from input scalars
    arrays = []
    scalars = []
    for value in args.values():
        if isinstance(value, ia.IArray):
            arrays.append(value)
        else:
            scalars.append(value)

    # From the string expression produce the udf function source
    tree = ast.parse(expr)  # AST of the input expression
    tree = Transformer(args).visit(tree)  # AST of the UDF function
    ast.fix_missing_locations(tree)
    source = decompile(tree)  # Source code of the UDF function
    if debug > 0:
        print(source)

    # The Python function
    exec(source, globals(), locals())
    py_func = locals()["f"]

    # The UDF function
    udf_func = udf.jit(py_func, ast=tree, debug=debug)

    # The IArray expression
    ia_expr = ia.expr_from_udf(udf_func, arrays, scalars, cfg=cfg, **kwargs)
    return ia_expr