From fc8d83a91c7b53a8d686205e247c761b29bf1f28 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 12 Aug 2026 16:58:25 +0000 Subject: [PATCH] Fix spectrum, tree, IO, and plotting bugs; add calculate_mutrate Keep RawMutSpec separate from scaled MutSpec, accept 12-component Mut values, and add calculate_mutrate. Fix CodonAnnotation (CDS length, ValueError, logger crash), GenomeStates/GenesStates IO, MutSpecExtractor serial extraction, plot_mutspec closing caller axes, and several pandas 3 compat issues. Add regression tests and bump to 0.0.16. Co-authored-by: kpotoh --- CHANGELOG.md | 29 ++- README.md | 9 +- scripts/calculate_mutspec.py | 6 +- scripts/collect_mutations.py | 28 +-- scripts/collect_mutations_parallel.py | 15 +- src/pymutspec/__init__.py | 4 +- src/pymutspec/annotation/__init__.py | 2 +- src/pymutspec/annotation/mut.py | 200 +++++++++++-------- src/pymutspec/annotation/phylo_tree.py | 34 +++- src/pymutspec/annotation/spectra.py | 165 +++++++++++----- src/pymutspec/annotation/tree.py | 20 +- src/pymutspec/draw/spectra.py | 13 +- src/pymutspec/io/gb.py | 14 +- src/pymutspec/io/states.py | 36 ++-- src/pymutspec/utils/logging.py | 8 +- tests/test_bugfixes.py | 255 +++++++++++++++++++++++++ tests/test_codon_ann.py | 11 +- tests/test_mutspec_calc.py | 45 ++++- 18 files changed, 694 insertions(+), 200 deletions(-) create mode 100644 tests/test_bugfixes.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a34708..5ee52a8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,14 +1,37 @@ # CHANGELOG - +- `_prepare_codontable` now raises `ValueError` for invalid genetic-code arguments +- `collect_exp_muts_proba` no longer requires a `logger` attribute on `CodonAnnotation` +- `GenomeStates` reads the states file instead of the gappy-sites path, and `path_to_gappy_sites` is optional +- `GenesStates` passes `states_fmt` through to alignment readers and uses parameterized SQL +- `GenesStates` rate-length mismatch raises a valid `RuntimeError` (previously passed an illegal `file=` keyword) +- `MutSpecExtractor._derive_mutspec` now processes `Branch` objects (serial extraction was broken after the parallel refactor) +- `plot_mutspec` no longer closes a figure the caller passed in via `ax=` +- `filter_short_exp_seqs` actually drops IQR outliers (previously kept every node) +- `scripts/calculate_mutspec.py` uses keyword arguments for `DataFrame.pivot` (required by pandas 3) +- `collapse_mutspec` accepts `ObsNum`/`ExpNum` as well as `ObsFr`/`ExpFr` +- `jackknife_spectra_sampling` / `calc_edgewise_spectra` no longer mutate caller-provided DataFrames +- `filter_outlier_branches` falls back to `ProbaFull` when `ProbaMut` is absent +- `read_genbank_ref` writes the detected gene qualifier instead of a leftover loop variable +- `basic_logger` no longer attaches duplicate handlers +- Spectrum scaling of an all-zero vector no longer produces NaNs +- `complete_sbs192_columns` fills missing columns in one concat (avoids pandas fragmentation warning) +- Tests for `calculate_mutspec` now actually assert `RawMutSpec` values + +**Full Changelog**: https://github.com/mitoclub/PyMutSpec/compare/0.0.15...0.0.16 ## 0.0.15 (2026-07-16) diff --git a/README.md b/README.md index 49f97af..5b13ae3 100644 --- a/README.md +++ b/README.md @@ -404,7 +404,7 @@ spectra_per_edge = calc_edgewise_spectra( obs, # observed mutations with RefNode/AltNode columns exp, # expected frequencies with Node column nmtypes_cutoff=10, # minimum distinct mutation types per branch - nobs_cuttof=10, # minimum observed mutations per branch + nobs_cutoff=10, # minimum observed mutations per branch scale=True, # normalise each branch spectrum ) # spectra_per_edge: DataFrame indexed by (RefNode, AltNode), 192 columns @@ -428,7 +428,8 @@ cossim = get_cossim(spectra_per_edge.loc[["nodeA"]], spectra_per_edge.loc[["node | `CodonAnnotation.is_fourfold(cdn)` | Test if a codon is fourfold-degenerate | | `CodonAnnotation.translate_codon(cdn)` | Translate a codon to a single-letter amino acid | | `CodonAnnotation.get_mut_type(cdn1, cdn2, pic)` | Return mutation type label, ref AA, alt AA | -| `calculate_mutspec(obs, exp, ...)` | Compute a 12- or 192-component mutational spectrum | +| `calculate_mutspec(obs, exp, ...)` | Compute a 12- or 192-component mutational spectrum (`RawMutSpec` = obs/exp, `MutSpec` = scaled) | +| `calculate_mutrate(obs, exp, ...)` | Compute unscaled mutation rates (`MutRate` = observed / expected) | | `complete_sbs192_columns(df)` | Ensure a DataFrame has all 192 SBS columns (fill missing with 0) | | `mutations_summary(mutations, ...)` | Tabulate mutation type counts from an observed-mutations table | | `rev_comp(sbs)` | Reverse-complement a 192-component SBS string, e.g. `A[C>A]T` → `A[G>T]T` | @@ -509,5 +510,5 @@ twine upload dist/* - [ ] New way of annotation from HGT - [ ] Separate scripts for mutations collection and annotation - [ ] Integrate parallelisation from HGT project -- [ ] Add feature to calculate mutation rate (MutRate) -- [ ] Rename some functions and variables for better readability +- [x] Add feature to calculate mutation rate (MutRate) +- [x] Rename some functions and variables for better readability diff --git a/scripts/calculate_mutspec.py b/scripts/calculate_mutspec.py index 3742087..9e88ce9 100644 --- a/scripts/calculate_mutspec.py +++ b/scripts/calculate_mutspec.py @@ -28,9 +28,9 @@ def dump_expected(exp, path): def filter_short_exp_seqs(exp_freqs: pd.DataFrame): gene_len_proxy = exp_freqs.set_index(['Label', 'Node'])[possible_sbs12]\ - .sum(1).unstack(0).iloc[:, 0] + .sum(axis=1).unstack(0).iloc[:, 0] lower_bound, upper_bound = get_iqr_bounds(gene_len_proxy) - used_nodes = gene_len_proxy.between(lower_bound, upper_bound).index.values + used_nodes = gene_len_proxy[gene_len_proxy.between(lower_bound, upper_bound)].index exp_freqs_flt = exp_freqs[exp_freqs.Node.isin(used_nodes)] return exp_freqs_flt @@ -115,7 +115,7 @@ def main( if branches: raise ValueError("For branch specific spectra expected mutations required for every internal tree node") - exp_mean = exp_raw.pivot("Label", "Mut", "Count") + exp_mean = exp_raw.pivot(index="Label", columns="Mut", values="Count") exp_freqs = None else: raise RuntimeError("Expected another columns in the table {}".format(path_to_exp)) diff --git a/scripts/collect_mutations.py b/scripts/collect_mutations.py index 0241dc6..54e0cde 100644 --- a/scripts/collect_mutations.py +++ b/scripts/collect_mutations.py @@ -172,8 +172,8 @@ def extract_mutspec_from_tree(self): ) add_header["freqs"] = False if self._save_exp_muts and node_expected_sbs is not None: - self.dump_table(node_expected_sbs, self.handle["exp"], add_header["exp"]) - add_header["exp"] = False + if self.dump_table(node_expected_sbs, self.handle["exp"], add_header["exp"]): + add_header["exp"] = False # summarize state frequencies over genome for lbl in self.mut_labels: @@ -212,8 +212,8 @@ def extract_mutspec_from_tree(self): mutspec12["Label"] = lbl mutspec12["Gene"] = gene # Dump gene mutspecs - self.dump_table(mutspec12, self.handle["ms12s"], add_header["ms12g"]) - add_header["ms12g"] = False + if self.dump_table(mutspec12, self.handle["ms12s"], add_header["ms12g"]): + add_header["ms12g"] = False if gene_mut_df.Mut.nunique() >= self.mnum192: if lbl == 'nonsyn': @@ -233,8 +233,8 @@ def extract_mutspec_from_tree(self): mutspec192["Label"] = lbl mutspec192["Gene"] = gene # Dump gene mutspecs - self.dump_table(mutspec192, self.handle["ms192g"], add_header["ms192g"]) - add_header["ms192g"] = False + if self.dump_table(mutspec192, self.handle["ms192g"], add_header["ms192g"]): + add_header["ms192g"] = False visited_nodes.add(ref_node.name) @@ -252,8 +252,8 @@ def extract_mutspec_from_tree(self): logger.warning(f"Observed too many mutations ({mut_num} > {aln_size} * 0.1) for branch ({ref_node.name} - {alt_node.name})") # dump mutations - self.dump_table(genome_mutations_df, self.handle["mut"], add_header["mut"]) - add_header["mut"] = False + if self.dump_table(genome_mutations_df, self.handle["mut"], add_header["mut"]): + add_header["mut"] = False # calculate full genome mutational spectra for all labels if self.derive_spectra: @@ -289,9 +289,10 @@ def extract_mutspec_from_tree(self): mutspec192["Label"] = lbl # Dump genome spectra - self.dump_table(mutspec12, self.handle["ms12"], add_header["ms"]) - self.dump_table(mutspec192, self.handle["ms192"], add_header["ms"]) - add_header["ms"] = False + if self.dump_table(mutspec12, self.handle["ms12"], add_header["ms12"]): + add_header["ms12"] = False + if self.dump_table(mutspec192, self.handle["ms192"], add_header["ms192"]): + add_header["ms192"] = False logger.info(f"Processed {ei} tree edges") logger.info(f"Observed {total_mut_num:.3f} substitutions") @@ -302,9 +303,12 @@ def extract_mutspec_from_tree(self): @staticmethod def dump_table(df: pd.DataFrame, handle, header=False): + if df is None or df.empty: + return False if header: - handle.write("\t".join(df.columns) + "\n") + handle.write("\t".join(map(str, df.columns)) + "\n") handle.write(df.to_csv(sep="\t", index=None, header=None, float_format='%g')) + return True def dump_expected_mutations(self, gene_exp_sbs12, gene_exp_sbs192, node, gene, handle, header=False): # TODO rewrite using self.dump_table diff --git a/scripts/collect_mutations_parallel.py b/scripts/collect_mutations_parallel.py index 30c3f3f..fcc983e 100644 --- a/scripts/collect_mutations_parallel.py +++ b/scripts/collect_mutations_parallel.py @@ -380,7 +380,7 @@ def _derive_mutspec(self): total_mut_num = 0 for edge_data in self.iter_branches(): - edge_mutations = self.process_branch(*edge_data[1:]) + edge_mutations = edge_data.process_branch() mut_num = edge_mutations['ProbaFull'].sum() if self.use_proba and \ 'ProbaFull' in edge_mutations.columns else len(edge_mutations) total_mut_num += mut_num @@ -390,7 +390,7 @@ def _derive_mutspec(self): add_header["mut"] = False self.close_handles() - logger.info(f"Processed {edge_data[1]} tree edges") + logger.info(f"Processed {edge_data.index} tree edges") logger.info(f"Observed {total_mut_num:.3f} substitutions") logger.info("Extraction of mutations from phylogenetic tree completed succesfully") @@ -400,7 +400,11 @@ def _derive_mutspec_parallel(self): with mp.Pool(processes=self.num_processes) as pool: genome_mutations_lst = pool.map(Branch.process_branch, self.iter_branches()) - genome_mutations = pd.concat(genome_mutations_lst) + frames = [df for df in genome_mutations_lst if df is not None and not df.empty] + if frames: + genome_mutations = pd.concat(frames, ignore_index=True) + else: + genome_mutations = pd.DataFrame() self.open_handles(self.outdir) self.dump_table(genome_mutations, self.handle["mut"], True) @@ -450,9 +454,12 @@ def get_edges_data(self): @staticmethod def dump_table(df: pd.DataFrame, handle, header=False): + if df is None or df.empty: + return False if header: - handle.write("\t".join(df.columns) + "\n") + handle.write("\t".join(map(str, df.columns)) + "\n") handle.write(df.to_csv(sep="\t", index=None, header=None, float_format='%g')) + return True def turn_to_MAP(self, states: np.ndarray): if isinstance(states, pd.DataFrame): diff --git a/src/pymutspec/__init__.py b/src/pymutspec/__init__.py index 55581a9..b54fdd9 100644 --- a/src/pymutspec/__init__.py +++ b/src/pymutspec/__init__.py @@ -1,8 +1,8 @@ from .draw import plot_mutspec, plot_mutspec12, plot_mutspec192 from .annotation import ( CodonAnnotation, mutations_summary, - calculate_mutspec, complete_sbs192_columns, + calculate_mutspec, calculate_mutrate, complete_sbs192_columns, rev_comp, transcriptor ) -__version__ = "0.0.15" \ No newline at end of file +__version__ = "0.0.16" \ No newline at end of file diff --git a/src/pymutspec/annotation/__init__.py b/src/pymutspec/annotation/__init__.py index be720b0..bfc72d5 100644 --- a/src/pymutspec/annotation/__init__.py +++ b/src/pymutspec/annotation/__init__.py @@ -2,7 +2,7 @@ from .phylo_tree import TreeNode, Tree from .tree import iter_tree_edges, node_parent, calc_phylocoefs, get_ingroup_root, get_tree_height from .spectra import ( - calculate_mutspec, collapse_mutspec, calc_edgewise_spectra, + calculate_mutspec, calculate_mutrate, collapse_mutspec, calc_edgewise_spectra, complete_sbs192_columns, jackknife_spectra_sampling, collapse_sbs192, get_cossim, get_eucdist, filter_outlier_branches, sample_spectrum, get_iqr_bounds, diff --git a/src/pymutspec/annotation/mut.py b/src/pymutspec/annotation/mut.py index 297ac96..125f972 100644 --- a/src/pymutspec/annotation/mut.py +++ b/src/pymutspec/annotation/mut.py @@ -28,12 +28,55 @@ def __init__(self, gencode: Union[NCBICodonTableDNA, int]): self.possible_syn_contexts = self.__extract_possible_syn_contexts() self.startcodons, self.stopcodons = self.read_start_stop_codons(gencode) + def _cds_bounds(self, cds): + """Return ``(full_length, coding_length)`` for a CDS-like sequence. + + Coding length is truncated to a multiple of 3. A warning is emitted + when the input length is not codon-aligned. + """ + n = len(cds) + n_coding = n - (n % 3) + if n_coding != n: + warn( + f"Sequence length ({n}) is not divisible by 3; " + "the incomplete last codon is ignored for codon-aware labels " + "(syn, ff, pos3, nonsyn)", + UserWarning, + stacklevel=3, + ) + return n, n_coding + + def _iter_cds_sites(self, cds, mask=None): + """Yield ``(pos, pic, nuc, cdn, cxt, codon_complete)`` for each site. + + ``pos`` is 0-based. Sites without a full trinucleotide context + (first and last nucleotide) are skipped. ``cdn`` is ``None`` when + the site falls in an incomplete trailing codon. + """ + n, n_coding = self._cds_bounds(cds) + if mask is not None and len(mask) != n: + raise ValueError(f"Mask must have same length as cds ({len(mask)} != {n})") + for pos in range(1, n - 1): + if mask is not None and not mask[pos]: + continue + pic = pos % 3 + codon_complete = pos < n_coding + nuc = cds[pos] + cxt = cds[pos - 1: pos + 2] + cxt = cxt if isinstance(cxt, str) else "".join(cxt) + if codon_complete: + cdn = cds[pos - pic: pos - pic + 3] + cdn = cdn if isinstance(cdn, str) else "".join(cdn) + else: + cdn = None + yield pos, pic, nuc, cdn, cxt, codon_complete + def is_fourfold(self, cdn: str): """Check if codon is neutral in 3rd position""" return cdn in self._ff_codons def translate_codon(self, cdn: str) -> str: - """Translate codon to animo acid""" + """Translate codon to an amino acid""" if isinstance(cdn, str): return self.codontable.forward_table.get(cdn, "*") else: @@ -75,7 +118,7 @@ def get_syn_codons(self, cdn: str, pic: int) -> Set[str]: possible synonymous codons """ assert 0 <= pic <= 2, "pic must be 0-based and less than 3" - syn_codons = self._syn_codons.get((cdn, pic), dict()) + syn_codons = self._syn_codons.get((cdn, pic), set()) return syn_codons def get_mut_type(self, cdn1: str, cdn2: str, pic: int): @@ -158,13 +201,14 @@ def extract_mutations_simple(self, g1: np.ndarray, g2: np.ndarray): - AltAa """ n, m = len(g1), len(g2) - assert n == m, f"genomes lengths are not equal: {n} != {m}" - assert n % 3 == 0, "genomes length must be divisible by 3 (codon structure)" + if n != m: + raise ValueError(f"genomes lengths are not equal: {n} != {m}") + _, n_coding = self._cds_bounds(g1) nucleotides = set("ACGTacgt") mutations = [] # pass initial codon and last nucleotide without right context - for pos in range(3, n - 1): + for pos in range(3, min(n - 1, n_coding)): pic = pos % 3 # 0-based cdn_start = pos - pic cdn1 = g1[cdn_start: cdn_start + 3] @@ -199,6 +243,11 @@ def extract_mutations_simple(self, g1: np.ndarray, g2: np.ndarray): mutations.append(sbs) mut_df = pd.DataFrame(mutations) + if mut_df.empty: + mut_df = pd.DataFrame(columns=[ + "Mut", "Label", "PosInGene", "PosInCodon", + "RefCodon", "AltCodon", "RefAa", "AltAa", + ]) return mut_df def collect_exp_mut_freqs(self, cds: Union[str, Iterable[str]], mask: Iterable[Union[int, bool]] = None, labels=["all", "syn", "ff"]): @@ -222,35 +271,21 @@ def collect_exp_mut_freqs(self, cds: Union[str, Iterable[str]], mask: Iterable[U sbs192_freqs: Dict[label, Dict[context, count]] for each label collected expected single nucleotide substitutions frequencies with contexts """ - n = len(cds) - if mask is not None and len(mask) != n: - raise ValueError("Mask must have same lenght as cds") - - assert n % 3 == 0, "genomes length must be divisible by 3 (codon structure)" - labels = set(labels) sbs12_freqs = {lbl: defaultdict(int) for lbl in labels} sbs192_freqs = {lbl: defaultdict(int) for lbl in labels} - for pos in range(1, n - 1): - if mask is not None and not mask[pos]: - continue - pic = pos % 3 - nuc = cds[pos] - cdn = cds[pos - pic: pos - pic + 3] - cdn = cdn if isinstance(cdn, str) else "".join(cdn) - cxt = cds[pos - 1: pos + 2] - cxt = cxt if isinstance(cxt, str) else "".join(cxt) + for pos, pic, nuc, cdn, cxt, codon_complete in self._iter_cds_sites(cds, mask): sbs12_pattern = nuc + ">" + "{}" sbs192_pattern = cxt[0] + "[" + nuc + ">{}]" + cxt[-1] - syn_codons = self.get_syn_codons(cdn, pic) + syn_codons = self.get_syn_codons(cdn, pic) if codon_complete else set() - if "syn" in labels: + if codon_complete and "syn" in labels: for alt_cdn in syn_codons: alt_nuc = alt_cdn[pic] sbs12_freqs["syn"][sbs12_pattern.format(alt_nuc)] += 1 sbs192_freqs["syn"][sbs192_pattern.format(alt_nuc)] += 1 - if "nonsyn" in labels: + if codon_complete and "nonsyn" in labels: syn_alt_nucs = [cdn[pic] for cdn in syn_codons] syn_alt_nucs.append(nuc) nonsyn_alt_nucs = set(self.nucl_order).difference(syn_alt_nucs) @@ -267,13 +302,13 @@ def collect_exp_mut_freqs(self, cds: Union[str, Iterable[str]], mask: Iterable[U if "all" in labels: sbs12_freqs["all"][cur_sbs12] += 1 sbs192_freqs["all"][cur_sbs192] += 1 - if "pos3" in labels and pic == 2: + if codon_complete and "pos3" in labels and pic == 2: sbs12_freqs["pos3"][cur_sbs12] += 1 sbs192_freqs["pos3"][cur_sbs192] += 1 - if "ff" in labels and pic == 2 and self.is_fourfold(cdn): + if codon_complete and "ff" in labels and pic == 2 and self.is_fourfold(cdn): sbs12_freqs["ff"][cur_sbs12] += 1 sbs192_freqs["ff"][cur_sbs192] += 1 - if "syn_c" in labels and len(syn_codons) > 0: + if codon_complete and "syn_c" in labels and len(syn_codons) > 0: sbs12_freqs["syn_c"][cur_sbs12] += 1 sbs192_freqs["syn_c"][cur_sbs192] += 1 @@ -296,27 +331,13 @@ def collect_exp_muts(self, cds, mask=None, labels=["syn"]): --------- sbs_table: pd.DataFrame """ - n = len(cds) - if mask is not None and len(mask) != n: - raise ValueError("Mask must have same lenght as cds") - - assert n % 3 == 0, "genomes length must be divisible by 3 (codon structure)" - labels = set(labels) data = [] - for pos in range(1, n - 1): - if mask is not None and not mask[pos]: - continue - pic = pos % 3 - nuc = cds[pos] - cdn = cds[pos - pic: pos - pic + 3] - cdn = cdn if isinstance(cdn, str) else "".join(cdn) - cxt = cds[pos - 1: pos + 2] - cxt = cxt if isinstance(cxt, str) else "".join(cxt) + for pos, pic, nuc, cdn, cxt, codon_complete in self._iter_cds_sites(cds, mask): sbs192_pattern = cxt[0] + "[" + nuc + ">{}]" + cxt[-1] - syn_codons = self.get_syn_codons(cdn, pic) + syn_codons = self.get_syn_codons(cdn, pic) if codon_complete else set() - if "syn" in labels: + if codon_complete and "syn" in labels: for alt_cdn in syn_codons: alt_nuc = alt_cdn[pic] data.append({ @@ -324,7 +345,7 @@ def collect_exp_muts(self, cds, mask=None, labels=["syn"]): "Mut": sbs192_pattern.format(alt_nuc), "Cdn": cdn, "Label": "syn", }) - if "nonsyn" in labels: + if codon_complete and "nonsyn" in labels: syn_alt_nucs = [cdn[pic] for cdn in syn_codons] syn_alt_nucs.append(nuc) nonsyn_alt_nucs = set(self.nucl_order).difference(syn_alt_nucs) @@ -346,19 +367,19 @@ def collect_exp_muts(self, cds, mask=None, labels=["syn"]): "Mut": cur_sbs192, "Cdn": cdn, "Label": "all", }) - if "pos3" in labels and pic == 2: + if codon_complete and "pos3" in labels and pic == 2: data.append({ "Pos": pos + 1, "Pic": pic + 1, "Mut": cur_sbs192, "Cdn": cdn, "Label": "pos3", }) - if "ff" in labels and pic == 2 and self.is_fourfold(cdn): + if codon_complete and "ff" in labels and pic == 2 and self.is_fourfold(cdn): data.append({ "Pos": pos + 1, "Pic": pic + 1, "Mut": cur_sbs192, "Cdn": cdn, "Label": "syn4f", }) - if "syn_c" in labels and len(syn_codons) > 0: + if codon_complete and "syn_c" in labels and len(syn_codons) > 0: data.append({ "Pos": pos + 1, "Pic": pic + 1, "Mut": cur_sbs192, @@ -386,12 +407,13 @@ def extract_mutations_proba(self, g1: np.ndarray, g2: np.ndarray, phylocoef: flo - mut - dataframe of mutations """ n, m = len(g1), len(g2) - assert n == m, f"genomes lengths are not equal: {n} != {m}" - assert n % 3 == 0, "genomes length must be divisible by 3 (codon structure)" + if n != m: + raise ValueError(f"genomes lengths are not equal: {n} != {m}") + _, n_coding = self._cds_bounds(g1) mutations = [] # pass initial codon and last nucleotide without right context - for pos in range(3, n - 1): + for pos in range(3, min(n - 1, n_coding)): pic = pos % 3 # 0-based for cdn1, mut_cxt1, proba1 in self.sample_context(pos, pic, g1, mut_proba_cutoff / phylocoef): cdn1_str = "".join(cdn1) @@ -428,6 +450,12 @@ def extract_mutations_proba(self, g1: np.ndarray, g2: np.ndarray, phylocoef: flo mutations.append(sbs) mut_df = pd.DataFrame(mutations) + if mut_df.empty: + mut_df = pd.DataFrame(columns=[ + "Mut", "Label", "PosInGene", "PosInCodon", + "RefCodon", "AltCodon", "RefAa", "AltAa", + "ProbaRef", "ProbaMut", "ProbaFull", + ]) return mut_df def collect_exp_mut_freqs_proba( @@ -435,14 +463,13 @@ def collect_exp_mut_freqs_proba( mask: Iterable[Union[int, bool]] = None, labels = ["all", "syn", "ff"], mut_proba_cutoff=0.05, ): - n = len(cds) + n, n_coding = self._cds_bounds(cds) if mask is not None and len(mask) != n: - msg = f"Mask (len = {len(mask)}) must have same lenght as cds (len = {n})" - self.logger.error(msg) - raise ValueError(msg) - - assert n % 3 == 0, "genomes length must be divisible by 3 (codon structure)" - assert 0 < phylocoef <= 1, "Evol coefficient must be between 0 and 1" + raise ValueError( + f"Mask (len = {len(mask)}) must have same length as cds (len = {n})" + ) + if phylocoef <= 0 or phylocoef > 1: + raise ValueError(f"Evol coefficient must be between 0 and 1, but got {phylocoef}") labels = set(labels) sbs12_freqs = {lbl: defaultdict(int) for lbl in labels} @@ -451,6 +478,8 @@ def collect_exp_mut_freqs_proba( for pos in range(1, n - 1): if mask is not None and not mask[pos]: continue + if pos >= n_coding: + continue # sample_context requires a complete codon pic = pos % 3 # 0-based for cdn_tuple, cxt, p in self.sample_context(pos, pic, cds, mut_proba_cutoff / phylocoef): # we don't use low-probability mutations by unite cutoff `mut_proba_cutoff / phylocoef` @@ -500,26 +529,21 @@ def collect_exp_muts_proba( mask: Iterable[Union[int, bool]] = None, labels = ["all", "syn", "ff"], mut_proba_cutoff=0.05, ): - n = len(cds) + n, n_coding = self._cds_bounds(cds) if mask is not None and len(mask) != n: - msg = f"Mask (len = {len(mask)}) must have same lenght as cds (len = {n})" - self.logger.error(msg) - raise ValueError(msg) - - if n % 3 != 0: - self.logger.warning(f"genomes length ({n}) is not divisible by 3 (codon structure). Last codon will be skipped in syn, syn4f and pos3 modes") - n = n - (n % 3) - + raise ValueError( + f"Mask (len = {len(mask)}) must have same length as cds (len = {n})" + ) if phylocoef <= 0 or phylocoef > 1: - msg = f"Evol coefficient must be between 0 and 1, but got {phylocoef}" - self.logger.error(msg) - raise ValueError(msg) + raise ValueError(f"Evol coefficient must be between 0 and 1, but got {phylocoef}") labels = set(labels) data = [] for pos in range(1, n - 1): if mask is not None and not mask[pos]: continue + if pos >= n_coding: + continue pic = pos % 3 # 0-based for cdn_tuple, cxt, p in self.sample_context(pos, pic, cds, mut_proba_cutoff / phylocoef): # we don't use low-probability mutations by unite cutoff `mut_proba_cutoff / phylocoef` @@ -696,7 +720,7 @@ def _prepare_codontable(codontable: Union[NCBICodonTableDNA, int]): elif isinstance(codontable, int): codontable = CodonTable.unambiguous_dna_by_id[codontable] else: - ValueError("passed codontable is not appropriate") + raise ValueError("passed codontable is not appropriate") return codontable def _codon_iterator(self, codon_states: np.ndarray, cutoff=0.01): @@ -844,6 +868,10 @@ def extract_mutations_proba(self, g1: pd.DataFrame, g2: pd.DataFrame, g1, g2, site, g1_most_prob, g2_most_prob, mut_proba_cutoff, phylocoef) mutations.extend(site_mutations) mut_df = pd.DataFrame(mutations) + if mut_df.empty: + mut_df = pd.DataFrame(columns=[ + "Mut", "Site", "ProbaRef", "ProbaMut", "ProbaFull", + ]) return mut_df def process_site(self, g1: pd.DataFrame, g2: pd.DataFrame, site: int, @@ -1010,21 +1038,22 @@ def extract_mutspec_from_tree(self): def _derive_mutspec(self): self.logger.info("Start mutation extraction from tree") self.open_handles(self.outdir) - add_header = defaultdict(lambda: True) + add_header = True total_mut_num = 0 + n_edges = 0 - for edge_data in self.iter_branches(): - edge_mutations = self.process_branch(*edge_data[1:]) + for branch in self.iter_branches(): + n_edges += 1 + edge_mutations = branch.process_branch() mut_num = edge_mutations['ProbaFull'].sum() if self.use_proba and \ 'ProbaFull' in edge_mutations.columns else len(edge_mutations) total_mut_num += mut_num - # dump current edge mutations - self.dump_table(edge_mutations, self.handle["mut"], add_header["mut"]) - add_header["mut"] = False + if self.dump_table(edge_mutations, self.handle["mut"], add_header): + add_header = False self.close_handles() - self.logger.info(f"Processed {edge_data[1]} tree edges") + self.logger.info(f"Processed {n_edges} tree edges") self.logger.info(f"Observed {total_mut_num:.3f} substitutions") self.logger.info("Extraction of mutations from phylogenetic tree completed succesfully") @@ -1034,7 +1063,13 @@ def _derive_mutspec_parallel(self): with mp.Pool(processes=self.num_processes) as pool: genome_mutations_lst = pool.map(Branch.process_branch, self.iter_branches()) - genome_mutations = pd.concat(genome_mutations_lst) + frames = [df for df in genome_mutations_lst if df is not None and not df.empty] + if frames: + genome_mutations = pd.concat(frames, ignore_index=True) + else: + genome_mutations = pd.DataFrame(columns=[ + "Mut", "Site", "ProbaRef", "ProbaMut", "ProbaFull", "RefNode", "AltNode", + ]) self.open_handles(self.outdir) self.dump_table(genome_mutations, self.handle["mut"], True) @@ -1079,9 +1114,12 @@ def get_edges_data(self): @staticmethod def dump_table(df: pd.DataFrame, handle, header=False): + if df is None or df.empty: + return False if header: - handle.write("\t".join(df.columns) + "\n") + handle.write("\t".join(map(str, df.columns)) + "\n") handle.write(df.to_csv(sep="\t", index=None, header=None, float_format='%g')) + return True def turn_to_MAP(self, states: np.ndarray): if isinstance(states, pd.DataFrame): @@ -1139,7 +1177,9 @@ def mutations_summary(mutations: pd.DataFrame, gene_col=None, proba_col=None, ge ].groupby(grp)[proba_col].sum().reset_index() mutations_descr["Label"] = mutations_descr.Label.map(label_mapper) - pivot_mutations = mutations_descr.pivot_table(proba_col, gene_col, "Label", fill_value=0) + pivot_mutations = mutations_descr.pivot_table( + values=proba_col, index=gene_col, columns="Label", fill_value=0 + ) pivot_mutations.columns = [x[1:] for x in pivot_mutations.columns] if gene_name_mapper is not None: diff --git a/src/pymutspec/annotation/phylo_tree.py b/src/pymutspec/annotation/phylo_tree.py index 8cfbae9..7664591 100644 --- a/src/pymutspec/annotation/phylo_tree.py +++ b/src/pymutspec/annotation/phylo_tree.py @@ -2,6 +2,7 @@ Custom phylogenetic tree classes replacing ete3 dependency. Uses BioPython for newick format parsing. """ +import os from io import StringIO from Bio import Phylo as _BioPhylo @@ -26,6 +27,21 @@ def __init__(self, name="", dist=0.0): def is_leaf(self): return len(self.children) == 0 + @property + def parent(self): + """Immediate ancestor, or ``None`` for the root.""" + return self._parent + + @property + def up(self): + """ete3-compatible alias of :attr:`parent`.""" + return self._parent + + def iter_edges(self): + """Yield ``(parent, child)`` for every directed edge under this node.""" + for node in self.iter_descendants(): + yield node.parent, node + def traverse(self): """Yield all nodes (self first, then descendants).""" yield self @@ -189,12 +205,13 @@ def _bio_clade_to_node(clade): class Tree(TreeNode): """ - Load a phylogenetic tree from a newick file. + Load a phylogenetic tree from a newick file or newick string. Parameters ---------- newick_path : str - Path to a newick-format tree file. + Path to a newick-format tree file, or a newick string + (must start with ``(`` or end with ``;``). format : int Newick format hint (kept for API compatibility with ete3; this parameter is currently ignored – BioPython's newick parser handles @@ -202,8 +219,17 @@ class Tree(TreeNode): """ def __init__(self, newick_path, format=1): # noqa: A002 - with open(newick_path) as fh: - tree_str = fh.read().strip() + if not isinstance(newick_path, str): + raise TypeError("newick_path must be a file path or newick string") + + stripped = newick_path.strip() + if os.path.isfile(newick_path): + with open(newick_path) as fh: + tree_str = fh.read().strip() + elif stripped.startswith("(") or stripped.endswith(";"): + tree_str = stripped + else: + raise FileNotFoundError(f"Tree file not found: {newick_path}") bio_tree = _BioPhylo.read(StringIO(tree_str), "newick") root = _bio_clade_to_node(bio_tree.root) diff --git a/src/pymutspec/annotation/spectra.py b/src/pymutspec/annotation/spectra.py index 853dcc0..4e8301e 100644 --- a/src/pymutspec/annotation/spectra.py +++ b/src/pymutspec/annotation/spectra.py @@ -11,6 +11,18 @@ from .auxiliary import rev_comp +_SBS12_RE = r"[ACGT]>[ACGT]" + + +def _as_sbs12(mut_series: pd.Series) -> pd.Series: + """Accept both 192-component (``A[C>T]G``) and 12-component (``C>T``) Mut values.""" + mut_str = mut_series.astype(str) + sbs12 = mut_str.str.slice(2, 5) + already_sbs12 = mut_str.str.fullmatch(_SBS12_RE) + sbs12 = sbs12.where(~already_sbs12, mut_str) + return sbs12 + + def calculate_mutspec( obs_muts: pd.DataFrame, exp_muts: Dict[str, float], @@ -30,16 +42,12 @@ def calculate_mutspec( --------- obs_muts: pd.DataFrame table containing mutations with annotation; table must contain 2 columns: - - Mut: str; Pattern: ``[ACGT]\\[[ACGT]>[ACGT]\\][ACGT]`` + - Mut: str; Pattern: ``[ACGT]\\[[ACGT]>[ACGT]\\][ACGT]`` or ``[ACGT]>[ACGT]`` - ProbaFull (optional, only for use_proba=True) - probability of mutation exp_muts: dict[str, float] dictionary that contains expected mutations frequencies of reference genome if use_context=False, else trinucleotide freqs - label: str - kind of needed mutspec, coulb be one of ['all', 'syn', 'ff'] - gencode: int - Number of genetic code to use in expected mutations collection, required if exp_muts_or_genome is genome use_context: bool To use trinucleotide context or not, in other words calculate 192 component mutspec use_proba: bool @@ -60,7 +68,8 @@ def calculate_mutspec( Return ------- mutspec: pd.DataFrame - table, containing extended mutspec values including observed mutations numbers. + table, containing observed/expected counts and both unscaled (``RawMutSpec``) + and optionally scaled (``MutSpec``) rates. If use_context=True len(mutspec) = 192, else len(mutspec) = 12 """ _cols = ["Mut", "ProbaFull"] if use_proba else ["Mut"] @@ -75,42 +84,87 @@ def calculate_mutspec( col_mut = "Mut" full_sbs = possible_sbs192_set else: - # TODO add support of sbs12 in Mut column - mut["Sbs12"] = mut["Mut"].str.slice(2, 5) + mut["Sbs12"] = _as_sbs12(mut["Mut"]) col_mut = "Sbs12" full_sbs = possible_sbs12_set if not use_proba: mut["ProbaFull"] = 1 - mutspec = mut.groupby(col_mut)["ProbaFull"].sum().reset_index() + mutspec = mut.groupby(col_mut, sort=False)["ProbaFull"].sum().reset_index() mutspec.columns = ["Mut", "ObsNum"] if fill_unobserved: - # fill unobserved mutations by zeros - mutspec_appendix = [] unobserved_sbs = full_sbs.difference(mutspec["Mut"].values) - for usbs in unobserved_sbs: - mutspec_appendix.append({"Mut": usbs, "ObsNum": 0}) - mutspec = pd.concat([mutspec, pd.DataFrame(mutspec_appendix)], ignore_index=True) + if unobserved_sbs: + mutspec = pd.concat( + [mutspec, pd.DataFrame({"Mut": list(unobserved_sbs), "ObsNum": 0})], + ignore_index=True, + ) mutspec["ExpNum"] = mutspec["Mut"].map(exp_muts) - mutspec["MutSpec"] = (mutspec["ObsNum"] / mutspec["ExpNum"]).fillna(0) + raw = mutspec["ObsNum"] / mutspec["ExpNum"] if verbose: - msg = mutspec[mutspec["MutSpec"] == np.inf] - if len(msg) > 0: - print(f"WARNING! Following substitutions are unexpected but observed:\n{msg}", file=stderr) - - mutspec.loc[mutspec["MutSpec"] == np.inf, "MutSpec"] = 0 - + unexpected = mutspec[(mutspec["ObsNum"] > 0) & (mutspec["ExpNum"].fillna(0) <= 0)] + if len(unexpected) > 0: + print(f"WARNING! Following substitutions are unexpected but observed:\n{unexpected}", file=stderr) + + raw = raw.replace([np.inf, -np.inf], np.nan).fillna(0) + mutspec["RawMutSpec"] = raw + mutspec["MutSpec"] = raw + if drop_underrepresented: - mutspec.loc[(mutspec["ObsNum"] < nobs_min) | (mutspec.ExpNum < nexp_min), "MutSpec"] = 0. + under = (mutspec["ObsNum"] < nobs_min) | (mutspec["ExpNum"] < nexp_min) + mutspec.loc[under.fillna(False), "MutSpec"] = 0. if scale: - mutspec["MutSpec"] = mutspec["MutSpec"] / mutspec["MutSpec"].sum() + total = mutspec["MutSpec"].sum() + if total > 0: + mutspec["MutSpec"] = mutspec["MutSpec"] / total + else: + mutspec["MutSpec"] = 0. return mutspec +def calculate_mutrate( + obs_muts: pd.DataFrame, + exp_muts: Dict[str, float], + use_context: bool = False, + use_proba: bool = False, + fill_unobserved=True, + verbose=False, +): + """ + Calculate mutation rates (observed / expected) without scaling to a spectrum. + + This is the un-normalised rate vector that ``calculate_mutspec`` stores in + ``RawMutSpec``. The returned table uses the column name ``MutRate``. + + Arguments + --------- + obs_muts, exp_muts, use_context, use_proba, fill_unobserved, verbose + Same meaning as in :func:`calculate_mutspec`. + + Return + ------ + mutrate: pd.DataFrame + Table with columns ``Mut``, ``ObsNum``, ``ExpNum``, ``RawMutSpec``, + ``MutSpec`` (unscaled) and ``MutRate`` (alias of ``RawMutSpec``). + """ + mutrate = calculate_mutspec( + obs_muts, + exp_muts, + use_context=use_context, + use_proba=use_proba, + scale=False, + fill_unobserved=fill_unobserved, + drop_underrepresented=False, + verbose=verbose, + ) + mutrate["MutRate"] = mutrate["RawMutSpec"] + return mutrate + + def sample_spectrum(obs_df: pd.DataFrame, exp_freqs, use_proba=True, use_context=False, frac=0.5, nreplics=100): @@ -155,9 +209,11 @@ def filter_outlier_branches(obs_df: pd.DataFrame, use_proba=True): --------- obs_df: pd.DataFrame Observed-mutations table containing at least the columns - ``'AltNode'``, ``'Mut'``, and optionally ``'ProbaMut'``. + ``'AltNode'``, ``'Mut'``, and optionally ``'ProbaMut'`` or + ``'ProbaFull'``. use_proba: bool - If ``True`` sum ``'ProbaMut'`` per branch; otherwise count rows. + If ``True`` sum ``'ProbaMut'`` (falling back to ``'ProbaFull'``) + per branch; otherwise count rows. Return ------ @@ -165,7 +221,15 @@ def filter_outlier_branches(obs_df: pd.DataFrame, use_proba=True): Filtered mutations table with outlier branches removed. """ if use_proba: - edge_nobs = obs_df.groupby('AltNode')['ProbaMut'].sum() + if "ProbaMut" in obs_df.columns: + proba_col = "ProbaMut" + elif "ProbaFull" in obs_df.columns: + proba_col = "ProbaFull" + else: + raise ValueError( + "use_proba=True requires a 'ProbaMut' or 'ProbaFull' column" + ) + edge_nobs = obs_df.groupby('AltNode')[proba_col].sum() else: edge_nobs = obs_df.groupby('AltNode')['Mut'].count() @@ -188,8 +252,9 @@ def collapse_mutspec(ms192: pd.DataFrame): Arguments --------- ms192: pd.DataFrame - 192-component spectrum table. Must contain columns ``'Mut'``, - ``'ObsFr'``, and ``'ExpFr'``, and must have exactly 192 rows. + 192-component spectrum table. Must contain column ``'Mut'`` and + either ``'ObsFr'``/``'ExpFr'`` or ``'ObsNum'``/``'ExpNum'``, and + must have exactly 192 rows. Return ------ @@ -204,11 +269,16 @@ def collapse_mutspec(ms192: pd.DataFrame): columns. """ assert ms192.shape[0] == 192, f"Expected 192 rows, got {ms192.shape[0]}" + ms192 = ms192.copy() + if "ObsFr" not in ms192.columns and "ObsNum" in ms192.columns: + ms192["ObsFr"] = ms192["ObsNum"] + if "ExpFr" not in ms192.columns and "ExpNum" in ms192.columns: + ms192["ExpFr"] = ms192["ExpNum"] for c in ["Mut", "ObsFr", "ExpFr"]: assert c in ms192.columns, f"Required column '{c}' not found in ms192" - ms1 = ms192[ms192["Mut"].str.get(2).isin(list("CT"))] - ms2 = ms192[ms192["Mut"].str.get(2).isin(list("AG"))] + ms1 = ms192.loc[ms192["Mut"].str.get(2).isin(list("CT"))].copy() + ms2 = ms192.loc[ms192["Mut"].str.get(2).isin(list("AG"))].copy() ms2["Mut"] = ms2["Mut"].apply(rev_comp) ms96 = pd.concat([ms1, ms2]).groupby("Mut")[["ObsFr", "ExpFr"]].sum() @@ -237,11 +307,13 @@ def complete_sbs192_columns(df: pd.DataFrame): DataFrame with exactly 192 columns in canonical order. """ df = df.copy() - if len(df.columns) != 192: - for sbs192 in possible_sbs192_set.difference(df.columns.values): - df[sbs192] = 0. - df = df[possible_sbs192] - return df + missing = [sbs for sbs in possible_sbs192 if sbs not in df.columns] + if missing: + df = pd.concat( + [df, pd.DataFrame(0.0, index=df.index, columns=missing)], + axis=1, + ) + return df[possible_sbs192] def collapse_sbs192(df: pd.DataFrame, to=12): @@ -326,8 +398,8 @@ def jackknife_spectra_sampling(obs: pd.DataFrame, exp: pd.DataFrame, frac=0.5, n assert obs.index.names == ["RefNode", "AltNode"] assert exp.index.names == ["Node"] altnodes = obs.index.get_level_values(1).values - obs_edges = obs - freqs_nodes = exp + obs_edges = obs.copy() + freqs_nodes = exp.copy() obs_edges.index = obs_edges.index.reorder_levels(order=["AltNode", "RefNode"]) freqs_nodes.index.name = "RefNode" else: @@ -342,7 +414,7 @@ def jackknife_spectra_sampling(obs: pd.DataFrame, exp: pd.DataFrame, frac=0.5, n spectra = [] for _ in range(n): altnodes_sample = np.random.choice(altnodes, edges_sample_size, False) - obs_sample = obs_edges.loc[altnodes_sample].reset_index(0, drop=True) + obs_sample = obs_edges.loc[altnodes_sample].reset_index(level=0, drop=True) exp_sample = freqs_nodes.loc[obs_sample.index] obs_sample_cnt = obs_sample.sum() @@ -358,9 +430,10 @@ def jackknife_spectra_sampling(obs: pd.DataFrame, exp: pd.DataFrame, frac=0.5, n def calc_edgewise_spectra( obs: pd.DataFrame, exp: pd.DataFrame, - nmtypes_cutoff=10, nobs_cuttof=10, + nmtypes_cutoff=10, nobs_cutoff=None, collapse_to_12=False, scale=True, - both_12_and_192=False + both_12_and_192=False, + nobs_cuttof=10, ): """ Calculate per-branch (edge-wise) mutational spectra. @@ -383,9 +456,11 @@ def calc_edgewise_spectra( nmtypes_cutoff: int Minimum number of distinct mutation types a branch must have to be retained (only applied when ``collapse_to_12=False``). - nobs_cuttof: int + nobs_cutoff: int Minimum total observed mutations a branch must have to be retained (only applied when ``collapse_to_12=False``). + nobs_cuttof: int + Deprecated alias of ``nobs_cutoff`` kept for backward compatibility. collapse_to_12: bool If ``True`` collapse the 192-component spectra to 12 components before returning. @@ -401,13 +476,15 @@ def calc_edgewise_spectra( Branch-wise spectrum DataFrame (or tuple of two DataFrames when ``both_12_and_192=True``). """ + if nobs_cutoff is None: + nobs_cutoff = nobs_cuttof if len(obs.columns) == 192 and \ (obs.columns == possible_sbs192).all() and \ (exp.columns == possible_sbs192).all(): assert obs.index.names == ["RefNode", "AltNode"] assert exp.index.names == ["Node"] - obs_edges = obs - freqs_nodes = exp + obs_edges = obs.copy() + freqs_nodes = exp.copy() freqs_nodes.index.name = "RefNode" else: obs_edges = obs.groupby(["RefNode", "AltNode", "Mut"]).ProbaFull.sum().unstack() @@ -418,7 +495,7 @@ def calc_edgewise_spectra( if not collapse_to_12: obs_edges = obs_edges[((obs_edges > 0).sum(axis=1) >= nmtypes_cutoff) & \ - (obs_edges.sum(axis=1) >= nobs_cuttof)] + (obs_edges.sum(axis=1) >= nobs_cutoff)] edges_df = obs_edges.index.to_frame(False) diff --git a/src/pymutspec/annotation/tree.py b/src/pymutspec/annotation/tree.py index 382ce8a..c066ec9 100644 --- a/src/pymutspec/annotation/tree.py +++ b/src/pymutspec/annotation/tree.py @@ -24,7 +24,7 @@ def node_parent(node): """ try: return next(node.iter_ancestors()) - except BaseException: + except StopIteration: return None @@ -102,9 +102,10 @@ def get_tree_height(tree, mode='geom_mean'): distances_to_leaves.append(d) if mode == 'mean': - md = np.mean(distances_to_leaves) + md = np.mean(distances_to_leaves) if distances_to_leaves else 0.0 elif mode == 'geom_mean': - md = geometric_mean(distances_to_leaves) + positive = [d for d in distances_to_leaves if d > 0] + md = geometric_mean(positive) if positive else 0.0 else: raise TypeError("mode must be 'mean', 'geom_mean' or 'max'") @@ -116,8 +117,9 @@ def get_ingroup_root(tree): """ Return the ingroup root of a binary rooted tree that contains an outgroup. - The function assumes the tree root has exactly two children, one of which - is a leaf (the outgroup). If no leaf child is found the tree root itself + The function assumes the tree root has exactly two children. If exactly + one child is a leaf, that leaf is treated as the outgroup and the other + child is returned as the ingroup root. Otherwise the tree root itself is returned. Arguments @@ -137,16 +139,16 @@ def get_ingroup_root(tree): """ assert len(tree.children) == 2, 'Tree must be binary' found_outgroup = False + ingrp = tree for node in tree.children: if node.is_leaf(): found_outgroup = True else: ingrp = node - if found_outgroup: + if found_outgroup and ingrp is not tree: return ingrp - else: - return tree + return tree def calc_phylocoefs(tree): @@ -171,6 +173,8 @@ def calc_phylocoefs(tree): """ ingroup = get_ingroup_root(tree) tree_height = get_tree_height(ingroup, 'geom_mean') + if tree_height <= 0: + return {node.name: 1.0 for node in ingroup.traverse()} root_phylocoef = 1 - min(0.999, ingroup.get_closest_leaf()[1] / tree_height) phylocoefs = {ingroup.name: root_phylocoef} for node in ingroup.iter_descendants(): diff --git a/src/pymutspec/draw/spectra.py b/src/pymutspec/draw/spectra.py index 9758b4d..7201b4a 100644 --- a/src/pymutspec/draw/spectra.py +++ b/src/pymutspec/draw/spectra.py @@ -231,9 +231,13 @@ def plot_mutspec( palette = color_mapping12 tick_rotation = 0 + created_fig = False if ax is None: fig = plt.figure(figsize=figsize) ax = fig.gca() + created_fig = True + else: + fig = ax.figure if style == "bar": _cols = set(ms.columns) @@ -286,14 +290,15 @@ def plot_mutspec( ax.set_xticks(range(len(order_styled))) ax.set_xticklabels(order_styled, rotation=tick_rotation, fontsize=ticksize, fontname=fontname) else: - plt.xticks(fontsize=ticksize, fontname=fontname) + ax.set_xticks(range(len(order))) + ax.set_xticklabels(order, fontsize=ticksize, fontname=fontname) if savepath is not None: - plt.savefig(savepath, dpi=dpi, bbox_inches="tight") + fig.savefig(savepath, dpi=dpi, bbox_inches="tight") if show: plt.show() - else: - plt.close() + elif created_fig: + plt.close(fig) return ax diff --git a/src/pymutspec/io/gb.py b/src/pymutspec/io/gb.py index d576831..8339c17 100644 --- a/src/pymutspec/io/gb.py +++ b/src/pymutspec/io/gb.py @@ -41,11 +41,19 @@ def read_genbank_ref(gb: Union[str, SeqRecord]): if gene_qualifier is None: raise RuntimeError(f"Cannot find any expected qualifier of feature: {ftr}; with following qualifiers: {ftr.qualifiers}") + if df is None: + raise ValueError("GenBank record must contain a source feature before other features") + for pos in list(ftr.location): df.at[pos, "Type"] = ftr.type - df.at[pos, "Strand"] = ftr.strand - if ftr.type in ftypes: - df.at[pos, qualifier] = ftr.qualifiers[qualifier][0] + strand = getattr(ftr, "strand", None) + if strand is None: + strand = ftr.location.strand + df.at[pos, "Strand"] = strand + if ftr.type in ftypes and gene_qualifier is not None: + values = ftr.qualifiers.get(gene_qualifier) + if values: + df.at[pos, gene_qualifier] = values[0] # add codon features df["PosInGene"] = -1 diff --git a/src/pymutspec/io/states.py b/src/pymutspec/io/states.py index 66da220..276d893 100644 --- a/src/pymutspec/io/states.py +++ b/src/pymutspec/io/states.py @@ -14,14 +14,15 @@ class GenomeStates: def __init__(self, path_to_states, path_to_gappy_sites=None, format='tsv', logger=None): - for path in list([path_to_states, path_to_gappy_sites]): - if not os.path.exists(path): - raise ValueError(f"Path to states doesn't exist: '{path}'") + if not os.path.exists(path_to_states): + raise ValueError(f"Path to states doesn't exist: '{path_to_states}'") + if path_to_gappy_sites is not None and not os.path.exists(path_to_gappy_sites): + raise ValueError(f"Path to gappy sites doesn't exist: '{path_to_gappy_sites}'") self.logger = logger or basic_logger() self.logger.info(f'Reading states "{path_to_states}"') - states = self.read_states(path, format) + states = self.read_states(path_to_states, format) nodes_arr = states['Node'].unique() self.node2id = dict(zip(nodes_arr, range(len(nodes_arr)))) @@ -51,20 +52,20 @@ def __contains__(self, item: str): def __getitem__(self, item: str): return self.get_genome(item) - def read_states(self, path, fmt='iqtree'): + def read_states(self, path, fmt='tsv'): fpt = np.float32 dtype = { "p_A": fpt, "p_C": fpt, "p_G": fpt, "p_T": fpt, "Site": np.int32, "Node": str, } - if fmt == 'iqtree': + if fmt in ('iqtree', 'tsv'): states = pd.read_csv(path, sep='\t', comment='#', dtype=dtype) - if fmt == 'csv': + elif fmt == 'csv': states = pd.read_csv(path, dtype=dtype) - elif fmt == 'tsv': - states = pd.read_csv(path, sep='\t', dtype=dtype) elif fmt == 'parquet': states = pd.read_parquet(path) + else: + raise ValueError(f"Unsupported states format: {fmt!r}") return states def read_nongappy_sites(self, path): @@ -170,7 +171,7 @@ def __init__( if states_fmt == "table": self._prepare_node2genome(path_states) else: - states = self.read_alignment(path_states) + states = self.read_alignment(path_states, fmt=states_fmt) self._prepare_node2genome(states) elif mode == "db": if states_fmt in ["fasta", "phylip"]: @@ -191,7 +192,9 @@ def __init__( if self.category: for part in genes_sizes: if genes_sizes[part] != len(self.category[part]): - raise RuntimeError("Wrong rates: number of positions in alignment are not equal to number of rates", file=sys.stderr) + raise RuntimeError( + "Wrong rates: number of positions in alignment are not equal to number of rates" + ) self.mask = self.get_mask(self.category, cat_cutoff) def get_genome(self, node: str): @@ -202,17 +205,17 @@ def get_genome(self, node: str): genome_raw = defaultdict(list) cur = self.con.cursor() if self.use_proba: - query = f"""SELECT Part, Site, p_A, p_C, p_G, p_T FROM states WHERE Node='{node}'""" + query = """SELECT Part, Site, p_A, p_C, p_G, p_T FROM states WHERE Node=?""" dtype = [ ("Site", np.int32), ("p_A", np.float32), ("p_C", np.float32), ("p_G", np.float32), ("p_T", np.float32), ] else: - query = f"""SELECT Part, Site, State FROM states WHERE Node='{node}'""" + query = """SELECT Part, Site, State FROM states WHERE Node=?""" dtype = [("Site", np.int32), ("State", np.object_)] - for row in cur.execute(query): + for row in cur.execute(query, (node,)): part = str(row[0]) state = row[1:] genome_raw[part].append(state) @@ -270,8 +273,7 @@ def _prepare_db(self, path_states, rewrite=False): for line in handle: row = line.strip().split() - query = "INSERT INTO states VALUES ('{}',{},{},'{}',{},{},{},{})".format(*row) - cur.execute(query) + cur.execute("INSERT INTO states VALUES (?,?,?,?,?,?,?,?)", row) con.commit() handle.close() @@ -288,7 +290,7 @@ def _get_nodes_from_db(self): def states2dct(self, states: pd.DataFrame, out=None): # all genes (genomes) must have same length - aln_sizes = states.groupby("Node").apply(len) + aln_sizes = states.groupby("Node").size() assert aln_sizes.nunique() == 1, "uncomplete state table: some site states absent in some node genomes" node2genome = defaultdict(dict) if out is None else out diff --git a/src/pymutspec/utils/logging.py b/src/pymutspec/utils/logging.py index 87a19b3..abe79d8 100644 --- a/src/pymutspec/utils/logging.py +++ b/src/pymutspec/utils/logging.py @@ -21,11 +21,13 @@ def load_logger(path=None, stream_level: str = None, filename=None): return logger -def basic_logger(): - logger = logging.getLogger(__name__) +def basic_logger(name="pymutspec"): + logger = logging.getLogger(name) + if logger.handlers: + return logger logger.setLevel(logging.DEBUG) + logger.propagate = False - # create console handler and set level to debug ch = logging.StreamHandler() ch.setLevel(logging.DEBUG) formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s', '%d-%m-%y %H:%M:%S') diff --git a/tests/test_bugfixes.py b/tests/test_bugfixes.py new file mode 100644 index 0000000..1429f57 --- /dev/null +++ b/tests/test_bugfixes.py @@ -0,0 +1,255 @@ +"""Regression tests for bugs fixed in 0.0.16.""" +import os +import tempfile +import warnings + +import numpy as np +import pandas as pd +import pytest +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +from pymutspec.annotation import ( + CodonAnnotation, + calculate_mutspec, + collapse_mutspec, + complete_sbs192_columns, + jackknife_spectra_sampling, + calc_edgewise_spectra, + get_tree_height, + calc_phylocoefs, +) +from pymutspec.annotation.phylo_tree import Tree, TreeNode +from pymutspec.constants import possible_sbs192 +from pymutspec.draw import plot_mutspec12 +from pymutspec.io.states import GenomeStates +from pymutspec.io.gb import read_genbank_ref +from pymutspec.utils.logging import basic_logger + + +def test_prepare_codontable_invalid_raises(): + with pytest.raises(ValueError, match="not appropriate"): + CodonAnnotation._prepare_codontable("not-a-table") + + +def test_get_syn_codons_returns_set(coda): + result = coda.get_syn_codons("ATA", 1) + assert isinstance(result, set) + assert len(result) == 0 + + +def test_cds_length_not_divisible_by_3(coda): + seq3 = "ATGAAATAA" # Met-Lys-Stop (9 nt); 3rd position of the stop is not synonymous + seq4 = seq3 + "C" # incomplete last codon + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + freqs3, _ = coda.collect_exp_mut_freqs(seq3, labels=["all", "syn"]) + freqs4, _ = coda.collect_exp_mut_freqs(seq4, labels=["all", "syn"]) + assert any(issubclass(w.category, UserWarning) for w in caught) + # codon-aware (syn) counts ignore the extra base + assert freqs3["syn"] == freqs4["syn"] + # "all" can use the extra nucleotide as right-hand context of the last codon + assert sum(freqs4["all"].values()) >= sum(freqs3["all"].values()) + + +def test_collect_exp_muts_proba_missing_logger(coda): + """CodonAnnotation has no logger; a bad mask must still raise ValueError.""" + genome = np.zeros((6, 4), dtype=float) + genome[:, 0] = 1.0 + with pytest.raises(ValueError, match="same length"): + coda.collect_exp_muts_proba(genome, phylocoef=1.0, mask=[1, 0]) + + +def test_extract_mutations_simple(coda): + g1 = np.array(list("ATGCTAGTA")) # Met-Leu-Val + g2 = np.array(list("ATGCTGGTA")) # CTA -> CTG (Leu, fourfold) + muts = coda.extract_mutations_simple(g1, g2) + assert len(muts) == 1 + row = muts.iloc[0] + assert row["Mut"] == "T[A>G]G" + assert int(row["Label"]) == 2 + assert row["RefCodon"] == "CTA" + assert row["AltCodon"] == "CTG" + + +def test_extract_mutations_simple_empty_has_columns(coda): + g = np.array(list("ATGCTAGTA")) + muts = coda.extract_mutations_simple(g, g) + assert muts.empty + assert "Mut" in muts.columns + + +def test_plot_mutspec_preserves_provided_axes(): + ms = pd.DataFrame({"Mut": ["C>A", "C>G", "C>T"], "MutSpec": [0.2, 0.3, 0.5]}) + fig, (ax1, ax2) = plt.subplots(1, 2) + plot_mutspec12(ms, ax=ax1, show=False, title="left") + plot_mutspec12(ms, ax=ax2, show=False, title="right") + assert fig.number in plt.get_fignums() + assert ax1.get_title() == "left" + assert ax2.get_title() == "right" + plt.close(fig) + + +def test_complete_sbs192_columns_no_fragmentation(): + partial = pd.DataFrame({"A[C>A]A": [1.0], "T[G>T]T": [2.0]}) + complete = complete_sbs192_columns(partial) + assert list(complete.columns) == possible_sbs192 + assert complete["A[C>A]A"].iloc[0] == 1.0 + assert complete["A[A>C]A"].iloc[0] == 0.0 + + +def test_collapse_mutspec_accepts_obsnum_expnum(): + rows = [] + for sbs in possible_sbs192: + rows.append({"Mut": sbs, "ObsNum": 1.0, "ExpNum": 1.0}) + ms192 = pd.DataFrame(rows) + ms96 = collapse_mutspec(ms192) + assert len(ms96) == 96 + assert "RawMutSpec" in ms96.columns + + +def test_jackknife_does_not_mutate_input(): + idx = pd.MultiIndex.from_tuples( + [("R1", "A1"), ("R1", "A2"), ("R2", "A3")], + names=["RefNode", "AltNode"], + ) + obs = pd.DataFrame(1.0, index=idx, columns=possible_sbs192) + exp = pd.DataFrame(1.0, index=pd.Index(["R1", "R2"], name="Node"), columns=possible_sbs192) + obs_index_before = list(obs.index.names) + exp_name_before = exp.index.name + jackknife_spectra_sampling(obs, exp, frac=0.5, n=2) + assert list(obs.index.names) == obs_index_before + assert exp.index.name == exp_name_before + + +def test_nobs_cutoff_alias(): + idx = pd.MultiIndex.from_tuples( + [("R1", "A1"), ("R1", "A2")], + names=["RefNode", "AltNode"], + ) + obs = pd.DataFrame(5.0, index=idx, columns=possible_sbs192) + exp = pd.DataFrame(5.0, index=pd.Index(["R1"], name="Node"), columns=possible_sbs192) + s1 = calc_edgewise_spectra(obs, exp, nmtypes_cutoff=0, nobs_cutoff=1, scale=False) + s2 = calc_edgewise_spectra(obs, exp, nmtypes_cutoff=0, nobs_cuttof=1, scale=False) + pd.testing.assert_frame_equal(s1, s2) + + +def test_tree_from_newick_string(): + t = Tree("(A:0.1,B:0.2)Root:0.0;") + assert t.name in ("Root", "") + leaves = sorted(n.name for n in t.iter_leaves()) + assert leaves == ["A", "B"] + assert t.parent is None + child = t.children[0] + assert child.parent is t + assert child.up is t + edges = list(t.iter_edges()) + assert len(edges) == 2 + + +def test_tree_missing_file_raises(): + with pytest.raises(FileNotFoundError): + Tree("definitely_not_a_tree_file.nwk") + + +def test_get_tree_height_zero_distances(): + root = TreeNode(name="R", dist=0.0) + a = TreeNode(name="A", dist=0.0) + b = TreeNode(name="B", dist=0.0) + a._parent = b._parent = root + root.children = [a, b] + assert get_tree_height(root, "geom_mean") == 0.0 + assert get_tree_height(root, "mean") == 0.0 + coefs = calc_phylocoefs(root) + assert coefs["R"] == 1.0 + + +def test_genome_states_reads_states_not_gappy_file(): + states_txt = ( + "Node\tSite\tp_A\tp_C\tp_G\tp_T\n" + "N1\t1\t1.0\t0.0\t0.0\t0.0\n" + "N1\t2\t0.0\t1.0\t0.0\t0.0\n" + "N1\t3\t0.0\t0.0\t1.0\t0.0\n" + ) + gappy_txt = "99\n" + with tempfile.TemporaryDirectory() as tmp: + states_path = os.path.join(tmp, "states.tsv") + gappy_path = os.path.join(tmp, "gappy.csv") + with open(states_path, "w") as fh: + fh.write(states_txt) + with open(gappy_path, "w") as fh: + fh.write(gappy_txt) + gs = GenomeStates(states_path, path_to_gappy_sites=gappy_path, format="tsv") + genome = gs.get_genome("N1") + assert list(genome.index) == [1, 2, 3] + assert gs.genome_size == 3 + + +def test_genome_states_gappy_optional(): + states_txt = ( + "Node\tSite\tp_A\tp_C\tp_G\tp_T\n" + "N1\t1\t1.0\t0.0\t0.0\t0.0\n" + "N1\t2\t0.0\t1.0\t0.0\t0.0\n" + ) + with tempfile.TemporaryDirectory() as tmp: + states_path = os.path.join(tmp, "states.tsv") + with open(states_path, "w") as fh: + fh.write(states_txt) + gs = GenomeStates(states_path, path_to_gappy_sites=None, format="tsv") + assert gs.genome_size == 2 + + +def test_basic_logger_does_not_duplicate_handlers(): + log1 = basic_logger() + n1 = len(log1.handlers) + log2 = basic_logger() + assert log1 is log2 + assert len(log2.handlers) == n1 + + +def test_read_genbank_ref_uses_gene_qualifier(): + gb = """\ +LOCUS TEST 12 bp DNA UNK 01-JAN-1980 +FEATURES Location/Qualifiers + source 1..12 + /organism="test" + CDS 1..12 + /gene="cox1" + /codon_start=1 +ORIGIN + 1 atgctagtaa tg +// +""" + with tempfile.NamedTemporaryFile("w", suffix=".gb", delete=False) as fh: + fh.write(gb) + path = fh.name + try: + df = read_genbank_ref(path) + finally: + os.remove(path) + assert "gene" in df.columns + assert (df.loc[df["Type"] == "CDS", "gene"] == "cox1").all() + + +def test_filter_short_exp_seqs_drops_outlier(): + import sys + sys.path.append("./scripts") + from calculate_mutspec import filter_short_exp_seqs + from pymutspec.constants import possible_sbs12, possible_sbs192 + + cols = possible_sbs12 + possible_sbs192 + rows = [] + for node in [f"n{i}" for i in range(8)] + ["outlier"]: + row = {c: 1.0 for c in cols} + row["Node"] = node + row["Label"] = "all" + row["Gene"] = "g" + if node == "outlier": + for c in possible_sbs12: + row[c] = 10000.0 + rows.append(row) + flt = filter_short_exp_seqs(pd.DataFrame(rows)) + assert "outlier" not in set(flt.Node) + assert len(flt) == 8 diff --git a/tests/test_codon_ann.py b/tests/test_codon_ann.py index 16ecd80..a8098ba 100644 --- a/tests/test_codon_ann.py +++ b/tests/test_codon_ann.py @@ -1,4 +1,5 @@ from collections import Counter +import numpy as np import pytest def test_get_syn_codons(coda): @@ -194,9 +195,13 @@ def test_collect_exp_muts_on_real_gene_proba(coda, states): assert exp_sbs192_freqs['nonsyn'][k] == pytest.approx(v, abs=1e-3) -def test_extract_mutations_simple(): - # TODO - pass +def test_extract_mutations_simple(coda): + g1 = np.array(list("ATGCTAGTA")) + g2 = np.array(list("ATGCTGGTA")) + muts = coda.extract_mutations_simple(g1, g2) + assert len(muts) == 1 + assert muts.iloc[0]["Mut"] == "T[A>G]G" + assert int(muts.iloc[0]["Label"]) == 2 # use to test collect_exp_muts_proba and collect_exp_mut_freqs_proba (results must be comparable and equal) # np.all(coda.collect_exp_muts_proba(states.get_genome("Node4705")["1"], 1, mut_proba_cutoff=0.05).groupby(["Label", "Mut"]).Proba.sum().unstack()[possible_sbs192].fillna(0) == \ diff --git a/tests/test_mutspec_calc.py b/tests/test_mutspec_calc.py index 339f201..9a0ef8f 100644 --- a/tests/test_mutspec_calc.py +++ b/tests/test_mutspec_calc.py @@ -3,7 +3,7 @@ import pandas as pd -from pymutspec.annotation import calculate_mutspec +from pymutspec.annotation import calculate_mutspec, calculate_mutrate from pymutspec.constants import possible_sbs12, possible_sbs192 @@ -56,7 +56,7 @@ def test_ms12_calc(mut, nucl_freqs, use_proba, lbl_id): cur_mut = mut[(mut.Label >= lbl_id)] ms = calculate_mutspec(cur_mut, nucl_freqs[lbl], use_context=False, use_proba=use_proba) for sbs in possible_sbs12: - divisor = nucl_freqs[lbl].get(sbs[0], 0) + divisor = nucl_freqs[lbl].get(sbs, 0) if divisor <= 0: continue if use_proba: @@ -81,9 +81,8 @@ def test_ms192_calc(mut, cxt_freqs, use_proba, lbl_id): ms = calculate_mutspec(cur_mut, cxt_freqs[lbl], use_context=True, use_proba=use_proba, fill_unobserved=False) - for sbs in mut['Mut'].unique(): - cxt = sbs[0] + sbs[2] + sbs[-1] - divisor = cxt_freqs[lbl].get(cxt, 0) + for sbs in cur_mut['Mut'].unique(): + divisor = cxt_freqs[lbl].get(sbs, 0) if divisor == 0: continue cond = cur_mut.Mut.str.fullmatch(sbs.replace("[", r"\[").replace("]", r"\]")) @@ -93,3 +92,39 @@ def test_ms192_calc(mut, cxt_freqs, use_proba, lbl_id): expected = cur_mut[cond].shape[0] / divisor observed = ms[ms.Mut == sbs].RawMutSpec.values[0] assert observed == expected + + +def test_calculate_mutspec_keeps_raw_and_scaled(mut, nucl_freqs): + ms = calculate_mutspec( + mut, nucl_freqs["all"], use_context=False, use_proba=True, + drop_underrepresented=False, + ) + assert "RawMutSpec" in ms.columns + assert "MutSpec" in ms.columns + assert pytest.approx(ms["MutSpec"].sum()) == 1.0 + positive = ms[ms["RawMutSpec"] > 0] + ratios = positive["MutSpec"] / positive["RawMutSpec"] + assert ratios.max() == pytest.approx(ratios.min()) + + +def test_calculate_mutspec_empty_does_not_nan(): + obs = pd.DataFrame({"Mut": pd.Series(dtype=str), "ProbaFull": pd.Series(dtype=float)}) + exp = {"A>C": 1.0, "C>T": 2.0} + ms = calculate_mutspec(obs, exp, use_context=False, use_proba=True) + assert ms["MutSpec"].isna().sum() == 0 + assert (ms["MutSpec"] == 0).all() + + +def test_calculate_mutspec_accepts_sbs12_mut_column(): + obs = pd.DataFrame({"Mut": ["C>T", "C>T", "A>G"], "ProbaFull": [1.0, 1.0, 0.5]}) + exp = {"C>T": 2.0, "A>G": 1.0} + ms = calculate_mutspec(obs, exp, use_context=False, use_proba=True, scale=False) + by_mut = ms.set_index("Mut") + assert by_mut.loc["C>T", "RawMutSpec"] == pytest.approx(1.0) + assert by_mut.loc["A>G", "RawMutSpec"] == pytest.approx(0.5) + + +def test_calculate_mutrate(mut, nucl_freqs): + rates = calculate_mutrate(mut, nucl_freqs["all"], use_context=False, use_proba=True) + assert "MutRate" in rates.columns + assert (rates["MutRate"] == rates["RawMutSpec"]).all()