Allow selecting heads to plot

hexviz/pages/Identify_interesting_Heads.py

@@ -32,11 +32,18 @@ slice_end = st.sidebar.number_input(f"Section end(1-{l})",value=50, min_value=1,
 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)
+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))
+
+
 st.markdown(f"Each tile is a heatmap of attention for a section of {pdb_id}(A) from residue {slice_start} to {slice_end}. Adjust the section length and starting point in the sidebar.")
 
 # TODO: Decide if you should get attention for the full sequence or just the truncated sequence
 # Attention values will change depending on what we do.
 attention = get_attention(sequence=truncated_sequence, model_type=selected_model.name)
 
-fig = plot_tiled_heatmap(attention, layer_count=selected_model.layers, head_count=selected_model.heads)
+fig = plot_tiled_heatmap(attention, layer_sequence=layer_sequence, head_sequence=head_sequence)
 st.pyplot(fig)

hexviz/plot.py

@@ -1,10 +1,12 @@
+from typing import List
+
 import matplotlib.pyplot as plt
 
 
-def plot_tiled_heatmap(tensor, layer_count=12, head_count=12):
-    tensor = tensor[:layer_count:2, :head_count:2, :, :]  # Slice the tensor according to the provided arguments
-    num_layers = layer_count // 2
-    num_heads = head_count // 2
+def plot_tiled_heatmap(tensor, layer_sequence: List[int], head_sequence: List[int]):
+    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))
     for i in range(num_layers):
         for j in range(num_heads):
@@ -13,11 +15,11 @@ def plot_tiled_heatmap(tensor, layer_count=12, head_count=12):
 
             # Enumerate the axes
             if i == 0:
-                axes[i, j].set_title(f'Head {j * 2 + 1}', fontsize=10, y=1.05)
+                axes[i, j].set_title(f'Head {head_sequence[j] + 1}', fontsize=10, y=1.05)
 
     # Add layer labels on the right Y-axis
     for i in range(num_layers):
-        fig.text(0.98, (num_layers - i - 1) / num_layers + 0.025, f'Layer {i * 2 + 1}', fontsize=10, rotation=0, ha='right', va='center')
+        fig.text(0.98, (num_layers - i - 1) / num_layers + 0.025, f'Layer {layer_sequence[i]+1}', fontsize=10, rotation=0, ha='right', va='center')
 
     plt.subplots_adjust(wspace=0.1, hspace=0.1)
-    return fig
+    return fig