add subset - ya ni sé

Author: zorbax

atg/tools/app.py

@@ -1,11 +1,12 @@
 from pathlib import Path
 
 import typer
-from psutil import cpu_count
+from multiprocessing import cpu_count
 
 from atg.tools.ena import ena_download, ena_fields, ena_retrieve
 from atg.tools.git import gitcheck
 from atg.tools.manifest import create_manifest
+from atg.tools.sample import subsample_fastq_seqkit
 from atg.utils import OrderCommands, count_fastq, get_abundance, one_liner
 
 tools_app = typer.Typer(
@@ -190,3 +191,55 @@ def count_tools_command(
 def abundance_tools_command():
     """Relative abundance tables"""
     print(get_abundance())
+
+
+@tools_app.command(name="subsample")
+def subsample_seqkit_tools_command(
+    input_path: str = typer.Option(
+        ...,
+        "--input",
+        "-i",
+        show_default=False,
+        help="FASTQ file or directory containing FASTQ files (supports both single-end and paired-end)",
+    ),
+    output_dir: str = typer.Option(
+        ...,
+        "--output",
+        "-o",
+        show_default=False,
+        help="Output directory for subsampled files",
+    ),
+    proportion: float = typer.Option(
+        0.1,
+        "--proportion",
+        "-p",
+        min=0.0,
+        max=1.0,
+        help="Proportion of reads to subsample (0.0 to 1.0)",
+    ),
+    seed: int = typer.Option(
+        42,
+        "--seed",
+        "-s",
+        help="Random seed for reproducible subsampling",
+    ),
+    max_workers: int = typer.Option(
+        cpu_count(),
+        "--workers",
+        "-w",
+        help="Maximum number of worker threads for parallel processing",
+    ),
+):
+    """Create subsamples of FASTQ files with specified proportion using seqkit.
+    
+    Supports both single-end (NAME.fastq.gz, NAME.fastq) and paired-end 
+    (NAME_R1.fastq.gz/NAME_R2.fastq.gz, NAME_1.fastq.gz/NAME_2.fastq.gz) files.
+    Requires seqkit to be installed.
+    """
+    subsample_fastq_seqkit(
+        input_path=input_path,
+        output_dir=output_dir,
+        proportion=proportion,
+        seed=seed,
+        max_workers=max_workers,
+    )

atg/tools/ena.py

@@ -8,8 +8,10 @@
 import requests
 from tabulate import tabulate
 from tqdm import tqdm
-from tqdm.contrib.concurrent import thread_map
 from urllib3 import util
+from loguru import logger
+
+from atg.utils import thread_map_parallel
 
 # ic.configureOutput(prefix=" -> ")
 
@@ -45,16 +47,16 @@ def download_url(input_file: Tuple[str, str, Path]) -> None:
                 shutil.copyfileobj(r_raw, file)
 
     except requests.exceptions.Timeout:
-        print(f"Timeout downloading {fname}. Please retry later.")
+        logger.error(f"Timeout downloading {fname}. Please retry later.")
         raise
     except requests.exceptions.HTTPError as e:
-        print(f"HTTP error downloading {fname}: {e}")
+        logger.error(f"HTTP error downloading {fname}: {e}")
         raise
     except requests.exceptions.ConnectionError:
-        print(f"Connection error downloading {fname}. Please check your internet connection.")
+        logger.error(f"Connection error downloading {fname}. Please check your internet connection.")
         raise
     except Exception as e:
-        print(f"Error downloading {fname}: {e}")
+        logger.error(f"Error downloading {fname}: {e}")
         raise
 
 
@@ -102,10 +104,10 @@ def request_get(api: str, pdict: str, fields: str, safe: str = ",") -> pd.DataFr
 
         return df
     except requests.exceptions.Timeout:
-        print("Connection to the server has timed out. Please retry.")
+        logger.error("Connection to the server has timed out. Please retry.")
         return None
     except requests.exceptions.HTTPError:
-        print("HTTPError: This is likely caused by an invalid search query")
+        logger.error("HTTPError: This is likely caused by an invalid search query")
         return None
 
 
@@ -129,7 +131,7 @@ def ena_fields(id_err: str, save: bool = True, fields: str = "") -> Dict[str, st
 
     if save:
         df.to_csv(f"{id_err}.tsv", sep="\t", index=False)
-        print(f"ENA metadata saved as {id_err}.tsv")
+        logger.info(f"ENA metadata saved as {id_err}.tsv")
     else:
         pass
 
@@ -155,8 +157,7 @@ def md5_hash(filename, block_size=2**20):
 
 def thread_map_urls(url_dict: Dict[str, str], outdir: Path, cpu: int) -> None:
     iter_url = [(k, v, outdir) for k, v in url_dict.items()]
-    if len(iter_url) > 0:
-        thread_map(download_url, iter_url, max_workers=cpu)
+    thread_map_parallel(download_url, iter_url, cpu)
 
 
 def checksums(id_err: str, output_dir: Path, file_lst: List[str], threads: int) -> None:
@@ -179,7 +180,7 @@ def compare_lists(df_md5: pd.DataFrame, test_list: List[str], outdir: Path, n_cp
     md5_failed = [k for k, v in urls_dict.items() if v[0] != md5_hash(output_dir / k)]
 
     if compare_lists(dfmd5, md5_failed, output_dir, threads) is None:
-        print("All files are already downloaded")
+        logger.info("All files are already downloaded")
 
 
 def ena_search(
@@ -192,7 +193,7 @@ def ena_search(
     df = request_get("browser/api/tsv/textsearch", params, fields="all", safe=safe)
 
     if df.empty:
-        print("Check your query")
+        logger.error("Check your query")
 
     return df
 
@@ -207,14 +208,14 @@ def ena_retrieve(keywords: str, save: bool, only_ids: bool):
 
     if save and not only_ids:
         df.to_csv(f"{keywords}.tsv", sep="\t", index=False)
-        print(f"ENA metadata saved as {keywords}.tsv")
+        logger.info(f"ENA metadata saved as {keywords}.tsv")
     elif only_ids and not save:
-        print(tabulate(df[["accession"]], headers="keys", showindex=False))
+        logger.info(tabulate(df[["accession"]], headers="keys", showindex=False))
     elif only_ids and save:
         df["accession"].to_csv(f"{keywords}_ids.tsv", index=False, header=False)
-        print(f"ENA metadata saved as {keywords}_ids.tsv")
+        logger.info(f"ENA metadata saved as {keywords}_ids.tsv")
     else:
-        print(tabulate(df, headers="keys", showindex=False, tablefmt="plain"))
+        logger.info(tabulate(df, headers="keys", showindex=False, tablefmt="plain"))
 
 
 def ena_download(bioproject: str, cpus: int, fields: str = None, output_base_dir: Path = None) -> None:
@@ -230,7 +231,7 @@ def ena_download(bioproject: str, cpus: int, fields: str = None, output_base_dir
     files = [x.name for x in Path.glob(out_dir, "*.fastq.gz")]
 
     if len(files) > 0:
-        print("Verifying MD5 File Checksums...")
+        logger.info("Verifying MD5 File Checksums...")
         checksums(err_id, out_dir, files, cpus)
     else:
         err_urls = ena_urls(ena_fields(id_err=err_id, fields=fields))

atg/utils.py

@@ -4,13 +4,14 @@
 from enum import Enum
 from functools import wraps
 from pathlib import Path
-from typing import Any
+from typing import Any, Union, Dict, Callable, List
 
 import anndata as ad
 import pyfastx
 from click import Context
 from loguru import logger
 from typer.core import TyperGroup
+from tqdm.contrib.concurrent import thread_map
 
 
 def timeit(f: Any) -> Any:
@@ -30,6 +31,18 @@ def wrapper(*args, **kargs):  # type: ignore
     return wrapper
 
 
+def thread_map_parallel(func: Callable, data: List, max_workers: int) -> None:
+    """Execute function in parallel using thread_map.
+    
+    Args:
+        func: Function to execute
+        data: List of data items to process
+        max_workers: Maximum number of worker threads
+    """
+    if len(data) > 0:
+        thread_map(func, data, max_workers=max_workers, desc="Processing")
+
+
 def one_liner(input_fasta: str) -> None:
     """
     Convert multiline FASTA to single line FASTA. The input file is overwritten.
@@ -103,13 +116,23 @@ def check_dir(fastq_dir: str) -> Path:
     return _fastq_dir
 
 
-def fastq_files(fastq: str, pattern: str) -> list:
+def fastq_files(fastq: str, pattern: str) -> dict:
+    """
+    Get FASTQ files from directory or single file.
+    
+    Args:
+        fastq: Path to FASTQ file or directory
+        pattern: Regex pattern to match files
+        
+    Returns:
+        dict: Dictionary mapping sample names to file paths
+    """
     if Path(fastq).is_file():
         fqfile = Path(fastq).name.partition(".")[0]
         return {fqfile: fastq}
 
     check_dir(fastq)
-    pattern = re.compile(r".*_([1-2]|R[1-2]).(fastq|fq)\.gz$")
+    pattern = re.compile(pattern)
     fqfiles = sorted([x for x in Path(fastq).glob("*") if pattern.match(str(x))])
     snames = sorted([str(x.name.partition(".")[0]) for x in fqfiles])
 
@@ -123,3 +146,55 @@ def count_fastq(fastq_file, pattern: str):
         index_file = Path(f"{str(v)}.fxi")
         if index_file.exists():
             index_file.unlink()
+
+
+def check_paired(fastq_files: dict) -> dict:
+    """
+    Check if FASTQ files are paired-end based on naming patterns.
+    
+    Args:
+        fastq_files: Dictionary of FASTQ files from fastq_files function
+        
+    Returns:
+        dict: Dictionary with 'paired' and 'single' keys containing file lists
+    """
+    paired_files = {}
+    single_files = {}
+    
+    # Common paired-end patterns (both compressed and uncompressed)
+    paired_patterns = [
+        (r'_R1\.(fastq|fq)(\.gz)?$', r'_R2\.(fastq|fq)(\.gz)?$'),
+        (r'_1\.(fastq|fq)(\.gz)?$', r'_2\.(fastq|fq)(\.gz)?$'),
+        (r'\.1\.(fastq|fq)(\.gz)?$', r'\.2\.(fastq|fq)(\.gz)?$'),
+    ]
+    
+    for sample_name, file_path in fastq_files.items():
+        file_str = str(file_path)
+        is_paired = False
+        
+        for pattern1, pattern2 in paired_patterns:
+            if re.search(pattern1, file_str):
+                # Find corresponding R2/2 file
+                potential_pair = re.sub(pattern1, pattern2, file_str)
+                if Path(potential_pair).exists():
+                    paired_files[sample_name] = {
+                        'R1': file_path,
+                        'R2': Path(potential_pair)
+                    }
+                    is_paired = True
+                    break
+            elif re.search(pattern2, file_str):
+                # Find corresponding R1/1 file
+                potential_pair = re.sub(pattern2, pattern1, file_str)
+                if Path(potential_pair).exists():
+                    paired_files[sample_name] = {
+                        'R1': Path(potential_pair),
+                        'R2': file_path
+                    }
+                    is_paired = True
+                    break
+        
+        if not is_paired:
+            single_files[sample_name] = file_path
+    
+    return {'paired': paired_files, 'single': single_files}