| from rdkit.Chem.rdmolfiles import MolToPDBBlock, MolToPDBFile | |
| import rdkit.Chem | |
| from rdkit import Geometry | |
| from collections import defaultdict | |
| import copy | |
| import numpy as np | |
| import torch | |
| class PDBFile: | |
| def __init__(self, mol): | |
| self.parts = defaultdict(dict) | |
| self.mol = copy.deepcopy(mol) | |
| [self.mol.RemoveConformer(j) for j in range(mol.GetNumConformers()) if j] | |
| def add(self, coords, order, part=0, repeat=1): | |
| if type(coords) in [rdkit.Chem.Mol, rdkit.Chem.RWMol]: | |
| block = MolToPDBBlock(coords).split('\n')[:-2] | |
| self.parts[part][order] = {'block': block, 'repeat': repeat} | |
| return | |
| elif type(coords) is np.ndarray: | |
| coords = coords.astype(np.float64) | |
| elif type(coords) is torch.Tensor: | |
| coords = coords.double().numpy() | |
| for i in range(coords.shape[0]): | |
| self.mol.GetConformer(0).SetAtomPosition(i, Geometry.Point3D(coords[i, 0], coords[i, 1], coords[i, 2])) | |
| block = MolToPDBBlock(self.mol).split('\n')[:-2] | |
| self.parts[part][order] = {'block': block, 'repeat': repeat} | |
| def write(self, path=None, limit_parts=None): | |
| is_first = True | |
| str_ = '' | |
| for part in sorted(self.parts.keys()): | |
| if limit_parts and part >= limit_parts: | |
| break | |
| part = self.parts[part] | |
| keys_positive = sorted(filter(lambda x: x >=0, part.keys())) | |
| keys_negative = sorted(filter(lambda x: x < 0, part.keys())) | |
| keys = list(keys_positive) + list(keys_negative) | |
| for key in keys: | |
| block = part[key]['block'] | |
| times = part[key]['repeat'] | |
| for _ in range(times): | |
| if not is_first: | |
| block = [line for line in block if 'CONECT' not in line] | |
| is_first = False | |
| str_ += 'MODEL\n' | |
| str_ += '\n'.join(block) | |
| str_ += '\nENDMDL\n' | |
| if not path: | |
| return str_ | |
| with open(path, 'w') as f: | |
| f.write(str_) |