com.py

ATOMIC_WEIGHTS = {'H':1.008, 'HE':4.002602, 'LI':6.94, 'BE':9.012182,
       'B':10.81, 'C':12.011, 'N':14.007, 'O':15.999, 'F':18.9984032,
       'NE':20.1797, 'NA':22.98976928, 'MG':24.305, 'AL':26.9815386,
       'SI':28.085, 'P':30.973762, 'S':32.06, 'CL':35.45, 'AR':39.948,
       'K':39.0983, 'CA':40.078, 'SC':44.955912, 'TI':47.867, 'V':50.9415,
       'CR':51.9961, 'MN':54.938045, 'FE':55.845, 'CO':58.933195,
       'NI':58.6934, 'CU':63.546, 'ZN':65.38, 'GA':69.723, 'GE':72.630,
       'AS':74.92160, 'SE':78.96, 'BR':79.904, 'RB':85.4678, 'SR':87.62,
       'Y':88.90585, 'ZR':91.224, 'NB':92.90638, 'MO':95.96, 'TC':98,
       'RU':101.07, 'RH':102.90550, 'PD':106.42, 'AG':107.8682, 'CD':112.411,
       'IN':114.818, 'SN':118.710, 'SB':121.760, 'TE':127.60, 'I':126.90447,
       'XE':131.293, 'CS':132.9054519, 'BA':137.327, 'LA':138.90547,
       'CE':140.116, 'PR':140.90765, 'ND':144.242, 'PM':145, 'SM':150.36,
       'EU':151.964, 'GD':157.25, 'TB':158.92535, 'DY':162.500, 'HO':164.93032,
       'ER':167.259, 'TM':168.93421, 'YB':173.054, 'LU':174.9668, 'HF':178.49,
       'TA':180.94788, 'W':183.84, 'RE':186.207, 'OS':190.23, 'IR':192.217,
       'PT':195.084, 'AU':196.966569, 'HG':200.592, 'TL':204.38, 'PB':207.2,
       'BI':208.98040, 'PO':209, 'AT':210, 'RN':222, 'FR':223, 'RA':226,
       'AC':227, 'TH':232.03806, 'PA':231.03588, 'U':238.02891, 'NP':237,
       'PU':244, 'AM':243, 'CM':247, 'BK':247, 'CF':251, 'ES':252, 'FM':257,
       'MD':258, 'NO':259, 'LR':262, 'RF':267, 'DB':268, 'SG':269, 'BH':270,
       'HS':269, 'MT':278, 'DS':281, 'RG':281, 'CN':285, 'UUT':286, 'FL':289,
       'UUP':288, 'LV':293, 'UUS':294}

def center_of_mass(pdbfile, include='ATOM,HETATM'):
    """
    Calculates center of mass of a protein and/or ligand structure.
    Returns:
        center (list): List of float coordinates [x,y,z] that represent the
        center of mass (precision 3).
    """
    
    center = [None, None, None]
    include = tuple(include.split(','))
    
    with open(pdbfile, 'r') as pdb:

        # extract coordinates [ [x1,y1,z1], [x2,y2,z2], ... ]
        coordinates = []
        masses = []    
        for line in pdb:
            if line.startswith(include):
                coordinates.append([float(line[30:38]),    # x_coord
                                    float(line[38:46]),    # y_coord
                                    float(line[46:54])     # z_coord
                                   ])
                element_name = line[76:].strip()
                if element_name not in ATOMIC_WEIGHTS:
                    element_name = line.split()[2].strip()[0]
                masses.append(ATOMIC_WEIGHTS[element_name])

        assert len(coordinates) == len(masses)

        # calculate relative weight of every atomic mass
        total_mass = sum(masses)
        weights = [float(atom_mass/total_mass) for atom_mass in masses]

        # calculate center of mass
        center = [sum([coordinates[i][j] * weights[i]
              for i in range(len(weights))]) for j in range(3)]
        center_rounded = [round(center[i], 3) for i in range(3)]
        return center_rounded
    
    
if __name__ == '__main__':
    
    import argparse
    
    parser = argparse.ArgumentParser(
        description='Calculates the weighted center of mass for structures in a PDB file.\n'\
                    'By default, all atoms in the PDB file are included in the calculation.',
        epilog='Example:\n'\
                'center_of_mass.py ~/Desktop/protein.pdb \n'\
                '[-8.125, 20.461, -10.438]'
                '\n\nNote that for the center of mass calculation, the relative\natomic'\
                ' weights are taken into account (atomic mass unit [u]).\n\n'\
                'A list of the atomic weights can be found, e.g., at\n'\
                'http://en.wikipedia.org/wiki/List_of_elements',
        formatter_class=argparse.RawTextHelpFormatter
        )
    
    parser.add_argument('PDBfile')
    parser.add_argument('-i', '--include', type=str, 
            default='ATOM,HETATM', 
            metavar='coordinate-ID',
            help='Coordinate lines to include (default: "ATOM,HETATM")')
            
    parser.add_argument('-v', '--version', action='version', version='center_of_mass v. 1.0')
            
            
    args = parser.parse_args()
    
    print(center_of_mass(args.PDBfile, include=args.include))