spacr 1.0.9__py3-none-any.whl → 1.1.1__py3-none-any.whl

spacr/sequencing.py

@@ -10,6 +10,23 @@ from IPython.display import display
 
 # Function to map sequences to names (same as your original)
 def map_sequences_to_names(csv_file, sequences, rc):
+    """
+    Maps DNA sequences to their corresponding names based on a CSV file.
+
+    Args:
+        csv_file (str): Path to the CSV file containing 'sequence' and 'name' columns.
+        sequences (list of str): List of DNA sequences to map.
+        rc (bool): If True, reverse complement the sequences in the CSV before mapping.
+
+    Returns:
+        list: A list of names corresponding to the input sequences. If a sequence is not found,
+              `pd.NA` is returned in its place.
+
+    Notes:
+        - The CSV file must contain columns named 'sequence' and 'name'.
+        - If `rc` is True, sequences in the CSV will be reverse complemented prior to mapping.
+        - Sequences in `sequences` are not altered—only sequences in the CSV are reverse complemented.
+    """
     def rev_comp(dna_sequence):
         complement_dict = {'A': 'T', 'T': 'A', 'C': 'G', 'G': 'C', 'N': 'N'}
         reverse_seq = dna_sequence[::-1]
@@ -24,6 +41,24 @@ def map_sequences_to_names(csv_file, sequences, rc):
 
 # Functions to save data (same as your original)
 def save_df_to_hdf5(df, hdf5_file, key='df', comp_type='zlib', comp_level=5):
+    """
+    Saves a pandas DataFrame to an HDF5 file, optionally appending to an existing dataset.
+
+    Args:
+        df (pd.DataFrame): The DataFrame to save.
+        hdf5_file (str): Path to the target HDF5 file.
+        key (str, optional): Key under which to store the DataFrame. Defaults to 'df'.
+        comp_type (str, optional): Compression algorithm to use (e.g., 'zlib', 'bzip2', 'blosc'). Defaults to 'zlib'.
+        comp_level (int, optional): Compression level (0–9). Higher values yield better compression at the cost of speed. Defaults to 5.
+
+    Returns:
+        None
+
+    Notes:
+        - If the specified key already exists in the HDF5 file, the new DataFrame is appended to it.
+        - The combined DataFrame is saved in 'table' format to support appending and querying.
+        - Errors encountered during saving are printed to standard output.
+    """
     try:
         with pd.HDFStore(hdf5_file, 'a', complib=comp_type, complevel=comp_level) as store:
             if key in store:
@@ -34,6 +69,23 @@ def save_df_to_hdf5(df, hdf5_file, key='df', comp_type='zlib', comp_level=5):
         print(f"Error while saving DataFrame to HDF5: {e}")
 
 def save_unique_combinations_to_csv(unique_combinations, csv_file):
+    """
+    Saves or appends a DataFrame of unique gRNA combinations to a CSV file, aggregating duplicates.
+
+    Args:
+        unique_combinations (pd.DataFrame): DataFrame containing 'rowID', 'columnID', and 'grna_name' columns,
+                                            along with associated count or metric columns.
+        csv_file (str): Path to the CSV file where data will be saved.
+
+    Returns:
+        None
+
+    Notes:
+        - If the file exists, it reads the existing contents and appends the new data.
+        - Duplicate combinations (same 'rowID', 'columnID', 'grna_name') are summed.
+        - The resulting DataFrame is saved with index included.
+        - Any exception during the process is caught and printed to stdout.
+    """
     try:
         try:
             existing_df = pd.read_csv(csv_file)
@@ -50,6 +102,22 @@ def save_unique_combinations_to_csv(unique_combinations, csv_file):
         print(f"Error while saving unique combinations to CSV: {e}")
 
 def save_qc_df_to_csv(qc_df, qc_csv_file):
+    """
+    Saves or appends a QC (quality control) DataFrame to a CSV file by summing overlapping entries.
+
+    Args:
+        qc_df (pd.DataFrame): DataFrame containing numeric QC metrics (e.g., counts, read stats).
+        qc_csv_file (str): Path to the CSV file where the QC data will be saved.
+
+    Returns:
+        None
+
+    Notes:
+        - If the file exists, it reads the existing QC data and adds the new values to it (element-wise).
+        - If the file doesn't exist, it creates a new one.
+        - The final DataFrame is saved without including the index.
+        - Any exception is caught and logged to stdout.
+    """
     try:
         try:
             existing_qc_df = pd.read_csv(qc_csv_file)
