diff --git a/bakta/features/t_rna.py b/bakta/features/t_rna.py index 083fe28d..6174fb2e 100644 --- a/bakta/features/t_rna.py +++ b/bakta/features/t_rna.py @@ -1,13 +1,16 @@ import logging -import subprocess as sp + +from concurrent.futures import ThreadPoolExecutor from collections import OrderedDict from pathlib import Path +from typing import Sequence from Bio import SeqIO import bakta.config as cfg import bakta.constants as bc +import bakta.io.fasta as fasta import bakta.so as so import bakta.utils as bu @@ -47,27 +50,35 @@ def predict_t_rnas(data: dict, sequences_path: Path): txt_output_path = cfg.tmp_path.joinpath('trna.tsv') fasta_output_path = cfg.tmp_path.joinpath('trna.fasta') - cmd = [ - 'tRNAscan-SE', - '-B', - '--output', str(txt_output_path), - '--fasta', str(fasta_output_path), - '--thread', str(cfg.threads), - str(sequences_path) - ] - log.debug('cmd=%s', cmd) - proc = sp.run( - cmd, - cwd=str(cfg.tmp_path), - env=cfg.env, - stdout=sp.PIPE, - stderr=sp.PIPE, - universal_newlines=True - ) - if(proc.returncode != 0): - log.debug('stdout=\'%s\', stderr=\'%s\'', proc.stdout, proc.stderr) - log.warning('tRNAs failed! tRNAscan-SE-error-code=%d', proc.returncode) - raise Exception(f'tRNAscan-SE error! error code: {proc.returncode}') + no_chunks = min(cfg.threads, len(data['sequences'])) + threads = cfg.threads if no_chunks == 1 else 0 + chunks = fasta.split_sequences(data['sequences'], no_chunks) + cmds, tsv_paths, fasta_paths = [], [], [] + for i, chunk in enumerate(chunks): + tsv_paths.append(cfg.tmp_path.joinpath(f'trna.{i}.tsv')) + fasta_paths.append(cfg.tmp_path.joinpath(f'trna.{i}.fasta')) + chunk_path = sequences_path + if(no_chunks > 1): + chunk_path = cfg.tmp_path.joinpath(f'trna.{i}.fna') + fasta.export_sequences(chunk, chunk_path) + cmds.append([ + 'tRNAscan-SE', + '-B', + '--output', str(tsv_paths[i]), + '--fasta', str(fasta_paths[i]), + '--thread', str(threads), + str(chunk_path) + ]) + log.debug('cmds=%s', cmds) + with ThreadPoolExecutor(max_workers=len(cmds)) as tpe: + procs = list(tpe.map(bu.run_tool, cmds)) + for proc in procs: + if(proc.returncode != 0): + log.debug('stdout=\'%s\', stderr=\'%s\'', proc.stdout, proc.stderr) + log.warning('tRNAs failed! tRNAscan-SE-error-code=%d', proc.returncode) + raise Exception(f'tRNAscan-SE error! error code: {proc.returncode}') + fasta.concat(tsv_paths, txt_output_path, skip=3) + fasta.concat(fasta_paths, fasta_output_path) trnas = {} sequences = {seq['id']: seq for seq in data['sequences']} diff --git a/bakta/io/fasta.py b/bakta/io/fasta.py index ed4ceed2..f2ecf35c 100644 --- a/bakta/io/fasta.py +++ b/bakta/io/fasta.py @@ -76,6 +76,31 @@ def export_sequences(sequences: Sequence[dict], fasta_path: Path, description: b fh.write('\n') +def split_sequences(sequences: Sequence[dict], no_chunks: int) -> Sequence[Sequence[dict]]: + """Split sequences into chunks of roughly equal size, keeping their order.""" + target = sum([seq['length'] for seq in sequences]) / no_chunks + split, chunk, size = [], [], 0 + for seq in sequences: + chunk.append(seq) + size += seq['length'] + if(size >= target and len(split) < no_chunks - 1): + split.append(chunk) + chunk, size = [], 0 + if(len(chunk) > 0): + split.append(chunk) + return split + + +def concat(parts: Sequence[Path], path: Path, skip: int=0): + """Concatenate tool output files in order, keeping one copy of the header.""" + with path.open('wt') as fh_out: + for i, part in enumerate(parts): + with part.open() as fh_in: + for j, line in enumerate(fh_in): + if(i == 0 or j >= skip): + fh_out.write(line) + + def wrap_sequence(sequence: str): lines = [] for i in range(0, len(sequence), FASTA_LINE_WRAPPING): diff --git a/bakta/utils.py b/bakta/utils.py index 37e9e5e2..303077a1 100644 --- a/bakta/utils.py +++ b/bakta/utils.py @@ -209,6 +209,18 @@ def check_version(tool, min: int, max:int ) -> bool: return True +def run_tool(cmd: Sequence[str]): + """Run one external tool process.""" + return sp.run( + cmd, + cwd=str(cfg.tmp_path), + env=cfg.env, + stdout=sp.PIPE, + stderr=sp.PIPE, + universal_newlines=True + ) + + def test_dependency(dependency): """Test the proper installation of the required 3rd party executable.""" version = read_tool_output(dependency)