test.py

import ray
import matplotlib.pyplot as plt
import numpy as np

import ray.array.distributed as rd
# from scipy.linalg import lu


def main():
    ray.init(start_ray_local=True, num_workers=10)

    test = np.random.rand(100, 100)
    # test = np.identity(400)

    # p, l, u = lu(test)

    pr, lr, ur = ray.get(rd.linalg.block_lu.remote(test, block_size=40))
    # print(pr)
    # print(lr)
    # print(ur)

    result = np.dot(pr, np.dot(lr, ur))
    print("If this is above 1, something went terribly wrong. {}".format(result.max()))

    residual = np.clip(result - test, 0, 10)
    print(test)
    print(result)
    print(residual)

    plt.imshow(residual, interpolation='nearest')
    plt.colorbar()
    plt.show()


if __name__ == "__main__":
    main()