source: rotdif/bin/ELM_ellipsoid.py @ 1705

Last change on this file since 1705 was 1705, checked in by yuexi, 2 years ago

finish ellipsoid

File size: 3.5 KB
Line 
1# Draw a rotated ellipsoid, given the values of center of mass, semiaxes, and eigenvectors
2# written by Matthew Casertano and David Fushman, University of Maryland, June 2019
3from pymol.cgo import BEGIN, COLOR, TRIANGLES, VERTEX, NORMAL, END
4from pymol import cmd
5
6def sin (theta): return math.sin(math.radians(theta)) # sine of an angle in degree form
7def cos (theta): return math.cos(math.radians(theta)) # cosine of an angle in degree form
8
9
10def ellipsoid(cmx, cmy, cmz, a1, a2, a3, u, v, R11, R12, R13, R21, R22, R23, R31, R32, R33):
11 
12    #coordinates of ellipsoid points and the normals in the ellipsoid coordinate frame
13        x0 = a1 * cos(u) * cos(v)
14        y0 = a2 * cos(u) * sin(v)
15        z0 = a3 * sin(u)
16       
17        nx0 = x0 / (a1 ** 2)
18        ny0 = y0 / (a2 ** 2)
19        nz0 = z0 / (a3 ** 2)
20
21    #rotate the coordinates of ellipsoid points and the normals into the protein coordinate frame
22    #and shift the center of ellipsoid to coincide with the COM of the protein molecule
23
24        x = x0 * R11 + y0 * R12 + z0 * R13 + cmx
25        y = x0 * R21 + y0 * R22 + z0 * R23 + cmy
26        z = x0 * R31 + y0 * R32 + z0 * R33 + cmz
27
28        nx = nx0 * R11 + ny0 * R12 + nz0 * R13
29        ny = nx0 * R21 + ny0 * R22 + nz0 * R23
30        nz = nx0 * R31 + ny0 * R32 + nz0 * R33
31
32        return x, y, z, nx, ny, nz
33
34def makeEllipsoid(color, cmx, cmy, cmz, a1, a2, a3, u1, u2, v1, v2, u_segs, v_segs, R11, R12, R13, R21, R22, R23, R31, R32, R33):
35
36        r, g, b = color
37
38        # Calculate delta variables
39        dU = (u2 - u1) / u_segs
40        dV = (v2 - v1) / v_segs
41
42        o = [BEGIN, TRIANGLES]
43
44        U = u1
45        for Y in range(0, u_segs):
46                # Initialize variables for loop
47                V = v1
48                for X in range(0, v_segs):
49                    x, y, z, nx, ny, nz = [0] * 5, [0] * 5, [0] * 5, [0] * 5, [0] * 5, [0] * 5 # this allows the array to be filled
50                    for i in range(1, 5): # creates four sets of ellipsoid points and normals, indexed from 1 through 4
51                        x[i], y[i], z[i], nx[i], ny[i], nz[i] = ellipsoid(cmx, cmy, cmz, a1, a2, a3, U + dU * (i == 2 or i == 3), V + dV * (i == 3 or i == 4), R11, R12, R13, R21, R22, R23, R31, R32, R33)
52                    for i in (1, 2, 4, 2, 3, 4): # extends the output list with the new points
53                        o.extend([COLOR, r, g, b, NORMAL, nx[i], ny[i], nz[i], VERTEX, x[i], y[i], z[i]])
54
55                    # Update variables for next loop
56                    V += dV
57                # Update variables for next loop
58                U += dU
59        o.append(END)
60        return o
61
62def drawEllipsoid(color, cmx, cmy, cmz, a1, a2, a3, R11, R12, R13, R21, R22, R23, R31, R32, R33):
63                return makeEllipsoid(color, cmx, cmy, cmz, a1, a2, a3, -90, 90, -180, 180, 10, 10, 
64                                R11, R12, R13, R21, R22, R23, R31, R32, R33) 
65
66def rotationMatrix(rotationInput): #Elements of the rotation matrix R = transpose of the rotation matrix defined in Arfken's book
67    if (len (rotationInput)) == 9: return rotationInput # simply return input if it is already a matrix
68    a, b, g = rotationInput[0], rotationInput[1], rotationInput[2]
69    R11, R12, R13 = cos(g) * cos(b) * cos(a) - sin(g) * sin(a), -sin(g) * cos(b) * cos(a) - cos(g) * sin(a), sin(b) * cos(a)
70    R21, R22, R23 = cos(g) * cos(b) * sin(a) + sin(g) * cos(a), -sin(g) * cos(b) * sin(a) + cos(g) * cos(a), sin(b) * sin(a)
71    R31, R32, R33 = -cos(g) * sin(b), sin(g) * sin(b), cos(b)
72    return R11, R12, R13, R21, R22, R23, R31, R32, R33
Note: See TracBrowser for help on using the repository browser.