Python에서 벡터 그리기 drawRobotics Ver0.01

최근 저는 Python에서 3D 공간상에서 벡터나 좌표계를 표현하는 걸 이야기한 적이 있는데요.[바로가기] 그게 뭐라고... 거기서 좀 더 수정을 했습니다. 뭐 이딴 기초적인걸 왜 자꾸 올리는거지??? 라고 생각하지 마세요.. 그냥~~~ 취미랍니다. ㅎ^^ 아무튼.. 아주 정확하게 수정된 사항은 GitHub에 있습니다.

drawRobotics ver0.01 변경사항 확인하기

그걸 좀 자세히 보면

먼저 [바로가기]에서도 이야기한 인터넷에서 주워왔다는 화살표 그리는 아이를 별개로 분리했습니다. 그래서 그 코드만 따로

<Arrow3D.py>

from matplotlib.patches import FancyArrowPatch
from mpl_toolkits.mplot3d import proj3d

class Arrow3D(FancyArrowPatch):
    def __init__(self, xs, ys, zs, *args, **kwargs):
        FancyArrowPatch.__init__(self, (0,0), (0,0), *args, **kwargs)
        self._verts3d = xs, ys, zs

    def draw(self, renderer):
        xs3d, ys3d, zs3d = self._verts3d
        xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
        self.set_positions((xs[0],ys[0]),(xs[1],ys[1]))
        FancyArrowPatch.draw(self, renderer)

이렇게 분리해 두구요.

그 다음은 사소하게 라인스타일을 지정할 수 있도록 하고, 3차원 행렬이든 4차이든지 다 x,y,z 성분을 잡도록 수정했습니다.

그리고, annotation을 추가했습니다. 이건 맨 마지막에 예제를 보시면 뭔지 아실겁니다^^

그리고 이부분은 좀 고민하다가 기준점을 원점으로 방향벡터를 그리도록 수정했습니다. 이 부분 코드의 형태가 마음에 안드는데 혹시 지나가던 고수님의 조언을 기다립니다.^^ 이렇게 해서 전체 코드는 아래와 같습니다.^^

<drawRobotics.py>

import matplotlib.pyplot as plt
import numpy as np
from Arrow3D import *

conv2Rad = lambda x : x*np.pi/180

def drawVector(fig, pointA, pointB, **kwargs):
	ms = kwargs.get('mutation_scale', 20)
	ars = kwargs.get('arrowstyle', '-|>')
	lc = kwargs.get('lineColor', 'k')
	pc = kwargs.get('projColor', 'k')
	pointEnable = kwargs.get('pointEnable', True)
	projOn = kwargs.get('proj', True)
	lineStyle = kwargs.get('lineStyle', '-')
	annotationString = kwargs.get('annotationString', '')

	if (3 <= pointA.size <= 4):
		xs = [pointA[0], pointB[0]]
		ys = [pointA[1], pointB[1]]
		zs = [pointA[2], pointB[2]]
	else:
		xs = [pointA[0,3], pointB[0,3]]
		ys = [pointA[1,3], pointB[1,3]]
		zs = [pointA[2,3], pointB[2,3]]

	out = Arrow3D(xs, ys, zs, mutation_scale=ms, arrowstyle=ars, color=lc, linestyle=lineStyle)
	fig.add_artist(out)

	if pointEnable: fig.scatter(xs[1], ys[1], zs[1], color='k', s=50)

	if projOn:
		fig.plot(xs, ys, [0, 0], color=pc, linestyle='--')
		fig.plot([xs[0], xs[0]], [ys[0], ys[0]], [0, zs[0]], color=pc, linestyle='--')
		fig.plot([xs[1], xs[1]], [ys[1], ys[1]], [0, zs[1]], color=pc, linestyle='--')

	if annotationString != '':
		fig.text(xs[1], ys[1], zs[1], annotationString, size=15, zorder=1, color='k') 
		
