Calculate top residue using attention to residue

Using attention from the residues as well does not makes much sense
as the attention from the residues will always sum to 1, so
it will only be adding 1 to all of the residues total attention.

hexviz/attention.py

@@ -292,11 +292,13 @@ def get_attention_pairs(
                         residue_attention[res - ec_tag_length] = (
                             residue_attention.get(res - ec_tag_length, 0) + attn_value
                         )
-
-        top_n_residues = sorted(residue_attention.items(), key=lambda x: x[1], reverse=True)[:top_n]
-
-        for res, attn_sum in top_n_residues:
-            coord = chain[res]["CA"].coord.tolist()
-            top_residues.append((attn_sum, coord, chain_ids[i], res))
+        if not ec_number:
+            attention_into_res = attention[head, layer].sum(dim=0)
+        else:
+            attention_into_res = attention[head, layer, ec_tag_length:, ec_tag_length:].sum(dim=0)
+        top_n_values, top_n_indexes = torch.topk(attention_into_res, top_n)
+
+        for res, attn_sum in zip(top_n_indexes, top_n_values):
+            top_residues.append((attn_sum.item(), chain_ids[i], res.item()))
 
     return attention_pairs, top_residues

hexviz/🧬Attention_Visualization.py

@@ -71,7 +71,6 @@ n_highest_resis = st.sidebar.number_input(
 )
 label_highest = st.sidebar.checkbox("Label highest attention residues", value=True)
 sidechain_highest = st.sidebar.checkbox("Show sidechains", value=True)
-# TODO add avg or max attention as params
 
 
 with st.sidebar.expander("Label residues manually"):
@@ -238,7 +237,7 @@ def get_3dview(pdb):
             )
 
     if label_highest:
-        for _, _, chain, res in top_residues:
+        for _, chain, res in top_residues:
             one_indexed_res = res + 1
             xyzview.addResLabels(
                 {"chain": chain, "resi": one_indexed_res},
@@ -265,7 +264,7 @@ Pick a PDB ID, layer and head to visualize attention from the selected protein l
 
 chain_dict = {f"{chain.id}": list(chain.get_residues()) for chain in list(structure.get_chains())}
 data = []
-for att_weight, _, chain, resi in top_residues:
+for att_weight, chain, resi in top_residues:
     try:
         res = chain_dict[chain][resi]
     except KeyError:
@@ -273,9 +272,9 @@ for att_weight, _, chain, resi in top_residues:
     el = (att_weight, f"{res.resname:3}{res.id[1]}({chain})")
     data.append(el)
 
-df = pd.DataFrame(data, columns=["Total attention (disregarding direction)", "Residue"])
+df = pd.DataFrame(data, columns=["Total attention to", "Residue"])
 st.markdown(
-    f"The {n_highest_resis} residues (per chain) with the highest attention sums are labeled in the visualization and listed here:"
+    f"The {n_highest_resis} residues (per chain) with the highest attention to them are labeled in the visualization and listed here:"
 )
 st.table(df)