Merge branch 'dev' into main

Author: dltamayo

bin/tcrdist3_matrix.py

@@ -8,20 +8,21 @@
 import pandas as pd
 from tcrdist.repertoire import TCRrep
 
-def reverse_transform_trbv(trbv):
-    """Convert TCRBV notation back to TRBV format, remove zero padding before *, and handle /OR cases."""
+def transform_trbv(trbv):
+    """Convert gene names from Adaptive ImmunoSEQ to IMGT format."""
     if not isinstance(trbv, str):
         return trbv  # Return as-is if not a string
 
-    trbv = trbv.replace("TCRBV", "TRBV")  # Convert TCRBV → TRBV
+    # Convert locus name
+    trbv = trbv.replace("TCRBV", "TRBV")
 
-    # Remove zero padding from main number (TCRBV07 → TRBV7)
+    # Remove zero padding from gene name (TCRBV07 to TRBV7)
     trbv = re.sub(r'(?<=TRBV)0*(\d+)', r'\1', trbv)  
 
-    # Remove zero padding from subgroup (TCRBV7-02 → TRBV7-2)
+    # Remove zero padding from subgroup (TCRBV7-02 to TRBV7-2)
     trbv = re.sub(r'-(0\d+)', lambda m: f'-{int(m.group(1))}', trbv)  
 
-    # Convert "-orXX_XX" format back to "/OR#-#"
+    # Convert "-orXX_XX" to "/OR#-#" for orphon genes
     trbv = re.sub(r'-or0?(\d+)_0?(\d+)', r'/OR\1-\2', trbv)
 
     # Add *01 if allele group not specified
@@ -31,16 +32,22 @@ def reverse_transform_trbv(trbv):
     return trbv
 
 def remove_locus(gene_name):
-    """If gene is in TCRBVXX-##*0# format, try removing the -##."""
-    return re.sub(r'-(\d+)\*', '*', gene_name)
+    """Remove the -## gene position from TRBV names unless the gene contains /OR."""
+    if '/OR' in gene_name:
+        return gene_name
+    else:
+        return re.sub(r'-(\d+)\*', '*', gene_name)
 
 def split_and_check_genes(gene_name):
-    """Handle cases where two genes are combined (TCRBVXX-YY/XX-ZZ*0#) and return both separately."""
-    if '/' in gene_name and not re.search(r'/OR\d+-\d+', gene_name):  # Ensure it's not an OR case
-        base, star_part = gene_name.split("*") if "*" in gene_name else (gene_name, "01")  
+    """Split combined TCRBV genes (e.g., TCRBV06-02/06-03*01 to TCRBV06-02*01 and TCRBV06-03*01."""
+    if '/' in gene_name and not re.search(r'/OR\d+-\d+', gene_name):  # Ensure it's not an orphon
+        base, allele = gene_name.split("*") if "*" in gene_name else (gene_name, "01")
+        prefix_match = re.match(r"(TCRBV\d+)", gene_name)
+        prefix = prefix_match.group(1) if prefix_match else "TCRBV"  # Fallback just in case
         genes = base.split("/")  # Split the genes
-        return [f"{g}*{star_part}" for g in genes]  # Reattach the *0# part to both genes
-    return [gene_name]  # Return as list for consistency
+        return [f"{prefix}-{g.split('-')[-1]}*{allele}" for g in genes]
+    return [gene_name]
+
 
 def find_matching_gene(row, db):
     # Collect all possible genes from vMaxResolved and vGeneNameTies
@@ -57,12 +64,12 @@ def find_matching_gene(row, db):
         if "/" in gene and not re.search(r"/OR\d+-\d+", gene):  # Avoid /OR cases
             sub_genes = split_and_check_genes(gene)
             for sub_gene in sub_genes:
-                sub_gene = reverse_transform_trbv(sub_gene)  # Ensure correct *0# format
+                sub_gene = transform_trbv(sub_gene)  # Ensure correct *0# format
                 if sub_gene in db["id"].values:
                     return sub_gene
 
         # Direct match in db
-        transform_gene = reverse_transform_trbv(gene)
+        transform_gene = transform_trbv(gene)
         if transform_gene in db["id"].values:
             return transform_gene
 
@@ -71,60 +78,62 @@ def find_matching_gene(row, db):
         if modified_gene in db["id"].values:
             return modified_gene
 
