An Overview of SymPy

Importing Sympy

from sympy import *

from sympy import *

Symbols

Symbols can be created in a few different ways in SymPy, for example,
sympy.Symbol, sympy.symbols, and sympy.var

x = Symbol('x')
x

$\displaystyle x$

The variable $x$ now represents an abstract mathematical symbol which could, for example, represent a real number, an integer, a complex number, a function, as well as a large number of other possibilities.
If we have a mathematical variable $y$ that is known to be a real number, we can use the real=True keyword argument when creating the corresponding symbol instance. We can verify that SymPy indeed recognizes that the symbol is real by using the is_real attribute of the Symbol class:

y = Symbol('y', real=True)
y.is_real

True

If, on the other hand we were to use is_real to query the previously defined symbol $x$, which was not explicitly specified to real, and therefore can represent both real and nonreal variables, we get None as result:

x.is_real is None

False

is_real returns True if the symbol is known to be real, False if the symbol is known to be not real, and None if it is not known if the symbol is real or not

Symbol('z', imaginary=True).is_real

False

Explicitly specifying when creating new symbols as real and positive or anything else as required, can help SymPy to simplify various expressions further than otherwise possible.

x = Symbol('x')
y = Symbol('y', positive=True)

sqrt(x**2)

$\displaystyle \sqrt{x^{2}}$

sqrt(y**2)

$\displaystyle y$

When working with mathematical symbols that represent integers, rather than real numbers, it is also useful to explicitly specify this when creating the corresponding SymPy symbols, using, for example, the integer=True, or even=True or odd=True, if applicable. This may also allow SymPy to analytically simplify certain expressions and function evaluations.

n1 = Symbol('n')
n2 = Symbol('n', integer=True)
n3 = Symbol('n', odd=True)

cos(n1*pi)

$\displaystyle \cos{\left(\pi n \right)}$

cos(n2*pi)

$\displaystyle \left(-1\right)^{n}$

cos(n3*pi)

$\displaystyle -1$

Using Python’s tuple unpacking syntax together with a call to sympy.symbols is a convenient way to create multiple symbols:

a, b, c = symbols('a, b, c', negative=True)
d, e, f = symbols('d, e, f', positive=True)

Constants and Special Symbols

Selected mathematical constants and special symbols and their corresponding symbols in SymPy:

Mathematical Symbol	SymPy Symbol	Description
$\pi$	pi	Ratio of the circumference to the diameter of a circle.
$e$	E	The base of the natural logarithm $e = exp (1)$.
$\gamma$	EulerGamma	Euler's constant
$i$	I	The imaginary unit.
$\infty$	oo	Infinity

Functions

Undefined Functions

x, y, z = symbols('x, y, z')
f = Function('f')
type(f)

sympy.core.function.UndefinedFunction

f(x)

$\displaystyle f{\left(x \right)}$

g = Function('g')(x, y, z)
g

$\displaystyle g{\left(x,y,z \right)}$

g.free_symbols

{x, y, z}

Defined Functions

type(sin)

sympy.core.function.FunctionClass

sin(1.5*pi)

$\displaystyle -1$

n = Symbol('n')
sin(n*pi)

$\displaystyle \sin{\left(\pi n \right)}$

n = Symbol('n', integer=True)
sin(n*pi)

$\displaystyle 0$

Lambda Functions

It can be created sympy.Lambda

f = Lambda(x, x**2)
f

$\displaystyle \left( x \mapsto x^{2} \right)$

f(1.5)

$\displaystyle 2.25$

f(1 + x)

$\displaystyle \left(x + 1\right)^{2}$

Expressions

In SymPy, mathematical expressions are represented as trees where leafs are symbols, and nodes are class instances that represent mathematical operations.
Examples of these classes are Add, Mul, and Pow for basic arithmetic operators, and Sum, Product, Integral, and Derivative for analytical mathematical operations.

x = symbols('x')
expr = 1 + x + x**2 + x**3
expr

$\displaystyle x^{3} + x^{2} + x + 1$

expr.args

(1, x, x**2, x**3)

expr.args[0]

$\displaystyle 1$

expr.args[1]

$\displaystyle x$

expr.args[2]

$\displaystyle x^{2}$

expr.args[3]

$\displaystyle x^{3}$

Manipulating Expressions

Simplification

expr = 2 * (x**2 - x) - x * (x + 1)
expr

$\displaystyle 2 x^{2} - x \left(x + 1\right) - 2 x$

simplify(expr)

