SymPyで辛い計算をちょっと辛い計算にする

この記事はPython Advent Calendar 2015 - Adventarの7日目の記事です。

SymPyというシンボル計算ライブラリを使ったことがありますか？最近、卒論のお手伝いをしてもらっているライブラリです。日本語の情報が少なくて困ったので自分用にという意味も込めて使い方を書きます。

インストール

Anacondaを使います。大人しくAnacondaを使います。

シンボル計算ライブラリ SymPy

シンボル計算ライブラリというのは「X」とか「Y」とか多項式の計算をしてくれるライブラリです。例えば ${Z = X \left(2 X + 4\right) + 5 Y - 7}$ なんて式が与えられたとして ${Z^{2}}$ を計算する必要があるとしましょう。 ${Z^{2} = \left(X \left(2 X + 4\right) + 5 Y - 7\right)^{2} = 4 X^{4} + 16 X^{3} + 20 X^{2} Y - 12 X^{2} + 40 X Y - 56 X + 25 Y^{2} - 70 Y + 49}$ となるのですが、こんな計算は高校の宿題までにしたいところで、SymPyの出番になるわけです。

使う記号を宣言して、計算式を入力するだけで計算結果を表示してくれるなんて、なんて素敵なんだ！

import sympy

sympy.init_printing()
# いい感じに出力するよう設定してくれるらしい

X = sympy.Symbol('X')
Y = sympy.Symbol('Y')
# 使用する記号を宣言して、変数に代入
# 左辺の変数Xと右辺の引数の文字Xは異なる文字列でも可
# 混乱を防ぐために同じ文字を使っている。

Z = 2 * (X + 2) * X + 5 * Y - 7

print(Z)
# => X*(2*X + 4) + 5*Y - 7

ZZ = Z * Z
print(ZZ)
# => (X*(2*X + 4) + 5*Y - 7)**2
print(sympy.simplify(ZZ))
# => (X*(2*X + 4) + 5*Y - 7)**2
print(sympy.expand(ZZ))
# => 4 X^{4} + 16 X^{3} + 20 X^{2} Y - 12 X^{2} + 40 X Y - 56 X + 25 Y^{2} - 70 Y + 49

代入　subs

多項式を計算したら、数値を代入したくなるのが人情というもの。 substituteのsubsだと思われるsubs関数を使って値を知ることが出来ます。引数はタプルのリストで与え、各タプルの第一要素は何「に」代入するか、第二要素は何「を」代入するかを与えます。

print(Z.subs([(X, 1), (Y, 2)]))
# => 9
print(ZZ.subs([(X, 1), (Y, 2)]))
# => 81
print(ZZ.subs([(Z, X)]))
# => X**2

latex

さらに便利なのがsympy.latexという関数。この子を使えば添字も分数もバシッと出力してくれます。

print(sympy.latex(Z))
# => X \left(2 X + 4\right) + 5 Y - 7

print(sympy.latex(sympy.expand(ZZ)))
# => 4 X^{4} + 16 X^{3} + 20 X^{2} Y - 12 X^{2} + 40 X Y - 56 X + 25 Y^{2} - 70 Y + 49

print(sympy.latex(ZZ/(Z+1)))
# => \frac{\left(X \left(2 X + 4\right) + 5 Y - 7\right)^{2}}{X \left(2 X + 4\right) + 5 Y - 6}

手打ちしたくない数式も出力してくれます。 ${\frac{Z^{2}}{Z+1} = \frac{\left(X \left(2 X + 4\right) + 5 Y - 7\right)^{2}}{X \left(2 X + 4\right) + 5 Y - 6}}$