plasmidhub 1.0.0__py3-none-any.whl

plasmidhub/__init__.py

@@ -0,0 +1 @@
+# plasmidhub package

plasmidhub/abricate.py

@@ -0,0 +1,46 @@
+import os
+import subprocess
+import shutil
+import glob
+import logging
+logger = logging.getLogger(__name__)
+
+def run_abricate_bulk(input_dir, results_dir, db_list, threads=None):
+    os.makedirs(results_dir, exist_ok=True)
+
+    # Use default thread count if not provided
+    if threads is None:
+        threads = 4
+
+    # Move into input_dir because wildcard expansion happens here
+    original_dir = os.getcwd()
+    os.chdir(input_dir)
+
+    # Collect all fasta-like files
+    fasta_files = sorted(
+        glob.glob("*.fna") +
+        glob.glob("*.fa") +
+        glob.glob("*.fasta")
+    )
+
+    if not fasta_files:
+        raise RuntimeError(f"No input files found in {input_dir} with .fna/.fa/.fasta extensions.")
+
+    for db in db_list:
+        logger.info(f"Running abricate on database: {db}")
+
+        # Build the shell command with all fasta file names
+        cmd = f"abricate {' '.join(fasta_files)} --db {db} -t {threads}"
+
+        # Output file path (temporary inside input_dir)
+        temp_output = f"{db}.abr"
+        with open(temp_output, "w") as out_f:
+            subprocess.run(cmd, shell=True, stdout=out_f, stderr=subprocess.DEVNULL)
+
+        # Move the output to results_dir
+        final_output_path = os.path.join(results_dir, f"{db}.abr")
+        shutil.move(temp_output, final_output_path)
+        logger.info(f"Saved: {final_output_path}")
+
+    # Return to original directory
+    os.chdir(original_dir)

plasmidhub/ani.py

