Python Define Piecewise Linear Function