@@ -64,9 +132,34 @@ def save_qc_df_to_csv(qc_df, qc_csv_file):
         print(f"Error while saving QC DataFrame to CSV: {e}")
 
 def extract_sequence_and_quality(sequence, quality, start, end):
+    """
+    Extracts a subsequence and its corresponding quality scores.
+
+    Args:
+        sequence (str): DNA sequence string.
+        quality (str): Quality string corresponding to the sequence.
+        start (int): Start index of the region to extract.
+        end (int): End index of the region to extract (exclusive).
+
+    Returns:
+        tuple: (subsequence, subquality) as strings.
+    """
     return sequence[start:end], quality[start:end]
 
 def create_consensus(seq1, qual1, seq2, qual2):
+    """
+    Constructs a consensus DNA sequence from two reads with associated quality scores.
+
+    Args:
+        seq1 (str): First DNA sequence.
+        qual1 (str): Quality scores for `seq1` (as ASCII characters or integer-encoded).
+        seq2 (str): Second DNA sequence.
+        qual2 (str): Quality scores for `seq2`.
+
+    Returns:
+        str: Consensus sequence, selecting the base with the highest quality at each position.
+             If one base is 'N', the non-'N' base is chosen regardless of quality.
+    """
     consensus_seq = []
     for i in range(len(seq1)):
         bases = [(seq1[i], qual1[i]), (seq2[i], qual2[i])]
@@ -74,6 +167,15 @@ def create_consensus(seq1, qual1, seq2, qual2):
     return ''.join(consensus_seq)
 
 def get_consensus_base(bases):
+    """
+    Selects the most reliable base from a list of two base-quality pairs.
+
+    Args:
+        bases (list of tuples): Each tuple contains (base, quality_score), expected length is 2.
+
+    Returns:
+        str: The consensus base. Prefers non-'N' bases and higher quality scores.
+    """
     # Prefer non-'N' bases, if 'N' exists, pick the other one.
     if bases[0][0] == 'N':
         return bases[1][0]
@@ -84,13 +186,78 @@ def get_consensus_base(bases):
         return bases[0][0] if bases[0][1] >= bases[1][1] else bases[1][0]
 
 def reverse_complement(seq):
+    """
+    Computes the reverse complement of a DNA sequence.
+
+    Args:
+        seq (str): Input DNA sequence.
+
+    Returns:
+        str: Reverse complement of the input sequence.
+    """
     return str(Seq(seq).reverse_complement())
 
 # Core logic for processing a chunk (same as your original)
 def process_chunk(chunk_data):
-    
+    """
+    Processes a chunk of sequencing reads to extract and map barcode sequences to corresponding names.
+
+    This function handles both single-end and paired-end FASTQ data. It searches for a target barcode
+    sequence in each read, extracts a consensus region around it, applies a regex to extract barcodes,
+    and maps those to known IDs using reference CSVs. Quality control data and unique combinations are
+    also computed.
+
+    Args:
+        chunk_data (tuple): Contains either 9 or 10 elements:
+
+            For paired-end mode (10 elements):
+                - r1_chunk (list): List of strings, each 4-line block from R1 FASTQ.
+                - r2_chunk (list): List of strings, each 4-line block from R2 FASTQ.
+                - regex (str): Regex pattern with named groups ('rowID', 'columnID', 'grna').
+                - target_sequence (str): Sequence to anchor barcode extraction.
+                - offset_start (int): Offset from target_sequence to start consensus extraction.
+                - expected_end (int): Length of the region to extract.
+                - column_csv (str): Path to column barcode reference CSV.
+                - grna_csv (str): Path to gRNA barcode reference CSV.
+                - row_csv (str): Path to row barcode reference CSV.
+                - fill_na (bool): Whether to fill unmapped names with raw barcode sequences.
+
+            For single-end mode (9 elements):
+                - Same as above, but r2_chunk is omitted.
+
+    Returns:
+        tuple:
+            - df (pd.DataFrame): Full dataframe with columns:
+              ['read', 'column_sequence', 'columnID', 'row_sequence', 'rowID',
+              'grna_sequence', 'grna_name']
+            - unique_combinations (pd.DataFrame): Count of each unique (rowID, columnID, grna_name) triplet.
+            - qc_df (pd.DataFrame): Summary of missing values and total reads.
+    """
     def paired_find_sequence_in_chunk_reads(r1_chunk, r2_chunk, target_sequence, offset_start, expected_end, regex):
