-
Notifications
You must be signed in to change notification settings - Fork 0
/
arrow.py
151 lines (128 loc) · 5.27 KB
/
arrow.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
from pymol import cmd
from pymol.cgo import *
from math import *
#
# Some functions to allow drawing arrows (vectors) in Pymol
# In need of proper documentation...
#
# Please don't distribute (parts of) this file, without credits
#
# (c)2006 Tsjerk A. Wassenaar, PhD, University of Utrecht
#
# t s j e r k w .at. g m a i l .dot. c o m
# http://nmr.chem.uu.nl/~tsjerk/
#
# BY Documentation
# Copy this file into the folder with the command script(s). Change the colour designations below. In the command script, insert for each arrow the following command after running this script.
# Syntax: cgo_arrow([startx, starty, startz], [endx, endy, endz], thickness of arrow, width of the wide end of the arrowhead, length of arrowhead, name of object (must be an integer))
# The start is the arrow tail, and the end is the arrow head.
# Arrows with arrowhead length one quarter of the vector length and with arrowhead width one seventh of the arrowhead length look reasonable. Vector length is: SQRT[(x1-x2)^2 + (y1-y2)^2 + (z1-z2)^2]
# Example with arrowhead: cgo_arrow( [ -5.07 , 28.02 , 3.06 ], [ -4.11 , 29.57 , 5.42 ],0.1, 0.11 , 0.74 , 1 )
# Example without arrowhead: cgo_arrow( [ -5.07 , 28.02 , 3.06 ], [ -4.11 , 29.57 , 5.42 ],0.1, 0 , 0 , 1 )
def t( X ):
if not X: return X
Y = []
for i in range( len( X[0] ) ):
Y.append( [] )
for j in X:
Y[i].append( j[i] )
return Y
def v_add( a, b ): return ( a[0]+b[0], a[1]+b[1], a[2]+b[2] )
def v_sub( a, b ): return ( a[0]-b[0], a[1]-b[1], a[2]-b[2] )
def vecprod( a, b ): return ( a[1]*b[2]-a[2]*b[1], a[2]*b[0]-a[0]*b[2], a[0]*b[1]-a[1]*b[0] )
def inprod( a, b=None ):
if b: return a[0]*b[0] + a[1]*b[1] + a[2]*b[2]
else: return a[0]*a[0] + a[1]*a[1] + a[2]*a[2]
def svmult( s, a ): return ( s*a[0], s*a[1], s*a[2] )
def norm( a ): return svmult( 1/sqrt(inprod( a )), a )
def mvmult( R, x ):
y = []
for i in R: y.append( inprod( i, x ) )
return tuple(y)
def mv_add( X, a ):
Y = []
for i in X:
Y.append( v_add( i, a ) )
return Y
def mmmult( R, X ):
Y = []
for i in X: Y.append( mvmult( R, i ) )
return Y
def smatrix( v ): return [[ v[0], 0, 0 ], [ 0, v[1], 0 ], [ 0, 0, v[2] ]]
def rmatrix( v ):
cosx, sinx = cos( v[0] ), sin( v[0] )
cosy, siny = cos( v[1] ), sin( v[1] )
cosz, sinz = cos( v[2] ), sin( v[2] )
return mmmult( mmmult( [[1,0,0],[0,cosx,-sinx],[0,sinx,cosx]],
[[cosy,0,-siny],[0,1,0],[siny,0,cosy]] ),
[[cosz,-sinz,0],[sinz,cosz,0],[0,0,1]] )
def block( i, dphi ):
ddphi = 0.25*dphi
phi0 = i*dphi
phi1 = phi0+ddphi
phi2 = phi1+ddphi
phi3 = phi2+ddphi
phi4 = phi3+ddphi
sqrt2 = sqrt(2)
return [ (-0.5*sqrt2,-0.5*sqrt2*cos(phi2),0.5*sqrt2*sin(phi2)),
(1,0,0),
(0,cos(phi0),-sin(phi0)),
(0,cos(phi1),-sin(phi1)),
(0,cos(phi2),-sin(phi2)),
(0,cos(phi3),-sin(phi3)),
(0,cos(phi4),-sin(phi4))
]
def cgo_triangle_fan( X ):
Y = []
while ( X ):
i = X.pop(0)
Y.extend( [ NORMAL, i[0], i[1], i[2], ] )
for i in range( 6 ): #BY changing the 6 to 0 will eliminate the arrowhead
i = X.pop(0)
Y.extend( [ VERTEX, i[0], i[1], i[2], ] )
return Y
def cgo_arrow1( S, E, r=0.2, hr=0.4, hl=1.0 ):
P0 = S
D = v_sub( E, S )
DL = inprod( D, D )
P1 = v_add( S, svmult( (DL-hl)/DL, D ) )
P2 = E
# Define a vector orthogonal to P1-P0
V = v_sub( P1, P0 )
V = norm( V )
if V[2] != 0:
A = ( 1, 1, -(V[0]+V[1])/V[2] )
elif V[1] != 0:
A = ( 1, -V[0]/V[1], 0 )
else:
A = ( 0, -V[0], 0 )
A = norm( A )
B = vecprod( V, A )
print (inprod(V), inprod(B), inprod(A))
R = t([ svmult( hl,V ), svmult( hr,A ), svmult( hr,B ) ])
# Define the transformation matrix (scale and rotation)
#C = v_sub( P2, P1 )
#scale = ( hl, hr, hr )
#rotate = ( 0, acos( C[0]/sqrt(C[0]**2+C[2]**2) ), acos( C[0]/sqrt(C[0]**2+C[1]**2) ) )
#R = mmmult( smatrix( scale ), rmatrix( rotate ) )
obj = [
CYLINDER, S[0], S[1], S[2], P1[0], P1[1], P1[2], r, 0, 1, 0, 0, 0, 1, #BY: after the r, the first three are the RGB numbers for the arrow tail, and the last three are the RGB near the arrow head.
COLOR, 0, 0, 1, #BY: after COLOR, the three numbers are the RGB numbers for the arrow head itself
BEGIN, TRIANGLE_FAN ]
N = 10
dphi = 2*pi/N
crds = []
for i in range(N+1):
crds.extend( block( i, dphi ) )
crds = mv_add( mmmult( R, crds ), P1 )
obj.extend( cgo_triangle_fan( crds ) )
obj.extend( [ END, ] )
return obj
def cgo_arrow( S, E, r=0.2, hr=0.4, hl=1.0, name="arrow", state=1 ):
obj = cgo_arrow1( S, E, r=r, hr=hr, hl=hl )
cmd.load_cgo( obj, name, state )
def cgo_arrows( X, r=0.2, hr=0.4, hl=1.0, name="arrow", state=1 ):
obj = []
for i in X:
obj.extend( cgo_arrow1( (i[0], i[1], i[2]), (i[3], i[4], i[5]), r=r, hr=hr, hl=hl ) )
cmd.load_cgo( obj, name, state )