Matplotlibで図形を組み合わせた複雑な車両のような図形を描画する

Matplotlibではpatchesを使って四角形や円、楕円などを描画できる。しかし四角形の左下を原点として指定しないといけないため、回転させたり、複数の基本図形を組み合わせて図を描画するときに座標の計算が面倒になる。

:                +------------------+
:                |                  |
:              height               |
:                |                  |
:               (xy)---- width -----+

本記事ではmatplotlib.transformsとmatplotlib.patches.Rectangle を使って、以下のような車両の図形を描画する。

車体の中心の位置が0.5, 0.5の位置にあり、車体もある程度旋回していて、ステアリングが切られている状態の図になる。それぞれの四角形の原点の計算を独自に実装して描画するなんてことをしたらそれなりに面倒になることが予想できる。
そこで図形(patches)のtransform optionにTransform objectを指定してやることで図形に座標変換を適用し、扱いやすくする。

##　前提条件
以下のようにライブラリがimportされていることを前提とする。Jupyternotebookで記述しているので、matplotlib inlineの指定をしてある。

from math import radians

%matplotlib inline
import matplotlib.pyplot as plt
from matplotlib import patches
from matplotlib.patches import Rectangle
from matplotlib.transforms import Affine2D

描画対象のfigとaxesは以下の通り生成してあるとする。

fig, ax = plt.subplots()

axesの詳細については以下の記事が詳しい。

車体までの座標変換

図の中で、座標軸の原点までの座標変換は axis.transDataに格納されている。ここを基準のtransformとして各図形までのtransformを生成していく。transDataの詳細は the-transformation-pipeline を読むといいのだろうけど、筆者はあまりしっかり読み込んでいないのはご容赦いただきたい。
さて、車体の中心まで、[0.5, 0.5]、右手系で正の方向に20degree傾いた位置を表すtransformを作りたい。Affine2Dを用いて、この座標変換は以下のように計算できる。

to_body_center_tf = Affine2D().rotate(radians(20)).translate(0.5, 0.5) + ax.transData

ax.transData が一番右にあることが重要。直感とは逆に思えるかもしれないが、座標変換が右から左に適用されていく。
Affine2D().rotate(radians(20)).translate(0.5, 0.5) は Affine2D().rotate(radians(20)) + Affine2D().translate(0.5, 0.5) と同じで、とにかく右から左に座標変換が合成される。
次に、四角形を描画するために、四角形の中心から左下の角(描画原点)までのtransformを計算したい。ただ、このように中心を指定して四角形を描画するのは共通処理にしたいので、描画図形を生成するfunctionにする。

def relative_rectangle(w: float, h: float, center_tf, **kwargs):
  rect_origin_to_center = Affine2D().translate(w / 2, h / 2)
  return Rectangle(
      (0, 0), w, h, transform=rect_origin_to_center.inverted() + center_tf, **kwargs)

rect_origin_to_center が図形の左下から中心の変換であるから、その逆変換である中心から左下までの変換は rect_origin_to_center.inverted() で得られる。図形の中心までの変換 center_tf にこれを合成することで、座標系原点から四角形の描画原点までの変換 rect_origin_to_center.inverted() + center_tf を得ることができる。これを Rectangle の引数 transorm に渡してやれば、その位置に描画される。ここでは位置や回転はtransformで扱うため、位置には (0, 0) が指定してある (角度はデフォルトで0)。
この関数を使って任意の位置に任意の角度の四角形が描画できる。

to_body_center_tf = Affine2D().rotate(radians(20)).translate(0.5, 0.5) + ax.transData
body = relative_rectangle(0.2, 0.4, to_body_center_tf, edgecolor='black')
ax.add_patch(body)

plt.axis('auto')
ax.set_aspect('equal')
plt.show()

車体から車輪までの座標変換

以上で計算した、車体の中心までの座標変換を表す to_body_center_tfを基準にして4つの車輪までの座標変換を計算していく。まず、車両の中心から各車輪までの相対座標変換を用意する。

body_center_to_left_front = Affine2D().rotate(radians(15)).translate(-0.1, 0.15)
body_center_to_right_front = Affine2D().rotate(radians(20)).translate(0.1, 0.15)
body_center_to_left_rear = Affine2D().translate(-0.1, -0.15)
body_center_to_right_rear = Affine2D().translate(0.1, -0.15)

車両の中心から一定距離移動し、前輪についてはステアリングを表現する角度ぶんだけ回転している。
それぞれに、座標原点から車体の中心までの変換 to_body_center_tf を合成して図形の位置を指定することで各車輪を表す四角形を生成する。

left_front_wheel = relative_rectangle(
    0.05, 0.1, body_center_to_left_front + to_body_center_tf,
    facecolor='black', edgecolor='black')
right_front_wheel = relative_rectangle(
    0.05, 0.1, body_center_to_right_front + to_body_center_tf,
    facecolor='black', edgecolor='black')
left_rear_wheel = relative_rectangle(
    0.05, 0.1, body_center_to_left_rear + to_body_center_tf,
    facecolor='black', edgecolor='black')
right_rear_wheel = relative_rectangle(
    0.05, 0.1, body_center_to_right_rear + to_body_center_tf,
    facecolor='black', edgecolor='black')

##　まとめ
以上をまとめて以下のコードで車両のような図形を描画できる。

from math import radians

%matplotlib inline
import matplotlib.pyplot as plt
from matplotlib import patches
from matplotlib.patches import Rectangle
from matplotlib.transforms import Affine2D


fig, ax = plt.subplots()


def relative_rectangle(w: float, h: float, center_tf, **kwargs):
  rect_origin_to_center = Affine2D().translate(w / 2, h / 2)
  return Rectangle(
      (0, 0), w, h, transform=rect_origin_to_center.inverted() + center_tf, **kwargs)

to_body_center_tf = Affine2D().rotate(radians(20)).translate(0.5, 0.5) + ax.transData
body = relative_rectangle(0.2, 0.4, to_body_center_tf, edgecolor='black')

body_center_to_left_front = Affine2D().rotate(radians(15)).translate(-0.1, 0.15)
body_center_to_right_front = Affine2D().rotate(radians(20)).translate(0.1, 0.15)
body_center_to_left_rear = Affine2D().translate(-0.1, -0.15)
body_center_to_right_rear = Affine2D().translate(0.1, -0.15)

left_front_wheel = relative_rectangle(
    0.05, 0.1, body_center_to_left_front + to_body_center_tf,
    facecolor='black', edgecolor='black')
right_front_wheel = relative_rectangle(
    0.05, 0.1, body_center_to_right_front + to_body_center_tf,
    facecolor='black', edgecolor='black')
left_rear_wheel = relative_rectangle(
    0.05, 0.1, body_center_to_left_rear + to_body_center_tf,
    facecolor='black', edgecolor='black')
right_rear_wheel = relative_rectangle(
    0.05, 0.1, body_center_to_right_rear + to_body_center_tf,
    facecolor='black', edgecolor='black')

ax.add_patch(left_front_wheel)
ax.add_patch(right_front_wheel)
ax.add_patch(left_rear_wheel)
ax.add_patch(right_rear_wheel)
ax.add_patch(body)

plt.axis('auto')
ax.set_aspect('equal')