added introgression genomes of origin; linting

nlbrown62 · nlbrown62 · commit da53114bb268 · 2024-12-11T12:51:17.000-05:00
diff --git a/panagram/experimentation.py b/panagram/experimentation.py
@@ -34,11 +34,13 @@
 # arrs[step] = np.concatenate(arr, axis=1)
 # arrs[step] = np.concatenate(np.array([4,5,6,7]), axis=1)
 
+
 def better_dir(item):
     # don't print hidden functions
     methods = dir(item)
     return [method for method in methods if not method.startswith("_")]
 
+
 # Reading in an index - you can do this after running snakemake
 index_dir = "/home/nbrown62/data_mschatz1/nbrown62/panagram/example_data"
 index = Index(index_dir)
@@ -68,17 +70,19 @@ def better_dir(item):
 # Get part of the bitmap from the genome
 # Query a chromosome at a certain position to get a piece of the bitmap
 # TODO: Not sure what step does yet? Sum of kmer occurrence across a window of size step?
-ecoli_bitmap = ecoli_genome.query('NZ_CP015023.1', 0, 5506781, step=1)
+ecoli_bitmap = ecoli_genome.query("NZ_CP015023.1", 0, 5506781, step=1)
 
 
 # You can also query using the index itself
 # print(index.query_bitmap(ecoli_genome, chrom, start=None, end=None, step=1))
 
 # what if you wanted to look at intergressions
 # set a threshold for what counts as an intergression
-threshold = 0.75 # 75% co-occurrence across the pangenome
+threshold = 0.75  # 75% co-occurrence across the pangenome
 ecoli_bitmap["frac_co_occurrence"] = ecoli_bitmap.sum(axis=1)
-ecoli_bitmap["frac_co_occurrence"] = ecoli_bitmap["frac_co_occurrence"] / len(index.genomes)
+ecoli_bitmap["frac_co_occurrence"] = ecoli_bitmap["frac_co_occurrence"] / len(
+    index.genomes
+)
 print(ecoli_bitmap)
 
 # the index on the left (+ starting point index) gives you the indices of intergressions
@@ -90,6 +94,5 @@ def better_dir(item):
 # TODO: how long are intergressions typically?
 # use a sliding window and take average co-occurance perhaps within a k-sized window
 
-diff_threshold = 10000 # consider different intergression if more than n away?
+diff_threshold = 10000  # consider different intergression if more than n away?
 # index_diff = df['index'].diff().fillna(0)
-
diff --git a/panagram/generate_corr_figure.py b/panagram/generate_corr_figure.py
@@ -1,5 +1,6 @@
 import sys
 import os
+
 # import os.path
 # from os import path
 from pathlib import Path
@@ -35,19 +36,20 @@ def visualize(pair, output_file, inverse=False):
     # take a look at what pair looks like after manipulation
     # pair[pair >= 1] = 10
     if inverse:
-        fig = px.imshow(pair,
-                        color_continuous_scale=px.colors.sequential.Plasma[::-1],
-                        x=pair.columns,
-                        y=pair.index)
+        fig = px.imshow(
+            pair,
+            color_continuous_scale=px.colors.sequential.Plasma[::-1],
+            x=pair.columns,
+            y=pair.index,
+        )
     else:
-        fig = px.imshow(pair,
-                        x=pair.columns,
-                        y=pair.index)
+        fig = px.imshow(pair, x=pair.columns, y=pair.index)
     fig.write_image(output_file)
     return
 
+
 index_dir = "/home/nbrown62/data_mschatz1/nbrown62/panagram_data/tomato"
-anchor = "SL5" #"SL5"
+anchor = "SL5"  # "SL5"
 # chr_name = "BGV006775_MAS2.0ch11"
 output_dir = "/home/nbrown62/data_mschatz1/nbrown62/panagram_data/tomato/introgression_analysis_v1/"
 
@@ -60,7 +62,7 @@ def visualize(pair, output_file, inverse=False):
 # # print(index.genomes)
 bitmap_step = 100
 max_chr_bins = 350
-k=31
+k = 31
 
 for chr_name in chrs:
     # get an entire chr's bitmap
@@ -75,4 +77,6 @@ def visualize(pair, output_file, inverse=False):
 
     pan, pair = index.bitmap_to_bins(chr_bitmap, bin_size)
 
-    visualize(pair, output_dir / f"{anchor}_{chr_name}_original_heatmap.png", inverse=True)
+    visualize(
+        pair, output_dir / f"{anchor}_{chr_name}_original_heatmap.png", inverse=True
+    )
diff --git a/panagram/introgressions.py b/panagram/introgressions.py
@@ -1,26 +1,7 @@
-import sys
-import os
-# import os.path
-# from os import path
 from pathlib import Path