@@ -0,0 +1,29 @@
+import os
+import subprocess
+import logging
+logger = logging.getLogger(__name__)
+
+def run_fastani(plasmid_list_file, fragLen=1000, minFrag=3, kmer=14, output_dir=".", threads=None):
+    if threads is None:
+        threads = 4
+
+    output_file = os.path.join(output_dir, "fastani_raw_results.tsv")
+    cmd = [
+        "fastANI",
+        "--ql", plasmid_list_file,
+        "--rl", plasmid_list_file,
+        "-o", output_file,
+        "--fragLen", str(fragLen),
+        "--minFraction", str(minFrag),
+        "--kmer", str(kmer),
+        "-t", str(threads)
+    ]
+    logger.info("Running FastANI with command:")
+    logger.info(" ".join(cmd))
+    result = subprocess.run(cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        logger.error("FastANI failed with error:")
+        logger.error(result.stderr)
+        exit(1)
+    else:
+        logger.info("FastANI completed successfully.")

plasmidhub/cluster_color.py

@@ -0,0 +1,48 @@
+import os
+import matplotlib.pyplot as plt
+import random
+import logging
+
+logger = logging.getLogger(__name__)
+
+def assign_cluster_colors(results_dir, mapping_file):
+    cluster_list_path = os.path.join(results_dir, "cluster_list.txt")
+    color_file = os.path.join(results_dir, "cluster_colours.txt")
+    
+    clusters = []
+    with open(cluster_list_path) as f:
+        next(f)  # Skip header
+        for line in f:
+            if line.strip():
+                cluster_file, _ = line.strip().split('\t')
+                cluster = cluster_file.replace('.txt', '')
+                clusters.append(cluster)
+
+    n_clusters = len(clusters)
+
+    # Start with base colors from tab20
+    cmap = plt.get_cmap('tab20')
+    base_colors = [
+        '#{:02x}{:02x}{:02x}'.format(int(r * 255), int(g * 255), int(b * 255))
+        for r, g, b in cmap.colors
+    ]
+
+    used_colors = set(base_colors[:min(n_clusters, len(base_colors))])
+    full_color_list = base_colors[:min(n_clusters, len(base_colors))]
+
+    # Generate additional distinct random colors if needed
+    while len(full_color_list) < n_clusters:
+        while True:
+            color = "#{:06x}".format(random.randint(0, 0xFFFFFF))
+            if color not in used_colors:
+                used_colors.add(color)
+                full_color_list.append(color)
+                break
+
+    color_map = dict(zip(clusters, full_color_list))
+
+    with open(color_file, 'w') as out:
+        for cluster, color in color_map.items():
+            out.write(f"{cluster}\t{color}\n")
+
+    logger.info(f"Cluster colors saved to: {color_file}")

plasmidhub/clustering.py

@@ -0,0 +1,143 @@
+import os
+import pandas as pd
+from collections import defaultdict
+import argparse
+import logging
+logger = logging.getLogger(__name__)
+
+def find_valid_subclusters(results_dir):
+    valid_subclusters = []
+
+    for filename in sorted(os.listdir(results_dir)):
+        if filename.startswith("subcluster_") and filename.endswith("_plasmids.txt"):
+            filepath = os.path.join(results_dir, filename)
+            with open(filepath, 'r') as f:
+                plasmid_count = sum(1 for _ in f)
+
+            if plasmid_count >= 3:  # Hardcoded rule
+                valid_subclusters.append((filename, plasmid_count))
+
+    valid_subclusters.sort(key=lambda x: x[1], reverse=True)
+    return valid_subclusters
+
+def write_subcluster_list(valid_subclusters, output_path):
+    with open(output_path, "w") as f:
+        f.write("Subcluster\tPlasmids\n")
+        for subcluster, count in valid_subclusters:
+            f.write(f"{subcluster}\t{count}\n")
+
+def extract_clusters(valid_subclusters, results_dir, fastani_path, output_dir):
+    fastani_df = pd.read_csv(fastani_path, sep="\t")
+
+    for subcluster_file, _ in valid_subclusters:
+        full_path = os.path.join(results_dir, subcluster_file)
+        try:
+            with open(full_path, "r") as f:
+                original_plasmids = set(line.strip() for line in f)
+        except FileNotFoundError:
+            logger.warning(f"File {subcluster_file} not found. Skipping.")
+            continue
+
+        subcluster_plasmids = original_plasmids.copy()
+
+        connections = defaultdict(set)
+        for _, row in fastani_df.iterrows():
+            q, r = row["Query"], row["Reference"]
+            if q in subcluster_plasmids and r in subcluster_plasmids:
+                connections[q].add(r)
+                connections[r].add(q)
+
+        # Iteratively remove nodes with the fewest connections until we get a complete subgraph
+        while True:
+            current_nodes = set(connections.keys())
+            if len(current_nodes) < 3:
+                subcluster_plasmids = set()
+                break
+
+            # Check if the current graph is a complete subgraph (clique)
+            complete = all(len(connections[node]) == len(current_nodes) - 1 for node in current_nodes)
+            if complete:
+                subcluster_plasmids = current_nodes
+                break
+
+            # Find the node with the fewest connections (lowest degree)
+            min_node = min(current_nodes, key=lambda x: len(connections[x]))
+
+            # Remove that node from the graph
+            del connections[min_node]
+            for conn in connections.values():
+                conn.discard(min_node)
+
+        # Step 7: Save the refined subcluster to a new file with the desired naming format
+        cluster_number = subcluster_file.split("_")[1]  # Extract the number from subcluster_XX_plasmids.txt
+        output_file = f"cluster_{cluster_number}.txt"
+
+        cluster_path = os.path.join(output_dir, output_file)
+        with open(cluster_path, "w") as f:
+
+            for plasmid in subcluster_plasmids:
+                f.write(plasmid + "\n")
+
+        
+def filter_clusters_by_size(output_dir, min_cluster_size):
+    for filename in os.listdir(output_dir):
+        if filename.startswith("cluster_") and filename.endswith(".txt"):
+            path = os.path.join(output_dir, filename)
+            with open(path, "r") as f:
+                lines = f.readlines()
+            if len(lines) < min_cluster_size:
+                os.remove(path)
+
+def write_cluster_list(output_dir, output_path):
+    cluster_files = []
+    for filename in os.listdir(output_dir):
+        if (
+            filename.startswith("cluster_") 
+            and filename.endswith(".txt") 
+            and filename not in {os.path.basename(output_path), "cluster_colours.txt"}  # exclude output file itself and cluster_colours.txt
+        ):
+            path = os.path.join(output_dir, filename)
+            with open(path, "r") as f:
+                count = sum(1 for _ in f)
+            cluster_files.append((filename, count))
+    cluster_files.sort(key=lambda x: x[1], reverse=True)
+    with open(output_path, "w") as f:
+        f.write("Cluster\tPlasmids\n")
+        for filename, count in cluster_files:
+            f.write(f"{filename}\t{count}\n")
+                
+def main(results_dir, min_cluster_size):
+    fastani_path = os.path.join(results_dir, "ANI_results_final.tsv")
+    subcluster_list_output = os.path.join(results_dir, "subcluster_list.txt")
+    cluster_list_output = os.path.join(results_dir, "cluster_list.txt")
+
+    logger.info("Finding valid subclusters (>=3 plasmids)...")
+    valid_subclusters = find_valid_subclusters(results_dir)
+
+    write_subcluster_list(valid_subclusters, subcluster_list_output)
+
+    logger.info("Identifying  clusters...")
+    extract_clusters(valid_subclusters, results_dir, fastani_path, results_dir)
+
+    logging.info(f"Keep only clusters with >={min_cluster_size} plasmids...")
+    filter_clusters_by_size(results_dir, min_cluster_size)
+
+    write_cluster_list(results_dir, cluster_list_output)
+
+    # Check: warn user if cluster_list.txt is empty
+    if os.path.exists(cluster_list_output):
+        with open(cluster_list_output, "r") as f:
+            lines = f.readlines()
+            if len(lines) <= 1:
+                logger.warning("No clusters detected with the given parameters!")
+
+
+    # logger.info("Done!")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Clustering Tool")
+    parser.add_argument("results_dir", help="Path to results directory created by main.py")
+    parser.add_argument("--min_cluster_size", type=int, default=3, help="Minimum number of plasmids in final cluster (default: 3)")
+    args = parser.parse_args()
+
+    main(args.results_dir, args.min_cluster_size)

plasmidhub/filtering.py

@@ -0,0 +1,54 @@
+import os
+import pandas as pd
+import logging
+logger = logging.getLogger(__name__)
+
+def strip_paths_in_fastani(input_file):
+    """Remove directory paths from plasmid names in FastANI output."""
+    df = pd.read_csv(input_file, sep='\t', header=None)
+    df.columns = ['Query', 'Reference', 'ANI', 'Matching_Frags_Query', 'Matching_Frags_Ref']
+    df['Query'] = df['Query'].apply(os.path.basename)
+    df['Reference'] = df['Reference'].apply(os.path.basename)
+    df.to_csv(input_file, sep='\t', header=False, index=False)
+    
+def strip_paths_in_plasmid_list(file_path):
+    """Remove directory paths from plasmid names in Plasmid_list.txt."""
+    with open(file_path, 'r') as f:
+        lines = f.readlines()
+    stripped_lines = [os.path.basename(line.strip()) + '\n' for line in lines]
+    with open(file_path, 'w') as f:
+        f.writelines(stripped_lines)
+
+def filter_self_comparisons(input_file, output_file):
+    df = pd.read_csv(input_file, sep='\t', header=None)
+    df.columns = ['Query', 'Reference', 'ANI', 'Matching_Frags_Query', 'Matching_Frags_Ref']
+    df_filtered = df[df['Query'] != df['Reference']]
+    df_filtered.to_csv(output_file, sep='\t', index=False)
+
+def add_plasmid_sizes(ani_file, sizes_file, output_file):
+    sizes_df = pd.read_csv(sizes_file, sep='\t')
+    size_dict = dict(zip(sizes_df['PlasmidID'], sizes_df['Size']))
+    ani_df = pd.read_csv(ani_file, sep='\t')
+    def get_size(plasmid_id):
+        return size_dict.get(plasmid_id, None)
+    ani_df['Query_size'] = ani_df['Query'].apply(get_size)
+    ani_df['Reference_size'] = ani_df['Reference'].apply(get_size)
+    ani_df.to_csv(output_file, sep='\t', index=False)
+    
+def apply_filters(input_file, output_file, frag_len=1000, coverage_threshold=0.5, ani_threshold=95.0):
+    df = pd.read_csv(input_file, sep='\t')
+
+    # Convert matching fragments to base pairs using user-defined fragment length
+    df['Matching_Frags_Query_bp'] = df['Matching_Frags_Query'] * frag_len
+    df['Matching_Frags_Ref_bp'] = df['Matching_Frags_Ref'] * frag_len
+
+    # Apply filtering logic
+    filtered = df[
+        (df['Matching_Frags_Query_bp'] > df['Reference_size'] * coverage_threshold) &
+        (df['Matching_Frags_Ref_bp'] > df['Query_size'] * coverage_threshold) &
+        (df['Matching_Frags_Query_bp'] > df['Query_size'] * coverage_threshold) &
+        (df['ANI'] >= ani_threshold)
+    ]
+
+    filtered.to_csv(output_file, sep='\t', index=False)
+    logger.info(f"[INFO] Applied filters (ANI = {ani_threshold}, coverage = {coverage_threshold*100:.1f}%). ")