forked from biopython/biopython
/
Geometry.py
52 lines (42 loc) · 1.96 KB
/
Geometry.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
# Copyright (C) 2010, Joao Rodrigues (anaryin@gmail.com)
# This code is part of the Biopython distribution and governed by its
# license. Please see the LICENSE file that should have been included
# as part of this package.
"""
Module with assorted geometrical functions on
macromolecules.
"""
from Bio.PDB import Entity
def center_of_mass(entity, geometric=False):
"""
Returns gravitic [default] or geometric center of mass of an Entity.
Geometric assumes all masses are equal (geometric=True)
"""
# Structure, Model, Chain, Residue
if isinstance(entity, Entity.Entity):
atom_list = entity.get_atoms()
# List of Atoms
elif hasattr(entity, '__iter__') and [x for x in entity if x.level == 'A']:
atom_list = entity
else: # Some other weirdo object
raise ValueError("Center of Mass can only be calculated from the following objects:\n"
"Structure, Model, Chain, Residue, list of Atoms.")
masses = []
positions = [ [], [], [] ] # [ [X1, X2, ..] , [Y1, Y2, ...] , [Z1, Z2, ...] ]
for atom in atom_list:
masses.append(atom.mass)
for i, coord in enumerate(atom.coord.tolist()):
positions[i].append(coord)
# If there is a single atom with undefined mass complain loudly.
if 'ukn' in set(masses) and not geometric:
raise ValueError("Some Atoms don't have an element assigned.\n"
"Try adding them manually or calculate the geometrical center of mass instead.")
if geometric:
return [sum(coord_list)/len(masses) for coord_list in positions]
else:
w_pos = [ [], [], [] ]
for atom_index, atom_mass in enumerate(masses):
w_pos[0].append(positions[0][atom_index]*atom_mass)
w_pos[1].append(positions[1][atom_index]*atom_mass)
w_pos[2].append(positions[2][atom_index]*atom_mass)
return [sum(coord_list)/sum(masses) for coord_list in w_pos]