Spaces:
Sleeping
Sleeping
File size: 2,134 Bytes
4a3f787 |
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 |
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_) |