Manhattan Plots

def manplot(df, p_col, chr_col, pos_col, log_col = None, gene_col = None, p_threshold = None, gene_highlights = None, size = 1, colors = ["darkgray", "black"]):
    #ASSERTIONS 
    assert type(df) == pd.DataFrame, "INPUT DATAFRAME MUST BE A PANDAS DATAFRAME"

    ####################
    ### RUN FUNCTION ###
    ####################

    # REPLACE X Y MT M by 23, 24, 25
    df[chr_col] = df[chr_col].replace("chr", "") # REPLACE chr if chr in contig name
    df[chr_col] = df[chr_col].replace(["X", "Y", "M", "MT"], [23, 24, 25, 25])
    df[chr_col] = df[chr_col].astype("int")


    # SORT BY COORDINATES
    df = df.sort_values(by = [chr_col, pos_col], ascending= (True, True))

    # GET SIGNIFICANCE COLUMNS
    if log_col:
        sig_col = df[log_col]
    else:
        sig_col =  - log10(df[p_col])
    df["-log10(p)"] = sig_col

    # GROUP BY CHROMOSOMES
    groupped = df.groupby(by = chr_col)
    new_start = 0

    #LOOP OVER ALL CHROMOSOMES
    for tup in groupped:
        chrom = tup[0]
        chrom_df = tup[1]

        position_series = chrom_df[pos_col]
        #GET START AND END OF CURRENT CHROMOSOME
        start = position_series.min()
        end = position_series.max()
        
        #UPDATE THE START WITH END OF PREVIOUS CHR
        x_series = position_series + new_start
        
        #GET LOG10 COLUMNS
        effect_series = chrom_df["-log10(p)"]
        
        # PLOT DIFFERENT COLORS FOR ALTERNATING CHROMOSOMES
        if chrom % 2 == 0:
            sns.scatterplot(x = x_series, y = effect_series, color = colors[0],  s=1*size)
        else:
            sns.scatterplot(x = x_series, y = effect_series, color = colors[1],  s=1*size)
            
        # FIND MIDDLE OF POSITION TO PRINT CHROMOSOME NUMBER 
        chrom_label_x =  (position_series.max() -position_series.min())/2 + new_start
        chrom_label_y = - 0.5
        
        plt.text(chrom_label_x,chrom_label_y, "{}".format(chrom).replace("23", "X").replace("24", "Y").replace("25", "M"))

        # HIGHLIGHT SIGNIFICANT ONES (-logp > 4)
        if p_threshold:
            sig_df = chrom_df[["-log10(p)", pos_col]]
            sig_df = sig_df[sig_df["-log10(p)"] >= -log10(p_threshold)]
            sig_effect = sig_df["-log10(p)"]
            sig_pos = sig_df[pos_col] + new_start
            sns.scatterplot(x = sig_pos, y = sig_effect, color = "red",  s=3*size)
        
        # HIGHLIGHT GENES
        if gene_highlights:
            for gene in gene_highlights:
                if gene in list(set(chrom_df[gene_col])):
                    gene_df = chrom_df[chrom_df[gene_col] == gene]

                    sns.scatterplot(x = gene_df[pos_col] + new_start, y = gene_df["-log10(p)"], color = "green", s= 3*size)
                    plt.text( list(gene_df[pos_col])[0] + new_start, gene_df["-log10(p)"].max(), gene , horizontalalignment='left', size=5*size, color='blue', weight='semibold')

        # UPDATE THE NEW STARTING POINT OF NEXT CHROMOSOME
        new_start = new_start + end
        
        # PAINT SIGNIFICANCE LINE
        if p_threshold:
            plt.hlines(-log10(p_threshold), 0, new_start, color = "k", alpha = 0.75, linestyles = "dotted")
            
    sns.despine(top = True, left = False, bottom = False, right = True, trim=True)
    plt.xticks([])
    plt.xlabel("Chromosome")
    plt.show()

df = pd.read_csv("input_manplot.csv", index_col = 0)

manplot(df = df, # Dataframe containing all data
        p_col = "unadjusted_p", # Column name containing unadjusted p values
        chr_col = "Chromosome", # Column name containing Chromosome names
        pos_col = "Position", # Column name containing position  in base pair
        size=2, # Relative size of all elements of plot
        p_threshold=0.0000000005, # SIgnificance threshold (unadjusted)
        gene_col="Gene", # Column name containing Gene Names
        gene_highlights=["HLA-DRB5"], # Highlight SNPs from thoses genes (list)
        colors=["blue", "red"]) # Alternating Colors between chromosomes