MeMoModelling/gene-expression-bootstrapper

1import argparse2import ast3import libsbml4import os5import re6import numpy as np7import pandas as pd8import sys9from utils.utils import *10 11def extract_genes_from_gpa(gpa_str):12    """Extract all gene IDs from a gpaAssociation string (e.g. 'G_A or (G_B and G_C)')."""13    if pd.isna(gpa_str) or str(gpa_str).strip() == "":14        return []15    tokens = re.findall(r'[A-Za-z_]\w*', str(gpa_str))16    keywords = {'or', 'and', 'not'}17    return [t for t in tokens if t.lower() not in keywords]18 19def parse_genes(gene_value):20    """Parse either a stringified gene list or a boolean gene expression."""21    if pd.isna(gene_value) or str(gene_value).strip() == "":22        return []23 24    if isinstance(gene_value, (list, tuple, set)):25        return list(gene_value)26 27    try:28        parsed = ast.literal_eval(str(gene_value))29    except (ValueError, SyntaxError):30        return extract_genes_from_gpa(gene_value)31 32    if isinstance(parsed, (list, tuple, set)):33        return list(parsed)34    return extract_genes_from_gpa(parsed)35 36def read_combined_geneExpr(combined_geneExpr_filename):37    if not os.path.isfile(combined_geneExpr_filename):38        raise FileNotFoundError(f"Missing combined normalized count file, expected at {combined_geneExpr_filename}")39    combined_geneExpr_df = pd.read_csv(combined_geneExpr_filename, index_col=0)40    return combined_geneExpr_df41 42def is_real_gene_with_mapping(gene):43    return not gene.startswith(("unmapped", "unknown", "Spontaneous", "Exchange", "Sink", "Diffusion"))44 45def read_system_gene(model_pre_filename):46    rxn_df = read_model_excel(model_pre_filename, "Reactions")47 48    # Support both stringified lists and boolean expressions in 'genes'.49    if "genes" in rxn_df.columns:50        gene_col = "genes"51    elif "gpaAssociation" in rxn_df.columns:52        gene_col = "gpaAssociation"53    else:54        raise ValueError(55            "Reactions sheet must have a 'genes' or 'gpaAssociation' column — "56            f"columns found: {list(rxn_df.columns)}"57        )58 59    system_genes_dict = {}60    missing_gene_system_dict = {}61    all_genes_set = set()62    for system, gene_list_raw in zip(rxn_df["system"], rxn_df[gene_col]):63        gene_list = parse_genes(gene_list_raw)64 65        if system != "" and system not in system_genes_dict:66            system_genes_dict[system] = set()67 68        for gene in gene_list:69            if system != "" and is_real_gene_with_mapping(gene):70                system_genes_dict[system].add(gene)71            if gene.startswith(("unmapped", "unknown")):72                missing_gene_system_dict[gene] = system73            if is_real_gene_with_mapping(gene):74                all_genes_set.add(gene)75 76    missing_gene_system_dict = dict(sorted(missing_gene_system_dict.items()))77    return system_genes_dict, missing_gene_system_dict, all_genes_set78 79def map_genes(mapping_dict, system_genes_dict, missing_gene_system_dict, all_genes_set):80    """Map model gene tags to expression IDs and retain unmapped tags for bootstrapping."""81    mapped_system_genes = {}82    missing_genes = dict(missing_gene_system_dict)83 84    for system, gene_set in system_genes_dict.items():85        mapped_genes = set()86        for gene in gene_set:87            if gene in mapping_dict:88                mapped_genes.add(mapping_dict[gene])89            else:90                missing_genes.setdefault(gene, system)91        if mapped_genes:92            mapped_system_genes[system] = mapped_genes93 94    mapped_all_genes = {mapping_dict[gene] for gene in all_genes_set if gene in mapping_dict}95    return mapped_system_genes, dict(sorted(missing_genes.items())), mapped_all_genes96 97def filter_geneExpr_df(combined_geneExpr_df, all_genes_set):98    geneExpr_df = combined_geneExpr_df.loc[combined_geneExpr_df.index.isin(all_genes_set)]99    return geneExpr_df100 101def get_system_gene_counts(system_genes_dict, gene_count_dict):102    system_gene_counts_dict = {}103    missing_genes = set()104    for system, gene_set in system_genes_dict.items():105        gene_counts = []106        for gene in sorted(gene_set):107            if gene in gene_count_dict:108                gene_counts.append(gene_count_dict[gene])109            else:110                missing_genes.add(gene)111        system_gene_counts_dict[system] = gene_counts112    if missing_genes:113        print("Warning   : These genes do not have normalized count values - {}".format(", ".join(missing_genes)))114    return system_gene_counts_dict115 116def initialize_df(geneExpr_sample, columns):117    new_geneExpr_df_sample = pd.concat([geneExpr_sample.rename(col) for col in columns], axis=1)118    return new_geneExpr_df_sample119 120def bootstrap_missing_genes(missing_gene_system_dict, system_gene_counts_dict, columns):121    rng = np.random.default_rng(seed=0)122    sample_pool_all_genes = [gc for gene_counts in system_gene_counts_dict.values() for gc in gene_counts]123    new_geneExpr_df_sample_missing = pd.DataFrame(columns=columns)124    for missing_gene, system in missing_gene_system_dict.items():125        if system in system_gene_counts_dict:126            sample_pool = system_gene_counts_dict[system]127            new_geneExpr_df_sample_missing.loc[missing_gene] = rng.choice(sample_pool, len(columns))128        else:129            new_geneExpr_df_sample_missing.loc[missing_gene] = rng.choice(sample_pool_all_genes, len(columns))130    return new_geneExpr_df_sample_missing131 132def filter_by_batch(new_geneExpr_df, start, end, columns):133    target_columns = [f"{col}_{i}" for col in columns for i in range(start, end)]134    new_geneExpr_df_batch_df = new_geneExpr_df[target_columns]135    return new_geneExpr_df_batch_df136 137def bootstrap_genes(model_pre_filenames, mapping_filenames, species_prefixes, combined_geneExpr_filename, geneExpr_folder, batch_count=1000):138    print("Read models from", ", ".join(model_pre_filenames))139    print("Read mapping tables from", ", ".join(mapping_filenames))140    print("Species prefixes:", ", ".join(species_prefixes))141    print("Read gene normalized counts from", combined_geneExpr_filename)142    os.makedirs(geneExpr_folder, exist_ok=True)143 144    total_to_sample = batch_count145 146    # column: sample_name147    # row: real genes for all species148    combined_geneExpr_df = read_combined_geneExpr(combined_geneExpr_filename)149 150    # column: A_1, A_2, ..., B_1, B_2, ...151    # row: (real genes with mapping + unmapped real genes + unknown genes) + Exchange + Sink for all species152    new_geneExpr_df = pd.DataFrame()153    154    for model_pre_filename, mapping_filename, species in zip(model_pre_filenames, mapping_filenames, species_prefixes): # loop by species155        print(f"Bootstrapping genes for {species}...")156        # mapping_dict: {model_tag in the model file: gene_id in the geneExpr file}157        mapping_dict = read_mapping(mapping_filename)158        159        # system_genes_dict: {system: set of real genes with mapping with that system}160        # missing_gene_system_dict: {unmapped real gene & unknown gene: system of the gene}161        # all_genes_set: set of all the real genes with mapping used in the model of that species162        system_genes_dict, missing_gene_system_dict, all_genes_set = read_system_gene(model_pre_filename)163 164        # Map model tags to the annotation IDs used in the expression file.165        system_genes_dict, missing_gene_system_dict, all_genes_set = map_genes(166            mapping_dict, system_genes_dict, missing_gene_system_dict, all_genes_set167        )168 169        # filter to get only geneExprs for that species170        geneExpr_df = filter_geneExpr_df(combined_geneExpr_df, all_genes_set)171 172        # column: A_1, A_2, ..., B_1, B_2, ...173        # row: (real genes with mapping + unmapped real genes + unknown genes) for one species174        new_geneExpr_df_species = pd.DataFrame()175        for sample in geneExpr_df.columns:176            # gene_count_dict: {gene: gene count}177            gene_count_dict = geneExpr_df[sample].to_dict()178            # system_gene_counts_dict: {system: list of gene counts of the genes with that system}179            system_gene_counts_dict = get_system_gene_counts(system_genes_dict, gene_count_dict)180 181            # column: (one sample) A_1, A_2, ..182            columns = [f"{sample}_{i+1}" for i in range(total_to_sample)]183            # row: real genes with mapping (same as geneExpr file)184            new_geneExpr_df_sample = initialize_df(geneExpr_df[sample], columns)185            # row: unmapped real genes + unknown genes 186            # (randomly bootstrap from the gene counts of the same species and the same sample and the same system)187            # (if the system of the missing gene do not have gene counts, randomly bootstrap from the gene counts of the same species and the same sample)188            new_geneExpr_df_sample_missing = bootstrap_missing_genes(missing_gene_system_dict, system_gene_counts_dict, columns)189            # concat by rows: real genes with mapping + unmapped real genes + unknown genes190            new_geneExpr_df_sample = pd.concat([new_geneExpr_df_sample, new_geneExpr_df_sample_missing], axis=0)191 192            # concat by columns: A_1, A_2, ..., B_1, B_2, ...193            new_geneExpr_df_species = pd.concat([new_geneExpr_df_species, new_geneExpr_df_sample], axis=1)194            195        # concat by rows: species_1 + species_2, ...196        new_geneExpr_df = pd.concat([new_geneExpr_df, new_geneExpr_df_species], axis=0)197 198    new_geneExpr_df.loc["Spontaneous"] = 0199    new_geneExpr_df.loc["Exchange"] = 0200    new_geneExpr_df.loc["Sink"] = 0201    new_geneExpr_df.loc["Growth"] = 0202    new_geneExpr_df.loc["Diffusion"] = 0203 204    # separate into batch to different files205    print("Writing to files...")206    new_geneExpr_df_batch_filename = os.path.join(geneExpr_folder, "geneExpr")207    for i in range(batch_count):208        # first file: A_1, B_1, ..209        # second file: A_2, B_2, 210        new_geneExpr_df_batch_df = filter_by_batch(new_geneExpr_df, i+1, (i+1)+1, combined_geneExpr_df.columns)211        new_geneExpr_df_batch_df.to_csv(new_geneExpr_df_batch_filename + f"_{i+1}.csv")    212    print(f"Write to {new_geneExpr_df_batch_filename}_<1_{batch_count}>.csv")213 214if __name__ == "__main__":215    # define arguments216    parser = argparse.ArgumentParser()217    parser.add_argument("--model_pre_filenames", nargs="+", required=True, help="List of parsed model files from construct_modularized_community_model tool")218    parser.add_argument("--mapping_filenames", nargs="+", required=True, help="List of mapping files to map gene in model to gene in annotation, from identifiers_mapping tool")219    parser.add_argument("--species_prefixes", nargs="+", required=True, help="List of species prefixes for the models, in the same order as input files")220    parser.add_argument("--combined_geneExpr_filename", required=True, help="The gene expression values for all species and for all samples")221    parser.add_argument("--geneExpr_folder", required=True, help="Folder containing gene expression files (geneExpr_<1-batch_count>.csv) with bootstrapped values added for unmapped and unknown genes")222    parser.add_argument("--batch_count", type=int, default=1000, help="Number of batch to bootstrap, write each batch to a file")223    args = parser.parse_args()224 225    # read arguments226    model_pre_filenames = args.model_pre_filenames227    mapping_filenames = args.mapping_filenames228    species_prefixes = args.species_prefixes229    combined_geneExpr_filename = args.combined_geneExpr_filename230    geneExpr_folder = args.geneExpr_folder231    batch_count = args.batch_count232    233    bootstrap_genes(model_pre_filenames, mapping_filenames, species_prefixes, combined_geneExpr_filename, geneExpr_folder, batch_count)234