Main.py

import os
from concurrent.futures import ProcessPoolExecutor
from tqdm import tqdm
import Checksum
import multiprocessing
import Rid_the_water
import Gromacs_gmx_bash_imports
import RMSD_calculator
import RMSF_calculator
import RMSD_visualizer
import RMSF_visualizer
import Radius_of_gyration_calculator
import dssp_to_index_file

# Lock for file writing
file_write_lock = multiprocessing.Lock()

def process_water(file_path):
    try:
        print("Processing:", file_path)
        Rid_the_water.remove_solvent_from_xtc(file_path)
        checksum = Checksum.compute_checksum_CRC32(file_path)
        with file_write_lock:
            Checksum.store_checksum(file_path, checksum, "process_water")
        print("Processed:", file_path)
    except Exception as e:
        print(f"Error processing file {file_path}: {e}")

def whole_traj(xtc_noWater):
    try:
        print("Processing:", xtc_noWater)
        Gromacs_gmx_bash_imports.PBC_unwrap(xtc_noWater)
        Gromacs_gmx_bash_imports.thousand_frames(xtc_noWater)
        checksum = Checksum.compute_checksum_CRC32(xtc_noWater)
        with file_write_lock:
            Checksum.store_checksum(xtc_noWater, checksum, "whole_traj")
        print("Processed:", xtc_noWater)
    except Exception as e:
        print(f"Error processing file {xtc_noWater}: {e}")

def process_RMSD(xtc_PBCunwrapped):
    try:
        print("Processing:", xtc_PBCunwrapped)
        RMSD_calculator.calculate_RMSD(xtc_PBCunwrapped)
        checksum = Checksum.compute_checksum_CRC32(xtc_PBCunwrapped)
        with file_write_lock:
            Checksum.store_checksum(xtc_PBCunwrapped, checksum, "process_RMSD")
        print("Processed:", xtc_PBCunwrapped)
    except Exception as e:
        print(f"Error processing file {xtc_PBCunwrapped}: {e}")

def process_RMSF(xtc_PBCunwrapped):
    try:
        print("Processing:", xtc_PBCunwrapped)
        RMSF_calculator.calculate_RMSF(xtc_PBCunwrapped)
        checksum = Checksum.compute_checksum_CRC32(xtc_PBCunwrapped)
        with file_write_lock:
            Checksum.store_checksum(xtc_PBCunwrapped, checksum, "process_RMSF")
        print("Processed:", xtc_PBCunwrapped)
    except Exception as e:
        print(f"Error processing file {xtc_PBCunwrapped}: {e}")

def process_rad_gyration(xtc_PBCunwrapped):
    try:
        print("Processing:", xtc_PBCunwrapped)
        Radius_of_gyration_calculator.comp_and_plot_rad_of_gyration(xtc_PBCunwrapped,
                                                                    destination_root_path="/ceph/users/akv16273/34_Elcock_13abundant_inEColi/21_Analysis_Datastructure")
        checksum = Checksum.compute_checksum_CRC32(xtc_PBCunwrapped)
        with file_write_lock:
            Checksum.store_checksum(xtc_PBCunwrapped, checksum, "process_rad_gyration")
        print("Processed:", xtc_PBCunwrapped)
    except Exception as e:
        print(f"Error processing file {xtc_PBCunwrapped}: {e}")

