Read a pdb filename using pymol and return a topology and coordinates. Save a pdb filename from topology and coordinates using pymol.

get_topology.py

#!/usr/bin/env python
# -*- coding: UTF8 -*-

# Author: Guillaume Bouvier -- guillaume.bouvier@pasteur.fr
# https://research.pasteur.fr/en/member/guillaume-bouvier/
# 2019-04-12 13:53:36 (UTC+0200)

import sys
import __main__
__main__.pymol_argv = ['pymol', '-qc'] # Optionnaly pass options to pymol
import pymol
stdout = sys.stdout # To get the stdout in the notebook instead of in the pymol console
import pymol.cmd as cmd
pymol.finish_launching()
sys.stdout = stdout

import numpy

def get_topology(filename):
    """
    • filename: pdb file name
    """
    cmd.load(filename, 'struct')
    selection = 'struct'
    coords = cmd.get_coords(selection=selection)
    topology = {'resnames': [],
                'names': [],
                'chains': [],
                'resids': []}
    cmd.iterate('name ca and struct',
                'resnames.append(resn);names.append(name);chains.append(chain);resids.append(resi)',
                space=topology)
    for k in topology:
        topology[k] = numpy.asarray(topology[k])
        if k == 'resids':
            topology[k] = numpy.int_(topology[k])
    topology['counts'] = numpy.arange(len(topology['resids']))
    return coords, topology

def save_coords(coords, topology, outfilename, selection=None):
    """
    Save the coordinates to a pdb file
    • coords: coordinates
    • topology: topology
    • outfilename: name of the oupyt pdb
    • selection: Boolean array to select atoms
    """
    object_name = 'struct_save_coords'
    cmd.delete(object_name)
    if selection is None:
        selection = numpy.ones(len(topology['resids']), dtype=bool)
    for i, coords_ in enumerate(coords):
        if selection[i]:
            name = topology['names'][i]
            resn = topology['resnames'][i]
            resi = topology['resids'][i]
            chain = topology['chains'][i]
            elem = name[0]
            cmd.pseudoatom(object_name,
                           name=name,
                           resn=resn,
                           resi=resi,
                           chain=chain,
                           elem=elem,
                           hetatm=0,
                           segi=chain,
                           pos=list(coords_))
    cmd.save(outfilename, selection=object_name)