Fix subplot sizing and bug

hexviz/Attention_Visualization.py

@@ -110,7 +110,6 @@ def get_3dview(pdb):
             {"backgroundColor": "lightgray","fontColor": "black","backgroundOpacity": 0.5})
     return xyzview
 
-
 xyzview = get_3dview(pdb_id)
 showmol(xyzview, height=500, width=800)
 st.markdown(f'PDB: [{pdb_id}](https://www.rcsb.org/structure/{pdb_id})', unsafe_allow_html=True)

hexviz/pages/Identify_interesting_Heads.py

@@ -27,14 +27,15 @@ chains = list(structure.get_chains())
 
 sequence = get_sequence(chains[0])
 l = len(sequence)
-slice_start= st.sidebar.number_input(f"Section start(1-{l})",value=1, min_value=1, max_value=l)
-slice_end = st.sidebar.number_input(f"Section end(1-{l})",value=50, min_value=1, max_value=l)
+st.sidebar.markdown("Sequence segment to plot")
+slice_start, slice_end = st.sidebar.slider("Sequence", min_value=1, max_value=l, value=(1, 50), step=1)
+# slice_start= st.sidebar.number_input(f"Section start(1-{l})",value=1, min_value=1, max_value=l)
+# slice_end = st.sidebar.number_input(f"Section end(1-{l})",value=50, min_value=1, max_value=l)
 truncated_sequence = sequence[slice_start-1:slice_end]
 
-
-layer_range = st.sidebar.slider("Heads", min_value=1, max_value=selected_model.layers, value=(1, selected_model.layers), step=1)
-head_range = st.sidebar.slider("Layers", min_value=1, max_value=selected_model.heads, value=(1, selected_model.heads), step=1)
-step_size = st.sidebar.number_input("Step size", value=2, min_value=1, max_value=selected_model.layers)
+head_range = st.sidebar.slider("Heads to plot", min_value=1, max_value=selected_model.heads, value=(1, selected_model.heads), step=1)
+layer_range = st.sidebar.slider("Layers to plot", min_value=1, max_value=selected_model.layers, value=(1, selected_model.layers), step=1)
+step_size = st.sidebar.number_input("Optional step size to skip heads", value=2, min_value=1, max_value=selected_model.layers)
 layer_sequence = list(range(layer_range[0]-1, layer_range[1], step_size))
 head_sequence = list(range(head_range[0]-1, head_range[1], step_size))
 

hexviz/plot.py

@@ -7,10 +7,13 @@ def plot_tiled_heatmap(tensor, layer_sequence: List[int], head_sequence: List[in
     tensor = tensor[layer_sequence, :][:, head_sequence, :, :]  # Slice the tensor according to the provided sequences and sequence_count
     num_layers = len(layer_sequence)
     num_heads = len(head_sequence)
-    fig, axes = plt.subplots(num_layers, num_heads, figsize=(12, 12))
+
+    x_size = num_heads * 2
+    y_size = num_layers * 2
+    fig, axes = plt.subplots(num_layers, num_heads, figsize=(x_size, y_size))
     for i in range(num_layers):
         for j in range(num_heads):
-            axes[i, j].imshow(tensor[i, j].detach().numpy(), cmap='viridis', aspect='auto')
+            axes[i, j].imshow(tensor[i, j].detach().numpy(), cmap='viridis', aspect='equal')
             axes[i, j].axis('off')
 
             # Enumerate the axes