Update app.py

app.py

@@ -330,76 +330,64 @@ def compute_shap_difference(shap1_norm, shap2_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
+    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.
     """
-    
-    # Prepare data for the heatmap
+    # Prepare data
     heatmap_data = shap_diff.reshape(1, -1)
-    extent = max(abs(np.min(shap_diff)), abs(np.max(shap_diff)))
+    extent = max(abs(shap_diff.min()), abs(shap_diff.max()))
     
-    # Create figure and main axis
+    # 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)
     
-    # Plot the main heatmap
-    cmap = get_zero_centered_cmap()
-    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)
-    
-    # Configure the top axis for relative (%) positions
+    # Bottom axis: percentage-based x-axis
     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)]
     ax.set_xticks(tick_positions)
-    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 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
+    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)
     
-    # Prepare positions for each sequence
+    # Top axis: actual sequence positions for both Seq1 and Seq2
+    ax_top = ax.twiny()
+    ax_top.set_xlim(ax.get_xlim())
     seq1_positions = np.linspace(0, seq1_length, num_ticks)
     seq2_positions = np.linspace(0, seq2_length, num_ticks)
     
-    # Format large numbers with 'K' or 'M'
     def format_position(x):
         if x >= 1e6:
-            return f"{x/1e6:.1f}M"
+            return f"{x / 1e6:.1f}M"
         elif x >= 1e3:
-            return f"{x/1e3:.0f}K"
+            return f"{x / 1e3:.0f}K"
         else:
-            return str(int(x))
+            return f"{int(x)}"
     
     seq1_labels = [format_position(x) for x in seq1_positions]
     seq2_labels = [format_position(x) for x in seq2_positions]
     
-    # 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(bottom_labels, fontsize=9)
-    ax2.set_xlabel('Sequence Positions', fontsize=10, labelpad=10)
+    ax_top.set_xticks(tick_positions)
+    ax_top.set_xticklabels(
+        [f"S1: {s1}\nS2: {s2}" for s1, s2 in zip(seq1_labels, seq2_labels)],
+        fontsize=8
+    )
+    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)
+    
+    # Aesthetics
+    ax.set_yticks([])  
+    ax.set_title(title, fontsize=12, pad=10)
+    
+    fig.tight_layout(rect=[0, 0, 0.88, 1])
     
-    # Use tight_layout to reduce overlap
-    plt.tight_layout()
     return fig
 
 def plot_shap_histogram(shap_array, title="SHAP Distribution", num_bins=30):