Build_PPIGraph.py

import argparse
import os
import Data_Processing as dp
import Pipeline_Functions as pf

def main():
    parser = argparse.ArgumentParser(description="Build GNN PPI graphs from merged networks.")
    parser.add_argument("--esm2_embeddings_dir", default="./data/esm2_embeddings/", help="Folder containing ESM2 embeddings (.pkl)")
    parser.add_argument("--raw_folder", default="./data/merged_networks/", help="Folder containing merged networks (.dat)")
    parser.add_argument("--save_folder", default="./data/ppigraphs/", help="Folder to save constructed graphs")
    parser.add_argument("--percentile", type=int, default=70, help="Percentile threshold for network edges")
    parser.add_argument("--score_saver_percentile", type=int, default=99, help="Percentile threshold for score saver")
    parser.add_argument("--gene_dict_percentile", type=int, default=90, help="Percentile threshold for gene dictionary")
    parser.add_argument("--min_interactors", type=int, default=20, help="Minimum number of interactors per node")
    parser.add_argument("--max_degree", type=int, default=300, help="Maximum degree per node")
    args = parser.parse_args()

    # Create save folder if it doesn't exist
    if not os.path.exists(args.save_folder):
        print(f"Creating missing output directory: {args.save_folder}")
        os.makedirs(args.save_folder, exist_ok=True)

    print("Loading ESM2 Embeddings...")
    esm_dict = {}
    pkl_files = [f for f in os.listdir(args.esm2_embeddings_dir) if f.endswith('.pkl')]
    if not pkl_files:
        print(f"Warning: No ESM2 embedding .pkl files found in {args.esm2_embeddings_dir}")
    else:
        for file in pkl_files:
            address = os.path.join(args.esm2_embeddings_dir, file[:-4])
            print(f"  - Loading tokens from: {address}")
            sub_dict = dp.load_esm_token_dict(address, weighted=True)
            esm_dict.update(sub_dict)

    pf.build_ppigraphs(
        esm_dict=esm_dict,
        raw_folder=args.raw_folder,
        save_folder=args.save_folder,
        percentile=args.percentile,
        score_saver_percentile=args.score_saver_percentile,
        gene_dict_percentile=args.gene_dict_percentile,
        min_interactors=args.min_interactors,
        max_degree=args.max_degree
    )

if __name__ == '__main__':
    main()