$\displaystyle x \left(x - 3\right)$

expr.simplify()

$\displaystyle x \left(x - 3\right)$

expr

$\displaystyle 2 x^{2} - x \left(x + 1\right) - 2 x$

expr = 2 * cos(x) * sin(x)
expr

$\displaystyle 2 \sin{\left(x \right)} \cos{\left(x \right)}$

expr.simplify()

$\displaystyle \sin{\left(2 x \right)}$

expr = exp(x)*exp(y)
expr

$\displaystyle e^{x} e^{y}$

simplify(expr)

$\displaystyle e^{x + y}$

Each specific type of simplification can also be carried out with more specialized functions, such as sympy.trigsimp and sympy.powsimp, for trigonometric and power simplifications, respectively.
Summary of selected SymPy functions for simplifying expressions:

Function	Description
sympy.simplify	Attempt various methods and approaches to obtain a simpler form of a given expression.
sympy.trigsimp	Attempt to simplify an expression using trigonometric identities.
sympy.powsimp	Attempt to simplify an expression using laws of powers.
sympy.compsimp	Simplify combinatorial expressions.
sympy.ratsimp	Simplify an expression by writing on a common denominator

Expand

expr = (x + 1) * (x + 2)
expr

$\displaystyle \left(x + 1\right) \left(x + 2\right)$

expand(expr)

$\displaystyle x^{2} + 3 x + 2$

sin(x + y).expand(trig=True)

$\displaystyle \sin{\left(x \right)} \cos{\left(y \right)} + \sin{\left(y \right)} \cos{\left(x \right)}$

log(x * y).expand(log=True)

$\displaystyle \log{\left(x y \right)}$

a, b = symbols('a, b', positive=True)
log(a*b).expand(log=True)

$\displaystyle \log{\left(a \right)} + \log{\left(b \right)}$

expr = exp(a + I*b)
expr

$\displaystyle e^{a + i b}$

expr.expand(complex=True)

$\displaystyle i e^{a} \sin{\left(b \right)} + e^{a} \cos{\left(b \right)}$

Factor, Collect and Combine

expr = x**2 - 1
factor(expr)

$\displaystyle \left(x - 1\right) \left(x + 1\right)$

expr = x*cos(y) + x*sin(z)
factor(expr)

$\displaystyle x \left(\sin{\left(z \right)} + \cos{\left(y \right)}\right)$

expr = log(a) - log(b)
logcombine(expr)

$\displaystyle \log{\left(\frac{a}{b} \right)}$

expr = x + y + x*z + x*y
expr.collect(x)

$\displaystyle x \left(y + z + 1\right) + y$

expr.collect(y)

$\displaystyle x z + x + y \left(x + 1\right)$

Apart, Together and Cancel

apart(1/(x**2 + 3*x + 2), x)

$\displaystyle - \frac{1}{x + 2} + \frac{1}{x + 1}$

together(1 / (y * x + y) + 1 / (1+x))

$\displaystyle \frac{y + 1}{y \left(x + 1\right)}$

cancel(y / (y * x + y))

$\displaystyle \frac{1}{x + 1}$

Substitutions

(x + y).subs(x, y)

$\displaystyle 2 y$

expr = sin(x*exp(x))
expr

$\displaystyle \sin{\left(x e^{x} \right)}$

expr.subs(x, y)

$\displaystyle \sin{\left(y e^{y} \right)}$

For muliple substitutions, we can pass a dictionary as first and only argument to subs, which maps old symbols or expressions to new symbols or expressions:

expr = sin(x * z)
expr

$\displaystyle \sin{\left(x z \right)}$

expr.subs({sin:cos, x:y, z:exp(x)})

$\displaystyle \cos{\left(y e^{x} \right)}$

expr = x*y + z**2* x

To substitute numerical values in place of symbolic number, for numerical evaluation, a convenient way of doing this is to define a dictionary that translates the symbols to numerical values, and passing this dictionary as argument to the subs method.

expr

$\displaystyle x y + x z^{2}$

values = {x:1.25,
         y:0.04,
         z:3.2}

expr.subs(values)

$\displaystyle 12.85$

Numerical Evaluation

Even when working with symbolic mathematics, it is almost invariably sooner or later required to evaluate the symbolic expressions numerically, for example, when producing plots or concrete numerical results. A SymPy expression can be evaluated using either the sympy.N function, or the evalf method of SymPy expression instances.
Both sympy.N and the evalf method take an optional argument that specifies the number of significant digits to which the expression is to be evaluated.

