smftools 0.2.4__py3-none-any.whl → 0.3.0__py3-none-any.whl

smftools/tools/position_stats.py

@@ -1,23 +1,61 @@
+from __future__ import annotations
+
+import os
+import warnings
+from contextlib import contextmanager
+from itertools import cycle
+from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple
+
+import numpy as np
+import pandas as pd
+from scipy.stats import chi2_contingency
+from tqdm import tqdm
+
+from smftools.optional_imports import require
+
+if TYPE_CHECKING:
+    import anndata as ad
+
+plt = require("matplotlib.pyplot", extra="plotting", purpose="position stats plots")
+
+# -----------------------------
+# Compute positionwise statistic (multi-method + simple site_types)
+# -----------------------------
+
+
 # ------------------------- Utilities -------------------------
-def random_fill_nans(X):
+def random_fill_nans(X: np.ndarray) -> np.ndarray:
+    """Fill NaNs with random values in-place.
+
+    Args:
+        X: Input array with NaNs.
+
+    Returns:
+        numpy.ndarray: Array with NaNs replaced by random values.
+    """
     import numpy as np
+
     nan_mask = np.isnan(X)
     X[nan_mask] = np.random.rand(*X[nan_mask].shape)
     return X
 
-def calculate_relative_risk_on_activity(adata, sites, alpha=0.05, groupby=None):
-    """
-    Perform Bayesian-style methylation vs activity analysis independently within each group.
 
-    Parameters:
-        adata (AnnData): Annotated data matrix.
-        sites (list of str): List of site keys (e.g., ['GpC_site', 'CpG_site']).
-        alpha (float): FDR threshold for significance.
-        groupby (str or list of str): Column(s) in adata.obs to group by.
+def calculate_relative_risk_on_activity(
+    adata: "ad.AnnData",
+    sites: Sequence[str],
+    alpha: float = 0.05,
+    groupby: str | Sequence[str] | None = None,
+) -> dict:
+    """Perform methylation vs. activity analysis within each group.
+
+    Args:
+        adata: Annotated data matrix.
+        sites: Site keys (e.g., ``["GpC_site", "CpG_site"]``).
+        alpha: FDR threshold for significance.
+        groupby: Obs column(s) to group by.
 
     Returns:
-        results_dict (dict): Dictionary with structure:
-            results_dict[ref][group_label] = (results_df, sig_df)
+        dict: Mapping of reference -> group label -> ``(results_df, sig_df)``.
     """
     import numpy as np
     import pandas as pd
@@ -25,30 +63,44 @@ def calculate_relative_risk_on_activity(adata, sites, alpha=0.05, groupby=None):
     from statsmodels.stats.multitest import multipletests
 
     def compute_risk_df(ref, site_subset, positions_list, relative_risks, p_values):
-        p_adj = multipletests(p_values, method='fdr_bh')[1] if p_values else []
+        """Build result and significant-data DataFrames for a reference.
+
+        Args:
+            ref: Reference name.
+            site_subset: AnnData subset restricted to sites.
+            positions_list: Positions tested.
+            relative_risks: Relative risk values.
+            p_values: Raw p-values.
+
+        Returns:
+            Tuple of (results_df, sig_df).
+        """
+        p_adj = multipletests(p_values, method="fdr_bh")[1] if p_values else []
 
         genomic_positions = np.array(site_subset.var_names)[positions_list]
         is_gpc_site = site_subset.var[f"{ref}_GpC_site"].values[positions_list]
         is_cpg_site = site_subset.var[f"{ref}_CpG_site"].values[positions_list]
 
-        results_df = pd.DataFrame({
-            'Feature_Index': positions_list,
-            'Genomic_Position': genomic_positions.astype(int),
-            'Relative_Risk': relative_risks,
-            'Adjusted_P_Value': p_adj,
-            'GpC_Site': is_gpc_site,
-            'CpG_Site': is_cpg_site
-        })
-
-        results_df['log2_Relative_Risk'] = np.log2(results_df['Relative_Risk'].replace(0, 1e-300))
-        results_df['-log10_Adj_P'] = -np.log10(results_df['Adjusted_P_Value'].replace(0, 1e-300))
-        sig_df = results_df[results_df['Adjusted_P_Value'] < alpha]
+        results_df = pd.DataFrame(
+            {
+                "Feature_Index": positions_list,
+                "Genomic_Position": genomic_positions.astype(int),
+                "Relative_Risk": relative_risks,
+                "Adjusted_P_Value": p_adj,
+                "GpC_Site": is_gpc_site,
+                "CpG_Site": is_cpg_site,
+            }
+        )
+
+        results_df["log2_Relative_Risk"] = np.log2(results_df["Relative_Risk"].replace(0, 1e-300))
+        results_df["-log10_Adj_P"] = -np.log10(results_df["Adjusted_P_Value"].replace(0, 1e-300))
+        sig_df = results_df[results_df["Adjusted_P_Value"] < alpha]
         return results_df, sig_df
 
     results_dict = {}
 
-    for ref in adata.obs['Reference_strand'].unique():
-        ref_subset = adata[adata.obs['Reference_strand'] == ref].copy()
+    for ref in adata.obs["Reference_strand"].unique():
+        ref_subset = adata[adata.obs["Reference_strand"] == ref].copy()
         if ref_subset.shape[0] == 0:
             continue
 
@@ -56,20 +108,22 @@ def calculate_relative_risk_on_activity(adata, sites, alpha=0.05, groupby=None):
         if groupby is not None:
             if isinstance(groupby, str):
                 groupby = [groupby]
+
             def format_group_label(row):
+                """Format a group label string from obs row values."""
                 return ",".join([f"{col}={row[col]}" for col in groupby])
 
-            combined_label = '__'.join(groupby)
+            combined_label = "__".join(groupby)
             ref_subset.obs[combined_label] = ref_subset.obs.apply(format_group_label, axis=1)
             groups = ref_subset.obs[combined_label].unique()
         else:
             combined_label = None
-            groups = ['all']
+            groups = ["all"]
 
         results_dict[ref] = {}
 
         for group in groups:
-            if group == 'all':
+            if group == "all":
                 group_subset = ref_subset
             else:
                 group_subset = ref_subset[ref_subset.obs[combined_label] == group]
@@ -85,7 +139,7 @@ def calculate_relative_risk_on_activity(adata, sites, alpha=0.05, groupby=None):
 
             # Matrix and labels
             X = random_fill_nans(site_subset.X.copy())
-            y = site_subset.obs['activity_status'].map({'Active': 1, 'Silent': 0}).values
+            y = site_subset.obs["activity_status"].map({"Active": 1, "Silent": 0}).values
             P_active = np.mean(y)
 
             # Analysis
@@ -104,7 +158,9 @@ def calculate_relative_risk_on_activity(adata, sites, alpha=0.05, groupby=None):
                     continue
 
                 P_active_given_methylated = (P_methylated_given_active * P_active) / P_methylated
-                P_active_given_unmethylated = ((1 - P_methylated_given_active) * P_active) / (1 - P_methylated)
+                P_active_given_unmethylated = ((1 - P_methylated_given_active) * P_active) / (
+                    1 - P_methylated
+                )
                 RR = P_active_given_methylated / P_active_given_unmethylated
 
                 _, p_value = fisher_exact(table)
@@ -112,49 +168,13 @@ def calculate_relative_risk_on_activity(adata, sites, alpha=0.05, groupby=None):
                 relative_risks.append(RR)
                 p_values.append(p_value)
 
-            results_df, sig_df = compute_risk_df(ref, site_subset, positions_list, relative_risks, p_values)
+            results_df, sig_df = compute_risk_df(
+                ref, site_subset, positions_list, relative_risks, p_values
+            )
             results_dict[ref][group] = (results_df, sig_df)
 
     return results_dict
 
-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
@@ -162,10 +182,15 @@ import warnings
 @contextmanager
 def tqdm_joblib(tqdm_object: tqdm):
     """Context manager to patch joblib to update a tqdm progress bar."""
+    joblib = require("joblib", extra="ml-base", purpose="parallel position statistics")
+
     old = joblib.parallel.BatchCompletionCallBack
 
     class TqdmBatchCompletionCallback(old):  # type: ignore
+        """Joblib callback that updates a tqdm progress bar."""
+
         def __call__(self, *args, **kwargs):
+            """Update the progress bar when a batch completes."""
             try:
                 tqdm_object.update(n=self.batch_size)
             except Exception:
@@ -183,6 +208,16 @@ def tqdm_joblib(tqdm_object: tqdm):
 # row workers (upper-triangle only)
 # ---------------------------
 def _chi2_row_job(i: int, X_bin: np.ndarray, min_count_for_pairwise: int) -> Tuple[int, np.ndarray]:
+    """Compute chi-squared statistics for one row of a pairwise matrix.
+
+    Args:
+        i: Row index.
+        X_bin: Binary matrix.
+        min_count_for_pairwise: Minimum count for valid comparison.
+
+    Returns:
+        Tuple of (row_index, row_values).
+    """
     n_pos = X_bin.shape[1]
     row = np.full((n_pos,), np.nan, dtype=float)
     xi = X_bin[:, i]
@@ -202,7 +237,19 @@ def _chi2_row_job(i: int, X_bin: np.ndarray, min_count_for_pairwise: int) -> Tup
     return (i, row)
 
 
-def _relative_risk_row_job(i: int, X_bin: np.ndarray, min_count_for_pairwise: int) -> Tuple[int, np.ndarray]:
+def _relative_risk_row_job(
+    i: int, X_bin: np.ndarray, min_count_for_pairwise: int
+) -> Tuple[int, np.ndarray]:
+    """Compute relative-risk values for one row of a pairwise matrix.
+
+    Args:
+        i: Row index.
+        X_bin: Binary matrix.
+        min_count_for_pairwise: Minimum count for valid comparison.
+
+    Returns:
+        Tuple of (row_index, row_values).
+    """
     n_pos = X_bin.shape[1]
     row = np.full((n_pos,), np.nan, dtype=float)
     xi = X_bin[:, i]
@@ -226,8 +273,9 @@ def _relative_risk_row_job(i: int, X_bin: np.ndarray, min_count_for_pairwise: in
             row[j] = np.nan
     return (i, row)
 
+
 def compute_positionwise_statistics(
-    adata,
+    adata: "ad.AnnData",
     layer: str,
     methods: Sequence[str] = ("pearson",),
     sample_col: str = "Barcode",
@@ -238,14 +286,24 @@ def compute_positionwise_statistics(
     min_count_for_pairwise: int = 10,
     max_threads: Optional[int] = None,
     reverse_indices_on_store: bool = False,
-):
+) -> None:
+    """Compute per-(sample, ref) positionwise matrices for selected methods.
+
+    Args:
+        adata: AnnData object to analyze.
+        layer: Layer name to use for statistics.
+        methods: Methods to compute (e.g., ``"pearson"``).
+        sample_col: Obs column containing sample identifiers.
+        ref_col: Obs column containing reference identifiers.
+        site_types: Optional site types to subset positions.
+        encoding: ``"signed"`` or ``"binary"`` encoding.
+        output_key: Key prefix for results stored in ``adata.uns``.
+        min_count_for_pairwise: Minimum counts for pairwise comparisons.
+        max_threads: Maximum number of threads.
+        reverse_indices_on_store: Whether to reverse indices on output storage.
     """
-    Compute per-(sample,ref) positionwise matrices for methods in `methods`.
+    joblib = require("joblib", extra="ml-base", purpose="parallel position statistics")
 
-    Results stored at:
-      adata.uns[output_key][method][ (sample, ref) ] = DataFrame
-      adata.uns[output_key + "_n"][method][ (sample, ref) ] = int(n_reads)
-    """
     if isinstance(methods, str):
         methods = [methods]
     methods = [m.lower() for m in methods]