-
+        """
+        Processes paired-end FASTQ read chunks to extract consensus barcode sequences and decode them
+        using a regex pattern.
+
+        For each R1–R2 read pair, this function identifies the `target_sequence`, extracts a window of 
+        defined length with an offset, computes a consensus sequence using base quality scores, and 
+        applies a regex to extract barcode components.
+
+        Args:
+            r1_chunk (list of str): List of 4-line strings for each R1 read in the chunk.
+            r2_chunk (list of str): List of 4-line strings for each R2 read in the chunk.
+            target_sequence (str): Nucleotide sequence used as anchor for barcode extraction.
+            offset_start (int): Position offset from `target_sequence` to begin extracting barcode.
+            expected_end (int): Total length of region to extract after offset.
+            regex (str): Regular expression with named groups ('rowID', 'columnID', 'grna') 
+                        to parse barcodes from the extracted consensus sequence.
+
+        Returns:
+            tuple:
+                consensus_sequences (list of str): Consensus DNA sequences extracted from read pairs.
+                columns (list of str): Extracted column barcode sequences.
+                grnas (list of str): Extracted gRNA barcode sequences.
+                rows (list of str): Extracted row barcode sequences.
+        """
         consensus_sequences, columns, grnas, rows = [], [], [], []
         consensus_seq = None
         
@@ -154,6 +321,26 @@ def process_chunk(chunk_data):
         return consensus_sequences, columns, grnas, rows
     
     def single_find_sequence_in_chunk_reads(r1_chunk, target_sequence, offset_start, expected_end, regex):
+        """
+        Processes single-end FASTQ read chunks to extract barcode sequences using a target motif and regex pattern.
+
+        For each R1 read, the function identifies the `target_sequence`, extracts a region starting at an offset 
+        and of fixed length, pads if necessary, and applies a regex with named groups to decode barcodes.
+
+        Args:
+            r1_chunk (list of str): List of 4-line strings for each R1 read in the chunk.
+            target_sequence (str): Anchor sequence to locate the barcode region in R1.
+            offset_start (int): Position offset from the end of `target_sequence` to start barcode extraction.
+            expected_end (int): Total length of the barcode region to extract.
+            regex (str): Regular expression with named groups ('rowID', 'columnID', 'grna') to extract barcodes.
+
+        Returns:
+            tuple:
+                consensus_sequences (list of str): Extracted sequences used as barcode consensus (R1 only).
+                columns (list of str): Extracted column barcode subsequences.
+                grnas (list of str): Extracted gRNA barcode subsequences.
+                rows (list of str): Extracted row barcode subsequences.
+        """
 
         consensus_sequences, columns, grnas, rows = [], [], [], []
 
@@ -251,6 +438,20 @@ def process_chunk(chunk_data):
 
 # Function to save data from the queue
 def saver_process(save_queue, hdf5_file, save_h5, unique_combinations_csv, qc_csv_file, comp_type, comp_level):
+    """
+    Continuously reads data from a multiprocessing queue and saves it to disk in various formats.
+
+    This function is intended to run in a separate process. It terminates when it receives the "STOP" sentinel value.
+
+    Args:
+        save_queue (multiprocessing.Queue): Queue containing tuples of (df, unique_combinations, qc_df).
+        hdf5_file (str): Path to the HDF5 file to store full reads (only used if save_h5 is True).
+        save_h5 (bool): Whether to save the full reads DataFrame to HDF5.
+        unique_combinations_csv (str): Path to the CSV file for aggregated barcode combinations.
+        qc_csv_file (str): Path to the CSV file for quality control statistics.
+        comp_type (str): Compression algorithm for HDF5 (e.g., 'zlib').
+        comp_level (int): Compression level for HDF5.
+    """
     while True:
         item = save_queue.get()
         if item == "STOP":