-    transform_row = reverse_transform_trbv(row["vMaxResolved"])
+    transform_row = transform_trbv(row["vMaxResolved"])
     print(f'No match found for {transform_row}')
 
     return transform_row  # Return original vMaxResolved if no match is found
 
-# Parse input arguments
-parser = argparse.ArgumentParser(description="Take positional args")
 
-parser.add_argument("sample_tsv")
-parser.add_argument("ref_database")
-parser.add_argument("cores", type=int)
+if __name__ == "__main__":
+    # Parse input arguments
+    parser = argparse.ArgumentParser(description="Take positional args")
+
+    parser.add_argument("sample_tsv")
+    parser.add_argument("ref_database")
+    parser.add_argument("cores", type=int)
 
-args = parser.parse_args()
+    args = parser.parse_args()
 
-print(f"sample_tsv: {args.sample_tsv}")
-print(f"ref_database: {args.ref_database}")
-print(f"cores: {args.cores}")
+    print(f"sample_tsv: {args.sample_tsv}")
+    print(f"ref_database: {args.ref_database}")
+    print(f"cores: {args.cores}")
 
-sample_tsv = args.sample_tsv
+    sample_tsv = args.sample_tsv
 
-# Get the basename
-basename = os.path.splitext(os.path.basename(sample_tsv))[0]
+    # Get the basename
+    basename = os.path.splitext(os.path.basename(sample_tsv))[0]
 
-# --- 1. Convert Adaptive output to tcrdist db format ---
-db = pd.read_table(args.ref_database, delimiter = '\t')
+    # --- 1. Convert Adaptive output to tcrdist db format ---
+    db = pd.read_table(args.ref_database, delimiter = '\t')
 
-db = db[db['organism']=='human']
+    db = db[db['organism']=='human']
 
-df = pd.read_table(sample_tsv, delimiter = '\t')
+    df = pd.read_table(sample_tsv, delimiter = '\t')
 
-df = df[['nucleotide', 'aminoAcid', 'vMaxResolved', 'vGeneNameTies', 'count (templates/reads)']]
-df["vMaxResolved"] = df.apply(lambda row: find_matching_gene(row, db), axis=1)
+    df = df[['nucleotide', 'aminoAcid', 'vMaxResolved', 'vGeneNameTies', 'count (templates/reads)']]
+    df["vMaxResolved"] = df.apply(lambda row: find_matching_gene(row, db), axis=1)
 
-df = df.rename(columns={'nucleotide': 'cdr3_b_nucseq',
-                    'aminoAcid': 'cdr3_b_aa',
-                    # 'CDR3a': 'cdr3_a_aa', 
-                    'vMaxResolved': 'v_b_gene',
-                    # 'TRBJ': 'j_b_gene',
-                    'count (templates/reads)': 'count'})
+    df = df.rename(columns={'nucleotide': 'cdr3_b_nucseq',
+                        'aminoAcid': 'cdr3_b_aa',
+                        # 'CDR3a': 'cdr3_a_aa', 
+                        'vMaxResolved': 'v_b_gene',
+                        # 'TRBJ': 'j_b_gene',
+                        'count (templates/reads)': 'count'})
 
-df = df[df['cdr3_b_aa'].notna()]
-df = df[df['v_b_gene'].notna()]
-df = df.drop('vGeneNameTies', axis=1)
+    df = df[df['cdr3_b_aa'].notna()]
+    df = df[df['v_b_gene'].notna()]
+    df = df.drop('vGeneNameTies', axis=1)
 
-# --- 2. Calculate sparse distance matrix ---
-tr = TCRrep(cell_df = df,
-            organism = 'human',
-            chains = ['beta'],
-            db_file = 'alphabeta_gammadelta_db.tsv',
-            compute_distances = False)
-tr.cpus = args.cores
-tr.compute_distances()
+    # --- 2. Calculate sparse distance matrix ---
+    tr = TCRrep(cell_df = df,
+                organism = 'human',
+                chains = ['beta'],
+                db_file = 'alphabeta_gammadelta_db.tsv',
+                compute_distances = False)
+    tr.cpus = args.cores
+    tr.compute_distances()
 
