Source code for graphein.construct_meshes

"""Class for creating Protein Meshes"""

# Graphein
# Author: Arian Jamasb <arian@jamasb.io>
# License: MIT
# Project Website:
# Code Repository: https://github.com/a-r-j/graphein

from pytorch3d.io import load_obj, save_obj
from ipymol import viewer as pymol


[docs]class ProteinMesh(object):
[docs] def __init__(self): """ Initialise ProteinGraph Generator Class """
[docs] def get_obj_file(self, pdb_file=None, pdb_code=None, out_dir=None): """ Produces .Obj file from PDB structure through IPyMol. pdb_code and pdb_file are optional arguments. Use one as suits your purposes :param pdb_file: Path to local .PDB file :type pdb_file: str :param pdb_code: 4 character PDB accession code :type pdb_code: str :param out_dir: Path to output directory :type out_dir: str :return: """ file_name = "a" pymol.start() if not pdb_code and not pdb_file: print('Please pass a pdb_file or pdb_code argument') print(pdb_file) if pdb_code and pdb_file: print('Do not pass both a PDB code and PDB file. Choose one.') if pdb_file: pymol.load(pdb_file) file_name = pdb_file[:-3] + '.obj' if pdb_code: pymol.fetch(pdb_code) file_name = out_dir + pdb_code + '.obj' pymol.do('show_as surface') pymol.do(f'save {file_name}') print(f'Saved file to: {file_name}') return file_name
[docs] def create_mesh(self, pdb_code=None, pdb_file=None, out_dir=None): """ Creates a PyTorch3D Mesh from an .Obj file. pdb_code and pdb_file are optional arguments. Use one as suits your purposes :param pdb_code: 4-character PDB accession code :type pdb_code: str :param pdb_file: Path to local .PDB file :type pdb_file: str :param out_dir: Path to output directory :type out_dir: str :return: verts, faces, aux """ obj_file = self.get_obj_file(pdb_code=pdb_code, pdb_file=pdb_file, out_dir=out_dir) verts, faces, aux = load_obj(obj_file) return verts, faces, aux
if __name__ == "__main__": p = ProteinMesh() print(p.create_mesh(pdb_code='3eiy', out_dir='../examples/meshes/')) #print(p.create_mesh(obj_file='../examples/meshes/3eiy.obj'))