Speed up `build_mask` (~100 times)

[chAwater]

Modifications

• Use np.ogrid instead of double for loop to speed up the build_mask
• Inspired by a StackOverflow post link

Tests & results

Code

# small image
%timeit -n 3 org_build_mask(256)
%timeit -n 1000 my_build_mask(256)

# large image
%timeit -n 3 org_build_mask(1800)
%timeit -n 100 my_build_mask(1800)

# diffs
(my_build_mask(256) - org_build_mask(256)).abs().sum()
(my_build_mask(1800) - org_build_mask(1800)).abs().sum()

Output

448 ms ± 44 ms per loop (mean ± std. dev. of 7 runs, 3 loops each)
2.21 ms ± 389 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

23.4 s ± 325 ms per loop (mean ± std. dev. of 7 runs, 3 loops each)
144 ms ± 4.81 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

tensor(4.9242e-08, dtype=torch.float64)
tensor(1.2405e-06, dtype=torch.float64)

BTW

I'm also trying to make the same animation using the escnn. But I have encountered some problems (I'll resolve them in the next few days, hopefully). Please let me know if you already have the escnn version of the animation.py code.