import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

# Load the coverage data for all 8 datasets
positive_strands = []
negative_strands = []
for i in range(1, 9):
    positive_strands.append(np.load(f'positive_strand_coverage{i}.npy'))
    negative_strands.append(np.load(f'negative_strand_coverage{i}.npy') * -1)  # Invert negative strand values

# Load filtered data
positive_filtered = np.load('positive_filtered.npy')
negative_filtered = np.load('negative_filtered.npy') * -1  # Invert for plotting

# Load sigma70 probability data
s70_probs = np.load('s70-probs.npy')
s70_probs_rev = np.load('s70-probs-rev.npy')

# Start-end coordinates
cus_start, cus_end = 1, 10000

# ---------- Prep x-axis and subset ----------
x_start, x_end = cus_start, cus_end
x = np.arange(x_start, x_end)

# Subset the data for all datasets
pos_subsets = [pos[x_start:x_end] for pos in positive_strands]
neg_subsets = [neg[x_start:x_end] for neg in negative_strands]

# Subset filtered data
pos_filtered_subset = positive_filtered[x_start:x_end]
neg_filtered_subset = negative_filtered[x_start:x_end]

# Subset sigma70 probability data
s70_probs_subset = s70_probs[x_start:x_end]
s70_probs_rev_subset = s70_probs_rev[x_start:x_end]

# Calculate global maximum for rectangle height
global_max_coverage = max(
    max([subset.max() for subset in pos_subsets]),
    max([abs(subset.min()) for subset in neg_subsets]),
    pos_filtered_subset.max(),
    abs(neg_filtered_subset.min())
) * 1.05

# ---------- Create the plot with two y-axes ----------
fig, ax1 = plt.subplots(figsize=(12, 6))

# Create a second y-axis that shares the same x-axis
ax2 = ax1.twinx()

# Define dataset groups with colors and labels
datasets = [
    {'index': 0, 'group': 'IVT (1)', 'color': '#0000ff'},  # Dark blue
    {'index': 1, 'group': 'IVT (2)', 'color': '#7755ff'},  # Light blue
    {'index': 2, 'group': 'IVT 10x DNA (3)', 'color': '#66aaee'},  # Green-blue
    {'index': 3, 'group': 'No polymerase (4)', 'color': '#d62728'},  # Dark red
    {'index': 4, 'group': 'No polymerase (5)', 'color': '#ff7f0e'},  # Medium red
    {'index': 5, 'group': 'No polymerase 10x DNA (6)', 'color': '#ff0088'},  # Light red
    {'index': 6, 'group': 'DNA (7)', 'color': '#17fe1f'},  # Dark green
    {'index': 7, 'group': 'DNA (8)', 'color': '#9efa15'}   # Light green
]

# Plot all datasets on the first y-axis
for dataset in datasets:
    i = dataset['index']
    ax1.plot(x, pos_subsets[i], color=dataset['color'], linewidth=1, label=dataset['group'])
    ax1.plot(x, neg_subsets[i], color=dataset['color'], linewidth=1)

regions = []

# Add region rectangles if needed
for r in regions:
    # Positive strand rectangle
    ax1.add_patch(
        Rectangle(
            (r['start'], 0),
            r['end'] - r['start'],
            global_max_coverage,
            facecolor=r['color'],
            alpha=0.75,
            label=r['label']
        )
    )
    # Negative strand rectangle
    ax1.add_patch(
        Rectangle(
            (r['start'], -global_max_coverage),
            r['end'] - r['start'],
            global_max_coverage,
            facecolor=r['color'],
            alpha=0.75
        )
    )

# Plot filtered data with grey fill on the first y-axis
ax1.plot(x, pos_filtered_subset, color='grey', linewidth=1, label='Filtered')
ax1.plot(x, neg_filtered_subset, color='grey', linewidth=1)
ax1.fill_between(x, pos_filtered_subset, neg_filtered_subset, color='grey', alpha=0.7)

# Plot sigma70 probability data on the second y-axis
ax2.plot(x, s70_probs_subset, color='purple', linewidth=1.5, label='σ70 probability (+)')
ax2.plot(x, -s70_probs_rev_subset, color='magenta', linewidth=1.5, label='σ70 probability (-)')

# ---------- Adjust plot cosmetics ----------
ax1.set_ylim(bottom=-global_max_coverage, top=global_max_coverage)
ax2.set_ylim(bottom=-1, top=1)  # Set second y-axis from -1 to 1

ax1.axhline(0, color='black', linewidth=0.5)  # Add zero line

ax1.ticklabel_format(style='plain', axis='x', useOffset=False)
ax1.set_xticks(np.arange(x_start, x_end+1, 500))
ax1.set_xticklabels(np.arange(x_start, x_end+1, 500), rotation=45)
ax1.set_xlabel('Position')
ax1.set_ylabel('Coverage (positive and negative strands)')
ax2.set_ylabel('σ70 probability')

plt.title('Strand-specific coverage and σ70 probability across genetic elements')

# Combine legends from both axes
handles1, labels1 = ax1.get_legend_handles_labels()
handles2, labels2 = ax2.get_legend_handles_labels()
all_handles = handles1 + handles2
all_labels = labels1 + labels2

# Deduplicate legend entries
unique_labels = dict(zip(all_labels, all_handles))
ax1.legend(unique_labels.values(), unique_labels.keys(), loc='upper right')

plt.tight_layout()
plt.show()


How to make s70-probs.npy and s70-probs-rev.npy to ^2?