Python turtleでカーボンナノチューブを描きたい

Pythonにあるturtleモジュールでは描画を簡単に行うことができます。プログラムが正確に描画するので、絵心のない人でも安心して描けます。

高専の頃にTurtleで何か描いてみようと思い、お題を友人に求めたところ「カーボンナノチューブ」と帰ってきました。チューリップとかが良かったです。せっかくだからとそのときに書いたプログラムについて、以下に記していきたいと思います。
ちなみに、あくまでお絵かきなので精密な構造の計算はしていません。

↓こんな感じで描きあがります。

#描画プログラム
以下がプログラム全文です。

carbon_nano_tube.py

import numpy as np
import turtle as tr
import math

#２次元回転行列
def rotate_2d(x0,y0,deg):
    x=x0*np.cos(deg)-y0*np.sin(deg)
    y=x0*np.sin(deg)+y0*np.cos(deg)
    return x,y

#3次元回転行列
def rotate_3d(vec0, rol=None, yaw=None, pit=None):
    mat=np.array([[1.,0.,0.],[0.,1.,0.],[0.,0.,1.]])
    if(pit is not None):
        mat=np.dot(mat,np.array([[1.,0.,0.],[0.,np.cos(pit),-np.sin(pit)],[0.,np.sin(pit),np.cos(pit)]]))
    if(yaw is not None):
        mat=np.dot(mat,np.array([[np.cos(yaw),0.,np.sin(yaw)],[0.,1.,0.],[-np.sin(yaw),0.,np.cos(yaw)]]))
    if(rol is not None):
        mat=np.dot(mat,np.array([[np.cos(rol),-np.sin(rol),0.],[np.sin(rol),np.cos(rol),0.],[0.,0.,1.]]))
    vec=np.dot(mat,vec0)
    return vec

#与えられた半径の円を描く。同じ段の他の原子と比べて奥行きが奥なら青、手前なら赤で円を描く。
def maru(r,coordinate,now_i):
    hed=tr.heading()
    pos=tr.pos()
    ave=0
    warp(pos[0],pos[1]-r)
    tr.seth(0.0)
    for i in range(atoms):
        ave += coordinate[i][2]/atoms
    if(coordinate[now_i][2] > ave):
        tr.color('Blue')
    else:
        tr.color('Red')
    tr.circle(r)
    tr.color('Black')
    warp(pos[0],pos[1])
    tr.seth(hed)

#線を引かずに場所を移動する。
def warp(x,y):
    tr.pu()
    tr.goto(x,y)
    tr.pd()
    
############## parameters ############
#rol, yaw, pitの回転方向はそれぞれ画面手前、下、左への軸に対して右手系。
atoms =32           #チューブ断面当たりの炭素原子数
steps =50           #チューブ断面の段数
space=30             #炭素原子同士の間隔
rol=0               
yaw=np.pi/9
pit=-np.pi/12

######################################
splits=atoms/2*3        
tube=np.array([[[0.,0.,j*space/2*math.sqrt(3)] for i in range(atoms)] for j in range(steps)])
ix=0
iy=space/2/np.sin(2*np.pi/splits/2)
iz=0
tr.reset()
tr.speed(0)
screen = tr.Screen()
screen.delay(0)

#２次元回転行列を利用して、各段の炭素原子の座標を決定していく。
for j in range(steps):
    vec=np.array([ix,iy,iz])    #座標指定用の位置ベクトル
    if(j%2==0):
        tube[j][0] = tube[j][0] + vec
        for i in range(atoms-1):
            if(i%2==0):
                vec[0],vec[1] = rotate_2d(vec[0],vec[1],-2*np.pi/splits*2)           
            else:
                vec[0],vec[1] = rotate_2d(vec[0],vec[1],-2*np.pi/splits)
            tube[j][i+1] = tube[j][i+1] + vec
    else:
        vec[0],vec[1] = rotate_2d(vec[0],vec[1],-2*np.pi/splits/2)
        tube[j][0] = tube[j][0] + vec
        for i in range(atoms-1):
            if(i%2==0):
                vec[0],vec[1] = rotate_2d(vec[0],vec[1],-2*np.pi/splits)
            else:
                vec[0],vec[1] = rotate_2d(vec[0],vec[1],-2*np.pi/splits*2)
            tube[j][i+1] = tube[j][i+1] + vec

#設定された回転角に応じてそれぞれの炭素原子座標をを3次元回転。
for j in range(steps):
    for i in range(atoms):
        tube[j][i]=rotate_3d(tube[j][i],rol ,yaw, pit)

#炭素原子の位置で丸を描き、接続する炭素原子へ直線を描く。
for j in range(steps-1):
    for i in range(atoms):
        warp(tube[j][i][0],tube[j][i][1])
        maru(5,tube[j],i)
        tr.goto(tube[j+1][i][0],tube[j+1][i][1])
        warp(tube[j][i][0],tube[j][i][1])
        if(j%2==0):
            if(i%2==0):
                if(i==0):
                    tr.goto(tube[j][atoms-1][0],tube[j][atoms-1][1])
                else:
                    tr.goto(tube[j][i-1][0],tube[j][i-1][1])
        else:
            if(i%2==0):
                tr.goto(tube[j][i+1][0],tube[j][i+1][1])

input('Finish! Please put Enter Key')     #待機用Input
  

チューブの長さや角度、炭素原子同士の間隔等はプログラム中のパラメータを変えることで任意に変更できます。
また、このプログラムでは各炭素原子の座標を配列に入れておき、座標を回転させたり、線で繋いだりしてチューブを描画しています。

##処理の流れ
全体の処理の流れとしては、

1. ３次元の座標を入れる[x,y,z]を、段数 × チューブ断面当たりの原子の個数で用意(３次元配列で作成)。

2. 設定した原子の間隔や個数から、それぞれの炭素原子の座標を計算。

3. 設定した3次元の回転角(rol, yaw, pit)から、すべての原子を原点を中心に回転(チューブが回転)。

4. 各原子の位置で丸を描画、同時に接続する原子に向かって線を描画。(この時に各原子の奥行に応じて丸の色を変えている。)

以上です。結果としてけっこう綺麗に描くことができたので満足しています。暇を見て座標の計算部分やプログラムの進み方についてより詳しく書きたいと思います。

##参考リンク
ベクトルの内積や行列の積を求めるnumpy.dot関数の使い方