def do_load():
from mbf_nested_intervals import IntervalSet
import itertools
df = gr.df
joined = []
for ii, (chr, start,
stop) in enumerate(zip(df["chr"], df["start"], df["stop"])):
joined.append(((chr), start, stop, ii))
joined.sort(key=lambda tup: tup[0])
out = []
chr_lengths = gr.genome.get_chromosome_lengths()
seen = set()
for chr, sub_group in itertools.groupby(joined, lambda tup: tup[0]):
args = [x[1:] for x in sub_group]
iv = IntervalSet.from_tuples_with_id(args)
new_order = iv.invert(0, chr_lengths[chr]).to_numpy()
out.append(
pd.DataFrame({
"start": new_order[0],
"stop": new_order[1],
"chr": chr
}))
seen.add(chr)
for chr in chr_lengths.keys() - seen:
out.append(
pd.DataFrame({
"start": [0],
"stop": [chr_lengths[chr]],
"chr": chr
}))
return pd.concat(out).reset_index(drop=True)
def merge_df_intervals(df, iv_func=lambda iv: iv.merge_hull()):
"""take a DataFrame {chr, start, end, *} and merge overlapping intervals.
* is from the last entry.
"""
if not "strand" in df.columns:
df = df.assign(strand=1)
strand_added = True
else:
strand_added = False
joined = _df_to_tup(df)
out = []
for chr_strand, sub_group in itertools.groupby(joined, lambda tup: tup[0]):
args = [x[1:] for x in sub_group]
iv = IntervalSet.from_tuples_with_id(args)
new_order = iv_func(iv).to_tuples_last_id()
new_df = df.iloc[[x[2] for x in new_order]].copy()
new_df.loc[:, "start"] = [x[0] for x in new_order]
new_df.loc[:, "stop"] = [x[1] for x in new_order]
out.append(new_df)
res = pd.concat(out)
if strand_added:
res = res.drop("strand", axis=1)
return res.sort_values(["chr", "start"])
def test_merge_hull(self):
i = IntervalSet.from_tuples_with_id([
(1,10, 100),
(7,15, 200),
(0,5, 333),
])
i2 = i.merge_hull()
assert i2.to_tuples_with_id() == [
(0, 15, [100, 200, 333])
]
def test_from_tuples_with_id2(self):
i = IntervalSet.from_tuples_with_id([
(1,10, 100),
(1,15, 200),
(0,5, 333),
])
assert i.to_tuples_with_id() == [
(0,5, [333]),
(1,15, [200]),
(1,10, [100]),
]
def merge_df_intervals_with_callback(df, callback):
"""take a {chr, start, end, *} dataframe and merge overlapping intervals, calling callback for group larger than one.."""
if not "strand" in df:
df = df.assign(strand=1)
strand_added = True
else:
strand_added = False
joined = _df_to_tup(df)
result = []
for chr, sub_group in itertools.groupby(joined, lambda tup: tup[0]):
args = [x[1:] for x in sub_group]
iv = IntervalSet.from_tuples_with_id(args)
subsets = iv.merge_hull().to_tuples_with_id()
for s in subsets:
sub_df = df.iloc[list(s[2])].copy()
sub_df.at[:, "start"] = s[0]
sub_df.at[:, "stop"] = s[1]
row_data = callback(sub_df)
if not isinstance(
row_data, dict
): # and not (isinstance(row_data, pd.core.series.Series) and len(row_data.shape) == 1):
print("type", type(row_data))
# print 'len(shape)', len(row_data.shape)
print(callback)
raise ValueError(
"Merge_function returned something other than dict (writing to the pandas series directly is very slow, call to_dict() on it, then modify it.)"
)
if set(row_data.keys()) != set(df.columns):
raise ValueError(
"Merge_function return wrong columns. Expected %s, was %s"
% (df.columns, list(row_data.keys()))
)
row_data["start"] = s[0]
row_data["stop"] = s[1]
result.append(row_data)
res = pd.DataFrame(result).sort_values(["chr", "start"])
if strand_added:
res = res.drop("strand", axis=1)
return res
def _get_interval_tuples_by_chr(self, genome):
from mbf_nested_intervals import IntervalSet
coll = {chr: [] for chr in genome.get_chromosome_lengths()}
for g in genome.genes.values():
exons = g.exons_overlapping
if len(exons[0]) == 0: # pragma: no cover
exons = g.exons_merged
for start, stop in zip(*exons):
coll[g.chr].append(
(start, stop, 0b0101 if g.strand == 1 else 0b0110))
for start, stop in zip(*g.introns_strict):
coll[g.chr].append(
(start, stop, 0b1001 if g.strand == 1 else 0b1010))
result = {}
for chr, tups in coll.items():
iset = IntervalSet.from_tuples_with_id(tups)
# iset = iset.merge_split()
iset = iset.merge_hull()
if iset.any_overlapping():
raise NotImplementedError("Should not be reached")
result[chr] = []
for start, stop, ids in iset.to_tuples_with_id():
ids = set(ids)
if len(ids) == 1:
id = list(ids)[0]
if id == 0b0101:
tag = "exon"
strand = +1
elif id == 0b0110:
tag = "exon"
strand = -1
elif id == 0b1001:
tag = "intron"
strand = +1
elif id == 0b1010:
tag = "intron"
strand = -1
else: # pragma: no cover
raise NotImplementedError(
"Should not be reached")
else:
down = 0
for i in ids:
down |= i
if down & 0b1100 == 0b1100:
tag = "both"
elif down & 0b0100 == 0b0100:
tag = "exon"
else: # pragma: no cover haven't observed this case in the wild yet.
tag = ( # pragma: no cover
"intron" # pragma: no cover
) # pragma: no cover haven't observed this case in the wild yet.
if down & 0b11 == 0b11:
tag += "_undecidable"
strand = (
1
) # doesn't matter, but must be one or the other
elif down & 0b01:
strand = 1
else:
strand -= 1
result[chr].append((tag, strand, [start], [stop]))
return result