From fb5bc229295bc40096227b03f90e299637cf14d2 Mon Sep 17 00:00:00 2001 From: Kevin Wu Date: Sun, 22 Oct 2023 16:28:08 -0700 Subject: [PATCH] Ability to specify threads --- scripts/faspr.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/scripts/faspr.py b/scripts/faspr.py index 981cba2..f41b7a9 100644 --- a/scripts/faspr.py +++ b/scripts/faspr.py @@ -71,6 +71,9 @@ def build_parser() -> argparse.ArgumentParser: "input_fasta_dir", type=str, help="Input dir containing .fasta files" ) parser.add_argument("output_pdb_dir", type=str, help="Output dir for .pdb files") + parser.add_argument( + "--threads", "-t", type=int, default=mp.cpu_count(), help="Num threads to use" + ) return parser @@ -105,7 +108,7 @@ def main(): ) ) - with mp.Pool(mp.cpu_count()) as pool: + with mp.Pool(args.threads) as pool: pool.starmap(run_faspr, arg_tuples, chunksize=10)