create_primer_lookup: added doc strings and improved var names

Author: ChrisgKent

align_trim/main.py

@@ -748,42 +748,72 @@ def read_pair_generator(bam, region_string=None):
 
 def create_primer_lookup(ref_len_tuple, amplicons: list[Amplicon], padding=35):
     """
-    Returns a dict of chroms, each containing a (N, chrom_len) shaped array
+    Create a lookup table for efficient primer position queries across reference genomes.
+
+    Each chromosome gets its own 2D lookup array where:
+    - Rows represent non-overlapping "pools"* of amplicons at their corresponding positions.
+    - Columns represent genomic positions
+    - Values are Amplicon objects or None
+
+    The function automatically determines the minimum number of rows needed to ensure
+    no amplicons overlap within the same row when accounting for padding.
+
+    * Amplicons are placed in the first available row where they don't overlap, not their pool index.
+
+    Parameters
+    ----------
+    ref_len_tuple : list[tuple[str, int]]
+        List of tuples containing (chromosome_name, chromosome_length) pairs
+        from the reference genome
+    amplicons : list[Amplicon]
+        List of Amplicon objects containing primer scheme information
+    padding : int, optional
+        Number of bases to extend amplicon boundaries on both sides to allow
+        for fuzzy matching of reads with barcodes/adapters (default: 35)
+
+    Returns
+    -------
+    dict[str, np.ndarray]
+        Dictionary mapping chromosome names to 2D numpy arrays of shape (N, chrom_len+1)
+        where N is the minimum number of rows needed to prevent amplicon overlap.
+        Array elements are either Amplicon objects or None.
+
+
     """
     lookups = {}
     for chrom, chromlen in ref_len_tuple:
-        a = np.empty_like(None, shape=(1, chromlen + 1))
+        lookup_array = np.empty_like(None, shape=(1, chromlen + 1))
         for amp in amplicons:
             added = False
             if amp.chrom == chrom:
                 # If amplicon clashes with any in same pool add new row
-                amp_slice = a[
+                amp_slice = lookup_array[
                     :,
                     max(amp.amplicon_start - padding, 0) : min(
                         amp.amplicon_end + padding, chromlen
                     ),
                 ]
                 for i, row in enumerate(amp_slice):  # Check each row for collision
                     if row[row != None].size == 0:
-                        a[
+                        lookup_array[
                             i,
                             max(amp.amplicon_start - padding, 0) : min(
                                 amp.amplicon_end + padding, chromlen
                             ),
                         ] = amp
                         added = True
-                # If not added, add new row to array and add to that.
+                # If not added, create new row, add the amplicon to that then add back to original array
                 if not added:
-                    b = np.empty_like(None, shape=(1, chromlen + 1))
-                    b[
+                    new_row = np.empty_like(None, shape=(1, chromlen + 1))
+                    new_row[
                         0,
                         max(amp.amplicon_start - padding, 0) : min(
                             amp.amplicon_end + padding, chromlen
                         ),
                     ] = amp
-                    a = np.vstack((a, b))
+                    lookup_array = np.vstack((lookup_array, new_row))
 
-        lookups[chrom] = a
+        lookups[chrom] = lookup_array
     return lookups