add_aligned_suffix_csv.py

import pandas as pd
import sys
import re
from os.path import exists, isdir
from paths import PDB_FILES, CSP_Rank_Scores

pdb_id = sys.argv[1].upper()
# Read the CSV file
dfs = [CSP_Rank_Scores + 'CSP_'+pdb_id.lower()+'_CSpred.csv']

for csv_file in dfs:
    if exists(csv_file):
        df = pd.read_csv(csv_file)
        
        aligned_directory = f'{PDB_FILES}{pdb_id}_aligned/'
        alt_aligned_directory = f'{PDB_FILES}{pdb_id}_alt_aligned/'

        if exists(aligned_directory) and isdir(aligned_directory):
            #df['holo_model_path'] = df['holo_model_path'].apply(lambda x: x.replace(r'^\./' + pdb_id + '/', './' + pdb_id + '_aligned/', 1) if '_aligned' not in x else x)
            #df['holo_model_path'] = df['holo_model_path'].apply(lambda x: re.sub(r'^\./' + pdb_id + '/', './' + pdb_id + '_aligned/', 1) if '_aligned' not in x and '_alt' not in x else x )
            df['holo_model_path'] = df['holo_model_path'].apply(
                lambda x: re.sub('/' + pdb_id + '/', '/' + pdb_id + '_aligned/', x)
                if isinstance(x, str) and '_aligned' not in x and '_alt' not in x else x
            )
        if exists(alt_aligned_directory) and isdir(alt_aligned_directory):
            df['holo_model_path'] = df['holo_model_path'].apply(lambda x: re.sub('/' + pdb_id + '_alt/', '/' + pdb_id + '_alt_aligned/', str(x)) 
            if isinstance(x, str) and '_aligned' not in x and '_alt' in x else x
            )
        print(df['holo_model_path'])
        # Save the modified DataFrame back to a new CSV file
        df.to_csv(csv_file, index=False)