N(1 + pi)

$\displaystyle 4.14159265358979$

N(pi, 10)

$\displaystyle 3.141592654$

(x + 1/pi).evalf()

$\displaystyle x + 0.318309886183791$

(x + 1/pi).evalf(5)

$\displaystyle x + 0.31831$

sympy.lambdify function takes a set of free symbols and an expression as arguments, and generates a function that efficiently evaluates the numerical value of the expression. The produced function takes the same number of arguments as the number of free symbols passed as first argument to sympy.lambdify.

x  = symbols('x')
expr = sin(pi*x*exp(x))
expr

$\displaystyle \sin{\left(\pi x e^{x} \right)}$

expr_num = lambdify(x, expr)
expr_num(10)

0.879393997597802

By passing the optional argument 'numpy' as third argument to sympy.lambdify SymPy creates a vectorized function that accepts NumPy arrays as input.

expr_num = lambdify(x, expr, 'numpy')
import numpy as np
xvalues = np.arange(1, 10)
expr_num(xvalues)

array([ 0.77394269,  0.64198244,  0.72163867,  0.94361635,  0.20523391,
        0.97398794,  0.97734066, -0.87034418, -0.69512687])

Calculas

Derivatives

In SymPy we can calculate the derivative of a function using sympy.diff

x = Symbol('x')
f = Function('f')(x)
diff(f, x)

$\displaystyle \frac{d}{d x} f{\left(x \right)}$

diff(f, x, x)

$\displaystyle \frac{d^{2}}{d x^{2}} f{\left(x \right)}$

diff(f, x, 3)

$\displaystyle \frac{d^{3}}{d x^{3}} f{\left(x \right)}$

For multivariate functions -

x, y, z = symbols('x, y, z')
g = Function('g')(x, y)
g.diff(x, y)

$\displaystyle \frac{\partial^{2}}{\partial y\partial x} g{\left(x,y \right)}$

g.diff(x, 3, y, 2)

$\displaystyle \frac{\partial^{5}}{\partial y^{2}\partial x^{3}} g{\left(x,y \right)}$

For defined functions -

expr = x**4 + x**3 + x**2 + x + 1
expr.diff(x)

$\displaystyle 4 x^{3} + 3 x^{2} + 2 x + 1$

expr.diff(x, x)

$\displaystyle 2 \left(6 x^{2} + 3 x + 1\right)$

expr = (x + 1)**3 * y**2 *(z - 1)
expr.diff(x, y, z)

$\displaystyle 6 y \left(x + 1\right)^{2}$

For trigonometric functions -

expr = sin(x * y) * cos(x / 2)
expr.diff(x)

$\displaystyle y \cos{\left(\frac{x}{2} \right)} \cos{\left(x y \right)} - \frac{\sin{\left(\frac{x}{2} \right)} \sin{\left(x y \right)}}{2}$

Alternative way (delayed evaluation) -

expr = exp(cos(x))
d = Derivative(expr, x)
d

$\displaystyle \frac{d}{d x} e^{\cos{\left(x \right)}}$

d.doit()

$\displaystyle - e^{\cos{\left(x \right)}} \sin{\left(x \right)}$

Integrals

a, b, x, y = symbols('a, b, x, y')
f = Function('f')(x)
integrate(f)

$\displaystyle \int f{\left(x \right)}\, dx$

integrate(f, (x, a, b))

$\displaystyle \int\limits_{a}^{b} f{\left(x \right)}\, dx$

integrate(sin(x))

$\displaystyle - \cos{\left(x \right)}$

integrate(sin(x), (x, a, b))

$\displaystyle \cos{\left(a \right)} - \cos{\left(b \right)}$

integrate(exp(-x**2), (x, 0, oo))

$\displaystyle \frac{\sqrt{\pi}}{2}$

SymPy will not be able to give symbolic results for any integral. When SymPy fails to evaluate an integral, an instance of sympy.Integral, representing the formal integral, is returned instead.

integrate(sin(x*cos(x)), x)

$\displaystyle \int \sin{\left(x \cos{\left(x \right)} \right)}\, dx$

Multibariable expressions -

integrate(sin(x*exp(y)), x)

$\displaystyle - e^{- y} \cos{\left(x e^{y} \right)}$

expr = (x + y)**2
integrate(expr, x)

$\displaystyle \frac{x^{3}}{3} + x^{2} y + x y^{2}$

