How to write a piecewise-defined function (in Python, using SymPy)

See all solutions.

In mathematics, we use the following notation for a “piecewise-defined” function.

$f(x) = \begin{cases} x^2 & \text{if } x>2 \\ 1+x & \text{if } x\leq 2 \end{cases}$

This means that for all $x$ values larger than 2, $f(x)=x^2$, but for $x$ values less than or equal to 2, $f(x)=1+x$.

How can we express this in mathematical software?

Solution

This answer assumes you have imported SymPy as follows.

1
2
from sympy import *                   # load all math functions
init_printing( use_latex='mathjax' )  # use pretty math output


SymPy has support for piecewise functions built in, using Piecewise. The function above would be written as follows.

1
2
3
var( 'x' )
formula = Piecewise( (x**2, x>2), (1+x, x<=2) )
formula


$\displaystyle \begin{cases} x^{2} & \text{for}\: x > 2 \\x + 1 & \text{otherwise} \end{cases}$

We can test to be sure the function works correctly by plugging in a few $x$ values and ensuring the correct $y$ values result. Here we’re using the method from how to substitute a value for a symbolic variable.

1
formula.subs(x,1), formula.subs(x,2), formula.subs(x,3)


$\displaystyle \left( 2, \ 3, \ 9\right)$

For $x=1$ we got $1+1=2$. For $x=2$ we got $2+1=3$. For $x=3$, we got $3^2=9$.