Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| import io, os, re, math, zipfile | |
| from typing import Dict, List, Tuple, Set, Optional | |
| import pandas as pd | |
| from Bio import SeqIO | |
| from statsmodels.stats.multitest import multipletests | |
| from scipy.stats import fisher_exact | |
| import matplotlib.pyplot as plt | |
| FA_EXT = (".fasta", ".fa", ".txt", ".fas", ".fna") | |
| def _read_fasta_bytes(name: str, data: bytes) -> List[Tuple[str, str, str]]: | |
| recs = [] | |
| with io.BytesIO(data) as bio: | |
| for rec in SeqIO.parse(io.TextIOWrapper(bio, encoding="utf-8"), "fasta"): | |
| header = str(rec.id) | |
| seq = str(rec.seq).upper().replace("\n", "").replace("\r", "") | |
| recs.append((name, header, seq)) | |
| return recs | |
| def read_uploaded_fasta_or_zip(uploaded_file) -> List[Tuple[str, str, str]]: | |
| if uploaded_file is None: | |
| return [] | |
| name = uploaded_file.name | |
| data = uploaded_file.read() | |
| if name.lower().endswith(".zip"): | |
| results = [] | |
| with zipfile.ZipFile(io.BytesIO(data)) as z: | |
| for zi in z.infolist(): | |
| if zi.is_dir(): continue | |
| if not any(zi.filename.lower().endswith(ext) for ext in FA_EXT): | |
| continue | |
| file_bytes = z.read(zi.filename) | |
| results.extend(_read_fasta_bytes(os.path.basename(zi.filename), file_bytes)) | |
| return results | |
| else: | |
| return _read_fasta_bytes(os.path.basename(name), data) | |
| def clean_protein(seq: str) -> str: | |
| return re.sub(r"[^ACDEFGHIKLMNPQRSTVWY]", "", seq.upper()) | |
| def clean_dna(seq: str) -> str: | |
| return re.sub(r"[^ACGTUN]", "", seq.upper()) | |
| def reverse_complement_dna(seq: str) -> str: | |
| """Return reverse complement of a cleaned DNA/RNA-like sequence. | |
| Supports A, C, G, T, U, and N. U is treated as T on the | |
| input strand and complements to A. | |
| """ | |
| seq = clean_dna(seq).replace("U", "T") | |
| table = str.maketrans({"A": "T", "C": "G", "G": "C", "T": "A", "N": "N"}) | |
| return seq.translate(table)[::-1] | |
| def get_kmers_noN(sequence: str, k: int) -> List[str]: | |
| s = sequence | |
| out = [] | |
| L = len(s) | |
| for i in range(L - k + 1): | |
| kmer = s[i:i+k] | |
| if "N" not in kmer: | |
| out.append(kmer) | |
| return out | |
| def parse_k_input(k_input: str, default_single: int) -> List[int]: | |
| k_input = (k_input or "").strip() | |
| if not k_input: | |
| return [default_single] | |
| if "-" in k_input: | |
| a, b = k_input.split("-", 1) | |
| a = int(a.strip()); b = int(b.strip()) | |
| if a > b: a, b = b, a | |
| return list(range(a, b+1)) | |
| if "," in k_input: | |
| return [int(x.strip()) for x in k_input.split(",") if x.strip()] | |
| return [int(k_input)] | |
| def derive_serotype_names_from_sources(known_records: List[Tuple[str, str, str]]) -> Dict[str, str]: | |
| counts: Dict[str, int] = {} | |
| for src, header, _ in known_records: | |
| counts[src] = counts.get(src, 0) + 1 | |
| name_map: Dict[str, str] = {} | |
| for src, header, _ in known_records: | |
| if counts.get(src, 0) == 1: | |
| sero = os.path.splitext(os.path.basename(src))[0] | |
| else: | |
| sero = header.split()[0] | |
| name_map[header] = sero | |
| return name_map | |
| def compute_unique_kmers_per_serotype(serotype_to_seq: Dict[str, str], is_protein: bool, k_values: List[int]) -> Dict[str, Dict[int, Set[str]]]: | |
| all_sets: Dict[str, Dict[int, Set[str]]] = {g: {} for g in serotype_to_seq} | |
| for g, seq in serotype_to_seq.items(): | |
| seq = clean_protein(seq) if is_protein else clean_dna(seq) | |
| for k in k_values: | |
| all_sets[g][k] = set(get_kmers_noN(seq, k)) | |
| unique: Dict[str, Dict[int, Set[str]]] = {g: {k: set() for k in k_values} for g in serotype_to_seq} | |
| for k in k_values: | |
| union_all = set().union(*(all_sets[g][k] for g in all_sets)) | |
| for g in all_sets: | |
| others_union = union_all - all_sets[g][k] | |
| unique[g][k] = all_sets[g][k] - others_union | |
| return unique | |
| def _score_sequence_against_uniques(seq2: str, unique_kmers: Dict[str, Dict[int, Set[str]]], k_values: List[int], vocab_by_sero: Dict[str, int], fdr_alpha: float) -> Dict: | |
| """Classify one already-cleaned sequence against the unique k-mer database.""" | |
| unk_kmers: Dict[int, Set[str]] = {} | |
| for k in k_values: | |
| unk_kmers[k] = set(get_kmers_noN(seq2, k)) | |
| match_counts: Dict[str, int] = {} | |
| total_matches = 0 | |
| for g in unique_kmers: | |
| mg = 0 | |
| for k in k_values: | |
| mg += len(unique_kmers[g][k].intersection(unk_kmers[k])) | |
| match_counts[g] = mg | |
| total_matches += mg | |
| if total_matches == 0: | |
| predicted = "NoMatch" | |
| conf_present = 0.0 | |
| conf_vocab = 0.0 | |
| else: | |
| predicted = max(match_counts, key=match_counts.get) | |
| conf_present = match_counts[predicted] / total_matches | |
| conf_vocab = match_counts[predicted] / max(1, vocab_by_sero[predicted]) | |
| fisher_p = {} | |
| if total_matches > 0: | |
| sum_vocab_all = sum(vocab_by_sero.values()) | |
| for g in unique_kmers: | |
| a = match_counts[g] | |
| b = vocab_by_sero[g] - a | |
| c = total_matches - a | |
| d = (sum_vocab_all - vocab_by_sero[g]) - c | |
| a = max(0, a); b = max(0, b); c = max(0, c); d = max(0, d) | |
| _, p = fisher_exact([[a, b], [c, d]], alternative="greater") | |
| fisher_p[g] = p | |
| groups = list(unique_kmers.keys()) | |
| pvals = [fisher_p[g] for g in groups] | |
| _, qvals, _, _ = multipletests(pvals, alpha=fdr_alpha, method="fdr_bh") | |
| fdr_map = {g: q for g, q in zip(groups, qvals)} | |
| else: | |
| fisher_p = {g: 1.0 for g in unique_kmers} | |
| fdr_map = {g: 1.0 for g in unique_kmers} | |
| best_fdr = fdr_map.get(predicted, 1.0) if predicted != "NoMatch" else 1.0 | |
| best_group_matches = match_counts.get(predicted, 0) if predicted != "NoMatch" else 0 | |
| return { | |
| "Predicted_serotype": predicted, | |
| "Matches_total": total_matches, | |
| "Confidence_by_present": conf_present, | |
| "Confidence_by_serotype_vocab": conf_vocab, | |
| "match_counts": match_counts, | |
| "fisher_p": fisher_p, | |
| "fdr_map": fdr_map, | |
| "best_fdr": best_fdr, | |
| "best_group_matches": best_group_matches, | |
| } | |
| def _is_better_orientation(candidate: Dict, current: Optional[Dict]) -> bool: | |
| """Choose the orientation with the strongest evidence.""" | |
| if current is None: | |
| return True | |
| # Prefer more total unique k-mer matches, then more matches for the predicted group, | |
| # then higher confidence, then lower FDR. | |
| cand_key = ( | |
| candidate["Matches_total"], | |
| candidate["best_group_matches"], | |
| candidate["Confidence_by_present"], | |
| -candidate["best_fdr"], | |
| ) | |
| curr_key = ( | |
| current["Matches_total"], | |
| current["best_group_matches"], | |
| current["Confidence_by_present"], | |
| -current["best_fdr"], | |
| ) | |
| return cand_key > curr_key | |
| def classify_unknown_sequences(unknown_records: List[Tuple[str, str, str]], unique_kmers: Dict[str, Dict[int, Set[str]]], is_protein: bool, fdr_alpha: float = 0.05) -> pd.DataFrame: | |
| vocab_by_sero: Dict[str, int] = {} | |
| k_values = sorted({k for g in unique_kmers for k in unique_kmers[g]}) | |
| for g in unique_kmers: | |
| vocab_by_sero[g] = sum(len(unique_kmers[g][k]) for k in k_values) | |
| results = [] | |
| for src, header, seq in unknown_records: | |
| if is_protein: | |
| seq2 = clean_protein(seq) | |
| best = _score_sequence_against_uniques(seq2, unique_kmers, k_values, vocab_by_sero, fdr_alpha) | |
| orientation_used = "input" | |
| else: | |
| forward_seq = clean_dna(seq).replace("U", "T") | |
| reverse_seq = reverse_complement_dna(forward_seq) | |
| forward = _score_sequence_against_uniques(forward_seq, unique_kmers, k_values, vocab_by_sero, fdr_alpha) | |
| forward["Orientation_used"] = "forward" | |
| reverse = _score_sequence_against_uniques(reverse_seq, unique_kmers, k_values, vocab_by_sero, fdr_alpha) | |
| reverse["Orientation_used"] = "reverse_complement" | |
| best = reverse if _is_better_orientation(reverse, forward) else forward | |
| orientation_used = best["Orientation_used"] | |
| row = { | |
| "Source": src, | |
| "Sequence": header, | |
| "Orientation_used": orientation_used, | |
| "Predicted_serotype": best["Predicted_serotype"], | |
| "Matches_total": best["Matches_total"], | |
| "Confidence_by_present": best["Confidence_by_present"], | |
| "Confidence_by_serotype_vocab": best["Confidence_by_serotype_vocab"], | |
| } | |
| for g in unique_kmers: | |
| row[f"Matches_{g}"] = best["match_counts"][g] | |
| row[f"FisherP_{g}"] = best["fisher_p"][g] | |
| row[f"FDR_{g}"] = best["fdr_map"][g] | |
| results.append(row) | |
| return pd.DataFrame(results) | |
| def plot_counts_by_serotype(simple_df: pd.DataFrame): | |
| fig = plt.figure(figsize=(8,5)) | |
| ax = fig.add_subplot(111) | |
| counts = simple_df["Predicted_serotype"].value_counts() | |
| ax.bar(counts.index.astype(str), counts.values) | |
| ax.set_xlabel("Predicted serotype") | |
| ax.set_ylabel("Number of sequences") | |
| ax.set_title("Predicted serotype counts") | |
| fig.tight_layout() | |
| return fig | |