load_data.py

import os.path
import csv
import random
## import traceback
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from bisect import bisect_right, bisect_left
from collections import namedtuple, Counter
from intervaltree import Interval, IntervalTree
from scipy.stats import beta
from pyfaidx import Fasta

SHARE_DIR = "/mnt/MS_Archive_one/Hendrik/temp"
DATA_DIR = "/mnt/bigdata/Hendrik/Data/GLORI-Predict"
CELL_LINES = {"A": "A549", "C": "CAOV3", "E": "EFO21", "F": "FADU", "H": "H1792", "O": "OE21", "S": "SKOV3",
               "V": "COV644"}
TREATMENTS = {"D": "untreated", "M": "low-dose", "N": "medium-dose", "T": "high-dose"}

Exon = namedtuple("Exon", ("start", "end", "attrs"))
SampleRegion = namedtuple("SampleRegion",
                          ("seed", "exon", "strand", "start", "end", "seq", "m6A_proportions", "read_counts"))

class SampleLists():
    "Manage the lists of samples that were processed in the two batches"

    def __init__(self):
        self.samples1 = self.get_sample_list(1)
        self.samples2 = self.get_sample_list(2)

    def get(self, batch=(1, 2)):
        if batch == 1:
            return self.samples1
        elif batch == 2:
            return self.samples2
        else:
            return (self.samples1, self.samples2)

    def find_batch(self, cell_line, treatment):
        replicate = f"{cell_line}_{treatment}_1"
        in1 = replicate in self.samples1
        in2 = replicate in self.samples2
        if in1 and in2:
            return (1, 2)
        elif in1:
            return 1
        elif in2:
            return 2
        else:
            return 0

    @staticmethod
    def translate_batch2(samples):
        result = []
        for sample in samples:
            parts = list(sample)
            cell_line = CELL_LINES[parts[0]]
            treatment = TREATMENTS[parts[1]]
            replicate = parts[2]
            result.append(f"{cell_line}_{treatment}_{replicate}")
        return result

    @staticmethod
    def get_sample_list(batch=1):
        assert batch in (1, 2)
        bam_list = os.path.join(SHARE_DIR, f"all_alignGenome_gtoa_bamfiles_batch{batch}.txt")
        samples = []
        with open(bam_list) as source:
            while line := source.readline().rstrip():
                filename = os.path.basename(line)
                samples.append(filename.replace("_alignGenome_gtoa.bam", ""))
        if batch == 1: # all treatments are medium dose
            samples = [sample.replace("DMSO", "untreated") for sample in samples]
            samples = [sample.replace("STM3675", "medium-dose") for sample in samples]
        else: # translate short codes
            samples = SampleLists.translate_batch2(samples)
        return samples


