Fixed potential bug

kiwijuice56 · Jan 26, 2024 · 8c3f22a · 8c3f22a
1 parent 1bfe3ca
commit 8c3f22a
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 8 deletions.
diff --git a/gui/button.py b/gui/button.py
@@ -5,7 +5,7 @@ class Button(pyglet.gui.WidgetBase):
     BACKGROUND_COLOR = (0, 0, 0, 140)
     HOVER_COLOR = (64, 64, 64, 140)
 
-    def __init__(self, bounding_box, text, index, window, batch, bg_batch):
+    def __init__(self, bounding_box, text, index, window, batch, bg_batch, text_width=21):
         pyglet.gui.WidgetBase.__init__(self, *bounding_box)
         self.batch = batch
         self.window = window
@@ -61,8 +61,8 @@ def on_mouse_motion(self, x, y, dx, dy):
 
 
 class DropDown(Button):
-    def __init__(self, bounding_box, title, options, window, batch, bg_batch):
-        Button.__init__(self, bounding_box, '%-21s ▾' % options[0], -1, window, batch, bg_batch)
+    def __init__(self, bounding_box, title, options, window, batch, bg_batch, text_width=21):
+        Button.__init__(self, bounding_box, f'%-{text_width}s ▾' % options[0], -1, window, batch, bg_batch)
         self.batch = batch
         self.window = window
         self.window.push_handlers(self.on_mouse_press, self.on_mouse_motion)

diff --git a/gui/user_interface.py b/gui/user_interface.py
@@ -49,7 +49,7 @@ def __init__(self, protein, window, pdb_renderer, embedding_renderer):
             go_titles.append(go_title)
 
         self.go_annotation = DropDown(bounding_box=[608, -80, 416, 32], title="GO Annotation",
-                                      options=go_titles,
+                                      options=go_titles, text_width=32,
                                       window=window, batch=self.batch, bg_batch=self.bg_batch)
 
         self.update_residue_label()

diff --git a/protein.py b/protein.py
@@ -39,7 +39,7 @@ def __init__(self, atoms, bio_residue, index):
         self.highlighted = False
 
         # A string [GO id] : float [0.0, 1.0] pair of how strongly this residue contributed to a certain GO prediction
-        self.go_map = {}
+        self.go_map = {"GO:XXXXXXX": 1.0}
 
 
 # Contains information about a protein, such as its 3D structure and node embeddings
@@ -164,9 +164,14 @@ def __init__(self, pdb_path, chain_id=None, verbose=False):
             self.cluster_index = data["cluster_indices"]
             self.cluster_count = len(set(self.cluster_index)) - (1 if -1 in self.cluster_index else 0)
             self.go_ids = data["GO_ids"]
-            self.go_names = data["GO_names"]
-            self.scores = data["confidence"]
-            self.current_go_id = self.go_ids[0]
+            if len(self.go_ids) == 0:
+                self.go_ids.append("GO:XXXXXXX")
+                self.go_names = ["none"]
+                self.scores = [0.0]
+            else:
+                self.go_names = data["GO_names"]
+                self.scores = data["confidence"]
+            self.current_go_id = self.go_ids[0] if len(self.go_ids) > 0 else "none"
 
             # Assign saliency to each residue in the protein sequence
             for i, annotation in enumerate(self.go_ids):