diff --git a/.gitignore b/.gitignore index e87154f..994d40b 100644 --- a/.gitignore +++ b/.gitignore @@ -15,7 +15,8 @@ result .DS_Store .env data - +nematodes +output # Logs logs *.log diff --git a/requirements.txt b/requirements.txt index f10e6af..2209e8e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,4 +5,5 @@ networkx==3.3 powerlaw==1.5 ete3==3.1.3 fastapi==0.111.0 -slowapi==0.1.9 \ No newline at end of file +slowapi==0.1.9 +polars==1.32.3 \ No newline at end of file diff --git a/scripts/simfin.py b/scripts/simfin.py new file mode 100755 index 0000000..86997c8 --- /dev/null +++ b/scripts/simfin.py @@ -0,0 +1,457 @@ +import argparse +import collections +import os +import random +import sys +import time +import traceback + +import numpy as np + + +def parse_args(): + parser = argparse.ArgumentParser( + prog="simfin.py", + description="""Simulate KinFin input datasets for benchmarking. + +Arguments: + -n : Number of clusters to simulate (e.g., 25000) + -m : Number of unique taxa (e.g., 20) + -lc: Array of integers defining attribute labels: + * Length = number of attributes + * Each value = number of unique labels for that attribute + Example: -lc 2 4 6 creates 3 attributes with 2, 4, and 6 unique labels + +Example commands: + # Generate 25k clusters for 20 taxa with taxid and outgroup info + python scripts/simfin.py -n 25000 -m 20 --with_taxid --with_outgroup -g /path/to/Orthogroups.txt -o ./test_data + + # Generate 50k clusters for 30 taxa with 3 attributes + python scripts/simfin.py -n 50000 -m 30 -lc 2 4 6 --with_taxid --with_outgroup -g /path/to/Orthogroups.txt -o ./test_data + +Notes: + - Change -n for cluster count + - Change -m for taxon count + - Adjust -lc to set attributes and label counts + - Use --with_taxid and --with_outgroup to include TAXID and OUT columns + - The output directory is specified with -o""", + epilog="Author: DRL (2025)", + formatter_class=argparse.RawTextHelpFormatter, + ) + parser.add_argument( + "-n", + "--sim_clusters", + required=False, + default=1_000_000, + type=int, + help="Number of clusters to simulate", + ) + parser.add_argument( + "-m", + "--taxon_counts", + required=False, + default=20, + type=int, + help="Number of unique taxa to simulate", + ) + parser.add_argument( + "-r", + "--random_seed", + required=False, + default=19, + type=int, + help="Random seed for reproducibility", + ) + parser.add_argument( + "-f", + "--fasta_sim", + action="store_true", + required=False, + default=False, + help="Also simulate protein sequences in FASTA format", + ) + parser.add_argument( + "-g", "--groups_raw", required=True, help="Path to input orthogroups file" + ) + + parser.add_argument( + "-wt", + "--with_taxid", + action="store_true", + required=False, + help="Include TAXID column in config", + ) + parser.add_argument( + "-wo", + "--with_outgroup", + action="store_true", + required=False, + help="Include OUT column in config", + ) + parser.add_argument( + "-lc", + "--label_counts", + type=int, + nargs="+", + required=False, + help="Number of unique labels per attribute (length = num attributes)", + ) + parser.add_argument( + "-o", + "--output_dir", + required=False, + default=".", + help="Directory to write output files", + ) + + return parser.parse_args() + + +def get_count_matrix(orthogroups_f): + taxon_ids = set() + taxon_counters = [] + with open(orthogroups_f) as fh: + for i, line in enumerate(fh): + cluster_id, protein_string = line.rstrip("\n").split(": ") + taxon_counter = collections.Counter( + [ + protein_name.split(".")[0] + for protein_name in protein_string.split(" ") + ] + ) + taxon_ids.update(taxon_counter.keys()) + taxon_counters.append(taxon_counter) + taxon_ids = sorted(taxon_ids) + matrix = np.zeros(shape=(len(taxon_counters), len(taxon_ids)), dtype=int) + for j, taxon_counter in enumerate(taxon_counters): + for i, taxon_id in enumerate(taxon_ids): + matrix[j][i] = taxon_counter[taxon_id] + return taxon_ids, matrix + + +def get_simulated_counts( + matrix, + sim_clusters, + seed, +): + rng = np.random.default_rng(seed) + simulated_counts = rng.choice( + matrix, size=sim_clusters, replace=True, axis=0, shuffle=True + ) + simulated_counts_sorted = simulated_counts[ + np.argsort(simulated_counts.sum(axis=1))[::-1], : + ] + return simulated_counts_sorted + + +def write_sequence_ids( + taxon_ids, + protein_ids_by_taxon_id, + seed, + sim_clusters, + output_dir, +): + lines = [] + for taxon_index, taxon_id in enumerate(taxon_ids): + protein_list = protein_ids_by_taxon_id.get(taxon_id, []) + for protein_index, protein_id in enumerate(protein_list): + lines.append(f"{taxon_index}_{protein_index}: {protein_id}") + + sequence_ids_f = os.path.join( + output_dir, f"sequence_ids.sim_{sim_clusters}.seed_{seed}.txt" + ) + + with open(sequence_ids_f, "w") as fh: + fh.write("\n".join(lines) + "\n") + print(f"[+] wrote {sequence_ids_f}") + + +def write_clusters(taxon_ids, simulated_counts, seed, output_dir): + rows = [] + padding = len(str(simulated_counts.shape[0])) + max_count_by_taxon = collections.Counter() + protein_ids_by_taxon_id = collections.defaultdict(list) + for n, cluster in enumerate(simulated_counts): + proteins_in_cluster = [] + for taxon_i, count in enumerate(cluster): + for _ in range(count): + taxon_id = taxon_ids[taxon_i] + max_count_by_taxon[taxon_id] += 1 + protein_id = f"{taxon_id}.seq{max_count_by_taxon[taxon_id]}" + protein_ids_by_taxon_id[taxon_id].append(protein_id) + proteins_in_cluster.append(protein_id) + proteins_in_cluster.sort() + row = [f"OG0{n:0{padding}}:"] + row.extend(proteins_in_cluster) + rows.append(" ".join(row)) + + cluster_f = os.path.join(output_dir, f"orthogroups.sim_{len(rows)}.seed_{seed}.txt") + with open(cluster_f, "w") as fh: + fh.write("\n".join(rows) + "\n") + print(f"[+] wrote {cluster_f}") + return protein_ids_by_taxon_id + + +def write_config_file( + taxon_ids, + sim_clusters, + seed, + with_taxid, + with_outgroup, + label_counts, + output_dir, +): + print("[+] Generating flexible config file...") + rng = np.random.default_rng(seed) + header = ["#IDX", "TAXON"] + if with_taxid: + header.append("TAXID") + if with_outgroup: + header.append("OUT") + + all_label_pools = [] + if label_counts: + num_attributes = len(label_counts) + attribute_names = [f"attr{i+1}" for i in range(num_attributes)] + header.extend(attribute_names) + for i, num_labels in enumerate(label_counts): + pool = [f"attr{i+1}_label{j+1}" for j in range(num_labels)] + all_label_pools.append(pool) + + rows = [",".join(header)] + taxid_pool = [ + 6238, + 318479, + 42156, + 6277, + 6239, + 1147741, + 42157, + 387005, + 7209, + 108094, + 6279, + 6287, + 6293, + 103827, + 42157, + 6280, + 6282, + 6293, + 7209, + ] + outgroup_idx = rng.choice(np.arange(len(taxon_ids))) if with_outgroup else -1 + + for idx, taxon_id in enumerate(taxon_ids): + row_data = [str(idx), taxon_id] + + if with_taxid: + row_data.append(str(rng.choice(taxid_pool))) + if with_outgroup: + row_data.append("1" if idx == outgroup_idx else "0") + if label_counts: + for i in range(len(label_counts)): + label_to_add = rng.choice(all_label_pools[i]) + row_data.append(label_to_add) + + rows.append(",".join(row_data)) + + config_f = os.path.join(output_dir, f"config.sim_{sim_clusters}.seed_{seed}.txt") + with open(config_f, "w") as fh: + fh.write("\n".join(rows) + "\n") + print(f"[+] wrote {config_f}") + + +def write_proteins( + protein_ids_by_taxon_id, + seed, + output_dir, + sim_clusters, + length_max=200, + length_min=10, +): + rng = np.random.default_rng(seed) # length rng + random.seed(seed) # aminoacid rng + alphabet = [ + "A", + "C", + "D", + "E", + "F", + "G", + "H", + "I", + "K", + "L", + "M", + "N", + "P", + "Q", + "R", + "S", + "T", + "V", + "W", + "Y", + ] + for taxon_id, protein_ids in protein_ids_by_taxon_id.items(): + lines = [] + fasta_f = os.path.join( + output_dir, f"{taxon_id}.sim_{sim_clusters}.seed_{seed}.fasta" + ) + + protein_lengths = rng.integers( + size=len(protein_ids), low=length_min, high=length_max + ) + # print(f"{taxon_id=} {protein_ids=} {protein_lengths=}") + for protein_id, protein_length in zip(protein_ids, protein_lengths): + lines.append( + f">{protein_id}\nM{''.join(random.choices(alphabet, k=(protein_length - 1)))}\n" + ) + with open(fasta_f, "w") as fh: + fh.write("".join(lines)) + print(f"[+] wrote {fasta_f}") + + +def adjust_matrix_dimensions(taxon_ids, matrix, desired_taxon_count, seed): + num_original_taxa = matrix.shape[1] + rng = np.random.default_rng(seed) + + if desired_taxon_count <= 0 or desired_taxon_count == num_original_taxa: + print(f"[+] Using all {num_original_taxa} available taxa.") + return taxon_ids, matrix + + if desired_taxon_count < num_original_taxa: + print( + f"[+] Subsampling from {num_original_taxa} down to {desired_taxon_count} taxa." + ) + all_indices = np.arange(num_original_taxa) + selected_indices = rng.choice( + all_indices, size=desired_taxon_count, replace=False + ) + selected_indices.sort() + new_taxon_ids = [taxon_ids[i] for i in selected_indices] + new_matrix = matrix[:, selected_indices] + return new_taxon_ids, new_matrix + + if desired_taxon_count > num_original_taxa: + num_to_add = desired_taxon_count - num_original_taxa + print( + f"[+] Extrapolating from {num_original_taxa} to {desired_taxon_count} by adding {num_to_add} synthetic taxa." + ) + template_indices = rng.choice( + np.arange(num_original_taxa), size=num_to_add, replace=True + ) + new_columns = matrix[:, template_indices] + new_matrix = np.hstack((matrix, new_columns)) + new_taxon_ids = list(taxon_ids) + for i in range(num_to_add): + new_taxon_ids.append(f"SimFin_{i+1}") + + return new_taxon_ids, new_matrix + + +""" + Feature request: + - simulate data with arg --taxon_ids M + - uses real counts to spit out clusters for more or fewer taxon_ids + +(n,m) matrix of counts +- n = number of clusters +- m = number of taxons + +[ +[1 1 1 1] +[2 2 2 2] +[1 0 1 0] +] + +- cluster counts: 10_000 advanced 100_000 1_000_000 +- taxon counts: 10 19 50 100 +- config: + attributes 2 5 10 + labels 2 5 half-of-taxon-count +""" +if __name__ == "__main__": + __version__ = 0.1 + t_0 = time.monotonic() + args = parse_args() + try: + os.makedirs(args.output_dir, exist_ok=True) + if args.label_counts: + for count in args.label_counts: + if count < 2: + print( + f"\n[X] FATAL ERROR: All values for --label_counts (-lc) must be 2 or greater. You provided: {count}" + ) + sys.exit(1) + + print(f"[+] For {args.output_dir} ...") + + print( + f"[+] Simulating based on {args.sim_clusters=} {args.random_seed=} {args.groups_raw=} ..." + ) + taxon_ids, matrix = get_count_matrix(args.groups_raw) + taxon_ids, matrix = adjust_matrix_dimensions( + taxon_ids, + matrix, + args.taxon_counts, + args.random_seed, + ) + + non_empty_matrix = matrix[matrix.sum(axis=1) > 0] + + if non_empty_matrix.shape[0] == 0: + print( + "\n[X] FATAL ERROR: No non-empty orthogroups exist for the selected taxa." + ) + print( + "[X] Cannot proceed with simulation. Try a different or larger taxon set." + ) + sys.exit() + + print( + f"[+] Found {non_empty_matrix.shape[0]} non-empty cluster patterns to use for simulation." + ) + + simulated_counts = get_simulated_counts( + non_empty_matrix, + args.sim_clusters, + args.random_seed, + ) + print(f"[+] Simulated counts into {simulated_counts.shape=}") + protein_ids_by_taxon_id = write_clusters( + taxon_ids, + simulated_counts, + args.random_seed, + args.output_dir, + ) + write_sequence_ids( + taxon_ids, + protein_ids_by_taxon_id, + args.random_seed, + args.sim_clusters, + args.output_dir, + ) + write_config_file( + taxon_ids, + args.sim_clusters, + args.random_seed, + args.with_taxid, + args.with_outgroup, + args.label_counts, + args.output_dir, + ) + if args.fasta_sim: + write_proteins( + protein_ids_by_taxon_id, + args.random_seed, + args.output_dir, + args.sim_clusters, + ) + print(f"[+] Done {args.output_dir} ...\n") + + except Exception as exc: + print(f"[X] Error {exc}\n{traceback.format_exc()}") + finally: + pass diff --git a/src/api/__init__.py b/src/api/__init__.py index 2603a62..3c22f44 100644 --- a/src/api/__init__.py +++ b/src/api/__init__.py @@ -2,34 +2,21 @@ from slowapi.errors import RateLimitExceeded from api.utils import load_clustering_datasets -from core.input import ServeArgs from .core.limiter import limiter def run_server( - args: ServeArgs, + port: int, nodesdb_f: str, - pfam_mapping_f: str, - ipr_mapping_f: str, - go_mapping_f: str, cluster_f: str, - taxon_idx_mapping_file: str, - sequence_ids_f: str, + ndb_f: str, + # pfam_mapping_f: str, + # ipr_mapping_f: str, + # go_mapping_f: str, + # taxon_idx_mapping_file: str, + # sequence_ids_f: str, ) -> None: - """ - Starts the uvicorn server - - Parameters: - - args [ServeArgs] : An object containing server configuration arguments, such as the port. - - nodesdb_f [str] : File path to the nodesDB file. - - pfam_mapping_f [str] : File path to the PFAM mapping file. - - ipr_mapping_f [str] : File path to the InterPro mapping file. - - go_mapping_f [str] : File path to the Gene Ontology mapping file. - - cluster_f [str] : File path to the clustering data file. - - taxon_idx_mapping_file [str] : File path to the taxon index mapping file. - - sequence_ids_f [str] : File path to the sequence IDs file. - """ import os import uvicorn @@ -44,12 +31,10 @@ def run_server( from api.sessions import query_manager query_manager.cluster_f = cluster_f - query_manager.sequence_ids_f = sequence_ids_f - query_manager.taxon_idx_mapping_file = taxon_idx_mapping_file query_manager.nodesdb_f = nodesdb_f - query_manager.pfam_mapping_f = pfam_mapping_f - query_manager.ipr_mapping_f = ipr_mapping_f - query_manager.go_mapping_f = go_mapping_f + + os.environ["NODESDB_F"] = nodesdb_f + os.environ["NDB_F"] = ndb_f app = FastAPI() app.state.limiter = limiter @@ -82,4 +67,4 @@ def hello(): app.include_router(router) - uvicorn.run(app=app, port=args.port) + uvicorn.run(app=app, port=port) diff --git a/src/api/endpoints.py b/src/api/endpoints.py index 9abe067..1b997c6 100644 --- a/src/api/endpoints.py +++ b/src/api/endpoints.py @@ -8,6 +8,7 @@ from functools import wraps from typing import Any, Dict, List, Optional +import polars as pl from fastapi import APIRouter, Body, Depends, HTTPException, Query, Request from fastapi.responses import FileResponse, JSONResponse, StreamingResponse from fastapi.security import APIKeyHeader @@ -32,7 +33,16 @@ run_cli_command, sort_and_paginate_result, ) -from core.utils import check_file +from internal import ( + attribute_metrics, + cluster_metrics, + cluster_summary, + counts, + parsers, + plots, + utils, +) +from internal.utils import check_file from .config.limits import LIMIT_INIT, LIMIT_STANDARD from .core.limiter import limiter @@ -51,6 +61,98 @@ router = APIRouter() +def get_session_data(session_dir: str): + config_file = os.path.join(session_dir, "config.txt") + if not os.path.exists(config_file): + raise FileNotFoundError(f"Config file not found: {config_file}") + + nodesdb_f = os.environ.get("NODESDB_F") + ndb_f = os.environ.get("NDB_F") + + if not nodesdb_f or not ndb_f: + raise RuntimeError("NODESDB_F and NDB_F environment variables must be set") + + nodesdb = parsers.nodesdb(filepath=nodesdb_f, outpath=ndb_f) + config_df, attributes = parsers.configfile(config_file, nodesdb) + + cluster_file = os.path.join(session_dir, "cluster_data.parquet") + if not os.path.exists(cluster_file): + raise FileNotFoundError(f"Cluster data not found: {cluster_file}") + + cluster_df = pl.read_parquet(cluster_file) + + label_columns = [ + col for col in config_df.columns if col not in ("#IDX", "TAXON", "TAXID", "OUT") + ] + + taxon_label_dict = { + row["TAXON"]: { + **{label: row[label] for label in label_columns}, + "TAXON": row["TAXON"], + "all": "all", + } + for row in config_df.to_dicts() + } + + return cluster_df, config_df, attributes, taxon_label_dict + + +async def run_analysis( + session_id: str, result_dir: str, cluster_file: str, config_file: str +): + """ + Args: + session_id: The session ID + result_dir: Path to the result directory + cluster_file: Path to the cluster file + config_file: Path to the config file + """ + status_file = os.path.join(result_dir, f"{session_id}.status") + try: + with open(status_file, "w") as f: + f.write("status=running\n") + + log_path = os.path.join(result_dir, "kinfin.log") + from internal.logger import setup_logger + + setup_logger(log_path) + + nodesdb_f = os.environ.get("NODESDB_F") + ndb_f = os.environ.get("NDB_F") + + if not nodesdb_f or not ndb_f: + raise RuntimeError("NODESDB_F and NDB_F environment variables must be set") + + nodesdb = parsers.nodesdb(filepath=nodesdb_f, outpath=ndb_f) + config_df, attributes = parsers.configfile(config_file, nodesdb) + + utils.setup_dirs(result_dir, attributes) + + cluster_df = parsers.clusterfile(cluster_file, config_df, result_dir) + + from internal import classify_clusters + + cluster_df = classify_clusters(cluster_df, config_df, attributes) + + cluster_data_path = os.path.join(result_dir, "cluster_data.parquet") + cluster_df.write_parquet(cluster_data_path) + + del cluster_df, config_df, nodesdb + + with open(status_file, "w") as f: + f.write("status=completed\n") + f.write("exit_code=0\n") + + LOGGER.info(f"[✓] Analysis completed for session {session_id}") + + except Exception as e: + with open(status_file, "w") as f: + f.write("status=error\n") + f.write("exit_code=1\n") + f.write(f"error={str(e)}\n") + LOGGER.error(f"[✗] Analysis failed for session {session_id}: {e}") + + class InputSchema(BaseModel): config: List[Dict[str, str]] clusterId: str @@ -203,7 +305,6 @@ async def initialize(input_data: InputSchema, request: Request): status_code=400, ) cluster_info = CLUSTERING_DATASETS.get(input_data.clusterId) - if not cluster_info: return JSONResponse( content=ResponseSchema( @@ -219,13 +320,16 @@ async def initialize(input_data: InputSchema, request: Request): cluster_path = os.path.join(KINFIN_WORKDIR, cluster_info["path"]) cluster_f = os.path.join(cluster_path, "Orthogroups.txt") - sequence_ids_f = os.path.join(cluster_path, "kinfin.SequenceIDs.txt") - taxon_idx_mapping_file = os.path.join(cluster_path, "taxon_idx_mapping.json") + # ! NOT NEEDED + # sequence_ids_f = os.path.join(cluster_path, "kinfin.SequenceIDs.txt") + # taxon_idx_mapping_file = os.path.join(cluster_path, "taxon_idx_mapping.json") try: + # TODO check_file(cluster_f, install_kinfin=True) - check_file(sequence_ids_f, install_kinfin=True) - check_file(taxon_idx_mapping_file, install_kinfin=True) + # ! NOT NEEDED + # check_file(sequence_ids_f, install_kinfin=True) + # check_file(taxon_idx_mapping_file, install_kinfin=True) except FileNotFoundError as e: return JSONResponse( content=ResponseSchema( @@ -238,30 +342,9 @@ async def initialize(input_data: InputSchema, request: Request): ) session_id, result_dir = query_manager.get_or_create_session(input_data.config) - config_f = os.path.join(result_dir, "config.json") - with open(config_f, "w") as file: - json.dump(input_data.config, file) - - command = [ - "python", - "src/main.py", - "analyse", - "-g", - cluster_f, - "-c", - config_f, - "-s", - sequence_ids_f, - "-m", - taxon_idx_mapping_file, - "-o", - result_dir, - "--plot_format", - "png", - ] - - status_file = os.path.join(result_dir, f"{session_id}.status") - asyncio.create_task(run_cli_command(command, status_file)) + config_f = os.path.join(result_dir, "config.txt") + + asyncio.create_task(run_analysis(session_id, result_dir, cluster_f, config_f)) return JSONResponse( content=ResponseSchema( @@ -413,15 +496,24 @@ async def get_counts_by_tanon( filepath = os.path.join(result_dir, COUNTS_FILEPATH) if not os.path.exists(filepath): - return JSONResponse( - content=ResponseSchema( - status="error", - message=f"{RUN_SUMMARY_FILEPATH} File Not Found", - error="File does not exist", - query=str(request.url), - ).model_dump(), - status_code=404, - ) + try: + cluster_df, config_df, _, _ = get_session_data(result_dir) + counts.get_cluster_counts_by_taxon( + cluster_df=cluster_df, + config_df=config_df, + base_output_dir=result_dir, + ) + del cluster_df, config_df + except Exception as e: + return JSONResponse( + content=ResponseSchema( + status="error", + message="Failed to generate counts by taxon", + error=str(e), + query=str(request.url), + ).model_dump(), + status_code=500, + ) result = parse_taxon_counts_file( filepath, @@ -487,7 +579,7 @@ async def get_cluster_summary( ) -> JSONResponse: try: result_dir = query_manager.get_session_dir(session_id) - config_f = os.path.join(result_dir, "config.json") + config_f = os.path.join(result_dir, "config.txt") if not os.path.exists(config_f): return JSONResponse( content=ResponseSchema( @@ -499,33 +591,54 @@ async def get_cluster_summary( status_code=428, ) - valid_endpoints = extract_attributes_and_taxon_sets(result_dir) - valid_attributes = valid_endpoints["attributes"] - - if attribute and attribute not in valid_attributes: + try: + cluster_df, config_df, attributes, _ = get_session_data(result_dir) + except Exception as e: return JSONResponse( content=ResponseSchema( status="error", - message=f"Invalid attribute: {attribute}. Must be one of {valid_attributes}.", - error="Invalid Input", + message="Failed to load session data", + error=str(e), + query=str(request.url), ).model_dump(), - status_code=400, + status_code=500, ) - filename = f"{attribute}/{attribute}.{CLUSTER_SUMMARY_FILENAME}" - filepath = os.path.join(result_dir, filename) - - if not os.path.exists(filepath): + if attribute and attribute not in attributes: return JSONResponse( content=ResponseSchema( status="error", - message=f"{COUNTS_FILEPATH} File Not Found", - error="File does not exist", - query=str(request.url), + message=f"Invalid attribute: {attribute}. Must be one of {attributes}.", + error="Invalid Input", ).model_dump(), - status_code=404, + status_code=400, ) + output_dir = os.path.join(result_dir, attribute) + filepath = os.path.join(output_dir, f"{attribute}.cluster_summary.txt") + + if not os.path.exists(filepath): + try: + summary_df = cluster_summary.get_cluster_summary( + cluster_df=cluster_df, + attribute=attribute, + config_df=config_df, + ) + + os.makedirs(output_dir, exist_ok=True) + summary_df.write_csv(filepath, separator="\t") + del summary_df, cluster_df, config_df + except Exception as e: + return JSONResponse( + content=ResponseSchema( + status="error", + message="Failed to generate cluster summary", + error=str(e), + query=str(request.url), + ).model_dump(), + status_code=500, + ) + result = parse_cluster_summary_file( filepath=filepath, include_clusters=include_clusters, @@ -565,7 +678,9 @@ async def get_cluster_summary( first_row = flattened_rows[0] buffer = io.StringIO() - writer = csv.DictWriter(buffer, fieldnames=first_row.keys(), delimiter="\t") + writer = csv.DictWriter( + buffer, fieldnames=first_row.keys(), delimiter="\t" + ) writer.writeheader() writer.writerows(flattened_rows) buffer.seek(0) @@ -676,7 +791,9 @@ async def get_valid_taxons_api( cluster_path = os.path.join(KINFIN_WORKDIR, cluster_info["path"]) taxon_idx_mapping_file = os.path.join(cluster_path, "taxon_idx_mapping.json") try: - check_file(taxon_idx_mapping_file, install_kinfin=True) + pass + # TODO + # check_file(taxon_idx_mapping_file, install_kinfin=True) except FileNotFoundError as e: return JSONResponse( content=ResponseSchema( @@ -816,7 +933,9 @@ async def get_column_descriptions_api( page: int = Query(1, ge=1), size: int = Query(40, ge=1, le=100), sort_by: str = Query(None, description="Comma-separated fields to sort by"), - sort_order: str = Query("asc", regex="^(asc|desc)$", description="Sort order: asc or desc"), + sort_order: str = Query( + "asc", regex="^(asc|desc)$", description="Sort order: asc or desc" + ), ): try: current_dir = os.path.dirname(os.path.abspath(__file__)) @@ -895,7 +1014,7 @@ async def get_attribute_summary( ): try: result_dir = query_manager.get_session_dir(session_id) - config_f = os.path.join(result_dir, "config.json") + config_f = os.path.join(result_dir, "config.txt") if not os.path.exists(config_f): return JSONResponse( content=ResponseSchema( @@ -907,33 +1026,54 @@ async def get_attribute_summary( status_code=428, ) - valid_endpoints = extract_attributes_and_taxon_sets(result_dir) - valid_attributes = valid_endpoints["attributes"] - - if attribute and attribute not in valid_attributes: + try: + cluster_df, config_df, attributes, _ = get_session_data(result_dir) + except Exception as e: return JSONResponse( content=ResponseSchema( status="error", - message=f"Invalid attribute: {attribute}. Must be one of {valid_attributes}.", - error="Invalid Input", + message="Failed to load session data", + error=str(e), + query=str(request.url), ).model_dump(), - status_code=400, + status_code=500, ) - filename = f"{attribute}/{attribute}.{ATTRIBUTE_METRICS_FILENAME}" - filepath = os.path.join(result_dir, filename) - - if not os.path.exists(filepath): + if attribute and attribute not in attributes: return JSONResponse( content=ResponseSchema( status="error", - message=f"{COUNTS_FILEPATH} File Not Found", - error="File does not exist", - query=str(request.url), + message=f"Invalid attribute: {attribute}. Must be one of {attributes}.", + error="Invalid Input", ).model_dump(), - status_code=404, + status_code=400, ) + output_dir = os.path.join(result_dir, attribute) + filepath = os.path.join(output_dir, f"{attribute}.attribute_metrics.txt") + + if not os.path.exists(filepath): + try: + attribute_df = attribute_metrics.get_attribute_metrics( + cluster_df=cluster_df, + config_df=config_df, + attribute=attribute, + ) + + os.makedirs(output_dir, exist_ok=True) + attribute_df.write_csv(filepath, separator="\t") + del attribute_df, cluster_df, config_df + except Exception as e: + return JSONResponse( + content=ResponseSchema( + status="error", + message="Failed to generate attribute metrics", + error=str(e), + query=str(request.url), + ).model_dump(), + status_code=500, + ) + result = parse_attribute_summary_file(filepath=filepath) if as_file: @@ -1015,7 +1155,7 @@ async def get_cluster_metrics( ): try: result_dir = query_manager.get_session_dir(session_id) - config_f = os.path.join(result_dir, "config.json") + config_f = os.path.join(result_dir, "config.txt") if not os.path.exists(config_f): return JSONResponse( content=ResponseSchema( @@ -1027,45 +1167,99 @@ async def get_cluster_metrics( status_code=428, ) - valid_endpoints = extract_attributes_and_taxon_sets(result_dir) - valid_attributes = valid_endpoints["attributes"] - - if attribute and attribute not in valid_attributes: + try: + cluster_df, config_df, attributes, taxon_label_dict = get_session_data( + result_dir + ) + except Exception as e: return JSONResponse( content=ResponseSchema( status="error", - message=f"Invalid attribute: {attribute}. Must be one of {valid_attributes}.", - error="Invalid Input", + message="Failed to load session data", + error=str(e), + query=str(request.url), ).model_dump(), - status_code=400, + status_code=500, ) - valid_taxon_sets = valid_endpoints["taxon_set"][attribute] - - if taxon_set and taxon_set not in valid_taxon_sets: + if attribute and attribute not in attributes: return JSONResponse( content=ResponseSchema( status="error", - message=f"Invalid taxon set: {taxon_set}. Must be one of {valid_taxon_sets}.", + message=f"Invalid attribute: {attribute}. Must be one of {attributes}.", error="Invalid Input", ).model_dump(), status_code=400, ) - filename = f"{attribute}/{attribute}.{taxon_set}.{CLUSTER_METRICS_FILENAME}" - filepath = os.path.join(result_dir, filename) + unique_label_values = { + col: config_df.select(pl.col(col)).unique().to_series().to_list() + for col in config_df.columns + if col not in ("#IDX", "TAXON", "TAXID", "OUT") + } + unique_label_values["all"] = ["all"] + unique_label_values["TAXON"] = config_df["TAXON"].to_list() - if not os.path.exists(filepath): + if taxon_set and taxon_set not in unique_label_values[attribute]: return JSONResponse( content=ResponseSchema( status="error", - message=f"{CLUSTER_METRICS_FILENAME} File Not Found", - error="File does not exist", - query=str(request.url), + message=f"Invalid taxon set: {taxon_set}. Must be one of {unique_label_values[attribute]}.", + error="Invalid Input", ).model_dump(), - status_code=404, + status_code=400, ) + output_dir = os.path.join(result_dir, attribute) + filepath = os.path.join( + output_dir, f"{attribute}.{taxon_set}.cluster_metrics.txt" + ) + + if not os.path.exists(filepath): + try: + label_to_taxons = { + attr: { + label: { + taxon + for taxon, labels in taxon_label_dict.items() + if labels.get(attr) == label + } + for label in unique_label_values[attr] + } + for attr in attributes + } + + base_metrics_df = ( + cluster_df.pipe( + cluster_metrics.add_status_and_TAXON_protein_count_columns, + label_to_taxons, + ) + .rename({"TAXON_count": "cluster_proteome_count"}) + .pipe(cluster_metrics.add_all_taxon_info_columns, label_to_taxons) + ) + + metrics_df = cluster_metrics.get_cluster_metrics( + attribute=attribute, + label_group=taxon_set, + label_to_taxons=label_to_taxons, + base_metrics_df=base_metrics_df, + ) + + # Write to file + os.makedirs(output_dir, exist_ok=True) + metrics_df.write_csv(filepath, separator="\t") + del metrics_df, base_metrics_df, cluster_df, config_df + except Exception as e: + return JSONResponse( + content=ResponseSchema( + status="error", + message="Failed to generate cluster metrics", + error=str(e), + query=str(request.url), + ).model_dump(), + status_code=500, + ) + result = parse_cluster_metrics_file(filepath, cluster_status, cluster_type) if as_file: @@ -1086,7 +1280,9 @@ async def get_cluster_metrics( first_row = flattened_rows[0] if flattened_rows else {} buffer = io.StringIO() - writer = csv.DictWriter(buffer, fieldnames=first_row.keys(), delimiter="\t") + writer = csv.DictWriter( + buffer, fieldnames=first_row.keys(), delimiter="\t" + ) writer.writeheader() writer.writerows(flattened_rows) buffer.seek(0) @@ -1159,7 +1355,7 @@ async def get_pairwise_analysis( ): try: result_dir = query_manager.get_session_dir(session_id) - config_f = os.path.join(result_dir, "config.json") + config_f = os.path.join(result_dir, "config.txt") if not os.path.exists(config_f): return JSONResponse( content=ResponseSchema( @@ -1171,14 +1367,26 @@ async def get_pairwise_analysis( status_code=428, ) - valid_endpoints = extract_attributes_and_taxon_sets(result_dir) - valid_attributes = valid_endpoints["attributes"] + try: + cluster_df, config_df, attributes, taxon_label_dict = get_session_data( + result_dir + ) + except Exception as e: + return JSONResponse( + content=ResponseSchema( + status="error", + message="Failed to load session data", + error=str(e), + query=str(request.url), + ).model_dump(), + status_code=500, + ) - if attribute and attribute not in valid_attributes: + if attribute and attribute not in attributes: return JSONResponse( content=ResponseSchema( status="error", - message=f"Invalid attribute: {attribute}. Must be one of {valid_attributes}.", + message=f"Invalid attribute: {attribute}. Must be one of {attributes}.", error="Invalid Input", ).model_dump(), status_code=400, @@ -1188,15 +1396,45 @@ async def get_pairwise_analysis( filepath = os.path.join(result_dir, filename) if not os.path.exists(filepath): - return JSONResponse( - content=ResponseSchema( - status="error", - message=f"{PAIRWISE_ANALYSIS_FILE} File Not Found", - error="File does not exist", - query=str(request.url), - ).model_dump(), - status_code=404, - ) + try: + unique_label_values = { + col: config_df.select(pl.col(col)).unique().to_series().to_list() + for col in config_df.columns + if col not in ("#IDX", "TAXON", "TAXID", "OUT") + } + unique_label_values["all"] = ["all"] + unique_label_values["TAXON"] = config_df["TAXON"].to_list() + + label_to_taxons = { + attr: { + label: { + taxon + for taxon, labels in taxon_label_dict.items() + if labels.get(attr) == label + } + for label in unique_label_values[attr] + } + for attr in attributes + } + + # Generate pairwise analysis + cluster_metrics.generate_pairwise_representation_test( + cluster_df=cluster_df, + attribute=attribute, + label_to_taxons=label_to_taxons, + output_dir=result_dir, + ) + + except Exception as e: + return JSONResponse( + content=ResponseSchema( + status="error", + message="Failed to generate pairwise analysis", + error=str(e), + query=str(request.url), + ).model_dump(), + status_code=500, + ) result = parse_pairwise_file(filepath, taxon_1, taxon_2) @@ -1269,36 +1507,49 @@ async def get_plot( ) result_dir = query_manager.get_session_dir(session_id) - filepath: str = "" - match plot_type: - case "cluster-size-distribution": - filepath = "cluster_size_distribution.png" - case "all-rarefaction-curve": - filepath = "all/all.rarefaction_curve.png" - case _: - return JSONResponse( - content=ResponseSchema( - status="error", - message="Invalid Plot Type", - error="invalid_plot_type", - query=str(request.url), - ).model_dump(), - status_code=404, - ) - - filepath = os.path.join(result_dir, filepath) - if not os.path.exists(filepath): + if plot_type == "cluster-size-distribution": + filepath = os.path.join(result_dir, "cluster_size_distribution.png") + elif plot_type == "all-rarefaction-curve": + filepath = os.path.join(result_dir, "all", "all.rarefaction_curve.png") + else: return JSONResponse( content=ResponseSchema( status="error", - message="Plot not found", - error="plot_not_found", + message=f"Unknown plot type: {plot_type}", + error="invalid_plot_type", query=str(request.url), ).model_dump(), - status_code=404, + status_code=400, ) + if not os.path.exists(filepath): + try: + cluster_df, config_df, attributes, _ = get_session_data(result_dir) + + if plot_type == "cluster-size-distribution": + plots.plot_cluster_sizes( + cluster_df=cluster_df, base_output_dir=result_dir + ) + elif plot_type == "all-rarefaction-curve": + plots.generate_kinfin_rarefaction_plots( + cluster_df=cluster_df, + config_df=config_df, + attributes=attributes, + base_output_dir=result_dir, + ) + + except Exception as e: + return JSONResponse( + content=ResponseSchema( + status="error", + message="Failed to generate plot", + error=str(e), + query=str(request.url), + ).model_dump(), + status_code=500, + ) + return FileResponse( filepath, media_type="image/png", diff --git a/src/api/sessions.py b/src/api/sessions.py index 66b5d8b..bebd322 100644 --- a/src/api/sessions.py +++ b/src/api/sessions.py @@ -10,6 +10,8 @@ from datetime import datetime, timedelta from typing import Dict, List, Optional, Tuple +import polars as pl + logger = logging.getLogger("kinfin_logger") @@ -63,8 +65,14 @@ def get_or_create_session(self, query: List[Dict[str, str]]) -> Tuple[str, str]: session_id = self.get_session_id(query) session_dir = os.path.join(self.results_base_dir, session_id) + config_file_path = os.path.join(session_dir, "config.txt") if not os.path.exists(session_dir): os.makedirs(session_dir) + + df = pl.DataFrame(query) + if "taxon" in df.columns: + df = df.rename({"taxon": "TAXON"}) + df.write_csv(config_file_path) else: os.utime(session_dir, None) diff --git a/src/api/utils.py b/src/api/utils.py index ffa4e55..6b18473 100644 --- a/src/api/utils.py +++ b/src/api/utils.py @@ -1,9 +1,12 @@ import asyncio -import glob import json -from collections import defaultdict +import os from typing import Any, Dict +import polars as pl + +from internal import parsers + def read_status(status_file): status_info = {} @@ -81,18 +84,33 @@ async def run_cli_command(command: list, status_file: str): return None -def extract_attributes_and_taxon_sets(filepath: str): - files = glob.glob(f"{filepath}/**/*.cluster_metrics.txt") - files = [file.split(filepath)[1] for file in files] - attributes = set() - result = {"attributes": [], "taxon_set": defaultdict(list)} - for file in files: - filename = file.split("/")[-1] - attribute = filename.split(".")[0] - taxon_set = filename.split(".")[1] - attributes.add(attribute) - result["taxon_set"][attribute].append(taxon_set) - result["attributes"] = sorted(attributes) +def extract_attributes_and_taxon_sets(session_dir: str): + """ + Extract attributes and taxon sets directly from the session's config file, + rather than relying on pre-generated *.cluster_metrics.txt files. + """ + config_file = os.path.join(session_dir, "config.txt") + if not os.path.exists(config_file): + raise FileNotFoundError(f"Config file not found: {config_file}") + + nodesdb_f = os.environ.get("NODESDB_F") + ndb_f = os.environ.get("NDB_F") + if not nodesdb_f or not ndb_f: + raise RuntimeError("NODESDB_F and NDB_F environment variables must be set") + + nodesdb = parsers.nodesdb(filepath=nodesdb_f, outpath=ndb_f) + config_df, attributes = parsers.configfile(config_file, nodesdb) + + label_columns = [ + col for col in attributes if col not in ("#IDX", "TAXON", "TAXID", "OUT") + ] + + result = { + "attributes": label_columns, + "taxon_set": { + label: sorted(config_df[label].unique()) for label in label_columns + }, + } return result diff --git a/src/cli/__init__.py b/src/cli/__init__.py deleted file mode 100644 index 5b52ebe..0000000 --- a/src/cli/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -import os - -from core.input import InputData -from core.logger import setup_logger -from core.results import analyse - - -def run_cli(args: InputData) -> None: - """ - Run the command-line interface to perform analysis based on the provided input data. - - Args: - args (InputData): An instance of InputData containing input parameters and data. - - Returns: - None - """ - log_path = os.path.join(args.output_path, "kinfin.log") - setup_logger(log_path) - analyse(args) diff --git a/src/cli/commands.py b/src/cli/commands.py deleted file mode 100644 index 6a8d176..0000000 --- a/src/cli/commands.py +++ /dev/null @@ -1,212 +0,0 @@ -import argparse -import sys -from typing import Union - -from cli.validate import validate_cli_args -from core.config import SUPPORTED_PLOT_FORMATS, SUPPORTED_TAXRANKS, SUPPORTED_TESTS -from core.input import InputData, ServeArgs - - -# TODO : --plotsize should take a tuple -# TODO : --taxranks should take multiple inputs -def parse_args( - nodesdb_f: str, - pfam_mapping_f: str, - ipr_mapping_f: str, - go_mapping_f: str, -) -> Union[ServeArgs, InputData]: - """Parse command-line arguments. - - Args: - nodesdb_f (str): filepath of nodesdb_f. - pfam_mapping_f (str): filepath of pfam_mapping_f. - ipr_mapping_f (str): filepath of ipr_mapping_f. - go_mapping_f (str): filepath of go_mapping_f. - - Returns: - ServeArgs or InputData: Parsed arguments based on the command. - - Raises: - SystemExit: If an invalid command is provided. - """ - - parser = argparse.ArgumentParser( - description="Kinfin proteome cluster analysis tool" - ) - - subparsers = parser.add_subparsers(title="command", required=True, dest="command") - api_parser = subparsers.add_parser("serve", help="Start the server") - api_parser.add_argument( - "-p", - "--port", - type=int, - default=8000, - help="Port number for the server (default: 8000)", - ) - - cli_parser = subparsers.add_parser("analyse", help="Perform analysis") - - # Required Arguments - required_group = cli_parser.add_argument_group("Required Arguments") - required_group.add_argument( - "-g", - "--cluster_file", - help="OrthologousGroups.txt produced by OrthoFinder", - required=True, - ) - required_group.add_argument( - "-c", "--config_file", help="Config file (in CSV format)", required=True - ) - required_group.add_argument( - "-s", - "--sequence_ids_file", - help="SequenceIDs.txt used in OrthoFinder", - required=True, - ) - - # Other Files - other_files_group = cli_parser.add_argument_group("Other Files") - other_files_group.add_argument( - "-p", "--species_ids_file", help="SpeciesIDs.txt used in OrthoFinder" - ) - other_files_group.add_argument( - "-m", "--taxon_idx_mapping", help="TAXON IDX Mapping File" - ) - other_files_group.add_argument( - "-f", - "--functional_annotation", - help="Mapping of ProteinIDs to GO/IPRS/SignalP/Pfam (can be generated through 'iprs_to_table.py')", - ) - other_files_group.add_argument("-a", "--fasta_dir", help="Directory of FASTA files") - other_files_group.add_argument( - "-t", - "--tree_file", - help="Tree file in Newick format (taxon names must be the same as TAXON in config file)", - ) - - # General Options - general_group = cli_parser.add_argument_group("General Options") - general_group.add_argument("-o", "--output_path", help="Output prefix") - general_group.add_argument( - "--infer_singletons", - help="Absence of proteins in clustering is interpreted as singleton (based on SequenceIDs.txt)", - action="store_true", - ) - general_group.add_argument( - "--plot_tree", - help="Plot PDF of annotated phylogenetic tree (requires -t, full ETE3 installation and X-server/xvfb-run)", - action="store_true", - ) - general_group.add_argument( - "--min_proteomes", - help="Required number of proteomes in a taxon-set to be used in rarefaction/representation-test computations [default: 2]", - default=2, - type=int, - ) - general_group.add_argument( - "--test", - help="Test to be used in representation-test computations [default: mannwhitneyu]. Options: ttest, welch, mannwhitneyu, ks, kruskal", - default="mannwhitneyu", - choices=SUPPORTED_TESTS, - ) - general_group.add_argument( - "-r", - "--taxranks", - help="Taxonomic ranks to be inferred from TaxIDs in config file [default: phylum,order,genus]", - # TODO : Add SUPPORTED_TAXRANKS here - default=["phylum", "order", "genus"], - nargs="+", - choices=SUPPORTED_TAXRANKS, - ) - general_group.add_argument( - "--repetitions", - help="Number of repetitions for rarefaction curves [default: 30]", - default=30, - type=int, - ) - - # Fuzzy Orthology Groups - fuzzy_group = cli_parser.add_argument_group("Fuzzy Orthology Groups") - fuzzy_group.add_argument( - "-n", - "--target_count", - help="Target number of copies per proteome [default: 1]", - default=1, - type=int, - ) - fuzzy_group.add_argument( - "-x", - "--target_fraction", - help="Min proportion of proteomes at target_count [default: 0.75]", - default=0.75, - type=float, - ) - fuzzy_group.add_argument( - "--min", - help="Min count of proteins for proteomes outside of target_fraction [default: 0]", - default=0, - type=int, - ) - fuzzy_group.add_argument( - "--max", - help="Max count of proteins for proteomes outside of target_fraction [default: 20]", - default=20, - type=int, - ) - - plotting_group = cli_parser.add_argument_group("Plotting Options") - plotting_group.add_argument( - "--fontsize", help="Fontsize for plots [default: 18]", default=18, type=int - ) - plotting_group.add_argument( - "--plotsize", - help="Size (WIDTH,HEIGHT) for plots [default: 24,12]", - default=(24, 12), - nargs=2, - ) - plotting_group.add_argument( - "--plot_format", - help="Plot formats [default: pdf]", - default="pdf", - choices=SUPPORTED_PLOT_FORMATS, - ) - - args = parser.parse_args() - - if args.command == "serve": - return ServeArgs(port=args.port) - elif args.command == "analyse": - validate_cli_args(args=args) - fuzzy_range = { - x for x in range(args.min, args.max + 1) if x != args.target_count - } - - return InputData( - cluster_file=args.cluster_file, - config_f=args.config_file, - sequence_ids_file=args.sequence_ids_file, - species_ids_file=args.species_ids_file, - functional_annotation_f=args.functional_annotation, - fasta_dir=args.fasta_dir, - tree_file=args.tree_file, - output_path=args.output_path, - infer_singletons=args.infer_singletons, - plot_tree=args.plot_tree, - min_proteomes=args.min_proteomes, - test=args.test, - taxranks=args.taxranks, - repetitions=args.repetitions + 1, - fuzzy_count=args.target_count, - fuzzy_fraction=args.target_fraction, - fuzzy_range=fuzzy_range, - fontsize=args.fontsize, - plotsize=args.plotsize, - plot_format=args.plot_format, - nodesdb_f=nodesdb_f, - pfam_mapping_f=pfam_mapping_f, - ipr_mapping_f=ipr_mapping_f, - go_mapping_f=go_mapping_f, - taxon_idx_mapping_file=args.taxon_idx_mapping, - ) - else: - sys.exit() diff --git a/src/cli/validate.py b/src/cli/validate.py deleted file mode 100644 index cc9bbe0..0000000 --- a/src/cli/validate.py +++ /dev/null @@ -1,83 +0,0 @@ -import logging -import sys - -from core.utils import check_file - -logger = logging.getLogger("kinfin_logger") - - -def validate_cli_args(args) -> None: - """Validate cli input arguments. - - This function checks if all required files exist and if the arguments meet specific conditions. - - Args: - args (InputData): Input arguments as a named tuple. - - Raises: - SystemExit: If there are any validation errors, exits the program with error messages. - """ - - error_msgs = [] - - try: - check_file(args.cluster_file) - except FileNotFoundError as e: - error_msgs.append(str(e)) - try: - if not isinstance(args.config_file, str): - raise ValueError("[ERROR] - Invalid config file data") - - check_file(args.config_file) - except (FileNotFoundError, ValueError) as e: - error_msgs.append(str(e)) - try: - check_file(args.sequence_ids_file) - except FileNotFoundError as e: - error_msgs.append(str(e)) - try: - check_file(args.species_ids_file) - except FileNotFoundError as e: - error_msgs.append(str(e)) - try: - check_file(args.tree_file) - except FileNotFoundError as e: - error_msgs.append(str(e)) - try: - check_file(args.functional_annotation) - except FileNotFoundError as e: - error_msgs.append(str(e)) - - if args.fasta_dir and not args.species_ids_file: - error_msgs.append( - "[ERROR] : You have provided a FASTA-dir using '--fasta-dir'. Please also provide a Species-ID file using ('--species_ids_file')." - ) - - if args.target_count < 0: - error_msgs.append( - f"[ERROR] : --target_count {args.target_count} must be greater than 0" - ) - - if args.target_fraction < 0 or args.target_fraction > 1: - error_msgs.append( - f"[ERROR] : --target_fraction {args.target_fraction} is not between 0.0 and 1.0" - ) - - if args.min > args.max: - error_msgs.append( - f"[ERROR] : --min {args.min} is greater than --max {args.max}" - ) - - if args.repetitions <= 0: - error_msgs.append( - "[ERROR] : Please specify a positive integer for the number of repetitions for the rarefaction curves" - ) - - if args.min_proteomes <= 0: - error_msgs.append( - "[ERROR] : Please specify a positive integer for the minimum number of proteomes to consider for computations" - ) - - if error_msgs: - logger.error("\n".join(error_msgs)) - sys.exit(1) diff --git a/src/core/__init__.py b/src/core/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/core/alo.py b/src/core/alo.py deleted file mode 100644 index 378c83c..0000000 --- a/src/core/alo.py +++ /dev/null @@ -1,250 +0,0 @@ -from typing import Dict, List, Literal, Optional, Set, Union - -from core.clusters import Cluster - - -class AttributeLevel: - """ - Definitions: - 'shared' : shared between one ALO and others - 'singleton' : cardinality of 1 ('specific', but separate) - 'specific' : only present within one ALO - """ - - def __init__(self, attribute: str, level: str, proteomes: Set[str]) -> None: - self.attribute: str = attribute - self.level: str = level - self.proteomes: Set[str] = set(proteomes) - self.proteomes_list: List[str] = list(proteomes) - self.proteome_count: int = len(proteomes) - - self.cluster_ids_by_cluster_type_by_cluster_status: Dict[ - str, Dict[str, List[str]] - ] = { - # sums up to cluster_count - "present": {"singleton": [], "specific": [], "shared": []}, - "absent": {"singleton": [], "specific": [], "shared": []}, - } - - self.protein_ids_by_cluster_type: Dict[str, List[str]] = { - # list of lists - "singleton": [], - "specific": [], - "shared": [], - } - - self.protein_span_by_cluster_type: Dict[str, List[Union[int, float]]] = { - "singleton": [], - "specific": [], - "shared": [], - } - - self.clusters_by_cluster_cardinality_by_cluster_type: Dict[ - str, Dict[str, List[str]] - ] = { - "shared": {"true": [], "fuzzy": []}, - "specific": {"true": [], "fuzzy": []}, - } - - self.cluster_status_by_cluster_id: Dict[str, Literal["absent", "present"]] = {} - self.cluster_type_by_cluster_id: Dict[ - str, Literal["singleton", "shared", "specific"] - ] = {} - - self.cluster_mwu_pvalue_by_cluster_id = {} - self.cluster_mwu_log2_mean_by_cluster_id = {} - self.cluster_mean_ALO_count_by_cluster_id = {} - self.cluster_mean_non_ALO_count_by_cluster_id = {} - - self.domain_counter_by_domain_source_by_cluster_type = None - self.protein_with_domain_count_by_domain_source_by_cluster_type = None - - self.protein_length_stats_by_cluster_id: Dict[ - str, Dict[str, Union[int, float]] - ] = {} - self.protein_count_by_cluster_id: Dict[str, int] = {} - - def add_cluster( - self, - cluster: Cluster, - attribute_cluster_type: Literal["singleton", "shared", "specific"], - ALO_cluster_status: Literal["absent", "present"], - ALO_protein_length_stats: Dict[str, Union[int, float]], - ALO_protein_ids_in_cluster: List[str], - ALO_cluster_cardinality: Optional[str], - mwu_pvalue: Optional[float], - mwu_log2_mean: Optional[float], - mean_ALO_count: Optional[float], - mean_non_ALO_count: Optional[float], - ) -> None: - """ - Adds a cluster to various data structures maintained by the class. - - Args: - cluster (Cluster): The cluster object to add. - attribute_cluster_type (Literal["singleton", "shared", "specific"]): - Type of the cluster as either 'singleton', 'shared', or 'specific'. - ALO_cluster_status (Literal["absent", "present"]): - Status of the cluster, either 'absent' or 'present'. - ALO_protein_length_stats (Dict[str, Union[int, float]]): - Length statistics of proteins in the cluster. - ALO_protein_ids_in_cluster (List[str]): - List of protein IDs present in the cluster. - ALO_cluster_cardinality (Optional[str]): - Cardinality of the cluster (if applicable). - mwu_pvalue (Optional[float]): - P-value from Mann-Whitney U test (if applicable). - mwu_log2_mean (Optional[float]): - Log2 transformed mean (if applicable). - mean_ALO_count (Optional[float]): - Mean count of ALO (if applicable). - mean_non_ALO_count (Optional[float]): - Mean count of non-ALO (if applicable). - - Returns: - None - """ - self.cluster_ids_by_cluster_type_by_cluster_status[ALO_cluster_status][ - attribute_cluster_type - ].append(cluster.cluster_id) - self.cluster_status_by_cluster_id[cluster.cluster_id] = ALO_cluster_status - self.cluster_type_by_cluster_id[cluster.cluster_id] = attribute_cluster_type - self.protein_length_stats_by_cluster_id[cluster.cluster_id] = ( - ALO_protein_length_stats - ) - - self.protein_count_by_cluster_id[cluster.cluster_id] = len( - ALO_protein_ids_in_cluster - ) - - if ALO_cluster_status == "present": - for ALO_protein_id in ALO_protein_ids_in_cluster: - self.protein_ids_by_cluster_type[attribute_cluster_type].append( - ALO_protein_id - ) - self.protein_span_by_cluster_type[attribute_cluster_type].append( - ALO_protein_length_stats["sum"] - ) - if attribute_cluster_type != "singleton" and ALO_cluster_cardinality: - self.clusters_by_cluster_cardinality_by_cluster_type[ - attribute_cluster_type - ][ALO_cluster_cardinality].append(cluster.cluster_id) - - self.cluster_mwu_pvalue_by_cluster_id[cluster.cluster_id] = mwu_pvalue - self.cluster_mwu_log2_mean_by_cluster_id[cluster.cluster_id] = mwu_log2_mean - self.cluster_mean_ALO_count_by_cluster_id[cluster.cluster_id] = mean_ALO_count - self.cluster_mean_non_ALO_count_by_cluster_id[cluster.cluster_id] = ( - mean_non_ALO_count - ) - - def get_protein_count_by_cluster_type(self, cluster_type: str) -> int: - """ - Return the count of proteins for a specific cluster type. - - Args: - cluster_type (str): Type of the cluster. Use "total" for the total count across all types. - - Returns: - int: Number of proteins in the specified cluster type. - - Raises: - KeyError: If 'cluster_type' is not found in self.protein_ids_by_cluster_type. - """ - if cluster_type == "total": - return sum( - len(protein_ids) - for _, protein_ids in list(self.protein_ids_by_cluster_type.items()) - ) - else: - return len(self.protein_ids_by_cluster_type[cluster_type]) - - def get_cluster_count_by_cluster_status_by_cluster_type( - self, - cluster_status: str, - cluster_type: str, - ) -> int: - """ - Get the count of clusters of a specific status and type. - - Args: - cluster_status (str): The status of clusters to count. - cluster_type (str): The type of cluster to count. Use "total" to get - the total count across all cluster types for the given status. - - Returns: - int: Number of clusters with the specified status and type. - - Raises: - KeyError: If 'cluster_status' or 'cluster_type' is not found in - self.cluster_ids_by_cluster_type_by_cluster_status. - """ - if cluster_type == "total": - return sum( - len(cluster_ids) - for _, cluster_ids in list( - self.cluster_ids_by_cluster_type_by_cluster_status[ - cluster_status - ].items() - ) - ) - else: - return len( - self.cluster_ids_by_cluster_type_by_cluster_status[cluster_status][ - cluster_type - ] - ) - - def get_protein_span_by_cluster_type(self, cluster_type: str) -> Union[int, float]: - """ - Get the total span of proteins for a specific cluster type. - - Args: - cluster_type (str): The type of cluster for which to retrieve protein span. - Use "total" to get the total span across all cluster types. - - Returns: - Union[int, float]: Total span of proteins in the specified cluster type. - If 'cluster_type' is "total", returns the sum of spans across all - cluster types. - """ - return ( - sum( - sum(protein_ids) - for _, protein_ids in list(self.protein_span_by_cluster_type.items()) - ) - if cluster_type == "total" - else sum(self.protein_span_by_cluster_type[cluster_type]) - ) - - def get_cluster_count_by_cluster_cardinality_by_cluster_type( - self, - cluster_type: str, - cluster_cardinality: str, - ) -> int: - """ - Return the count of clusters of a specific type and cardinality. - - Args: - cluster_type (str): Type of the cluster. - cluster_cardinality (str): Cardinality of the clusters. - - Returns: - int: Number of clusters with the specified type and cardinality. - - Raises: - KeyError: If 'cluster_type' or 'cluster_cardinality' is not found. - """ - return len( - self.clusters_by_cluster_cardinality_by_cluster_type[cluster_type][ - cluster_cardinality - ] - ) - - def get_proteomes(self) -> str: - """ - Get a sorted string representation of proteome IDs. - - Returns: - str: Comma-separated and sorted list of proteome IDs. - """ - return ", ".join(sorted([str(proteome_id) for proteome_id in self.proteomes])) diff --git a/src/core/alo_collections.py b/src/core/alo_collections.py deleted file mode 100644 index f459696..0000000 --- a/src/core/alo_collections.py +++ /dev/null @@ -1,526 +0,0 @@ -import logging -import os -import random -from typing import Any, Dict, List, Optional, Set - -import ete3 -import matplotlib as mat -import matplotlib.pyplot as plt -import numpy as np -from ete3 import Tree - -from core.alo import AttributeLevel -from core.config import ATTRIBUTE_RESERVED - -logger = logging.getLogger("kinfin_logger") - -mat.use("agg") - - -plt.style.use("ggplot") -mat.rc("ytick", labelsize=20) -mat.rc("xtick", labelsize=20) -axis_font = {"size": "20"} -mat.rcParams.update({"font.size": 22}) - - -class AloCollection: - def __init__( - self, - proteomes: Set[str], - attributes: List[str], - proteome_id_by_species_id: Dict[str, str], - level_by_attribute_by_proteome_id: Dict[str, Dict[str, str]], - node_idx_by_proteome_ids: Optional[Dict[Any, Any]], - tree_ete: Optional[Tree], - ) -> None: - self.proteomes = proteomes - self.attributes_verbose = attributes - self.attributes = [ - # list of attributes - attribute - for attribute in attributes - if attribute not in ATTRIBUTE_RESERVED - ] - self.proteome_id_by_species_id = proteome_id_by_species_id - self.level_by_attribute_by_proteome_id = level_by_attribute_by_proteome_id - self.node_idx_by_proteome_ids = node_idx_by_proteome_ids - self.tree_ete = tree_ete - self.proteome_ids_by_level_by_attribute = ( - self.compute_proteomes_by_level_by_attribute() - ) - self.fastas_parsed: bool = False - self.ALO_by_level_by_attribute = self.create_ALOs() - - def compute_proteomes_by_level_by_attribute( - self, - ) -> Dict[str, Dict[str, Set[str]]]: - """ - Compute proteomes grouped by levels for each attribute. - - Args: - attributes (List[str]): A list of strings representing attributes. - level_by_attribute_by_proteome_id (Dict[str, Dict[str, str]]): A dictionary where keys - are proteome IDs (strings), and values are dictionaries with keys representing - attributes (strings) and values representing levels (strings). - - Returns: - Dict[str, Dict[str, Set[str]]]: A dictionary where keys are attributes (strings), - and values are dictionaries. The inner dictionaries have keys representing - levels (strings) and values representing sets of proteome IDs (strings). - """ - proteomes_by_level_by_attribute: Dict[str, Dict[str, Set[str]]] = { - attribute: {} for attribute in self.attributes - } - for proteome_id in self.level_by_attribute_by_proteome_id: - for attribute in self.attributes: - level = self.level_by_attribute_by_proteome_id[proteome_id][attribute] - if level not in proteomes_by_level_by_attribute[attribute]: - proteomes_by_level_by_attribute[attribute][level] = set() - proteomes_by_level_by_attribute[attribute][level].add(proteome_id) - return proteomes_by_level_by_attribute - - def create_ALOs(self) -> Dict[str, Dict[str, Optional[AttributeLevel]]]: - """ - Creates Attribute Level Objects (ALOs) for each attribute and level based on - proteome IDs. - - Returns: - Dict[str, Dict[str, Optional[AttributeLevel]]]: - A dictionary where each key is an attribute name (str), - and the corresponding value is a dictionary mapping level (str) - to an AttributeLevel instance or None. - """ - ALO_by_level_by_attribute: Dict[str, Dict[str, Optional[AttributeLevel]]] = { - attribute: {} for attribute in self.attributes - } - for attribute in self.proteome_ids_by_level_by_attribute: - for level in self.proteome_ids_by_level_by_attribute[attribute]: - proteome_ids = self.proteome_ids_by_level_by_attribute[attribute][level] - ALO = AttributeLevel( - # - attribute=attribute, - level=level, - proteomes=proteome_ids, - ) - if level not in ALO_by_level_by_attribute[attribute]: - ALO_by_level_by_attribute[attribute][level] = None - ALO_by_level_by_attribute[attribute][level] = ALO - return ALO_by_level_by_attribute - - def generate_header_for_node(self, node: ete3.TreeNode, dirs: Dict[str, str]): - """ - Generates a header image for a given node of a tree with specified statistics. - - Args: - node (ete3.TreeNode): The TreeNode object representing the node for which the header is generated. - dirs (Dict[str, str]): A dictionary containing directory paths, including 'tree_headers' where the header image will be saved. - - Returns: - str: File path to the generated header image. - - Notes: - - The method generates a header image in PNG format displaying various statistics (apomorphies and synapomorphies) for the given tree node. - - The statistics include counts of singletons, non-singletons, complete presence synapomorphies, and partial absence synapomorphies. - - The generated image is saved in the specified directory under 'tree_headers' with the node's name as the filename. - - Raises: - Any exceptions that might occur during file saving or table rendering. - """ - - node_header_f = os.path.join(dirs["tree_headers"], f"{node.name}.header.png") - data = [ - ( - "Apomorphies (size=1)", - "{:,}".format( - node.apomorphic_cluster_counts["singletons"] # type:ignore - ), - ), - ( - "Apomorphies (size>1)", - "{:,}".format( - node.apomorphic_cluster_counts["non_singletons"] # type:ignore - ), - ), - ( - "Synapomorphies (all)", - "{:,}".format( - node.synapomorphic_cluster_counts[ # type:ignore - "complete_presence" - ] - + node.synapomorphic_cluster_counts[ # type:ignore - "partial_absence" - ] - ), - ), - ( - "Synapomorphies (cov=100%)", - "{:,}".format( - node.synapomorphic_cluster_counts[ # type:ignore - "complete_presence" - ] - ), - ), - ( - "Synapomorphies (cov<100%)", - "{:,}".format( - node.synapomorphic_cluster_counts[ # type: ignore - "partial_absence" - ] # type:ignore - ), - ), - ] - col_labels = ("Type", "Count") - fig, ax = plt.subplots(figsize=(2, 0.5)) - ax.set_facecolor("white") - table = ax.table( - cellText=data, - colLabels=col_labels, - loc="bottom", - fontsize=24, - colLoc="center", - rowLoc="right", - edges="", - ) - table.set_fontsize(24) - table.scale(2, 1) - for key, cell in list(table.get_celld().items()): - row, col = key - cell._text.set_color("grey") # type:ignore - cell.set_edgecolor("darkgrey") - cell.visible_edges = "T" if row > 0 else "B" - if row == len(data) - 2: - cell.set_edgecolor("darkgrey") - cell.visible_edges = "T" - ax.axis("tight") - ax.axis("off") - logger.info(f"[STATUS]\t- Plotting {node_header_f}") - fig.savefig(node_header_f, pad=0, bbox_inches="tight", format="png") - plt.close() - return node_header_f - - def generate_chart_for_node( - self, - node, - dirs: Dict[str, str], - plot_format: str, - fontsize: int, - ) -> Optional[str]: - """ - Generate and save a histogram chart for a given node's synapomorphies. - - Args: - - node: The node object containing synapomorphic cluster strings. - - dirs: A dictionary containing directory paths, specifically 'tree_charts' for saving charts. - - plot_format: The format in which to save the chart ('png' or 'pdf'). - - fontsize: Font size for axis labels and ticks. - - Returns: - - Optional[str]: Path to the saved chart file if successful, None otherwise. - """ - - if proteome_coverages := [ - float(synapomorphic_cluster_string[3]) - for synapomorphic_cluster_string in node.synapomorphic_cluster_strings - ]: - chart_f = os.path.join(dirs["tree_charts"], f"{node.name}.barchart.png") - f, ax = plt.subplots(figsize=(3.0, 3.0)) - ax.set_facecolor("white") - x_values = np.array(proteome_coverages) - ax.hist( - x_values, - histtype="stepfilled", - align="mid", - bins=np.arange(0.0, 1.0 + 0.1, 0.1), - ) - ax.set_xlim(-0.1, 1.1) - for tick in ax.xaxis.get_majorticklabels(): - tick.set_fontsize(fontsize - 2) - tick.set_rotation("vertical") - for tick in ax.yaxis.get_majorticklabels(): - tick.set_fontsize(fontsize - 2) - ax.set_frame_on(False) - ax.xaxis.grid(True, linewidth=1, which="major", color="lightgrey") - ax.yaxis.grid(True, linewidth=1, which="major", color="lightgrey") - f.suptitle("Synapomorphies", y=1.1) - ax.set_ylabel("Count", fontsize=fontsize) - ax.set_xlabel("Proteome coverage", fontsize=fontsize) - logger.info(f"[STATUS]\t- Plotting {chart_f}") - f.savefig(chart_f, bbox_inches="tight", format="png") - if plot_format == "pdf": - pdf_chart_f = os.path.join( - dirs["tree_charts"], - f"{node.name}.barchart.pdf", - ) - logger.info(f"[STATUS]\t- Plotting {pdf_chart_f}") - f.savefig(pdf_chart_f, bbox_inches="tight", format="pdf") - plt.close() - return chart_f - - def plot_text_tree(self, dirs: Dict[str, str]) -> None: - """ - Plot and save the textual representation of the tree. - - This method uses the `tree_ete` attribute of the class to generate and save - both a Newick format (.nwk) and a text format (.txt) representation of the tree. - - Args: - - dirs: A dictionary containing directory paths, specifically 'tree' for saving tree files. - - Returns: - - None - """ - if self.tree_ete: - tree_nwk_f = os.path.join(dirs["tree"], "tree.nwk") - self.tree_ete.write(format=1, outfile=tree_nwk_f) - tree_txt_f = os.path.join(dirs["tree"], "tree.txt") - with open(tree_txt_f, "w") as tree_txt_fh: - tree_txt_fh.write( - f"{self.tree_ete.get_ascii(show_internal=True, compact=False)}\n" - ) - - def plot_tree( - self, - header_f_by_node_name, - charts_f_by_node_name, - dirs: Dict[str, str], - ) -> None: - """ - Plot and save a tree visualization with custom header and chart images for nodes. - - This method uses the `self.tree_ete` attribute of the class to visualize the tree - in a hierarchical manner, with customized header and chart images for each node. - - Args: - - header_f_by_node_name: Dictionary mapping node names to header image file paths (must be PNG). - - charts_f_by_node_name: Dictionary mapping node names to chart image file paths (must be PNG). - - dirs: A dictionary containing directory paths, specifically 'tree' for saving the tree PDF. - - Returns: - - None - """ - tree_f = os.path.join( - dirs["tree"], "tree.pdf" - ) # must be PDF! (otherwise it breaks) - style = ete3.NodeStyle() - style["vt_line_width"] = 5 - style["hz_line_width"] = 5 - style["fgcolor"] = "darkgrey" - for node in self.tree_ete.traverse("levelorder"): # type: ignore - node.set_style(style) - if header_f_by_node_name[node.name]: - # must be PNG! (ETE can't do PDF Faces) - node_header_face = ete3.faces.ImgFace(header_f_by_node_name[node.name]) - node.add_face(node_header_face, column=0, position="branch-top") - if charts_f_by_node_name[node.name]: - # must be PNG! (ETE can't do PDF Faces) - node_chart_face = ete3.faces.ImgFace(charts_f_by_node_name[node.name]) - node.add_face(node_chart_face, column=0, position="branch-bottom") - node_name_face = ete3.TextFace(node.name, fsize=64) - node.img_style["size"] = 10 - node.img_style["shape"] = "sphere" - node.img_style["fgcolor"] = "black" - if not node.is_leaf(): - node.add_face(node_name_face, column=0, position="branch-right") - node.add_face(node_name_face, column=0, position="aligned") - ts = ete3.TreeStyle() - ts.draw_guiding_lines = True - ts.show_scale = False - ts.show_leaf_name = False - ts.allow_face_overlap = True - ts.guiding_lines_color = "lightgrey" - logger.info(f"[STATUS] - Writing tree {tree_f}... ") - self.tree_ete.render( # type: ignore - tree_f, dpi=600, h=1189, units="mm", tree_style=ts - ) - - def write_tree( - self, - dirs: Dict[str, str], - render_tree: bool, - plot_format: str, - fontsize: int, - ) -> None: - """ - Write tree data to files and optionally render a graphical tree representation. - - This method generates and saves various metrics and data related to the tree structure, - including node statistics and cluster metrics. It can also render a graphical tree - representation if specified. - - Args: - - dirs: A dictionary containing directory paths, including 'tree' for saving tree-related files. - - render_tree: Boolean flag indicating whether to render a graphical tree representation. - - plot_format: Format for saving plots ('png', 'pdf', etc.). - - fontsize: Font size used for plotting. - - Returns: - - None - """ - if not self.tree_ete: - return - logger.info("[STATUS] - Writing data for tree ... ") - # Node stats - node_stats_f = os.path.join(dirs["tree"], "tree.node_metrics.txt") - node_stats_header: List[str] = [ - "#nodeID", - "taxon_specific_apomorphies_singletons", - "taxon_specific_apomorphies_non_singletons", - "node_specific_synapomorphies_total", - "node_specific_synapomorphies_complete_presence", - "node_specific_synapomorphies_partial_absence", - "proteome_count", - ] - node_stats: List[str] = ["\t".join(node_stats_header)] - # Cluster node stats - node_clusters_f = os.path.join(dirs["tree"], "tree.cluster_metrics.txt") - node_clusters_header = [ - "#clusterID", - "nodeID", - "synapomorphy_type", - "node_taxon_coverage", - "children_coverage", - "node_taxa_present", - ] - node_clusters = ["\t".join(node_clusters_header)] - # header_f_by_node_name - header_f_by_node_name = {} - charts_f_by_node_name = {} - for node in self.tree_ete.traverse("levelorder"): # type: ignore - node_clusters.extend( - "\t".join( - [str(string) for string in list(synapomorphic_cluster_string)] - ) - for synapomorphic_cluster_string in node.synapomorphic_cluster_strings - ) - node_stats_line = [ - node.name, - node.apomorphic_cluster_counts["singletons"], # type: ignore - node.apomorphic_cluster_counts["non_singletons"], # type: ignore - ( - # type: ignore - node.synapomorphic_cluster_counts["complete_presence"] # type: ignore - # type: ignore - + node.synapomorphic_cluster_counts["partial_absence"] # type: ignore - ), - # type: ignore - node.synapomorphic_cluster_counts["complete_presence"], # type: ignore - node.synapomorphic_cluster_counts["partial_absence"], # type: ignore - len(node.proteome_ids), # type: ignore - ] - node_stats.append("\t".join([str(string) for string in node_stats_line])) - if render_tree: - header_f_by_node_name[node.name] = self.generate_header_for_node( - node, dirs - ) - charts_f_by_node_name[node.name] = self.generate_chart_for_node( - node, dirs, plot_format, fontsize - ) - logger.info(f"[STATUS] - Writing {node_stats_f} ... ") - with open(node_stats_f, "w") as node_stats_fh: - node_stats_fh.write("\n".join(node_stats) + "\n") - logger.info(f"[STATUS] - Writing {node_clusters_f} ... ") - with open(node_clusters_f, "w") as node_clusters_fh: - node_clusters_fh.write("\n".join(node_clusters) + "\n") - if render_tree: - self.plot_tree(header_f_by_node_name, charts_f_by_node_name, dirs) - else: - self.plot_text_tree(dirs) - - def compute_repetition_for_rarefaction_curve( - self, - ALO: AttributeLevel, - attribute: str, - level: str, - rarefaction_by_samplesize_by_level_by_attribute: Dict[ - str, Dict[str, Dict[int, List[int]]] - ], - ): - seen_cluster_ids = set() - random_list_of_proteome_ids = list(ALO.proteomes) - random.shuffle(random_list_of_proteome_ids) - for idx, proteome_id in enumerate(random_list_of_proteome_ids): - key = "TAXON" if "TAXON" in self.ALO_by_level_by_attribute else "taxon" - if proteome_ALO := self.ALO_by_level_by_attribute[key][proteome_id]: - seen_cluster_ids.update( - proteome_ALO.cluster_ids_by_cluster_type_by_cluster_status[ - "present" - ]["specific"] - ) - seen_cluster_ids.update( - proteome_ALO.cluster_ids_by_cluster_type_by_cluster_status[ - "present" - ]["shared"] - ) - sample_size = idx + 1 - if ( - sample_size - not in rarefaction_by_samplesize_by_level_by_attribute[attribute][ - level - ] - ): - rarefaction_by_samplesize_by_level_by_attribute[attribute][level][ - sample_size - ] = [] - rarefaction_by_samplesize_by_level_by_attribute[attribute][level][ - sample_size - ].append(len(seen_cluster_ids)) - - def compute_rarefaction_data( - self, repetitions: int - ) -> Dict[str, Dict[str, Dict[int, List[int]]]]: - """ - Compute rarefaction data and generate rarefaction curves for proteome clusters. - - This method computes rarefaction curves to analyze the accumulation of non-singleton - clusters as proteome samples increase. It generates plots for each attribute based on - the specified parameters. - - Args: - - repetitions: Number of repetitions to shuffle proteome lists for random sampling. - - Returns: - - Dict[str, Dict[str, Dict[int, List[int]]]] - """ - rarefaction_by_samplesize_by_level_by_attribute: Dict[ - str, Dict[str, Dict[int, List[int]]] - ] = {} - logger.info("[STATUS] - Generating rarefaction data ...") - try: - for attribute in self.attributes: - if attribute not in rarefaction_by_samplesize_by_level_by_attribute: - rarefaction_by_samplesize_by_level_by_attribute[attribute] = {} - for level, proteome_ids in self.proteome_ids_by_level_by_attribute[ - attribute - ].items(): - logger.info( - f"[STATUS] - Processing {attribute} at level {level} ..." - ) - if len(proteome_ids) == 1: - logger.info( - f"[STATUS] - ... skipping {attribute} at level {level}" - ) - continue - if ( - level - not in rarefaction_by_samplesize_by_level_by_attribute[ - attribute - ] - ): - rarefaction_by_samplesize_by_level_by_attribute[attribute][ - level - ] = {} - ALO = self.ALO_by_level_by_attribute[attribute].get(level) - if ALO is None: - continue - for _ in range(repetitions): - self.compute_repetition_for_rarefaction_curve( - ALO=ALO, - attribute=attribute, - level=level, - rarefaction_by_samplesize_by_level_by_attribute=rarefaction_by_samplesize_by_level_by_attribute, - ) - except Exception as e: - logger.error(f"[ERROR] - {e}") - return {} - return rarefaction_by_samplesize_by_level_by_attribute diff --git a/src/core/build.py b/src/core/build.py deleted file mode 100644 index 025194b..0000000 --- a/src/core/build.py +++ /dev/null @@ -1,403 +0,0 @@ -import json -import logging -import os -from collections import Counter, OrderedDict, defaultdict -from typing import Dict, List, Optional, Set - -from core.alo_collections import AloCollection -from core.clusters import Cluster, ClusterCollection -from core.logic import ( - add_taxid_attributes, - parse_attributes_from_config_data, - parse_fasta_dir, - parse_go_mapping, - parse_ipr_mapping, - parse_pfam_mapping, - parse_tree_from_file, -) -from core.proteins import Protein, ProteinCollection -from core.utils import progress, yield_file_lines - -logger = logging.getLogger("kinfin_logger") - - -def get_singletons( - proteinCollection: ProteinCollection, - cluster_list: List[Cluster], -) -> int: - """ - Identify and create singleton clusters for unclustered proteins in a protein collection. - - Args: - - proteinCollection (ProteinCollection): An instance of ProteinCollection class. - - cluster_list (List[Cluster]): A list to which new singleton Cluster objects will be appended. - - Returns: - - int: Number of singleton clusters created and appended to cluster_list. - - This function iterates through proteins in the given protein collection that are not yet clustered. - For each unclustered protein, it creates a new singleton cluster and appends it to cluster_list. - """ - logger.info("[STATUS] - Inferring singletons ...") - singleton_idx = 0 - for protein in proteinCollection.proteins_list: - if protein.clustered is False: - cluster_id = f"singleton_{singleton_idx}" - cluster = Cluster( - cluster_id, - [protein.protein_id], - proteinCollection, - ) - cluster_list.append(cluster) - singleton_idx += 1 - return singleton_idx - - -def parse_cluster_file( - output_dir: str, - cluster_f: str, - proteinCollection: ProteinCollection, - available_proteomes: Set[str], -) -> List[Cluster]: - """ - Parses a cluster file to create Cluster objects and updates protein information. - Saves the filtered clustering data and stats to files. - - Args: - output_dir (str): Base directory path for saving files. - cluster_f (str): Path to the cluster file. - proteinCollection (ProteinCollection): Collection of Protein objects. - available_proteomes (Set[str]): Set of all available proteomes. - - Returns: - Tuple[List[Cluster], Dict[str, any]]: List of Cluster objects and stats. - - Raises: - FileNotFoundError: If the cluster file `cluster_f` does not exist. - """ - cluster_list: List[Cluster] = [] - stats = { - "total_clusters": 0, - "total_proteins": 0, - "total_proteomes": len(available_proteomes), - "filtered_clusters": 0, - "filtered_proteins": 0, - "included_proteins": [], - "excluded_proteins": [], - "included_proteomes": defaultdict(int), - "excluded_proteomes": defaultdict(int), - } - - output_filtered_file = os.path.join(output_dir, "orthogroups.filtered.txt") - stats_file = os.path.join(output_dir, "summary.json") - - logger.info(f"[STATUS] - Available proteomes: {available_proteomes}") - - try: - - with open(cluster_f) as fh, open(output_filtered_file, "w") as ofh: - for line in fh: - stats["total_clusters"] += 1 - temp: List[str] = line.rstrip("\n").split(" ") - cluster_id, protein_ids = temp[0].replace(":", ""), temp[1:] - protein_ids = [protein_id for protein_id in protein_ids if protein_id] - - filtered_protein_ids = [] - for protein_id in protein_ids: - proteome_id = protein_id.split(".")[0] # Extract proteome ID - if proteome_id in available_proteomes: - filtered_protein_ids.append(protein_id) - stats["included_proteins"].append(protein_id) - stats["included_proteomes"][proteome_id] += 1 - else: - stats["excluded_proteins"].append(protein_id) - stats["excluded_proteomes"][proteome_id] += 1 - - stats["total_proteins"] += len(protein_ids) - stats["filtered_proteins"] += len(filtered_protein_ids) - - if filtered_protein_ids: - # Only create a cluster if there are proteins left after filtering - cluster = Cluster( - cluster_id, filtered_protein_ids, proteinCollection - ) - for protein_id in filtered_protein_ids: - protein = proteinCollection.proteins_by_protein_id[protein_id] - protein.clustered = True - cluster_list.append(cluster) - filtered_protein_ids.sort() - ofh.write(f"{cluster_id}: {', '.join(filtered_protein_ids)}\n") - stats["filtered_clusters"] += 1 - except Exception as e: - logger.error("[ERROR] - Something has gone wrong in build.py") - logger.error(f"[ERROR] - Error parsing cluster file: {e}") - raise e - - stats["included_proteins_count"] = len(set(stats["included_proteins"])) - stats["excluded_proteins_count"] = len(set(stats["excluded_proteins"])) - - # Convert proteome counts to lists of counts for JSON serialization - stats["included_proteomes"] = dict(stats["included_proteomes"]) - stats["excluded_proteomes"] = dict(stats["excluded_proteomes"]) - - # Reorder stats - ordered_stats = OrderedDict( - [ - ("total_clusters", stats["total_clusters"]), - ("total_proteins", stats["total_proteins"]), - ("total_proteomes", stats["total_proteomes"]), - ("filtered_clusters", stats["filtered_clusters"]), - ("filtered_proteins", stats["filtered_proteins"]), - ("included_proteins_count", stats["included_proteins_count"]), - ("excluded_proteins_count", stats["excluded_proteins_count"]), - ("included_proteomes", stats["included_proteomes"]), - ("excluded_proteomes", stats["excluded_proteomes"]), - ("included_proteins", stats["included_proteins"]), - ("excluded_proteins", stats["excluded_proteins"]), - ] - ) - - with open(stats_file, "w") as mf: - json.dump( - ordered_stats, - mf, - separators=(", ", ": "), - indent=4, - ) - - return cluster_list - - -# cli -def parse_domains_from_functional_annotations_file( - functional_annotation_f: str, - proteinCollection: ProteinCollection, -) -> None: - """ - Parse functional annotations from a file and populate ProteinCollection with parsed data. - - Parameters: - - functional_annotation_f (str): Path to the functional annotation file. - - proteinCollection (ProteinCollection): Instance of ProteinCollection class to store parsed data. - - pfam_mapping (bool): Flag indicating whether to parse Pfam mappings. - - ipr_mapping (bool): Flag indicating whether to parse InterPro mappings. - - pfam_mapping_f (str): File path to the Pfam mapping file. - - ipr_mapping_f (str): File path to the InterPro mapping file. - - go_mapping_f (str): File path to the GO mapping file. - - Raises: - - ValueError: If the functional annotation file lacks a header. - - Notes: - - The function reads each line of the functional annotation file, parses relevant data, - and populates the proteinCollection with domain annotations and GO terms. - - It also optionally parses additional mappings (Pfam, InterPro, GO) based on provided flags. - - Updates proteinCollection.functional_annotation_parsed and proteinCollection.domain_desc_by_id_by_source. - """ - - logger.info( - f"[STATUS] - Parsing {functional_annotation_f} ... this may take a while" - ) - - for line in yield_file_lines(functional_annotation_f): - temp: List[str] = line.split() - if temp[0].startswith("#"): - proteinCollection.domain_sources = temp[1:] - - else: - if not proteinCollection.domain_sources: - error_msg = f"[ERROR] - {functional_annotation_f} does not seem to have a header." - raise ValueError(error_msg) - - domain_protein_id: str = temp.pop(0) - go_terms: List[str] = [] - domain_counter_by_domain_source: Dict[str, Counter[str]] = {} - for idx, field in enumerate(temp): - if field != "None": - domain_source: str = proteinCollection.domain_sources[idx] - domain_string: List[str] = field.split(";") - domain_counts_by_domain_id: Dict[str, int] = {} - for domain_id_count in domain_string: - domain_id: str - domain_count: int = 1 - if domain_source == "GO": - domain_id = domain_id_count - else: - domain_id, domain_count_str = domain_id_count.rsplit(":", 2) - domain_count = int(domain_count_str) - domain_counts_by_domain_id[domain_id] = domain_count - domain_counter: Counter[str] = Counter(domain_counts_by_domain_id) - domain_counter_by_domain_source[domain_source] = domain_counter - proteinCollection.add_annotation_to_protein( - domain_protein_id=domain_protein_id, - domain_counter_by_domain_source=domain_counter_by_domain_source, - go_terms=go_terms, - ) - - proteinCollection.functional_annotation_parsed = True - - -# common -def build_AloCollection( - config_f: str, - nodesdb_f: str, - taxranks: List[str], - tree_f: Optional[str], - taxon_idx_mapping_file: Optional[str], -) -> AloCollection: - """ - Builds an AloCollection object from command-line interface (CLI) inputs. - - Args: - config_f (str): Path to the configuration file containing proteome attributes. - nodesdb_f (str): Path to the nodes database file for inferring taxonomic ranks. - taxranks (List[str]): List of taxonomic ranks to be inferred. - tree_f (Optional[str]): Path to the tree file. If provided, ALOs are added from the tree. - - Returns: - AloCollection: An instance of the AloCollection class containing parsed data. - """ - ( - proteomes, - proteome_id_by_species_id, - attributes, - level_by_attribute_by_proteome_id, - ) = parse_attributes_from_config_data(config_f, taxon_idx_mapping_file) - # Add taxonomy if needed - if "TAXID" in set(attributes): - logger.info( - "[STATUS] - Attribute 'TAXID' found, inferring taxonomic ranks from nodesDB" - ) - attributes, level_by_attribute_by_proteome_id = add_taxid_attributes( - attributes=attributes, - level_by_attribute_by_proteome_id=level_by_attribute_by_proteome_id, - nodesdb_f=nodesdb_f, - taxranks=taxranks, - ) - - # Add ALOs from tree if provided - tree_ete, node_idx_by_proteome_ids = parse_tree_from_file( - tree_f, - attributes, - level_by_attribute_by_proteome_id, - proteomes, - ) - - logger.info("[STATUS] - Building AloCollection ...") - return AloCollection( - proteomes=proteomes, - attributes=attributes, - proteome_id_by_species_id=proteome_id_by_species_id, - level_by_attribute_by_proteome_id=level_by_attribute_by_proteome_id, - node_idx_by_proteome_ids=node_idx_by_proteome_ids, - tree_ete=tree_ete, - ) - - -def get_protein_list_from_seq_f(sequence_ids_f: str, aloCollection: AloCollection): - logger.info(f"[STATUS] - Parsing sequence IDs: {sequence_ids_f} ...") - - proteins_list: List[Protein] = [] - for line in yield_file_lines(sequence_ids_f): - temp = line.split(": ") - sequence_id = temp[0] - protein_id = ( - temp[1] - .split(" ")[0] - .replace(":", "_") - .replace(",", "_") - .replace("(", "_") - .replace(")", "_") - ) # orthofinder replaces characters - species_id = sequence_id.split("_")[0] - if proteome_id := aloCollection.proteome_id_by_species_id.get(species_id, None): - protein = Protein(protein_id, proteome_id, species_id, sequence_id) - proteins_list.append(protein) - return proteins_list - - -# common -def build_ProteinCollection( - sequence_ids_f: str, - aloCollection: AloCollection, - fasta_dir: Optional[str], - species_ids_f: Optional[str], - functional_annotation_f: Optional[str], - pfam_mapping: bool, - ipr_mapping: bool, - pfam_mapping_f: str, - go_mapping_f: str, - ipr_mapping_f: str, -) -> ProteinCollection: - proteins_list = get_protein_list_from_seq_f( - sequence_ids_f=sequence_ids_f, - aloCollection=aloCollection, - ) - proteinCollection = ProteinCollection(proteins_list) - - logger.info(f"[STATUS]\t - Proteins found = {proteinCollection.protein_count}") - - if fasta_dir is not None and species_ids_f is not None: - fasta_len_by_protein_id = parse_fasta_dir( - fasta_dir=fasta_dir, - species_ids_f=species_ids_f, - ) - logger.info("[STATUS] - Adding FASTAs to ProteinCollection ...") - parse_steps: float = proteinCollection.protein_count / 100 - for idx, protein in enumerate(proteinCollection.proteins_list): - protein.update_length(fasta_len_by_protein_id[protein.protein_id]) - progress(idx + 1, parse_steps, proteinCollection.protein_count) - aloCollection.fastas_parsed = True - proteinCollection.fastas_parsed = True - else: - logger.info( - "[STATUS] - No Fasta-Dir given, no AA-span information will be reported ..." - ) - - if functional_annotation_f is not None: - parse_domains_from_functional_annotations_file( - functional_annotation_f=functional_annotation_f, - proteinCollection=proteinCollection, - ) - domain_desc_by_id_by_source = {} - - if pfam_mapping and "Pfam" in proteinCollection.domain_sources: - domain_desc_by_id_by_source["Pfam"] = parse_pfam_mapping(pfam_mapping_f) - - if ipr_mapping and "IPR" in proteinCollection.domain_sources: - domain_desc_by_id_by_source["IPR"] = parse_ipr_mapping(ipr_mapping_f) - - if go_mapping_f: - domain_desc_by_id_by_source["GO"] = parse_go_mapping(go_mapping_f) - - proteinCollection.domain_desc_by_id_by_source = domain_desc_by_id_by_source - - return proteinCollection - - -def build_ClusterCollection( - output_dir: str, - cluster_f: str, - proteinCollection: ProteinCollection, - infer_singletons: Optional[bool], - available_proteomes: Set[str], -) -> ClusterCollection: - logger.info(f"[STATUS] - Parsing {cluster_f} ... this may take a while") - cluster_list: List[Cluster] = parse_cluster_file( - output_dir, - cluster_f, - proteinCollection, - available_proteomes, - ) - - inferred_singletons_count = 0 - if infer_singletons: - inferred_singletons_count = get_singletons(proteinCollection, cluster_list) - - return ClusterCollection( - cluster_list, - inferred_singletons_count, - proteinCollection.functional_annotation_parsed, - proteinCollection.fastas_parsed, - proteinCollection.domain_sources, - ) diff --git a/src/core/clusters.py b/src/core/clusters.py deleted file mode 100644 index 79cff89..0000000 --- a/src/core/clusters.py +++ /dev/null @@ -1,196 +0,0 @@ -from collections import Counter -from math import log -from typing import DefaultDict, Dict, FrozenSet, List, Literal, Optional, Set - -from core.logic import compute_protein_ids_by_proteome -from core.proteins import ProteinCollection -from core.utils import mean, median, sd - - -class Cluster: - def __init__( - self, - cluster_id: str, - protein_ids: List[str], - proteinCollection: ProteinCollection, - ) -> None: - self.cluster_id: str = cluster_id - self.protein_ids = set(protein_ids) - self.protein_count: int = len(protein_ids) - try: - - self.proteomes_by_protein_id: Dict[str, str] = { - _id: proteinCollection.proteins_by_protein_id[_id].proteome_id - for _id in protein_ids - } - except KeyError as e: - error_msg = f"[ERROR] - Protein {e.args[0]} in clustering belongs to proteomes that are not present in the config-file." - error_msg += ( - "Please add those proteomes or recluster by omitting these proteomes." - ) - raise KeyError(error_msg) from e - - self.proteome_ids_list: List[str] = list(self.proteomes_by_protein_id.values()) - self.protein_count_by_proteome_id: Counter[str] = Counter( - self.proteome_ids_list - ) - self.proteome_ids: FrozenSet[str] = frozenset(self.proteome_ids_list) - self.proteome_count: int = len(self.proteome_ids) - self.singleton: bool = self.protein_count <= 1 - self.apomorphy: bool = self.proteome_count <= 1 - - self.protein_ids_by_proteome_id: DefaultDict[str, Set[str]] = ( - compute_protein_ids_by_proteome(self.proteomes_by_protein_id) - ) - self.protein_counts_of_proteomes_by_level_by_attribute: Dict[ - str, Dict[str, List[int]] - ] = {} - self.proteome_coverage_by_level_by_attribute: Dict[str, Dict[str, float]] = {} - self.implicit_protein_ids_by_proteome_id_by_level_by_attribute: Dict[ - str, Dict[str, Dict[str, List[str]]] - ] = {} - self.cluster_type_by_attribute: Dict[ - str, - Literal["singleton", "shared", "specific"], - ] = {} - self.protein_median: Optional[float] = None - self.protein_length_stats: Optional[Dict[str, float]] = ( - self.compute_protein_length_stats(proteinCollection, self.protein_ids) - ) - self.secreted_cluster_coverage: float = self.compute_secreted_cluster_coverage( - proteinCollection, self.protein_ids, self.protein_count - ) - self.domain_counter_by_domain_source: Dict[str, Counter[str]] = ( - self.compute_domain_counter_by_domain_source( - proteinCollection, self.protein_ids - ) - ) - self.domain_entropy_by_domain_source: Dict[str, float] = ( - self.compute_domain_entropy_by_domain_source() - ) - - def compute_protein_length_stats( - self, - proteinCollection: ProteinCollection, - protein_ids: Set[str], - ) -> Optional[Dict[str, float]]: - """ - Computes statistics (mean, median, standard deviation) of protein lengths. - - Parameters: - - proteinCollection: A ProteinCollection object containing protein data. - - protein_ids: A set of protein IDs for which lengths are to be computed. - - Returns: - - Optional[Dict[str, float]]: A dictionary containing 'mean', 'median', and 'sd' - (standard deviation) of protein lengths, if all lengths are available and at least - one protein ID is provided. Returns None if no valid protein lengths are found. - """ - protein_lengths: List[Optional[int]] = [ - proteinCollection.proteins_by_protein_id[protein_id].length - for protein_id in protein_ids - ] - if all(protein_lengths): - protein_length_stats: Dict[str, float] = {"mean": mean(protein_lengths)} - protein_length_stats["median"] = median(protein_lengths) - protein_length_stats["sd"] = sd(protein_lengths) - return protein_length_stats - - def compute_secreted_cluster_coverage( - self, - proteinCollection: ProteinCollection, - protein_ids: Set[str], - protein_count: int, - ) -> float: - """ - Computes the fraction of secreted proteins in a given set of protein IDs. - - Parameters: - - proteinCollection: A ProteinCollection object containing protein data. - - protein_ids: A set of protein IDs to compute secreted protein coverage. - - protein_count: Total count of proteins in the cluster. - - Returns: - - float: Fraction of secreted proteins in the provided set of protein IDs. - """ - secreted = sum( - bool(proteinCollection.proteins_by_protein_id[protein_id].secreted) - for protein_id in protein_ids - ) - return secreted / protein_count - - def compute_domain_counter_by_domain_source( - self, - proteinCollection: ProteinCollection, - protein_ids: Set[str], - ) -> Dict[str, Counter[str]]: - """ - Computes the aggregated domain counts by domain source for a set of protein IDs. - - Parameters: - - proteinCollection: A ProteinCollection object containing protein data. - - protein_ids: A set of protein IDs for which domain counts are computed. - - Returns: - - Dict[str, Counter[str]]: A dictionary where keys are domain sources and values are - Counters mapping domain IDs to their respective counts. - """ - cluster_domain_counter_by_domain_source: Dict[str, Counter[str]] = {} - for protein_id in protein_ids: - if protein_domain_counter_by_domain_source := proteinCollection.proteins_by_protein_id[ - protein_id - ].domain_counter_by_domain_source: - for domain_source, protein_domain_counter in list( - protein_domain_counter_by_domain_source.items() - ): - if domain_source not in cluster_domain_counter_by_domain_source: - cluster_domain_counter_by_domain_source[domain_source] = ( - Counter() - ) - cluster_domain_counter_by_domain_source[ - domain_source - ] += protein_domain_counter - return cluster_domain_counter_by_domain_source - - def compute_domain_entropy_by_domain_source(self) -> Dict[str, float]: - """ - Computes entropy for domains grouped by different sources. - - Returns: - - Dict[str, float]: Dictionary where keys are domain sources and values are computed entropy values. - """ - self.domain_entropy_by_domain_source: Dict[str, float] = {} - for domain_source, domain_counter in list( - self.domain_counter_by_domain_source.items() - ): - total_count: int = len(list(domain_counter.elements())) - domain_entropy: float = -sum( - i / total_count * log(i / total_count, 2) - for i in list(domain_counter.values()) - ) - if str(domain_entropy) == "-0.0": - self.domain_entropy_by_domain_source[domain_source] = 0.0 - else: - self.domain_entropy_by_domain_source[domain_source] = domain_entropy - return self.domain_entropy_by_domain_source - - -class ClusterCollection: - def __init__( - self, - cluster_list: List[Cluster], - inferred_singletons_count: int, - functional_annotation_parsed: bool, - fastas_parsed: bool, - domain_sources: List[str], - ): - self.cluster_list: List[Cluster] = cluster_list - self.cluster_list_by_cluster_id: Dict[str, Cluster] = { - cluster.cluster_id: cluster for cluster in cluster_list - } # only for testing - self.cluster_count: int = len(cluster_list) - self.inferred_singletons_count: int = inferred_singletons_count - self.functional_annotation_parsed: bool = functional_annotation_parsed - self.fastas_parsed: bool = fastas_parsed - # self.domain_sources = [domain_source for domain_source in domain_sources if not domain_source == "GO"] - self.domain_sources: List[str] = domain_sources diff --git a/src/core/datastore.py b/src/core/datastore.py deleted file mode 100644 index b49004f..0000000 --- a/src/core/datastore.py +++ /dev/null @@ -1,2116 +0,0 @@ -import logging -import os -import time -from collections import Counter, defaultdict -from typing import Any, Dict, FrozenSet, Generator, List, Set, Tuple, Union - -import matplotlib as mat -import matplotlib.pyplot as plt -import numpy as np -from matplotlib.lines import Line2D -from matplotlib.ticker import FormatStrFormatter, NullFormatter - -from core.alo import AttributeLevel -from core.alo_collections import AloCollection -from core.build import ( - build_AloCollection, - build_ClusterCollection, - build_ProteinCollection, -) -from core.clusters import Cluster, ClusterCollection -from core.input import InputData -from core.logic import get_ALO_cluster_cardinality, get_attribute_cluster_type -from core.proteins import ProteinCollection -from core.utils import median, progress, statistic - -logger = logging.getLogger("kinfin_logger") -mat.use("agg") - -plt.style.use("ggplot") -mat.rc("ytick", labelsize=20) -mat.rc("xtick", labelsize=20) -axis_font = {"size": "20"} -mat.rcParams.update({"font.size": 22}) - - -class DataFactory: - def __init__(self, inputData: InputData) -> None: - self.dirs = {} - self.inputData: InputData = inputData - self.aloCollection: AloCollection = build_AloCollection( - config_f=self.inputData.config_f, - nodesdb_f=self.inputData.nodesdb_f, - tree_f=self.inputData.tree_f, - taxranks=self.inputData.taxranks, - taxon_idx_mapping_file=self.inputData.taxon_idx_mapping_file, - ) - self.proteinCollection: ProteinCollection = build_ProteinCollection( - aloCollection=self.aloCollection, - fasta_dir=self.inputData.fasta_dir, - go_mapping_f=self.inputData.go_mapping_f, - functional_annotation_f=self.inputData.functional_annotation_f, - ipr_mapping=self.inputData.ipr_mapping, - ipr_mapping_f=self.inputData.ipr_mapping_f, - pfam_mapping=self.inputData.pfam_mapping, - pfam_mapping_f=self.inputData.pfam_mapping_f, - sequence_ids_f=self.inputData.sequence_ids_f, - species_ids_f=self.inputData.species_ids_f, - ) - self.clusterCollection: ClusterCollection = build_ClusterCollection( - cluster_f=self.inputData.cluster_f, - output_dir=self.inputData.output_path, - proteinCollection=self.proteinCollection, - infer_singletons=self.inputData.infer_singletons, - available_proteomes=self.aloCollection.proteomes, - ) - - def setup_dirs(self) -> None: - """ - Set up output directories for storing results and attributes. - """ - output_path: str = self.inputData.output_path - - self.dirs["main"] = output_path - logger.info("[STATUS] - Output directories in") - logger.info(f"\t{output_path}") - if not os.path.exists(output_path): - logger.info("[STATUS] - Creating main output directory...") - os.makedirs(output_path) - - logger.info("[STATUS] - Creating directories ...") - for attribute in self.aloCollection.attributes: - attribute_path = os.path.join(output_path, attribute) - self.dirs[attribute] = attribute_path - if not os.path.exists(attribute_path): - logger.info( - f"[STATUS] - Creating directory for attribute: {attribute_path}" - ) - os.makedirs(attribute_path) - - if self.aloCollection.tree_ete is not None: - self._create_tree_directories(output_path) - logger.info("[STATUS] - Creating directories ... Done") - - def _create_tree_directories(self, output_path): - tree_path = os.path.join(output_path, "tree") - node_chart_path = os.path.join(tree_path, "charts") - node_header_path = os.path.join(tree_path, "headers") - - if not os.path.exists(tree_path): - self._extracted_from_setup_dirs_30( - "[STATUS] - Creating tree directory: ", tree_path, "tree" - ) - if not os.path.exists(node_chart_path): - self._extracted_from_setup_dirs_30( - "[STATUS] - Creating node charts directory: ", - node_chart_path, - "tree_charts", - ) - if self.inputData.plot_tree and not os.path.exists(node_header_path): - self._create_directory( - "[STATUS] - Creating node headers directory: ", - node_header_path, - "tree_headers", - ) - - def _create_directory(self, log_message, directory_path, dir_key): - logger.info(f"{log_message}{directory_path}") - os.makedirs(directory_path) - self.dirs[dir_key] = directory_path - - def analyse_clusters(self) -> None: - """ - Analyses clusters within the cluster collection. - - Then proceeds to analyse each cluster individually, - logging progress and timing information. - - Returns: - None - """ - if self.clusterCollection.inferred_singletons_count: - logger.info( - f"[STATUS]\t - Clusters found = {self.clusterCollection.cluster_count} (of which {self.clusterCollection.inferred_singletons_count} were inferred singletons)") # fmt:skip - - else: - logger.info( - f"[STATUS]\t - Clusters found = {self.clusterCollection.cluster_count}" - ) - - parse_steps = self.clusterCollection.cluster_count / 100 - - logger.info("[STATUS] - Analysing clusters ...") - analyse_clusters_start = time.time() - for idx, cluster in enumerate(self.clusterCollection.cluster_list): - self.__analyse_cluster(cluster) - progress(idx + 1, parse_steps, self.clusterCollection.cluster_count) - analyse_clusters_end = time.time() - analyse_clusters_elapsed = analyse_clusters_end - analyse_clusters_start - logger.info(f"[STATUS] - Took {analyse_clusters_elapsed}s to analyse clusters") - - def plot_rarefaction_data( - self, - rarefaction_by_samplesize_by_level_by_attribute: Dict[ - str, Dict[str, Dict[int, List[int]]] - ], - dirs: Dict[str, str], - plotsize: Tuple[float, float], - plot_format: str, - fontsize: int, - ) -> None: - """ - Plot rarefaction curves based on provided data. - - Args: - rarefaction_by_samplesize_by_level_by_attribute (dict): A nested dictionary - where keys are attribute names, and values are dictionaries where keys - are level names and values are dictionaries mapping sample sizes to - lists of non-singleton cluster counts. - dirs (dict): A dictionary mapping attribute names to directory paths where - plots will be saved. - plotsize (tuple): A tuple specifying the size of the plot (width, height) in inches. - plot_format (str): The format of the plot to save (e.g., 'png', 'pdf'). - fontsize (int): Font size for plot labels and legend. - - Returns: - None - """ - for ( - attribute, - rarefaction_by_samplesize_by_level, - ) in rarefaction_by_samplesize_by_level_by_attribute.items(): - rarefaction_plot_f = os.path.join( - dirs[attribute], f"{attribute}.rarefaction_curve.{plot_format}" - ) - f, ax = plt.subplots(figsize=plotsize) - ax.set_facecolor("white") - max_number_of_samples = 0 - for idx, level in enumerate(rarefaction_by_samplesize_by_level): - number_of_samples = len(rarefaction_by_samplesize_by_level[level]) - if number_of_samples > max_number_of_samples: - max_number_of_samples = number_of_samples - colour = plt.cm.Paired( # type: ignore - idx / len(rarefaction_by_samplesize_by_level) - ) # type: ignore - x_values = [] - y_mins = [] - y_maxs = [] - median_y_values = [] - median_x_values = [] - for x, y_reps in list( - rarefaction_by_samplesize_by_level[level].items() - ): - x_values.append(x) - y_mins.append(min(y_reps)) - y_maxs.append(max(y_reps)) - median_y_values.append(median(y_reps)) - median_x_values.append(x) - x_array = np.array(x_values) - y_mins_array = np.array(y_mins) - y_maxs_array = np.array(y_maxs) - ax.plot( - median_x_values, - median_y_values, - "-", - color=colour, - label=level, - ) - ax.fill_between( - x_array, - y_mins_array, # type:ignore - y_maxs_array, # type:ignore - color=colour, - alpha=0.5, - ) - ax.set_xlim([0, max_number_of_samples + 1]) - ax.set_ylabel("Count of non-singleton clusters", fontsize=fontsize) - ax.set_xlabel("Sampled proteomes", fontsize=fontsize) - - ax.grid(True, linewidth=1, which="major", color="lightgrey") - legend = ax.legend( - ncol=1, - numpoints=1, - loc="lower right", - frameon=True, - fontsize=fontsize, - ) - legend.get_frame().set_facecolor("white") - logger.info(f"[STATUS]\t- Plotting {rarefaction_plot_f}") - f.savefig(rarefaction_plot_f, format=plot_format) - plt.close() - - def write_output(self) -> None: - """ - Executes various methods to generate and write output files related to cluster analysis. - - This method sequentially calls private methods to: - - Plot cluster sizes. - - Write cluster counts by taxon. - - Write cluster metrics related to domains. - - Write detailed cluster metrics related to domains. - - Write attribute metrics. - - Write a summary of cluster metrics. - - Write cluster metrics related to ALO (Additive Log Ratio) transformation. - - Write cluster 1-to-1 ALO metrics. - - Write pairwise representation metrics. - - Each private method is responsible for generating specific outputs based on internal data. - - Returns: - None - """ - self.__plot_cluster_sizes() - self.__write_cluster_counts_by_taxon() - self.__write_cluster_metrics_domains() - self.__write_cluster_metrics_domains_detailed() - self.__write_attribute_metrics() - self.__write_cluster_summary() - self.__write_cluster_metrics_ALO() - self.__write_cluster_1to1_ALO() - self.__write_pairwise_representation() - - # analyse cluster - def __analyse_ete_for_specific_cluster( - self, - cluster: Cluster, - intersection: FrozenSet[str], - node, - ) -> None: - """ - Analyzes a specific cluster within an evolutionary tree node. - - Updates various counts and attributes of the node based on the characteristics - of the given cluster and its intersection with proteome IDs. - - Args: - cluster (Cluster): The cluster to analyze. - intersection (FrozenSet[str]): The intersection of proteome IDs between - the cluster and the current node. - node: The evolutionary tree node to update. - - Returns: - None - """ - node.counts["specific"] += 1 # type: ignore - if cluster.proteome_count == 1: - # But it only belongs to one proteome - node.apomorphic_cluster_counts["non_singletons"] += 1 # type: ignore - else: - # It has more than one proteome - child_nodes_covered = [] - child_node_proteome_coverage_strings = [] - child_node_proteome_ids_covered_count = 0 - for child_node in node.get_children(): - if child_node.proteome_ids.isdisjoint(cluster.proteome_ids): - # No child node proteomes are not in cluster - child_nodes_covered.append(False) - else: - # At least on child node proteome in cluster - child_nodes_covered.append(True) - child_node_proteome_ids_covered_count = len( - cluster.proteome_ids.intersection(child_node.proteome_ids) - ) - child_node_proteome_coverage_strings.append( - f"{child_node.name}=({child_node_proteome_ids_covered_count}/{len(child_node.proteome_ids)})" - ) - if all(child_nodes_covered): - # At least one proteome of each child node in cluster - # => SYNAPOMORPHY - node_proteome_coverage = len(intersection) / len( - node.proteome_ids - ) # type: ignore - node_cluster_type = "" - node_cluster_type = ( - "complete_presence" - if node_proteome_coverage == 1.0 - else "partial_absence" - ) - # type: ignore - node.synapomorphic_cluster_counts[node_cluster_type] += 1 - - node.synapomorphic_cluster_strings.append( # type: ignore - ( - cluster.cluster_id, - node.name, - node_cluster_type, - "{0:.3}".format(node_proteome_coverage), - ";".join(child_node_proteome_coverage_strings), - ",".join(sorted(intersection)), - ) - ) - - def __analyse_tree_ete(self, cluster: Cluster) -> None: - """ - Analyzes a cluster within an ETE Tree if available in the ALO collection. - - Traverses the ETE Tree in level order, comparing proteome IDs of each node - with the cluster's proteome IDs. Updates counts and attributes of nodes - based on the analysis results. - - Args: - cluster (Cluster): The cluster to analyze. - - Returns: - None - """ - if not self.aloCollection.tree_ete: - return - - for node in self.aloCollection.tree_ete.traverse("levelorder"): # type: ignore - intersection = cluster.proteome_ids.intersection( - node.proteome_ids # type: ignore - ) # type: ignore - difference = cluster.proteome_ids.difference( - node.proteome_ids # type: ignore - ) # type: ignore - - if len(intersection) == 0: - # Nothing to see here ... - node.counts["absent"] += 1 # type: ignore - - elif cluster.singleton is True: - # This is a singleton - node.counts["singleton"] += 1 # type: ignore - node.apomorphic_cluster_counts["singletons"] += 1 # type: ignore - - elif len(difference) > 0: - # This is a 'shared' cluster - node.counts["shared"] += 1 # type: ignore - - elif len(difference) == 0: - # This is a node 'specific' cluster - self.__analyse_ete_for_specific_cluster( - cluster=cluster, - intersection=intersection, - node=node, - ) - - def __process_level( - self, - cluster: Cluster, - attribute: str, - level: str, - protein_ids_by_level: Dict[str, List[str]], - protein_length_stats_by_level: Dict[str, Dict[str, Union[int, float]]], - explicit_protein_count_by_proteome_id_by_level: Dict[str, Dict[str, int]], - ) -> None: - """ - Processes a specific level within an attribute for a given cluster. - - Retrieves protein IDs and their counts associated with the specified level - from the ALO collection and updates various attributes and collections within - the cluster and the class instance. - - Args: - cluster (Cluster): The cluster for which to process the level. - attribute (str): The attribute associated with the level. - level (str): The specific level to process. - protein_ids_by_level (dict): A dictionary to store protein IDs by level. - protein_length_stats_by_level (dict): A dictionary to store protein length statistics - by level. - explicit_protein_count_by_proteome_id_by_level (dict): A dictionary to store explicit - protein counts by proteome ID for each level. - - Returns: - None - """ - ALO = self.aloCollection.ALO_by_level_by_attribute[attribute][level] - if ALO is None: - return - - protein_ids_by_proteome_id = {} - protein_count_by_proteome_id = {} - protein_ids_by_level[level] = [] - - for proteome_id in ALO.proteomes_list: - protein_ids = list(cluster.protein_ids_by_proteome_id.get(proteome_id, [])) - protein_ids_by_level[level].extend(protein_ids) - protein_count_by_proteome_id[proteome_id] = len(protein_ids) - if protein_count_by_proteome_id[proteome_id] != 0: - protein_ids_by_proteome_id[proteome_id] = protein_ids - - if protein_ids_by_proteome_id: - cluster.implicit_protein_ids_by_proteome_id_by_level_by_attribute[ - attribute - ][level] = protein_ids_by_proteome_id - - explicit_protein_count_by_proteome_id_by_level[level] = ( - protein_count_by_proteome_id - ) - - protein_length_stats_by_level[level] = ( - self.proteinCollection.get_protein_length_stats(protein_ids_by_level[level]) - ) - - cluster.protein_counts_of_proteomes_by_level_by_attribute[attribute][level] = ( - list(protein_count_by_proteome_id.values()) - ) - - def __update_ALO_data( - self, - cluster: Cluster, - attribute: str, - protein_ids_by_level: Dict[str, List[str]], - protein_length_stats_by_level: Dict[str, Dict[str, Union[int, float]]], - explicit_protein_count_by_proteome_id_by_level: Dict[str, Dict[str, int]], - ) -> None: - """ - Updates ALO (Additive Log Ratio) data for a given cluster and attribute. - - Iterates through each level of the ALO collection corresponding to the attribute, - calculates various metrics based on the cluster's protein IDs and attributes, and - updates the ALO object with this information. - - Args: - cluster (Cluster): The cluster to update ALO data for. - attribute (str): The attribute associated with the ALO data. - protein_ids_by_level (dict): A dictionary mapping level names to lists of protein IDs. - protein_length_stats_by_level (dict): A dictionary mapping level names to dictionaries - containing protein length statistics. - explicit_protein_count_by_proteome_id_by_level (dict): A dictionary mapping level names - to dictionaries where keys are proteome IDs and values are explicit protein counts. - - Returns: - None - """ - for level in self.aloCollection.ALO_by_level_by_attribute[attribute]: - ALO = self.aloCollection.ALO_by_level_by_attribute[attribute][level] - if ALO is None: - continue - - cluster.proteome_coverage_by_level_by_attribute[attribute][level] = ( - len( - cluster.implicit_protein_ids_by_proteome_id_by_level_by_attribute[ - attribute - ].get(level, []) - ) - / ALO.proteome_count - ) - - ALO_cluster_status = ( - "present" - if level - in cluster.implicit_protein_ids_by_proteome_id_by_level_by_attribute[ - attribute - ] - else "absent" - ) - - ALO_cluster_cardinality = None - mwu_pvalue = None - mwu_log2_mean = None - mean_ALO_count = None - mean_non_ALO_count = None - - if ( - ALO_cluster_status == "present" - and cluster.cluster_type_by_attribute[attribute] != "singleton" - ): - ALO_proteome_counts_in_cluster = list( - explicit_protein_count_by_proteome_id_by_level[level].values() - ) - ALO_cluster_cardinality = get_ALO_cluster_cardinality( - ALO_proteome_counts_in_cluster=ALO_proteome_counts_in_cluster, - fuzzy_count=self.inputData.fuzzy_count, - fuzzy_fraction=self.inputData.fuzzy_fraction, - fuzzy_range=self.inputData.fuzzy_range, - ) - - if cluster.cluster_type_by_attribute[attribute] == "shared": - non_ALO_proteome_counts_in_cluster = [ - count - for non_ALO_level in explicit_protein_count_by_proteome_id_by_level - if non_ALO_level != level - for count in explicit_protein_count_by_proteome_id_by_level[ - non_ALO_level - ].values() - ] - mwu_pvalue, mwu_log2_mean, mean_ALO_count, mean_non_ALO_count = ( - statistic( - count_1=ALO_proteome_counts_in_cluster, - count_2=non_ALO_proteome_counts_in_cluster, - test=self.inputData.test, - min_proteomes=self.inputData.min_proteomes, - ) - ) - - ALO.add_cluster( - cluster=cluster, - attribute_cluster_type=cluster.cluster_type_by_attribute[attribute], - ALO_cluster_status=ALO_cluster_status, - ALO_protein_length_stats=protein_length_stats_by_level[level], - ALO_protein_ids_in_cluster=protein_ids_by_level[level], - ALO_cluster_cardinality=ALO_cluster_cardinality, - mwu_pvalue=mwu_pvalue, - mwu_log2_mean=mwu_log2_mean, - mean_ALO_count=mean_ALO_count, - mean_non_ALO_count=mean_non_ALO_count, - ) - - def __process_single_attribute(self, cluster: Cluster, attribute: str) -> None: - """ - Processes a single attribute for a given cluster. - - Retrieves and processes each level associated with the attribute from the ALO - collection, updating various protein and cluster metrics within the cluster object. - - Args: - cluster (Cluster): The cluster to process the attribute for. - attribute (str): The attribute to process. - - Returns: - None - """ - protein_ids_by_level: Dict[str, List[str]] = {} - protein_length_stats_by_level: Dict[str, Dict[str, Union[int, float]]] = {} - explicit_protein_count_by_proteome_id_by_level: Dict[str, Dict[str, int]] = {} - - cluster.protein_counts_of_proteomes_by_level_by_attribute[attribute] = {} - cluster.proteome_coverage_by_level_by_attribute[attribute] = {} - cluster.implicit_protein_ids_by_proteome_id_by_level_by_attribute[attribute] = ( - {} - ) - - for level in self.aloCollection.ALO_by_level_by_attribute[attribute]: - self.__process_level( - cluster, - attribute, - level, - protein_ids_by_level, - protein_length_stats_by_level, - explicit_protein_count_by_proteome_id_by_level, - ) - - cluster.cluster_type_by_attribute[attribute] = get_attribute_cluster_type( - cluster.singleton, - cluster.implicit_protein_ids_by_proteome_id_by_level_by_attribute[ - attribute - ], - ) - - self.__update_ALO_data( - cluster, - attribute, - protein_ids_by_level, - protein_length_stats_by_level, - explicit_protein_count_by_proteome_id_by_level, - ) - - def __process_attributes(self, cluster: Cluster) -> None: - """ - Processes all attributes in the ALO collection for a given cluster. - - Iterates through each attribute in the ALO collection and processes it - using the __process_single_attribute method. - - Args: - cluster (Cluster): The cluster to process attributes for. - - Returns: - None - """ - for attribute in self.aloCollection.attributes: - self.__process_single_attribute(cluster, attribute) - - def __finalize_cluster_analysis(self, cluster: Cluster) -> None: - """ - Finalizes the cluster analysis by calculating the median protein count. - - Calculates the median protein count for the given cluster using the protein - counts of proteomes from specific levels and attributes. - - Args: - cluster (Cluster): The cluster for which to finalize the analysis. - - Returns: - None - """ - cluster.protein_median = median( - [ - count - for count in cluster.protein_counts_of_proteomes_by_level_by_attribute[ - "all" - ]["all"] - if count != 0 - ] - ) - - def __analyse_cluster(self, cluster: Cluster) -> None: - """ - Analyzes a cluster by performing various analysis steps. - - Executes the analysis steps for the given cluster: - 1. If an ETE tree is available in aloCollection, analyzes the tree structure. - 2. Processes attributes associated with the cluster. - 3. Finalizes the cluster analysis by calculating median protein counts. - - Args: - cluster (Cluster): The cluster to be analyzed. - - Returns: - None - """ - if self.aloCollection.tree_ete: - self.__analyse_tree_ete(cluster=cluster) - - self.__process_attributes(cluster) - self.__finalize_cluster_analysis(cluster) - - # write output - # 0. __get_header_line - def __get_header_line(self, filetype: str, attribute: str) -> str: - """ - Generates a header line for different types of file formats based on the provided - `filetype` and `attribute`. - - Args: - filetype (str): The type of file for which the header line is generated. Valid values: - - "attribute_metrics": Header line for attribute metrics. - - "cafe": Header line for CAFE analysis. - - "cluster_1to1s_ALO": Header line for cluster 1-to-1 relationships with ALO. - - "cluster_metrics": Header line for cluster metrics. - - "cluster_metrics_ALO": Header line for cluster metrics with ALO. - - "cluster_metrics_domains": Header line for cluster metrics with domains. - - "cluster_metrics_domains_detailed": Header line for detailed cluster metrics with domains. - - "pairwise_representation_test": Header line for pairwise representation test. - - attribute (str): The attribute associated with the cluster, used in certain file types. - - Returns: - str: The generated header line as a tab-separated string. - - Raises: - ValueError: If `filetype` is not recognized. - """ - if filetype == "attribute_metrics": - attribute_metrics_header = [ - "#attribute", - "taxon_set", - "cluster_total_count", - "protein_total_count", - "protein_total_span", - "singleton_cluster_count", - "singleton_protein_count", - "singleton_protein_span", - "specific_cluster_count", - "specific_protein_count", - "specific_protein_span", - "shared_cluster_count", - "shared_protein_count", - "shared_protein_span", - "specific_cluster_true_1to1_count", - "specific_cluster_fuzzy_count", - "shared_cluster_true_1to1_count", - "shared_cluster_fuzzy_count", - "absent_cluster_total_count", - "absent_cluster_singleton_count", - "absent_cluster_specific_count", - "absent_cluster_shared_count", - "TAXON_count", - "TAXON_taxa", - ] - return "\t".join(attribute_metrics_header) - elif filetype == "cafe": - cafe_header = ["#ID"] - cafe_header.extend( - iter(sorted(self.aloCollection.ALO_by_level_by_attribute["taxon"])) - ) - return "\t".join(cafe_header) - elif filetype == "cluster_1to1s_ALO": - cluster_1to1s_ALO_header = [ - "#cluster_id", - "cluster_type", - "1to1_type", - "proteome_count", - "percentage_at_target_count", - ] - return "\t".join(cluster_1to1s_ALO_header) - elif filetype == "cluster_metrics": - cluster_metrics_header = [ - "#cluster_id", - "cluster_protein_count", - "protein_median_count", - "TAXON_count", - "attribute", - "attribute_cluster_type", - "protein_span_mean", - "protein_span_sd", - ] - cluster_metrics_header += [ - f"{level}_count" - for level in sorted( - self.aloCollection.ALO_by_level_by_attribute[attribute] - ) - ] - if attribute.lower() != "taxon": - cluster_metrics_header += [ - f"{level}_median" - for level in sorted( - self.aloCollection.ALO_by_level_by_attribute[attribute] - ) - ] - cluster_metrics_header += [ - f"{level}_cov" - for level in sorted( - self.aloCollection.ALO_by_level_by_attribute[attribute] - ) - ] - return "\t".join(cluster_metrics_header) - elif filetype == "cluster_metrics_ALO": - cluster_metrics_ALO_header = [ - "#cluster_id", - "cluster_status", - "cluster_type", - "cluster_protein_count", - "cluster_proteome_count", - "TAXON_protein_count", - "TAXON_mean_count", - "non_taxon_mean_count", - "representation", - "log2_mean(TAXON/others)", - "pvalue(TAXON vs. others)", - "TAXON_coverage", - "TAXON_count", - "non_TAXON_count", - "TAXON_taxa", - "non_TAXON_taxa", - ] - # for domain_source in clusterCollection.domain_sources: - # cluster_metrics_ALO_header.append(domain_source) - return "\t".join(cluster_metrics_ALO_header) - elif filetype == "cluster_metrics_domains": - cluster_metrics_domains_header = [ - "#cluster_id", - "cluster_protein_count", - "TAXON_count", - "protein_span_mean", - "protein_span_sd", - "fraction_secreted", - ] - for domain_source in self.clusterCollection.domain_sources: - cluster_metrics_domains_header.extend( - (domain_source, f"{domain_source}_entropy") - ) - return "\t".join(cluster_metrics_domains_header) - elif filetype == "cluster_metrics_domains_detailed": - cluster_metrics_domains_detailed_header = [ - "#cluster_id", - "domain_source", - "domain_id", - "domain_description", - "protein_count", - "protein_count_with_domain", - "TAXA_with_domain_fraction", - "TAXA_with_domain", - "TAXA_without_domain", - ] - return "\t".join(cluster_metrics_domains_detailed_header) - elif filetype == "pairwise_representation_test": - pairwise_representation_test_header = [ - "#cluster_id", - "TAXON_1", - "TAXON_1_mean", - "TAXON_2", - "TAXON_2_mean", - "log2_mean(TAXON_1/TAXON_2)", - "mwu_pvalue(TAXON_1 vs. TAXON_2)", - ] - # pairwise_representation_test_header.append("go_terms") - # for domain_source in clusterCollection.domain_sources: - # pairwise_representation_test_header.append(domain_source) - return "\t".join(pairwise_representation_test_header) - else: - error_msg = f"[ERROR] {filetype} is not a valid header 'filetype'" - raise ValueError(error_msg) - - # 1. plot_cluster_sizes - def __plot_cluster_sizes(self) -> None: - """ - Plot the distribution of cluster sizes based on the protein counts in each cluster. - - Saves the plot as a figure in the directory specified by self.dirs["main"]. - - Returns: - None - - Raises: - ValueError: If self.inputData.plot_format is not a valid file format. - """ - cluster_protein_count = [ - cluster.protein_count for cluster in self.clusterCollection.cluster_list - ] - cluster_protein_counter = Counter(cluster_protein_count) - count_plot_f = os.path.join( - self.dirs["main"], - f"cluster_size_distribution.{self.inputData.plot_format}", - ) - f, ax = plt.subplots(figsize=self.inputData.plotsize) - ax.set_facecolor("white") - x_values = [] - y_values = [] - for value, count in list(cluster_protein_counter.items()): - x_values.append(value) - y_values.append(count) - x_array = np.array(x_values) # type: ignore - y_array = np.array(y_values) - ax.scatter(x_array, y_array, marker="o", alpha=0.8, s=100) # type: ignore - ax.set_xlabel("Cluster size", fontsize=self.inputData.fontsize) - ax.set_ylabel("Count", fontsize=self.inputData.fontsize) - ax.set_yscale("log") - ax.set_xscale("log") - plt.margins(0.8) - plt.gca().set_ylim(bottom=0.8) - plt.gca().set_xlim(left=0.8) - ax.xaxis.set_major_formatter(FormatStrFormatter("%.0f")) - ax.yaxis.set_major_formatter(FormatStrFormatter("%.0f")) - f.tight_layout() - - ax.grid(True, linewidth=1, which="major", color="lightgrey") - ax.grid(True, linewidth=0.5, which="minor", color="lightgrey") - logger.info(f"[STATUS] - Plotting {count_plot_f}") - f.savefig(count_plot_f, format=self.inputData.plot_format) - plt.close() - - # 2. write_cluster_counts_by_taxon - def __write_cluster_counts_by_taxon(self) -> None: - """ - Write cluster counts by taxon attribute to a text file. - - This method iterates through attributes in self.aloCollection.attributes, - retrieves protein counts by level for clusters in self.clusterCollection.cluster_list - that match the attribute "taxon", and writes the data to a text file named - 'cluster_counts_by_taxon.txt' in the directory specified by self.dirs["main"]. - - Raises: - ValueError: If the header type 'cafe' is not recognized. - """ - cafe_f = os.path.join(self.dirs["main"], "cluster_counts_by_taxon.txt") - for attribute in self.aloCollection.attributes: - levels = sorted( - list(self.aloCollection.ALO_by_level_by_attribute[attribute]) - ) - cafe_output = [] - for cluster in self.clusterCollection.cluster_list: - if attribute.lower() == "taxon": - cafe_line = f"{cluster.cluster_id}" - # cafe_line.append("None") - for _level in levels: - total_proteins = sum( - cluster.protein_counts_of_proteomes_by_level_by_attribute[ - attribute - ][_level] - ) - cafe_line += f"\t{total_proteins}" - cafe_output.append(cafe_line) - if cafe_output: - with open(cafe_f, "w") as cafe_fh: - logger.info(f"[STATUS] - Writing {cafe_f}") - cafe_output.sort() - cafe_output.insert(0, self.__get_header_line("cafe", "taxon")) - cafe_fh.write("\n".join(cafe_output) + "\n") - cafe_output = [] - - # 3. write_cluster_metrics_domains - def __write_cluster_metrics_domains(self) -> None: - """ - Write cluster metrics to a file 'cluster_metrics_domains.txt'. - - This method constructs and writes cluster metrics data to a text file, - including cluster IDs, protein counts, taxon counts, domain statistics, - and entropy for each domain source present in the cluster collection. - - Raises: - IOError: If there is an issue writing to the output file. - - """ - cluster_metrics_domains_f = os.path.join( - self.dirs["main"], "cluster_metrics_domains.txt" - ) - header = self.__get_header_line("cluster_metrics_domains", "taxon").split("\t") - cluster_metrics_domains_output = [] - - if self.clusterCollection.functional_annotation_parsed: - for cluster in self.clusterCollection.cluster_list: - line_parts = { - "#cluster_id": cluster.cluster_id, - "cluster_protein_count": str(cluster.protein_count), - "TAXON_count": str(cluster.proteome_count), - "protein_span_mean": "N/A", - "protein_span_sd": "N/A", - "fraction_secreted": "N/A", - } - - if ( - self.clusterCollection.fastas_parsed - and cluster.protein_length_stats - ): - line_parts["protein_span_mean"] = str( - cluster.protein_length_stats["mean"] - ) - line_parts["protein_span_sd"] = str( - cluster.protein_length_stats["sd"] - ) - - if "SignalP_EUK" in self.clusterCollection.domain_sources: - line_parts["fraction_secreted"] = "{0:.2f}".format( - cluster.secreted_cluster_coverage - ) - - for domain_source in self.clusterCollection.domain_sources: - if domain_source in cluster.domain_counter_by_domain_source: - sorted_counts = sorted( - [ - f"{domain_id}:{count}" - for domain_id, count in cluster.domain_counter_by_domain_source[ - domain_source - ].most_common() - ], - key=lambda x: (x.split(":")[-1], x.split(":")[-2]), - ) - line_parts[domain_source] = ";".join(sorted_counts) - line_parts[f"{domain_source}_entropy"] = "{0:.3f}".format( - cluster.domain_entropy_by_domain_source[domain_source] - ) - else: - line_parts[domain_source] = "N/A" - line_parts[f"{domain_source}_entropy"] = "N/A" - - # Ensure we're following the correct order from the header - ordered_line = [line_parts.get(col, "N/A") for col in header] - cluster_metrics_domains_output.append("\t".join(ordered_line)) - - if cluster_metrics_domains_output: - with open(cluster_metrics_domains_f, "w") as cluster_metrics_domains_fh: - logger.info(f"[STATUS] - Writing {cluster_metrics_domains_f}") - cluster_metrics_domains_output.sort() - cluster_metrics_domains_output.insert(0, "\t".join(header)) - cluster_metrics_domains_fh.write( - "\n".join(cluster_metrics_domains_output) + "\n" - ) - - # 4. write_cluster_metrics_domains_detailed - def __count_proteins_with_domain( - self, cluster: Cluster, domain_source: str, domain_id: str - ) -> Tuple[Dict[str, int], Dict[str, int]]: - """ - Count proteins with and without a specific domain in each proteome of a cluster. - - Args: - cluster (Cluster): The cluster object containing proteins to be analyzed. - domain_source (str): The source of the domain to be counted (e.g., "Pfam", "InterPro"). - domain_id (str): The ID of the specific domain to be counted. - - Returns: - Tuple[Dict[str, int], Dict[str, int]]: A tuple containing: - - A dictionary where keys are proteome IDs and values are counts of proteins - in the proteome that have the specified domain (`with_domain`). - - A dictionary where keys are proteome IDs and values are counts of proteins - in the proteome that do not have the specified domain (`without_domain`). - - """ - with_domain = defaultdict(int) - without_domain = defaultdict(int) - - for proteome_id, protein_ids in cluster.protein_ids_by_proteome_id.items(): - for protein_id in protein_ids: - protein = self.proteinCollection.proteins_by_protein_id[protein_id] - if ( - domain_source in protein.domain_counter_by_domain_source - and domain_id - in protein.domain_counter_by_domain_source[domain_source] - ): - with_domain[proteome_id] += 1 - else: - without_domain[proteome_id] += 1 - - return with_domain, without_domain - - def __format_proteome_counts( - self, count_dict: Dict[str, int], cluster: Cluster - ) -> str: - """ - Format proteome counts into a string representation. - - Args: - count_dict (Dict[str, int]): A dictionary where keys are proteome IDs and values are counts. - cluster (Cluster): The cluster object associated with the counts. - - Returns: - str: A string representation of proteome counts formatted as "proteome_id:count/total" - for each proteome ID in sorted order, separated by commas. If count_dict is empty, - returns "N/A". - - """ - return ( - ",".join( - f"{proteome_id}:{count}/{len(cluster.protein_ids_by_proteome_id[proteome_id])}" - for proteome_id, count in sorted(count_dict.items()) - ) - or "N/A" - ) - - def __get_domain_description(self, domain_source: str, domain_id: str) -> str: - """ - Get the description of a domain based on its source and ID. - - Args: - domain_source (str): The source of the domain (e.g., "SignalP_EUK", "Pfam"). - domain_id (str): The ID of the domain whose description is to be retrieved. - - Returns: - str: The description of the domain if found in `self.proteinCollection.domain_desc_by_id_by_source`, - otherwise returns "N/A". - - """ - if domain_source == "SignalP_EUK": - return domain_id - return self.proteinCollection.domain_desc_by_id_by_source.get( - domain_source, {} - ).get(domain_id, "N/A") - - def __process_cluster_domains( - self, cluster: Cluster, output_by_domain_source: Dict[str, List[str]] - ) -> None: - """ - Process domain statistics for a cluster and populate the output dictionary. - - Args: - cluster (Cluster): The cluster object containing domain statistics to process. - output_by_domain_source (Dict[str, List[str]]): A dictionary where keys are domain sources - and values are lists of output lines to be populated with processed domain statistics. - - Returns: - None - - """ - for ( - domain_source, - domain_counter, - ) in cluster.domain_counter_by_domain_source.items(): - for domain_id, count in domain_counter.most_common(): - with_domain, without_domain = self.__count_proteins_with_domain( - cluster, domain_source, domain_id - ) - proteome_count_with_domain = sum( - count > 0 for count in with_domain.values() - ) - - with_domain_str = self.__format_proteome_counts(with_domain, cluster) - without_domain_str = self.__format_proteome_counts( - without_domain, cluster - ) - - domain_description = self.__get_domain_description( - domain_source, domain_id - ) - - output_line = ( - f"{cluster.cluster_id}\t{domain_source}\t{domain_id}\t" - f"{domain_description}\t{cluster.protein_count}\t" - f"{sum(with_domain.values())}\t" - f"{proteome_count_with_domain / cluster.proteome_count:.3f}\t" - f"{with_domain_str}\t{without_domain_str}" - ) - - output_by_domain_source[domain_source].append(output_line) - - def __write_domain_outputs( - self, - output_by_domain_source: Dict[str, List[str]], - output_files: Dict[str, str], - ) -> None: - """ - Write domain outputs to respective output files. - - Args: - output_by_domain_source (Dict[str, List[str]]): A dictionary where keys are domain sources - and values are lists of output lines to be written to output files. - output_files (Dict[str, str]): A dictionary where keys are domain sources and values are - corresponding output file paths. - - Returns: - None - - """ - for domain_source, output_lines in output_by_domain_source.items(): - if len(output_lines) > 1: - output_file = output_files[domain_source] - logger.info(f"[STATUS] - Writing {output_file}") - with open(output_file, "w") as fh: - fh.write("\n".join(output_lines) + "\n") - - def __write_cluster_metrics_domains_detailed(self) -> None: - """ - Write detailed cluster metrics for domain annotations to respective output files. - - This method constructs detailed cluster metrics for domain annotations and writes - them to individual output files for each domain source specified in the cluster - collection. - - Returns: - None - """ - output_by_domain_source: Dict[str, List[str]] = { - source: [] for source in self.clusterCollection.domain_sources - } - - output_files: Dict[str, str] = { - source: os.path.join( - self.dirs["main"], f"cluster_domain_annotation.{source}.txt" - ) - for source in self.clusterCollection.domain_sources - } - - if self.clusterCollection.functional_annotation_parsed: - for cluster in self.clusterCollection.cluster_list: - self.__process_cluster_domains(cluster, output_by_domain_source) - - self.__write_domain_outputs(output_by_domain_source, output_files) - - # 5. write attribute metrics - def __get_attribute_metrics(self, ALO: AttributeLevel) -> str: - """ - Retrieve attribute metrics as a formatted string. - - Args: - ALO (AttributeLevel): An instance of AttributeLevel containing the attribute metrics. - - Returns: - str: A tab-separated string containing various attribute metrics: - - Attribute name - - Attribute level - - Cluster counts and protein counts/span for different cluster types and statuses. - - Proteome count and other relevant metrics. - - """ - attribute_metrics = [ - ALO.attribute, - ALO.level, - ALO.get_cluster_count_by_cluster_status_by_cluster_type("present", "total"), - ALO.get_protein_count_by_cluster_type("total"), - ALO.get_protein_span_by_cluster_type("total"), - ALO.get_cluster_count_by_cluster_status_by_cluster_type( - "present", "singleton" - ), - ALO.get_protein_count_by_cluster_type("singleton"), - ALO.get_protein_span_by_cluster_type("singleton"), - ALO.get_cluster_count_by_cluster_status_by_cluster_type( - "present", "specific" - ), - ALO.get_protein_count_by_cluster_type("specific"), - ALO.get_protein_span_by_cluster_type("specific"), - ALO.get_cluster_count_by_cluster_status_by_cluster_type( - "present", "shared" - ), - ALO.get_protein_count_by_cluster_type("shared"), - ALO.get_protein_span_by_cluster_type("shared"), - ALO.get_cluster_count_by_cluster_cardinality_by_cluster_type( - "specific", "true" - ), - ALO.get_cluster_count_by_cluster_cardinality_by_cluster_type( - "specific", "fuzzy" - ), - ALO.get_cluster_count_by_cluster_cardinality_by_cluster_type( - "shared", "true" - ), - ALO.get_cluster_count_by_cluster_cardinality_by_cluster_type( - "shared", "fuzzy" - ), - ALO.get_cluster_count_by_cluster_status_by_cluster_type("absent", "total"), - ALO.get_cluster_count_by_cluster_status_by_cluster_type( - "absent", "singleton" - ), - ALO.get_cluster_count_by_cluster_status_by_cluster_type( - "absent", "specific" - ), - ALO.get_cluster_count_by_cluster_status_by_cluster_type("absent", "shared"), - ALO.proteome_count, - ALO.get_proteomes(), - ] - - return "\t".join(map(str, attribute_metrics)) - - def __write_attribute_metrics(self) -> None: - """ - Write attribute metrics for each attribute to respective output files. - - This method iterates over each attribute in self.aloCollection.attributes, - retrieves attribute metrics for each level of the attribute, and writes them - to individual output files named after the attribute. - - Returns: - None - - """ - for attribute in self.aloCollection.attributes: - attribute_metrics_f = os.path.join( - self.dirs[attribute], f"{attribute}.attribute_metrics.txt" - ) - attribute_metrics_output = [] - levels = sorted( - list(self.aloCollection.ALO_by_level_by_attribute[attribute]) - ) - for level in levels: - if ALO := self.aloCollection.ALO_by_level_by_attribute[attribute][ - level - ]: - attribute_metrics_output.append(self.__get_attribute_metrics(ALO)) - - if attribute_metrics_output: - with open(attribute_metrics_f, "w") as attribute_metrics_fh: - logger.info(f"[STATUS] - Writing {attribute_metrics_f}") - attribute_metrics_output.sort() - header_line = self.__get_header_line("attribute_metrics", attribute) - attribute_metrics_output.insert(0, header_line) - attribute_metrics_fh.write( - "\n".join(attribute_metrics_output) + "\n" - ) - - # 6. write cluster summary - def __write_cluster_summary(self) -> None: - """ - Write cluster summary metrics for each attribute to respective output files. - - This method iterates over each attribute in self.aloCollection.attributes, - retrieves cluster summary metrics for each cluster in self.clusterCollection.cluster_list, - and writes them to individual output files named after the attribute. - - Returns: - None - - """ - for attribute in self.aloCollection.attributes: - cluster_metrics_f = os.path.join( - self.dirs[attribute], f"{attribute}.cluster_summary.txt" - ) - - levels = sorted( - list(self.aloCollection.ALO_by_level_by_attribute[attribute]) - ) - cluster_metrics_output = [] - for cluster in self.clusterCollection.cluster_list: - cluster_metrics_line = [ - str(cluster.cluster_id), - str(cluster.protein_count), - str(cluster.protein_median), - str(cluster.proteome_count), - str(attribute), - str(cluster.cluster_type_by_attribute[attribute]), - ] - if ( - self.clusterCollection.fastas_parsed - and cluster.protein_length_stats - ): - cluster_metrics_line.extend( - [ - str(cluster.protein_length_stats.get("mean", "N/A")), - str(cluster.protein_length_stats.get("sd", "N/A")), - ] - ) - else: - cluster_metrics_line.extend(["N/A", "N/A"]) - - cluster_metrics_line.extend( - str( - sum( - cluster.protein_counts_of_proteomes_by_level_by_attribute[ - attribute - ][_level] - ) - ) - for _level in levels - ) - - if attribute.lower() != "taxon": - cluster_metrics_line.extend( - [ - str( - median( - cluster.protein_counts_of_proteomes_by_level_by_attribute[ - attribute - ][ - _level - ] - ) - ) - for _level in levels - ] - ) - cluster_metrics_line.extend( - [ - "{0:.2f}".format( - cluster.proteome_coverage_by_level_by_attribute[ - attribute - ][_level] - ) - for _level in levels - ] - ) - - cluster_metrics_output.append("\t".join(cluster_metrics_line)) - - if cluster_metrics_output: - with open(cluster_metrics_f, "w") as cluster_metrics_fh: - logger.info(f"[STATUS] - Writing {cluster_metrics_f}") - cluster_metrics_output.sort() - header_line = self.__get_header_line("cluster_metrics", attribute) - cluster_metrics_output.insert(0, header_line) - cluster_metrics_fh.write("\n".join(cluster_metrics_output) + "\n") - cluster_metrics_output = [] - - # 7. Write cluster ALO metrics - def __get_enrichment_data(self, ALO: AttributeLevel, cluster: Cluster) -> List[str]: - """ - Retrieve enrichment data for a given AttributeLevel and Cluster. - - Args: - ALO (AttributeLevel): An instance of AttributeLevel containing enrichment data. - cluster (Cluster): An instance of Cluster for which enrichment data is retrieved. - - Returns: - List[str]: A list containing enrichment data: - - Enrichment status ("enriched", "depleted", "equal" or "N/A" if unavailable) - - Log2 mean value - - p-value - - """ - if ( - ALO - and ALO.cluster_type_by_cluster_id[cluster.cluster_id] == "shared" - and ALO.cluster_mwu_log2_mean_by_cluster_id[cluster.cluster_id] - ): - log2_mean = ALO.cluster_mwu_log2_mean_by_cluster_id[cluster.cluster_id] - enrichment = ( - "enriched" - if log2_mean > 0 - else "depleted" if log2_mean < 0 else "equal" - ) - return [ - enrichment, - f"{log2_mean}", - f"{ALO.cluster_mwu_pvalue_by_cluster_id[cluster.cluster_id]}", - ] - return ["N/A", "N/A", "N/A"] - - def __get_proteome_data(self, ALO: AttributeLevel, cluster: Cluster) -> List[str]: - """ - Retrieve proteome data for a given AttributeLevel and Cluster. - - Args: - ALO (AttributeLevel): An instance of AttributeLevel containing proteome data. - cluster (Cluster): An instance of Cluster for which proteome data is retrieved. - - Returns: - List[str]: A list containing proteome data: - - Number of proteomes present in both ALO and cluster - - Number of proteomes present only in cluster - - Sorted list of proteome IDs present in both ALO and cluster, or "N/A" if none - - Sorted list of proteome IDs present only in cluster, or "N/A" if none - - """ - ALO_proteomes_present = cluster.proteome_ids.intersection( - ALO.proteomes if ALO else set() - ) - non_ALO_proteomes_present = cluster.proteome_ids.difference( - ALO.proteomes if ALO else set() - ) - return [ - f"{len(ALO_proteomes_present)}", - f"{len(non_ALO_proteomes_present)}", - ( - f"{','.join(sorted(list(ALO_proteomes_present)))}" - if ALO_proteomes_present - else "N/A" - ), - ( - f"{','.join(sorted(list(non_ALO_proteomes_present)))}" - if non_ALO_proteomes_present - else "N/A" - ), - ] - - def __write_cluster_metrics_ALO(self) -> None: - """ - Write cluster metrics for each attribute level object (ALO) to separate files. - - For each attribute in self.aloCollection.attributes, this method writes cluster metrics - to a file named '{attribute}.{level}.cluster_metrics.txt' in the corresponding directory - under self.dirs[attribute]. - - Metrics include cluster ID, status, type, protein count, proteome count, counts by level, - mean ALO counts, mean non-ALO counts, enrichment data, and proteome coverage. - - Returns: - None - """ - for attribute in self.aloCollection.attributes: - levels = sorted( - list(self.aloCollection.ALO_by_level_by_attribute[attribute]) - ) - - for level in levels: - ALO = self.aloCollection.ALO_by_level_by_attribute[attribute][level] - cluster_metrics_ALO_f = os.path.join( - self.dirs[attribute], f"{attribute}.{level}.cluster_metrics.txt" - ) - if ALO is None: - continue - cluster_metrics_ALO_output = [ - "\t".join( - [ - f"{cluster.cluster_id}", - ( - f"{ALO.cluster_status_by_cluster_id[cluster.cluster_id]}" - if ALO - else "N/A" - ), - ( - f"{ALO.cluster_type_by_cluster_id[cluster.cluster_id]}" - if ALO - else "N/A" - ), - f"{cluster.protein_count}", - f"{cluster.proteome_count}", - f"{sum(cluster.protein_counts_of_proteomes_by_level_by_attribute[attribute][level])}", - ( - f"{ALO.cluster_mean_ALO_count_by_cluster_id[cluster.cluster_id]}" - if ALO - and ALO.cluster_mean_ALO_count_by_cluster_id[ - cluster.cluster_id - ] - else "N/A" - ), - ( - f"{ALO.cluster_mean_non_ALO_count_by_cluster_id[cluster.cluster_id]}" - if ALO - and ALO.cluster_mean_non_ALO_count_by_cluster_id[ - cluster.cluster_id - ] - else "N/A" - ), - *self.__get_enrichment_data(ALO, cluster), - "{0:.2f}".format( - cluster.proteome_coverage_by_level_by_attribute[ - attribute - ][level] - ), - *self.__get_proteome_data(ALO, cluster), - ] - ) - for cluster in self.clusterCollection.cluster_list - ] - if cluster_metrics_ALO_output: - with open(cluster_metrics_ALO_f, "w") as cluster_metrics_ALO_fh: - logger.info(f"[STATUS] - Writing {cluster_metrics_ALO_f}") - cluster_metrics_ALO_output.sort() - - header_line = self.__get_header_line( - "cluster_metrics_ALO", attribute - ) - cluster_metrics_ALO_output.insert(0, header_line) - cluster_metrics_ALO_fh.write( - "\n".join(cluster_metrics_ALO_output) + "\n" - ) - - # 8. write cluster 1to1 ALO - def __write_cluster_1to1_ALO(self) -> None: - """ - Write cluster 1-to-1 relationships for each attribute level object (ALO) to separate files. - - For each attribute in self.aloCollection.attributes, this method writes cluster 1-to-1 - relationships to a file named '{attribute}.{level}.cluster_1to1s.txt' in the corresponding - directory under self.dirs[attribute]. - - Relationships include cluster ID, type, cardinality, proteome count, and fuzzy count ratio. - - Returns: - None - """ - for attribute in self.aloCollection.attributes: - levels = sorted( - list(self.aloCollection.ALO_by_level_by_attribute[attribute]) - ) - for level in levels: - cluster_1to1_ALO_f = os.path.join( - self.dirs[attribute], f"{attribute}.{level}.cluster_1to1s.txt" - ) - cluster_1to1_ALO_output = [] - - ALO = self.aloCollection.ALO_by_level_by_attribute[attribute][level] - - if attribute.lower() != "taxon" and ALO: - for ( - cluster_type - ) in ALO.clusters_by_cluster_cardinality_by_cluster_type: - for ( - cluster_cardinality - ) in ALO.clusters_by_cluster_cardinality_by_cluster_type[ - cluster_type - ]: - for ( - cluster_id - ) in ALO.clusters_by_cluster_cardinality_by_cluster_type[ - cluster_type - ][ - cluster_cardinality - ]: - cluster = ( - self.clusterCollection.cluster_list_by_cluster_id[ - cluster_id - ] - ) - protein_count_by_proteome = ( - cluster.protein_count_by_proteome_id - ) - proteome_count = cluster.proteome_count - - fuzzy_proteome_ratio = ( - len( - [ - protein_count - for _, protein_count in protein_count_by_proteome.items() - if protein_count - == self.inputData.fuzzy_count - ] - ) - / proteome_count - ) - - cluster_1to1_ALO_line = "\t".join( - [ - str(cluster_id), - str(cluster_type), - str(cluster_cardinality), - str(proteome_count), - "{0:.2f}".format(fuzzy_proteome_ratio), - ] - ) - - cluster_1to1_ALO_output.append(cluster_1to1_ALO_line) - - if cluster_1to1_ALO_output: - with open(cluster_1to1_ALO_f, "w") as cluster_1to1_ALO_fh: - logger.info(f"[STATUS] - Writing {cluster_1to1_ALO_f}") - cluster_1to1_ALO_output.sort() - header_line = self.__get_header_line( - "cluster_1to1s_ALO", attribute - ) - cluster_1to1_ALO_output.insert(0, header_line) - cluster_1to1_ALO_fh.write( - "\n".join(cluster_1to1_ALO_output) + "\n" - ) - cluster_1to1_ALO_output = [] - - # 9. write_pairwise_representation - def __process_background_representation( - self, - attribute: str, - level: str, - ALO: AttributeLevel, - cluster: Cluster, - background_representation_test_by_pair_by_attribute, - ) -> None: - """ - Process and append background representation test results for a cluster and attribute level. - - Args: - attribute (str): The attribute name. - level (str): The attribute level. - ALO (AttributeLevel): The AttributeLevel object for the attribute and level. - cluster (Cluster): The Cluster object representing the cluster. - background_representation_test_by_pair_by_attribute (Dict[str, Dict[str, Any]]): - A nested dictionary to store background representation test results, - structured as [attribute][background_pair] = list of test results. - - Returns: - None - """ - background_pair = (level, "background") - if attribute not in background_representation_test_by_pair_by_attribute: - background_representation_test_by_pair_by_attribute[attribute] = {} - if ( - background_pair - not in background_representation_test_by_pair_by_attribute[attribute] - ): - background_representation_test_by_pair_by_attribute[attribute][ - background_pair - ] = [] - - background_representation_test = [ - cluster.cluster_id, - level, - "background", - ALO.cluster_mean_ALO_count_by_cluster_id[cluster.cluster_id], - ALO.cluster_mean_non_ALO_count_by_cluster_id[cluster.cluster_id], - ALO.cluster_mwu_log2_mean_by_cluster_id[cluster.cluster_id], - ALO.cluster_mwu_pvalue_by_cluster_id[cluster.cluster_id], - ] - background_representation_test_by_pair_by_attribute[attribute][ - background_pair - ].append(background_representation_test) - - def __get_pairwise_representation_test( - self, - cluster: Cluster, - attribute: str, - level: str, - levels_seen: Set[str], - levels: List[str], - ) -> Generator[List[Any], None, None]: - """ - Generate pairwise representation test results for a cluster and attribute level. - - Args: - cluster (Cluster): The Cluster object representing the cluster. - attribute (str): The attribute name. - level (str): The current attribute level. - levels_seen (Set[str]): A set of attribute levels already processed. - levels (List[str]): A list of all attribute levels. - - Yields: - Generator[List[Any], None, None]: A generator yielding lists containing pairwise representation test results. - Each list includes: - - cluster.cluster_id: ID of the cluster. - - level: Current attribute level. - - other_level: Another attribute level being compared with `level`. - - mean_ALO_count: Mean count of ALOs in the cluster at `level`. - - mean_non_ALO_count: Mean count of non-ALOs in the cluster at `level`. - - mwu_log2_mean: Log2 mean of the Mann-Whitney U test results between `level` and `other_level`. - - mwu_pvalue: P-value of the Mann-Whitney U test results between `level` and `other_level`. - """ - for other_level in set(levels).difference(levels_seen): - if other_level != level: - other_ALO = self.aloCollection.ALO_by_level_by_attribute[attribute][ - other_level - ] - if ( - other_ALO - and len(cluster.proteome_ids.intersection(other_ALO.proteomes)) >= 2 - ): - protein_counts_level = [ - count - for count in cluster.protein_counts_of_proteomes_by_level_by_attribute[ - attribute - ][ - level - ] - if count > 0 - ] - protein_counts_other_level = [ - count - for count in cluster.protein_counts_of_proteomes_by_level_by_attribute[ - attribute - ][ - other_level - ] - if count > 0 - ] - if protein_counts_level and protein_counts_other_level: - ( - mwu_pvalue, - mwu_log2_mean, - mean_ALO_count, - mean_non_ALO_count, - ) = statistic( - protein_counts_level, - protein_counts_other_level, - self.inputData.test, - self.inputData.min_proteomes, - ) - yield [ - cluster.cluster_id, - level, - other_level, - mean_ALO_count, - mean_non_ALO_count, - mwu_log2_mean, - mwu_pvalue, - ] - # pvalue = None - # try: - # pvalue = scipy.stats.mannwhitneyu(protein_counts_level, protein_counts_other_level, alternative="two-sided")[1] - # except: - # pvalue = 1.0 - # mean_level = mean(protein_counts_level) - # mean_other_level = mean(protein_counts_other_level) - # log2fc_mean = log((mean_level/mean_other_level), 2) - # yield [cluster.cluster_id, level, other_level, mean_level, - # mean_other_level, log2fc_mean, pvalue] - - def __process_pairwise_representation( - self, - attribute: str, - level: str, - levels_seen: Set[str], - levels: List[str], - cluster: Cluster, - pairwise_representation_test_by_pair_by_attribute, - pairwise_representation_test_output: List[str], - ) -> None: - """ - Process pairwise representation tests for a specific attribute level and cluster. - - Args: - attribute (str): The attribute name. - level (str): The current attribute level. - levels_seen (Set[str]): A set of attribute levels already processed. - levels (List[str]): A list of all attribute levels. - cluster (Cluster): The Cluster object representing the cluster. - pairwise_representation_test_by_pair_by_attribute (Dict[str, Dict[Tuple[str, str], List[List[Any]]]]): - Dictionary storing pairwise representation test results by attribute and pair of levels. - pairwise_representation_test_output (List[str]): List to store formatted output lines of pairwise tests. - - Returns: - None - """ - for result in self.__get_pairwise_representation_test( - cluster, attribute, level, levels_seen, levels - ): - if attribute not in pairwise_representation_test_by_pair_by_attribute: - pairwise_representation_test_by_pair_by_attribute[attribute] = {} - pair = (result[1], result[2]) - if pair not in pairwise_representation_test_by_pair_by_attribute[attribute]: - pairwise_representation_test_by_pair_by_attribute[attribute][pair] = [] - pairwise_representation_test_by_pair_by_attribute[attribute][pair].append( - result - ) - - pairwise_representation_test_output.append( - f"{result[0]}\t{result[1]}\t{result[3]}\t{result[2]}\t{result[4]}\t{result[5]}\t{result[6]}" - ) - - # 9.5 __plot_count_comparisons_volcano - def __prepare_data(self, pair_data: List[str]) -> Tuple[List[float], List[float]]: - """ - Prepare data from pair_data into lists of p-values and log2 fold change (log2fc) values. - - Args: - pair_data (List[str]): List of strings containing data for each pair. - - Returns: - Tuple[List[float], List[float]]: Tuple containing: - - List[float]: p-values extracted from pair_data. - - List[float]: log2 fold change (log2fc) values extracted from pair_data. - """ - pair_data_count = len(pair_data) - p_values: List[float] = [] - log2fc_values: List[float] = [] - - for data in pair_data: - log2fc_values.append(float(data[5])) - pvalue = data[6] if data[6] != 0.0 else 0.01 / (pair_data_count + 1) - p_values.append(float(pvalue)) - - return p_values, log2fc_values - - def __get_output_filename(self, attribute: str, pair_list: List[str]) -> str: - """ - Generate an output filename based on attribute, pair_list, and plot_format. - - Args: - attribute (str): Attribute name used in the filename. - pair_list (List[str]): List of strings used to form part of the filename. - - Returns: - str: Generated output filename. - """ - return os.path.join( - self.dirs[attribute], - f"{attribute}.pairwise_representation_test.{'_'.join(pair_list)}.{self.inputData.plot_format}", - ) - - def __create_volcano_plot( - self, - p_values: List[float], - log2fc_values: List[float], - pair_list: List[str], - output_file: str, - ) -> None: - """ - Create a volcano plot to visualize differential expression analysis results. - - Parameters: - - p_values (List[float]): List of p-values for each comparison. - - log2fc_values (List[float]): List of log2 fold change values for each comparison. - - pair_list (List[str]): List of pairs or labels corresponding to each comparison. - - output_file (str): Filepath where the plot will be saved. - - Returns: - - None - """ - plt.figure(1, figsize=self.inputData.plotsize) - - axScatter, axHistx = self.__setup_plot_axes() - - p_array = np.array(p_values) - log2fc_array = np.array(log2fc_values) - - log2fc_percentile = self.__plot_data(axScatter, axHistx, log2fc_array, p_array) - self.__set_plot_properties( - axScatter, axHistx, log2fc_array, p_array, pair_list, log2fc_percentile - ) - - logger.info(f"[STATUS] - Plotting {output_file}") - plt.savefig(output_file, format=self.inputData.plot_format) - plt.close() - - def __setup_plot_axes(self) -> Tuple[Any, Any]: - """ - Set up the axes for a combined scatter plot and histogram. - - Returns: - - Tuple of matplotlib.axes.Axes: Tuple containing the scatter plot axes (`axScatter`) - and the histogram axes (`axHistx`). - """ - left, width = 0.1, 0.65 - bottom, height = 0.1, 0.65 - bottom_h = left + width + 0.02 - axScatter = self.__extracted_from___setup_plot_axes_( - left, bottom, width, height - ) - axHistx = self.__extracted_from___setup_plot_axes_(left, bottom_h, width, 0.2) - axHistx.xaxis.set_major_formatter(NullFormatter()) - axHistx.yaxis.set_major_formatter(NullFormatter()) - - return axScatter, axHistx - - # TODO Rename this here and in `__setup_plot_axes` - def __extracted_from___setup_plot_axes_(self, left, arg1, width, arg3): - rect_scatter = left, arg1, width, arg3 - result = plt.axes(rect_scatter) - result.set_facecolor("white") - return result - - def __plot_data( - self, - axScatter: Any, - axHistx: Any, - log2fc_array: np.ndarray, - p_array: np.ndarray, - ) -> Any: - """ - Plot data on scatter and histogram axes. - - Parameters: - - axScatter (Any): Axes for the scatter plot. - - axHistx (Any): Axes for the histogram plot. - - log2fc_array (np.ndarray): Array of log2 fold change values. - - p_array (np.ndarray): Array of p-values. - - Returns: - - float: 95th percentile of log2 fold change values. - """ - # Plot histogram - binwidth = 0.05 - xymax = np.max(np.fabs(log2fc_array)) # type: ignore - lim = (int(xymax / binwidth) + 1) * binwidth - bins = np.arange(-lim, lim + binwidth, binwidth) - axHistx.hist( - log2fc_array, bins=bins, histtype="stepfilled", color="grey", align="mid" - ) - - # Plot scatter - axScatter.scatter( - log2fc_array, p_array, alpha=0.8, edgecolors="none", s=25, c="grey" - ) - - # Add reference lines - ooFive, ooOne = 0.05, 0.01 - log2fc_percentile = np.percentile(log2fc_array, 95) - - axScatter.axhline(y=ooFive, linewidth=2, color="orange", linestyle="--") - axScatter.axhline(y=ooOne, linewidth=2, color="red", linestyle="--") - axScatter.axvline(x=1.0, linewidth=2, color="purple", linestyle="--") - axScatter.axvline( - x=log2fc_percentile, linewidth=2, color="blue", linestyle="--" - ) - axScatter.axvline(x=-1.0, linewidth=2, color="purple", linestyle="--") - axScatter.axvline( - x=-log2fc_percentile, linewidth=2, color="blue", linestyle="--" - ) - - return log2fc_percentile - - def __set_plot_properties( - self, - axScatter: Any, - axHistx: Any, - log2fc_array: np.ndarray, - p_array: np.ndarray, - pair_list: List[str], - log2fc_percentile: Any, - ) -> None: - """ - Set properties and customize the appearance of the volcano plot. - - Parameters: - - axScatter (Any): Axes for the scatter plot. - - axHistx (Any): Axes for the histogram plot. - - log2fc_array (np.ndarray): Array of log2 fold change values. - - p_array (np.ndarray): Array of p-values. - - pair_list (List[str]): List of pairs or labels corresponding to each comparison. - - log2fc_percentile (Any): 95th percentile of log2 fold change values. - - Returns: - - None - """ - # Set axis limits and properties - x_min = -max(abs(np.min(log2fc_array)), abs(np.max(log2fc_array))) - x_max = -x_min - axScatter.set_xlim(x_min - 1, x_max + 1) - axScatter.grid(True, linewidth=1, which="major", color="lightgrey") - axScatter.grid(True, linewidth=0.5, which="minor", color="lightgrey") - axScatter.set_ylim(1.1, np.min(p_array) * 0.1) - axScatter.set_xlabel( - f"log2(mean({pair_list[0]})/mean({pair_list[1]}))", - fontsize=self.inputData.fontsize, - ) - axScatter.set_ylabel("p-value", fontsize=self.inputData.fontsize) - axScatter.set_yscale("log") - axHistx.set_xlim(axScatter.get_xlim()) - - # Add legend - legend_elements = [ - Line2D([0], [0], color="orange", linestyle="--", label="p-value = 0.05"), - Line2D([0], [0], color="red", linestyle="--", label="p-value = 0.01"), - Line2D([0], [0], color="purple", linestyle="--", label="|log2FC| = 1"), - Line2D( - [0], - [0], - color="blue", - linestyle="--", - label=f"|log2FC-95%ile| = {log2fc_percentile:.2f}", - ), - ] - legend = axScatter.legend( - handles=legend_elements, fontsize=self.inputData.fontsize, frameon=True - ) - legend.get_frame().set_facecolor("white") - - def __plot_count_comparisons_volcano( - self, - pairwise_representation_test_by_pair_by_attribute, - ) -> None: - """ - Generate volcano plots for count comparisons based on pairwise representation test results. - - Parameters: - - pairwise_representation_test_by_pair_by_attribute (Dict[str, Dict[Tuple[str, str], Any]]): - Dictionary containing test results organized by attribute and pair. - - Returns: - - None - """ - for attribute in pairwise_representation_test_by_pair_by_attribute: - for pair in pairwise_representation_test_by_pair_by_attribute[attribute]: - pair_list = list(pair) - pair_data = pairwise_representation_test_by_pair_by_attribute[ - attribute - ][pair] - - p_values, log2fc_values = self.__prepare_data(pair_data) - - if p_values: - output_file = self.__get_output_filename(attribute, pair_list) - self.__create_volcano_plot( - p_values, log2fc_values, pair_list, output_file - ) - - def __write_pairwise_representation(self) -> None: - """ - Process pairwise representation tests, write results, and generate volcano plots. - - Iterates through attributes in `self.aloCollection.attributes` and performs the - following steps for each attribute: - 1. Initializes dictionaries `pairwise_representation_test_by_pair_by_attribute` - and `background_representation_test_by_pair_by_attribute`. - 2. Prepares output file path (`pairwise_representation_test_f`) and header line - (`pairwise_representation_test_output`) for pairwise representation test results. - 3. Retrieves sorted levels from `self.aloCollection.ALO_by_level_by_attribute[attribute]`. - 4. Iterates through each level and processes pairwise and background representation - tests for each cluster in `self.clusterCollection.cluster_list`. - 5. Generates volcano plots using `__plot_count_comparisons_volcano` for - `background_representation_test_by_pair_by_attribute` if available. - 6. Writes pairwise representation test results to `pairwise_representation_test_f` - if data is available. - 7. Generates volcano plots using `__plot_count_comparisons_volcano` for - `pairwise_representation_test_by_pair_by_attribute` if data is available. - - Returns: - - None - """ - for attribute in self.aloCollection.attributes: - pairwise_representation_test_by_pair_by_attribute: Dict[ - str, Dict[str, str] - ] = {} - background_representation_test_by_pair_by_attribute = {} - pairwise_representation_test_output = [] - pairwise_representation_test_f = os.path.join( - self.dirs[attribute], f"{attribute}.pairwise_representation_test.txt" - ) - levels = sorted( - list(self.aloCollection.ALO_by_level_by_attribute[attribute]) - ) - levels_seen: Set[str] = set() - - for level in levels: - ALO = self.aloCollection.ALO_by_level_by_attribute[attribute][level] - - for cluster in self.clusterCollection.cluster_list: - if ( - ALO - and ALO.cluster_type_by_cluster_id[cluster.cluster_id] - == "shared" - and ALO.cluster_mwu_log2_mean_by_cluster_id[cluster.cluster_id] - ): - self.__process_background_representation( - attribute, - level, - ALO, - cluster, - background_representation_test_by_pair_by_attribute, - ) - - ALO_proteomes_present = cluster.proteome_ids.intersection( - ALO.proteomes if ALO else set("") - ) - - if ( - len(levels) > 1 - and len(ALO_proteomes_present) >= self.inputData.min_proteomes - ): - self.__process_pairwise_representation( - attribute, - level, - levels_seen, - levels, - cluster, - pairwise_representation_test_by_pair_by_attribute, - pairwise_representation_test_output, - ) - - levels_seen.add(level) - - if background_representation_test_by_pair_by_attribute: - self.__plot_count_comparisons_volcano( - background_representation_test_by_pair_by_attribute - ) - - if pairwise_representation_test_output: - with open( - pairwise_representation_test_f, "w" - ) as pairwise_representation_test_fh: - logger.info(f"[STATUS] - Writing {pairwise_representation_test_f}") - pairwise_representation_test_output.sort() - header_line = self.__get_header_line( - "pairwise_representation_test", attribute - ) - pairwise_representation_test_output.insert(0, header_line) - pairwise_representation_test_fh.write( - "\n".join(pairwise_representation_test_output) + "\n" - ) - - if pairwise_representation_test_by_pair_by_attribute: - self.__plot_count_comparisons_volcano( - pairwise_representation_test_by_pair_by_attribute - ) diff --git a/src/core/input.py b/src/core/input.py deleted file mode 100644 index 1ab4636..0000000 --- a/src/core/input.py +++ /dev/null @@ -1,77 +0,0 @@ -import os -from typing import List, Optional, Set, Tuple - - -class ServeArgs: - def __init__(self, port: int = 8000): - self.port = port - - -class InputData: - def __init__( - self, - nodesdb_f: str, - pfam_mapping_f: str, - ipr_mapping_f: str, - go_mapping_f: str, - cluster_file: str, - config_f: str, - sequence_ids_file: str, - species_ids_file: Optional[str] = None, - functional_annotation_f: Optional[str] = None, - fasta_dir: Optional[str] = None, - tree_file: Optional[str] = None, - output_path: Optional[str] = None, - infer_singletons: Optional[bool] = False, - plot_tree: bool = False, - min_proteomes: int = 2, - test: str = "mannwhitneyu", - taxranks: List[str] = None, - repetitions: int = 30, - fuzzy_count: int = 1, - fuzzy_fraction: float = 0.75, - fuzzy_range: Set[int] = {x for x in range(20 + 1) if x != 1}, - fontsize: int = 18, - plotsize: Tuple[float, float] = (24, 12), - plot_format: str = "pdf", - taxon_idx_mapping_file: Optional[str] = None, - ) -> None: - if taxranks is None: - taxranks = ["phylum", "order", "genus"] - if output_path: - if not os.path.isabs(output_path): - output_path = os.path.abspath(output_path) - else: - output_path = os.path.join(os.getcwd(), "kinfin_results") - - self.cluster_f = cluster_file - self.config_f = config_f - self.sequence_ids_f = sequence_ids_file - self.species_ids_f = species_ids_file - self.tree_f = tree_file - self.functional_annotation_f = functional_annotation_f - if config_f.endswith(".json") and not taxon_idx_mapping_file: - raise ValueError("[ERROR] - taxon_idx_mapping not present") - self.taxon_idx_mapping_file = taxon_idx_mapping_file - self.nodesdb_f = nodesdb_f - self.pfam_mapping_f = pfam_mapping_f - self.ipr_mapping_f = ipr_mapping_f - self.go_mapping_f = go_mapping_f - - self.test = test - self.plot_tree = plot_tree - self.fasta_dir = fasta_dir - self.output_path = output_path - self.infer_singletons = infer_singletons - self.fuzzy_count = fuzzy_count - self.fuzzy_fraction = fuzzy_fraction - self.fuzzy_range = fuzzy_range - self.repetitions = repetitions - self.min_proteomes = min_proteomes - self.plot_format = plot_format - self.fontsize = fontsize - self.taxranks = taxranks - self.plotsize = plotsize - - self.pfam_mapping = True - self.ipr_mapping = True diff --git a/src/core/logic.py b/src/core/logic.py deleted file mode 100644 index 93f7bca..0000000 --- a/src/core/logic.py +++ /dev/null @@ -1,500 +0,0 @@ -import contextlib -import logging -import os -from collections import defaultdict -from typing import DefaultDict, Dict, List, Literal, Optional, Set, Tuple - -import ete3 -from ete3 import Tree, TreeNode - -from core.utils import progress, read_fasta_len, yield_config_lines, yield_file_lines - -logger = logging.getLogger("kinfin_logger") - - -# common -def parse_nodesdb(filepath: str) -> Dict[str, Dict[str, str]]: - """ - Parses the nodes database file. - - Args: - filepath (str): The path to the nodes database file. - - Returns: - Dict[str, Dict[str, str]]: A dictionary containing node information. - Keys are node identifiers, and values are dictionaries with keys: - - - 'rank': The rank of the node. - - 'name': The name of the node. - - 'parent': The parent of the node. - """ - logger.info(f"[STATUS] - Parsing nodesDB {filepath}") - - nodesdb: Dict[str, Dict[str, str]] = {} - nodesdb_count = 0 - nodes_count = 0 - - for line in yield_file_lines(filepath): - if line.startswith("#"): - nodesdb_count = int(line.lstrip("# nodes_count = ").rstrip("\n")) - elif line.strip(): - nodes_count += 1 - with contextlib.suppress(Exception): - node, rank, name, parent = line.rstrip("\n").split("\t") - nodesdb[node] = {"rank": rank, "name": name, "parent": parent} - if nodesdb_count: - progress(nodes_count, 1000, nodesdb_count) - return nodesdb - - -# cli -def get_lineage( - taxid: str, - nodesdb: Dict[str, Dict[str, str]], - taxranks: List[str], -) -> Dict[str, str]: - """ - Get the lineage of a taxonomic identifier. - - Args: - taxid (str): The taxonomic identifier. - nodesdb (Dict[str, Dict[str, str]]): A dictionary containing information about nodes. - taxranks (List[str]): A list of taxonomic ranks to include in the lineage. - - Returns: - Dict[str, str]: A dictionary containing the lineage information, with taxonomic ranks as keys - and corresponding names as values. - """ - lineage = {taxrank: "undef" for taxrank in taxranks} - parent = "" - node = taxid - while parent != "1": - taxrank = nodesdb[node]["rank"] - parent = nodesdb[node]["parent"] - if taxrank in taxranks: - name = nodesdb[node]["name"] - lineage[taxrank] = name - node = parent - return lineage - - -# cli -def parse_attributes_from_config_data( - config_f: str, - taxon_idx_mapping_file: Optional[str], -) -> Tuple[Set[str], Dict[str, str], List[str], Dict[str, Dict[str, str]]]: - """ - Parses attributes from a configuration file. - - Args: - config_f (str): The path to the configuration file. - - Returns: - Tuple[Set[str], Dict[str, str], List[str], Dict[str, Dict[str, str]]]: A tuple containing: - - A set of proteome IDs. - - A dictionary mapping species IDs to proteome IDs. - - A list of attributes. - - A dictionary mapping proteome IDs to dictionaries, where each inner dictionary - maps attributes to their corresponding levels. - - Raises: - FileNotFoundError: If the specified configuration file is not found. - ValueError: If there are errors in the configuration file format or content. - - Note: - - The configuration file is expected to have a header line starting with '#', - where the first element is 'IDX' and the second element is 'TAXON'. - - Each subsequent non-empty line in the configuration file should contain - comma-separated values corresponding to the attributes defined in the header line. - - The 'TAXON' attribute is expected to be unique for each line. - """ - - try: - logger.info("[STATUS] - Parsing config data ...") - attributes: List[str] = [] - level_by_attribute_by_proteome_id: Dict[str, Dict[str, str]] = {} - proteomes: Set[str] = set() - proteome_id_by_species_id: Dict[str, str] = {} - - for line in yield_config_lines(config_f, taxon_idx_mapping_file): - if line.startswith("#"): - if not attributes: - attributes = [x.strip() for x in line.lstrip("#").split(",")] - if ( - attributes[0].upper() != "IDX" - or attributes[1].upper() != "TAXON" - ): - error_msg = f"[ERROR] - First/second element have to be IDX/TAXON.\n\t{attributes}" - logger.info(error_msg) - raise ValueError(error_msg) - elif line.strip(): - temp = line.split(",") - - if len(temp) != len(attributes): - error_msg = f"[ERROR] - number of columns in line differs from header\n\t{attributes}\n\t{temp}" - logger.info(error_msg) - raise ValueError(error_msg) - - if temp[1] in proteomes: - error_msg = f"[ERROR] - 'TAXON' should be unique. {temp[0]} was encountered multiple times" # fmt:skip - logger.info(error_msg) - raise ValueError(error_msg) - - species_id = temp[0] - proteome_id = temp[1] - proteomes.add(proteome_id) - proteome_id_by_species_id[species_id] = proteome_id - - level_by_attribute_by_proteome_id[proteome_id] = dict( - zip(attributes, temp) - ) - level_by_attribute_by_proteome_id[proteome_id]["all"] = "all" - attributes.insert(0, "all") # append to front - return ( - proteomes, - proteome_id_by_species_id, - attributes, - level_by_attribute_by_proteome_id, - ) - except Exception as e: - logger.error(f"[ERROR] - {e}") - - -# common -def add_taxid_attributes( - nodesdb_f: str, - taxranks: List[str], - attributes: List[str], - level_by_attribute_by_proteome_id: Dict[str, Dict[str, str]], -) -> Tuple[List[str], Dict[str, Dict[str, str]]]: - """ - Adds taxonomic attributes to the dictionary of attributes indexed by proteome ID. - - Parameters: - - - nodesdb_f (str): File path to the nodes database. - - taxranks (List[str]): List of taxonomic ranks to be included as attributes. - - attributes (List[str]): List of existing attributes. - - level_by_attribute_by_proteome_id (Dict[str, Dict[str, str]]): Dictionary where keys - are proteome IDs and values are dictionaries of attributes for each proteome ID, - including at least the "TAXID" attribute. - - Returns: - Tuple[List[str], Dict[str, Dict[str, str]]]: A tuple containing: - - - Updated list of attributes with taxonomic ranks added and "TAXID" removed. - - Updated dictionary of attributes indexed by proteome ID, with taxonomic attributes added and "TAXID" removed. - """ - NODESDB = parse_nodesdb(nodesdb_f) - for proteome_id in level_by_attribute_by_proteome_id: - taxid = level_by_attribute_by_proteome_id[proteome_id]["TAXID"] - lineage = get_lineage(taxid=taxid, nodesdb=NODESDB, taxranks=taxranks) - - # add lineage attribute/levels - for taxrank in taxranks: - level_by_attribute_by_proteome_id[proteome_id][taxrank] = lineage[taxrank] - - # remove taxid-levels - del level_by_attribute_by_proteome_id[proteome_id]["TAXID"] - - # remove taxid-attribute - attributes.remove("TAXID") - - # add taxranks to rank - attributes.extend(iter(taxranks)) - return attributes, level_by_attribute_by_proteome_id - - -# cli -def parse_tree_from_file( - tree_f: Optional[str], - attributes: List[str], - level_by_attribute_by_proteome_id: Dict[str, Dict[str, str]], - proteomes: Set[str], -) -> Tuple[Optional[Tree], Optional[Dict[frozenset[str], str]]]: - """ - Parse a phylogenetic tree from nwk file and set specified outgroups. - - Args: - tree_f (str): Path to the nwk tree file. - outgroups (List[str]): List of outgroup taxa names. - - Returns: - tuple[ete3.Tree, Dict[str, int]]: A tuple containing the parsed phylogenetic tree - and a dictionary mapping proteome IDs to node indices. - """ - if not tree_f: - return None, None - outgroups: List[str] = [] - if "OUT" not in attributes: - error_msg = "[ERROR] - Please specify one of more outgroup taxa" - ValueError(error_msg) - outgroups = [ - proteome_id - for proteome_id in proteomes - if level_by_attribute_by_proteome_id[proteome_id]["OUT"] == "1" - ] - logger.info(f"[STATUS] - Parsing Tree file : {tree_f} ...") - tree_ete: TreeNode = ete3.Tree(tree_f) - if len(outgroups) > 1: - outgroup_node: TreeNode = tree_ete.get_common_ancestor( - outgroups - ) # type: ignore - try: - logger.info( - f"[STATUS] - Setting LCA of {', '.join(outgroups)} as outgroup : ..." - ) - tree_ete.set_outgroup(outgroup_node) # type: ignore - except ete3.coretype.tree.TreeError: # type: ignore - logger.info("[STATUS] - Tree seems to be rooted already : ...") - else: - logger.info(f"[STATUS] - Setting {','.join(outgroups)} as outgroup : ...") - tree_ete.set_outgroup(outgroups[0]) # type: ignore - logger.info(tree_ete) - node_idx_by_proteome_ids: Dict[frozenset[str], str] = {} - for idx, node in enumerate(tree_ete.traverse("levelorder")): # type: ignore - proteome_ids = frozenset(leaf.name for leaf in node) - if not node.name: - node.add_features( - name=f"n{idx}", - nodetype="node", - proteome_ids=proteome_ids, - apomorphic_cluster_counts={"singletons": 0, "non_singletons": 0}, - synapomorphic_cluster_counts={ - "complete_presence": 0, - "partial_absence": 0, - }, - synapomorphic_cluster_strings=[], - counts={"specific": 0, "shared": 0, "absent": 0, "singleton": 0}, - ) - else: - node.add_features( - nodetype="tip", - proteome_ids=proteome_ids, - apomorphic_cluster_counts={"singletons": 0, "non_singletons": 0}, - synapomorphic_cluster_counts={ - "complete_presence": 0, - "partial_absence": 0, - }, - synapomorphic_cluster_strings=[], - counts={"specific": 0, "shared": 0, "absent": 0, "singleton": 0}, - ) - node_idx_by_proteome_ids[proteome_ids] = node.name - return tree_ete, node_idx_by_proteome_ids - - -def parse_fasta_dir(species_ids_f: str, fasta_dir: str) -> Dict[str, int]: - """ - Parse a species IDs file to retrieve fasta file names and then calculate - lengths of sequences from corresponding FASTA files. - - Args: - - species_ids_f (str): Path to the species IDs file, where each line contains - an index and a corresponding FASTA file name separated by ': '. - - fasta_dir (str): Directory path where the FASTA files are located. - - Returns: - - Dict[str, int]: A dictionary mapping header strings (protein IDs) to their - corresponding sequence lengths extracted from the FASTA files. - """ - logger.info("[STATUS] - Parsing FASTAs ...") - fasta_file_by_species_id: Dict[str, str] = {} - - for line in yield_file_lines(species_ids_f): - if not line.startswith("#"): - idx, fasta = line.split(": ") - fasta_file_by_species_id[idx] = fasta - - fasta_len_by_protein_id: Dict[str, int] = {} - for _, fasta_f in list(fasta_file_by_species_id.items()): - fasta_f = os.path.join(fasta_dir, fasta_f) - - for header, length in read_fasta_len(fasta_f): - fasta_len_by_protein_id[header] = length - - return fasta_len_by_protein_id - - -def parse_pfam_mapping(pfam_mapping_f: str) -> Dict[str, str]: - """ - Parse a PFAM mapping file to create a dictionary mapping PFAM domain IDs to their descriptions. - - Args: - - pfam_mapping_f (str): Path to the PFAM mapping file, where each line contains tab-separated values - with the domain ID in the first column and its description in the fifth column. - - Returns: - - Dict[str, str]: A dictionary mapping PFAM domain IDs to their corresponding descriptions. - - Raises: - - ValueError: If conflicting descriptions are found for the same domain ID. - """ - logger.info(f"[STATUS] - Parsing {pfam_mapping_f} ... ") - - pfam_mapping_dict: Dict[str, str] = {} - for line in yield_file_lines(pfam_mapping_f): - temp: List[str] = line.split("\t") - domain_id: str = temp[0] - domain_desc: str = temp[4] - if domain_id not in pfam_mapping_dict: - pfam_mapping_dict[domain_id] = domain_desc - elif domain_desc != pfam_mapping_dict[domain_id]: - error_msg = f"[ERROR] : Conflicting descriptions for {domain_id}" - raise ValueError(error_msg) - - return pfam_mapping_dict - - -def parse_ipr_mapping(ipr_mapping_f: str) -> Dict[str, str]: - """ - Parse an InterPro (IPR) mapping file to create a dictionary mapping InterPro IDs to their descriptions. - - Args: - - ipr_mapping_f (str): Path to the InterPro mapping file, where each line contains an InterPro ID and its description. - Lines starting with "Active_site" are skipped as they are not relevant to mapping. - - Returns: - - Dict[str, str]: A dictionary mapping InterPro IDs to their corresponding descriptions. - - Raises: - - ValueError: If conflicting descriptions are found for the same InterPro ID. - """ - logger.info(f"[STATUS] - Parsing {ipr_mapping_f} ... ") - - ipr_mapping_dict: Dict[str, str] = {} - for line in yield_file_lines(ipr_mapping_f): - if not line.startswith("Active_site"): - temp: List[str] = line.split() - ipr_id: str = temp[0] - ipr_desc: str = " ".join(temp[1:]) - if ipr_id not in ipr_mapping_dict: - ipr_mapping_dict[ipr_id] = ipr_desc - elif ipr_desc != ipr_mapping_dict[ipr_id]: - error_msg = f"[ERROR] : Conflicting descriptions for {ipr_id}" - raise ValueError(error_msg) - return ipr_mapping_dict - - -def parse_go_mapping(go_mapping_f: str) -> Dict[str, str]: - """ - Parse a Gene Ontology (GO) mapping file to create a dictionary mapping GO IDs to their descriptions. - - Args: - - go_mapping_f (str): Path to the GO mapping file, where each line contains a GO ID and its description. - Lines starting with '!' are skipped as they are comments. - - Returns: - - Dict[str, str]: A dictionary mapping GO IDs (without 'GO:' prefix) to their corresponding descriptions. - - Raises: - - ValueError: If conflicting descriptions are found for the same GO ID. - """ - logger.info(f"[STATUS] - Parsing {go_mapping_f} ... ") - go_mapping_dict: Dict[str, str] = {} - for line in yield_file_lines(go_mapping_f): - if not line.startswith("!"): - temp: List[str] = line.replace(" > ", "|").split("|") - go_string: List[str] = temp[1].split(";") - go_desc, go_id = go_string[0].replace("GO:", ""), go_string[1].lstrip(" ") - - if go_id not in go_mapping_dict: - go_mapping_dict[go_id] = go_desc - elif go_desc != go_mapping_dict[go_id]: - error_msg = f"[ERROR] : Conflicting descriptions for {go_id}" - raise ValueError(error_msg) - return go_mapping_dict - - -def compute_protein_ids_by_proteome( - proteomes_by_protein_id: Dict[str, str], -) -> DefaultDict[str, Set[str]]: - """ - Compute protein IDs grouped by proteome IDs. - - Args: - proteomes_by_protein_id (Dict[str, str]): A dictionary mapping protein IDs to proteome IDs. - - Returns: - DefaultDict[str, Set[str]]: A defaultdict where keys are proteome IDs and values are sets - of protein IDs belonging to each proteome ID. - """ - protein_ids_by_proteome_id: DefaultDict[str, Set[str]] = defaultdict(set) - for protein_id, proteome_id in list(proteomes_by_protein_id.items()): - protein_ids_by_proteome_id[proteome_id].add(protein_id) - return protein_ids_by_proteome_id - - -# common -def get_attribute_cluster_type( - singleton, - implicit_protein_ids_by_proteome_id_by_level, -) -> Literal["singleton", "shared", "specific"]: - """ - Determines the type of cluster based on the parameters. - - Parameters: - - singleton: A boolean indicating whether the cluster is a singleton. - - implicit_protein_ids_by_proteome_id_by_level: A dictionary representing protein ids - grouped by proteome id at different levels. - - Returns: - - One of the following strings: - - "singleton": If `singleton` is True. - - "shared": If there are protein ids grouped under multiple proteome ids. - - "specific": If there is only one proteome id with protein ids. - - """ - if singleton: - return "singleton" - if len(implicit_protein_ids_by_proteome_id_by_level) > 1: - return "shared" - else: - return "specific" - - -def get_ALO_cluster_cardinality( - ALO_proteome_counts_in_cluster: List[int], - fuzzy_range: Set[int], - fuzzy_count: int = 1, - fuzzy_fraction: float = 0.75, -) -> Optional[str]: - """ - Determine the cardinality type of a cluster based on ALO proteome counts. - - Args: - ALO_proteome_counts_in_cluster (List[int]): List of ALO proteome counts in the cluster. - fuzzy_range (Set[int]): Set of integers representing the range of fuzzy counts. - fuzzy_count (int, optional): Specific count considered as fuzzy. Default is 1. - fuzzy_fraction (float, optional): Fraction threshold for considering a cluster as 'fuzzy'. Default is 0.75. - - Returns: - Optional[str]: Returns "true" (str) if all counts are 1, "fuzzy" (str) if the cluster meets fuzzy criteria, - and None otherwise. - """ - if len(ALO_proteome_counts_in_cluster) > 2: - length = len(ALO_proteome_counts_in_cluster) - if all(count == 1 for count in ALO_proteome_counts_in_cluster): - return "true" - fuzzycount_count = len( - [ - ALO_proteome_counts - for ALO_proteome_counts in ALO_proteome_counts_in_cluster - if ALO_proteome_counts == fuzzy_count - ] - ) - - fuzzyrange_count = len( - [ - ALO_proteome_counts - for ALO_proteome_counts in ALO_proteome_counts_in_cluster - if ALO_proteome_counts in fuzzy_range - ] - ) - - if fuzzycount_count + fuzzyrange_count == length: - fuzzy_fr = fuzzycount_count / length - - if fuzzy_fr >= fuzzy_fraction: - return "fuzzy" - - return None diff --git a/src/core/proteins.py b/src/core/proteins.py deleted file mode 100644 index c47439d..0000000 --- a/src/core/proteins.py +++ /dev/null @@ -1,108 +0,0 @@ -from collections import Counter -from typing import Dict, List, Optional, Union - -from core.utils import mean, median, sd - - -class Protein: - def __init__( - self, - protein_id: str, - proteome_id: str, - species_id: str, - sequence_id: str, - ) -> None: - - self.protein_id = protein_id - self.proteome_id = proteome_id - self.species_id = species_id - self.sequence_id = sequence_id - self.length: Optional[int] = None - self.clustered: bool = False - self.secreted: bool = False - self.domain_counter_by_domain_source: Dict[str, Counter[str]] = {} - self.go_terms: List[str] = [] - - def update_length(self, length: int) -> None: - self.length = length - - -class ProteinCollection: - def __init__(self, proteins_list: List[Protein]) -> None: - self.proteins_list: List[Protein] = proteins_list - self.proteins_by_protein_id: Dict[str, Protein] = { - protein.protein_id: protein for protein in proteins_list - } - self.protein_count: int = len(proteins_list) - self.domain_sources: List[str] = [] - self.fastas_parsed: bool = False - self.functional_annotation_parsed: bool = False - self.domain_desc_by_id_by_source: Dict[str, Dict[str, str]] = {} - - def add_annotation_to_protein( - self, - domain_protein_id: str, - domain_counter_by_domain_source: Dict[str, Counter], - go_terms: List[str], - ): - """ - Updates a protein object with domain counters and GO terms. - - Args: - - domain_protein_id (str): Identifier of the protein to annotate. - - domain_counter_by_domain_source (Dict[str, Counter]): Domain sources mapped to counters of domains. - - go_terms (List[str]): Gene Ontology (GO) terms associated with the protein. - - This method sets domain counters, assigns GO terms, and checks if the protein is secreted - based on domain information ('SignalP_EUK' source). - - Note: If 'SignalP_EUK' indicates 'SignalP-noTM', sets protein.secreted = True. - """ - protein: Optional[Protein] = self.proteins_by_protein_id.get( - domain_protein_id, None - ) - if protein is not None: - protein.domain_counter_by_domain_source = domain_counter_by_domain_source - signalp_notm = protein.domain_counter_by_domain_source.get( - "SignalP_EUK", None - ) - if signalp_notm and "SignalP-noTM" in signalp_notm: - protein.secreted = True - protein.go_terms = go_terms - - def get_protein_length_stats( - self, protein_ids: List[str] - ) -> Dict[str, Union[int, float]]: - """ - Calculate statistics (sum, mean, median, standard deviation) of protein lengths. - - Args: - protein_ids (List[str]): List of protein IDs for which to calculate statistics. - - Returns: - Dict[str, Union[int, float]): A dictionary containing the calculated statistics: - - 'sum': Sum of lengths of proteins in the input list. - - 'mean': Mean length of proteins in the input list. - - 'median': Median length of proteins in the input list. - - 'sd': Standard deviation of lengths of proteins in the input list. - - If no valid protein lengths could be calculated (e.g., if protein_ids is empty or no lengths - are available for the provided protein IDs), the values in the dictionary will default to 0 or 0.0. - """ - protein_length_stats = {"sum": 0, "mean": 0.0, "median": 0, "sd": 0.0} - if protein_ids and self.fastas_parsed: - protein_lengths: List[int] = [ - length - for length in [ - self.proteins_by_protein_id[protein_id].length - for protein_id in protein_ids - if protein_id in self.proteins_by_protein_id - ] - if length is not None - ] - protein_length_stats["sum"] = sum(protein_lengths) - protein_length_stats["mean"] = mean(protein_lengths) - protein_length_stats["median"] = median(protein_lengths) - protein_length_stats["sd"] = sd(protein_lengths) - - return protein_length_stats diff --git a/src/core/results.py b/src/core/results.py deleted file mode 100644 index 9aa0f60..0000000 --- a/src/core/results.py +++ /dev/null @@ -1,46 +0,0 @@ -import logging -import time - -from core.datastore import DataFactory -from core.input import InputData - -logger = logging.getLogger("kinfin_logger") - - -def analyse(input_data: InputData) -> None: - """ - Performs KinFin analysis based on the provided input data using DataFactory. - - Args: - input_data (InputData): An instance of InputData containing input parameters and data. - - Returns: - None - - Raises: - Any exceptions raised by DataFactory methods. - """ - overall_start = time.time() - dataFactory = DataFactory(input_data) - dataFactory.setup_dirs() - dataFactory.analyse_clusters() - dataFactory.aloCollection.write_tree( - dataFactory.dirs, - dataFactory.inputData.plot_tree, - dataFactory.inputData.plot_format, - dataFactory.inputData.fontsize, - ) - rarefaction_data = dataFactory.aloCollection.compute_rarefaction_data( - repetitions=dataFactory.inputData.repetitions - ) - dataFactory.plot_rarefaction_data( - dirs=dataFactory.dirs, - plotsize=dataFactory.inputData.plotsize, - plot_format=dataFactory.inputData.plot_format, - fontsize=dataFactory.inputData.fontsize, - rarefaction_by_samplesize_by_level_by_attribute=rarefaction_data, - ) - dataFactory.write_output() - overall_end = time.time() - overall_elapsed = overall_end - overall_start - logger.info(f"[STATUS] - Took {overall_elapsed}s to run kinfin.") diff --git a/src/core/utils.py b/src/core/utils.py deleted file mode 100644 index 7cdc693..0000000 --- a/src/core/utils.py +++ /dev/null @@ -1,285 +0,0 @@ -import gzip -import json -import logging -import os -import sys -from math import log, sqrt -from typing import Any, Generator, List, Optional, Tuple, Union - -import scipy - -logger = logging.getLogger("kinfin_logger") - - -def progress(iteration: int, steps: Union[int, float], max_value: int) -> None: - """ - Print progress in percentage based on the current iteration, steps, and maximum value. - - Parameters: - - iteration (int): Current iteration or step number. - - steps (int | float): Number of steps or intervals after which progress is updated. - - max_value (int): Maximum value or total number of iterations. - - Returns: - - None - - Example: - >>> progress(5, 2, 10) - [PROGRESS] - 50% - """ - if iteration == max_value: - sys.stdout.write("\r") - print("[PROGRESS]\t- %d%%" % (100)) - elif iteration % int(steps + 1) == 0: - sys.stdout.write("\r") - print("[PROGRESS]\t- %d%%" % (float(iteration / max_value) * 100), end=" ") - sys.stdout.flush() - - -def check_file(filepath: Optional[str], install_kinfin: bool = False) -> None: - """ - Check if a file exists. - - Args: - filepath (str): Path to the file to be checked. - - Raises: - FileNotFoundError: If the file does not exist. - """ - - if filepath is not None and not os.path.isfile(filepath): - error_msg = f"[ERROR] - file {filepath} not found." - if install_kinfin: - error_msg += " Please run the install script to download kinfin." - raise FileNotFoundError(error_msg) - - -def yield_file_lines(filepath: str) -> Generator[str, Any, None]: - """ - Args: - filepath (str): Path to the file. - - Yields: - str: Each line from the file. - """ - check_file(filepath) - if filepath.endswith(".gz"): - with gzip.open(filepath, "rb") as fh: - for line in fh: - line = line.decode("utf-8") - if line.startswith("nodesDB.txt"): - line = f'#{line.split("#")[1]}' - yield line.rstrip("\n") - else: - with open(filepath) as fh: - for line in fh: - yield line.rstrip("\n") - - -def yield_config_lines( - config_f: str, - taxon_idx_mapping_file: Optional[str], -): - if config_f.endswith(".json"): - if not taxon_idx_mapping_file: - raise ValueError("[ERROR] - taxon_idx_mapping not present") - - with ( - open(taxon_idx_mapping_file, "r") as f_mapping, - open(config_f, "r") as f_config, - ): - taxon_idx_mapping = json.load(f_mapping) - config_data = json.load(f_config) - headers = ["IDX"] + list(config_data[0].keys()) - yield "#" + ",".join(headers) - - for item in config_data: - idx = taxon_idx_mapping[item.get("taxon") or item.get("TAXON")] - row = [idx] + [item[key] for key in headers[1:]] - yield ",".join(row) - else: - yield from yield_file_lines(config_f) - - return - - -def read_fasta_len(fasta_file: str) -> Generator[Tuple[str, int], Any, None]: - """ - Generator function to parse a FASTA file and yield tuples of header and sequence length. - - Args: - - fasta_file (str): Path to the FASTA file to be parsed. - - Yields: - Tuple[str, int]: A tuple containing the header and the length of the sequence. - - Raises: - FileNotFoundError: If the specified FASTA file does not exist. - """ - check_file(fasta_file) - with open(fasta_file) as fh: - logger.info(f"[STATUS]\t - Parsing FASTA {fasta_file}") - header: str = "" - seqs: List[str] = [] - for line in fh: - if line[0] == ">": - if header: - header = ( - header.replace(":", "_") - .replace(",", "_") - .replace("(", "_") - .replace(")", "_") - ) # orthofinder replaces chars - yield header, len("".join(seqs)) - header, seqs = ( - line[1:-1].split()[0], - [], - ) # Header is split at first whitespace - else: - seqs.append(line[:-1]) - header = ( - header.replace(":", "_") - .replace(",", "_") - .replace("(", "_") - .replace(")", "_") - ) # orthofinder replaces chars - yield header, len("".join(seqs)) - - -def median(lst) -> float: - """ - Calculate the median of a list of numbers. - - Args: - - lst (list): List of numerical values. - - Returns: - - float: Median of the list. - """ - list_sorted = sorted(lst) - list_length = len(lst) - index = (list_length - 1) // 2 - if list_length % 2: - return list_sorted[index] / 1.0 - else: - return (list_sorted[index] + list_sorted[index + 1]) / 2.0 - - -def mean(lst) -> float: - """ - Calculate the mean (average) of a list of numbers. - - Args: - - lst (list): List of numerical values. - - Returns: - - float: Mean of the list. - """ - return float(sum(lst)) / len(lst) if lst else 0.0 - - -def sd(lst, population=True) -> float: - """ - Calculate the standard deviation of a list of numbers. - - Args: - - lst (list): List of numerical values. - - population (bool, optional): If True, calculates population standard deviation, - otherwise calculates sample standard deviation. Default is True. - - Returns: - - float: Standard deviation of the list. - """ - n = len(lst) - differences = [x_ - mean(lst) for x_ in lst] - sq_differences = [d**2 for d in differences] - ssd = sum(sq_differences) - variance = ssd / n if population is True else ssd / (n - 1) - return sqrt(variance) - - -def statistic( - count_1: List[int], - count_2: List[int], - test: str, - min_proteomes: int, -) -> Tuple[ - Optional[float], - Optional[float], - Optional[float], - Optional[float], -]: - """ - Perform statistical tests and calculate relevant statistics between two lists of counts. - - Args: - - count_1 (list): List of counts (integers). - - count_2 (list): Another list of counts (integers). - - test (str): Type of statistical test to perform, one of "welch", "mannwhitneyu", "ttest", "ks", "kruskal". - - min_proteomes (int): Minimum number of proteomes required for valid analysis. - - Returns: - - Tuple[Optional[float], Optional[float], Optional[float], Optional[float]]: - Tuple containing: - - pvalue: p-value of the statistical test (or None if test is not applicable). - - log2_mean: Logarithm base 2 of the mean of count_1 divided by count_2. - - mean_count_1: Mean of count_1. - - mean_count_2: Mean of count_2. - """ - pvalue: Optional[float] = None - log2_mean: Optional[float] = None - mean_count_1: Optional[float] = None - mean_count_2: Optional[float] = None - - implicit_count_1: List[float] = [count for count in count_1 if count > 0] - implicit_count_2: List[float] = [count for count in count_2 if count > 0] - - if len(implicit_count_1) < min_proteomes or len(implicit_count_2) < min_proteomes: - return None, None, None, None - - mean_count_1 = mean(implicit_count_1) - mean_count_2 = mean(implicit_count_2) - log2_mean = log(mean_count_1 / mean_count_2, 2) - - if ( - len(set(implicit_count_1)) == 1 - and len(set(implicit_count_2)) == 1 - and set(implicit_count_1) == set(implicit_count_2) - ): # equal - pvalue = 1.0 - elif test == "welch": - # try: - # Welch's t-test - pvalue = scipy.stats.ttest_ind( - implicit_count_1, - implicit_count_2, - equal_var=False, - )[1] - - if pvalue != pvalue: # testing for "nan" - pvalue = 1.0 - elif test == "mannwhitneyu": - try: - pvalue = scipy.stats.mannwhitneyu( - implicit_count_1, - implicit_count_2, - alternative="two-sided", - )[1] - except ValueError: # throws ValueError when all numbers are equal - pvalue = 1.0 - elif test == "ttest": - # try: - pvalue = scipy.stats.ttest_ind(implicit_count_1, implicit_count_2)[1] # t-test - if pvalue != pvalue: # testing for "nan" - pvalue = 1.0 - elif test == "ks": - # H0 that they are drawn from the same distribution - pvalue = scipy.stats.ks_2samp(implicit_count_1, implicit_count_2)[1] - if pvalue != pvalue: # testing for "nan" - pvalue = 1.0 - elif test == "kruskal": - # H0 is that population median is equal - pvalue = scipy.stats.kruskal(implicit_count_1, implicit_count_2)[1] - if pvalue != pvalue: # testing for "nan" - pvalue = 1.0 - return pvalue, log2_mean, mean_count_1, mean_count_2 diff --git a/src/internal/__init__.py b/src/internal/__init__.py new file mode 100644 index 0000000..7583b54 --- /dev/null +++ b/src/internal/__init__.py @@ -0,0 +1,122 @@ +import os + +import polars as pl + +from internal import ( + attribute_metrics, + cluster_metrics, + cluster_summary, + counts, + logger, + parsers, + plots, + utils, +) + + +def classify_clusters( + cluster_df: pl.DataFrame, config_df: pl.DataFrame, attributes: list[str] +) -> pl.DataFrame: + exploded_df = cluster_df.select(["cluster_id", "taxons"]).explode("taxons").unique() + + config_subset = config_df.select(attributes) + merged_df = exploded_df.join(config_subset, left_on="taxons", right_on="TAXON") + + aggs = [ + (pl.col("taxons") if attr == "TAXON" else pl.col(attr)) + .n_unique() + .alias(f"{attr}_n_unique") + for attr in attributes + ] + + agg_df = merged_df.group_by("cluster_id").agg(aggs) + + result_df = cluster_df.join(agg_df, on="cluster_id", how="left") + + classification_exprs = [] + for attr in attributes: + class_col = f"{attr}_cluster_type" + + expr = ( + pl.when(pl.col("cluster_protein_count") == 1) + .then(pl.lit("singleton")) + .when(pl.col(f"{attr}_n_unique") == 1) + .then(pl.lit("specific")) + .otherwise(pl.lit("shared")) + .alias(class_col) + ) + + classification_exprs.append(expr) + + return result_df.with_columns(classification_exprs).drop( + [f"{attr}_n_unique" for attr in attributes] + ) + + +def analyse(args, nodesdb_f, ndb_f): + log_path = os.path.join(args.output_path, "kinfin.log") + logger.setup_logger(log_path) + cluster_file = args.cluster_file + config_file = args.config_file + output_dir = args.output_path + + nodesdb = parsers.nodesdb(filepath=nodesdb_f, outpath=ndb_f) + config_df, attributes = parsers.configfile(config_file, nodesdb) + + utils.setup_dirs(output_dir, attributes) + + cluster_df = parsers.clusterfile(cluster_file, config_df, output_dir) + cluster_df = classify_clusters(cluster_df, config_df, attributes) + + label_columns = [ + col for col in config_df.columns if col not in ("#IDX", "TAXON", "TAXID", "OUT") + ] + + taxon_label_dict = { + row["TAXON"]: { + **{label: row[label] for label in label_columns}, + "TAXON": row["TAXON"], + "all": "all", + } + for row in config_df.to_dicts() + } + + unique_label_values = { + col: config_df.select(pl.col(col)).unique().to_series().to_list() + for col in config_df.columns + if col not in ("#IDX", "TAXON", "TAXID", "OUT") + } + unique_label_values["all"] = ["all"] + unique_label_values["TAXON"] = config_df["TAXON"].to_list() + + counts.get_cluster_counts_by_taxon( + cluster_df=cluster_df, + config_df=config_df, + base_output_dir=output_dir, + ) + cluster_summary.get_all_cluster_summaries( + attributes=attributes, + cluster_df=cluster_df, + config_df=config_df, + base_output_dir=output_dir, + ) + cluster_metrics.get_all_cluster_metrics( + cluster_df=cluster_df, + taxon_label_dict=taxon_label_dict, + attributes=attributes, + unique_label_values=unique_label_values, + base_output_dir=output_dir, + ) + attribute_metrics.get_all_attribute_metrics( + config_df=config_df, + cluster_df=cluster_df, + attributes=attributes, + base_output_dir=output_dir, + ) + plots.plot_cluster_sizes(cluster_df=cluster_df, base_output_dir=output_dir) + plots.generate_kinfin_rarefaction_plots( + cluster_df=cluster_df, + config_df=config_df, + attributes=attributes, + base_output_dir=output_dir, + ) diff --git a/src/internal/attribute_metrics.py b/src/internal/attribute_metrics.py new file mode 100644 index 0000000..9d1c247 --- /dev/null +++ b/src/internal/attribute_metrics.py @@ -0,0 +1,433 @@ +import logging +import os +from typing import List, Optional + +import polars as pl + +logger = logging.getLogger("kinfin_logger") + + +def precompute_cluster_info( + cluster_df: pl.DataFrame, config_df: pl.DataFrame, attribute: str +) -> pl.DataFrame: + taxon_to_label_df = config_df.select( + pl.col("TAXON").alias("taxon"), + pl.col(attribute).alias("taxon_set"), + ) + + return ( + cluster_df.select(["cluster_id", "taxons", "protein_cluster"]) + .lazy() + .explode("taxons") + .rename({"taxons": "taxon"}) + .join(taxon_to_label_df.lazy(), on="taxon", how="left") + .group_by("cluster_id") + .agg( + pl.col("taxon_set").drop_nulls().unique().alias("taxon_sets"), + pl.col("protein_cluster").first().alias("protein_cluster"), + pl.col("protein_cluster").first().list.len().alias("protein_cluster_len"), + ) + .with_columns( + pl.when(pl.col("protein_cluster_len") == 1) + .then(pl.lit("singleton")) + .when(pl.col("taxon_sets").list.len() == 1) + .then(pl.lit("specific")) + .otherwise(pl.lit("shared")) + .alias("cluster_type") + ) + .collect() + ) + + +def add_cluster_and_protein_counts( + attribute_df: pl.DataFrame, + cluster_info_df: pl.DataFrame, + config_df: pl.DataFrame, + attribute: str, +) -> pl.DataFrame: + cluster_counts = ( + cluster_info_df.lazy() + .explode("taxon_sets") + .group_by("taxon_sets", "cluster_type") + .agg(pl.count().alias("count")) + .collect() + .pivot( + index="taxon_sets", + columns="cluster_type", + values="count", + aggregate_function="sum", + ) + .rename({"taxon_sets": "taxon_set"}) + ) + + for col_name in ["singleton", "specific", "shared"]: + if col_name not in cluster_counts.columns: + cluster_counts = cluster_counts.with_columns( + pl.lit(0, dtype=pl.UInt32).alias(col_name) + ) + + cluster_counts = cluster_counts.rename( + { + "singleton": "singleton_cluster_count", + "specific": "specific_cluster_count", + "shared": "shared_cluster_count", + } + ).with_columns( + cluster_total_count=pl.col("singleton_cluster_count") + + pl.col("specific_cluster_count") + + pl.col("shared_cluster_count") + ) + + taxon_to_label_df = config_df.select( + pl.col("TAXON").alias("taxon"), + pl.col(attribute).alias("taxon_set"), + ).filter(pl.col("taxon_set").is_not_null()) + + protein_level_info = ( + cluster_info_df.lazy() + .select(["cluster_id", "cluster_type", "protein_cluster"]) + .explode("protein_cluster") + .with_columns( + pl.col("protein_cluster").str.split(".").list.get(0).alias("taxon") + ) + .join(taxon_to_label_df.lazy(), on="taxon", how="left") + ) + + protein_counts = ( + protein_level_info.group_by("taxon_set", "cluster_type") + .agg(pl.count().alias("protein_count")) + .collect() + .pivot( + index="taxon_set", + columns="cluster_type", + values="protein_count", + aggregate_function="sum", + ) + ) + for col_name in ["singleton", "specific", "shared"]: + if col_name not in protein_counts.columns: + protein_counts = protein_counts.with_columns( + pl.lit(0, dtype=pl.UInt32).alias(col_name) + ) + + protein_counts = protein_counts.rename( + { + "singleton": "singleton_protein_count", + "specific": "specific_protein_count", + "shared": "shared_protein_count", + } + ).with_columns( + protein_total_count=pl.col("singleton_protein_count") + + pl.col("specific_protein_count") + + pl.col("shared_protein_count") + ) + + attribute_df = attribute_df.join(cluster_counts, on="taxon_set", how="left") + attribute_df = attribute_df.join(protein_counts, on="taxon_set", how="left") + + return attribute_df.fill_null(0) + + +def add_protein_spans( + attribute_df: pl.DataFrame, + cluster_info_df: pl.DataFrame, + protein_lengths_df: Optional[pl.DataFrame] = None, +) -> pl.DataFrame: + span_columns = [ + "protein_total_span", + "singleton_protein_span", + "specific_protein_span", + "shared_protein_span", + ] + if protein_lengths_df is None or protein_lengths_df.is_empty(): + return attribute_df.with_columns([pl.lit(0).alias(col) for col in span_columns]) + + protein_spans = ( + cluster_info_df.lazy() + .select(["cluster_id", "taxon_sets", "protein_cluster", "cluster_type"]) + .explode("protein_cluster") + .join( + protein_lengths_df.lazy(), + left_on="protein_cluster", + right_on="protein_id", + how="left", + ) + .with_columns(pl.col("length").fill_null(0)) + .explode("taxon_sets") + .group_by("taxon_sets", "cluster_type") + .agg(pl.col("length").sum().alias("span")) + .collect() + .pivot( + index="taxon_sets", + columns="cluster_type", + values="span", + aggregate_function="sum", + ) + .rename({"taxon_sets": "taxon_set"}) + ) + + for col_name in ["singleton", "specific", "shared"]: + if col_name not in protein_spans.columns: + protein_spans = protein_spans.with_columns(pl.lit(0).alias(col_name)) + + protein_spans = protein_spans.rename( + { + "singleton": "singleton_protein_span", + "specific": "specific_protein_span", + "shared": "shared_protein_span", + } + ).with_columns( + protein_total_span=pl.col("singleton_protein_span") + + pl.col("specific_protein_span") + + pl.col("shared_protein_span") + ) + + return attribute_df.join(protein_spans, on="taxon_set", how="left").fill_null(0) + + +def add_special_cluster_counts( + attribute_df: pl.DataFrame, + cluster_df: pl.DataFrame, + config_df: pl.DataFrame, + cluster_info_df: pl.DataFrame, + attribute: str, + fuzzy_count: int = 1, + fuzzy_min: int = 0, + fuzzy_max: int = 20, + fuzzy_fraction: float = 0.75, +) -> pl.DataFrame: + taxon_to_label_df = config_df.select( + pl.col("TAXON").alias("taxon"), + pl.col(attribute).alias("taxon_set"), + ).filter(pl.col("taxon_set").is_not_null()) + + present_taxa_long = ( + cluster_df.lazy() + .explode("protein_cluster") + .with_columns( + pl.col("protein_cluster").str.split(".").list.get(0).alias("taxon") + ) + .join(taxon_to_label_df.lazy(), on="taxon", how="left") + .select(["cluster_id", "taxon_set", "taxon"]) + .filter(pl.col("taxon_set").is_not_null()) + ) + + expected_taxa_long = taxon_to_label_df.select(["taxon_set", "taxon"]).rename( + {"taxon": "expected_taxon"} + ) + + counts_df = ( + present_taxa_long.join(expected_taxa_long.lazy(), on="taxon_set") + .with_columns(is_match=(pl.col("taxon") == pl.col("expected_taxon"))) + .group_by("cluster_id", "taxon_set", "expected_taxon") + .agg(pl.col("is_match").sum().alias("count")) + .group_by("cluster_id", "taxon_set") + .agg(pl.col("count").alias("taxon_counts")) + ) + + counts_with_type = counts_df.join( + cluster_info_df.lazy().select(["cluster_id", "cluster_type"]), + on="cluster_id", + how="left", + ) + + special_counts = ( + counts_with_type.with_columns( + num_expected=pl.col("taxon_counts").list.len(), + is_true_1to1=(pl.col("taxon_counts").list.len() > 2) + & (pl.col("taxon_counts").list.eval(pl.element() == 1).list.all()), + fraction_at_fuzzy=( + pl.col("taxon_counts").list.count_matches(fuzzy_count) + / pl.col("taxon_counts").list.len() + ), + all_in_range=pl.col("taxon_counts") + .list.eval((pl.element() >= fuzzy_min) & (pl.element() <= fuzzy_max)) + .list.all(), + ) + .with_columns( + is_fuzzy=(pl.col("num_expected") > 2) + & (pl.col("fraction_at_fuzzy") >= fuzzy_fraction) + & pl.col("all_in_range") + & ~(pl.col("is_true_1to1")) + ) + .group_by("taxon_set") + .agg( + pl.col("is_true_1to1") + .filter(pl.col("cluster_type") == "specific") + .sum() + .alias("specific_cluster_true_1to1_count"), + pl.col("is_true_1to1") + .filter(pl.col("cluster_type") == "shared") + .sum() + .alias("shared_cluster_true_1to1_count"), + pl.col("is_fuzzy") + .filter(pl.col("cluster_type") == "specific") + .sum() + .alias("specific_cluster_fuzzy_count"), + pl.col("is_fuzzy") + .filter(pl.col("cluster_type") == "shared") + .sum() + .alias("shared_cluster_fuzzy_count"), + ) + .collect() + ) + + return attribute_df.join(special_counts, on="taxon_set", how="left").fill_null(0) + + +def add_absent_cluster_counts( + attribute_df: pl.DataFrame, cluster_info_df: pl.DataFrame +) -> pl.DataFrame: + total_counts = cluster_info_df.group_by("cluster_type").agg( + pl.count().alias("total_count") + ) + + total_singleton = total_counts.filter(pl.col("cluster_type") == "singleton")[ + "total_count" + ].sum() + total_specific = total_counts.filter(pl.col("cluster_type") == "specific")[ + "total_count" + ].sum() + total_shared = total_counts.filter(pl.col("cluster_type") == "shared")[ + "total_count" + ].sum() + + present_counts_df = ( + cluster_info_df.lazy() + .explode("taxon_sets") + .group_by("taxon_sets", "cluster_type") + .agg(pl.count().alias("present_count")) + .collect() + .pivot( + index="taxon_sets", + columns="cluster_type", + values="present_count", + aggregate_function="sum", + ) + .rename({"taxon_sets": "taxon_set"}) + ) + + for col_type in ["singleton", "specific", "shared"]: + if col_type not in present_counts_df.columns: + present_counts_df = present_counts_df.with_columns( + pl.lit(0, dtype=pl.UInt32).alias(col_type) + ) + + return ( + attribute_df.join(present_counts_df, on="taxon_set", how="left") + .fill_null(0) + .with_columns( + absent_cluster_singleton_count=(total_singleton - pl.col("singleton")), + absent_cluster_specific_count=(total_specific - pl.col("specific")), + absent_cluster_shared_count=(total_shared - pl.col("shared")), + ) + .with_columns( + absent_cluster_total_count=pl.col("absent_cluster_singleton_count") + + pl.col("absent_cluster_specific_count") + + pl.col("absent_cluster_shared_count") + ) + .drop(["singleton", "specific", "shared"]) + ) + + +def get_attribute_metrics( + config_df: pl.DataFrame, + attribute: str, + cluster_df, + protein_lengths_df: Optional[pl.DataFrame] = None, +) -> pl.DataFrame: + attribute_df = ( + config_df.group_by(pl.col(attribute).alias("taxon_set")) + .agg( + pl.col("TAXON").unique().sort().alias("TAXON_taxa_list"), + pl.col("TAXON").n_unique().alias("TAXON_count"), + ) + .with_columns( + pl.lit(attribute).alias("#attribute"), + pl.col("TAXON_taxa_list").list.join(", ").alias("TAXON_taxa"), + ) + .select(["#attribute", "taxon_set", "TAXON_count", "TAXON_taxa"]) + .sort("taxon_set") + ) + + cluster_info_df = precompute_cluster_info(cluster_df, config_df, attribute) + + final_df = ( + attribute_df.pipe( + add_cluster_and_protein_counts, + cluster_info_df=cluster_info_df, + config_df=config_df, + attribute=attribute, + ) + .pipe( + add_protein_spans, + cluster_info_df=cluster_info_df, + protein_lengths_df=protein_lengths_df, + ) + .pipe( + add_special_cluster_counts, + cluster_df=cluster_df, + config_df=config_df, + attribute=attribute, + cluster_info_df=cluster_info_df, + ) + .pipe(add_absent_cluster_counts, cluster_info_df=cluster_info_df) + ) + + final_cols = [ + "#attribute", + "taxon_set", + "cluster_total_count", + "protein_total_count", + "protein_total_span", + "singleton_cluster_count", + "singleton_protein_count", + "singleton_protein_span", + "specific_cluster_count", + "specific_protein_count", + "specific_protein_span", + "shared_cluster_count", + "shared_protein_count", + "shared_protein_span", + "specific_cluster_true_1to1_count", + "specific_cluster_fuzzy_count", + "shared_cluster_true_1to1_count", + "shared_cluster_fuzzy_count", + "absent_cluster_total_count", + "absent_cluster_singleton_count", + "absent_cluster_specific_count", + "absent_cluster_shared_count", + "TAXON_count", + "TAXON_taxa", + ] + + for col in final_cols: + if col not in final_df.columns: + final_df = final_df.with_columns(pl.lit(0).alias(col)) + + return final_df.select(final_cols) + + +def get_all_attribute_metrics( + config_df: pl.DataFrame, + cluster_df: pl.DataFrame, + attributes: List[str], + base_output_dir: str, + protein_lengths_df: Optional[pl.DataFrame] = None, +) -> None: + config_df = config_df.with_columns(pl.lit("all").alias("all")) + + for attribute in attributes: + attribute_df = get_attribute_metrics( + cluster_df=cluster_df, + config_df=config_df, + protein_lengths_df=protein_lengths_df, + attribute=attribute, + ) + + out_dir = os.path.join(base_output_dir, attribute) + os.makedirs(out_dir, exist_ok=True) + + out_path = os.path.join(out_dir, f"{attribute}.attribute_metrics.txt") + attribute_df.write_csv(out_path, separator="\t") + logger.info(f"[✓] Writing: {out_path}") diff --git a/src/internal/cluster_metrics.py b/src/internal/cluster_metrics.py new file mode 100644 index 0000000..9cf49b4 --- /dev/null +++ b/src/internal/cluster_metrics.py @@ -0,0 +1,729 @@ +import logging +import warnings +from typing import Dict, List, Set + +import matplotlib as mat +import matplotlib.pyplot as plt +import numpy as np +import polars as pl +from matplotlib.ticker import NullFormatter +from scipy import stats + +logger = logging.getLogger("kinfin_logger") + + +warnings.filterwarnings("ignore", message="Precision loss occurred*") + +mat.use("agg") + +plt.style.use("ggplot") + +mat.rc("ytick", labelsize=20) +mat.rc("xtick", labelsize=20) +mat.rcParams.update({"font.size": 22}) + + +def generate_background_plots( + results_df: pl.DataFrame, attribute: str, output_dir: str, plot_format: str = "png" +): + if results_df.height == 0: + return + + for taxon_key, data in results_df.group_by("TAXON_1"): + taxon_1 = taxon_key[0] + log2fc_values = ( + data.get_column("log2_mean(TAXON_1/background)") + .drop_nulls() + .to_numpy() + .copy() + ) + p_values = ( + data.get_column("mwu_pvalue(TAXON_1 vs. background)") + .drop_nulls() + .to_numpy() + .copy() + ) + + if len(p_values) == 0: + continue + + p_values[p_values == 0] = 0.01 / (len(p_values) + 1) + + plt.figure(1, figsize=(24, 12)) + left, width = 0.1, 0.65 + bottom, height = 0.1, 0.65 + rect_scatter = [left, bottom, width, height] + axScatter = plt.axes(rect_scatter) + axScatter.set_facecolor("white") + + ooFive = 0.05 + ooOne = 0.01 + log2fc_percentile = np.percentile(np.abs(log2fc_values), 95) + + axScatter.axhline(y=ooFive, linewidth=2, color="orange", linestyle="--") + ooFive_artist = plt.Line2D((0, 1), (0, 0), color="orange", linestyle="--") + axScatter.axhline(y=ooOne, linewidth=2, color="red", linestyle="--") + ooOne_artist = plt.Line2D((0, 1), (0, 0), color="red", linestyle="--") + + axScatter.axvline(x=1.0, linewidth=2, color="purple", linestyle="--") + axScatter.axvline(x=-1.0, linewidth=2, color="purple", linestyle="--") + v1_artist = plt.Line2D((0, 1), (0, 0), color="purple", linestyle="--") + axScatter.axvline( + x=log2fc_percentile, linewidth=2, color="blue", linestyle="--" + ) + axScatter.axvline( + x=-log2fc_percentile, linewidth=2, color="blue", linestyle="--" + ) + nine_five_percentile_artist = plt.Line2D( + (0, 1), (0, 0), color="blue", linestyle="--" + ) + + axScatter.scatter( + log2fc_values, p_values, alpha=0.8, edgecolors="none", s=25, c="grey" + ) + + legend = axScatter.legend( + [ooFive_artist, ooOne_artist, v1_artist, nine_five_percentile_artist], + [ + f"p-value = {ooFive}", + f"p-value = {ooOne}", + "|log2FC| = 1", + f"|log2FC-95%ile| = {log2fc_percentile:.2f}", + ], + fontsize=18, + frameon=True, + ) + legend.get_frame().set_facecolor("white") + + x_max_abs = np.max(np.abs(log2fc_values)) if len(log2fc_values) > 0 else 1.0 + axScatter.set_xlim(-x_max_abs - 1, x_max_abs + 1) + axScatter.set_ylim(1.1, np.min(p_values) * 0.1) + axScatter.set_yscale("log") + axScatter.set_xlabel(f"log2(mean({taxon_1})/mean(background))", fontsize=18) + axScatter.set_ylabel("p-value", fontsize=18) + + axScatter.grid(True, linewidth=1, which="major", color="lightgrey") + axScatter.grid(True, linewidth=0.5, which="minor", color="lightgrey") + + plot_file = f"{output_dir}/{attribute}/{attribute}.pairwise_representation_test.{taxon_1}_background.{plot_format}" + print(f"[✓] Plotting: {plot_file}") + plt.savefig(plot_file, format=plot_format) + plt.close() + + +def generate_background_representation_test( + cluster_df: pl.DataFrame, + attribute: str, + label_to_taxons: Dict[str, Dict[str, Set[str]]], + output_dir: str, + min_proteomes: int = 2, +): + """ + Performs a statistical test for each label against all other labels (background). + """ + all_results = [] + taxon_map_df = pl.DataFrame( + [ + (taxon, label) + for label, taxons in label_to_taxons[attribute].items() + for taxon in taxons + ], + schema=["taxons", "label"], + orient="row", + ) + + protein_counts_per_taxon = ( + cluster_df.explode("taxons") + .group_by("cluster_id", "taxons") + .agg(pl.count().alias("count")) + .join(taxon_map_df, on="taxons") + ) + + for current_label in label_to_taxons[attribute]: + label_counts = ( + protein_counts_per_taxon.group_by("cluster_id") + .agg( + pl.col("count") + .filter(pl.col("label") == current_label) + .alias("inside_counts"), + pl.col("count") + .filter(pl.col("label") != current_label) + .alias("outside_counts"), + ) + .filter( + (pl.col("inside_counts").list.len() >= min_proteomes) + & (pl.col("outside_counts").list.len() >= min_proteomes) + ) + ) + + if label_counts.height == 0: + continue + + stats_df = label_counts.with_columns( + TAXON_1_mean=pl.col("inside_counts").list.mean(), + background_mean=pl.col("outside_counts").list.mean(), + ).with_columns( + log2_mean=pl.when( + (pl.col("TAXON_1_mean") > 0) & (pl.col("background_mean") > 0) + ) + .then((pl.col("TAXON_1_mean") / pl.col("background_mean")).log(base=2)) + .when(pl.col("TAXON_1_mean") == pl.col("background_mean")) + .then(0.0) + .otherwise(None) + ) + + p_values = [] + for row in stats_df.iter_rows(named=True): + c1, c2 = row["inside_counts"], row["outside_counts"] + if len(set(c1)) == 1 and len(set(c2)) == 1 and c1[0] == c2[0]: + p_values.append(1.0) + continue + try: + pvalue = stats.mannwhitneyu(c1, c2, alternative="two-sided").pvalue + p_values.append(pvalue) + except ValueError: + p_values.append(1.0) + + final_df = stats_df.with_columns(pl.Series("p_value", p_values)).select( + pl.col("cluster_id").alias("#cluster_id"), + pl.lit(current_label).alias("TAXON_1"), + pl.col("TAXON_1_mean"), + pl.lit("background").alias("TAXON_2"), + pl.col("background_mean").alias("TAXON_2_mean"), + pl.col("log2_mean").alias("log2_mean(TAXON_1/background)"), + pl.col("p_value").alias("mwu_pvalue(TAXON_1 vs. background)"), + ) + all_results.append(final_df) + + if not all_results: + return + + full_results_df = pl.concat(all_results) + generate_background_plots(full_results_df, attribute, output_dir) + + +def generate_pairwise_plots( + pairwise_df: pl.DataFrame, + attribute: str, + output_dir: str, + plot_format: str = "png", +): + + if pairwise_df.height == 0: + return + + pair_groups = pairwise_df.group_by("TAXON_1", "TAXON_2") + + for (taxon_1, taxon_2), data in pair_groups: + log2fc_values = ( + data.get_column("log2_mean(TAXON_1/TAXON_2)").drop_nulls().to_numpy().copy() + ) + p_values = ( + data.get_column("mwu_pvalue(TAXON_1 vs. TAXON_2)") + .drop_nulls() + .to_numpy() + .copy() + ) + + if len(p_values) == 0: + continue + + p_values[p_values == 0] = 0.01 / (len(p_values) + 1) + + plt.figure(1, figsize=(24, 12)) + + left, width = 0.1, 0.65 + bottom, height = 0.1, 0.65 + bottom_h = left + width + 0.02 + rect_scatter = [left, bottom, width, height] + rect_histx = [left, bottom_h, width, 0.2] + + axScatter = plt.axes(rect_scatter) + axScatter.set_facecolor("white") + axHistx = plt.axes(rect_histx) + axHistx.set_facecolor("white") + + nullfmt = NullFormatter() + axHistx.xaxis.set_major_formatter(nullfmt) + axHistx.yaxis.set_major_formatter(nullfmt) + + binwidth = 0.05 + xymax = np.max([np.max(np.fabs(log2fc_values)), np.max(np.fabs(p_values))]) + lim = (int(xymax / binwidth) + 1) * binwidth + bins = np.arange(-lim, lim + binwidth, binwidth) + axHistx.hist( + log2fc_values, bins=bins, histtype="stepfilled", color="grey", align="mid" + ) + + ooFive = 0.05 + ooOne = 0.01 + log2fc_percentile = np.percentile(np.abs(log2fc_values), 95) + + axScatter.axhline(y=ooFive, linewidth=2, color="orange", linestyle="--") + ooFive_artist = plt.Line2D((0, 1), (0, 0), color="orange", linestyle="--") + axScatter.axhline(y=ooOne, linewidth=2, color="red", linestyle="--") + ooOne_artist = plt.Line2D((0, 1), (0, 0), color="red", linestyle="--") + + axScatter.axvline(x=1.0, linewidth=2, color="purple", linestyle="--") + axScatter.axvline(x=-1.0, linewidth=2, color="purple", linestyle="--") + v1_artist = plt.Line2D((0, 1), (0, 0), color="purple", linestyle="--") + axScatter.axvline( + x=log2fc_percentile, linewidth=2, color="blue", linestyle="--" + ) + axScatter.axvline( + x=-log2fc_percentile, linewidth=2, color="blue", linestyle="--" + ) + nine_five_percentile_artist = plt.Line2D( + (0, 1), (0, 0), color="blue", linestyle="--" + ) + + axScatter.scatter( + log2fc_values, p_values, alpha=0.8, edgecolors="none", s=25, c="grey" + ) + + legend = axScatter.legend( + [ooFive_artist, ooOne_artist, v1_artist, nine_five_percentile_artist], + [ + f"p-value = {ooFive}", + f"p-value = {ooOne}", + "|log2FC| = 1", + f"|log2FC-95%ile| = {log2fc_percentile:.2f}", + ], + fontsize=18, + frameon=True, + ) + legend.get_frame().set_facecolor("white") + + x_max_abs = np.max(np.abs(log2fc_values)) + axScatter.set_xlim(-x_max_abs - 1, x_max_abs + 1) + axScatter.set_ylim(1.1, np.min(p_values) * 0.1) + axScatter.set_yscale("log") + axScatter.set_xlabel(f"log2(mean({taxon_1})/mean({taxon_2}))", fontsize=18) + axScatter.set_ylabel("p-value", fontsize=18) + + axScatter.grid(True, linewidth=1, which="major", color="lightgrey") + axScatter.grid(True, linewidth=0.5, which="minor", color="lightgrey") + + axHistx.set_xlim(axScatter.get_xlim()) + + plot_file = f"{output_dir}/{attribute}/{attribute}.pairwise_representation_test.{taxon_1}_{taxon_2}.{plot_format}" + print(f"[✓] Plotting: {plot_file}") + plt.savefig(plot_file, format=plot_format) + plt.close() + + +def generate_pairwise_representation_test( + cluster_df: pl.DataFrame, + attribute: str, + label_to_taxons: Dict[str, Dict[str, Set[str]]], + output_dir: str, + min_proteomes: int = 2, +): + taxon_to_label_mapping = [ + (taxon, label) + for label, taxons in label_to_taxons[attribute].items() + for taxon in taxons + ] + taxon_map_df = pl.DataFrame( + taxon_to_label_mapping, schema=["taxons", "label"], orient="row" + ) + + protein_counts_per_taxon = ( + cluster_df.explode("taxons") + .group_by("cluster_id", "taxons") + .agg(pl.count().alias("count")) + ) + + counts_by_label = ( + protein_counts_per_taxon.join(taxon_map_df, on="taxons") + .group_by("cluster_id", "label") + .agg(pl.col("count").alias("counts_list")) + ) + + counts_by_label_filtered = counts_by_label.filter( + pl.col("counts_list").list.len() >= min_proteomes + ) + + pairwise_df = counts_by_label_filtered.join( + counts_by_label_filtered, on="cluster_id", suffix="_2" + ).filter(pl.col("label") < pl.col("label_2")) + + if pairwise_df.height == 0: + print( + f"[!] No valid pairs for pairwise test in attribute '{attribute}'. Skipping." + ) + return + + stats_df = pairwise_df.with_columns( + TAXON_1_mean=pl.col("counts_list").list.mean(), + TAXON_2_mean=pl.col("counts_list_2").list.mean(), + ).with_columns( + pl.when((pl.col("TAXON_1_mean") > 0) & (pl.col("TAXON_2_mean") > 0)) + .then((pl.col("TAXON_1_mean") / pl.col("TAXON_2_mean")).log(base=2)) + .when(pl.col("TAXON_1_mean") == pl.col("TAXON_2_mean")) + .then(0.0) + .otherwise(None) + .alias("log2_mean(TAXON_1/TAXON_2)"), + ) + + p_values = [] + counts1_series = stats_df["counts_list"] + counts2_series = stats_df["counts_list_2"] + + for i in range(stats_df.height): + c1 = counts1_series[i] + c2 = counts2_series[i] + + if len(c1.unique()) == 1 and len(c2.unique()) == 1 and c1[0] == c2[0]: + p_values.append(1.0) + continue + try: + pvalue = stats.mannwhitneyu(c1, c2, alternative="two-sided").pvalue + p_values.append(pvalue) + except ValueError: + p_values.append(1.0) + + final_df = stats_df.with_columns( + pl.Series("mwu_pvalue(TAXON_1 vs. TAXON_2)", p_values) + ).drop("counts_list", "counts_list_2") + + output_df = final_df.select( + pl.col("cluster_id").alias("#cluster_id"), + pl.col("label").alias("TAXON_1"), + pl.col("TAXON_1_mean"), + pl.col("label_2").alias("TAXON_2"), + pl.col("TAXON_2_mean"), + pl.col("log2_mean(TAXON_1/TAXON_2)"), + pl.col("mwu_pvalue(TAXON_1 vs. TAXON_2)"), + ).sort("#cluster_id", "TAXON_1", "TAXON_2") + + file_path = f"{output_dir}/{attribute}/{attribute}.pairwise_representation_test.txt" + print(f"[✓] Writing: {file_path}") + output_df.write_csv(file_path, separator="\t") + + generate_pairwise_plots( + pairwise_df=output_df, + attribute=attribute, + output_dir=output_dir, + ) + + +def add_status_and_TAXON_protein_count_columns( + df: pl.DataFrame, + label_to_taxons: Dict[str, Dict[str, set]], +) -> pl.DataFrame: + expressions = [] + + for attribute, label_dict in label_to_taxons.items(): + for label_group, valid_taxa in label_dict.items(): + in_valid = pl.col("taxons").list.eval(pl.element().is_in(valid_taxa)) + + status_col_name = f"{attribute}_{label_group}_cluster_status" + status_expr = ( + pl.when(in_valid.list.any()) + .then(pl.lit("present")) + .otherwise(pl.lit("absent")) + .alias(status_col_name) + ) + + count_col_name = f"{attribute}_{label_group}_TAXON_protein_count" + count_expr = in_valid.list.sum().alias(count_col_name) + + expressions.extend([status_expr, count_expr]) + + return df.with_columns(expressions) + + +def add_coverage_column( + df: pl.DataFrame, + attribute: str, + level: str, + label_to_taxons: dict, +) -> pl.DataFrame: + expected_taxons = set(label_to_taxons[attribute][level]) + expected_taxons_len = len(expected_taxons) + + if expected_taxons_len == 0: + return df.with_columns(pl.lit("N/A").alias("TAXON_coverage")) + + numerator = ( + pl.col("taxons").list.unique().list.set_intersection(expected_taxons).list.len() + ) + + coverage_numeric = numerator / expected_taxons_len + + coverage_str_padded = ( + (coverage_numeric * 100).round(0).cast(pl.Int64).cast(pl.Utf8).str.zfill(3) + ) + + coverage_expr = ( + coverage_str_padded.str.slice(0, length=coverage_str_padded.str.len_chars() - 2) + + "." + + coverage_str_padded.str.slice(-2) + ) + + return df.with_columns(coverage_expr.alias("TAXON_coverage")) + + +def add_taxon_split_columns( + df: pl.DataFrame, + attribute: str, + level: str, + label_to_taxons: dict, + min_proteomes: int = 2, + test: str = "mannwhitneyu", +) -> pl.DataFrame: + """ + Calculates taxon statistics using a fully vectorized Polars-native approach. + """ + valid_taxa = set(label_to_taxons[attribute][level]) + + if "cluster_id" not in df.columns: + df = df.with_row_count(name="cluster_id") + + partitioned_counts = ( + df.select(["cluster_id", "taxons"]) + .explode("taxons") + .group_by("cluster_id", "taxons") + .agg(pl.count().alias("count")) + .with_columns(pl.col("taxons").is_in(valid_taxa).alias("is_inside")) + ) + + grouped_lists = partitioned_counts.group_by("cluster_id", maintain_order=True).agg( + pl.col("count").filter(pl.col("is_inside")).alias("inside_counts"), + pl.col("count").filter(~pl.col("is_inside")).alias("outside_counts"), + ) + + stats_results = ( + grouped_lists.with_columns( + inside_len=pl.col("inside_counts").list.len(), + outside_len=pl.col("outside_counts").list.len(), + ) + .with_columns( + is_valid=pl.when( + (pl.col("inside_len") >= min_proteomes) + & (pl.col("outside_len") >= min_proteomes) + ) + .then(True) + .otherwise(False), + mean_inside=pl.col("inside_counts").list.mean(), + mean_outside=pl.col("outside_counts").list.mean(), + ) + .with_columns( + log2_ratio=pl.when(pl.col("mean_outside") > 0) + .then((pl.col("mean_inside") / pl.col("mean_outside")).log(base=2)) + .otherwise(None), + ) + .with_columns( + representation=pl.when(pl.col("log2_ratio") > 0) + .then(pl.lit("enriched")) + .otherwise(pl.lit("depleted")), + ) + ) + + def compute_p_value(struct: dict) -> str | None: + """A small, focused function just for the p-value.""" + inside = struct.get("inside_counts") + outside = struct.get("outside_counts") + try: + if not struct.get("is_valid"): + return None + p = stats.mannwhitneyu(inside, outside, alternative="two-sided").pvalue + return f"{p:.16f}".rstrip("0").rstrip(".") + except Exception: + return None + + final_stats = stats_results.with_columns( + p_value=pl.struct(["inside_counts", "outside_counts", "is_valid"]).map_elements( + compute_p_value, return_dtype=pl.Utf8 + ) + ).select( + pl.col("cluster_id"), + pl.when(pl.col("is_valid")) + .then(pl.col("mean_inside")) + .otherwise(pl.lit("N/A")) + .alias("TAXON_mean_count"), + pl.when(pl.col("is_valid")) + .then(pl.col("mean_outside")) + .otherwise(pl.lit("N/A")) + .alias("non_taxon_mean_count"), + pl.when(pl.col("is_valid")) + .then(pl.col("log2_ratio").cast(pl.Utf8)) + .otherwise(pl.lit("N/A")) + .alias("log2_mean(TAXON/others)"), + pl.col("p_value").fill_null(pl.lit("N/A")).alias("pvalue(TAXON vs. others)"), + pl.when(pl.col("is_valid")) + .then(pl.col("representation")) + .otherwise(pl.lit("N/A")) + .alias("representation"), + ) + + return df.join(final_stats, on="cluster_id", how="left") + + +def add_all_taxon_info_columns( + df: pl.DataFrame, + label_to_taxons: Dict[str, Dict[str, set]], +) -> pl.DataFrame: + expressions = [] + + for attribute, label_dict in label_to_taxons.items(): + for label_group, valid_taxa in label_dict.items(): + prefix = f"{attribute}_{label_group}" + taxons_sorted_unique = pl.col("taxons").list.unique().list.sort() + + taxon_mask = taxons_sorted_unique.list.eval(pl.element().is_in(valid_taxa)) + non_taxon_mask = taxons_sorted_unique.list.eval( + ~pl.element().is_in(valid_taxa) + ) + + taxon_list_str = taxons_sorted_unique.list.eval( + pl.element().filter(pl.element().is_in(valid_taxa)) + ) + non_taxon_list_str = taxons_sorted_unique.list.eval( + pl.element().filter(~pl.element().is_in(valid_taxa)) + ) + + expressions.extend( + [ + taxon_mask.list.sum().alias(f"{prefix}_TAXON_count"), + non_taxon_mask.list.sum().alias(f"{prefix}_non_TAXON_count"), + pl.when(taxon_list_str.list.len() == 0) + .then(pl.lit("N/A")) + .otherwise(taxon_list_str.list.join(",")) + .alias(f"{prefix}_TAXON_taxa"), + pl.when(non_taxon_list_str.list.len() == 0) + .then(pl.lit("N/A")) + .otherwise(non_taxon_list_str.list.join(",")) + .alias(f"{prefix}_non_TAXON_taxa"), + ] + ) + + return df.with_columns(expressions) + + +def get_cluster_metrics( + attribute: str, + label_group: str, + label_to_taxons: Dict[str, Dict[str, set]], + base_metrics_df: pl.DataFrame, +) -> pl.DataFrame: + cols_for_loop = [ + "cluster_id", + "cluster_protein_count", + "cluster_proteome_count", + "taxons", + f"{attribute}_cluster_type", + f"{attribute}_{label_group}_TAXON_protein_count", + f"{attribute}_{label_group}_cluster_status", + f"{attribute}_{label_group}_TAXON_count", + f"{attribute}_{label_group}_non_TAXON_count", + f"{attribute}_{label_group}_TAXON_taxa", + f"{attribute}_{label_group}_non_TAXON_taxa", + ] + + metrics_df = ( + base_metrics_df.select(cols_for_loop) + .pipe( + add_coverage_column, + attribute=attribute, + level=label_group, + label_to_taxons=label_to_taxons, + ) + .pipe( + add_taxon_split_columns, + attribute=attribute, + level=label_group, + label_to_taxons=label_to_taxons, + ) + ) + + final_order = [ + "#cluster_id", + "cluster_status", + "cluster_type", + "cluster_protein_count", + "cluster_proteome_count", + "TAXON_protein_count", + "TAXON_mean_count", + "non_taxon_mean_count", + "representation", + "log2_mean(TAXON/others)", + "pvalue(TAXON vs. others)", + "TAXON_coverage", + "TAXON_count", + "non_TAXON_count", + "TAXON_taxa", + "non_TAXON_taxa", + ] + + rename_map = { + "cluster_id": "#cluster_id", + f"{attribute}_cluster_type": "cluster_type", + f"{attribute}_{label_group}_cluster_status": "cluster_status", + f"{attribute}_{label_group}_TAXON_protein_count": "TAXON_protein_count", + f"{attribute}_{label_group}_TAXON_count": "TAXON_count", + f"{attribute}_{label_group}_non_TAXON_count": "non_TAXON_count", + f"{attribute}_{label_group}_TAXON_taxa": "TAXON_taxa", + f"{attribute}_{label_group}_non_TAXON_taxa": "non_TAXON_taxa", + } + + final_df = metrics_df.rename(rename_map) + final_df = final_df.select(final_order).sort("#cluster_id") + return final_df + + +def get_all_cluster_metrics( + attributes: List[str], + unique_label_values: Dict[str, List[str]], + cluster_df: pl.DataFrame, + taxon_label_dict: Dict[str, Dict[str, str]], + base_output_dir: str, + min_proteomes=2, +) -> None: + label_to_taxons: Dict[str, Dict[str, set]] = { + attr: { + label: { + taxon + for taxon, labels in taxon_label_dict.items() + if labels.get(attr) == label + } + for label in unique_label_values[attr] + } + for attr in attributes + } + + base_metrics_df = ( + cluster_df.pipe( + add_status_and_TAXON_protein_count_columns, + label_to_taxons, + ) + .rename({"TAXON_count": "cluster_proteome_count"}) + .pipe(add_all_taxon_info_columns, label_to_taxons) + ) + + for attribute in attributes: + for label_group in unique_label_values[attribute]: + metrics_df = get_cluster_metrics( + attribute=attribute, + label_group=label_group, + label_to_taxons=label_to_taxons, + base_metrics_df=base_metrics_df, + ) + + file_path = f"{base_output_dir}/{attribute}/{attribute}.{label_group}.cluster_metrics.txt" + logger.info(f"[✓] Writing: {file_path}") + metrics_df.write_csv(file_path, separator="\t") + + # generate_pairwise_representation_test( + # cluster_df=cluster_df, + # attribute=attribute, + # label_to_taxons=label_to_taxons, + # output_dir=base_output_dir, + # ) + # generate_background_representation_test( + # cluster_df=cluster_df, + # attribute=attribute, + # label_to_taxons=label_to_taxons, + # output_dir=base_output_dir, + # ) diff --git a/src/internal/cluster_summary.py b/src/internal/cluster_summary.py new file mode 100644 index 0000000..a5b9676 --- /dev/null +++ b/src/internal/cluster_summary.py @@ -0,0 +1,222 @@ +import logging +from typing import List + +import polars as pl + +logger = logging.getLogger("kinfin_logger") + + +def add_attribute_count_median_cov( + cluster_df: pl.DataFrame, + config_df: pl.DataFrame, + attribute: str, +) -> pl.DataFrame: + if attribute == "TAXON": + non_empty_clusters = cluster_df.filter(pl.col("taxons").list.len() > 0) + if non_empty_clusters.is_empty(): + return cluster_df + + count_wide = ( + non_empty_clusters.select(["cluster_id", "taxons"]) + .explode("taxons") + .pivot( + values="taxons", + index="cluster_id", + columns="taxons", + aggregate_function="count", + ) + .fill_null(0) + ) + + sorted_cols = ["cluster_id"] + sorted( + [c for c in count_wide.columns if c != "cluster_id"] + ) + count_wide = count_wide.select(sorted_cols) + + rename_dict = { + col: f"{col}_count" for col in count_wide.columns if col != "cluster_id" + } + count_wide = count_wide.rename(rename_dict) + + result = cluster_df.join(count_wide, on="cluster_id", how="left") + + count_cols_to_fill = [c for c in rename_dict.values() if c in result.columns] + fill_expressions = [pl.col(c).fill_null(0) for c in count_cols_to_fill] + + if not fill_expressions: + return result + return result.with_columns(fill_expressions) + + labels_df = config_df.select( + pl.col("TAXON").alias("taxon"), + pl.col(attribute).alias("label"), + ).drop_nulls() + + taxon_counts = ( + cluster_df.select(["cluster_id", "taxons"]) + .explode("taxons") + .group_by(["cluster_id", "taxons"]) + .agg(pl.count().alias("count_taxon")) + ) + + scaffold = cluster_df.select("cluster_id").unique().join(labels_df, how="cross") + + data_for_agg = scaffold.join( + taxon_counts, + left_on=["cluster_id", "taxon"], + right_on=["cluster_id", "taxons"], + how="left", + ).with_columns(pl.col("count_taxon").fill_null(0)) + + total_taxons_per_label = labels_df.group_by("label").agg( + pl.count("taxon").alias("total_taxons_in_label") + ) + + agg = data_for_agg.group_by(["cluster_id", "label"]).agg( + pl.sum("count_taxon").alias("count"), + pl.median("count_taxon").alias("median"), + pl.col("count_taxon") + .filter(pl.col("count_taxon") > 0) + .count() + .alias("present_taxons_count"), + ) + + agg_with_cov = agg.join(total_taxons_per_label, on="label", how="left") + + agg_with_cov = ( + agg_with_cov.with_columns( + pl.when(pl.col("total_taxons_in_label") > 0) + .then(pl.col("present_taxons_count") / pl.col("total_taxons_in_label")) + .otherwise(0.0) + .alias("coverage_numeric") + ) + .with_columns( + (pl.col("coverage_numeric") * 100) + .round(0) + .cast(pl.Int64) + .cast(pl.Utf8) + .str.zfill(3) + .alias("coverage_str") + ) + .with_columns( + ( + pl.col("coverage_str").str.slice( + 0, length=pl.col("coverage_str").str.len_chars() - 2 + ) + + "." + + pl.col("coverage_str").str.slice(-2) + ).alias("coverage") + ) + .drop(["coverage_numeric", "coverage_str"]) + ) + + def pivot_and_prepare(df, value_col, suffix, fill_value): + pivoted = df.pivot( + values=value_col, index="cluster_id", columns="label" + ).fill_null(fill_value) + + sorted_cols = ["cluster_id"] + sorted( + [c for c in pivoted.columns if c != "cluster_id"] + ) + pivoted = pivoted.select(sorted_cols) + + rename_dict = { + col: f"{col}{suffix}" for col in pivoted.columns if col != "cluster_id" + } + return pivoted.rename(rename_dict) + + count_wide = pivot_and_prepare(agg_with_cov, "count", "_count", 0) + median_wide = pivot_and_prepare(agg_with_cov, "median", "_median", 0.0) + coverage_wide = pivot_and_prepare(agg_with_cov, "coverage", "_cov", "0.00") + + result = ( + cluster_df.join(count_wide, on="cluster_id", how="left") + .join(median_wide, on="cluster_id", how="left") + .join(coverage_wide, on="cluster_id", how="left") + ) + + fill_expressions = [] + all_labels = labels_df.get_column("label").unique().to_list() + for label in all_labels: + count_col, median_col, cov_col = ( + f"{label}_count", + f"{label}_median", + f"{label}_cov", + ) + if count_col in result.columns: + fill_expressions.append(pl.col(count_col).fill_null(0)) + if median_col in result.columns: + fill_expressions.append(pl.col(median_col).fill_null(0.0)) + if cov_col in result.columns: + fill_expressions.append(pl.col(cov_col).fill_null("0.00")) + + if not fill_expressions: + return result + + return result.with_columns(fill_expressions) + + +def get_cluster_summary( + attribute: str, + config_df: pl.DataFrame, + cluster_df: pl.DataFrame, +): + summary_df = cluster_df.select( + [ + "cluster_id", + "cluster_protein_count", + "protein_median_count", + "TAXON_count", + "taxons", + f"{attribute}_cluster_type", + ] + ) + summary_df = summary_df.with_columns(pl.lit(attribute).alias("attribute")) + summary_df = summary_df.select( + [ + "cluster_id", + "cluster_protein_count", + "protein_median_count", + "TAXON_count", + "taxons", + "attribute", + f"{attribute}_cluster_type", + ] + ) + summary_df = ( + summary_df.with_columns(pl.lit("N/A").alias("protein_span_mean")) + .with_columns(pl.lit("N/A").alias("protein_span_sd")) + .pipe( + add_attribute_count_median_cov, + config_df=config_df, + attribute=attribute, + ) + .drop("taxons") + .rename( + { + "cluster_id": "#cluster_id", + f"{attribute}_cluster_type": "attribute_cluster_type", + } + ) + .sort("#cluster_id") + ) + + return summary_df + + +def get_all_cluster_summaries( + attributes: List[str], + cluster_df: pl.DataFrame, + config_df: pl.DataFrame, + base_output_dir: str, +) -> None: + + for attribute in attributes: + summary_df = get_cluster_summary( + cluster_df=cluster_df, + attribute=attribute, + config_df=config_df, + ) + file_path = f"{base_output_dir}/{attribute}/{attribute}.cluster_summary.txt" + logger.info(f"[✓] Writing: {file_path}") + summary_df.write_csv(file_path, separator="\t") diff --git a/src/core/config.py b/src/internal/config.py similarity index 65% rename from src/core/config.py rename to src/internal/config.py index 098c545..b84eca6 100644 --- a/src/core/config.py +++ b/src/internal/config.py @@ -1,6 +1,3 @@ -ATTRIBUTE_RESERVED = ["IDX", "OUT", "TAXID"] -SUPPORTED_TESTS = {"welch", "mannwhitneyu", "ttest", "ks", "kruskal"} -SUPPORTED_PLOT_FORMATS = {"png", "pdf", "svg"} SUPPORTED_TAXRANKS = { "superkingdom", "kingdom", @@ -13,3 +10,4 @@ "genus", "species", } +ATTRIBUTE_RESERVED = ["IDX", "OUT", "TAXID"] diff --git a/src/internal/counts.py b/src/internal/counts.py new file mode 100644 index 0000000..e0b6aa3 --- /dev/null +++ b/src/internal/counts.py @@ -0,0 +1,35 @@ +import os + +import polars as pl + + +def get_cluster_counts_by_taxon( + config_df: pl.DataFrame, + cluster_df: pl.DataFrame, + base_output_dir: str, +) -> None: + taxa = config_df["TAXON"].unique().sort() + exploded = cluster_df.select(["cluster_id", "taxons"]).explode("taxons") + + counts = ( + exploded.group_by(["cluster_id", "taxons"]) + .count() + .rename({"count": "taxon_count"}) + ) + + result_df = counts.pivot( + values="taxon_count", + index="cluster_id", + columns="taxons", + aggregate_function=None, + ).fill_null(0) + + result_df = result_df.rename({"cluster_id": "#ID"}) + result_df = result_df.select(["#ID"] + taxa.to_list()) + + result_df = result_df.sort("#ID") + + output_path = f"{base_output_dir}/cluster_counts_by_taxon.txt" + os.makedirs(os.path.dirname(output_path), exist_ok=True) + result_df.write_csv(output_path, separator="\t") + return result_df diff --git a/src/core/logger.py b/src/internal/logger.py similarity index 100% rename from src/core/logger.py rename to src/internal/logger.py diff --git a/src/internal/parsers.py b/src/internal/parsers.py new file mode 100644 index 0000000..a6e1860 --- /dev/null +++ b/src/internal/parsers.py @@ -0,0 +1,297 @@ +import json +import os +from collections import OrderedDict +from typing import List, Tuple + +import polars as pl + +from internal.config import ATTRIBUTE_RESERVED + + +def nodesdb(filepath: str, outpath: str) -> pl.DataFrame: + if os.path.exists(outpath): + return pl.read_parquet(outpath) + + db = pl.read_csv( + filepath, + has_header=False, + separator="\t", + comment_prefix="#", + new_columns=["node", "rank", "name", "parent"], + ) + os.makedirs("db", exist_ok=True) + db.write_parquet(outpath) + return db + + +def clusterfile( + cluster_path: str, + config_df: pl.DataFrame, + output_dir: str, +) -> pl.DataFrame: + with open(cluster_path, "r") as file: + content = file.read().strip().split("\n") + + cluster_ids = [] + protein_clusters = [] + + for line in content: + cluster_id, proteins = line.split(": ") + prots = proteins.split() + cluster_ids.extend([cluster_id] * len(prots)) + protein_clusters.extend(prots) + + exploded = pl.DataFrame( + {"cluster_id": cluster_ids, "protein_cluster": protein_clusters} + ) + exploded = exploded.with_columns( + pl.col("protein_cluster") + .str.split_exact(".", 1) + .struct.field("field_0") + .alias("taxon") + ) + + available_proteomes = set(config_df["TAXON"].to_list()) + total_proteins = len(exploded) + total_clusters = exploded["cluster_id"].n_unique() + total_proteomes = len(available_proteomes) + + included_df = exploded.filter(pl.col("taxon").is_in(available_proteomes)) + excluded_df = exploded.filter(~pl.col("taxon").is_in(available_proteomes)) + + filtered_proteins = len(included_df) + filtered_clusters = included_df["cluster_id"].n_unique() + included_proteins_count = included_df["protein_cluster"].n_unique() + excluded_proteins_count = excluded_df["protein_cluster"].n_unique() + + included_proteomes = ( + included_df.group_by("taxon") + .agg(pl.len()) + .rename({"len": "count"}) + .to_dict(as_series=False) + ) + included_proteomes_dict = dict( + sorted(zip(included_proteomes["taxon"], included_proteomes["count"])) + ) + + excluded_proteomes = ( + excluded_df.group_by("taxon") + .agg(pl.len()) + .rename({"len": "count"}) + .to_dict(as_series=False) + ) + excluded_proteomes_dict = dict( + sorted(zip(excluded_proteomes["taxon"], excluded_proteomes["count"])) + ) + + stats = OrderedDict( + [ + ("total_clusters", total_clusters), + ("total_proteins", total_proteins), + ("total_proteomes", total_proteomes), + ("filtered_clusters", filtered_clusters), + ("filtered_proteins", filtered_proteins), + ("included_proteins_count", included_proteins_count), + ("excluded_proteins_count", excluded_proteins_count), + ("included_proteomes", included_proteomes_dict), + ("excluded_proteomes", excluded_proteomes_dict), + ("included_proteins", included_df["protein_cluster"].sort().to_list()), + ("excluded_proteins", excluded_df["protein_cluster"].sort().to_list()), + ] + ) + + with open(f"{output_dir}/summary.json", "w") as mf: + json.dump(stats, mf, separators=(", ", ": "), indent=4) + + agg_df = included_df.group_by("cluster_id").agg( + [ + pl.col("taxon").sort().alias("taxons"), + pl.col("protein_cluster") + .str.split_exact(".", 1) + .struct.field("field_1") + .sort() + .alias("sequences"), + pl.col("protein_cluster").sort().alias("protein_cluster"), + pl.len().alias("cluster_protein_count"), + pl.col("taxon").n_unique().alias("TAXON_count"), + ] + ) + + median_counts = ( + included_df.group_by(["cluster_id", "taxon"]) + .len() + .group_by("cluster_id") + .agg(pl.median("len").alias("protein_median_count")) + ) + return agg_df.join(median_counts, on="cluster_id", how="left") + + +def add_taxid_attributes( + config_df: pl.DataFrame, + nodesdb: pl.DataFrame, + taxranks: List[str] = ["phylum", "order", "genus"], +) -> pl.DataFrame: + if "TAXID" in config_df.columns: + nodes_lookup = { + row["node"]: { + "rank": row["rank"], + "name": row["name"], + "parent": row["parent"], + } + for row in nodesdb.iter_rows(named=True) + } + + def get_lineage(taxid: int) -> dict: + lineage = {} + current = taxid + while current in nodes_lookup: + entry = nodes_lookup[current] + rank = entry["rank"] + if rank in taxranks and rank not in lineage: + lineage[rank] = entry["name"] + if len(lineage) == len(taxranks): + break + current = entry["parent"] + return lineage + + taxid_series = config_df.get_column("TAXID") + taxid_lineages = [get_lineage(tid) for tid in taxid_series.to_list()] + + new_columns = {} + for rank in taxranks: + new_columns[rank] = [ + lineage.get(rank, "not_available") for lineage in taxid_lineages + ] + + for rank in taxranks: + config_df = config_df.with_columns(pl.Series(rank, new_columns[rank])) + + return config_df + + +def configfile(filepath: str, nodesdb: pl.DataFrame) -> Tuple[pl.DataFrame, List[str]]: + config_df = ( + pl.read_csv(filepath, has_header=True, separator=",") + .with_columns(pl.lit("all").alias("all")) + .pipe(add_taxid_attributes, nodesdb=nodesdb) + ) + attributes = config_df.columns + attributes = [ + attribute + for attribute in attributes + if attribute.strip("#") not in ATTRIBUTE_RESERVED + ] + return config_df, attributes + + +# def entrylist(filepath: str = "data/entry.list") -> pl.DataFrame: +# return pl.read_csv( +# filepath, +# has_header=True, +# separator="\t", +# ) + + +# def interpro2go(filepath: str = "data/interpro2go") -> pl.DataFrame: +# go_mapping_dict: Dict[str, str] = {} + +# with open(filepath) as go_mapping_f: +# for line in go_mapping_f: +# if not line.startswith("!"): +# temp: List[str] = line.replace(" > ", "|").split("|") +# go_string: List[str] = temp[1].split(";") +# go_desc: str = go_string[0].replace("GO:", "").strip() +# go_id: str = go_string[1].strip() + +# if go_id not in go_mapping_dict: +# go_mapping_dict[go_id] = go_desc +# elif go_desc != go_mapping_dict[go_id]: +# error_msg: str = f"[ERROR] : Conflicting descriptions for {go_id}" +# raise ValueError(error_msg) + +# go_ids: List[str] = list(go_mapping_dict.keys()) +# go_descriptions: List[str] = list(go_mapping_dict.values()) + +# return pl.DataFrame({"GO_ID": go_ids, "GO_Description": go_descriptions}) + + +# def cluster_file(filepath: str) -> pl.DataFrame: +# with open(filepath, "r") as file: +# content = file.read() + +# data: List[Tuple[str, List[str], List[str], List[str]]] = [] +# for line in content.strip().split("\n"): +# cluster_id, proteins = line.strip().split(": ") +# protein_list = proteins.split() +# taxons = [p.split(".")[0] for p in protein_list] +# sequences = [p.split(".")[1] for p in protein_list] +# data.append((cluster_id, taxons, sequences, protein_list)) + +# df = pl.DataFrame( +# data, +# schema=["cluster_id", "taxons", "sequences", "protein_cluster"], +# schema_overrides={ +# "taxons": pl.List(pl.Utf8), +# "sequences": pl.List(pl.Utf8), +# }, +# orient="row", +# ) +# df = df.with_columns( +# [pl.col("taxons").list.sort(), pl.col("sequences").list.sort()] +# ) +# return df + + +# def clusterfile_basic() -> pl.DataFrame: +# with open(CLUSTER_PATH, "r") as file: +# content = file.read() + +# data: List[Tuple[str, List[str]]] = [] +# for line in content.strip().split("\n"): +# cluster_id, proteins = line.strip().split(": ") +# protein_list = proteins.split() +# data.append((cluster_id, protein_list)) + + +# df = pl.DataFrame( +# data, +# schema=["cluster_id", "protein_cluster"], +# schema_overrides={ +# "taxons": pl.List(pl.Utf8), +# "sequences": pl.List(pl.Utf8), +# }, +# orient="row", +# ) + +# return df + + +# def sequence_ids(filepath: str) -> pl.DataFrame: +# with open(filepath, "r") as sequence_id_f: +# lines = sequence_id_f.readlines() + +# data: List[Dict[str, str]] = [] +# for line in lines: +# col = line.strip().split(": ") +# sequence_id = col[0].split("_")[0] +# species_id = sequence_id.split("_")[0] +# protein_id = ( +# col[1] +# .replace(":", "_") +# .replace(",", "_") +# .replace("(", "_") +# .replace(")", "_") +# ) +# proteome_id = protein_id.split(".")[0] +# data.append( +# { +# "sequence_id": sequence_id, +# "proteome_id": proteome_id, +# "protein_id": protein_id, +# "species_id": species_id, +# } +# ) + +# df = pl.DataFrame(data) +# df.write_parquet("db/sequenceids.parquet") +# return df diff --git a/src/internal/plots.py b/src/internal/plots.py new file mode 100644 index 0000000..74311cc --- /dev/null +++ b/src/internal/plots.py @@ -0,0 +1,207 @@ +import logging +import os +import random +import warnings +from collections import Counter, defaultdict +from typing import Tuple + +import matplotlib as mat +import matplotlib.pyplot as plt +import numpy as np +import polars as pl +from matplotlib.ticker import FormatStrFormatter + +logger = logging.getLogger("kinfin_logger") + +warnings.filterwarnings( + "ignore", message="No artists with labels found to put in legend.*" +) + + +mat.use("agg") + +plt.style.use("ggplot") + +mat.rc("ytick", labelsize=20) +mat.rc("xtick", labelsize=20) +mat.rcParams.update({"font.size": 22}) + + +def plot_cluster_sizes( + cluster_df: pl.DataFrame, + base_output_dir: str, + plot_format: str = "png", + plotsize: Tuple[int, int] = (24, 12), + fontsize: int = 18, +) -> None: + counts = cluster_df.get_column("cluster_protein_count").to_list() + cluster_size_counter = Counter(counts) + x_values = np.array(list(cluster_size_counter.keys())) + y_values = np.array(list(cluster_size_counter.values())) + + output_path = os.path.join(base_output_dir, "cluster_size_distribution.tsv") + with open(output_path, "w") as f: + f.write("cluster_size\tcount\n") + for size, count in sorted(cluster_size_counter.items()): + f.write(f"{size}\t{count}\n") + logger.info(f"[✓] Writing: {output_path}") + + fig, ax = plt.subplots(figsize=plotsize) + ax.set_facecolor("white") + + ax.scatter(x_values, y_values, marker="o", alpha=0.8, s=100) + + ax.set_xlabel("Cluster size", fontsize=fontsize) + ax.set_ylabel("Count", fontsize=fontsize) + ax.set_yscale("log") + ax.set_xscale("log") + + plt.margins(0.8) + plt.gca().set_ylim(bottom=0.8) + plt.gca().set_xlim(left=0.8) + + ax.xaxis.set_major_formatter(FormatStrFormatter("%.0f")) + ax.yaxis.set_major_formatter(FormatStrFormatter("%.0f")) + + fig.tight_layout() + + ax.grid(True, which="major", color="lightgrey", linewidth=1) + ax.grid(True, which="minor", color="lightgrey", linewidth=0.5) + + plot_out_path = os.path.join( + base_output_dir, f"cluster_size_distribution.{plot_format}" + ) + logger.info(f"[✓] Plotting: {plot_out_path}") + + fig.savefig(plot_out_path, format=plot_format) + plt.close() + + +def generate_kinfin_rarefaction_plots( + cluster_df: pl.DataFrame, + config_df: pl.DataFrame, + attributes: list[str], + base_output_dir: str, + repetitions: int = 30, + plot_size: tuple = (24, 12), + font_size: int = 18, + plot_format: str = "png", +): + os.makedirs(base_output_dir, exist_ok=True) + + non_singleton_clusters = cluster_df.filter(pl.col("TAXON_count") > 1) + clusters_by_taxon = defaultdict(set) + for taxon, cid in ( + non_singleton_clusters.select(["taxons", "cluster_id"]) + .explode("taxons") + .iter_rows() + ): + clusters_by_taxon[taxon].add(cid) + + for attribute in attributes: + attribute_output_path = os.path.join(base_output_dir, attribute) + os.makedirs(attribute_output_path, exist_ok=True) + + fig, ax = plt.subplots(figsize=plot_size) + ax.set_facecolor("white") + + color_map = plt.cm.Paired(np.linspace(0, 1, config_df[attribute].n_unique())) + + plot_data_by_level = {} + max_samples_on_plot = 0 + + # This list will store the raw data for the TSV file + data_for_saving = [ + "level\tsample_size\tmedian_clusters\tmin_clusters\tmax_clusters" + ] + + all_levels = config_df[attribute].unique().sort() + for i, level in enumerate(all_levels): + taxa_for_level = config_df.filter(pl.col(attribute) == level)[ + "TAXON" + ].to_list() + + if len(taxa_for_level) <= 1: + continue + + max_samples_on_plot = max(max_samples_on_plot, len(taxa_for_level)) + + rarefaction_data = defaultdict(list) + for _ in range(repetitions): + random.shuffle(taxa_for_level) + seen_clusters = set() + for j, taxon in enumerate(taxa_for_level): + seen_clusters.update(clusters_by_taxon[taxon]) + rarefaction_data[j + 1].append(len(seen_clusters)) + + sample_sizes = sorted(rarefaction_data.keys()) + medians = [np.median(rarefaction_data[s]) for s in sample_sizes] + mins = [np.min(rarefaction_data[s]) for s in sample_sizes] + maxs = [np.max(rarefaction_data[s]) for s in sample_sizes] + + plot_data_by_level[level] = { + "samples": sample_sizes, + "medians": medians, + "mins": mins, + "maxs": maxs, + "color": color_map[i], + } + + # Populate the data for saving + for k, sample_size in enumerate(sample_sizes): + row = f"{level}\t{sample_size}\t{medians[k]}\t{mins[k]}\t{maxs[k]}" + data_for_saving.append(row) + + # Plotting section + for level, data in plot_data_by_level.items(): + ax.plot( + data["samples"], + data["medians"], + color=data["color"], + label=level, + linewidth=2, + ) + ax.fill_between( + data["samples"], + data["mins"], + data["maxs"], + color=data["color"], + alpha=0.4, + ) + + ax.set_title( + f"Rarefaction Curve by {attribute}", + fontsize=font_size + 2, + fontweight="bold", + ) + ax.set_xlabel("Number of Proteomes Sampled", fontsize=font_size) + ax.set_ylabel("Count of Non-Singleton Clusters", fontsize=font_size) + ax.grid(True, which="both", linestyle="--", linewidth=1, color="lightgrey") + ax.set_xlim(0, max_samples_on_plot + 1) + legend = ax.legend( + title=attribute.capitalize(), + fontsize=font_size, + loc="lower right", + frameon=True, + ) + legend.get_frame().set_facecolor("white") + + # Save the plot + plot_filename = os.path.join( + attribute_output_path, f"{attribute}.rarefaction_curve.{plot_format}" + ) + try: + fig.savefig(plot_filename, format=plot_format, bbox_inches="tight") + except Exception as e: + logger.info(f"[ERROR] Could not save plot. {e}") + plt.close(fig) + + # *** NEW: Save the aggregated data to a TSV file *** + data_filename = os.path.join( + attribute_output_path, f"{attribute}.rarefaction_data.tsv" + ) + try: + with open(data_filename, "w") as f: + f.write("\n".join(data_for_saving)) + except Exception as e: + logger.info(f"[ERROR] Could not save data file. {e}") diff --git a/src/internal/utils.py b/src/internal/utils.py new file mode 100644 index 0000000..b712e60 --- /dev/null +++ b/src/internal/utils.py @@ -0,0 +1,39 @@ +import os +import shutil +from typing import List, Optional + + +def check_file(filepath: Optional[str], install_kinfin: bool = False) -> None: + """ + Check if a file exists. + + Args: + filepath (str): Path to the file to be checked. + + Raises: + FileNotFoundError: If the file does not exist. + """ + + if filepath is not None and not os.path.isfile(filepath): + error_msg = f"[ERROR] - file {filepath} not found." + if install_kinfin: + error_msg += " Please run the install script to download kinfin." + raise FileNotFoundError(error_msg) + + +def setup_dirs(output_dir, attributes: List[str]) -> None: + files_to_keep = { + f for f in os.listdir(output_dir) if f.endswith(".status") or f == "config.txt" + } + + if os.path.exists(output_dir): + for item in os.listdir(output_dir): + item_path = os.path.join(output_dir, item) + if item not in files_to_keep: + if os.path.isdir(item_path): + shutil.rmtree(item_path) + else: + os.remove(item_path) + os.makedirs(output_dir, exist_ok=True) + for attribute in attributes: + os.makedirs(os.path.join(output_dir, attribute), exist_ok=True) diff --git a/src/main.py b/src/main.py index dc1737a..61593c5 100755 --- a/src/main.py +++ b/src/main.py @@ -1,45 +1,75 @@ #!/usr/bin/env python3 +import argparse +import logging import os import sys +import time from api import run_server -from cli import run_cli -from cli.commands import parse_args -from core.input import InputData, ServeArgs -from core.utils import check_file +from internal import analyse +from internal.utils import check_file -if __name__ == "__main__": - # Without these files, application won't start +logger = logging.getLogger("kinfin_logger") + + +def main(): + # ---- Initial Setup ---- base_dir = os.getcwd() nodesdb_f = os.path.join(base_dir, "data/nodesdb.txt") - pfam_mapping_f = os.path.join(base_dir, "data/Pfam-A.clans.tsv.gz") - ipr_mapping_f = os.path.join(base_dir, "data/entry.list") - go_mapping_f = os.path.join(base_dir, "data/interpro2go") + ndb_f = os.path.join(base_dir, "db/ndb.parquet") + + # TODO: Handle pfam, ipr, go + # pfam_mapping_f = os.path.join(base_dir, "data/Pfam-A.clans.tsv.gz") + # ipr_mapping_f = os.path.join(base_dir, "data/entry.list") + # go_mapping_f = os.path.join(base_dir, "data/interpro2go") try: check_file(nodesdb_f, install_kinfin=True) - check_file(pfam_mapping_f, install_kinfin=True) - check_file(ipr_mapping_f, install_kinfin=True) - check_file(go_mapping_f, install_kinfin=True) except FileNotFoundError as e: sys.exit(str(e)) - args = parse_args(nodesdb_f, pfam_mapping_f, ipr_mapping_f, go_mapping_f) - - if isinstance(args, ServeArgs): - # cluster_f, sequence_ids_f, and taxon_idx_mapping_file will be set dynamically at runtime (from /init) - run_server( - args=args, - nodesdb_f=nodesdb_f, - go_mapping_f=go_mapping_f, - ipr_mapping_f=ipr_mapping_f, - pfam_mapping_f=pfam_mapping_f, - cluster_f="", # dummy - sequence_ids_f="", # dummy - taxon_idx_mapping_file="", # dummy - ) - elif isinstance(args, InputData): - run_cli(args) - else: - sys.exit("[ERROR] - Invalid input provided.") + # ---- Parse Arguments ---- + parser = argparse.ArgumentParser( + description="Kinfin proteome cluster analysis tool" + ) + + subparsers = parser.add_subparsers(title="command", required=True, dest="command") + api_parser = subparsers.add_parser("serve", help="Start the server") + api_parser.add_argument( + "-p", + "--port", + type=int, + default=8000, + help="Port number for the server (default: 8000)", + ) + + cli_parser = subparsers.add_parser("analyse", help="Perform analysis") + required_group = cli_parser.add_argument_group("Required Arguments") + required_group.add_argument( + "-g", + "--cluster_file", + help="OrthologousGroups.txt produced by OrthoFinder", + required=True, + ) + required_group.add_argument( + "-c", "--config_file", help="Config file (in CSV format)", required=True + ) + general_group = cli_parser.add_argument_group("General Options") + general_group.add_argument("-o", "--output_path", help="Output prefix") + + args = parser.parse_args() + if args.command == "serve": + # ---- RUN SERVER ---- + run_server(port=args.port, nodesdb_f=nodesdb_f, ndb_f=ndb_f, cluster_f="") + elif args.command == "analyse": + # ---- RUN POLARS ANALYSIS ---- + analyse(args, nodesdb_f=nodesdb_f, ndb_f=ndb_f) + + +if __name__ == "__main__": + overall_start = time.time() + main() + overall_end = time.time() + overall_elapsed = overall_end - overall_start + logger.info(f"Took {overall_elapsed}s to run kinfin.")