class PileupLoader:
    "Load read pileup data from one chromosome"
    def __init__(self, chromosome, cell_line, treatment, batch=1, samples=SampleLists()):
        self.chromosome = chromosome
        self.cell_line = cell_line
        self.treatment = treatment
        self.samples = samples
        if treatment != "untreated":
            ## every cell line/treatment combination was only measured in one batch - look it up:
            batch = samples.find_batch(cell_line, treatment)
        self.batch = batch
        ## read "mpileup" results:
        pileup_dir = os.path.join(SHARE_DIR, f"batch{batch}_mpileup", "allsamples_alignGenome_mpileup")
        replicates = [f"{cell_line}_{treatment}_{rep}" for rep in (1, 2)]
        batch_samples = samples.get(batch)
        sample_indexes = [batch_samples.index(rep) for rep in replicates]
        ## results from positive strand:
        pos_file = f"allsamples_alignGenome_all_mpileup_apositions_exons_chrp_{chromosome}.txt"
        self.pileup_pos = self.read_pileup(os.path.join(pileup_dir, pos_file), sample_indexes)
        ## results from negative strand (positions are counting "from the end"!):
        neg_file = f"allsamples_alignGenome_all_mpileup_apositions_exons_chrn_{chromosome}.txt"
        self.pileup_neg = self.read_pileup(os.path.join(pileup_dir, neg_file), sample_indexes)

    @staticmethod
    def get_AG_counts(item, a_index):
        parts = item.split(",")
        g_count = int(parts[0])
        if (a_index < 0) or ((a_count := int(parts[a_index + 1])) == 0):
            return g_count # just the G count (A count is 0)
        return (a_count, g_count)

    @staticmethod
    def read_pileup(path, sample_indexes):
        ## columns: 0. chromosome/strand, 1. position, 2. reference base, 3. variant bases,
        ## 4.-?. (per sample): total counts, counts per base
        result = []
        with open(path) as source:
            reader = csv.reader(source, delimiter="\t", quoting=csv.QUOTE_NONE)
            for row in reader:
                if row[2] != "G": # reference must be "G" (not "GT...")
                    continue
                totals = [int(row[2 * si + 4]) for si in sample_indexes]
                if sum(totals) == 0: # no reads for the sample(s) of interest
                    continue
                try:
                    a_index = row[3].split(",").index("A") # e.g. "A,<*>"
                except ValueError: # no A variant at this position
                    a_index = -1
                ## note: both A and G count can still be zero if a variant (C/T) was detected at the site
                ag_counts = [PileupLoader.get_AG_counts(row[2 * si + 5], a_index) for si in sample_indexes]
                result.append((int(row[1]), ag_counts))
            result.sort(key=lambda x: x[0]) # should be already sorted by chromosome position, but make sure
        return result

    def get_sample(self, strand=None):
        if not strand: # pick from either strand
            ## use 'choices' instead of 'choice' for the 'weights' argument:
            strand = random.choices(["p", "n"], weights=(len(self.pileup_pos), len(self.pileup_neg)))[0]
        if strand in ("p", "+"):
            return ("p", random.choice(self.pileup_pos))
        elif strand in ("n", "-"):
            return ("n", random.choice(self.pileup_neg))

    def get_sample_size(self, strand=None):
        if not strand:
            return (len(self.pileup_pos), len(self.pileup_neg))
        if strand in ("p", "+"):
            return len(self.pileup_pos)
        elif strand in ("n", "-"):
            return len(self.pileup_neg)

    def get_item(self, strand, index):
        if strand in ("p", "+"):
            return self.pileup_pos[index]
        elif strand in ("n", "-"):
            return self.pileup_neg[index]

    def get_range(self, strand, position, delta): # TODO: needed?
        assert strand in ("p", "n", "+", "-")
        pileup = self.pileup_pos if strand in ("p", "+") else self.pileup_neg
        index = bisect_right(pileup, (position, [])) # to the right of "position"
        ## NOTE: this would be so much easier with one-based indexing!
        ## find left (lower) end of region:
        for lower in range(index - 1, -1, -1):
            if pileup[lower][0] < position - delta:
                lower += 1 # went a step too far
                break
        ## find right (upper) end of region - here we want to go "past the end":
        upper = index
        while (upper < len(pileup)) and (pileup[upper][0] <= position + delta):
            upper += 1
        return pileup[lower:upper]

    def get_region(self, strand, start, end):
        assert strand in ("p", "n", "+", "-")
        pileup = self.pileup_pos if strand in ("p", "+") else self.pileup_neg
        left = bisect_left(pileup, (start, []))
        right = bisect_right(pileup, (end, []))
        return pileup[left:right]

    @staticmethod
    def sum_replicate_counts(site_data):
        ## 'site_data' is a tuple: (position, [rep1_counts, rep2_counts])
        a_counts = 0
        g_counts = 0
        for rep in site_data[1]:
            if isinstance(rep, int): # just G count (A count is 0)
                g_counts += rep
            else:
                g_counts += rep[1]
                a_counts += rep[0]
        return a_counts, g_counts

    def get_stats(self, pileup=None):
        if pileup is None:
            return self.get_stats(self.pileup_pos), self.get_stats(self.pileup_neg)
        counts = [self.sum_replicate_counts(site) for site in pileup]
        totals = np.array([sum(x) for x in counts])
        proportions = np.empty(len(counts)) # m6A proportions
        for i in range(len(counts)):
            if totals[i] == 0: # rare case with zero A/G counts (C/T variant)
                proportions[i] = 0
            else:
                proportions[i] = counts[i][0] / float(totals[i])
        return proportions, totals

    def plot_total_counts(self):
        stats_pos, stats_neg = self.get_stats()
        totals_pos = Counter(stats_pos[1].tolist())
        totals_neg = Counter(stats_neg[1].tolist())
        plt.bar(totals_pos.keys(), totals_pos.values(), width=0.4, align="edge")
        plt.bar(totals_neg.keys(), totals_neg.values(), color="r", width=-0.4, align="edge")
        plt.show()

    def plot_proportions(self, min_counts=0):
        stats_pos, stats_neg = self.get_stats()
        props = np.concatenate((stats_pos[0], stats_neg[0]))
        if min_counts > 0:
            totals = np.concatenate((stats_pos[1], stats_neg[1]))
            props = props[totals >= min_counts]
        pd.Series(props).plot(kind="density", xlim=(-0.1, 1.1), grid=True, bw_method=0.1)
        plt.show()