@@ -256,7 +314,7 @@ def compute_positionwise_statistics(
 
     # workers
     if max_threads is None or max_threads <= 0:
-        n_jobs = max(1, cpu_count() or 1)
+        n_jobs = max(1, joblib.cpu_count() or 1)
     else:
         n_jobs = max(1, int(max_threads))
 
@@ -349,7 +407,10 @@ def compute_positionwise_statistics(
                             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, :])
+                            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":
@@ -366,10 +427,15 @@ def compute_positionwise_statistics(
                         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)
+                        tasks = (
+                            joblib.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)
+                            results = joblib.Parallel(n_jobs=n_jobs, prefer="processes")(tasks)
                         pbar_rows.close()
                         for i, row in results:
                             out[int(i), :] = row
@@ -406,6 +472,7 @@ def compute_positionwise_statistics(
 # Plotting function
 # ---------------------------
 
+
 def plot_positionwise_matrices(
     adata,
     methods: List[str],
@@ -427,9 +494,10 @@ def plot_positionwise_matrices(
     """
     Plot grids of matrices for each method with pagination and rotated sample-row labels.
 
-    New args:
+    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).
     """
@@ -474,7 +542,12 @@ def plot_positionwise_matrices(
         # 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):
+                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]
@@ -486,7 +559,10 @@ def plot_positionwise_matrices(
         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)
+            warnings.warn(
+                f"No results found for method '{method}' in adata.uns['{output_key}']. Skipping.",
+                stacklevel=2,
+            )
             saved_files_by_method[method] = []
             continue
 
@@ -507,21 +583,41 @@ def plot_positionwise_matrices(
 
         # 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 (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 (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:
-            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)
+            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)
+                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
@@ -536,7 +632,9 @@ def plot_positionwise_matrices(
             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)
+            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)
@@ -548,13 +646,24 @@ def plot_positionwise_matrices(
                     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.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)
+                        im = ax.imshow(
+                            mat, origin=origin, aspect="auto", vmin=vmn, vmax=vmx, cmap=cmap
+                        )
                         any_plotted = True
                         ax.set_xticks([])
                         ax.set_yticks([])
@@ -569,9 +678,21 @@ def plot_positionwise_matrices(
                 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.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)
@@ -587,7 +708,7 @@ def plot_positionwise_matrices(
 
             # save or show
             if output_dir:
-                fname = f"positionwise_{method}_page{page_idx+1}.png"
+                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)

smftools/tools/read_stats.py

@@ -1,36 +1,53 @@
 # ------------------------- Utilities -------------------------
-def random_fill_nans(X):
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+
+if TYPE_CHECKING:
+    import anndata as ad
+    import numpy as np
+
+
+def random_fill_nans(X: "np.ndarray") -> "np.ndarray":
+    """Fill NaNs with random values in-place.
+
+    Args:
+        X: Input array with NaNs.
+
+    Returns:
+        numpy.ndarray: Array with NaNs replaced by random values.
+    """
     import numpy as np
+
     nan_mask = np.isnan(X)
     X[nan_mask] = np.random.rand(*X[nan_mask].shape)
     return X
 
+
 def calculate_row_entropy(
-    adata,
-    layer,
-    output_key="entropy",
-    site_config=None,
-    ref_col="Reference_strand",
-    encoding="signed",
-    max_threads=None):
-    """
-    Adds an obs column to the adata that calculates entropy within each read from a given layer
-    when looking at each site type passed in the site_config list.
-
-    Parameters:
-        adata (AnnData): The annotated data matrix.
-        layer (str): Name of the layer to use for entropy calculation.
-        method (str): Unused currently. Placeholder for potential future methods.
-        output_key (str): Base name for the entropy column in adata.obs.
-        site_config (dict): {ref: [site_types]} for masking relevant sites.
-        ref_col (str): Column in adata.obs denoting reference strands.
-        encoding (str): 'signed' (1/-1/0) or 'binary' (1/0/NaN).
-        max_threads (int): Number of threads for parallel processing.
+    adata: "ad.AnnData",
+    layer: str,
+    output_key: str = "entropy",
+    site_config: dict[str, Sequence[str]] | None = None,
+    ref_col: str = "Reference_strand",
+    encoding: str = "signed",
+    max_threads: int | None = None,
+) -> None:
+    """Add per-read entropy values to ``adata.obs``.
+
+    Args:
+        adata: Annotated data matrix.
+        layer: Layer name to use for entropy calculation.
+        output_key: Base name for the entropy column in ``adata.obs``.
+        site_config: Mapping of reference to site types for masking.
+        ref_col: Obs column containing reference strands.
+        encoding: ``"signed"`` (1/-1/0) or ``"binary"`` (1/0/NaN).
+        max_threads: Number of threads for parallel processing.
     """
     import numpy as np
     import pandas as pd
-    from scipy.stats import entropy
     from joblib import Parallel, delayed
+    from scipy.stats import entropy
     from tqdm import tqdm
 
     entropy_values = []
@@ -55,12 +72,14 @@ def calculate_row_entropy(
             X_bin = np.where(X == 1, 1, np.where(X == 0, 0, np.nan))
 
         def compute_entropy(row):
+            """Compute Shannon entropy for a row with NaNs ignored."""
             counts = pd.Series(row).value_counts(dropna=True).sort_index()
             probs = counts / counts.sum()
             return entropy(probs, base=2)
 
         entropies = Parallel(n_jobs=max_threads)(
-            delayed(compute_entropy)(X_bin[i, :]) for i in tqdm(range(X_bin.shape[0]), desc=f"Entropy: {ref}")
+            delayed(compute_entropy)(X_bin[i, :])
+            for i in tqdm(range(X_bin.shape[0]), desc=f"Entropy: {ref}")
         )
 
         entropy_values.extend(entropies)
@@ -69,6 +88,7 @@ def calculate_row_entropy(
     entropy_key = f"{output_key}_entropy"
     adata.obs.loc[row_indices, entropy_key] = entropy_values
 
+
 def binary_autocorrelation_with_spacing(row, positions, max_lag=1000, assume_sorted=True):
     """
     Fast autocorrelation over real genomic spacing.
@@ -125,13 +145,13 @@ def binary_autocorrelation_with_spacing(row, positions, max_lag=1000, assume_sor
             j += 1
         # consider pairs (i, i+1...j-1)
         if j - i > 1:
-            diffs = pos[i+1:j] - pos[i]                 # 1..max_lag
-            contrib = xc[i] * xc[i+1:j]                 # contributions for each pair
+            diffs = pos[i + 1 : j] - pos[i]  # 1..max_lag
+            contrib = xc[i] * xc[i + 1 : j]  # contributions for each pair
             # accumulate weighted sums and counts per lag
-            lag_sums[:max_lag+1] += np.bincount(diffs, weights=contrib,
-                                                minlength=max_lag+1)[:max_lag+1]
-            lag_counts[:max_lag+1] += np.bincount(diffs,
-                                                  minlength=max_lag+1)[:max_lag+1]
+            lag_sums[: max_lag + 1] += np.bincount(diffs, weights=contrib, minlength=max_lag + 1)[
+                : max_lag + 1
+            ]
+            lag_counts[: max_lag + 1] += np.bincount(diffs, minlength=max_lag + 1)[: max_lag + 1]
 
     autocorr = np.full(max_lag + 1, np.nan, dtype=np.float64)
     nz = lag_counts > 0
@@ -140,6 +160,7 @@ def binary_autocorrelation_with_spacing(row, positions, max_lag=1000, assume_sor
 
     return autocorr.astype(np.float32, copy=False)
 
+
 # def binary_autocorrelation_with_spacing(row, positions, max_lag=1000):
 #     """
 #     Compute autocorrelation within a read using real genomic spacing from `positions`.