TCRfeatureCal.py

#!/users/miniconda3/envs/bin/python

# coding: utf-8

from glob import glob
import sys
import os
import pandas as pd
import numpy as np
import argparse
from typing import Dict, Union, List, Optional,Tuple
from dataclasses import dataclass
from collections import Counter, defaultdict


try:
    import Levenshtein
    import parasail
except ImportError:
    raise ImportError(
        "Required packages not found. Please install with:\n"
        "pip install python-levenshtein parasail"
    )






###
@dataclass
class AnalysisResult:
    """Data class for storing analysis results"""
    v_gene: pd.DataFrame
    j_gene: pd.DataFrame
    aa_composition: pd.DataFrame  #tcr sequence aa composition
    length_group: pd.DataFrame  # tcr length group 
    convergence: pd.DataFrame  # 
    freq_group: pd.DataFrame  # clone frequency group
    enriched_score: pd.DataFrame  # cancer-enriched score

def validate_file(path):
    if not os.path.exists(path):
        raise argparse.ArgumentTypeError(f"File not exists: {path}")
    return path

def ensure_output_dir(outdir: str) -> None:
    
    if not os.path.exists(outdir):
        os.makedirs(outdir)

def get_script_dir():
    """locate the path of main script"""
    return os.path.dirname(os.path.abspath(__file__))



def main():
    # 
    parser = argparse.ArgumentParser(description="This script process TCR clonotype file (generated by MIXCR and format converted by vdjtools), and calculate multiple TCR features and output the table format.",
                                     usage="python %(prog)s [options] <-i inputFile | -m metadata> <-o outdir>",
                                     epilog="""
    Example: 
    Single Sample: python %(prog)s -i Sample1_clonotype.txt -o workdir 
    Multiple Samples: python %(prog)s -m metadata.txt -o workdir
                                     """,
                                     formatter_class=argparse.RawDescriptionHelpFormatter
                                     )
        # Create a mutually exclusive parameter group (- i and - m cannot be used simultaneously)
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument(
        '-i', '--infile',
        type=validate_file,
        help="Input tcr clonotype file which conforms vdjtools format."
    )
    group.add_argument(
        '-m', '--metadata',
        type=validate_file,
        help="Input metadata which contains multiple clonotype files path and id (colnames as sample_id and file_path)."
    )
    
    parser.add_argument('-o','--outdir', default="./Result", help="Output directory.")
    args = parser.parse_args()


    #
    if args.infile:
        # single pipeline
        print(f"Input file: {args.infile}")
        result=analyze_tcr_data(args.infile,is_metadata=False,method="score")  #

    elif args.metadata:
        # multiple samples pipeline
        print(f"Input metadata: {args.metadata}")
        result=analyze_tcr_data(args.metadata,is_metadata=True,method="score")   #  
    

    ensure_output_dir(args.outdir)

    result.v_gene.to_csv(os.path.join(args.outdir, 'v_freq.csv'),index=False)
    result.j_gene.to_csv(os.path.join(args.outdir, 'j_freq.csv'),index=False)
    result.aa_composition.to_csv(os.path.join(args.outdir, 'aa_composition.csv'), index=False)
    result.length_group.to_csv(os.path.join(args.outdir, 'length_group.csv'), index=False)  # 
    result.freq_group.to_csv(os.path.join(args.outdir, 'freq_group.csv'), index=False)  # 
    result.convergence.to_csv(os.path.join(args.outdir, 'convergence.csv'), index=False) 
    result.enriched_score.to_csv(os.path.join(args.outdir, 'enriched_score.csv'), index=False)  # 

####


def read_single_file(file_path: str, sample_id: str = None) -> pd.DataFrame:
    """read single file and use name as id"""
    if sample_id is None:
        
        sample_id = os.path.splitext(os.path.basename(file_path))[0]
    
    df = pd.read_csv(file_path, sep='\t')
    df['Sample_ID'] = sample_id
    return df