@@ -262,7 +463,33 @@ def saver_process(save_queue, hdf5_file, save_h5, unique_combinations_csv, qc_cs
         save_qc_df_to_csv(qc_df, qc_csv_file)
 
 def paired_read_chunked_processing(r1_file, r2_file, regex, target_sequence, offset_start, expected_end, column_csv, grna_csv, row_csv, save_h5, comp_type, comp_level, hdf5_file, unique_combinations_csv, qc_csv_file, chunk_size=10000, n_jobs=None, test=False, fill_na=False):
-
+    """
+    Processes paired-end FASTQ files in chunks to extract barcoded sequences and generate consensus reads.
+
+    This function identifies sequences matching a regular expression in both R1 and R2 reads, extracts barcodes,
+    and maps them to user-defined identifiers. Processed data is saved incrementally using a separate process.
+
+    Args:
+        r1_file (str): Path to the gzipped R1 FASTQ file.
+        r2_file (str): Path to the gzipped R2 FASTQ file.
+        regex (str): Regular expression with named capture groups: 'rowID', 'columnID', and 'grna'.
+        target_sequence (str): Anchor sequence to align from.
+        offset_start (int): Offset from anchor to start consensus extraction.
+        expected_end (int): Length of the consensus region to extract.
+        column_csv (str): Path to CSV file mapping column barcode sequences to IDs.
+        grna_csv (str): Path to CSV file mapping gRNA barcode sequences to names.
+        row_csv (str): Path to CSV file mapping row barcode sequences to IDs.
+        save_h5 (bool): Whether to save the full reads DataFrame to HDF5.
+        comp_type (str): Compression algorithm for HDF5 (e.g., 'zlib').
+        comp_level (int): Compression level for HDF5.
+        hdf5_file (str): Path to the HDF5 output file.
+        unique_combinations_csv (str): Path to CSV file for saving unique row/column/gRNA combinations.
+        qc_csv_file (str): Path to CSV file for saving QC summary (e.g., NaN counts).
+        chunk_size (int, optional): Number of reads per batch. Defaults to 10000.
+        n_jobs (int, optional): Number of parallel workers. Defaults to cpu_count() - 3.
+        test (bool, optional): If True, processes only a single chunk and prints the result. Defaults to False.
+        fill_na (bool, optional): If True, fills unmapped IDs with raw barcode sequences. Defaults to False.
+    """
     from .utils import count_reads_in_fastq, print_progress
 
     # Use cpu_count minus 3 cores if n_jobs isn't specified
@@ -344,7 +571,34 @@ def paired_read_chunked_processing(r1_file, r2_file, regex, target_sequence, off
     save_process.join()
 
 def single_read_chunked_processing(r1_file, r2_file, regex, target_sequence, offset_start, expected_end, column_csv, grna_csv, row_csv, save_h5, comp_type, comp_level, hdf5_file, unique_combinations_csv, qc_csv_file, chunk_size=10000, n_jobs=None, test=False, fill_na=False):
-
+    """
+    Processes single-end FASTQ data in chunks to extract barcoded sequences and map them to identifiers.
+
+    This function reads gzipped R1 FASTQ data, detects barcode-containing sequences using a target anchor and regex,
+    and maps row, column, and gRNA barcodes to user-defined identifiers. Results are processed in parallel
+    and saved incrementally via a background process.
+
+    Args:
+        r1_file (str): Path to gzipped R1 FASTQ file.
+        r2_file (str): Placeholder for interface consistency; not used in single-end mode.
+        regex (str): Regular expression with named capture groups: 'rowID', 'columnID', and 'grna'.
+        target_sequence (str): Anchor sequence used to locate barcode region.
+        offset_start (int): Offset from anchor to start barcode parsing.
+        expected_end (int): Length of the barcode region to extract.
+        column_csv (str): Path to CSV file mapping column barcode sequences to IDs.
+        grna_csv (str): Path to CSV file mapping gRNA barcode sequences to names.
+        row_csv (str): Path to CSV file mapping row barcode sequences to IDs.
+        save_h5 (bool): Whether to save the full reads DataFrame to HDF5 format.
+        comp_type (str): Compression algorithm for HDF5 (e.g., 'zlib').
+        comp_level (int): Compression level for HDF5.
+        hdf5_file (str): Output HDF5 file path.
+        unique_combinations_csv (str): Output path for CSV summarizing row/column/gRNA combinations.
+        qc_csv_file (str): Output path for CSV summarizing missing values and total reads.
+        chunk_size (int, optional): Number of reads per batch. Defaults to 10,000.
+        n_jobs (int, optional): Number of parallel worker processes. Defaults to cpu_count() - 3.
+        test (bool, optional): If True, processes only the first chunk and prints its result. Defaults to False.
+        fill_na (bool, optional): If True, fills missing mapped IDs with their corresponding barcode sequences. Defaults to False.
+    """
     from .utils import count_reads_in_fastq, print_progress
 
     # Use cpu_count minus 3 cores if n_jobs isn't specified
@@ -422,7 +676,40 @@ def single_read_chunked_processing(r1_file, r2_file, regex, target_sequence, off
     save_process.join()
 
 def generate_barecode_mapping(settings={}):
-
+    """
+    Orchestrates barcode extraction and mapping from gzipped sequencing data using user-defined or default settings.
+
+    This function parses sequencing reads from single-end or paired-end FASTQ (.gz) files, extracts barcode regions
+    using a regular expression, maps them to row, column, and gRNA identifiers, and saves the results to disk.
+    Results include the full annotated reads (optional), barcode combination counts, and a QC summary.
+
+    Args:
+        settings (dict, optional): Dictionary containing parameters required for barcode mapping. If not provided,
+            default values will be applied. Important keys include:
+            - 'src' (str): Source directory containing gzipped FASTQ files.
+            - 'mode' (str): Either 'single' or 'paired' for single-end or paired-end processing.
+            - 'single_direction' (str): If 'single', specifies which read to use ('R1' or 'R2').
+            - 'regex' (str): Regular expression with capture groups 'rowID', 'columnID', and 'grna'.
+            - 'target_sequence' (str): Anchor sequence to locate barcode start position.
+            - 'offset_start' (int): Offset from the anchor to the barcode start.
+            - 'expected_end' (int): Expected barcode region length.
+            - 'column_csv' (str): CSV file mapping column barcodes to names.
+            - 'grna_csv' (str): CSV file mapping gRNA barcodes to names.
+            - 'row_csv' (str): CSV file mapping row barcodes to names.
+            - 'save_h5' (bool): Whether to save annotated reads to HDF5.
+            - 'comp_type' (str): Compression algorithm for HDF5.
+            - 'comp_level' (int): Compression level for HDF5.
+            - 'chunk_size' (int): Number of reads to process per batch.
+            - 'n_jobs' (int): Number of parallel processes for barcode mapping.
+            - 'test' (bool): If True, only processes the first chunk for testing.
+            - 'fill_na' (bool): If True, fills unmapped barcodes with raw sequence instead of NaN.
+
+    Side Effects:
+        Saves the following files in the output directory:
+        - `annotated_reads.h5` (optional): Annotated read information in HDF5 format.
+        - `unique_combinations.csv`: Count table of (rowID, columnID, grna_name) triplets.
+        - `qc.csv`: Summary of missing values and read counts.
+    """
     from .settings import set_default_generate_barecode_mapping
     from .utils import save_settings
     from .io import parse_gz_files
@@ -490,7 +777,23 @@ def generate_barecode_mapping(settings={}):
 
 # Function to read the CSV, compute reverse complement, and save it
 def barecodes_reverse_complement(csv_file):
+    """
+    Reads a barcode CSV file, computes the reverse complement of each sequence, and saves the result to a new CSV.
+
+    This function assumes the input CSV contains a column named 'sequence' with DNA barcodes. It computes the
+    reverse complement for each sequence and saves the modified DataFrame to a new file with '_RC' appended
+    to the original filename.
+
+    Args:
+        csv_file (str): Path to the input CSV file. Must contain a column named 'sequence'.
+
+    Side Effects:
+        - Saves a new CSV file in the same directory with reverse-complemented sequences.
+        - Prints the path of the saved file.
 
+    Output:
+        New file path format: <original_filename>_RC.csv
+    """
     def reverse_complement(sequence):
         complement = {'A': 'T', 'T': 'A', 'G': 'C', 'C': 'G', 'N': 'N'}
         return ''.join(complement[base] for base in reversed(sequence))
@@ -512,10 +815,48 @@ def barecodes_reverse_complement(csv_file):
     print(f"Reverse complement file saved as {new_filename}")
 
 def graph_sequencing_stats(settings):
-
+    """
+    Analyze and visualize sequencing quality metrics to determine an optimal fraction threshold
+    that maximizes unique gRNA representation per well across plates.
+
+    This function reads one or more CSV files containing count data, filters out control wells,
+    calculates the fraction of reads per gRNA in each well, and identifies the minimum fraction 
+    required to recover a target average number of unique gRNAs per well. It generates plots to 
+    help visualize the chosen threshold and spatial distribution of unique gRNA counts.
+
+    Args:
+        settings (dict): Dictionary containing the following keys:
+            - 'count_data' (str or list of str): Paths to CSV file(s) with 'grna', 'count', 'rowID', 'columnID' columns.
+            - 'target_unique_count' (int): Target number of unique gRNAs per well to recover.
+            - 'filter_column' (str): Column name to filter out control wells.
+            - 'control_wells' (list): List of control well labels to exclude.
+            - 'log_x' (bool): Whether to log-scale the x-axis in the threshold plot.
+            - 'log_y' (bool): Whether to log-scale the y-axis in the threshold plot.
+
+    Returns:
+        float: Closest fraction threshold that approximates the target unique gRNA count per well.
+
+    Side Effects:
+        - Saves a PDF plot of unique gRNA count vs fraction threshold.
+        - Saves a spatial plate map of unique gRNA counts.
+        - Prints threshold and summary statistics.
+        - Displays intermediate DataFrames for inspection.
+    """
     from .utils import correct_metadata_column_names, correct_metadata
 
     def _plot_density(df, dependent_variable, dst=None):
+        """
+        Plot a kernel density estimate (KDE) of a specified variable from a DataFrame.
+
+        Args:
+            df (pd.DataFrame): DataFrame containing the data.
+            dependent_variable (str): Name of the column to plot.
+            dst (str, optional): Directory to save the plot. If None, the plot is not saved.
+
+        Side Effects:
+            - Displays the KDE plot.
+            - Saves the plot as 'dependent_variable_density.pdf' in the specified directory if dst is provided.
+        """
         """Plot a density plot of the dependent variable."""
         plt.figure(figsize=(10, 10))
         sns.kdeplot(df[dependent_variable], fill=True, alpha=0.6)
@@ -530,8 +871,22 @@ def graph_sequencing_stats(settings):
 
     def find_and_visualize_fraction_threshold(df, target_unique_count=5, log_x=False, log_y=False, dst=None):
         """
-        Find the fraction threshold where the recalculated unique count matches the target value,
-        and visualize the relationship between fraction thresholds and unique counts.
+        Identify the optimal fraction threshold that yields an average number of unique gRNAs per well 
+        closest to a specified target, and visualize the relationship between threshold and unique count.
+
+        Args:
+            df (pd.DataFrame): Input DataFrame containing 'fraction', 'plateID', 'rowID', 'columnID', and 'grna' columns.
+            target_unique_count (int, optional): Desired average number of unique gRNAs per well. Default is 5.
+            log_x (bool, optional): Whether to apply a log scale to the x-axis in the plot.
+            log_y (bool, optional): Whether to apply a log scale to the y-axis in the plot.
+            dst (str, optional): Directory where the plot will be saved. If None, the plot is not saved.
+
+        Returns:
+            float: The fraction threshold value closest to achieving the target_unique_count.
+
+        Side Effects:
+            - Displays a line plot of unique gRNA counts vs. fraction thresholds.
+            - Saves the plot as 'fraction_threshold.pdf' in a subdirectory 'results/' under `dst` if provided.
         """
 
         def _line_plot(df, x='fraction_threshold', y='unique_count', log_x=False, log_y=False):