diff --git a/mdpath/mdpath.py b/mdpath/mdpath.py index 63c903b..d747407 100644 --- a/mdpath/mdpath.py +++ b/mdpath/mdpath.py @@ -16,12 +16,8 @@ import os import argparse -import pandas as pd import MDAnalysis as mda import json -import multiprocessing -from tqdm import tqdm -from collections import defaultdict import pickle from mdpath.src.structure import StructureCalculations, DihedralAngles @@ -33,6 +29,39 @@ from mdpath.src.confidence import EdgeConfidenceCalculator +def _select_subsystem(topology, trajectories, selection, prefix, label): + """Restrict topology and every replica trajectory to an atom selection. + + Writes ``selected_.pdb`` and ``selected__.dcd`` per replica. + + Returns: + tuple: (new_topology, new_trajectories, universe) for the selected subsystem. + """ + universe = mda.Universe(topology, trajectories[0]) + atoms = universe.select_atoms(selection) + if len(atoms) == 0: + raise ValueError(f"No atoms found for {selection}") + merged = mda.Merge(atoms) + out_topology = f"selected_{prefix}.pdb" + merged.atoms.write(out_topology) + + new_trajectories = [] + for i, traj_file in enumerate(trajectories): + u = mda.Universe(topology, traj_file) + sel = u.select_atoms(selection) + out_name = f"selected_{prefix}_{i}.dcd" + with mda.Writer(out_name, sel.n_atoms) as W: + for ts in u.trajectory: + merged.atoms.positions = sel.positions + W.write(merged.atoms) + new_trajectories.append(out_name) + + print( + f"{label} selected for {len(new_trajectories)} trajectory file(s) and will now be analyzed." + ) + return out_topology, new_trajectories, mda.Universe(out_topology, new_trajectories[0]) + + def main(): """Main function for running MDPath from the command line. It can be called using 'mdpath' after installation. @@ -164,14 +193,6 @@ def main(): default=False, ) - parser.add_argument( - "-water", - dest="water", - help="Allows for the tracking of stable waters medigating allosteric communication. Only include if the trajectory includes the water model.", - required=False, - default=False, - ) - args = parser.parse_args() if not args.topology or not args.trajectory: print("Both trajectory and topology files are required!") @@ -193,7 +214,6 @@ def main(): numpath = int(args.numpath) invert = bool(args.invert) spline = bool(args.spline) - water = float(args.water) confidence = bool(args.confidence) if confidence and len(trajectories) < 2: print( @@ -205,57 +225,15 @@ def main(): # Chain selection if args.chain: chain = str(args.chain) - chain_atoms = traj.select_atoms(f"chainID {chain}") - if len(chain_atoms) == 0: - raise ValueError(f"No atoms found for chain {chain}") - chain_universe = mda.Merge(chain_atoms) - # Write new topology - chain_universe.atoms.write("selected_chain.pdb") - - # Write trajectory for each replica - new_trajectories = [] - for i, traj_file in enumerate(trajectories): - u = mda.Universe(topology, traj_file) - chain_sel = u.select_atoms(f"chainID {chain}") - out_name = f"selected_chain_{i}.dcd" - with mda.Writer(out_name, chain_sel.n_atoms) as W: - for ts in u.trajectory: - chain_universe.atoms.positions = chain_sel.positions - W.write(chain_universe.atoms) - new_trajectories.append(out_name) - - topology = "selected_chain.pdb" - trajectories = new_trajectories - traj = mda.Universe(topology, trajectories[0]) - print( - f"Chain {chain} selected for {len(trajectories)} trajectory file(s) and will now be analyzed." + topology, trajectories, traj = _select_subsystem( + topology, trajectories, f"chainID {chain}", "chain", f"Chain {chain}" ) # Segment ID selection (CHARMM compatibility) if args.segid: segid = str(args.segid) - segid_atoms = traj.select_atoms(f"segid {segid}") - if len(segid_atoms) == 0: - raise ValueError(f"No atoms found for segid {segid}") - segid_universe = mda.Merge(segid_atoms) - segid_universe.atoms.write("selected_segid.pdb") - - new_trajectories = [] - for i, traj_file in enumerate(trajectories): - u = mda.Universe(topology, traj_file) - segid_sel = u.select_atoms(f"segid {segid}") - out_name = f"selected_segid_{i}.dcd" - with mda.Writer(out_name, segid_sel.n_atoms) as W: - for ts in u.trajectory: - segid_universe.atoms.positions = segid_sel.positions - W.write(segid_universe.atoms) - new_trajectories.append(out_name) - - topology = "selected_segid.pdb" - trajectories = new_trajectories - traj = mda.Universe(topology, trajectories[0]) - print( - f"Segment {segid} selected for {len(trajectories)} trajectory file(s) and will now be analyzed." + topology, trajectories, traj = _select_subsystem( + topology, trajectories, f"segid {segid}", "segid", f"Segment {segid}" ) # Write first frame PDB after all selections @@ -337,7 +315,9 @@ def main(): ) # Exports image of the Graph to PNG # Calculate paths - path_total_weights = graph_builder.collect_path_total_weights(df_distant_residues) + path_total_weights = graph_builder.collect_path_total_weights_parallel( + df_distant_residues, num_parallel_processes + ) sorted_paths = sorted(path_total_weights, key=lambda x: x[1], reverse=True) with open("output.txt", "w") as file: for path, total_weight in sorted_paths[:numpath]: @@ -425,19 +405,5 @@ def main(): MDPathVisualize.create_splines("quick_precomputed_clusters_paths.json") -# if water: -# from mdpath.src.water_tracing import WaterTracer -# water_tracer = WaterTracer( -# topology=topology, -# trajectory=trajectory, -# path_total_weights=path_total_weights, -# occurrence_threshold=float(water) / 100.0, -# ) -# water_tracer.pathway_water_data.save("pathway_water_data.pkl") -# water_tracer.pathway_water_data.to_dataframe().to_csv( -# "pathway_water_bridges.csv", index=False -# ) - - if __name__ == "__main__": main() diff --git a/mdpath/mdpath_tools.py b/mdpath/mdpath_tools.py index 1897973..058582c 100644 --- a/mdpath/mdpath_tools.py +++ b/mdpath/mdpath_tools.py @@ -7,8 +7,6 @@ import os import argparse import pandas as pd -import numpy as np -import MDAnalysis as mda import json from tqdm import tqdm from collections import defaultdict @@ -484,10 +482,6 @@ def domain_mi_analysis(): $ mdpath_domain_mi -graph -config -output """ - import pickle - import networkx as nx - import pandas as pd - parser = argparse.ArgumentParser( prog="mdpath_domain_mi", description="Analyze mutual information within and between protein domains.", @@ -550,93 +544,32 @@ def get_residues_in_domain(domain_row): for domain, residues in domains.items(): print(f" {domain}: {len(residues)} residues ({min(residues)}-{max(residues)})") - intra_domain_mi = {} - inter_domain_mi = {} - - for domain in domains.keys(): - intra_domain_mi[domain] = [] - print("\nAnalyzing graph edges...") - for edge in tqdm(Graph.edges(data=True), desc="Processing edges"): - node1, node2, edge_data = edge - - mi_value = edge_data.get('nmi', edge_data.get('mi', edge_data.get('weight', None))) - - if mi_value is None: - continue - + records = [] + for node1, node2, edge_data in tqdm(Graph.edges(data=True), desc="Processing edges"): + mi_value = edge_data.get("nmi", edge_data.get("mi", edge_data.get("weight"))) domain1 = res_to_domain.get(node1) domain2 = res_to_domain.get(node2) - - if domain1 is None or domain2 is None: + if mi_value is None or domain1 is None or domain2 is None: continue - if domain1 == domain2: - intra_domain_mi[domain1].append(mi_value) - else: - domain_pair = tuple(sorted([domain1, domain2])) - if domain_pair not in inter_domain_mi: - inter_domain_mi[domain_pair] = [] - inter_domain_mi[domain_pair].append(mi_value) - - print("\n" + "="*60) - print("INTRA-DOMAIN MUTUAL INFORMATION") - print("="*60) - - for domain, mi_values in intra_domain_mi.items(): - if mi_values: - total_mi = sum(mi_values) - avg_mi = total_mi / len(mi_values) - print(f"{domain}:") - print(f" Total MI: {total_mi:.4f}") - print(f" Average MI: {avg_mi:.6f}") - print(f" Number of edges: {len(mi_values)}") + records.append(("Intra-domain", domain1, domain1, mi_value)) else: - print(f"{domain}: No edges found") - - print("\n" + "="*60) - print("INTER-DOMAIN MUTUAL INFORMATION") - print("="*60) - - for domain_pair, mi_values in inter_domain_mi.items(): - if mi_values: - total_mi = sum(mi_values) - avg_mi = total_mi / len(mi_values) - print(f"{domain_pair[0]} <-> {domain_pair[1]}:") - print(f" Total MI: {total_mi:.4f}") - print(f" Average MI: {avg_mi:.6f}") - print(f" Number of edges: {len(mi_values)}") - - summary_data = [] - - for domain, mi_values in intra_domain_mi.items(): - if mi_values: - summary_data.append({ - 'Type': 'Intra-domain', - 'Domain1': domain, - 'Domain2': domain, - 'Total_MI': sum(mi_values), - 'Average_MI': sum(mi_values) / len(mi_values), - 'Num_Edges': len(mi_values) - }) - - for domain_pair, mi_values in inter_domain_mi.items(): - if mi_values: - summary_data.append({ - 'Type': 'Inter-domain', - 'Domain1': domain_pair[0], - 'Domain2': domain_pair[1], - 'Total_MI': sum(mi_values), - 'Average_MI': sum(mi_values) / len(mi_values), - 'Num_Edges': len(mi_values) - }) - - summary_df = pd.DataFrame(summary_data) - print("\n" + "="*60) + d1, d2 = sorted([domain1, domain2]) + records.append(("Inter-domain", d1, d2, mi_value)) + + edges_df = pd.DataFrame(records, columns=["Type", "Domain1", "Domain2", "MI"]) + summary_df = ( + edges_df.groupby(["Type", "Domain1", "Domain2"])["MI"] + .agg(Total_MI="sum", Average_MI="mean", Num_Edges="count") + .reset_index() + ) + + print("\n" + "=" * 60) print("SUMMARY") - print("="*60) + print("=" * 60) print(summary_df.to_string(index=False)) - + summary_df.to_csv(args.output, index=False) print(f"\n\033[1mResults saved to '{args.output}'\033[0m") exit(0) diff --git a/mdpath/src/bootstrap.py b/mdpath/src/bootstrap.py index a91dacd..3e30ac0 100644 --- a/mdpath/src/bootstrap.py +++ b/mdpath/src/bootstrap.py @@ -14,7 +14,7 @@ import numpy as np from mdpath.src.graph import GraphBuilder from mdpath.src.mutual_information import NMICalculator -from typing import Dict, Set, Tuple, List +from typing import Tuple import os @@ -179,15 +179,7 @@ def bootstrap_write(self, file_name: str) -> None: Returns: None write the path confidence intervals to a file. """ - for path, (mean, lower, upper) in self.path_confidence_intervals.items(): - path_str = " -> ".join(map(str, path)) - with open(file_name, "w") as file: - for path, ( - mean, - lower, - upper, - ) in self.path_confidence_intervals.items(): - path_str = " -> ".join(map(str, path)) - file.write( - f"{path_str}: Mean={mean}, 2.5%={lower}, 97.5%={upper}\n" - ) + with open(file_name, "w") as file: + for path, (mean, lower, upper) in self.path_confidence_intervals.items(): + path_str = " -> ".join(map(str, path)) + file.write(f"{path_str}: Mean={mean}, 2.5%={lower}, 97.5%={upper}\n") diff --git a/mdpath/src/cluster.py b/mdpath/src/cluster.py index 013f620..5541f55 100644 --- a/mdpath/src/cluster.py +++ b/mdpath/src/cluster.py @@ -18,7 +18,6 @@ from scipy.cluster import hierarchy from sklearn.metrics import silhouette_score import matplotlib.pyplot as plt -import seaborn as sns class PatwayClustering: diff --git a/mdpath/src/graph.py b/mdpath/src/graph.py index fffc499..2312671 100644 --- a/mdpath/src/graph.py +++ b/mdpath/src/graph.py @@ -17,7 +17,7 @@ import pandas as pd from scipy.spatial import cKDTree from Bio import PDB -from typing import Tuple, List +from typing import Tuple from multiprocessing import Pool from tqdm import tqdm from mdpath.src.structure import StructureCalculations diff --git a/mdpath/src/mutual_information.py b/mdpath/src/mutual_information.py index 74971b3..b4eec09 100644 --- a/mdpath/src/mutual_information.py +++ b/mdpath/src/mutual_information.py @@ -16,9 +16,7 @@ from tqdm import tqdm from itertools import combinations from sklearn.metrics import mutual_info_score -from sklearn.mixture import GaussianMixture from scipy.stats import entropy -from scipy.special import digamma class NMICalculator: diff --git a/mdpath/src/structure.py b/mdpath/src/structure.py index 6edc8e9..6955c51 100644 --- a/mdpath/src/structure.py +++ b/mdpath/src/structure.py @@ -62,23 +62,6 @@ def res_num_from_pdb(self) -> tuple: last_res_num = res_num return int(first_res_num), int(last_res_num) - def calculate_distance(self, atom1: tuple, atom2: tuple) -> float: - """Calculates the distance between two atoms. - - Args: - atom1 (tuple): Coordinates of the first atom. - - atom2 (tuple): Coordinates of the second atom. - - Returns: - distance (float): Normalized distance between the two atoms. - """ - distance_vector = [ - atom1[i] - atom2[i] for i in range(min(len(atom1), len(atom2))) - ] - distance = np.linalg.norm(distance_vector) - return distance - def _build_kdtree(self, dist: float): """Builds a KDTree from heavy atoms and returns close residue pairs. diff --git a/mdpath/src/visualization.py b/mdpath/src/visualization.py index e02bcfa..692925c 100644 --- a/mdpath/src/visualization.py +++ b/mdpath/src/visualization.py @@ -66,9 +66,6 @@ class MDPathVisualize: Attributes: None (only static methods)""" - def __init__(self) -> None: - pass - @staticmethod def residue_CA_coordinates(pdb_file: str, end: int) -> dict: """Collects CA atom coordinates for residues. @@ -97,37 +94,6 @@ def residue_CA_coordinates(pdb_file: str, end: int) -> dict: residue_coordinates_dict[res_id].append(atom.coord) return residue_coordinates_dict - @staticmethod - def cluster_prep_for_visualisation(cluster: list, pdb_file: str) -> list: - """Prepares pathway clusters for visualisation. - - Args: - cluster (list): Cluster of pathways. - - pdb_file (str): Path to PDB file. - - Returns: - cluster (list): Cluster of pathways with CA atom coordinates. - """ - new_cluster = [] - parser = PDB.PDBParser(QUIET=True) - structure = parser.get_structure("pdb_structure", pdb_file) - - for pathway in cluster: - pathways = [] - for residue in pathway: - res_id = ("", residue, "") - try: - res = structure[0][res_id] - atom = res["CA"] - coord = tuple(atom.get_coord()) - pathways.append(coord) - except KeyError: - print(f"Residue {res_id} not found.") - new_cluster.append(pathways) - - return new_cluster - @staticmethod def apply_backtracking(original_dict: dict, translation_dict: dict) -> dict: """Backtracks the original dictionary with a translation dictionary. @@ -204,18 +170,22 @@ def visualise_graph(graph: nx.Graph, k=0.1, node_size=200) -> None: plt.savefig("graph.png", dpi=300, bbox_inches="tight") @staticmethod - def precompute_path_properties(json_data: dict) -> list: - """Precomputes path properties for quicker visualization in Jupyter notebook. + def _precompute_properties(json_data: dict, detailed: bool) -> list: + """Walks cluster pathways and emits one record per valid coordinate segment. Args: json_data (dict): Cluster data with pathways and CA atom coordinates. + detailed (bool): If True, include pathway_index, path_segment_index and + path_number (used by :meth:`precompute_path_properties`). If False, + emit only clusterid/coords/color/radius. + Returns: - path_properties (list): List of path properties. Contains clusterid, pathway index, path segment index, coordinates, color, radius, and path number. + list: Per-segment property dictionaries. """ cluster_colors = {} color_index = 0 - path_properties = [] + properties = [] for clusterid, cluster in json_data.items(): cluster_colors[clusterid] = Colors[color_index % len(Colors)] @@ -234,31 +204,41 @@ def precompute_path_properties(json_data: dict) -> list: and len(coord2) == 3 ): coord_pair = (tuple(coord1), tuple(coord2)) - if coord_pair not in coord_pair_counts: - coord_pair_counts[coord_pair] = 0 - coord_pair_counts[coord_pair] += 1 - radius = 0.015 + 0.015 * (coord_pair_counts[coord_pair] - 1) - color = cluster_colors[clusterid] - - path_properties.append( - { - "clusterid": clusterid, - "pathway_index": pathway_index, - "path_segment_index": i, - "coord1": coord1, - "coord2": coord2, - "color": color, - "radius": radius, - "path_number": path_number, - } + coord_pair_counts[coord_pair] = ( + coord_pair_counts.get(coord_pair, 0) + 1 ) + radius = 0.015 + 0.015 * (coord_pair_counts[coord_pair] - 1) - path_number += 1 + entry = { + "clusterid": clusterid, + "coord1": coord1, + "coord2": coord2, + "color": cluster_colors[clusterid], + "radius": radius, + } + if detailed: + entry["pathway_index"] = pathway_index + entry["path_segment_index"] = i + entry["path_number"] = path_number + path_number += 1 + properties.append(entry) else: print( f"Ignoring pathway {pathway} as it does not fulfill the coordinate format." ) - return path_properties + return properties + + @staticmethod + def precompute_path_properties(json_data: dict) -> list: + """Precomputes path properties for quicker visualization in Jupyter notebook. + + Args: + json_data (dict): Cluster data with pathways and CA atom coordinates. + + Returns: + path_properties (list): List of path properties. Contains clusterid, pathway index, path segment index, coordinates, color, radius, and path number. + """ + return MDPathVisualize._precompute_properties(json_data, detailed=True) @staticmethod def precompute_cluster_properties_quick(json_data: dict) -> list: @@ -270,46 +250,7 @@ def precompute_cluster_properties_quick(json_data: dict) -> list: Returns: cluster_properties (list): List of cluster properties. Contains clusterid,coordinates, color, and radius. """ - cluster_colors = {} - color_index = 0 - cluster_properties = [] - - for clusterid, cluster in json_data.items(): - cluster_colors[clusterid] = Colors[color_index % len(Colors)] - color_index += 1 - coord_pair_counts = {} - - for pathway_index, pathway in enumerate(cluster): - for i in range(len(pathway) - 1): - coord1 = pathway[i][0] - coord2 = pathway[i + 1][0] - if ( - isinstance(coord1, list) - and isinstance(coord2, list) - and len(coord1) == 3 - and len(coord2) == 3 - ): - coord_pair = (tuple(coord1), tuple(coord2)) - if coord_pair not in coord_pair_counts: - coord_pair_counts[coord_pair] = 0 - coord_pair_counts[coord_pair] += 1 - radius = 0.015 + 0.015 * (coord_pair_counts[coord_pair] - 1) - color = cluster_colors[clusterid] - - cluster_properties.append( - { - "clusterid": clusterid, - "coord1": coord1, - "coord2": coord2, - "color": color, - "radius": radius, - } - ) - else: - print( - f"Ignoring pathway {pathway} as it does not fulfill the coordinate format." - ) - return cluster_properties + return MDPathVisualize._precompute_properties(json_data, detailed=False) @staticmethod def remove_non_protein(input_pdb: str, output_pdb: str) -> None: diff --git a/mdpath/tests/test_structure.py b/mdpath/tests/test_structure.py index 05503ce..8c4bc38 100644 --- a/mdpath/tests/test_structure.py +++ b/mdpath/tests/test_structure.py @@ -41,13 +41,6 @@ def test_res_num_from_pdb(pdb_file): assert last_res == 303 -def test_calculate_distance(pdb_file): - calc = StructureCalculations(pdb_file) - # Test valid distance calculation - distance = calc.calculate_distance((1.0, 1.0, 1.0), (4.0, 5.0, 6.0)) - assert np.isclose(distance, 7.071) # Example assertion with expected distance - - def test_calculate_residue_surroundings(pdb_file): calc = StructureCalculations(pdb_file) df = calc.calculate_residue_suroundings(dist=2.0, mode="close") diff --git a/mdpath/tests/test_visualization.py b/mdpath/tests/test_visualization.py index ca631d4..dc73039 100644 --- a/mdpath/tests/test_visualization.py +++ b/mdpath/tests/test_visualization.py @@ -92,42 +92,6 @@ def test_apply_backtracking(): assert result == expected -def test_cluster_prep_for_visualisation(): - pdb_file = "mock_pdb_file.pdb" - input_cluster = [[1, 2], [3]] - mock_coordinates = {1: (1.0, 1.0, 1.0), 2: (2.0, 2.0, 2.0), 3: (3.0, 3.0, 3.0)} - - with patch("Bio.PDB.PDBParser") as mock_parser: - mock_structure = MagicMock() - - def get_structure(name, file): - return mock_structure - - mock_parser.return_value.get_structure.side_effect = get_structure - - mock_residues = {} - for residue_id, coord in mock_coordinates.items(): - mock_residue = MagicMock() - mock_atom = MagicMock() - mock_atom.get_coord.return_value = coord - mock_residue.__getitem__.return_value = mock_atom - mock_residues[("", residue_id, "")] = mock_residue - - def getitem(res_id): - if res_id in mock_residues: - return mock_residues[res_id] - else: - raise KeyError - - mock_structure[0].__getitem__.side_effect = getitem - - result = MDPathVisualize.cluster_prep_for_visualisation(input_cluster, pdb_file) - - expected_result = [[(1.0, 1.0, 1.0), (2.0, 2.0, 2.0)], [(3.0, 3.0, 3.0)]] - - assert result == expected_result - - def test_format_dict(): input_dict = {"array": np.array([1, 2, 3]), "nested_list": [1, 2, np.array([3, 4])]} expected_output = {"array": [1, 2, 3], "nested_list": [1, 2, [3, 4]]} diff --git a/pyproject.toml b/pyproject.toml index dd3cb89..9544c4f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,6 @@ dependencies = [ "scipy", "biopython", "tqdm", - "seaborn", "nglview", "pillow", "numpy-stl"