def read_metadata(metadata_path: str) -> List[Tuple[str, str]]:
    """read metadata file"""
    meta_df = pd.read_csv(metadata_path, sep='\t')
    return [(row['sample_id'], row['file_path']) for _, row in meta_df.iterrows()]


def analyze_tcr_data(input_arg: str, is_metadata: bool, method: str = 'raw') -> AnalysisResult:
   
    # input
    if is_metadata:
        # 
        samples = read_metadata(input_arg)
        dfs = [read_single_file(file_path, sample_id) for sample_id, file_path in samples]
    else:
        # 
        dfs = [read_single_file(input_arg)]
    
    # combine
    combined_df = pd.concat(dfs)
    

    # load enriched tcr db
    db_dict = load_enriched_databases()

    # calculate tcr features
    vj_result = get_vj_from_input_data(combined_df, method)
    aa_result = calculate_aa_composition(combined_df)
    length_result = calculate_length_group(combined_df)  # 
    freq_group_result = calculate_freq_groups(combined_df)  # 
    convergence_result = calculate_convergence(combined_df)  # 
    enriched_result = calculate_enriched_scores(combined_df, db_dict)  # 
    
    
    return AnalysisResult(
        v_gene=vj_result['v'],
        j_gene=vj_result['j'],
        aa_composition=aa_result,
        length_group=length_result,  
        freq_group=freq_group_result,  
        convergence=convergence_result,
        enriched_score=enriched_result  
    )


###########################################################
#V/J gene usages and score calculate.
def apply_log_transform(input_tbl: pd.DataFrame) -> pd.DataFrame:
    """logarithmic conversion on count data"""
    input_tbl = input_tbl.copy()
    input_tbl['count'] = input_tbl['count'].map(np.log)  # 
    total_count = input_tbl['count'].sum()
    input_tbl['freq'] = input_tbl['count'] / total_count
    return input_tbl        

def get_vj_from_input_data(input_data: pd.DataFrame, method: str = 'raw') -> Dict[str, pd.DataFrame]:
    # group by sample
    grouped = input_data.groupby('Sample_ID')
    v_results = []
    j_results = []
    
    for sample_id, sample_df in grouped:
        # 
        if method == 'score':
            sample_df = apply_log_transform(sample_df.copy())
    
    
        # v
        v_freq = sample_df.groupby('v')['freq'].sum().reset_index()
        v_freq['Sample_ID'] = sample_id
        v_results.append(v_freq)
        
        # j
        j_freq = sample_df.groupby('j')['freq'].sum().reset_index()
        j_freq['Sample_ID'] = sample_id
        j_results.append(j_freq)
    
    # combine
    v_gene = pd.concat(v_results).pivot(index='Sample_ID', columns='v', values='freq').fillna(0)
    j_gene = pd.concat(j_results).pivot(index='Sample_ID', columns='j', values='freq').fillna(0)

    v_gene = v_gene.reset_index().rename_axis(None, axis=1)
    j_gene = j_gene.reset_index().rename_axis(None, axis=1)
    
    return {'v': v_gene, 'j': j_gene}

###############################
#aa composition analysis.
#default value
LEN_LOW = 10
LEN_UP = 24
TRIM_LENGTH = 3
AA_LIST = list('ARNDCEQGHILKMFPSTWYV')
DEFAULT_COUNTS_THRESHOLD = 10

def trim_seq(seq: str) -> str:
    """Trim first and last amino acid (C and F)"""
    return seq[1:-1] if len(seq) > 2 else seq

