Pythonで数学ガールのリサみたいにベクトルを描画したい

最近少しずつですが、集合論や線型代数など大学数学を基礎から勉強し直しています。

そんな折に、結城浩さんの「数学ガールの秘密ノート/行列が描くもの」が刊行されました。 数学ガールシリーズは、高校生の「僕」と魅力的な女の子たちが一緒に数学の世界を旅する物語です（ちくしょう！うらやましいぞ！）

数学ガールの秘密ノート/行列が描くもの

上記「行列が描くもの」には「リサ」というコンピュータの扱いに長けたクールビューティな女の子が登場します。

物語の途中に、リサがコンピュータを使ってベクトルや線型変換を可視化してみせる重要なシーンがあるのですが、Pythonを使って同じことができないかやってみました。

行列演算にはnumpyを、グラフ描画にはmatplotlibを使用します。このmatplotlibですが、ベクトル描画にうってつけのquiver()が実装されています。

matplotlib.pyplot.quiver — Matplotlib 3.0.2 documentation

algorithm.joho.info

試しにベクトル[2, 1]と[1, 3]の和[3, 4]をそれぞれ赤、緑、青でプロットするコードを書いてみました。

# coding:utf-8

import numpy as np
import matplotlib.pyplot as plt

def main():

    a = np.array([2, 1])
    b = np.array([1, 3])
    c = a + b
    origin = [0, 0, 0] # 原点(0, 0)

    d = np.array([a, b, c])
    U = d.T[0] # 転置を取ってx成分取り出し
    V = d.T[1] # 転置を取ってy成分取り出し
    
    plt.quiver(origin, origin, U, V, angles='xy', scale=1, scale_units='xy', color=['r', 'g', 'b'])

    plt.xlim([-1, 5])
    plt.ylim([-1, 5])

    plt.grid()
    plt.show()

if __name__ == '__main__':
    main()

結果がこちら。期待どおりのグラフが出来上がりました。

次は、{(x, y) | x, y ∈ Z, -5 ≦ x ≦ 5, -5 ≦ y ≦ 5}の各格子点が行列array([2, 0], [0, 2])でどう移動するかプロットしてみます。 「行列が描くもの」P.147の再現です。

# coding:utf-8

import numpy as np
import matplotlib.pyplot as plt

def main():

    # 格子点を生成
    l = []
    for x in range(-5, 6):
        for y in range(-5, 6):
            l.append([x, y])
    vectors = np.array(l)

    A = np.array([[2, 0], [0, 2]])
    B = np.dot(A, vectors.T)

    U = B[0] # 転置を取ってx成分取り出し
    V = B[1] # 転置を取ってy成分取り出し
    
    plt.quiver(vectors.T[0], vectors.T[1], U, V, angles='xy', scale=1, scale_units='xy')

    plt.xlim([-20, 20])
    plt.ylim([-20, 20])

    plt.grid()
    plt.show()

if __name__ == '__main__':
    main()

変換前の格子点を生成するために愚直にforループを回してしまいましたが、もうちょっとスマートな書き方ができないかな…。

（追記）普通にリスト内包表記で書けました。

[[x, y] for x in range(-5, 6) for y in range(-5, 6)] 

実行結果がこちら。

無事にPythonでリサみたいにグラフを描くことができました。

それにしても、リサは秒でこの程度のコードを書いていたわけで…。

リサ…恐ろしい子！