calc_rmsf.py

import numpy as np
from Bio import PDB
import pymol
from pymol import cmd, util
import os

def calculate_rmsf(pdb_filename):
    parser = PDB.PDBParser(QUIET=True)
    structure = parser.get_structure('structure', pdb_filename)

    model_list = list(structure)
    n_models = len(model_list)
    residue_positions = {}
    
    for model in model_list:
        for chain in model:
            for residue in chain:
                if PDB.is_aa(residue):
                    res_id = (chain.id, residue.id[1])
                    if res_id not in residue_positions:
                        residue_positions[res_id] = []
                    ca = residue['CA']
                    residue_positions[res_id].append(ca.coord)

    rmsf_values = {}
    for res_id, coords in residue_positions.items():
        coords = np.array(coords)
        mean_coord = np.mean(coords, axis=0)
        sq_diff = np.square(coords - mean_coord)
        mean_sq_diff = np.mean(sq_diff, axis=0)
        rmsf = np.sqrt(np.mean(mean_sq_diff))
        rmsf_values[res_id] = rmsf

    return rmsf_values

def load_and_color_pymol(pdb_filename, rmsf_values):
    # Initialize PyMOL
    pymol.finish_launching()

    # Load the PDB file
    object_name = os.path.basename(pdb_filename).split('.')[0]
    cmd.load(pdb_filename, object_name)
    cmd.remove('not alt A+')

    # Create a color gradient from blue (low RMSF) to red (high RMSF)
    min_rmsf = min(rmsf_values.values())
    max_rmsf = max(rmsf_values.values())

    for (chain, resi), rmsf in rmsf_values.items():
        color = cmd.get_color_tuple('blue')
        if max_rmsf > min_rmsf:  # Avoid division by zero
            scale = (rmsf - min_rmsf) / (max_rmsf - min_rmsf)
            red = scale
            blue = 1 - scale
            color = (red, 0, blue)

        # Create a color name
        color_name = f"color_{chain}_{resi}"
        cmd.set_color(color_name, color)

        # Apply the color to the residue
        cmd.color(color_name, f"/{object_name}//{chain}/{resi}")

    # Show the cartoon representation
    cmd.show("cartoon")
    cmd.cartoon("tube")

    # Center on the structure
    cmd.orient()


pdb_filename = '7jyn_max_NLDR_consensus_files_1.pdb'
pdb_filename = '2mnu_max_NLDR_consensus_files_1.pdb'
pdb_filename = '2kwv_max_NLDR_consensus_files_1.pdb'
pdb_filename = '7jq8_max_NLDR_consensus_files_1.pdb'
rmsf_values = calculate_rmsf(pdb_filename)
load_and_color_pymol(pdb_filename, rmsf_values)

# Printing the RMSF values for each residue
for res_id, rmsf in rmsf_values.items():
    print(f'Residue {res_id}: RMSF = {rmsf:.3f}')

pymol.cmd.viewport(1920, 1080)  # Set the viewport size to 1920x1080 for HD frames