def filter_tbl(input_table: pd.DataFrame, 
              count_threshold: int = DEFAULT_COUNTS_THRESHOLD,
              len_low: int = LEN_LOW,
              len_up: int = LEN_UP) -> pd.DataFrame:
    
    filtered = input_table.copy()
    
    # filter by minimum detectable count 
    filtered = filtered[filtered['count'] >= count_threshold]
    
    # calculate tcr cdr3aa length
    filtered['cdr3_len'] = filtered['cdr3aa'].str.len()
    filtered = filtered[(filtered['cdr3_len'] >= len_low) & 
                       (filtered['cdr3_len'] <= len_up)]
    
    # advanced filter
    filtered = filtered[
        # Without termination or exception characters
        ~filtered['cdr3aa'].str.contains(r'[_*]', regex=True) &
        # start with C
        filtered['cdr3aa'].str.startswith('C') &
        # end with FFF
        filtered['cdr3aa'].str.endswith('F') &
        # Only contains standard amino acid characters
        filtered['cdr3aa'].str.match(f'^[{"".join(AA_LIST)}]+$')
    ]
    # Sequence after trim (for alignment)
    filtered['cdr3aa_trimmed'] = filtered['cdr3aa'].map(trim_seq)

    # frequency and log score
    total_count = filtered['count'].sum()
    filtered['cdr3aa_freq'] = filtered['count'] / total_count if total_count > 0 else 0
    filtered['score'] = filtered['count'].map(np.log)
    
    return filtered

def calculate_aa_composition(input_tbl: pd.DataFrame) -> pd.DataFrame:
    """Calculate the frequency of amino acid composition"""
    aa_results = []
    
    for sample_id, sample_df in input_tbl.groupby('Sample_ID'):
        # filter clonotypes
        filtered_df = filter_tbl(sample_df)
        
        # Calculate amino acid frequency
        aa_dict = Counter()
        for _, row in filtered_df.iterrows():
            seq, val = row['cdr3aa'], row['score']
            for aa in seq[TRIM_LENGTH:-TRIM_LENGTH]:  # Remove 3 amino acids from each end
                aa_dict[aa] += val 
        
        # normalization
        total = sum(aa_dict.values())
        if total > 0:
            aa_dict = {k: v/total for k, v in aa_dict.items()}
        
        # result data
        result_row = {'Sample_ID': sample_id}
        result_row.update({f"{aa}_score": aa_dict.get(aa, 0) for aa in AA_LIST})
        aa_results.append(result_row)
    
    output_df = pd.DataFrame(aa_results)
    # Supplement possible missing amino acid columns
    for aa in AA_LIST:
        col_name = f"{aa}_score"
        if col_name not in output_df.columns:
            output_df[col_name] = 0.0



    
    return output_df

###############################
#TCR length group.
#default value
LENGTH_GROUPS = {
    'A': (0, 10),
    'B': (11, 13),
    'C': (14, 16),
    'D': (17, 19),
    'E': (20, 24)
}
LENGTH_GROUP_LST = list(LENGTH_GROUPS.keys())

def calculate_length_group(input_tbl: pd.DataFrame) -> pd.DataFrame:
    """Calculate TCR length grouping frequency"""
    length_results = []
    
    for sample_id, sample_df in input_tbl.groupby('Sample_ID'):
        # Record the number of original clonotypes
        raw_records = len(sample_df)
        # filter
        filtered_df = filter_tbl(sample_df)
        filtered_records = len(filtered_df)
        
        # Calculate TCR length group
        filtered_df['length_group'] = pd.cut(
            filtered_df['cdr3aa'].str.len(),
            bins=[0, 10, 13, 16, 19, 24],
            labels=LENGTH_GROUP_LST,
            right=True
        )
        
        # Calculate grouping frequency
        group_freq = filtered_df.groupby('length_group')['score'].sum()
        total_score = group_freq.sum()
        
        #
        result_row = {
            'Sample_ID': sample_id,
            'raw_records': raw_records,
            'filtered_records': filtered_records
        }
        
        if total_score > 0:
            result_row.update({
                f"{group}_len": group_freq.get(group, 0) / total_score
                for group in LENGTH_GROUP_LST
            })
        else:
            result_row.update({f"{group}_len": 0 for group in LENGTH_GROUP_LST})
        
        length_results.append(result_row)
    
    # 
    output_df = pd.DataFrame(length_results)
    for group in LENGTH_GROUP_LST:
        col_name = f"{group}_len"
        if col_name not in output_df.columns:
            output_df[col_name] = 0.0
    
    # Fixed column order
    col_order = ['Sample_ID', 'raw_records', 'filtered_records'] + [f"{g}_len" for g in LENGTH_GROUP_LST]
    return output_df[col_order]

