Set `axis = ...` earlier in AFM.move_just_onto_surface().