Port to_3d_grid to cython. Add unit test