###############################
#TCR convergence.
def calculate_convergence(input_tbl: pd.DataFrame) -> pd.DataFrame:
    """Calculate TCR convergence index"""
    convergence_results = []
    
    for sample_id, sample_df in input_tbl.groupby('Sample_ID'):
        # filter
        filtered_df = filter_tbl(sample_df)
        
        # Calculate TCR convergence
        if len(filtered_df) == 0:
            result_row = {
                'Sample_ID': sample_id,
                'Convergence_type': 0,
                'Convergence_freq': 0,
                'Convergence_score': 0
            }
        else:
            # 
            filtered_df['v_cdr3aa'] = filtered_df['v'] + ';' + filtered_df['cdr3aa']
            
            # repeated combinations of v+cdr3aa
            v_cdr3aa_counts = Counter(filtered_df['v_cdr3aa'])
            convergence_set = {k for k, v in v_cdr3aa_counts.items() if v > 1}
            convergence_df = filtered_df[filtered_df['v_cdr3aa'].isin(convergence_set)]
            
            # calculate
            total_type = len(filtered_df)
            total_count = filtered_df['count'].sum()
            total_score = filtered_df['score'].sum()
            
            conv_type = len(convergence_df)
            conv_count = convergence_df['count'].sum()
            conv_score = convergence_df['score'].sum()
            
            result_row = {
                'Sample_ID': sample_id,
                'Convergence_type': conv_type / total_type if total_type > 0 else 0,
                'Convergence_freq': conv_count / total_count if total_count > 0 else 0,
                'Convergence_score': conv_score / total_score if total_score > 0 else 0
            }
        
        convergence_results.append(result_row)
    
    return pd.DataFrame(convergence_results)

########################
#TCR clone class by frequency group.
#default value
FREQ_GROUPS = {
    'Hyperexpanded': (0.01, 1.0),      # >1%
    'Large': (0.0001, 0.01),           # 0.01%~1%
    'Medium': (0.00001, 0.0001),       # 0.001%~0.01%
    'Small': (0.000001, 0.00001),      # 0.0001%~0.001%
    'Rare': (0, 0.000001)              # <0.0001%
}
FREQ_GROUP_LST = list(FREQ_GROUPS.keys())

def calculate_freq_groups(input_tbl: pd.DataFrame) -> pd.DataFrame:
    """Calculate cloning frequency grouping"""
    freq_results = []
    
    for sample_id, sample_df in input_tbl.groupby('Sample_ID'):
        #filter
        filtered_df = filter_tbl(sample_df)
        total_count = filtered_df['count'].sum()
        
        #calculate freq
        filtered_df['freq'] = filtered_df['count'] / total_count
        
        # result
        result_row = {
            'Sample_ID': sample_id,
            'raw_records': len(sample_df),
            'filtered_records': len(filtered_df)
        }
        
        # result for each group
        for group, (lower, upper) in FREQ_GROUPS.items():
            mask = (filtered_df['freq'] > lower) & (filtered_df['freq'] <= upper)
            group_df = filtered_df[mask]
            
            # clone types and total frequency
            result_row[f'{group}_records'] = len(group_df)
            result_row[f'{group}_freq'] = group_df['count'].sum() / total_count if total_count > 0 else 0
        
        freq_results.append(result_row)
    
    # 
    output_df = pd.DataFrame(freq_results)
    for group in FREQ_GROUP_LST:
        for suffix in ['_records', '_freq']:
            col_name = f"{group}{suffix}"
            if col_name not in output_df.columns:
                output_df[col_name] = 0 if suffix == '_records' else 0.0
    
    # Fixed column order
    base_cols = ['Sample_ID', 'raw_records', 'filtered_records']
    metric_cols = [f"{g}_{m}" for g in FREQ_GROUP_LST for m in ['records', 'freq']]
    return output_df[base_cols + metric_cols]

