Extract code from the main part of `trajectories.py` into new function
`assign_traj_edge`. Because we are going to make the script more
complex.