Skip to content

Instantly share code, notes, and snippets.

View FrancescAlted's full-sized avatar

Francesc Alted FrancescAlted

View GitHub Profile
@udf.jit()
def filter_func(out: udf.Array(udf.float32, 2),
vals: udf.Array(udf.float32, 2),
nrows: udf.int64, ncols: udf.int64,
iscircle: udf.bool) -> udf.int32:
n = out.window_shape[0]
m = out.window_shape[1]
row_start = out.window_start[0]
col_start = out.window_start[1]
for i in range(n):
expr = ia.expr_from_udf(filter_func,
[rand_data],
[shape[0], shape[1], False])
square = expr.eval()
expr = ia.expr_from_udf(filter_func,
[rand_data],
[shape[0], shape[1], True])
circle = expr.eval()
area_circle = ia.nansum(circle)
area_square = ia.nansum(square)
print(f"PI estimate: {4 * area_circle / area_square}")