class GTFLoader:
    "Load information about exons and their boundaries from pre-filtered GTF files"
    def __init__(self, chromosome, gtf_dir=os.path.join(DATA_DIR, "genome")):
        path_pos = os.path.join(gtf_dir, f"human_exons_filtered_chrp_{chromosome}.gtf") # plus strand
        path_neg = os.path.join(gtf_dir, f"human_exons_filtered_chrn_{chromosome}.gtf") # minus strand
        exons_pos = self.read_exons(path_pos)
        exons_neg = self.read_exons(path_neg)
        ## build interval trees (add 1 to end pos. because intervals exclude upper bounds!):
        self.tree_pos = IntervalTree(Interval(x.start, x.end + 1, x.attrs) for x in exons_pos)
        self.tree_neg = IntervalTree(Interval(x.start, x.end + 1, x.attrs) for x in exons_neg)

    @staticmethod
    def attributes_to_dict(string):
        parts = string.rstrip(";").split("; ")
        result = {}
        for part in parts:
            key, value = part.split(" ", 1)
            result[key] = value.strip('"')
        return result

    @staticmethod
    def read_exons(path):
        ## columns: 0. chromosome, 1. source, 2. feature ("exon"), 3. start position, 4. end position,
        ## 5. score ("."), 6. strand, 7. frame ("."), 8. attributes
        exons = []
        with open(path) as source:
            reader = csv.reader(source, delimiter="\t", quoting=csv.QUOTE_NONE)
            for row in reader:
                attrs = GTFLoader.attributes_to_dict(row[8])
                ## 'attrs' includes "transcript_source", so no need to store "source" column as well...
                exons.append(Exon(int(row[3]), int(row[4]), attrs))
        exons.sort(key=lambda exon: (exon.start, exon.end)) # can't compare the 'attrs' dict
        return exons

    def find_exons(self, position, strand):
        if strand in ("p", "+"):
            return self.tree_pos[position]
        elif strand in ("n", "-"):
            return self.tree_neg[position]
        raise ValueError(f"unsupported value for 'strand': {strand}")


class SequenceLoader:
    "Load nucleotide sequence given genomic coordinates"
    def __init__(self, chromosome, fasta_dir=os.path.join(DATA_DIR, "genome")):
        fasta_path = os.path.join(fasta_dir, f"Homo_sapiens.GRCh38.dna.chromosome.{chromosome}.fa")
        self.fasta = Fasta(fasta_path) # 'as_raw': return only sequence string, no header
        assert self.fasta[0].name == chromosome

    def get_sequence(self, strand, start, end=None):
        assert strand in ("p", "n", "+", "-")
        if end is None:
            end = start
        assert start <= end
        ## slices are zero-based and exclusive of upper bound, but coordinates are one-based and inclusive:
        if strand in ("p", "+"):
            return self.fasta[0][(start - 1):end].seq
        else: # minus strand - count from the end and reverse-complement:
            return (-self.fasta[0][-end:-(start - 1)]).seq

    def get_reverse_position(self, pos):
        return self.fasta[0][-pos].start


class DataLoader:
    "Load data for model training/testing"
    def __init__(self, chromosome, cell_line, treatment, batch=1, samples=SampleLists(),
                 gtf_dir=os.path.join(DATA_DIR, "genome"), fasta_dir=os.path.join(DATA_DIR, "genome"), window=200):
        self.chromosome = chromosome
        self.cell_line = cell_line
        self.treatment = treatment
        self.pileup_loader = PileupLoader(chromosome, cell_line, treatment, batch, samples)
        self.batch = self.pileup_loader.batch
        self.gtf_loader = GTFLoader(chromosome, gtf_dir)
        self.seq_loader = SequenceLoader(chromosome, fasta_dir)
        self.window = window

    def get_region(self, position, start, end):
        assert start < end
        assert (position >= start) and (position <= end)
        if end - start < self.window:
            return start, end
        region_start = random.randint(max(start, position - self.window), position)
        ## subtract 1 because end position is inclusive:
        region_end = region_start + self.window - 1
        if region_end > end:
            region_end = end
            region_start = region_end - self.window - 1
        return region_start, region_end

    def get_sample(self, strand=None, index=None):
        ## 1. pick an m6A site...
        if not index: ## ... at random
            strand, site = self.pileup_loader.get_sample(strand)
        else: ## ... at a specified list position
            site = self.pileup_loader.get_item(strand, index)
        ## 2. find the exon(s) that contains the site:
        pos = site[0]
        if strand in ("n", "-"):
            ## negative strand - translate between reverse (pileup) and forward (GTF) genomic coordinates:
            pos = self.seq_loader.get_reverse_position(pos)
        exons = self.gtf_loader.find_exons(pos, strand)
        if not exons:
            raise LookupError(f"no exon found for position {site[0]} on strand '{strand}'")
        if len(exons) == 1:
            exon = exons.pop()
        else: ## if multiple exons, pick the longest
            exon_lens = np.array([x.end - x.begin for x in exons]) # class of 'x': Interval
            exon = list(exons)[exon_lens.argmax()]
        ## 3. pick a random window containing the m6A site, subject to exon boundaries:
        start, end = self.get_region(pos, exon.begin, exon.end - 1) # change end from exclusive to inclusive
        if strand in ("n", "-"): # translate forward coordinates back to reverse (see above), swap start/end
            start, end = self.seq_loader.get_reverse_position(end), self.seq_loader.get_reverse_position(start)
        ## 4. get genomic sequence for that region:
        seq = self.seq_loader.get_sequence(strand, start, end)
        ## 5. get m6A proportions for sites in that region:
        sites = self.pileup_loader.get_region(strand, start, end)
        proportions = np.zeros(len(seq))
        total_counts = np.zeros(len(seq))
        for site in sites:
            offset = site[0] - start
            a_counts, g_counts = PileupLoader.sum_replicate_counts(site)
            total = a_counts + g_counts
            if total > 0: # can be zero in case of non-A variant!
                proportions[offset] = a_counts / float(a_counts + g_counts)
            total_counts[offset] = total
        ## consistency check:
        a_sites = np.argwhere(total_counts > 0).flatten().tolist()
        assert all([seq[i] == "A" for i in a_sites])
        sample = SampleRegion(site, exon.data.get("exon_id", ""), strand, start, end, seq, proportions, total_counts)
        return sample


