Sympy - plotting the result of a matrix calculation

Viewed 126

In my program, I'm using sympy to get an analytical expression for a set of constants (r, a, b and tau) then plot the graph of r as a function of l. Firstly, I define my constants as such:

import sympy
import numpy
import matplotlib.pyplot as plt

[z1, z2, z3, k2] = [1, 2, 4, 2 * numpy.pi]
l = sympy.Symbol('l')
r = sympy.Symbol('r')
tau = sympy.Symbol('tau')
base = numpy.linspace(0, 10, 10000)

I then define and calculate my matrices (where M2 is the transpose of [r, a, b, tau]):

M1 = sympy.Matrix([[-1, 1, 1, 0],
               [1, z1/z2, - z1/z2, 0],
               [0, sympy.exp(- 1j * k2 * l), sympy.exp(1j * k2 * l), -1],
               [0, - sympy.exp(- 1j * k2 * l), sympy.exp(1j * k2 * l), z2/z3],
               ])


M3 = sympy.Matrix([[1], [1], [0], [0]])

M2 = (M1 ** -1) * M3

I then try to plot this using:

sympy.plot(M2[0].subs(l, base), base)

Which results in the following error:

Traceback (most recent call last):
  File "<input>", line 1, in <module>
  File "C:\Program Files\JetBrains\PyCharm 2020.2.3\plugins\python\helpers\pydev    \_pydev_bundle\pydev_umd.py", line 197, in runfile
pydev_imports.execfile(filename, global_vars, local_vars)  # execute the script
  File "C:\Program Files\JetBrains\PyCharm 2020.2.3\plugins\python\helpers\pydev\_pydev_imps\_pydev_execfile.py", line 18, in execfile
exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "matrix solver.py", line 47, in <module>
sympy.plot(M2[0].subs(l, base), base)
  File "AppData\Local\Programs\Python\Python39\lib\site-packages\sympy\plotting\plot.py", line 1641, in plot
plot_expr = check_arguments(args, 1, 1)
  File "AppData\Local\Programs\Python\Python39\lib\site-packages\sympy\plotting\plot.py", line 2302, in check_arguments
assert all(isinstance(e, Expr) for expr in exprs for e in expr)
AssertionError
1 Answers

(The post is updated to include @OscarBenjamin's insightful comment.)

.subs(l, base) doesn't work. In general, sympy and numpy are quite hard to mix. Sympy works best well all expressions contain sympy symbols, integers and fractions. Floats, by their imprecise nature, often make it very hard for sympy trying to achieve exact symbolic solutions. Further, due to how sympy is incorporated into Python, an annoying gotcha is Python converting an expression such 1/2 to 0.5 instead of to a sympy fraction.

Therefore, it is best to not import numpy in the part of the code that works with sympy. To avoid the problem with fractions, the tutorials suggest writing z1 = sympy.S(1). np.pi is a float (about 16 digits), while sympy.pi is the symbolic constant pi. Replacing 1j by sympy.I lets sympy do the calculations symbolically.

import sympy

[z1, z2, z3, k2] = [sympy.S(1), sympy.S(2), sympy.S(4), 2 * sympy.pi]
l = sympy.Symbol('l')
r = sympy.Symbol('r')
tau = sympy.Symbol('tau')

M1 = sympy.Matrix([[-1, 1, 1, 0],
                   [1, z1 / z2, - z1 / z2, 0],
                   [0, sympy.exp(- sympy.I * k2 * l), sympy.exp(sympy.I * k2 * l), -1],
                   [0, - sympy.exp(- sympy.I * k2 * l), sympy.exp(sympy.I * k2 * l), z2 / z3],
                   ])

M3 = sympy.Matrix([[1], [1], [0], [0]])
M2 = (M1 ** -1) * M3

As the result contains complex numbers, the sympy way of plotting would use sympy.re(M2[0]). Sympy's plotting default uses an adaptive strategy to decide on the detail level, which can be slow for a plot with a lot of fine detail.

sympy.plot(sympy.re(M2[0]), (l, 0, 10), adaptive=False)

sympy plot of real part

Similartly, the imaginary part could be plotted using sympy.im(M2[0]).

sympy plot of imaginary part

To work with numpy and matplotlib, lambdify can convert sympy expressions to numpy functions:

M2_0_np = sympy.lambdify(l, M2[0])

From then on, M2_0_np() can be used as a numpy function.

import numpy
import matplotlib.pyplot as plt

base = numpy.linspace(0, 10, 10000)
plt.plot(M2_0_np(base), base)

Plotting the function issues a warning "ComplexWarning: Casting complex values to real discards the imaginary part". So, only the real part of the function is plotted:

plotting a sympy function with matplotlib

More fanciful plots could be created combining real and imaginary part, e.g. via color.

Related