By passing more than one symbol, or more than one tuple that contain symbols and their integration limits, we can carry out multiple integration:

expr = (x + y)**2
integrate(expr, (x, 0, 1), (y, 0, 1))

$\displaystyle \frac{7}{6}$

Series

x = Symbol('x')
f = Function('f')(x)
series(f, x)

$\displaystyle f{\left(0 \right)} + x \left. \frac{d}{d x} f{\left(x \right)} \right|_{\substack{ x=0 }} + \frac{x^{2} \left. \frac{d^{2}}{d x^{2}} f{\left(x \right)} \right|_{\substack{ x=0 }}}{2} + \frac{x^{3} \left. \frac{d^{3}}{d x^{3}} f{\left(x \right)} \right|_{\substack{ x=0 }}}{6} + \frac{x^{4} \left. \frac{d^{4}}{d x^{4}} f{\left(x \right)} \right|_{\substack{ x=0 }}}{24} + \frac{x^{5} \left. \frac{d^{5}}{d x^{5}} f{\left(x \right)} \right|_{\substack{ x=0 }}}{120} + O\left(x^{6}\right)$

To change the point around which the function is expanded, we specify $x_0$ argument -

x0 = Symbol('{x_0}')
f.series(x, x0, n=2)

$\displaystyle f{\left({x_0} \right)} + \left(x - {x_0}\right) \left. \frac{d}{d \xi_{1}} f{\left(\xi_{1} \right)} \right|_{\substack{ \xi_{1}={x_0} }} + O\left(\left(x - {x_0}\right)^{2}; x\rightarrow {x_0}\right)$

f.series(x, x0, n=2).removeO()

$\displaystyle \left(x - {x_0}\right) \left. \frac{d}{d \xi_{1}} f{\left(\xi_{1} \right)} \right|_{\substack{ \xi_{1}={x_0} }} + f{\left({x_0} \right)}$

For specified functions -

sin(x).series()

$\displaystyle x - \frac{x^{3}}{6} + \frac{x^{5}}{120} + O\left(x^{6}\right)$

cos(x).series()

$\displaystyle 1 - \frac{x^{2}}{2} + \frac{x^{4}}{24} + O\left(x^{6}\right)$

exp(x).series()

$\displaystyle 1 + x + \frac{x^{2}}{2} + \frac{x^{3}}{6} + \frac{x^{4}}{24} + \frac{x^{5}}{120} + O\left(x^{6}\right)$

(1/(1 + x)).series()

$\displaystyle 1 - x + x^{2} - x^{3} + x^{4} - x^{5} + O\left(x^{6}\right)$

For arbitrary expressions of symbols and functions, which in general can also be evaluated.

expr = cos(x)/(1 + sin(x*y))
expr.series(x, n=4)

$\displaystyle 1 - x y + x^{2} \left(y^{2} - \frac{1}{2}\right) + x^{3} \left(- \frac{5 y^{3}}{6} + \frac{y}{2}\right) + O\left(x^{4}\right)$

expr.series(y, n=4)

$\displaystyle \cos{\left(x \right)} - x y \cos{\left(x \right)} + x^{2} y^{2} \cos{\left(x \right)} - \frac{5 x^{3} y^{3} \cos{\left(x \right)}}{6} + O\left(y^{4}\right)$

Limits

In SymPy, limits can be evaluated using the sympy.limit function, which takes an expression, a symbol it depends on, as well as the value that the symbol approaches in the limit.

limit(sin(x)/x, x, 0)

$\displaystyle 1$

f = Function('f')
x, h = symbols('x, h')
diff_limit = (f(x + h) - f(x))/h
limit(diff_limit.subs(f, cos), h, 0)

$\displaystyle - \sin{\left(x \right)}$

limit(diff_limit.subs(f, sin), h, 0)

$\displaystyle \cos{\left(x \right)}$

Sums and Products

n = symbols('n', integer=True)
x = Sum(1/n**2, (n, 1, oo))
x

$\displaystyle \sum_{n=1}^{\infty} \frac{1}{n^{2}}$

x.doit()

$\displaystyle \frac{\pi^{2}}{6}$

x = Product(n, (n, 1, 7))
x

$\displaystyle \prod_{n=1}^{7} n$

x.doit()

$\displaystyle 5040$

Equations

x = Symbol('x')
solve(x**2 + 2*x -3)

[-3, 1]

a, b, c, x = symbols('a, b, c, x')
solve(a*x**2 + b*x + c, x)

