Update rotate function

[jimnel]

The rotate function could do with an update. Firstly the documentation should be in numpy doc format. Also the way we do rotations for 2D systems isn't ideal (we extend to 3D, then remove the dimension at the end).