bioscripts/paring.py at main · zeqri/bioscripts · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
import random
import pickle
import os
from scipy.spatial import KDTree
import multiprocessing as mp
from functools import partial
import numpy as np
import pandas as pd


def process_single_file(dill_file, dill_data_path, save_path, distance_cutoff=7):
    """Process a single dill file to find contact chains"""
    try:
        data_path = os.path.join(dill_data_path, dill_file)
        find_contact_chains_linear(data_path, distance_cutoff, save_path)
    except Exception as e:
        print(f"Error processing {dill_file}: {str(e)}")

def find_contact_chains_linear(data_path_dill, distance_cutoff=7, save_path=None):
    # Load data
    with open(data_path_dill, 'rb') as f:
        data = pickle.load(f)

    # Filter RNA and protein atoms
    rna_residues = data[(data['type'] == 'rna') & (data['atom_name'] == 'P')].reset_index()
    protein_residues = data[(data['type'] == 'protein') & (data['atom_name'] == 'CA')].reset_index()

    # Build KD-tree for protein coordinates
    protein_coords = protein_residues[['x', 'y', 'z']].to_numpy()
    protein_tree = KDTree(protein_coords)

    # RNA coordinates
    rna_coords = rna_residues[['x', 'y', 'z']].to_numpy()

    # Find pairs within the distance cutoff
    indices = protein_tree.query_ball_point(rna_coords, r=distance_cutoff)

    # Collect close contacts
    close_contacts = []
    for rna_index, protein_indices in enumerate(indices):
        if protein_indices:
            # RNA chain and protein chains in contact
            rna_chain = rna_residues.loc[rna_index, 'chain']
            for protein_index in protein_indices:
                protein_chain = protein_residues.loc[protein_index, 'chain']

                # Append contact only if unique
                contact = (protein_chain, rna_chain)
                if contact not in close_contacts:
                    close_contacts.append(contact)

                    # Extract RNA-protein pair
                    rna_protein_pair = data[data['chain'].isin([protein_chain, rna_chain])]
                    rna_filtered = rna_protein_pair[(rna_protein_pair['type'] == 'rna') & (rna_protein_pair['atom_name'] == 'P')]
                    protein_filtered = rna_protein_pair[(rna_protein_pair['type'] == 'protein') & (rna_protein_pair['atom_name'] == 'CA')]

                    # Skip invalid lengths
                    if not (6 <= len(rna_filtered) <= 1022 and 6 <= len(protein_filtered) <= 1022):
                        continue

                    # Generate file name
                    pdb_id_name = data_path_dill.split('/')[-1].split('.')[0]
                    file_name = f"{pdb_id_name}_{protein_chain}_{rna_chain}.dill"

                    # Save RNA-protein pair
                    if save_path:
                        output_file = os.path.join(save_path, file_name)
                        with open(output_file, 'wb') as f:
                            pickle.dump(rna_protein_pair, f)


def find_contact_chains_linear_all(data_path_dill, distance_cutoff=7, save_path=None):

    #return all protein chains around a single RNA
    # Load data
    with open(data_path_dill, 'rb') as f:
        data = pickle.load(f)

    # Filter RNA and protein atoms
    rna_residues = data[(data['type'] == 'rna') & (data['atom_name'] == 'P')].reset_index()
    protein_residues = data[(data['type'] == 'protein') & (data['atom_name'] == 'CA')].reset_index()

    # Build KD-tree for protein coordinates
    protein_coords = protein_residues[['x', 'y', 'z']].to_numpy()
    protein_tree = KDTree(protein_coords)

    # RNA coordinates
    rna_coords = rna_residues[['x', 'y', 'z']].to_numpy()

    # Find pairs within the distance cutoff
    indices = protein_tree.query_ball_point(rna_coords, r=distance_cutoff)

    # Collect all unique RNA and protein chains in contact
    contacts = []
    for rna_index, protein_indices in enumerate(indices):
        if protein_indices:
            # RNA chain
            rna_chain = rna_residues.loc[rna_index, 'chain']

            # Gather all protein chains in contact with this RNA chain
            for protein_index in protein_indices:
                protein_chain = protein_residues.loc[protein_index, 'chain']

                # Record unique RNA-protein chain pairs
                contacts.append((rna_chain, protein_chain))

    # Deduplicate contacts and aggregate all proteins near each RNA chain
    rna_to_protein_map = {}
    for rna_chain, protein_chain in contacts:
        if rna_chain not in rna_to_protein_map:
            rna_to_protein_map[rna_chain] = []
        rna_to_protein_map[rna_chain].append(protein_chain)

    # Save RNA with all nearby proteins as one unit
    pdb_id_name = data_path_dill.split('/')[-1].split('.')[0]
    for rna_chain, protein_chains in rna_to_protein_map.items():
        rna_protein_pair = data[data['chain'].isin([rna_chain] + protein_chains)]

        # if save_path:
        #     file_name = f"{pdb_id_name}_RNA_{rna_chain}_all_nearby_proteins.dill"
        #     output_file = os.path.join(save_path, file_name)

        #     with open(output_file, 'wb') as f:
        #         pickle.dump(rna_protein_pair, f)
    return rna_protein_pair

def main():
    # Set your paths
    dill_data_path = "path to cached files in data frame format"
    save_path = "path to save the data"
    csv_path = "csv file containing PDB IDs"
    df=pd.read_csv(csv_path)
    dill_files = df['path'].to_list()

    # Ensure save directory exists
    os.makedirs(save_path, exist_ok=True)

    random.shuffle(dill_files)

    # Set up multiprocessing
    num_processes = mp.cpu_count() - 1  # Leave one CPU free
    print(f"Found {mp.cpu_count()} CPUs")
    print(f"Starting processing with {num_processes} processes")

    # Create partial function with fixed arguments
    process_file = partial(process_single_file,
                         dill_data_path=dill_data_path,
                         save_path=save_path,
                         distance_cutoff=7)

    # Create pool and process files
    with mp.Pool(processes=num_processes) as pool:
        pool.map(process_file, dill_files)


if __name__ == '__main__':
    main()