Spaces:
Runtime error
Runtime error
File size: 5,283 Bytes
4063c61 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | import gzip
import random
from Bio import SeqIO
from mint.utils.logging import get_logger
logger = get_logger(__name__)
logger.info("===Reading seqs====")
seqs = {}
for seq in SeqIO.parse(open("../protein.sequences.v12.0.fa"), "fasta"):
seqs[seq.name] = str(seq.seq)
if len(seqs) % 1000000 == 0:
logger.info(f"{len(seqs) / 1e6} million seqs read")
logger.info(f"Done, {len(seqs)} seqs total")
logger.info("===Reading reps====")
reps = {}
for line in open("../clu50.tsv"):
rep, seq = line.strip().split()
reps[seq] = rep
if len(reps) % 1000000 == 0:
logger.info(f"{len(reps) / 1e6} million reps read")
logger.info(f"Done, {len(reps)} reps total, {len(set(reps.values()))} clusters")
logger.info("===Reading links====")
f = gzip.open("../protein.physical.links.full.v12.0.txt.gz", "rt")
f = iter(f)
next(f) # skip first line
links = []
i = 0
while True:
try:
line = next(f).strip()
except StopIteration:
break
i += 1
links.append(line)
if i % 1000000 == 0:
logger.info(f"{i / 1e6} million links read")
# if i / 1e6 == 10: break
logger.info(f"Done, {len(links)} links total")
logger.info("===Shuffling links===")
random.seed(137)
random.shuffle(links)
logger.info("Done shuffling links")
logger.info("===Filtering links===")
linked_clusters = set()
filtered_links = []
i = 0
for link in links:
i += 1
name1, name2 = link.split()[:2]
clu1, clu2 = reps[name1], reps[name2]
clu1, clu2 = tuple(sorted((clu1, clu2)))
if (clu1, clu2) not in linked_clusters:
linked_clusters.add((clu1, clu2))
filtered_links.append(link)
if i % 1000000 == 0:
logger.info(f"{i / 1e6} million links filtered, {len(filtered_links) / 1e6} million kept")
links = filtered_links
logger.info(f"Done, {i} links filtered, {len(links)} kept")
logger.info("===Shuffling links===")
random.seed(731)
random.shuffle(links)
logger.info("Done shuffling links")
with gzip.open("../filtered.links.txt.gz", "wt") as links_file:
for link in links:
links_file.write(link + "\n")
num_val = 250000
validation = links[:num_val]
training = links[num_val:]
logger.info("===Writing validation links===")
written_seqs = set()
with gzip.open("validation.links.txt.gz", "wt") as links_file:
with gzip.open("validation.seqs.txt.gz", "wt") as seqs_file:
for link in validation:
links_file.write(link + "\n")
name1, name2 = link.split()[:2]
if name1 not in written_seqs:
seqs_file.write(name1 + " " + seqs[name1] + "\n")
written_seqs.add(name1)
if name2 not in written_seqs:
seqs_file.write(name2 + " " + seqs[name2] + "\n")
written_seqs.add(name2)
logger.info(f"Done, {num_val} validation links written, {len(written_seqs)} seqs")
logger.info("===Writing training links===")
i = 0
written_seqs = set()
with gzip.open("training.links.txt.gz", "wt") as links_file:
with gzip.open("training.seqs.txt.gz", "wt") as seqs_file:
for link in training:
i += 1
links_file.write(link + "\n")
name1, name2 = link.split()[:2]
if name1 not in written_seqs:
seqs_file.write(name1 + " " + seqs[name1] + "\n")
written_seqs.add(name1)
if name2 not in written_seqs:
seqs_file.write(name2 + " " + seqs[name2] + "\n")
written_seqs.add(name2)
if i % 1000000 == 0:
logger.info(
f"{i / 1e6} million training links written, {len(written_seqs) / 1e6} million seqs"
)
logger.info(f"Done, {i} training links written, {len(written_seqs)} seqs")
logger.info("===Extracting validation clusters===")
val_clus = []
for link in validation:
name1, name2 = link.split()[:2]
val_clus.append(reps[name1])
val_clus.append(reps[name2])
val_clus = set(val_clus)
logger.info(f"Done, {len(val_clus)} validation clusters")
i, j = 0, 0
logger.info("===Writing filtered training links===")
written_seqs = set()
with gzip.open("training_filtered.links.txt.gz", "wt") as links_file:
with gzip.open("training_filtered.seqs.txt.gz", "wt") as seqs_file:
for link in training:
i += 1
name1, name2 = link.split()[:2]
clu1, clu2 = reps[name1], reps[name2]
if clu1 not in val_clus and clu2 not in val_clus:
j += 1
links_file.write(link + "\n")
if name1 not in written_seqs:
seqs_file.write(name1 + " " + seqs[name1] + "\n")
written_seqs.add(name1)
if name2 not in written_seqs:
seqs_file.write(name2 + " " + seqs[name2] + "\n")
written_seqs.add(name2)
if i % 1000000 == 0:
logger.info(
f"{i / 1e6} million training links filtered, {j / 1e6} million written, {len(written_seqs) / 1e6} million seqs"
)
logger.info(f"{i} training links filtered, {j} kept, {len(written_seqs)} seqs")
|