-import subprocess
+import math
 import numpy as np
 import pandas as pd
-import bgzip
-import gzip
-import csv
-import glob
-import pysam
-from collections import defaultdict, Counter
-from time import time
-from Bio import bgzf, SeqIO
-import yaml
-import multiprocessing as mp
-from types import SimpleNamespace
-import shutil
-import snakemake
-import re
-import logging
 from panagram.index import Index
 import plotly.express as px
 
@@ -35,19 +16,20 @@ def visualize(pair, output_file, inverse=False):
     # take a look at what pair looks like after manipulation
     # pair[pair >= 1] = 10
     if inverse:
-        fig = px.imshow(pair,
-                        color_continuous_scale=px.colors.sequential.Plasma[::-1],
-                        x=pair.columns,
-                        y=pair.index)
+        fig = px.imshow(
+            pair,
+            color_continuous_scale=px.colors.sequential.Plasma[::-1],
+            x=pair.columns,
+            y=pair.index,
+        )
     else:
-        fig = px.imshow(pair,
-                        x=pair.columns,
-                        y=pair.index)
+        fig = px.imshow(pair, x=pair.columns, y=pair.index)
     fig.write_image(output_file)
     return
 
 
 def fill_gaps(row, rounds=2):
+    row = row.values
     for j in range(rounds):
         # Find gaps of 0s surrounded by 1s
         for i in range(1, len(row) - 1):
@@ -56,7 +38,17 @@ def fill_gaps(row, rounds=2):
                 row[i] = (row[i - 1] + row[i + 1]) / 2
     return row
 
-def run_introgression_finder(index, anchor, chr_name, bitmap_step, max_chr_bins, k, output_dir):
+
+def run_introgression_finder(
+    index,
+    anchor,
+    chr_name,
+    bitmap_step,
+    max_chr_bins,
+    k,
+    set_difference_threshold,
+    output_dir,
+):
     # Step 1 - choose an anchor and re-create pairwise correlation matrix for it
     # kmer size is 20-30ish; kmer at position X starts at position X (not centered at position X)
     # there are multiple positions in a bin; there are <bin size> - k + 1 kmers in a bin
