''' http://www.pymolwiki.org/index.php/inertia_tensor (c) August 2010 by Mateusz Maciejewski matt (at) mattmaciejewski . com License: MIT ''' from pymol.cgo import * from pymol import cmd def tensor(selection, name="tensor", state=1, scaling=0, quiet=1): """ DESCRIPTION This script will draw the inertia tensor of the selection. ARGUMENTS selection = string: selection for the atoms included in the tensor calculation name = string: name of the tensor object to be created {default: "tensor"} state = int: state/model in the molecule object used in the tensor calculation scaling = int {0, 1, or 2}: 0 for no scaling of the inertia axes, 1 for scaling according to the molecular shape, 2 for scaling according to the eigenvalues {default: 0} EXAMPLE PyMOL> run inertia_tensor.py PyMOL> tensor molecule_object & i. 2-58+63-120 & n. n+ca+c, "tensor_model5_dom2", 5, 1 NOTES Requires numpy. """ import numpy def draw_axes(start, ends, cone_ends, radius=.2, name_obj="tensor"): radius = float(radius) size = radius * 15. obj = [ CYLINDER, start[0], start[1], start[2], ends[0][0] + start[0], ends[0][1] + start[1], ends[0][2] + start[2], radius, 1.0, 1.0, 1.0, 1.0, 0.0, 0., CYLINDER, start[0], start[1], start[2], (-1) * ends[0][0] + start[0], (-1) * ends[0][1] + start[1], (-1) * ends[0][2] + start[2], radius, 1.0, 1.0, 1.0, 1.0, 0.0, 0., CYLINDER, start[0], start[1], start[2], ends[1][0] + start[0], ends[1][1] + start[1], ends[1][2] + start[2], radius, 1.0, 1.0, 1.0, 0., 1.0, 0., CYLINDER, start[0], start[1], start[2], (-1) * ends[1][0] + start[0], (-1) * ends[1][1] + start[1], (-1) * ends[1][2] + start[2], radius, 1.0, 1.0, 1.0, 0., 1.0, 0., CYLINDER, start[0], start[1], start[2], ends[2][0] + start[0], ends[2][1] + start[1], ends[2][2] + start[2], radius, 1.0, 1.0, 1.0, 0., 0.0, 1.0, CYLINDER, start[0], start[1], start[2], (-1) * ends[2][0] + start[0], (-1) * ends[2][1] + start[1], (-1) * ends[2][2] + start[2], radius, 1.0, 1.0, 1.0, 0., 0.0, 1.0, CONE, ends[0][0] + start[0], ends[0][1] + start[1], ends[0][2] + start[2], cone_ends[0][0] + start[0], cone_ends[0][1] + start[1], cone_ends[0][2] + start[2], 0.5, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, CONE, (-1) * ends[0][0] + start[0], (-1) * ends[0][1] + start[1], (-1) * ends[0][2] + start[2], (-1) * cone_ends[0][0] + start[0], (-1) * cone_ends[0][1] + start[1], (-1) * cone_ends[0][2] + start[2], 0.5, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, CONE, ends[1][0] + start[0], ends[1][1] + start[1], ends[1][2] + start[2], cone_ends[1][0] + start[0], cone_ends[1][1] + start[1], cone_ends[1][2] + start[2], 0.5, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 1.0, CONE, (-1) * ends[1][0] + start[0], (-1) * ends[1][1] + start[1], (-1) * ends[1][2] + start[2], (-1) * cone_ends[1][0] + start[0], (-1) * cone_ends[1][1] + start[1], (-1) * cone_ends[1][2] + start[2], 0.5, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 1.0, CONE, ends[2][0] + start[0], ends[2][1] + start[1], ends[2][2] + start[2], cone_ends[2][0] + start[0], cone_ends[2][1] + start[1], cone_ends[2][2] + start[2], 0.5, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0, CONE, (-1) * ends[2][0] + start[0], (-1) * ends[2][1] + start[1], (-1) * ends[2][2] + start[2], (-1) * cone_ends[2][0] + start[0], (-1) * cone_ends[2][1] + start[1], (-1) * cone_ends[2][2] + start[2], 0.5, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0, ] cmd.load_cgo(obj, name_obj) totmass = 0.0 x_com, y_com, z_com = 0, 0, 0 model = cmd.get_model(selection, state) for a in model.atom: x_com += a.coord[0] * a.get_mass() y_com += a.coord[1] * a.get_mass() z_com += a.coord[2] * a.get_mass() totmass += a.get_mass() x_com /= totmass y_com /= totmass z_com /= totmass if not int(quiet): print() print("Center of mass: ") print() print(x_com, y_com, z_com) I = [] for index in range(9): I.append(0) for a in model.atom: temp_x, temp_y, temp_z = a.coord[0], a.coord[1], a.coord[2] temp_x -= x_com temp_y -= y_com temp_z -= z_com I[0] += a.get_mass() * (temp_y ** 2 + temp_z ** 2) I[4] += a.get_mass() * (temp_x ** 2 + temp_z ** 2) I[8] += a.get_mass() * (temp_x ** 2 + temp_y ** 2) I[1] -= a.get_mass() * temp_x * temp_y I[3] -= a.get_mass() * temp_x * temp_y I[2] -= a.get_mass() * temp_x * temp_z I[6] -= a.get_mass() * temp_x * temp_z I[5] -= a.get_mass() * temp_y * temp_z I[7] -= a.get_mass() * temp_y * temp_z tensor = numpy.array([(I[0:3]), (I[3:6]), (I[6:9])]) vals, vects = numpy.linalg.eig(tensor) # they come out unsorted, so the command below is needed eig_ord = numpy.argsort(vals) # a thing to note is that here COLUMN i corrensponds to eigenvalue i. ord_vals = vals[eig_ord] ord_vects = vects[:, eig_ord].T if not int(quiet): print() print("Inertia tensor z, y, x eigenvalues:") print() print(ord_vals) print() print("Inertia tensor z, y, x eigenvectors:") print() print(ord_vects) if int(scaling) == 0: norm_vals = [sum(numpy.sqrt(ord_vals / totmass)) / 3 for i in range(3)] elif int(scaling) == 1: normalizer = numpy.sqrt(max(ord_vals) / totmass) norm_vals = normalizer / numpy.sqrt(ord_vals / totmass) * normalizer norm_vals = norm_vals / (max(norm_vals) / min(norm_vals)) elif int(scaling) == 2: normalizer = numpy.sqrt(max(ord_vals) / totmass) norm_vals = numpy.sqrt(ord_vals / totmass) start = [x_com, y_com, z_com] ends = [[(norm_vals[0] - 1) * ord_vects[0][0], (norm_vals[0] - 1) * ord_vects[0][1], (norm_vals[0] - 1) * ord_vects[0][2]], [(norm_vals[1] - 1) * ord_vects[1][0], (norm_vals[1] - 1) * ord_vects[1][1], (norm_vals[1] - 1) * ord_vects[1][2]], [(norm_vals[2] - 1) * ord_vects[2][0], (norm_vals[2] - 1) * ord_vects[2][1], (norm_vals[2] - 1) * ord_vects[2][2]]] cone_ends = [[norm_vals[0] * ord_vects[0][0], norm_vals[0] * ord_vects[0][1], norm_vals[0] * ord_vects[0][2]], [norm_vals[1] * ord_vects[1][0], norm_vals[1] * ord_vects[1][1], norm_vals[1] * ord_vects[1][2]], [norm_vals[2] * ord_vects[2][0], norm_vals[2] * ord_vects[2][1], norm_vals[2] * ord_vects[2][2]]] draw_axes(start, ends, cone_ends, name_obj=name) cmd.extend("tensor", tensor)