-np.savetxt(f"{basename}_distance_matrix.csv", tr.pw_beta, delimiter=",", fmt="%d")
+    np.savetxt(f"{basename}_distance_matrix.csv", tr.pw_beta, delimiter=",", fmt="%d")
 
-clone_df = tr.clone_df
-clone_df.to_csv(f"{basename}_clone_df.csv", index=False)
+    clone_df = tr.clone_df
+    clone_df.to_csv(f"{basename}_clone_df.csv", index=False)

bin/test.py

@@ -0,0 +1,38 @@
+import unittest
+import tcrdist3_matrix as t3m
+
+class TestNameConversion(unittest.TestCase):
+    def setUp(self):
+        # This method is called before each test
+        self.test_cases = [
+            ("TCRBV07", "TRBV7*01", "TRBV7*01"),
+            ("TCRBV27", "TRBV27*01", "TRBV27*01"),
+            ("TCRBV07-02", "TRBV7-2*01", "TRBV7*01"),
+            ("TCRBV17-02", "TRBV17-2*01", "TRBV17*01"),
+            ("TCRBV10-03*02", "TRBV10-3*02", "TRBV10*02"),
+            ("TCRBV07-02*01", "TRBV7-2*01", "TRBV7*01"),
+            ("TCRBV10-or09_02*01", "TRBV10/OR9-2*01", "TRBV10/OR9-2*01"),
+        ]
+        self.split_cases = [
+            ("TCRBV06-02/06-03*01", ["TCRBV06-02*01", "TCRBV06-03*01"]),
+            ("TCRBV12-03/12-04", ["TCRBV12-03*01","TCRBV12-04*01"])
+     ]
+
+    def test_transform_trbv(self):
+        for trbv_input, expected_output, _ in self.test_cases:
+            result = t3m.transform_trbv(trbv_input)
+            self.assertEqual(result, expected_output, f"Failed for input: {trbv_input}")
+    
+    def test_remove_locus(self):
+        for _, trbv_output, remove_locus_output in self.test_cases:
+            result = t3m.remove_locus(trbv_output)
+            self.assertEqual(result, remove_locus_output, f"Failed for input: {trbv_output}")
+    
+    def test_split_and_check_genes(self):
+        for split_input, split_output in self.split_cases:
+            result = t3m.split_and_check_genes(split_input)
+            self.assertEqual(result, split_output, f"Failed for input: {split_input}")
+
+
+if __name__ == '__main__':
+    unittest.main()

env.yml

@@ -18,6 +18,7 @@ dependencies:
   - s3fs=2024.3.1
   - python-igraph=0.11.8
   - scikit-learn=1.6.1
+  - olga=1.2.4
 
   # R and R packages
   - r-base=4.4.2

modules/local/olga.nf

@@ -0,0 +1,31 @@
+process OLGA {
+    tag "${sample_meta[0]}"
+    label 'process_low'
+    container "ghcr.io/break-through-cancer/bulktcr:latest"
+    
+    input:
+    tuple val(sample_meta), path(count_table)
+    
+    output:
+    path "${count_table.baseName}_tcr_generation_probabilities.tsv", emit: "olga_output"
+    
+    script:
+    """
+    # Extract vector of cdr3 aa, dropping null values
+    
+    cat > dropAA.py <<EOF
+    
+    import pandas as pd
+    
+    df = pd.read_csv("${count_table}", sep="\t")
+    df = df.dropna(subset=["aminoAcid"])
+    df = df["aminoAcid"]
+    df.to_csv("output.tsv", sep="\t", index=False, header=False)
+    
+    EOF
+    
+    python dropAA.py
+    
+    olga-compute_pgen --humanTRB -i output.tsv -o "${count_table.baseName}_tcr_generation_probabilities.tsv"
+    """
+}

subworkflows/local/sample.nf

@@ -8,6 +8,7 @@
 include { SAMPLE_CALC } from '../../modules/local/sample_calc'
 include { SAMPLE_PLOT } from '../../modules/local/sample_plot'
 include { TCRDIST3_MATRIX } from '../../modules/local/tcrdist3_matrix'
+include { OLGA } from '../../modules/local/olga'
 
 /*
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -61,6 +62,8 @@ workflow SAMPLE {
         v_family_csv
         )
 
+    OLGA ( sample_map )
+    
     // emit:
     // sample_stats_csv
     // v_family_csv