[(-b + sqrt(-4*a*c + b**2))/(2*a), -(b + sqrt(-4*a*c + b**2))/(2*a)]

Trigonometric Functions -

solve(sin(x) - cos(x), x)

[-3*pi/4, pi/4]

Solving a system of equations for more than one unknown variable in SymPy is a straightforward generalization of the procedure used for univariate equations. Instead of passing a single expression as first argument to sympy.solve, a list of expressions that represent the system of equations is used, and in this case the second argument should be a list of symbols to solve for.

eq1 = x + 2*y -1
eq2 = x - y + 1

solve([eq1, eq2], [x, y], dict=True)

[{x: -1/3, y: 2/3}]

eq1 = x**2 - y
eq2 = y**2 - x

sols = solve([eq1, eq2], [x, y], dict=True)
sols

[{x: 0, y: 0},
 {x: 1, y: 1},
 {x: (-1/2 - sqrt(3)*I/2)**2, y: -1/2 - sqrt(3)*I/2},
 {x: (-1/2 + sqrt(3)*I/2)**2, y: -1/2 + sqrt(3)*I/2}]

Verification -

[eq1.subs(sol).simplify() == 0 and eq2.subs(sol).simplify() == 0 for sol in sols]

[True, True, True, True]

Linear Algebra

Matrix([1, 2])

$\displaystyle \left[\begin{matrix}1\\2\end{matrix}\right]$

Matrix([ [1, 2 ] ])

$\displaystyle \left[\begin{matrix}1 & 2\end{matrix}\right]$

Matrix( [[1, 2], [3, 4] ])

$\displaystyle \left[\begin{matrix}1 & 2\\3 & 4\end{matrix}\right]$

Matrix(3, 4, lambda m, n: 10*m + n)

$\displaystyle \left[\begin{matrix}0 & 1 & 2 & 3\\10 & 11 & 12 & 13\\20 & 21 & 22 & 23\end{matrix}\right]$

a, b, c, d = symbols('a, b, c, d')
M = Matrix([[a, b], [c, d] ] )
M

$\displaystyle \left[\begin{matrix}a & b\\c & d\end{matrix}\right]$

M * M

$\displaystyle \left[\begin{matrix}a^{2} + b c & a b + b d\\a c + c d & b c + d^{2}\end{matrix}\right]$

x = Matrix(symbols("x_1, x_2"))
x

$\displaystyle \left[\begin{matrix}x_{1}\\x_{2}\end{matrix}\right]$

M * x

$\displaystyle \left[\begin{matrix}a x_{1} + b x_{2}\\c x_{1} + d x_{2}\end{matrix}\right]$

Practical Examples

1. Newton's Law of Cooling

By Analytical Way using SymPy

$$\frac{dT(t)}{dt} = k(T(t) - T_a)$$ at $t=0$, $T(0)=T_0$.

from sympy import *

T, Ta, T0, k, t = symbols('T, T_a, T_0, k, t')
T  = Function('T')
diff_eq = Eq(T(t).diff(t),-k*(T(t) - Ta))
diff_eq

$\displaystyle \frac{d}{d t} T{\left(t \right)} = - k \left(- T_{a} + T{\left(t \right)}\right)$

sol = dsolve(diff_eq)
sol

$\displaystyle T{\left(t \right)} = C_{1} e^{- k t} + T_{a}$

ic = Eq(sol.rhs.subs(t, 0), T0)
ic

$\displaystyle C_{1} + T_{a} = T_{0}$

const = solve(ic)
const

[{C1: T_0 - T_a}]

final_sol = sol.subs(const[0])
final_sol

$\displaystyle T{\left(t \right)} = T_{a} + \left(T_{0} - T_{a}\right) e^{- k t}$

T =lambdify([Ta, T0, k, t], final_sol.rhs)
T

<function _lambdifygenerated(T_a, T_0, k, t)>

import numpy as np
import matplotlib.pyplot as plt

k = 0.5
Ta = 22
T0 = 100
t = np.linspace(0, 10, 100)

fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(t, T(Ta, T0, k, t), color='red', label="""$T_a$ = %0.1f,
$T_0$ = %0.1f,
 k  = %0.2f"""%(Ta, T0, k))
ax.set_xlabel('Time ($t$) in sec.', fontsize=16)
ax.set_ylabel('Temperature ($T$) in $^{\circ}C$', fontsize=16)
ax.set_title('Cooling Curve', fontsize=20)

# removing top and right spines
ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')

