Refactor and optimize code:

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_<prefix>.pdb`` and ``selected_<prefix>_<i>.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()

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 <path_to_graph.pkl> -config <path_to_config.csv> -output <output_path>
     """
-    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)

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")

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:

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

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:

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.