Coverage for barbet/embedding.py: 32.82%
326 statements
« prev ^ index » next coverage.py v7.9.1, created at 2026-08-17 06:22 +0000
1import gzip
2import os
3from pathlib import Path
4from abc import ABC, abstractmethod
5from Bio import SeqIO
6import random
7import numpy as np
8from rich.progress import track
9from hierarchicalsoftmax import SoftmaxNode
10from hierarchicalsoftmax import TreeDict
11import tarfile
12import torch
13from io import StringIO
14from torchapp.cli import CLIApp, tool, method
15import typer
16from dataclasses import dataclass
18from .data import read_memmap, RANKS
21def _open(path, mode='rt', **kwargs):
22 """
23 Open a file normally, or with gzip if it ends in .gz.
25 Args:
26 path (str or Path): The path to the file.
27 mode (str): The mode to open the file with (default 'rt' for reading text).
28 **kwargs: Additional arguments passed to open or gzip.open.
30 Returns:
31 A file object.
32 """
33 path = Path(path)
34 if path.suffix == '.gz':
35 return gzip.open(path, mode, **kwargs)
36 return open(path, mode, **kwargs)
39def set_validation_rank_to_treedict(
40 treedict:TreeDict,
41 validation_rank:str="species",
42 partitions:int=5,
43) -> TreeDict:
44 # find the taxonomic rank to use for the validation partition
45 validation_rank = validation_rank.lower()
46 assert validation_rank in RANKS
47 validation_rank_index = RANKS.index(validation_rank)
49 partitions_dict = {}
50 for key in treedict:
51 node = treedict.node(key)
52 # Assign validation partition at set rank
53 partition_node = node.ancestors[validation_rank_index]
54 if partition_node not in partitions_dict:
55 partitions_dict[partition_node] = random.randint(0,partitions-1)
57 treedict[key].partition = partitions_dict[partition_node]
59 return treedict
62def get_key(accession:str, gene:str) -> str:
63 """ Returns the standard format of a key """
64 key = f"{accession}/{gene}"
65 return key
68def get_node(lineage:str, lineage_to_node:dict[str,SoftmaxNode]) -> SoftmaxNode:
69 if lineage in lineage_to_node:
70 return lineage_to_node[lineage]
72 assert ";" in lineage, f"Semi-colon ';' not found in lineage '{lineage}'"
73 split_point = lineage.rfind(";")
74 parent_lineage = lineage[:split_point]
75 name = lineage[split_point+1:]
76 parent = get_node(parent_lineage, lineage_to_node)
77 node = SoftmaxNode(name=name, parent=parent)
78 lineage_to_node[lineage] = node
79 return node
82def generate_overlapping_intervals(total: int, interval_size: int, min_overlap: int, check:bool=True, variable_size:bool=False) -> list[tuple[int, int]]:
83 """
84 Creates a list of overlapping intervals within a specified range, adjusting the interval size to ensure
85 that the overlap is approximately the same across all intervals.
87 Args:
88 total (int): The total range within which intervals are to be created.
89 max_interval_size (int): The maximum size of each interval.
90 min_overlap (int): The minimum number of units by which consecutive intervals overlap.
91 check (bool): If True, checks are performed to ensure that the intervals meet the specified conditions.
93 Returns:
94 list[tuple[int, int]]: A list of tuples where each tuple represents the start (inclusive)
95 and end (exclusive) of an interval.
97 Example:
98 >>> generate_overlapping_intervals(20, 5, 2)
99 [(0, 5), (3, 8), (6, 11), (9, 14), (12, 17), (15, 20)]
100 """
101 intervals = []
102 start = 0
104 if total == 0:
105 return intervals
107 max_interval_size = interval_size
108 assert interval_size
109 assert min_overlap is not None
110 assert interval_size > min_overlap, f"Max interval size of {interval_size} must be greater than min overlap of {min_overlap}"
112 # Calculate the number of intervals needed to cover the range
113 num_intervals, remainder = divmod(total - min_overlap, interval_size - min_overlap)
114 if remainder > 0:
115 num_intervals += 1
117 # Calculate the exact interval size to ensure consistent overlap
118 overlap = min_overlap
119 if variable_size:
120 if num_intervals > 1:
121 interval_size, remainder = divmod(total + (num_intervals - 1) * overlap, num_intervals)
122 if remainder > 0:
123 interval_size += 1
124 else:
125 # If the size is fixed, then vary the overlap to keep it even
126 if num_intervals > 1:
127 overlap, remainder = divmod( num_intervals * interval_size - total, num_intervals - 1)
128 if overlap < min_overlap:
129 overlap = min_overlap
131 while True:
132 end = start + interval_size
133 if end > total:
134 end = total
135 start = max(end - interval_size,0)
136 intervals.append((start, end))
137 start += interval_size - overlap
138 if end >= total:
139 break
141 if check:
142 assert intervals[0][0] == 0
143 assert intervals[-1][1] == total
144 assert len(intervals) == num_intervals, f"Expected {num_intervals} intervals, got {len(intervals)}"
146 assert interval_size <= max_interval_size, f"Interval size of {interval_size} exceeds max interval size of {max_interval_size}"
147 for interval in intervals:
148 assert interval[1] - interval[0] == interval_size, f"Interval size of {interval[1] - interval[0]} is not the expected size {interval_size}"
150 for i in range(1, len(intervals)):
151 overlap = intervals[i - 1][1] - intervals[i][0]
152 assert overlap >= min_overlap, f"Min overlap condition of {min_overlap} not met for intervals {intervals[i - 1]} and {intervals[i]} (overlap {overlap})"
154 return intervals
157@dataclass
158class Embedding(CLIApp, ABC):
159 """ A class for embedding protein sequences. """
160 max_length:int|None=None
161 overlap:int=64
163 def __post_init__(self):
164 super().__init__()
165 if hasattr(self, "main_app") and hasattr(self.main_app, "info"):
166 self.main_app.info.context_settings = {"help_option_names": ["-h", "--help"]}
167 if hasattr(self, "tools_app") and hasattr(self.tools_app, "info"):
168 self.tools_app.info.context_settings = {"help_option_names": ["-h", "--help"]}
170 @abstractmethod
171 def embed(self, seq:str) -> torch.Tensor:
172 """ Takes a protein sequence as a string and returns an embedding vector. """
173 raise NotImplementedError
175 def reduce(self, tensor:torch.Tensor) -> torch.Tensor:
176 if tensor.ndim == 2:
177 tensor = tensor.mean(dim=0)
178 assert tensor.ndim == 1
179 return tensor
181 def __call__(self, seq:str) -> torch.Tensor:
182 """ Takes a protein sequence as a string and returns an embedding vector. """
183 if not self.max_length or len(seq) <= self.max_length:
184 tensor = self.embed(seq)
185 return self.reduce(tensor)
187 epsilon = 0.1
188 intervals = generate_overlapping_intervals(len(seq), self.max_length, self.overlap)
189 weights = torch.zeros( (len(seq),), device="cpu" )
190 tensor = None
191 for start,end in intervals:
192 result = self.embed(seq[start:end]).cpu()
194 assert result.shape[0] == end-start
195 embedding_size = result.shape[1]
197 if tensor is None:
198 tensor = torch.zeros( (len(seq), embedding_size ), device="cpu")
200 assert tensor.shape[-1] == embedding_size
202 interval_indexes = torch.arange(end-start)
203 distance_from_ends = torch.min( interval_indexes-start, end-interval_indexes-1 )
205 weight = epsilon + torch.minimum(distance_from_ends, torch.tensor(self.overlap))
207 tensor[start:end] += result * weight.unsqueeze(1)
208 weights[start:end] += weight
210 tensor = tensor/weights.unsqueeze(1)
212 return self.reduce(tensor)
214 @method
215 def setup(self, **kwargs):
216 pass
218 def build_treedict(self, taxonomy:Path) -> tuple[TreeDict,dict[str,SoftmaxNode]]:
219 # Create root of tree
220 lineage_to_node = {}
221 root = None
223 # Fill out tree with taxonomy
224 accession_to_node = {}
225 with _open(taxonomy) as f:
226 for line in f:
227 accesssion, lineage = line.split("\t")
229 if not root:
230 root_name = lineage.split(";")[0]
231 root = SoftmaxNode(root_name)
232 lineage_to_node[root_name] = root
234 node = get_node(lineage, lineage_to_node)
235 accession_to_node[accesssion] = node
237 treedict = TreeDict(classification_tree=root)
238 return treedict, accession_to_node
240 @tool("setup")
241 def test_lengths(
242 self,
243 end:int=5_000,
244 start:int=1000,
245 retries:int=5,
246 **kwargs,
247 ):
248 def random_amino_acid_sequence(k):
249 amino_acids = "ACDEFGHIKLMNPQRSTVWY" # standard 20 amino acids
250 return ''.join(random.choice(amino_acids) for _ in range(k))
252 self.max_length = None
253 self.setup(**kwargs)
254 for ii in track(range(start,end)):
255 for _ in range(retries):
256 seq = random_amino_acid_sequence(ii)
257 try:
258 self(seq)
259 except Exception as err:
260 print(f"{ii}: {err}")
261 return
264 @tool("setup")
265 def build_gene_array(
266 self,
267 marker_genes:Path=typer.Option(default=..., help="The path to the marker genes tarball (e.g. bac120_msa_marker_genes_all_r220.tar.gz)."),
268 family_index:int=typer.Option(default=..., help="The index for the gene family to use. E.g. if there are 120 gene families then this should be a number from 0 to 119."),
269 output_dir:Path=typer.Option(default=..., help="A directory to store the output which includes the memmap array, the listing of accessions and an error log."),
270 flush_every:int=typer.Option(default=5_000, help="An interval to flush the memmap array as it is generated."),
271 max_length:int=None,
272 **kwargs,
273 ):
274 self.max_length = max_length
275 self.setup(**kwargs)
277 assert marker_genes is not None
278 assert family_index is not None
279 assert output_dir is not None
281 dtype = 'float16'
283 memmap_wip_array = None
284 output_dir.mkdir(parents=True, exist_ok=True)
285 memmap_wip_path = output_dir / f"{family_index}-wip.npy"
286 error = output_dir / f"{family_index}-errors.txt"
287 accessions_wip = output_dir / f"{family_index}-accessions-wip.txt"
289 accessions = []
291 print(f"Loading {marker_genes} file.")
292 with tarfile.open(marker_genes, "r:gz") as tar, open(error, "w") as error_file, open(accessions_wip, "w") as accessions_wip_file:
293 members = [member for member in tar.getmembers() if member.isfile() and member.name.endswith(".faa")]
294 prefix_length = len(os.path.commonprefix([Path(member.name).with_suffix("").name for member in members]))
296 member = members[family_index]
297 print(f"Processing file {family_index} in {marker_genes}")
299 f = tar.extractfile(member)
300 marker_id = Path(member.name).with_suffix("").name[prefix_length:]
302 fasta_io = StringIO(f.read().decode('ascii'))
304 total = sum(1 for _ in SeqIO.parse(fasta_io, "fasta"))
305 fasta_io.seek(0)
306 print(marker_id, total)
308 for record in track(SeqIO.parse(fasta_io, "fasta"), total=total):
309 # for record in SeqIO.parse(fasta_io, "fasta"):
310 species_accession = record.id
312 key = get_key(species_accession, marker_id)
314 seq = str(record.seq).replace("-","").replace("*","")
315 try:
316 vector = self(seq)
317 except Exception as err:
318 print(f"{key} ({len(seq)}): {err}", file=error_file)
319 print(f"{key} ({len(seq)}): {err}")
320 continue
322 if vector is None:
323 print(f"{key} ({len(seq)}): Embedding is None", file=error_file)
324 print(f"{key} ({len(seq)}): Embedding is None")
325 continue
327 if torch.isnan(vector).any():
328 print(f"{key} ({len(seq)}): Embedding contains NaN", file=error_file)
329 print(f"{key} ({len(seq)}): Embedding contains NaN")
330 continue
332 if memmap_wip_array is None:
333 size = len(vector)
334 shape = (total,size)
335 memmap_wip_array = np.memmap(memmap_wip_path, dtype=dtype, mode='w+', shape=shape)
337 index = len(accessions)
338 memmap_wip_array[index,:] = vector.cpu().half().numpy()
339 if index % flush_every == 0:
340 memmap_wip_array.flush()
342 accessions.append(key)
343 print(key, file=accessions_wip_file)
345 memmap_wip_array.flush()
347 accessions_path = output_dir / f"{family_index}.txt"
348 with open(accessions_path, "w") as f:
349 for accession in accessions:
350 print(accession, file=f)
352 # Save final memmap array now that we now the final size
353 memmap_path = output_dir / f"{family_index}.npy"
354 shape = (len(accessions),size)
355 print(f"Writing final memmap array of shape {shape}: {memmap_path}")
356 memmap_array = np.memmap(memmap_path, dtype=dtype, mode='w+', shape=shape)
357 memmap_array[:len(accessions),:] = memmap_wip_array[:len(accessions),:]
358 memmap_array.flush()
360 # Clean up
361 memmap_array._mmap.close()
362 memmap_array._mmap = None
363 memmap_array = None
364 memmap_wip_path.unlink()
365 accessions_wip.unlink()
367 @tool
368 def set_validation_rank(
369 self,
370 treedict:Path=typer.Option(default=..., help="The path to the treedict file."),
371 output:Path=typer.Option(default=..., help="The path to save the adapted treedict file."),
372 validation_rank:str=typer.Option(default="species", help="The rank to hold out for cross-validation."),
373 partitions:int=typer.Option(default=5, help="The number of cross-validation partitions."),
374 ) -> TreeDict:
375 treedict = TreeDict.load(treedict)
376 set_validation_rank_to_treedict(treedict, validation_rank=validation_rank, partitions=partitions)
377 treedict.save(output)
378 return treedict
380 @tool
381 def preprocess(
382 self,
383 taxonomy:Path=typer.Option(default=..., help="The path to the TSV taxonomy file (e.g. bac120_taxonomy_r220.tsv)."),
384 marker_genes:Path=typer.Option(default=..., help="The path to the marker genes tarball (e.g. bac120_msa_marker_genes_all_r220.tar.gz)."),
385 output_dir:Path=typer.Option(default=..., help="A directory to store the output which includes the memmap array, the listing of accessions and an error log."),
386 partitions:int=typer.Option(default=5, help="The number of cross-validation partitions."),
387 seed:int=typer.Option(default=42, help="The random seed."),
388 treedict_only:bool=typer.Option(default=False, help="Only output TreeDict file and then exit before concatenating memmap array"),
389 ):
390 treedict, accession_to_node = self.build_treedict(taxonomy)
392 dtype = 'float16'
394 random.seed(seed)
396 print(f"Loading {marker_genes} file.")
397 with tarfile.open(marker_genes, "r:gz") as tar:
398 members = [member for member in tar.getmembers() if member.isfile() and member.name.endswith(".faa")]
399 family_count = len(members)
400 print(f"{family_count} gene families found.")
402 # Read and collect accessions
403 print(f"Building treedict")
404 keys = []
405 counts = []
406 node_to_partition_dict = dict()
407 for family_index in track(range(family_count)):
408 keys_path = output_dir / f"{family_index}.txt"
410 if not keys_path.exists():
411 counts.append(0)
412 continue
414 with open(keys_path) as f:
415 family_index_keys = [line.strip() for line in f]
416 keys += family_index_keys
417 counts.append(len(family_index_keys))
419 for key in family_index_keys:
420 genome_accession = key.split("/")[0]
421 node = accession_to_node[genome_accession]
422 partition = node_to_partition_dict.setdefault(node, random.randint(0, partitions - 1))
424 # Add to treedict
425 treedict.add(key, node, partition)
427 assert len(counts) == family_count
429 # Save treedict
430 treedict_path = output_dir / f"{output_dir.name}.td"
431 print(f"Saving TreeDict to {treedict_path}")
432 treedict.save(treedict_path)
434 if treedict_only:
435 return
437 # Concatenate numpy memmap arrays
438 memmap_array = None
439 memmap_array_path = output_dir / f"{output_dir.name}.npy"
440 print(f"Saving memmap to {memmap_array_path}")
441 current_index = 0
442 for family_index, family_count in track(enumerate(counts), total=len(counts)):
443 my_memmap_path = output_dir / f"{family_index}.npy"
445 # Build memmap for gene family if it doesn't exist
446 if not my_memmap_path.exists():
447 continue
448 # print("Building", my_memmap_path)
449 # self.build_gene_array(marker_genes=marker_genes, family_index=family_index, output_dir=output_dir)
450 # assert my_memmap_path.exists()
452 my_memmap = read_memmap(my_memmap_path, family_count)
454 # Build memmap for output if it doesn't exist
455 if memmap_array is None:
456 size = my_memmap.shape[1]
457 shape = (len(keys),size)
458 memmap_array = np.memmap(memmap_array_path, dtype=dtype, mode='w+', shape=shape)
460 # Copy memmap for gene family into output memmap
461 memmap_array[current_index:current_index+family_count,:] = my_memmap[:,:]
463 current_index += family_count
465 assert len(keys) == current_index
467 memmap_array.flush()
469 # Save keys
470 keys_path = output_dir / f"{output_dir.name}.txt"
471 print(f"Saving keys to {keys_path}")
472 with open(keys_path, "w") as f:
473 for key in keys:
474 print(key, file=f)
476 @tool
477 def prune_to_representatives(treedict:Path, representatives:Path, output:Path):
478 print("Getting list of representatives from", representatives)
479 keys_to_keep = []
480 with tarfile.open(representatives, "r:gz") as tar:
481 members = [member for member in tar.getmembers() if member.isfile() and member.name.endswith(".faa")]
483 print(f"Processing {len(members)} files in {representatives}")
485 for member in track(members):
486 f = tar.extractfile(member)
487 marker_id = Path(member.name.split("_")[-1]).with_suffix("").name
489 fasta_io = StringIO(f.read().decode('ascii'))
491 for record in SeqIO.parse(fasta_io, "fasta"):
492 species_accession = record.id
493 key = get_key(species_accession, marker_id)
494 keys_to_keep.append(key)
496 # keys_to_keep = set(keys_to_keep)
497 print(f"Keeping {len(keys_to_keep)} representatives")
499 print(f"Loading treedict {treedict}")
501 treedict = TreeDict.load(treedict)
502 print("Total", len(treedict))
503 missing = []
504 for key in track(keys_to_keep):
505 if key not in treedict:
506 missing.append(key)
508 print(f"{len(missing)} representatives missing output {len(keys_to_keep)} (total: {len(treedict)})")
509 if len(missing):
510 keys_to_keep = [k for k in keys_to_keep if k not in missing]
512 new_treedict = TreeDict(treedict.classification_tree)
513 new_treedict.update({k:treedict[k] for k in keys_to_keep})
514 print("Total after pruning", len(new_treedict))
516 print("Saving treedict to", output)
517 new_treedict.save(output)