Update app.py

app.py

@@ -330,41 +330,70 @@ def compute_shap_difference(shap1_norm, shap2_norm):
 
 def plot_comparative_heatmap(shap_diff, seq1_length, seq2_length, title="SHAP Difference Heatmap"):
     """
-    Plots a comparative heatmap of SHAP differences between two sequences. 
-    - The bottom x-axis shows relative positions (%) across the normalized dimension.
-    - The top x-axis shows actual positions for each sequence (S1 and S2).
-    - A vertical colorbar is placed to the right.
-    - Negative (blue) indicates Seq1 is more human-like in that region, 
-      positive (red) indicates Seq2 is more human-like,
-      white indicates no substantial difference.
+    Plots a comparative heatmap of SHAP differences between two sequences.
+
+    Parameters
+    ----------
+    shap_diff     : 1D array of differences (Seq2 SHAP - Seq1 SHAP).
+                    Negative = Seq1 more human-like, Positive = Seq2 more human-like.
+    seq1_length   : int, length of sequence 1 (for labeling).
+    seq2_length   : int, length of sequence 2 (for labeling).
+    title         : str, figure title.
+
+    Figure Layout
+    -------------
+    - Bottom X-Axis: Relative positions in percent (0% to 100%).
+    - Top X-Axis   : Actual positions for both sequences (S1 and S2).
+    - Y-Axis       : Hidden (this is effectively a 1D heatmap).
+    - Colorbar     : Vertical, placed on the right.
+    
+    The final layout uses `tight_layout` with a rectangular constraint to ensure 
+    nothing overlaps while still providing clear labeling.
     """
-    # Prepare data
+    # Reshape the 1D differences into a 1 x N image
     heatmap_data = shap_diff.reshape(1, -1)
+    # Determine the symmetric range around zero
     extent = max(abs(shap_diff.min()), abs(shap_diff.max()))
     
-    # Create figure and axis
-    fig, ax = plt.subplots(figsize=(12, 3))
-    cmap = get_zero_centered_cmap()  # Ensure this function is defined above
-    cax = ax.imshow(heatmap_data, aspect='auto', cmap=cmap, vmin=-extent, vmax=extent)
+    # Create the figure (width x height in inches)
+    fig, ax = plt.subplots(figsize=(10, 3))
+    
+    # Main heatmap
+    cmap = get_zero_centered_cmap()
+    cax = ax.imshow(
+        heatmap_data, 
+        aspect='auto', 
+        cmap=cmap, 
+        vmin=-extent, 
+        vmax=extent
+    )
     
-    # Bottom axis: percentage-based x-axis
+    # Number of major ticks along x-axis
     num_ticks = 5
     tick_positions = np.linspace(0, shap_diff.shape[0] - 1, num_ticks)
+    
+    # ----------------- Bottom Axis: Percentage ----------------- #
     ax.set_xticks(tick_positions)
-    ax.set_xticklabels([f"{int(x * 100)}%" for x in np.linspace(0, 1, num_ticks)], fontsize=9)
-    ax.set_xlabel("Relative Position (%)", fontsize=10)
+    ax.set_xticklabels(
+        [f"{int(x * 100)}%" for x in np.linspace(0, 1, num_ticks)], 
+        fontsize=9
+    )
+    ax.set_xlabel("Relative Position (%)", fontsize=10, labelpad=10)
     
-    # Top axis: actual sequence positions for both Seq1 and Seq2
+    # ----------------- Top Axis: Actual Positions ----------------- #
     ax_top = ax.twiny()
-    ax_top.set_xlim(ax.get_xlim())
+    ax_top.set_xlim(ax.get_xlim())  # Match the bottom axis
+    
+    # Create position arrays for both sequences
     seq1_positions = np.linspace(0, seq1_length, num_ticks)
     seq2_positions = np.linspace(0, seq2_length, num_ticks)
     
+    # Helper function to format large positions nicely
     def format_position(x):
         if x >= 1e6:
             return f"{x / 1e6:.1f}M"
         elif x >= 1e3:
-            return f"{x / 1e3:.0f}K"
+            return f"{int(x / 1e3)}K"
         else:
             return f"{int(x)}"
     
@@ -372,24 +401,33 @@ def plot_comparative_heatmap(shap_diff, seq1_length, seq2_length, title="SHAP Di
     seq2_labels = [format_position(x) for x in seq2_positions]
     
     ax_top.set_xticks(tick_positions)
+    # Each tick label shows the corresponding position in Seq1 and Seq2
     ax_top.set_xticklabels(
         [f"S1: {s1}\nS2: {s2}" for s1, s2 in zip(seq1_labels, seq2_labels)],
-        fontsize=8
+        fontsize=9
     )
     ax_top.set_xlabel("Sequence Positions", fontsize=10, labelpad=15)
     
-    # Colorbar on the right
-    cbar = fig.colorbar(cax, ax=ax, orientation='vertical', fraction=0.02, pad=0.02)
-    cbar.set_label("SHAP Difference (Seq2 - Seq1)", fontsize=10)
+    # ----------------- Colorbar (Vertical, on the right) ----------------- #
+    # 'fraction' = thickness of colorbar, 'pad' = gap from the right edge
+    cbar = fig.colorbar(
+        cax, ax=ax, orientation='vertical', fraction=0.03, pad=0.07
+    )
+    cbar.set_label("SHAP Difference\n(Seq2 - Seq1)", fontsize=10, labelpad=5)
+    
+    # Hide the y-axis (not needed in a 1D heatmap)
+    ax.set_yticks([])
     
-    # Aesthetics
-    ax.set_yticks([])  
+    # Set a descriptive title
     ax.set_title(title, fontsize=12, pad=10)
     
-    fig.tight_layout(rect=[0, 0, 0.88, 1])
+    # Adjust layout so everything fits without overlapping
+    # The rect parameter leaves space on the right for the colorbar
+    fig.tight_layout(rect=[0, 0, 0.9, 1])  # Adjust as necessary
     
     return fig
 
+
 def plot_shap_histogram(shap_array, title="SHAP Distribution", num_bins=30):
     """
     Plot histogram of SHAP values with configurable number of bins