Draw HSV color wheel

# color pallet
grid_2D = grid_3D[0, :2].view(2, -1)         # 2 (h w)
radius = grid_2D.pow(2).sum(0).unsqueeze(0)  # 1 (h w)
radius = radius / radius.max()
radian = torch.atan2(grid_2D[1], grid_2D[0])
chk_m = radian < 0
radian[chk_m] += 2 * np.pi                   # from (-pi, pi) to (0, 2pi)

h = (radian-radian.min()) / (radian.max()-radian.min()) # normalized from 0-1
s = radius # saturation is set as a function of radias
v = torch.ones_like(s) # value is constant

# convert the np arrays to lists. This actually speeds up the colorsys call
import colorsys
h,s,v = h.flatten().tolist(), s.flatten().tolist(), v.flatten().tolist()
c = [colorsys.hsv_to_rgb(*x) for x in zip(h,s,v)]
c = np.array(c)