Spaces:
Running on Zero
Running on Zero
File size: 4,753 Bytes
b1d2f45 | 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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | import os
from rdkit import Chem
from rdkit.Chem import Draw, AllChem
from rdkit.Geometry import Point3D
from rdkit import RDLogger
import imageio # What is this?
import networkx as nx
import numpy as np
import rdkit.Chem
import matplotlib.pyplot as plt
class MolecularVisualization:
def __init__(self, remove_h, dataset_infos):
self.remove_h = remove_h
self.dataset_infos = dataset_infos
def mol_from_graphs(self, node_list, adjacency_matrix):
"""
Convert graphs to rdkit molecules
node_list: the nodes of a batch of nodes (bs x n)
adjacency_matrix: the adjacency_matrix of the molecule (bs x n x n)
"""
# dictionary to map integer value to the char of atom
atom_decoder = self.dataset_infos.atom_decoder
# create empty editable mol object
mol = Chem.RWMol()
# add atoms to mol and keep track of index
node_to_idx = {}
for i in range(len(node_list)):
if node_list[i] == -1:
continue
a = Chem.Atom(atom_decoder[int(node_list[i])])
molIdx = mol.AddAtom(a)
node_to_idx[i] = molIdx
for ix, row in enumerate(adjacency_matrix):
for iy, bond in enumerate(row):
# only traverse half the symmetric matrix
if iy <= ix:
continue
if bond == 1:
bond_type = Chem.rdchem.BondType.SINGLE
elif bond == 2:
bond_type = Chem.rdchem.BondType.DOUBLE
elif bond == 3:
bond_type = Chem.rdchem.BondType.TRIPLE
elif bond == 4:
bond_type = Chem.rdchem.BondType.AROMATIC
else:
continue
mol.AddBond(node_to_idx[ix], node_to_idx[iy], bond_type)
try:
mol = mol.GetMol()
except rdkit.Chem.KekulizeException:
print("Can't kekulize molecule")
mol = None
return mol
def visualize(self, path: str, molecules: list, num_molecules_to_visualize: int):
# define path to save figures
if not os.path.exists(path):
os.makedirs(path)
# visualize the final molecules
print(f"Visualizing {num_molecules_to_visualize} of {len(molecules)}")
if num_molecules_to_visualize > len(molecules):
print(f"Shortening to {len(molecules)}")
num_molecules_to_visualize = len(molecules)
for i in range(num_molecules_to_visualize):
file_path = os.path.join(path, 'molecule_{}.png'.format(i))
mol = self.mol_from_graphs(molecules[i][0].numpy(), molecules[i][1].numpy())
try:
Draw.MolToFile(mol, file_path)
except rdkit.Chem.KekulizeException:
print("Can't kekulize molecule")
def visualize_chain(self, path, nodes_list, adjacency_matrix):
RDLogger.DisableLog('rdApp.*')
# convert graphs to the rdkit molecules
mols = [self.mol_from_graphs(nodes_list[i], adjacency_matrix[i]) for i in range(nodes_list.shape[0])]
# find the coordinates of atoms in the final molecule
final_molecule = mols[-1]
AllChem.Compute2DCoords(final_molecule)
coords = []
for i, atom in enumerate(final_molecule.GetAtoms()):
positions = final_molecule.GetConformer().GetAtomPosition(i)
coords.append((positions.x, positions.y, positions.z))
# align all the molecules
for i, mol in enumerate(mols):
AllChem.Compute2DCoords(mol)
conf = mol.GetConformer()
for j, atom in enumerate(mol.GetAtoms()):
x, y, z = coords[j]
conf.SetAtomPosition(j, Point3D(x, y, z))
# draw gif
save_paths = []
num_frams = nodes_list.shape[0]
for frame in range(num_frams):
file_name = os.path.join(path, 'fram_{}.png'.format(frame))
Draw.MolToFile(mols[frame], file_name, size=(300, 300), legend=f"Frame {frame}")
save_paths.append(file_name)
imgs = [imageio.imread(fn) for fn in save_paths]
gif_path = os.path.join(os.path.dirname(path), '{}.gif'.format(path.split('/')[-1]))
imgs.extend([imgs[-1]] * 5)
imageio.mimsave(gif_path, imgs, subrectangles=True, duration=20)
# draw grid image
try:
img = Draw.MolsToGridImage(mols, molsPerRow=10, subImgSize=(200, 200))
img.save(os.path.join(path, '{}_grid_image.png'.format(path.split('/')[-1])))
except Chem.rdchem.KekulizeException:
print("Can't kekulize molecule")
return mols
|