Save a sqrtf() call per particle in drift task

Any reason why this was not written like this in the first place ? Looks like we can save one sqrtf() call per particle in the drift tasks.