Refactor align_rotation_axis

Also addresses #240.