search

inaos / iron-array

2 stars 0 forks source link

Being able to call scalar udf functions from vectorial udf ones #539

Closed FrancescAlted closed 2 years ago

FrancescAlted commented 2 years ago

The idea is that things like:

import numpy as np
import iarray as ia
from iarray import udf

@udf.scalar(lib="lib2")
def fmult(a: udf.float64, b: udf.float64) -> float:
    return a * b

@udf.jit
def fmult2(out: udf.Array(udf.float64, 1),
           a: udf.Array(udf.float64, 1)) -> int:
    n = out.window_shape[0]
    for i in range(n):
        out[i] = lib.fmult(4, a[i])
    return 0

print("** vector udf evaluation ...")
a1 = ia.linspace([10], 0, 10)
expr = ia.expr_from_udf(fmult2, [a1])
b1 = expr.eval()
print(b1.data)

print("** numpy evaluation ...")
b2 = 4 * a1.data
print(b2)
np.testing.assert_array_almost_equal(b1.data, b2)

would work.

jdavid commented 2 years ago

Done in https://github.com/inaos/iron-array-python/pull/160