calc_DockQ.py

import sys
from tqdm import tqdm
from util import compute_structure_similarity
from util import compute_DockQ_score

def calculate_dockq_score(ref_pdb, target_pdb):
    iRMS, LRMS, DockQ = compute_DockQ_score(target_pdb, ref_pdb)
    return iRMS, LRMS, DockQ


if len(sys.argv) < 2:
    print("Usage: python script.py <pdb_id>")
    sys.exit(1)

pdb_id = sys.argv[1].upper()

# Get CSP rank score file path
csp_rank_score_file = f'./CSP_Rank_Scores/CSP_{pdb_id.lower()}_CSpred.csv'

import os
from os import listdir
from os.path import isfile, join

# Get path to PDB files directory
pdb_files_dir = f'./PDB_FILES/{pdb_id.upper()}_aligned/'

# Check if directory exists
if not os.path.exists(pdb_files_dir):
    print(f"Error: Directory {pdb_files_dir} does not exist")
    sys.exit(1)

# Get list of all files in directory
aligned_files = [f for f in listdir(pdb_files_dir) if isfile(join(pdb_files_dir, f))]

# Get path to experimental reference structure
ref_file = f'./PDB_FILES/experimental_structures/exp_{pdb_id.lower()}.pdb'

# Check if reference file exists
if not os.path.exists(ref_file):
    print(f"Error: Reference file {ref_file} does not exist")
    sys.exit(1)


# Initialize arrays to store scores
irms_scores = []
lrms_scores = []
dockq_scores = []

# Loop through each file
for aligned_file in tqdm(aligned_files):
    aligned_file_path = os.path.join(pdb_files_dir, aligned_file)
    irms, lrms, dockq = calculate_dockq_score(ref_file, aligned_file_path)
    irms_scores.append(irms)
    lrms_scores.append(lrms)
    dockq_scores.append(dockq)
    print(f"Processing {aligned_file}")

# Print scores
print(f"iRMS scores: {irms_scores}")
print(f"LRMS scores: {lrms_scores}")
print(f"DockQ scores: {dockq_scores}")

import pandas as pd

# Read the CSP rank scores file
try:
    df = pd.read_csv(csp_rank_score_file)
except Exception as e:
    print(f"Error reading CSP rank scores file: {e}")
    sys.exit(1)

# Create dictionaries mapping filenames to scores
irms_dict = {}
lrms_dict = {}
dockq_dict = {}
for irms, lrms, dockq, aligned_file in zip(irms_scores, lrms_scores, dockq_scores, aligned_files):
    # Extract just the filename without path
    filename = os.path.basename(aligned_file)
    irms_dict[filename] = irms
    lrms_dict[filename] = lrms
    dockq_dict[filename] = dockq

# Add scores to dataframe by matching filenames
df['irms_score'] = df['holo_model_path'].apply(lambda x: irms_dict.get(os.path.basename(x), None))
df['lrms_score'] = df['holo_model_path'].apply(lambda x: lrms_dict.get(os.path.basename(x), None))
df['dockq_score'] = df['holo_model_path'].apply(lambda x: dockq_dict.get(os.path.basename(x), None))

# Save updated dataframe
try:
    df.to_csv(csp_rank_score_file, index=False)
    print(f"Successfully added iRMS, LRMS and DockQ scores to {csp_rank_score_file}")
except Exception as e:
    print(f"Error saving updated CSP rank scores file: {e}")


# Create scatter plot of consensus vs dockq scores
import matplotlib.pyplot as plt

plt.figure(figsize=(10, 6))
plt.scatter(df['consensus'], df['dockq_score'], alpha=0.5)
plt.xlabel('Consensus Score')
plt.ylabel('DockQ Score') 
plt.title('Consensus Score vs DockQ Score')
plt.grid(True, linestyle='--', alpha=0.7)

# Save the plot
plt.savefig(f'consensus_dockq_scatter_{pdb_id.lower()}.png')
plt.close()