class BayesModel:
    "Simple model for m6A proportions (to sample from posterior distribution for low read counts)"
    def __init__(self, all_proportions: np.array, all_totals: np.array, count_threshold=10):
        self.count_threshold = count_threshold
        proportions = all_proportions[all_totals >= count_threshold] # reliable proportions (enough counts)
        ## model prior as beta distribution:
        mean = proportions.mean()
        ## use 'count_threshold' to define width (strength of prior):
        self.alpha = mean * count_threshold # TODO: round to integer?
        self.beta_ = count_threshold - self.alpha
        self.prior = beta(self.alpha, self.beta_)

    def sample(self, proportion: float, total: int) -> float:
        if total == 0: # no data - sample from prior distribution
            return self.prior.rvs()
        if total >= self.count_threshold: # if we have enough counts, use measured proportion directly
            return proportion
        ## otherwise sample from the posterior distribution:
        alpha = proportion * total
        beta_ = total - alpha
        return beta.rvs(alpha + self.alpha, beta_ + self.beta_)

    def sample_posterior(self, sample_region: SampleRegion) -> np.array:
        result = np.zeros(sample_region.m6A_proportions.shape)
        for i in range(len(sample_region.seq)):
            if sample_region.seq[i] == "A":
                result[i] = self.sample(sample_region.m6A_proportions[i], sample_region.read_counts[i])
        return result


if __name__ == "__main__":
    ## test:
    loader = DataLoader("Y", "CAOV3", "untreated", 1, window=100)
    sample = loader.get_sample()
    # try:
    #     sample = loader.get_sample("-")
    # except:
    #     traceback.print_exc()

    stats_pos, stats_neg = loader.pileup_loader.get_stats()
    props = np.concatenate((stats_pos[0], stats_neg[0]))
    totals = np.concatenate((stats_pos[1], stats_neg[1]))
    bm = BayesModel(props, totals)
    bm.sample_posterior(sample)

    pl = loader.pileup_loader
    gl = loader.gtf_loader
    sl = loader.seq_loader

    ## m6A data exploration:
    stats_pos, stats_neg = pl.get_stats()
    props = np.concatenate((stats_pos[0], stats_neg[0]))
    totals = np.concatenate((stats_pos[1], stats_neg[1]))
    ## m6A proportions for sites with at least 10 total counts:
    part = props[totals >= 10]
    part.mean() # overall average proprotion is 7%
    (part == 0).mean() # no m6A detected in 26% of cases (proportion is zero)
    (part[part > 0]).mean() # where proportion is non-zero, average is 9%

    np.quantile(totals, np.arange(0, 1.1, 0.1))
    bounds = [(1, 1), (2, 2), (3, 3), (4, 5), (6, 11), (12, 27), (28, 381)] # boundaries along 10% quantiles
    zero_props = [(props[np.logical_and((totals >= b[0]), (totals <= b[1]))] == 0).mean() for b in bounds]
    ## -> probability of m6A proportion being exactly zero decreases with read count
    half_props = [(props[np.logical_and((totals >= b[0]), (totals <= b[1]))] < 0.5).mean() for b in bounds]
    ## -> probability of m6A proportion being <50% stays roughly constant (95% +/- 5%)