def drawPointWithAxis(fig, *args, **kwargs):
	ms = kwargs.get('mutation_scale', 20)
	ars = kwargs.get('arrowstyle', '->')
	pointEnable = kwargs.get('pointEnable', True)
	axisEnable = kwargs.get('axisEnable', True)
	lineStyle = kwargs.get('lineStyle', '-')
	vectorLength = kwargs.get('vectorLength', 1)

	if len(args) == 4:
		ORG = args[0]
		hat_X = args[1]
		hat_Y = args[2]
		hat_Z = args[3]
		xs_n = [ORG[0], ORG[0] + hat_X[0]*vectorLength]
		ys_n = [ORG[1], ORG[1] + hat_X[1]*vectorLength]
		zs_n = [ORG[2], ORG[2] + hat_X[2]*vectorLength]
		xs_o = [ORG[0], ORG[0] + hat_Y[0]*vectorLength]
		ys_o = [ORG[1], ORG[1] + hat_Y[1]*vectorLength]
		zs_o = [ORG[2], ORG[2] + hat_Y[2]*vectorLength]
		xs_a = [ORG[0], ORG[0] + hat_Z[0]*vectorLength]
		ys_a = [ORG[1], ORG[1] + hat_Z[1]*vectorLength]
		zs_a = [ORG[2], ORG[2] + hat_Z[2]*vectorLength]
	else:
		tmp = args[0]
		ORG = tmp[:3,3:]
		hat_X = tmp[:3,0:1]
		hat_Y = tmp[:3,1:2]
		hat_Z = tmp[:3,2:3]
		xs_n = [ORG[0, 0], ORG[0, 0] + hat_X[0, 0]*vectorLength]
		ys_n = [ORG[1, 0], ORG[1, 0] + hat_X[1, 0]*vectorLength]
		zs_n = [ORG[2, 0], ORG[2, 0] + hat_X[2, 0]*vectorLength]
		xs_o = [ORG[0, 0], ORG[0, 0] + hat_Y[0, 0]*vectorLength]
		ys_o = [ORG[1, 0], ORG[1, 0] + hat_Y[1, 0]*vectorLength]
		zs_o = [ORG[2, 0], ORG[2, 0] + hat_Y[2, 0]*vectorLength]
		xs_a = [ORG[0, 0], ORG[0, 0] + hat_Z[0, 0]*vectorLength]
		ys_a = [ORG[1, 0], ORG[1, 0] + hat_Z[1, 0]*vectorLength]
		zs_a = [ORG[2, 0], ORG[2, 0] + hat_Z[2, 0]*vectorLength]

	if pointEnable: fig.scatter(xs_n[0], ys_n[0], zs_n[0], color='k', s=50)

	if axisEnable:
		n = Arrow3D(xs_n, ys_n, zs_n, mutation_scale=ms, arrowstyle=ars, color='r', linestyle=lineStyle)
		o = Arrow3D(xs_o, ys_o, zs_o, mutation_scale=ms, arrowstyle=ars, color='g', linestyle=lineStyle)
		a = Arrow3D(xs_a, ys_a, zs_a, mutation_scale=ms, arrowstyle=ars, color='b', linestyle=lineStyle)
		fig.add_artist(n)
		fig.add_artist(o)
		fig.add_artist(a)

def RotX(phi):
    return np.array([[1, 0, 			0],
                     [0, np.cos(phi), 	-np.sin(phi)],
                     [0, np.sin(phi), 	np.cos(phi)]])

def RotY(theta):
    return np.array([[np.cos(theta), 	0, np.sin(theta)],
                     [0, 				1, 0],
                     [-np.sin(theta), 	0, np.cos(theta)]])

def RotZ(psi):
    return np.array([[np.cos(psi), 	-np.sin(psi), 	0],
                     [np.sin(psi), 	np.cos(psi), 	0],
                     [0, 			0, 				1]])

이제 예제 하나 그려보죠^^. 

# Introduction to Robotics 3rd Edition by Craig
# Example 2-1.
# Figure 2.6 shows a frame {B} that is rotated relative to frame {A} about Z by 30 degrees.

from matplotlib import pyplot as plt
import drawRobotics as dR
import numpy as np

P_atB = np.array([0,2,0])

Rot_AtoB = dR.RotZ(dR.conv2Rad(30))

P_atA = np.dot(Rot_AtoB, P_atB)

print('P_atB = ', P_atB)
print('Rot_AtoB = ', Rot_AtoB)
print('P_atA = Rot_AtoB * P_atA = ', P_atA)

AORG = np.array([0,0,0])
hat_X_atA = np.array([1,0,0])
hat_Y_atA = np.array([0,1,0])
hat_Z_atA = np.array([0,0,1])

BORG = np.dot(Rot_AtoB, AORG)
hat_X_atB = np.dot(Rot_AtoB, hat_X_atA)
hat_Y_atB = np.dot(Rot_AtoB, hat_Y_atA)
hat_Z_atB = np.dot(Rot_AtoB, hat_Z_atA)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

dR.drawPointWithAxis(ax, AORG, hat_X_atA, hat_Y_atA, hat_Z_atA, pointEnable=False)
dR.drawPointWithAxis(ax, BORG, hat_X_atB, hat_Y_atB, hat_Z_atB, pointEnable=False, lineStyle='--')
dR.drawVector(ax, AORG, P_atA, arrowstyle='-|>', annotationString=' $ ^{A}P $ ')

ax.set_xlim([-1.5,1.5]), ax.set_ylim([-0.1,2]), ax.set_zlim([-0.1,1.5])
ax.set_xlabel('X axis'), ax.set_ylabel('Y axis'), ax.set_zlabel('Z axis')
ax.view_init(azim=-90, elev=90)
plt.show()

위 예제는 Craig의 로보틱스 책의 2-1번 예제를 그려본 것입니다. 하단부 fig = 부터 보면... 아주 쉽게...

dR.drawPointWithAxis(ax, AORG, hat_X_atA, hat_Y_atA, hat_Z_atA, pointEnable=False)
dR.drawPointWithAxis(ax, BORG, hat_X_atB, hat_Y_atB, hat_Z_atB, pointEnable=False, lineStyle='--')
dR.drawVector(ax, AORG, P_atA, arrowstyle='-|>', annotationString=' $ ^{A}P $ ')

이 세 줄이 중요한 부분인데요. 좌표계를 두 개 그린거고 , 벡터를 하나 그린겁니다. 그리고, 좌표계중에서 {B}계는 점선으로 그렸고, 벡터에는 annotation을 표시했습니다.

그 결과입니다. view_init을 적용해서 x-y평면만 보이도록 했는데요. 그건 Craig의 2-1예제가 저렇게 책에 그림이 있어서 비교해 볼려고 한거구요^^

괜찮죠?? ㅎㅎ 이제 각 종 로보틱스 예제를 쉽게 눈으로 확인해 볼 수 있을 것 같습니다.%^^ 이상 아무도 필요없어 하는 이상한 함수하나를 업데이트한 내용을 올리는 PinkWink였습니다.^^

GitHub : drawRobotics 바로가기