Merge pull request #205 from wjm41/main

Added ESMFold as structure import option

MolecularNodes/__init__.py

@@ -88,6 +88,14 @@ def register():
         subtype = 'NONE',
         default = 0
     )
+    bpy.types.Scene.mol_esmfold_sequence = bpy.props.StringProperty(
+        name = 'amino_acid_sequence', 
+        description = 'Amino acid sequence of the structure to open', 
+        options = {'TEXTEDIT_UPDATE'}, 
+        default = '', 
+        subtype = 'FILE_PATH', 
+        maxlen = 0
+        )
     bpy.types.Scene.mol_import_local_path = bpy.props.StringProperty(
         name = 'path_pdb', 
         description = 'File path of the structure to open', 
@@ -144,6 +152,14 @@ def register():
         subtype = 'NONE', 
         maxlen = 0
         )
+    bpy.types.Scene.mol_esmfold_name = bpy.props.StringProperty(
+        name = 'mol_name', 
+        description = 'Name of the molecule on import', 
+        options = {'TEXTEDIT_UPDATE'}, 
+        default = 'NewMolecule', 
+        subtype = 'NONE', 
+        maxlen = 0
+        )
     bpy.types.Scene.mol_import_md_name = bpy.props.StringProperty(
         name = 'mol_md_name', 
         description = 'Name of the molecule on import', 
@@ -210,6 +226,7 @@ def register():
     bpy.utils.register_class(MOL_OT_Style_Surface_Custom)
 
     bpy.utils.register_class(MOL_OT_Import_Protein_RCSB)
+    bpy.utils.register_class(MOL_OT_Import_Protein_ESMFold)
 
     bpy.utils.register_class(MOL_OT_Import_Method_Selection)
     bpy.utils.register_class(MOL_OT_Import_Protein_Local)
@@ -251,6 +268,9 @@ def unregister():
     del bpy.types.Scene.mol_import_md_frame_end
     del bpy.types.Scene.mol_import_default_style
 
+    del bpy.types.Scene.mol_esmfold_name
+    del bpy.types.Scene.mol_esmfold_sequence
+
     del bpy.types.Scene.trajectory_selection_list
     del bpy.types.Scene.list_index
 
@@ -281,6 +301,7 @@ def unregister():
     bpy.utils.unregister_class(MOL_OT_Import_Protein_RCSB)
     bpy.utils.unregister_class(MOL_OT_Import_Method_Selection)
     bpy.utils.unregister_class(MOL_OT_Import_Protein_Local)
+    bpy.utils.unregister_class(MOL_OT_Import_Protein_ESMFold)
     bpy.utils.unregister_class(MOL_OT_Import_Protein_MD)
     bpy.utils.unregister_class(MOL_OT_Import_Map)
     bpy.utils.unregister_class(MOL_OT_Import_Star_File)

MolecularNodes/load.py

@@ -1,3 +1,5 @@
+import requests
+import io
 import bpy
 import numpy as np
 from . import coll
@@ -40,6 +42,38 @@ def molecule_rcsb(
 
     return mol_object
 
+def molecule_esmfold(
+    amino_acid_sequence,               
+    mol_name = "Name",                   
+    center_molecule = False,               
+    del_solvent = True,               
+    include_bonds = True,   
+    starting_style = 0,               
+    setup_nodes = True              
+    ):
+    mol, file = open_structure_esm_fold(
+        amino_acid_sequence = amino_acid_sequence, 
+        include_bonds=include_bonds
+        )
+
+    mol_object, coll_frames = create_molecule(
+        mol_array = mol,
+        mol_name = mol_name,
+        file = file,
+        calculate_ss = True,
+        center_molecule = center_molecule,
+        del_solvent = del_solvent, 
+        include_bonds = True
+        )
+
+    if setup_nodes:
+        nodes.create_starting_node_tree(
+            obj = mol_object, 
+            coll_frames=coll_frames, 
+            starting_style = starting_style
+            )    
+    return mol_object
+
 def molecule_local(
     file_path,                    
     mol_name = "Name",                   
@@ -113,7 +147,40 @@ def open_structure_rcsb(pdb_code, include_bonds = True):
     mol = mmtf.get_structure(file, extra_fields = ["b_factor", "charge"], include_bonds = include_bonds) 
     return mol, file
 
+def open_structure_esm_fold(amino_acid_sequence, include_bonds=True):
+    import biotite.structure.io.pdb as pdb
+
+
+
+    # output_of_subprocess = subprocess.Popen([
+    # 'curl',
+    # '-X',
+    # 'POST',
+    # '--data',
+    # amino_acid_sequence,
+    # 'https://api.esmatlas.com/foldSequence/v1/pdb/'
+    # ], stdout=subprocess.PIPE)
+
+    # (esm_folded_pdb_str, err) = output_of_subprocess.communicate()
 
+    r = requests.post('https://api.esmatlas.com/foldSequence/v1/pdb/', data=amino_acid_sequence)
+
+    if r.ok:
+
+        esm_folded_pdb_str = r.text
+        with io.StringIO() as f:
+            f.write(esm_folded_pdb_str)
+            f.seek(0)
+            file = pdb.PDBFile.read(f)
+
+        # returns a numpy array stack, where each array in the stack is a model in the 
+        # the file. The stack will be of length = 1 if there is only one model in the file
+        mol = pdb.get_structure(file, extra_fields = ['b_factor', 'charge'], include_bonds = include_bonds)
+        return mol, file
+
+    else:
+        raise ValueError(f'ESMFold returned an error for the amino acid sequence input. This is the error message: {r.text}')
+
 def open_structure_local_pdb(file_path, include_bonds = True):
     import biotite.structure.io.pdb as pdb
 

MolecularNodes/ui.py

@@ -36,7 +36,39 @@ def execute(self, context):
 
     def invoke(self, context, event):
         return self.execute(context)
+
+# operator that calls the function to import the structure from ESMFold
+class MOL_OT_Import_Protein_ESMFold(bpy.types.Operator):
+    bl_idname = "mol.import_protein_esmfold"
+    bl_label = "import_protein_esmfold"
+    bl_description = "Generate structure from ESMFold"
+    bl_options = {"REGISTER", "UNDO"}
+
+    @classmethod
+    def poll(cls, context):
+        return not False
 
+    def execute(self, context):
+        amino_acid_sequence = bpy.context.scene.mol_esmfold_sequence
+
+        mol_object = load.molecule_esmfold(
+            amino_acid_sequence=amino_acid_sequence, 
+            mol_name=bpy.context.scene.mol_esmfold_name,
+            include_bonds=bpy.context.scene.mol_import_include_bonds, 
+            center_molecule=bpy.context.scene.mol_import_center, 
+            del_solvent=bpy.context.scene.mol_import_del_solvent, 
+            starting_style=bpy.context.scene.mol_import_default_style, 
+            setup_nodes=True
+            )
+
+        # return the good news!
+        bpy.context.view_layer.objects.active = mol_object
+        self.report({'INFO'}, message=f"Generated protein '{amino_acid_sequence}' as {mol_object.name}")
+        return {"FINISHED"}
+
+    def invoke(self, context, event):
+        return self.execute(context)
+
 # operator that calls the function to import the structure from a local file
 class MOL_OT_Import_Protein_Local(bpy.types.Operator):
     bl_idname = "mol.import_protein_local"
@@ -134,6 +166,25 @@ def MOL_PT_panel_rcsb(layout_function, ):
     row_import.prop(bpy.context.scene, 'mol_pdb_code', text='PDB ID')
     row_import.operator('mol.import_protein_rcsb', text='Download', icon='IMPORT')
 
+def MOL_PT_panel_esmfold(layout_function, ):
+    col_main = layout_function.column(heading = '', align = False)
+    col_main.alert = False
+    col_main.enabled = True
+    col_main.active = True
+    col_main.label(text = "Generate Structure from ESMFold")
+    row_name = col_main.row(align = False)
+    row_name.prop(bpy.context.scene, 'mol_esmfold_name', 
+                    text = "Name", icon_value = 0, emboss = True)
+    row_name.operator('mol.import_protein_esmfold', text='Generate', icon='IMPORT')
+
+    row_seq = col_main.row()
+    row_seq.prop(
+        bpy.context.scene, 'mol_esmfold_sequence', 
+        text = "Sequence", 
+        icon_value = 0, 
+        emboss = True
+    )
+
 def MOL_PT_panel_local(layout_function, ):
     col_main = layout_function.column(heading = '', align = False)
     col_main.alert = False
@@ -417,10 +468,11 @@ def MOL_PT_panel_ui(layout_function, scene):
 
 
     MOL_change_import_interface(row, 'PDB',           0,  "URL")
-    MOL_change_import_interface(row, 'Local File',    1, 108)
-    MOL_change_import_interface(row, 'MD Trajectory', 2, 487)
-    MOL_change_import_interface(row, 'EM Map', 3, 'LIGHTPROBE_CUBEMAP')
-    MOL_change_import_interface(row, 'Star File',     4, 487)
+    MOL_change_import_interface(row, 'ESMFold',       1,  "URL")
+    MOL_change_import_interface(row, 'Local File',    2, 108)
+    MOL_change_import_interface(row, 'MD Trajectory', 3, 487)
+    MOL_change_import_interface(row, 'EM Map', 4, 'LIGHTPROBE_CUBEMAP')
+    MOL_change_import_interface(row, 'Star File',     5, 487)
 
     panel_selection = bpy.context.scene.mol_import_panel_selection
     col = panel.column()
@@ -439,21 +491,27 @@ def MOL_PT_panel_ui(layout_function, scene):
             box.enabled = False
             box.alert = True
             box.label(text = "Please install biotite in the addon preferences.")
-        MOL_PT_panel_local(box)
+        MOL_PT_panel_esmfold(box)
     elif panel_selection == 2:
+        if not pkg.is_current('biotite'):
+            box.enabled = False
+            box.alert = True
+            box.label(text = "Please install biotite in the addon preferences.")
+        MOL_PT_panel_local(box)
+    elif panel_selection == 3:
         if not pkg.is_current('MDAnalysis'):
             box.enabled = False
             box.alert = True
             box.label(text = "Please install MDAnalysis in the addon preferences.")
 
         MOL_PT_panel_md_traj(box, scene)
-    elif panel_selection == 3:
+    elif panel_selection == 4:
         if not pkg.is_current('mrcfile'):
             box.enabled = False
             box.alert = True
             box.label(text = "Please intall 'mrcfile' in the addon preferences.")
         MOL_PT_panel_map(box, scene)
-    elif panel_selection == 4:
+    elif panel_selection == 5:
         for name in ['starfile', 'eulerangles']:
             if not pkg.is_current(name):
                 box.enabled = False