smftools 0.1.7__py3-none-any.whl → 0.2.3__py3-none-any.whl

smftools/tools/position_stats.py

@@ -117,123 +117,485 @@ def calculate_relative_risk_on_activity(adata, sites, alpha=0.05, groupby=None):
 
     return results_dict
 
-def compute_positionwise_statistic(
+import copy
+import warnings
+from typing import Dict, Any, List, Optional, Tuple, Union
+
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+# optional imports
+try:
+    from joblib import Parallel, delayed
+    JOBLIB_AVAILABLE = True
+except Exception:
+    JOBLIB_AVAILABLE = False
+
+try:
+    from scipy.stats import chi2_contingency
+    SCIPY_STATS_AVAILABLE = True
+except Exception:
+    SCIPY_STATS_AVAILABLE = False
+
+# -----------------------------
+# Compute positionwise statistic (multi-method + simple site_types)
+# -----------------------------
+import numpy as np
+import pandas as pd
+from typing import List, Optional, Sequence, Dict, Any, Tuple
+from contextlib import contextmanager
+from joblib import Parallel, delayed, cpu_count
+import joblib
+from tqdm import tqdm
+from scipy.stats import chi2_contingency
+import warnings
+import matplotlib.pyplot as plt
+from itertools import cycle
+import os
+import warnings
+
+
+# ---------------------------
+# joblib <-> tqdm integration
+# ---------------------------
+@contextmanager
+def tqdm_joblib(tqdm_object: tqdm):
+    """Context manager to patch joblib to update a tqdm progress bar."""
+    old = joblib.parallel.BatchCompletionCallBack
+
+    class TqdmBatchCompletionCallback(old):  # type: ignore
+        def __call__(self, *args, **kwargs):
+            try:
+                tqdm_object.update(n=self.batch_size)
+            except Exception:
+                tqdm_object.update(1)
+            return super().__call__(*args, **kwargs)
+
+    joblib.parallel.BatchCompletionCallBack = TqdmBatchCompletionCallback
+    try:
+        yield tqdm_object
+    finally:
+        joblib.parallel.BatchCompletionCallBack = old
+
+
+# ---------------------------
+# row workers (upper-triangle only)
+# ---------------------------
+def _chi2_row_job(i: int, X_bin: np.ndarray, min_count_for_pairwise: int) -> Tuple[int, np.ndarray]:
+    n_pos = X_bin.shape[1]
+    row = np.full((n_pos,), np.nan, dtype=float)
+    xi = X_bin[:, i]
+    for j in range(i, n_pos):
+        xj = X_bin[:, j]
+        mask = (~np.isnan(xi)) & (~np.isnan(xj))
+        if int(mask.sum()) < int(min_count_for_pairwise):
+            continue
+        try:
+            table = pd.crosstab(xi[mask], xj[mask])
+            if table.shape != (2, 2):
+                continue
+            chi2, _, _, _ = chi2_contingency(table, correction=False)
+            row[j] = float(chi2)
+        except Exception:
+            row[j] = np.nan
+    return (i, row)
+
+
+def _relative_risk_row_job(i: int, X_bin: np.ndarray, min_count_for_pairwise: int) -> Tuple[int, np.ndarray]:
+    n_pos = X_bin.shape[1]
+    row = np.full((n_pos,), np.nan, dtype=float)
+    xi = X_bin[:, i]
+    for j in range(i, n_pos):
+        xj = X_bin[:, j]
+        mask = (~np.isnan(xi)) & (~np.isnan(xj))
+        if int(mask.sum()) < int(min_count_for_pairwise):
+            continue
+        a = np.sum((xi[mask] == 1) & (xj[mask] == 1))
+        b = np.sum((xi[mask] == 1) & (xj[mask] == 0))
+        c = np.sum((xi[mask] == 0) & (xj[mask] == 1))
+        d = np.sum((xi[mask] == 0) & (xj[mask] == 0))
+        try:
+            if (a + b) > 0 and (c + d) > 0 and (c > 0):
+                p1 = a / float(a + b)
+                p2 = c / float(c + d)
+                row[j] = float(p1 / p2) if p2 > 0 else np.nan
+            else:
+                row[j] = np.nan
+        except Exception:
+            row[j] = np.nan
+    return (i, row)
+
+def compute_positionwise_statistics(
     adata,
-    layer,
-    method="pearson",
-    groupby=["Reference_strand"],
-    output_key="positionwise_result",
-    site_config=None,
-    encoding="signed",
-    max_threads=None
+    layer: str,
+    methods: Sequence[str] = ("pearson",),
+    sample_col: str = "Barcode",
+    ref_col: str = "Reference_strand",
+    site_types: Optional[Sequence[str]] = None,
+    encoding: str = "signed",
+    output_key: str = "positionwise_result",
+    min_count_for_pairwise: int = 10,
+    max_threads: Optional[int] = None,
+    reverse_indices_on_store: bool = False,
 ):
     """
-    Computes a position-by-position matrix (correlation, RR, or Chi-squared) from an adata layer.
+    Compute per-(sample,ref) positionwise matrices for methods in `methods`.
 
-    Parameters:
-        adata (AnnData): Annotated data matrix.
-        layer (str): Name of the adata layer to use.
-        method (str): 'pearson', 'binary_covariance', 'relative_risk', or 'chi_squared'.
-        groupby (str or list): Column(s) in adata.obs to group by.
-        output_key (str): Key in adata.uns to store results.
-        site_config (dict): Optional {ref: [site_types]} to restrict sites per reference.
-        encoding (str): 'signed' (1/-1/0) or 'binary' (1/0/NaN).
-        max_threads (int): Number of parallel threads to use (joblib).
+    Results stored at:
+      adata.uns[output_key][method][ (sample, ref) ] = DataFrame
+      adata.uns[output_key + "_n"][method][ (sample, ref) ] = int(n_reads)
     """
-    import numpy as np
-    import pandas as pd
-    from scipy.stats import chi2_contingency
-    from joblib import Parallel, delayed
-    from tqdm import tqdm
+    if isinstance(methods, str):
+        methods = [methods]
+    methods = [m.lower() for m in methods]
+
+    # prepare containers
+    adata.uns[output_key] = {m: {} for m in methods}
+    adata.uns[output_key + "_n"] = {m: {} for m in methods}
+
+    # workers
+    if max_threads is None or max_threads <= 0:
+        n_jobs = max(1, cpu_count() or 1)
+    else:
+        n_jobs = max(1, int(max_threads))
+
+    # samples / refs
+    sseries = adata.obs[sample_col]
+    if not pd.api.types.is_categorical_dtype(sseries):
+        sseries = sseries.astype("category")
+    samples = list(sseries.cat.categories)
+
+    rseries = adata.obs[ref_col]
+    if not pd.api.types.is_categorical_dtype(rseries):
+        rseries = rseries.astype("category")
+    references = list(rseries.cat.categories)
+
+    total_tasks = len(samples) * len(references)
+    pbar_outer = tqdm(total=total_tasks, desc="positionwise (sample x ref)", unit="cell")
+
+    for sample in samples:
+        for ref in references:
+            label = (sample, ref)
+            try:
+                mask = (adata.obs[sample_col] == sample) & (adata.obs[ref_col] == ref)
+                subset = adata[mask]
+                n_reads = subset.shape[0]
+
+                # nothing to do -> store empty placeholders
+                if n_reads == 0:
+                    for m in methods:
+                        adata.uns[output_key][m][label] = pd.DataFrame()
+                        adata.uns[output_key + "_n"][m][label] = 0
+                    pbar_outer.update(1)
+                    continue
 
-    if isinstance(groupby, str):
-        groupby = [groupby]
+                # select var columns based on site_types and reference
+                if site_types:
+                    col_mask = np.zeros(subset.shape[1], dtype=bool)
+                    for st in site_types:
+                        colname = f"{ref}_{st}"
+                        if colname in subset.var.columns:
+                            col_mask |= np.asarray(subset.var[colname].values, dtype=bool)
+                        else:
+                            # if mask not present, warn once (but keep searching)
+                            # user may pass generic site types
+                            pass
+                    if not col_mask.any():
+                        selected_var_idx = np.arange(subset.shape[1])
+                    else:
+                        selected_var_idx = np.nonzero(col_mask)[0]
+                else:
+                    selected_var_idx = np.arange(subset.shape[1])
+
+                if selected_var_idx.size == 0:
+                    for m in methods:
+                        adata.uns[output_key][m][label] = pd.DataFrame()
+                        adata.uns[output_key + "_n"][m][label] = int(n_reads)
+                    pbar_outer.update(1)
+                    continue
 
-    adata.uns[output_key] = {}
-    adata.uns[output_key + "_n"] = {}
+                # extract matrix
+                if (layer in subset.layers) and (subset.layers[layer] is not None):
+                    X = subset.layers[layer]
+                else:
+                    X = subset.X
+                X = np.asarray(X, dtype=float)
+                X = X[:, selected_var_idx]  # (n_reads, n_pos)
+
+                # binary encoding
+                if encoding == "signed":
+                    X_bin = np.where(X == 1, 1.0, np.where(X == -1, 0.0, np.nan))
+                else:
+                    X_bin = np.where(X == 1, 1.0, np.where(X == 0, 0.0, np.nan))
+
+                n_pos = X_bin.shape[1]
+                if n_pos == 0:
+                    for m in methods:
+                        adata.uns[output_key][m][label] = pd.DataFrame()
+                        adata.uns[output_key + "_n"][m][label] = int(n_reads)
+                    pbar_outer.update(1)
+                    continue
 
-    label_col = "__".join(groupby)
-    adata.obs[label_col] = adata.obs[groupby].astype(str).agg("_".join, axis=1)
+                var_names = list(subset.var_names[selected_var_idx])
+
+                # compute per-method
+                for method in methods:
+                    m = method.lower()
+                    if m == "pearson":
+                        # pairwise Pearson with column demean (nan-aware approximation)
+                        with np.errstate(invalid="ignore"):
+                            col_mean = np.nanmean(X_bin, axis=0)
+                            Xc = X_bin - col_mean  # nan preserved
+                            Xc0 = np.nan_to_num(Xc, nan=0.0)
+                            cov = Xc0.T @ Xc0
+                            denom = (np.sqrt((Xc0**2).sum(axis=0))[:, None] * np.sqrt((Xc0**2).sum(axis=0))[None, :])
+                            with np.errstate(divide="ignore", invalid="ignore"):
+                                mat = np.where(denom != 0.0, cov / denom, np.nan)
+                    elif m == "binary_covariance":
+                        binary = (X_bin == 1).astype(float)
+                        valid = (~np.isnan(X_bin)).astype(float)
+                        with np.errstate(divide="ignore", invalid="ignore"):
+                            numerator = binary.T @ binary
+                            denominator = valid.T @ valid
+                            mat = np.true_divide(numerator, denominator)
+                            mat[~np.isfinite(mat)] = 0.0
+                    elif m in ("chi_squared", "relative_risk"):
+                        if m == "chi_squared":
+                            worker = _chi2_row_job
+                        else:
+                            worker = _relative_risk_row_job
+                        out = np.full((n_pos, n_pos), np.nan, dtype=float)
+                        tasks = (delayed(worker)(i, X_bin, min_count_for_pairwise) for i in range(n_pos))
+                        pbar_rows = tqdm(total=n_pos, desc=f"{m}: rows ({sample}__{ref})", leave=False)
+                        with tqdm_joblib(pbar_rows):
+                            results = Parallel(n_jobs=n_jobs, prefer="processes")(tasks)
+                        pbar_rows.close()
+                        for i, row in results:
+                            out[int(i), :] = row
+                        iu = np.triu_indices(n_pos, k=1)
+                        out[iu[1], iu[0]] = out[iu]
+                        mat = out
+                    else:
+                        raise ValueError(f"Unsupported method: {method}")
+
+                    # optionally reverse order at store-time
+                    if reverse_indices_on_store:
+                        mat_store = np.flip(np.flip(mat, axis=0), axis=1)
+                        idx_names = var_names[::-1]
+                    else:
+                        mat_store = mat
+                        idx_names = var_names
+
+                    # make dataframe with labels
+                    df = pd.DataFrame(mat_store, index=idx_names, columns=idx_names)
+
+                    adata.uns[output_key][m][label] = df
+                    adata.uns[output_key + "_n"][m][label] = int(n_reads)
+
+            except Exception as exc:
+                warnings.warn(f"Failed computing positionwise for {sample}__{ref}: {exc}")
+            finally:
+                pbar_outer.update(1)
+
+    pbar_outer.close()
+    return None
+
+
+# ---------------------------
+# Plotting function
+# ---------------------------
+
+def plot_positionwise_matrices(
+    adata,
+    methods: List[str],
+    cmaps: Optional[List[str]] = None,
+    sample_col: str = "Barcode",
+    ref_col: str = "Reference_strand",
+    output_dir: Optional[str] = None,
+    vmin: Optional[Dict[str, float]] = None,
+    vmax: Optional[Dict[str, float]] = None,
+    figsize_per_cell: Tuple[float, float] = (3.5, 3.5),
+    dpi: int = 160,
+    cbar_shrink: float = 0.9,
+    output_key: str = "positionwise_result",
+    show_colorbar: bool = True,
+    flip_display_axes: bool = False,
+    rows_per_page: int = 6,
+    sample_label_rotation: float = 90.0,
+):
+    """
+    Plot grids of matrices for each method with pagination and rotated sample-row labels.
 
-    for group in adata.obs[label_col].unique():
-        subset = adata[adata.obs[label_col] == group].copy()
-        if subset.shape[0] == 0:
+    New args:
+      - rows_per_page: how many sample rows per page/figure (pagination)
+      - sample_label_rotation: rotation angle (deg) for the sample labels placed in the left margin.
+    Returns:
+      dict mapping method -> list of saved filenames (empty list if figures were shown).
+    """
+    if isinstance(methods, str):
+        methods = [methods]
+    if cmaps is None:
+        cmaps = ["viridis"] * len(methods)
+    cmap_cycle = cycle(cmaps)
+
+    # canonicalize sample/ref order
+    sseries = adata.obs[sample_col]
+    if not pd.api.types.is_categorical_dtype(sseries):
+        sseries = sseries.astype("category")
+    samples = list(sseries.cat.categories)
+
+    rseries = adata.obs[ref_col]
+    if not pd.api.types.is_categorical_dtype(rseries):
+        rseries = rseries.astype("category")
+    references = list(rseries.cat.categories)
+
+    # ensure directories
+    if output_dir:
+        os.makedirs(output_dir, exist_ok=True)
+
+    saved_files_by_method = {}
+
+    def _get_df_from_store(store, sample, ref):
+        """
+        try multiple key formats: (sample, ref) tuple, 'sample__ref' string,
+        or str(sample)+'__'+str(ref). Return None if not found.
+        """
+        if store is None:
+            return None
+        # try tuple key
+        key_t = (sample, ref)
+        if key_t in store:
+            return store[key_t]
+        # try string key
+        key_s = f"{sample}__{ref}"
+        if key_s in store:
+            return store[key_s]
+        # try stringified tuple keys (some callers store differently)
+        for k in store.keys():
+            try:
+                if isinstance(k, tuple) and len(k) == 2 and str(k[0]) == str(sample) and str(k[1]) == str(ref):
+                    return store[k]
+                if isinstance(k, str) and key_s == k:
+                    return store[k]
+            except Exception:
+                continue
+        return None
+
+    for method, cmap in zip(methods, cmap_cycle):
+        m = method.lower()
+        method_store = adata.uns.get(output_key, {}).get(m, {})
+        if not method_store:
+            warnings.warn(f"No results found for method '{method}' in adata.uns['{output_key}']. Skipping.", stacklevel=2)
+            saved_files_by_method[method] = []
             continue
 
-        ref = subset.obs["Reference_strand"].unique()[0] if "Reference_strand" in groupby else None
-
-        if site_config and ref in site_config:
-            site_mask = np.zeros(subset.shape[1], dtype=bool)
-            for site in site_config[ref]:
-                site_mask |= subset.var[f"{ref}_{site}"]
-            subset = subset[:, site_mask].copy()
-
-        X = subset.layers[layer].copy()
-
-        if encoding == "signed":
-            X_bin = np.where(X == 1, 1, np.where(X == -1, 0, np.nan))
+        # gather numeric values to pick sensible vmin/vmax when not provided
+        vals = []
+        for s in samples:
+            for r in references:
+                df = _get_df_from_store(method_store, s, r)
+                if isinstance(df, pd.DataFrame) and df.size > 0:
+                    a = df.values
+                    a = a[np.isfinite(a)]
+                    if a.size:
+                        vals.append(a)
+        if vals:
+            allvals = np.concatenate(vals)
         else:
-            X_bin = np.where(X == 1, 1, np.where(X == 0, 0, np.nan))
-
-        n_pos = subset.shape[1]
-        mat = np.zeros((n_pos, n_pos))
-
-        if method == "pearson":
-            with np.errstate(invalid='ignore'):
-                mat = np.corrcoef(np.nan_to_num(X_bin).T)
-
-        elif method == "binary_covariance":
-            binary = (X_bin == 1).astype(float)
-            valid = (X_bin == 1) | (X_bin == 0)  # Only consider true binary (ignore NaN)
-            valid = valid.astype(float)
-
-            numerator = np.dot(binary.T, binary)
-            denominator = np.dot(valid.T, valid)
-
-            with np.errstate(divide='ignore', invalid='ignore'):
-                mat = np.true_divide(numerator, denominator)
-                mat[~np.isfinite(mat)] = 0
-
-        elif method in ["relative_risk", "chi_squared"]:
-            def compute_row(i):
-                row = np.zeros(n_pos)
-                xi = X_bin[:, i]
-                for j in range(n_pos):
-                    xj = X_bin[:, j]
-                    mask = ~np.isnan(xi) & ~np.isnan(xj)
-                    if np.sum(mask) < 10:
-                        row[j] = np.nan
-                        continue
-                    if method == "relative_risk":
-                        a = np.sum((xi[mask] == 1) & (xj[mask] == 1))
-                        b = np.sum((xi[mask] == 1) & (xj[mask] == 0))
-                        c = np.sum((xi[mask] == 0) & (xj[mask] == 1))
-                        d = np.sum((xi[mask] == 0) & (xj[mask] == 0))
-                        if (a + b) > 0 and (c + d) > 0 and c > 0:
-                            p1 = a / (a + b)
-                            p2 = c / (c + d)
-                            row[j] = p1 / p2 if p2 > 0 else np.nan
-                        else:
-                            row[j] = np.nan
-                    elif method == "chi_squared":
-                        table = pd.crosstab(xi[mask], xj[mask])
-                        if table.shape != (2, 2):
-                            row[j] = np.nan
-                        else:
-                            chi2, _, _, _ = chi2_contingency(table, correction=False)
-                            row[j] = chi2
-                return row
-
-            mat = np.array(
-                Parallel(n_jobs=max_threads)(
-                    delayed(compute_row)(i) for i in tqdm(range(n_pos), desc=f"{method}: {group}")
-                )
-            )
-
+            allvals = np.array([])
+
+        # decide per-method defaults
+        if m == "pearson":
+            vmn = -1.0 if (vmin is None or (isinstance(vmin, dict) and m not in vmin)) else (vmin.get(m) if isinstance(vmin, dict) else vmin)
+            vmx = 1.0 if (vmax is None or (isinstance(vmax, dict) and m not in vmax)) else (vmax.get(m) if isinstance(vmax, dict) else vmax)
+            vmn = -1.0 if vmn is None else vmn
+            vmx = 1.0 if vmx is None else vmx
+        elif m == "binary_covariance":
+            vmn = 0.0 if (vmin is None or (isinstance(vmin, dict) and m not in vmin)) else (vmin.get(m) if isinstance(vmin, dict) else vmin)
+            vmx = 1.0 if (vmax is None or (isinstance(vmax, dict) and m not in vmax)) else (vmax.get(m) if isinstance(vmax, dict) else vmax)
+            vmn = 0.0 if vmn is None else vmn
+            vmx = 1.0 if vmx is None else vmx
         else:
-            raise ValueError(f"Unsupported method: {method}")
+            vmn = 0.0 if (vmin is None or (isinstance(vmin, dict) and m not in vmin)) else (vmin.get(m) if isinstance(vmin, dict) else vmin)
+            if (vmax is None) or (isinstance(vmax, dict) and m not in vmax):
+                vmx = float(np.nanpercentile(allvals, 99.0)) if allvals.size else 1.0
+            else:
+                vmx = (vmax.get(m) if isinstance(vmax, dict) else vmax)
+            vmn = 0.0 if vmn is None else vmn
+            if vmx is None:
+                vmx = 1.0
+
+        # prepare pagination over sample rows
+        saved_files = []
+        n_pages = max(1, int(np.ceil(len(samples) / float(max(1, rows_per_page)))))
+        for page_idx in range(n_pages):
+            start = page_idx * rows_per_page
+            chunk = samples[start : start + rows_per_page]
+            nrows = len(chunk)
+            ncols = max(1, len(references))
+            fig_w = ncols * figsize_per_cell[0]
+            fig_h = nrows * figsize_per_cell[1]
+            fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(fig_w, fig_h), dpi=dpi, squeeze=False)
+
+            # leave margin for rotated sample labels
+            plt.subplots_adjust(left=0.12, right=0.88, top=0.95, bottom=0.05)
+
+            any_plotted = False
+            im = None
+            for r_idx, sample in enumerate(chunk):
+                for c_idx, ref in enumerate(references):
+                    ax = axes[r_idx][c_idx]
+                    df = _get_df_from_store(method_store, sample, ref)
+                    if not isinstance(df, pd.DataFrame) or df.size == 0:
+                        ax.text(0.5, 0.5, "No data", ha="center", va="center", transform=ax.transAxes, fontsize=10, color="gray")
+                        ax.set_xticks([])
+                        ax.set_yticks([])
+                    else:
+                        mat = df.values.astype(float)
+                        origin = "upper" if flip_display_axes else "lower"
+                        im = ax.imshow(mat, origin=origin, aspect="auto", vmin=vmn, vmax=vmx, cmap=cmap)
+                        any_plotted = True
+                        ax.set_xticks([])
+                        ax.set_yticks([])
+
+                    # top title is reference (only for top-row)
+                    if r_idx == 0:
+                        ax.set_title(str(ref), fontsize=9)
+
+                # draw rotated sample label into left margin centered on the row
+                # compute vertical center of this row's axis in figure coords
+                ax0 = axes[r_idx][0]
+                ax_y0, ax_y1 = ax0.get_position().y0, ax0.get_position().y1
+                y_center = 0.5 * (ax_y0 + ax_y1)
+                # place text at x=0.01 (just inside left margin); rotation controls orientation
+                fig.text(0.01, y_center, str(chunk[r_idx]), va="center", ha="left", rotation=sample_label_rotation, fontsize=9)
+
+            fig.suptitle(f"{method} — per-sample x per-reference matrices (page {page_idx+1}/{n_pages})", fontsize=12, y=0.99)
+            fig.tight_layout(rect=[0.05, 0.02, 0.9, 0.96])
+
+            # colorbar (shared)
+            if any_plotted and show_colorbar and (im is not None):
+                try:
+                    cbar_ax = fig.add_axes([0.9, 0.15, 0.02, 0.7])
+                    fig.colorbar(im, cax=cbar_ax, shrink=cbar_shrink)
+                except Exception:
+                    try:
+                        fig.colorbar(im, ax=axes.ravel().tolist(), fraction=0.02, pad=0.02)
+                    except Exception:
+                        pass
+
+            # save or show
+            if output_dir:
+                fname = f"positionwise_{method}_page{page_idx+1}.png"
+                outpath = os.path.join(output_dir, fname)
+                plt.savefig(outpath, bbox_inches="tight")
+                saved_files.append(outpath)
+                plt.close(fig)
+            else:
+                plt.show()
+                saved_files.append("")  # placeholder to indicate a figure was shown
+
+        saved_files_by_method[method] = saved_files
 
-        var_names = subset.var_names.astype(int)
-        mat_df = pd.DataFrame(mat, index=var_names, columns=var_names)
-        adata.uns[output_key][group] = mat_df
-        adata.uns[output_key + "_n"][group] = subset.shape[0]
+    return saved_files_by_method