Computing π using ironArray

circle_filter.py

@udf.scalar()
def circle_filter(val: udf.float32, row: udf.int64, col: udf.int64,
                  nrows: udf.int64, ncols: udf.int64)
                      -> udf.float32:
    x = (2. * row / nrows) - 1.
    y = (2. * col / ncols) - 1.
    if ((x ** 2 + y ** 2) <= 1) and val >= 0.5:
        return 1.
    return math.nan  

create_circle_board.py

expr = ia.expr_from_udf(filter_func,
                        [rand_data],
                        [shape[0], shape[1], True])
circle = expr.eval()

create_random.py

import iarray as ia

shape = (40_000, 40_000)
ia.set_config_defaults(dtype=np.float32, fp_mantissa_bits=15)
rand_data = ia.random.random_sample(shape)

create_square_board.py

expr = ia.expr_from_udf(filter_func,
                        [rand_data],
                        [shape[0], shape[1], False])
square = expr.eval()

filter_func.py

@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):
        for j in range(m):
            if iscircle:
                out[i, j] = ulib.circle_filter(
                    vals[i, j], row_start + i, col_start + j,
                    nrows, ncols)
            else:
                out[i, j] = ulib.square_filter(vals[i, j])
    return 0

pi_estimate.py

area_circle = ia.nansum(circle)
area_square = ia.nansum(square)
print(f"PI estimate: {4 * area_circle / area_square}")

square_filter.py

@udf.scalar()
def square_filter(val: udf.float32) -> udf.float32:
    if val >= 0.5:
        return 1.
    return math.nan