blog snippit 3

snippit3.py

for i, x in enumerate(np.linspace(S[0], S[2], w)):
    for j, y in enumerate(np.linspace(S[1], S[3], h)):
        coords.append((i, j))
        next_task_xys.append((x, y))
        if len(next_task_xys) == CHUNK_SIZE:
            results.append(
                trace_rays_with_bounces.remote(next_task_xys)
            )
            next_task_xys = []
if next_task_xys:
    results.append(trace_rays_with_bounces.remote(next_task_xys))

flat_results = chain.from_iterable(ray.get(results))
for coord, result in zip(coords, flat_results):
    i, j = coord
    img[h - j - 1, i, :] = result