########################
#TCR align against reference sequences and calculate enriched score.
#default value

SIMILARITY_THRESHOLD = 4.5
EDIT_DISTANCE_THRESHOLD = 0.3



def load_enriched_databases() -> Dict[str, pd.DataFrame]:
    """Load all enriched tcr db"""
    script_dir = get_script_dir()
    data_dir = os.path.join(script_dir, 'data')
    
    ENRICHED_DB_PATHS = {
        'KRAS': os.path.join(data_dir, 'filtered.enriched.kras_gdna.csv'),
        'EGFR': os.path.join(data_dir, 'filtered.enriched.egfr_gdna.csv'),
        'REF': os.path.join(data_dir, 'filtered.healthy_gdna_150.csv'),
        'LUNG_CANCER_BLOOD': os.path.join(data_dir, 'filtered.enriched.lung_cancer_gdna.csv'),
        'LUNG_CANCER_TISSUE': os.path.join(data_dir, 'filtered.enriched.lung_cancer_tissue.csv')
    }

    db_dict = {}
    for name, path in ENRICHED_DB_PATHS.items():
        if not os.path.exists(path):
            raise FileNotFoundError(f"Cancer-enriched TCR db not found: {path}")
        db = pd.read_csv(path)
        db = db.copy()

        # trim processing on tcr sequences
        db['cdr3aa_trimmed'] = db['cdr3aa'].map(trim_seq)
        

        db_dict[name] = db
    return db_dict

def calculate_enriched_scores(input_tbl: pd.DataFrame, db_dict: Dict[str, pd.DataFrame]) -> pd.DataFrame:
    """Calculate tcr sequence enrichment score"""
    enriched_results = []
    
    for sample_id, sample_df in input_tbl.groupby('Sample_ID'):
        # filter
        filtered_df = filter_tbl(sample_df)
        sample_records = len(filtered_df)
        
        result_row = {'Sample_ID': sample_id}
        
        # load db
        for db_name, db_df in db_dict.items():
            # Compare only sequences of the same length
            merged = pd.merge(
                filtered_df[['cdr3aa_trimmed', 'cdr3_len', 'score']], 
                db_df,
                on='cdr3_len',
                suffixes=('_query', '_db')
            )
            
            # edit distance filtering
            merged['edit_dist'] = merged.apply(
                lambda x: Levenshtein.distance(
                    x['cdr3aa_trimmed_query'], 
                    x['cdr3aa_trimmed_db']
                ) / x['cdr3_len'],
                axis=1
            )
            merged = merged[merged['edit_dist'] <= EDIT_DISTANCE_THRESHOLD]
            
            # similarity score filtering
            merged['similarity'] = merged.apply(
                lambda x: parasail.nw(
                    x['cdr3aa_trimmed_query'],
                    x['cdr3aa_trimmed_db'],
                    10, 1, parasail.blosum62
                ).score / x['cdr3_len'],
                axis=1
            )
            merged = merged[merged['similarity'] >= SIMILARITY_THRESHOLD]
            
            # final enrichment score
            if len(merged) > 0:
                total_score = (merged['similarity'] * 
                             merged['score_query'] * 
                             merged['score_db']).sum()
                adj_score = total_score / sample_records
                adj_fraction = len(merged) / sample_records
            else:
                adj_score = 0
                adj_fraction = 0
                
            result_row[f'{db_name}_enriched_fraction'] = adj_fraction
            result_row[f'{db_name}_enriched_score'] = adj_score
        
        enriched_results.append(result_row)
    
    return pd.DataFrame(enriched_results)


###########################################################

# the main function
 
if __name__ == "__main__":
    main()