YuWang0103/LGGM-Text2Graph

1import os2 3from rdkit import Chem4from rdkit.Chem import Draw, AllChem5from rdkit.Geometry import Point3D6from rdkit import RDLogger7import imageio8import networkx as nx9import numpy as np10import rdkit.Chem11import wandb12import matplotlib.pyplot as plt13 14 15 16 17 18class MolecularVisualization:19    def __init__(self, remove_h, dataset_infos):20        self.remove_h = remove_h21        self.dataset_infos = dataset_infos22 23    def mol_from_graphs(self, node_list, adjacency_matrix):24        """25        Convert graphs to rdkit molecules26        node_list: the nodes of a batch of nodes (bs x n)27        adjacency_matrix: the adjacency_matrix of the molecule (bs x n x n)28        """29        # dictionary to map integer value to the char of atom30        atom_decoder = self.dataset_infos.atom_decoder31 32        # create empty editable mol object33        mol = Chem.RWMol()34 35        # add atoms to mol and keep track of index36        node_to_idx = {}37        for i in range(len(node_list)):38            if node_list[i] == -1:39                continue40            a = Chem.Atom(atom_decoder[int(node_list[i])])41            molIdx = mol.AddAtom(a)42            node_to_idx[i] = molIdx43 44        for ix, row in enumerate(adjacency_matrix):45            for iy, bond in enumerate(row):46                # only traverse half the symmetric matrix47                if iy <= ix:48                    continue49                if bond == 1:50                    bond_type = Chem.rdchem.BondType.SINGLE51                elif bond == 2:52                    bond_type = Chem.rdchem.BondType.DOUBLE53                elif bond == 3:54                    bond_type = Chem.rdchem.BondType.TRIPLE55                elif bond == 4:56                    bond_type = Chem.rdchem.BondType.AROMATIC57                else:58                    continue59                mol.AddBond(node_to_idx[ix], node_to_idx[iy], bond_type)60 61        try:62            mol = mol.GetMol()63        except rdkit.Chem.KekulizeException:64            print("Can't kekulize molecule")65            mol = None66        return mol67 68    def visualize(self, path: str, molecules: list, num_molecules_to_visualize: int, log='graph'):69        # define path to save figures70        if not os.path.exists(path):71            os.makedirs(path)72 73        # visualize the final molecules74        print(f"Visualizing {num_molecules_to_visualize} of {len(molecules)}")75        if num_molecules_to_visualize > len(molecules):76            print(f"Shortening to {len(molecules)}")77            num_molecules_to_visualize = len(molecules)78        79        for i in range(num_molecules_to_visualize):80            file_path = os.path.join(path, 'molecule_{}.png'.format(i))81            mol = self.mol_from_graphs(molecules[i][0].numpy(), molecules[i][1].numpy())82            try:83                Draw.MolToFile(mol, file_path)84                if wandb.run and log is not None:85                    print(f"Saving {file_path} to wandb")86                    wandb.log({log: wandb.Image(file_path)}, commit=True)87            except rdkit.Chem.KekulizeException:88                print("Can't kekulize molecule")89 90 91    def visualize_chain(self, path, nodes_list, adjacency_matrix, trainer=None):92        RDLogger.DisableLog('rdApp.*')93        # convert graphs to the rdkit molecules94        mols = [self.mol_from_graphs(nodes_list[i], adjacency_matrix[i]) for i in range(nodes_list.shape[0])]95 96        # find the coordinates of atoms in the final molecule97        final_molecule = mols[-1]98        AllChem.Compute2DCoords(final_molecule)99 100        coords = []101        for i, atom in enumerate(final_molecule.GetAtoms()):102            positions = final_molecule.GetConformer().GetAtomPosition(i)103            coords.append((positions.x, positions.y, positions.z))104 105        # align all the molecules106        for i, mol in enumerate(mols):107            AllChem.Compute2DCoords(mol)108            conf = mol.GetConformer()109            for j, atom in enumerate(mol.GetAtoms()):110                x, y, z = coords[j]111                conf.SetAtomPosition(j, Point3D(x, y, z))112 113        # draw gif114        save_paths = []115        num_frams = nodes_list.shape[0]116 117        for frame in range(num_frams):118            file_name = os.path.join(path, 'fram_{}.png'.format(frame))119            Draw.MolToFile(mols[frame], file_name, size=(300, 300), legend=f"Frame {frame}")120            save_paths.append(file_name)121 122        imgs = [imageio.imread(fn) for fn in save_paths]123        gif_path = os.path.join(os.path.dirname(path), '{}.gif'.format(path.split('/')[-1]))124        imgs.extend([imgs[-1]] * 10)125        imageio.mimsave(gif_path, imgs, subrectangles=True, duration=20)126 127        if wandb.run:128            print(f"Saving {gif_path} to wandb")129            wandb.log({"chain": wandb.Video(gif_path, fps=5, format="gif")}, commit=True)130 131        # draw grid image132        try:133            img = Draw.MolsToGridImage(mols, molsPerRow=10, subImgSize=(200, 200))134            img.save(os.path.join(path, '{}_grid_image.png'.format(path.split('/')[-1])))135        except Chem.rdchem.KekulizeException:136            print("Can't kekulize molecule")137        return mols138 139 140class NonMolecularVisualization:141    def to_networkx(self, node_list, adjacency_matrix):142        """143        Convert graphs to networkx graphs144        node_list: the nodes of a batch of nodes (bs x n)145        adjacency_matrix: the adjacency_matrix of the molecule (bs x n x n)146        """147        graph = nx.Graph()148 149        for i in range(len(node_list)):150            if node_list[i] == -1:151                continue152            graph.add_node(i, number=i, symbol=node_list[i], color_val=node_list[i])153 154        rows, cols = np.where(adjacency_matrix >= 1)155        edges = zip(rows.tolist(), cols.tolist())156        for edge in edges:157            edge_type = adjacency_matrix[edge[0]][edge[1]]158            graph.add_edge(edge[0], edge[1], color=float(edge_type), weight=3 * edge_type)159 160        return graph161 162    def visualize_non_molecule(self, graph, pos, path, iterations=100, node_size=100, largest_component=False):163        if largest_component:164            CGs = [graph.subgraph(c) for c in nx.connected_components(graph)]165            CGs = sorted(CGs, key=lambda x: x.number_of_nodes(), reverse=True)166            graph = CGs[0]167 168        # Plot the graph structure with colors169        if pos is None:170            pos = nx.spring_layout(graph, iterations=iterations)171 172        # Set node colors based on the eigenvectors173        w, U = np.linalg.eigh(nx.normalized_laplacian_matrix(graph).toarray())174        vmin, vmax = np.min(U[:, 1]), np.max(U[:, 1])175        m = max(np.abs(vmin), vmax)176        vmin, vmax = -m, m177 178        plt.figure()179        nx.draw(graph, pos, font_size=5, node_size=node_size, with_labels=False, node_color=U[:, 1],180                cmap=plt.cm.coolwarm, vmin=vmin, vmax=vmax, edge_color='grey')181 182        plt.tight_layout()183        plt.savefig(path)184        plt.close("all")185 186    def visualize(self, path: str, graphs: list, num_graphs_to_visualize: int, log='graph'):187        # define path to save figures188        if not os.path.exists(path):189            os.makedirs(path)190 191        # visualize the final molecules192        for i in range(num_graphs_to_visualize):193            file_path = os.path.join(path, 'graph_{}.png'.format(i))194            graph = self.to_networkx(graphs[i][0].numpy(), graphs[i][1].numpy())195            self.visualize_non_molecule(graph=graph, pos=None, path=file_path)196            im = plt.imread(file_path)197            if wandb.run and log is not None:198                wandb.log({log: [wandb.Image(im, caption=file_path)]})199 200    def visualize_chain(self, path, nodes_list, adjacency_matrix):201        # convert graphs to networkx202        graphs = [self.to_networkx(nodes_list[i], adjacency_matrix[i]) for i in range(nodes_list.shape[0])]203        # find the coordinates of atoms in the final molecule204        final_graph = graphs[-1]205        final_pos = nx.spring_layout(final_graph, seed=0)206 207        # draw gif208        save_paths = []209        num_frams = nodes_list.shape[0]210 211        for frame in range(num_frams):212            file_name = os.path.join(path, 'fram_{}.png'.format(frame))213            self.visualize_non_molecule(graph=graphs[frame], pos=final_pos, path=file_name)214            save_paths.append(file_name)215 216        imgs = [imageio.imread(fn) for fn in save_paths]217        gif_path = os.path.join(os.path.dirname(path), '{}.gif'.format(path.split('/')[-1]))218        imgs.extend([imgs[-1]] * 10)219        imageio.mimsave(gif_path, imgs, subrectangles=True, duration=20)220        if wandb.run:221            wandb.log({'chain': [wandb.Video(gif_path, caption=gif_path, format="gif")]})222