Update app.py

app.py

@@ -329,64 +329,77 @@ def compute_shap_difference(shap1_norm, shap2_norm):
     return shap2_norm - shap1_norm
 
 def plot_comparative_heatmap(shap_diff, seq1_length, seq2_length, title="SHAP Difference Heatmap"):
+    """
+    Plot a comparative SHAP-difference heatmap, showing how Sequence 2 differs
+    from Sequence 1 with respect to 'human-likeness'.
+    
+    shap_diff:   1D array of (Seq2 - Seq1) SHAP differences, already normalized to the same length
+    seq1_length: length of sequence 1 (in bases)
+    seq2_length: length of sequence 2 (in bases)
+    title:       title of the plot
+    """
+    
+    # Prepare data for the heatmap
     heatmap_data = shap_diff.reshape(1, -1)
     extent = max(abs(np.min(shap_diff)), abs(np.max(shap_diff)))
     
-    # Create figure with additional space
-    fig, ax = plt.subplots(figsize=(12, 3.2))
+    # Create figure and main axis
+    fig, ax = plt.subplots(figsize=(12, 3))
     
-    # Plot main heatmap
+    # Plot the main heatmap
     cmap = get_zero_centered_cmap()
-    cax = ax.imshow(heatmap_data, aspect='auto', cmap=cmap, vmin=-extent, vmax=extent)
+    cax = ax.imshow(
+        heatmap_data, 
+        aspect='auto', 
+        cmap=cmap, 
+        vmin=-extent, 
+        vmax=extent
+    )
+    
+    # Add a vertical colorbar on the right
+    cbar = plt.colorbar(cax, ax=ax, orientation='vertical', fraction=0.025, pad=0.03)
+    cbar.ax.tick_params(labelsize=9)
+    cbar.set_label('SHAP Difference (Seq2 - Seq1)', fontsize=9, labelpad=5)
     
-    # Create percentage-based x-axis ticks (top)
+    # Configure the top axis for relative (%) positions
     num_ticks = 5
-    tick_positions = np.linspace(0, shap_diff.shape[0]-1, num_ticks)
-    tick_labels = [f"{int(x*100)}%" for x in np.linspace(0, 1, num_ticks)]
+    tick_positions = np.linspace(0, shap_diff.shape[0] - 1, num_ticks)
+    tick_labels = [f"{int(x * 100)}%" for x in np.linspace(0, 1, num_ticks)]
     ax.set_xticks(tick_positions)
-    ax.set_xticklabels(tick_labels)
+    ax.set_xticklabels(tick_labels, fontsize=9)
+    
+    ax.set_xlabel('Relative Position (%)', fontsize=10, labelpad=8)
+    ax.set_title(title, fontsize=12, pad=10)
+    ax.set_yticks([])  # Hide the y-axis ticks for a 1D heatmap
     
-    # Create second x-axis for actual positions (bottom)
-    ax2 = ax.twiny()
-    ax2.set_xlim(ax.get_xlim())
+    # Create a second (bottom) x-axis for actual positions
+    ax2 = ax.secondary_xaxis('bottom')
+    ax2.set_xlim(ax.get_xlim())  # Match the same data range as the top axis
     
-    # Calculate actual positions for both sequences
+    # Prepare positions for each sequence
     seq1_positions = np.linspace(0, seq1_length, num_ticks)
     seq2_positions = np.linspace(0, seq2_length, num_ticks)
     
-    # Format position labels with appropriate scaling
+    # Format large numbers with 'K' or 'M'
     def format_position(x):
         if x >= 1e6:
             return f"{x/1e6:.1f}M"
         elif x >= 1e3:
             return f"{x/1e3:.0f}K"
         else:
-            return f"{int(x)}"
+            return str(int(x))
     
     seq1_labels = [format_position(x) for x in seq1_positions]
     seq2_labels = [format_position(x) for x in seq2_positions]
     
-    # Set positions for bottom axis
+    # Combine the two label sets in a single tick label
+    bottom_labels = [f"S1: {s1}\nS2: {s2}" for s1, s2 in zip(seq1_labels, seq2_labels)]
     ax2.set_xticks(tick_positions)
-    ax2.set_xticklabels([f"S1: {s1}\nS2: {s2}" for s1, s2 in zip(seq1_labels, seq2_labels)])
-    
-    # Add colorbar at the bottom with more spacing
-    cbar = plt.colorbar(cax, orientation='horizontal', pad=0.6, aspect=40, shrink=0.8)
-    cbar.ax.tick_params(labelsize=8)
-    cbar.set_label('SHAP Difference (Seq2 - Seq1)', fontsize=9, labelpad=5)
-    
-    # Remove percentage labels from colorbar
-    cbar.ax.set_xticklabels([f'{x:.3f}' for x in cbar.ax.get_xticks()])
-    
-    # Adjust labels and layout
-    ax.set_yticks([])
-    ax.set_xlabel('Relative Position (%)', fontsize=10, labelpad=10)
+    ax2.set_xticklabels(bottom_labels, fontsize=9)
     ax2.set_xlabel('Sequence Positions', fontsize=10, labelpad=10)
-    ax.set_title(title, pad=10)
-    
-    # Increase bottom margin to prevent overlap
-    plt.subplots_adjust(bottom=0.55, left=0.05, right=0.95, top=0.85)
     
+    # Use tight_layout to reduce overlap
+    plt.tight_layout()
     return fig
 
 def plot_shap_histogram(shap_array, title="SHAP Distribution", num_bins=30):