def traverse_directory_tree_parallel(root_path, remove_solvent=False, calculate_RMSD=False, calculate_RMSF=False, radius_gyration=False):
    cpu_avail = (os.cpu_count() or 1)
    print(f"Maximum available cpu used: {cpu_avail}")

    ## Removing Solvent to _noWater.xtc
    if remove_solvent:
        file_list = []
        for root, dirs, files in os.walk(root_path):
            for file in files:
                file_path = os.path.join(root, file)
                if file_path.endswith(".xtc") and ("_noWater.xtc" not in file_path) and ("_PBCunwrapped.xtc" not in file_path):
                    file_list.append(file_path)
        with (ProcessPoolExecutor(max_workers=cpu_avail) as executor, tqdm(total=len(file_list), desc="Remove Solvent", colour='yellow') as pbar):
            futures = []
            for file_path in file_list:
                if not Checksum.checksum_exists("process_water", Checksum.compute_checksum_CRC32(file_path)):
                    # Process the file only if it hasn't been processed or has changed
                    future = executor.submit(process_water, file_path)
                    future.add_done_callback(lambda p: pbar.update(1))
                    futures.append(future)
            for future in futures:
                future.result()
        print("Solvent removing complete.")

    # make trajectory whole and generate pdb from tpr with Gromacs
    if calculate_RMSD or calculate_RMSF or radius_gyration:
        noWater_list = []
        for root, dirs, files in os.walk(root_path):
            for file in files:
                file_path = os.path.join(root, file)
                if file_path.endswith("_noWater.xtc"):
                    noWater_list.append(file_path)
        with ProcessPoolExecutor(max_workers=cpu_avail) as executor, tqdm(total=len(noWater_list), desc="Making trajectory whole", colour='yellow') as pbar:
            futures = []
            for file_path in noWater_list:
                if not Checksum.checksum_exists("whole_traj", Checksum.compute_checksum_CRC32(file_path)):
                    # Process the file only if it hasn't been processed or has changed
                    future = executor.submit(whole_traj, file_path)
                    future.add_done_callback(lambda p: pbar.update(1))
                    futures.append(future)
            for future in futures:
                future.result()
        print("Make trajectories whole complete.")

    ## Calculate RMSD
    if calculate_RMSD:
        PBCunwrapped_list = []
        for root, dirs, files in os.walk(root_path):
            for file in files:
                file_path = os.path.join(root, file)
                if file_path.endswith("_noWater_PBCunwrapped.xtc"):
                    PBCunwrapped_list.append(os.path.abspath(file_path))
        with ProcessPoolExecutor(max_workers=cpu_avail) as executor, tqdm(total=len(PBCunwrapped_list), desc="Calculating RMSD", colour='yellow') as pbar:
            futures = []
            for file_path in PBCunwrapped_list:
                if not Checksum.checksum_exists("process_RMSD", Checksum.compute_checksum_CRC32(file_path)):
                    # Process the file only if it hasn't been processed or has changed
                    future = executor.submit(process_RMSD, file_path)
                    future.add_done_callback(lambda p: pbar.update(1))
                    futures.append(future)
            for future in futures:
                future.result()
        print("Computing RMSD complete.")

    ## Calculate RMSF
    if calculate_RMSF:
        PBCunwrapped_list2 = []
        for root, dirs, files in os.walk(root_path):
            for file in files:
                file_path = os.path.join(root, file)
                if file_path.endswith("_noWater_PBCunwrapped.xtc"):
                    PBCunwrapped_list2.append(os.path.abspath(file_path))
        # Use only one worker, otherwise the memory load of trajectory in RMSF align breaks (not in memory doesn't work either)
        with ProcessPoolExecutor(max_workers=2) as executor, tqdm(total=len(PBCunwrapped_list2), desc="Calculating RMSF", colour='yellow') as pbar:
            futures = []
            for file_path in PBCunwrapped_list2:
                if not Checksum.checksum_exists("process_RMSF", Checksum.compute_checksum_CRC32(file_path)):
                    # Process the file only if it hasn't been processed or has changed
                    future = executor.submit(process_RMSF, file_path)
                    future.add_done_callback(lambda p: pbar.update(1))
                    futures.append(future)
            for future in futures:
                future.result()
        print("Computing RMSF complete.")

    ## Calculate radius of gyrations and save the plots
    if radius_gyration:
        PBCunwrapped_list3 = []
        for root, dirs, files in os.walk(root_path):
            for file in files:
                file_path = os.path.join(root, file)
                if file_path.endswith("_noWater_PBCunwrapped.xtc"):
                    PBCunwrapped_list3.append(os.path.abspath(file_path))
        with ProcessPoolExecutor(max_workers=cpu_avail) as executor, tqdm(total=len(PBCunwrapped_list3), desc="Calculating Radius of Gyration", colour='yellow') as pbar:
            futures = []
            for file_path in PBCunwrapped_list3:
                if not Checksum.checksum_exists("process_rad_gyration", Checksum.compute_checksum_CRC32(file_path)):
                    # Process the file only if it hasn't been processed or has changed
                    future = executor.submit(process_rad_gyration, file_path)
                    future.add_done_callback(lambda p: pbar.update(1))
                    futures.append(future)
            for future in futures:
                future.result()
        print("Computing & Plotting Radius of Gyration complete.")

# Example usage
if __name__ == "__main__":
    traverse_directory_tree_parallel(root_path="/ceph/users/akv16273/34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/26_RbsB_only/1urp_actual_SIRAH/", remove_solvent=True, calculate_RMSD=True, calculate_RMSF=True,
                                     radius_gyration=True)

   # RMSD_visualizer.plt_dat_files(root_path="34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/9_Martini3/4_RbsB_elastic_net_edit/",
   #                               destination_root_path="34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/9_Martini3/4_RbsB_elastic_net_edit/",
   #                               results_folder="rmsd_results", all_takes_rmsd=True, rmsd_minus_250ns=True)   #no checksum so far
   #
   # RMSF_visualizer.plt_dat_files(root_path="34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/9_Martini3/4_RbsB_elastic_net_edit/",
   #                               destination_root_path="34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/9_Martini3/4_RbsB_elastic_net_edit/",
   #                               results_folder="rmsf_results")        #no checksum so far
   #
   # RMSF_calculator.calculate_b_factor(root_path="34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/9_Martini3/4_RbsB_elastic_net_edit/") #no checksum so far