| | import os |
| | import py3Dmol |
| | from rdkit import Chem |
| |
|
| |
|
| | def visualize_complex(pdb_block, sdf_block, show_protein_surface=True, show_ligand=True, show_ligand_surface=True): |
| | view = py3Dmol.view() |
| |
|
| | |
| | view.addModel(pdb_block, 'pdb') |
| | if show_protein_surface: |
| | view.addSurface(py3Dmol.VDW, {'opacity': 0.7, 'color': 'white'}, {'model': -1}) |
| | else: |
| | view.setStyle({'model': -1}, {'cartoon': {'color': 'spectrum'}, 'line': {}}) |
| | view.setStyle({'model': -1}, {"cartoon": {"style": "edged", 'opacity': 0}}) |
| |
|
| | |
| | if show_ligand: |
| | view.addModel(sdf_block, 'sdf') |
| | view.setStyle({'model': -1}, {'stick': {}}) |
| | |
| | if show_ligand_surface: |
| | view.addSurface(py3Dmol.VDW, {'opacity': 0.8}, {'model': -1}) |
| |
|
| | view.zoomTo() |
| | return view |
| |
|
| |
|
| | def visualize_data(data, root, show_ligand=True, show_surface=True): |
| | protein_path = os.path.join(root, data.protein_filename) |
| | ligand_path = os.path.join(root, data.ligand_filename) |
| | with open(protein_path, 'r') as f: |
| | pdb_block = f.read() |
| | with open(ligand_path, 'r') as f: |
| | sdf_block = f.read() |
| | return visualize_complex(pdb_block, sdf_block, show_ligand=show_ligand, show_surface=show_surface) |
| |
|
| |
|
| | def visualize_generated_mol(protein_filename, mol, root, show_surface=False, opacity=0.5): |
| | protein_path = os.path.join(root, protein_filename) |
| | with open(protein_path, 'r') as f: |
| | pdb_block = f.read() |
| |
|
| | view = py3Dmol.view() |
| |
|
| | |
| | view.addModel(pdb_block, 'pdb') |
| | view.setStyle({'model': -1}, {'cartoon': {'color': 'spectrum'}, 'line': {}}) |
| |
|
| | mblock = Chem.MolToMolBlock(mol) |
| | view.addModel(mblock, 'mol') |
| | view.setStyle({'model': -1}, {'stick': {}, 'sphere': {'radius': 0.35}}) |
| | if show_surface: |
| | view.addSurface(py3Dmol.SAS, {'opacity': opacity}, {'model': -1}) |
| |
|
| | view.zoomTo() |
| | return view |
| |
|
| |
|
| | def MolTo3DView(mol, size=(300, 300), style="stick", surface=False, opacity=0.5): |
| | """Draw molecule in 3D |
| | |
| | Args: |
| | ---- |
| | mol: rdMol, molecule to show |
| | size: tuple(int, int), canvas size |
| | style: str, type of drawing molecule |
| | style can be 'line', 'stick', 'sphere', 'carton' |
| | surface, bool, display SAS |
| | opacity, float, opacity of surface, range 0.0-1.0 |
| | Return: |
| | ---- |
| | viewer: py3Dmol.view, a class for constructing embedded 3Dmol.js views in ipython notebooks. |
| | """ |
| | assert style in ('line', 'stick', 'sphere', 'carton') |
| |
|
| | viewer = py3Dmol.view(width=size[0], height=size[1]) |
| | if isinstance(mol, list): |
| | for i, m in enumerate(mol): |
| | mblock = Chem.MolToMolBlock(m) |
| | viewer.addModel(mblock, 'mol' + str(i)) |
| | elif len(mol.GetConformers()) > 1: |
| | for i in range(len(mol.GetConformers())): |
| | mblock = Chem.MolToMolBlock(mol, confId=i) |
| | viewer.addModel(mblock, 'mol' + str(i)) |
| | else: |
| | mblock = Chem.MolToMolBlock(mol) |
| | viewer.addModel(mblock, 'mol') |
| | viewer.setStyle({style: {}}) |
| | if surface: |
| | viewer.addSurface(py3Dmol.SAS, {'opacity': opacity}) |
| | viewer.zoomTo() |
| | return viewer |
| |
|
