Show EC number in single attention plot and add option to show special tokens

hexviz/pages/1_🗺️Identify_Interesting_Heads.py

@@ -1,6 +1,8 @@
+import re
+
 import streamlit as st
 
-from hexviz.attention import clean_and_validate_sequence, get_attention, get_sequence
+from hexviz.attention import clean_and_validate_sequence, get_attention, res_to_1letter
 from hexviz.config import URL
 from hexviz.models import Model, ModelType
 from hexviz.plot import plot_single_heatmap, plot_tiled_heatmap
@@ -51,11 +53,36 @@ chain_selection = st.sidebar.selectbox(
 )
 
 selected_chain = next(chain for chain in chains if chain.id == chain_selection)
-sequence = get_sequence(selected_chain)
+
+ec_number = ""
+if selected_model.name == ModelType.ZymCTRL:
+    st.sidebar.markdown(
+        """
+        ZymCTRL EC number
+        ---
+        """
+    )
+    try:
+        ec_number = structure.header["compound"]["1"]["ec"]
+    except KeyError:
+        pass
+    ec_number = st.sidebar.text_input("Enzyme Comission number (EC)", ec_number)
+
+    # Validate EC number
+    if not re.match(r"^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$", ec_number):
+        st.sidebar.error(
+            """Please enter a valid Enzyme Commission number in the format of 4
+            integers separated by periods (e.g., 1.2.3.21)"""
+        )
+
+
+residues = [res for res in selected_chain.get_residues()]
+sequence = res_to_1letter(residues)
 
 l = len(sequence)
 slice_start, slice_end = select_sequence_slice(l)
 truncated_sequence = sequence[slice_start - 1 : slice_end]
+remove_special_tokens = st.sidebar.checkbox("Remove special tokens", key="remove_special_tokens")
 
 
 layer_sequence, head_sequence = select_heads_and_layers(st.sidebar, selected_model)
@@ -68,11 +95,10 @@ st.markdown(
 
 # 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(
 attention, tokens = get_attention(
     sequence=truncated_sequence,
     model_type=selected_model.name,
-    remove_special_tokens=True,
+    remove_special_tokens=remove_special_tokens,
     ec_number=ec_number,
 )
 
@@ -111,8 +137,5 @@ with right:
         unsafe_allow_html=True,
     )
 
-
-single_head_fig = plot_single_heatmap(
-    attention, layer, head, slice_start, slice_end, max_labels=10
-)
+single_head_fig = plot_single_heatmap(attention, layer, head, tokens=tokens)
 st.pyplot(single_head_fig)

hexviz/plot.py

@@ -2,7 +2,7 @@ from typing import List
 
 import matplotlib.pyplot as plt
 import numpy as np
-from matplotlib.ticker import MaxNLocator, MultipleLocator
+from matplotlib.ticker import FixedLocator
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 
 
@@ -40,42 +40,24 @@ def plot_single_heatmap(
     tensor,
     layer: int,
     head: int,
-    slice_start: int,
-    slice_end: int,
-    max_labels: int = 40,
+    tokens: list[str],
 ):
     single_heatmap = tensor[layer, head, :, :].detach().numpy()
 
     fig, ax = plt.subplots(figsize=(10, 10))
     heatmap = ax.imshow(single_heatmap, cmap="viridis", aspect="equal")
 
-    # Set the x and y axis major ticks and labels
-    ax.xaxis.set_major_locator(
-        MaxNLocator(integer=True, steps=[1, 2, 5], prune="both", nbins=max_labels)
-    )
-    ax.yaxis.set_major_locator(
-        MaxNLocator(integer=True, steps=[1, 2, 5], prune="both", nbins=max_labels)
-    )
-
-    tick_indices_x = np.clip((ax.get_xticks()).astype(int), 0, slice_end - slice_start)
-    tick_indices_y = np.clip((ax.get_yticks()).astype(int), 0, slice_end - slice_start)
-    ax.set_xticklabels(
-        np.arange(slice_start, slice_end + 1)[tick_indices_x], fontsize=8
-    )
-    ax.set_yticklabels(
-        np.arange(slice_start, slice_end + 1)[tick_indices_y], fontsize=8
-    )
-
-    # Set the x and y axis minor ticks
-    ax.xaxis.set_minor_locator(MultipleLocator(1))
-    ax.yaxis.set_minor_locator(MultipleLocator(1))
+    # Set the x and y axis ticks
+    ax.xaxis.set_major_locator(FixedLocator(np.arange(0, len(tokens))))
+    ax.yaxis.set_major_locator(FixedLocator(np.arange(0, len(tokens))))
 
-    # Set the axis labels
-    ax.set_xlabel("Residue Number")
-    ax.set_ylabel("Residue Number")
+    # Set tick labels as sequence values
+    ax.set_xticklabels(tokens, fontsize=8, rotation=45, ha="right", rotation_mode="anchor")
+    ax.set_yticklabels(tokens, fontsize=8)
 
-    # Rotate the tick labels and set their alignment
-    plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
+    # Set the axis labels
+    ax.set_xlabel("Sequence tokens")
+    ax.set_ylabel("Sequence tokens")
 
     # Create custom colorbar axes with the desired dimensions
     divider = make_axes_locatable(ax)