An efficient method to calculate depth-integrated, phase-averaged momentum balances in non-hydrostatic models