@@ -83,7 +75,9 @@ def run_introgression_finder(index, anchor, chr_name, bitmap_step, max_chr_bins,
     print("# Kmers in a bin", num_kmers_in_bin)
 
     pan, pair = index.bitmap_to_bins(chr_bitmap, bin_size)
-    visualize(pair, output_dir / f"{anchor}_{chr_name}_original_heatmap.png", inverse=True)
+    visualize(
+        pair, output_dir / f"{anchor}_{chr_name}_original_heatmap.png", inverse=True
+    )
 
     # # sanity check
     # print(len(pair.columns))
@@ -110,64 +104,121 @@ def run_introgression_finder(index, anchor, chr_name, bitmap_step, max_chr_bins,
     print("Chosen dissimilarity threshold", dissimilarity_threshold)
 
     # Histogram of correlations
-    fig = px.histogram(values, title='Histogram of Correlation Values')
-    fig.add_vline(x=dissimilarity_threshold, line_width=2, line_dash="dash", line_color="red", annotation_text="Threshold", annotation_position="top left")
+    fig = px.histogram(values, title="Histogram of Correlation Values")
+    fig.add_vline(
+        x=dissimilarity_threshold,
+        line_width=2,
+        line_dash="dash",
+        line_color="red",
+        annotation_text="Threshold",
+        annotation_position="top left",
+    )
     fig.write_image(output_dir / f"{anchor}_{chr_name}_corr_hist.png")
 
     # NOTE: if there are no outliers, or threshold is very high, there are likely no introgressions
     # test for this scenario
 
     # flip matrix so we can use pandas rolling operation
     transposed_pair = pair.transpose().drop(columns=[anchor])
+    pangenome_names = list(transposed_pair.columns)  # save for reference later
     # print(transposed_pair)
 
     # convert to binary array by applying dissimilarity threshold
     transposed_pair[transposed_pair > dissimilarity_threshold] = 0
     transposed_pair[transposed_pair != 0] = 1
 
     # Step 3 - create an introgression score for each position based on total dissimilarity
-    transposed_pair["introgression_score"] = transposed_pair.sum(axis=1)
-    # transposed_pair["genomes"] = transposed_pair.apply(lambda row: {col for col in transposed_pair.columns if row[col] > 0}, axis=1)
-    # print(transposed_pair)
-
-    # combine nearby locations as the same introgression
-    transposed_pair["introgression_score"] = fill_gaps(transposed_pair["introgression_score"].values)#, transposed_pair["genomes"].values)
-    visualize(transposed_pair.transpose(), output_dir / f"{anchor}_{chr_name}_introgressions_heatmap.png")
+    # for each genome combine nearby locations as the same introgression - column by column
+    transposed_pair = transposed_pair.apply(fill_gaps)
+
+    # For each position, figure out which subset of genomes share an introgression at the position
+    transposed_pair["genomes"] = transposed_pair.apply(
+        lambda row: {col for col in pangenome_names if row[col] > 0}, axis=1
+    )
+    # Calculate set difference with the next row and store the length
+    transposed_pair["genomes_overlap"] = (
+        transposed_pair["genomes"]
+        .shift(1)
+        .combine(
+            transposed_pair["genomes"],
+            lambda next_row, current_row: (
+                len(current_row & (next_row or set()))
+                / len(current_row | (next_row or set()))
+                if current_row | (next_row or set())
+                else 0
+            ),
+        )
+    )
+
+    # introgression score - counter of number of genomes that share the introgression
+    # TODO: normalize by number of genomes in the pangenome?
+    # transposed_pair["introgression_score"] = transposed_pair[pangenome_names].sum(axis=1)
+    # transposed_pair["introgression_score"] = fill_gaps(transposed_pair["introgression_score"].values, transposed_pair["genomes"].values, set_difference_threshold)
+
+    transposed_pair["introgression_score"] = transposed_pair["genomes"].apply(len)
+    visualize(
+        transposed_pair[pangenome_names + ["introgression_score"]].transpose(),
+        output_dir / f"{anchor}_{chr_name}_introgressions_heatmap.png",
+    )
+    pair.loc["intro_score"] = 1 - (
+        transposed_pair["introgression_score"].values
+        / max(transposed_pair["introgression_score"].values)
+    )
+    visualize(
+        pair, output_dir / f"{anchor}_{chr_name}_hybrid_heatmap.png", inverse=True
+    )
 
     # Step 4 - report, sorted by size and then by score
-    transposed_pair['introgression_starts'] = (transposed_pair['introgression_score'] != 0) & (transposed_pair['introgression_score'].shift(1, fill_value=0) == 0)
+    # transposed_pair['introgression_starts'] = (transposed_pair['introgression_score'] != 0) & (transposed_pair['introgression_score'].shift(1, fill_value=0) == 0)
+    transposed_pair["introgression_starts"] = (
+        (transposed_pair["genomes_overlap"] > 0)
+        & (transposed_pair["genomes_overlap"] < 0.9)
+    ) | (
+        (transposed_pair["genomes_overlap"] == 0)
+        & (transposed_pair["introgression_score"] == 1)
+    )
 
     # Create a group identifier for each set of non-zeros
-    transposed_pair['introgression_group'] = transposed_pair['introgression_starts'].cumsum()
+    transposed_pair["introgression_group"] = transposed_pair[
+        "introgression_starts"
+    ].cumsum()
+
     # sum of all introgression scores
-    total_introgression_scores = transposed_pair[transposed_pair['introgression_score'] != 0].groupby('introgression_group')['introgression_score'].mean()
+    introgression_groupby = transposed_pair[
+        transposed_pair["introgression_score"] != 0
+    ].groupby("introgression_group")
+    total_introgression_scores = introgression_groupby["introgression_score"].mean()
     # end index for finding introgression length in bps
-    last_indices = transposed_pair[transposed_pair['introgression_score'] != 0].groupby('introgression_group').tail(1).index.values
+    last_indices = introgression_groupby.tail(1).index.values
+    # genomes that introgression belongs to (union of all sets in the group)
+    introgression_genomes = (
+        introgression_groupby["genomes"].agg(lambda sets: set.union(*sets)).values
+    )
+
+    introgressions = transposed_pair[
+        transposed_pair.introgression_starts == True
+    ].copy()
 
-    introgressions = transposed_pair[transposed_pair.introgression_starts == True].copy()
+    # print(transposed_pair[transposed_pair['introgression_score'] != 0][transposed_pair.columns[10:20]])
     introgressions["introgression_score"] = total_introgression_scores.values
     introgressions["introgression_end"] = last_indices
-    introgressions["introgression_end"] = introgressions["introgression_end"] + bin_size # account for the fact that the index is the start and not the end of a bin
-    introgressions = introgressions[["introgression_end", "introgression_score"]].reset_index().rename(columns={"index":"introgression_start"})
-    introgressions["introgression_length"] = introgressions["introgression_end"] - introgressions["introgression_start"]
-    print(introgressions.sort_values(by=["introgression_length", "introgression_score"], ascending=False))
-    introgressions.sort_values(by=["introgression_length", "introgression_score"], ascending=False).to_csv(output_dir / f"{anchor}_{chr_name}_introgressions.csv", index=False)
-    # NOTE: don't need this loop anymore; df operations are faster
-    # introgression_locations = {}
-    # for location in transposed_pair["introgression_score"]:
-    #     print(location)
-        # if bin >= 1 and not in_introgression:
-        #     # start new introgression
-        #     in_introgression = True
-        #     introgression_locations = [locatio]
-        # elif bin >= 1 and in_introgression:
-        #     # append
-        # else:
-            # in_introgression = False
-
-    # TODO: turn into function that can repeat across all chrs/all genomes if requested
-    # Look at underlying sequence in found introgressions; compare/cluster? align with annotations?
-    # we're using dissimilarity to find them though, so by definition, won't these mostly look unique?
+    introgressions["introgression_end"] = (
+        introgressions["introgression_end"] + bin_size
+    )  # account for the fact that the index is the start and not the end of a bin
+    introgressions = (
+        introgressions[["introgression_end", "introgression_score"]]
+        .reset_index()
+        .rename(columns={"index": "introgression_start"})
+    )
+    introgressions["introgression_length"] = (
+        introgressions["introgression_end"] - introgressions["introgression_start"]
+    )
+    introgressions["introgression_genomes"] = introgression_genomes
+    print(introgressions.sort_values(by=["introgression_start"], ascending=True))
+    introgressions.sort_values(by=["introgression_start"], ascending=True).to_csv(
+        output_dir / f"{anchor}_{chr_name}_introgressions.csv", index=False
+    )
+
     return
 
 
@@ -177,29 +228,40 @@ def run_introgression_finder(index, anchor, chr_name, bitmap_step, max_chr_bins,
 # chr_name = "BGV006775_MAS2.0ch11"
 bitmap_step = 100
 max_chr_bins = 350
-size_threshold = 3000000 # NOTE: unused, minimum size in bps of the introgression
-k = 31 # TODO: k should be defined somewhere else; don't need from the user
+size_threshold = 3000000  # NOTE: unused, minimum size in bps of the introgression
+k = 31  # TODO: k should be defined somewhere else; don't need from the user
 output_dir = "/home/nbrown62/data_mschatz1/nbrown62/panagram_data/tomato/introgression_analysis_v1/"
 
 index = Index(index_dir)
+set_difference_threshold = 2  # NOTE: could also make this proportionate to pangenome size int(len(index.genomes) / 10)
 
-# For testing
-# for anchor in ["SL5"]:#index.genomes.keys():
-#     genome = index.genomes[anchor]
-#     print(genome.sizes.keys())
-#     for chr_name in [11]:#genome.sizes.keys():
-#         print(anchor, chr_name)
-#         run_introgression_finder(index, anchor, chr_name, bitmap_step, max_chr_bins, k, output_dir)
-#         break
-#     break
-
-for anchor in index.genomes.keys():
+# For testing with tomato pangenome
+for anchor in ["SL5"]:
     genome = index.genomes[anchor]
     print(genome.sizes.keys())
-    for chr_name in genome.sizes.keys():
+    for chr_name in [11]:
         print(anchor, chr_name)
-        run_introgression_finder(index, anchor, chr_name, bitmap_step, max_chr_bins, k, output_dir)
+        run_introgression_finder(
+            index,
+            anchor,
+            chr_name,
+            bitmap_step,
+            max_chr_bins,
+            k,
+            set_difference_threshold,
+            output_dir,
+        )
+        break
+    break
+
+# for anchor in index.genomes.keys():
+#     genome = index.genomes[anchor]
+#     for chr_name in genome.sizes.keys():
+#         print("Now running introgression analysis for", anchor, chr_name)
+#         run_introgression_finder(index, anchor, chr_name, bitmap_step, max_chr_bins, k, output_dir)
 
 # TODO: column for each position with a python set of genomes contributing to the introgression score
 # column for each position with a python set of genomes that may share the same introgression (similarity above threshold)
 # only merge adjacent areas where the set difference is <=2
+# TODO: Look at underlying sequence in found introgressions; compare/cluster? align with annotations?
+# we're using dissimilarity to find them though, so by definition, won't these mostly look unique?
