OneScience-Group/SurfDock

1from rdkit.Chem.rdmolfiles import MolToPDBBlock, MolToPDBFile2import rdkit.Chem 3from rdkit import Geometry4from collections import defaultdict5import copy6import numpy as np7import torch8 9    10class PDBFile:11    def __init__(self, mol):12        self.parts = defaultdict(dict)13        self.mol = copy.deepcopy(mol)14        [self.mol.RemoveConformer(j) for j in range(mol.GetNumConformers()) if j]        15    def add(self, coords, order, part=0, repeat=1):16        if type(coords) in [rdkit.Chem.Mol, rdkit.Chem.RWMol]:17            block = MolToPDBBlock(coords).split('\n')[:-2]18            self.parts[part][order] = {'block': block, 'repeat': repeat}19            return20        elif type(coords) is np.ndarray:21            coords = coords.astype(np.float64)22        elif type(coords) is torch.Tensor:23            coords = coords.double().numpy()24        for i in range(coords.shape[0]):25            self.mol.GetConformer(0).SetAtomPosition(i, Geometry.Point3D(coords[i, 0], coords[i, 1], coords[i, 2]))26        block = MolToPDBBlock(self.mol).split('\n')[:-2]27        self.parts[part][order] = {'block': block, 'repeat': repeat}28        29    def write(self, path=None, limit_parts=None):30        is_first = True31        str_ = ''32        for part in sorted(self.parts.keys()):33            if limit_parts and part >= limit_parts:34                break35            part = self.parts[part]36            keys_positive = sorted(filter(lambda x: x >=0, part.keys()))37            keys_negative = sorted(filter(lambda x: x < 0, part.keys()))38            keys = list(keys_positive) + list(keys_negative)39            for key in keys:40                block = part[key]['block']41                times = part[key]['repeat']42                for _ in range(times):43                    if not is_first:44                        block = [line for line in block if 'CONECT' not in line]45                    is_first = False46                    str_ += 'MODEL\n'47                    str_ += '\n'.join(block)48                    str_ += '\nENDMDL\n'49        if not path:50            return str_51        with open(path, 'w') as f:52            f.write(str_)