# move bottom and left spine to x = 0 and y = 0
ax.spines['bottom'].set_position(('data', 0))
ax.spines['left'].set_position(('data', 0))

ax.set_xticks(range(0, 11, 1))
ax.set_yticks(range(0, 110, 10))

ax.legend(frameon=False)
ax.grid(False)

2. Damped harmonic Oscillator

$$\frac{d^2 x(t)}{dt^2} + 2\gamma \omega_0 \frac{d x(t)}{dt} + \omega_0^2 x(t) = 0$$ where $x(t)$ is the position of the oscillator at time $t$, $\omega_0$ is the frequency of the oscillator in the undamped case and $\gamma$ is the damping ratio.

t, omega0, gamma = symbols('t, omega_0, gamma', positive=True)
x                = Function('x')

ode = x(t).diff(t, 2) + 2*gamma*omega0*x(t).diff(t) + omega0**2 *x(t)
ode

$\displaystyle 2 \gamma \omega_{0} \frac{d}{d t} x{\left(t \right)} + \omega_{0}^{2} x{\left(t \right)} + \frac{d^{2}}{d t^{2}} x{\left(t \right)}$

Eq(ode)

$\displaystyle 2 \gamma \omega_{0} \frac{d}{d t} x{\left(t \right)} + \omega_{0}^{2} x{\left(t \right)} + \frac{d^{2}}{d t^{2}} x{\left(t \right)} = 0$

ode_sol = dsolve(Eq(ode))
ode_sol

$\displaystyle x{\left(t \right)} = C_{1} e^{\omega_{0} t \left(- \gamma - \sqrt{\gamma - 1} \sqrt{\gamma + 1}\right)} + C_{2} e^{\omega_{0} t \left(- \gamma + \sqrt{\gamma - 1} \sqrt{\gamma + 1}\right)}$

Initial conditions:

1. at time $t=0$, displacement of the oscillator, $x(t) = 1$
2. at time $t=0$, velocity of the oscillator, $\frac{d x(t)}{dt} = 0$

ic1 = x(t).subs(t, 0)
ic2 = x(t).diff(t).subs(t, 0)
ics = {ic1:1, ic2:0}
ics

{x(0): 1, Subs(Derivative(x(t), t), t, 0): 0}

eq1 = Eq(ode_sol.rhs.subs(t, 0), 1)
eq2 = Eq(ode_sol.rhs.diff(t).subs(t, 0), 0)
constants = solve([eq1, eq2])
x_t_sol = ode_sol.subs(constants[0])
x_t_sol

$\displaystyle x{\left(t \right)} = \left(- \frac{\gamma}{2 \sqrt{\gamma^{2} - 1}} + \frac{1}{2}\right) e^{\omega_{0} t \left(- \gamma - \sqrt{\gamma - 1} \sqrt{\gamma + 1}\right)} + \left(\frac{\gamma}{2 \sqrt{\gamma^{2} - 1}} + \frac{1}{2}\right) e^{\omega_{0} t \left(- \gamma + \sqrt{\gamma - 1} \sqrt{\gamma + 1}\right)}$

x_t_critical = limit(x_t_sol.rhs, gamma, 1)
x_t_critical

$\displaystyle \left(\omega_{0} t + 1\right) e^{- \omega_{0} t}$

fig, ax = plt.subplots(figsize=(14, 8))
tt = np.linspace(0, 3, 300)
w0 = 2* np.pi
for g in [0.1, 0.5, 1, 2.0, 5.0]:
    if g == 1:
        x_t = lambdify(t, x_t_critical.subs({omega0:w0, gamma:g}), 'numpy')
    else:
        x_t = lambdify(t, x_t_sol.rhs.subs({omega0:w0, gamma:g}), 'numpy')
    ax.plot(tt, x_t(tt).real,  label="$\gamma = %.1f$" % g )
ax.set_xlabel(r"$t$", fontsize=18)
ax.set_ylabel(r"$x(t)$", fontsize=18)
ax.set_title('Damped Harmonic Oscillator', fontsize=20)

# removing top and right spines
ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')

# move bottom and left spine to x = 0 and y = 0
ax.spines['bottom'].set_position(('data', 0))
ax.spines['left'].set_position(('data', 0))

ax.set_xticks(np.arange(0, 4, 1))
ax.set_yticks(np.arange(-1, 1.5, 0.5))

ax.legend(frameon=False)
ax.grid(False)
plt.legend(frameon=False, loc='upper right')
plt.grid(False)