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

smftools/tools/spatial_autocorrelation.py

@@ -1,48 +1,72 @@
 # ------------------------- Utilities -------------------------
-import pandas as pd
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 import numpy as np
+import pandas as pd
+from numpy.fft import rfft, rfftfreq
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+from smftools.logging_utils import get_logger
+
+logger = get_logger(__name__)
+
+# optional parallel backend
+try:
+    from joblib import Parallel, delayed
+
+    _have_joblib = True
+except Exception:
+    _have_joblib = False
+
+
+# optionally use scipy for find_peaks (more robust)
+try:
+    from scipy.signal import find_peaks
+
+    _have_scipy = True
+except Exception:
+    _have_scipy = False
 
-def random_fill_nans(X):
+
+def random_fill_nans(X: NDArray[np.floating]) -> NDArray[np.floating]:
+    """Fill NaNs with random values in-place.
+
+    Args:
+        X: Input array containing NaNs.
+
+    Returns:
+        numpy.ndarray: Array with NaNs replaced by random values.
+    """
     nan_mask = np.isnan(X)
     X[nan_mask] = np.random.rand(*X[nan_mask].shape)
     return X
 
+
 def binary_autocorrelation_with_spacing(
-    row, 
-    positions, 
-    max_lag=1000, 
-    assume_sorted=True,
-    normalize: str = "sum", 
-    return_counts: bool = False
+    row: NDArray[np.floating],
+    positions: NDArray[np.integer],
+    max_lag: int = 1000,
+    assume_sorted: bool = True,
+    normalize: str = "sum",
+    return_counts: bool = False,
 ):
-    """
-    Fast autocorrelation over real genomic spacing.
-
-    Parameters
-    ----------
-    row : 1D array (float)
-        Values per position (NaN = missing). Works for binary or real-valued.
-    positions : 1D array (int)
-        Genomic coordinates for each column of `row`.
-    max_lag : int
-        Max genomic lag (inclusive).
-    assume_sorted : bool
-        If True, assumes `positions` are strictly non-decreasing.
-    normalize : {"sum", "pearson"}
-        "sum": autocorr[l] = sum_{pairs at lag l} (xc_i * xc_j) / sum(xc^2)
-               (fast; comparable across lags and molecules).
-        "pearson": autocorr[l] = (mean_{pairs at lag l} (xc_i * xc_j)) / (mean(xc^2))
-                   i.e., an estimate of Pearson-like correlation at that lag.
-    return_counts : bool
-        If True, return (autocorr, lag_counts). Otherwise just autocorr.
-
-    Returns
-    -------
-    autocorr : 1D array, shape (max_lag+1,)
-        Normalized autocorrelation; autocorr[0] = 1.0.
-        Lags with no valid pairs are NaN.
-    (optionally) lag_counts : 1D array, shape (max_lag+1,)
-        Number of pairs contributing to each lag.
+    """Compute autocorrelation over genomic spacing.
+
+    Args:
+        row: Values per position (NaN = missing).
+        positions: Genomic coordinates for each column of ``row``.
+        max_lag: Max genomic lag (inclusive).
+        assume_sorted: Whether ``positions`` are sorted.
+        normalize: ``"sum"`` or ``"pearson"`` normalization.
+        return_counts: Whether to return lag counts alongside autocorrelation.
+
+    Returns:
+        numpy.ndarray | tuple[numpy.ndarray, numpy.ndarray]: Autocorrelation values and
+        optionally counts per lag.
     """
 
     # mask valid entries
@@ -82,12 +106,12 @@ def binary_autocorrelation_with_spacing(
             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
             # bincount returns length >= max(diffs)+1; we request minlength
-            bc_vals = np.bincount(diffs, weights=contrib, minlength=max_lag+1)[:max_lag+1]
-            bc_counts = np.bincount(diffs, minlength=max_lag+1)[:max_lag+1]
+            bc_vals = np.bincount(diffs, weights=contrib, minlength=max_lag + 1)[: max_lag + 1]
+            bc_counts = np.bincount(diffs, minlength=max_lag + 1)[: max_lag + 1]
             lag_sums += bc_vals
             lag_counts += bc_counts
 
@@ -113,20 +137,17 @@ def binary_autocorrelation_with_spacing(
     return autocorr.astype(np.float32, copy=False)
 
 
-from numpy.fft import rfft, rfftfreq
-
-# optionally use scipy for find_peaks (more robust)
-try:
-    from scipy.signal import find_peaks
-    _have_scipy = True
-except Exception:
-    _have_scipy = False
-
 # ---------- helpers ----------
 def weighted_mean_autocorr(ac_matrix, counts_matrix, min_count=20):
-    """
-    Weighted mean across molecules: sum(ac * counts) / sum(counts) per lag.
-    Mask lags with total counts < min_count (set NaN).
+    """Compute weighted mean autocorrelation per lag.
+
+    Args:
+        ac_matrix: Autocorrelation matrix per molecule.
+        counts_matrix: Pair counts per lag.
+        min_count: Minimum total count required to keep a lag.
+
+    Returns:
+        tuple[numpy.ndarray, numpy.ndarray]: Mean autocorrelation and total counts.
     """
     counts_total = counts_matrix.sum(axis=0)
     # replace NaNs in ac_matrix with 0 for weighted sum
@@ -138,7 +159,22 @@ def weighted_mean_autocorr(ac_matrix, counts_matrix, min_count=20):
     mean_ac[counts_total < min_count] = np.nan
     return mean_ac, counts_total
 
-def psd_from_autocorr(mean_ac, lags, pad_factor=4):
+
+def psd_from_autocorr(
+    mean_ac: NDArray[np.floating],
+    lags: NDArray[np.floating],
+    pad_factor: int = 4,
+) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
+    """Compute a power spectral density from autocorrelation.
+
+    Args:
+        mean_ac: Mean autocorrelation values.
+        lags: Lag values in base pairs.
+        pad_factor: Padding factor for FFT resolution.
+
+    Returns:
+        tuple[numpy.ndarray, numpy.ndarray]: Frequencies and power values.
+    """
     n = len(mean_ac)
     pad_n = int(max(2**10, pad_factor * n))  # pad to at least some min to stabilize FFT res
     ac_padded = np.zeros(pad_n, dtype=np.float64)
@@ -149,7 +185,24 @@ def psd_from_autocorr(mean_ac, lags, pad_factor=4):
     freqs = rfftfreq(pad_n, d=df)
     return freqs, power
 
-def find_peak_in_nrl_band(freqs, power, nrl_search_bp=(120,260), prominence_frac=0.05):
+
+def find_peak_in_nrl_band(
+    freqs: NDArray[np.floating],
+    power: NDArray[np.floating],
+    nrl_search_bp: tuple[int, int] = (120, 260),
+    prominence_frac: float = 0.05,
+) -> tuple[float | None, int | None]:
+    """Find the peak frequency in the nucleosome repeat length band.
+
+    Args:
+        freqs: Frequency bins.
+        power: Power values.
+        nrl_search_bp: Search band in base pairs.
+        prominence_frac: Fraction of peak power for prominence.
+
+    Returns:
+        tuple[float | None, int | None]: Peak frequency and index, or ``(None, None)``.
+    """
     fmin = 1.0 / nrl_search_bp[1]
     fmax = 1.0 / nrl_search_bp[0]
     band_mask = (freqs >= fmin) & (freqs <= fmax)
@@ -170,7 +223,22 @@ def find_peak_in_nrl_band(freqs, power, nrl_search_bp=(120,260), prominence_frac
     idx = band_indices[rel]
     return freqs[idx], idx
 
-def fwhm_freq_to_bp(freqs, power, peak_idx):
+
+def fwhm_freq_to_bp(
+    freqs: NDArray[np.floating],
+    power: NDArray[np.floating],
+    peak_idx: int,
+) -> tuple[float, float, float]:
+    """Estimate FWHM in base pairs for a spectral peak.
+
+    Args:
+        freqs: Frequency bins.
+        power: Power values.
+        peak_idx: Index of the peak.
+
+    Returns:
+        tuple[float, float, float]: FWHM in bp and left/right frequencies.
+    """
     # find half power
     pk = power[peak_idx]
     half = pk / 2.0
@@ -182,39 +250,71 @@ def fwhm_freq_to_bp(freqs, power, peak_idx):
     if left == peak_idx:
         left_f = freqs[peak_idx]
     else:
-        x0, x1 = freqs[left], freqs[left+1]
-        y0, y1 = power[left], power[left+1]
-        left_f = x0 if y1 == y0 else x0 + (half - y0)*(x1-x0)/(y1-y0)
+        x0, x1 = freqs[left], freqs[left + 1]
+        y0, y1 = power[left], power[left + 1]
+        left_f = x0 if y1 == y0 else x0 + (half - y0) * (x1 - x0) / (y1 - y0)
     # move right
     right = peak_idx
-    while right < len(power)-1 and power[right] > half:
+    while right < len(power) - 1 and power[right] > half:
         right += 1
     if right == peak_idx:
         right_f = freqs[peak_idx]
     else:
-        x0, x1 = freqs[right-1], freqs[right]
-        y0, y1 = power[right-1], power[right]
-        right_f = x1 if y1 == y0 else x0 + (half - y0)*(x1-x0)/(y1-y0)
+        x0, x1 = freqs[right - 1], freqs[right]
+        y0, y1 = power[right - 1], power[right]
+        right_f = x1 if y1 == y0 else x0 + (half - y0) * (x1 - x0) / (y1 - y0)
     # convert to bp approximating delta_NRL = |1/left_f - 1/right_f|
     left_NRL = 1.0 / right_f if right_f > 0 else np.nan
     right_NRL = 1.0 / left_f if left_f > 0 else np.nan
     fwhm_bp = abs(left_NRL - right_NRL)
     return fwhm_bp, left_f, right_f
 
-def estimate_snr(power, peak_idx, exclude_bins=5):
+
+def estimate_snr(
+    power: NDArray[np.floating],
+    peak_idx: int,
+    exclude_bins: int = 5,
+) -> tuple[float, float, float]:
+    """Estimate signal-to-noise ratio around a spectral peak.
+
+    Args:
+        power: Power values.
+        peak_idx: Index of the peak.
+        exclude_bins: Bins to exclude around the peak when estimating background.
+
+    Returns:
+        tuple[float, float, float]: SNR, peak power, and background median.
+    """
     pk = power[peak_idx]
     mask = np.ones_like(power, dtype=bool)
-    lo = max(0, peak_idx-exclude_bins)
-    hi = min(len(power), peak_idx+exclude_bins+1)
+    lo = max(0, peak_idx - exclude_bins)
+    hi = min(len(power), peak_idx + exclude_bins + 1)
     mask[lo:hi] = False
     bg = power[mask]
     bg_med = np.median(bg) if bg.size else np.median(power)
     return pk / (bg_med if bg_med > 0 else np.finfo(float).eps), pk, bg_med
 
-def sample_autocorr_at_harmonics(mean_ac, lags, nrl_bp, max_harmonics=6):
+
+def sample_autocorr_at_harmonics(
+    mean_ac: NDArray[np.floating],
+    lags: NDArray[np.floating],
+    nrl_bp: float,
+    max_harmonics: int = 6,
+) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
+    """Sample autocorrelation heights at NRL harmonics.
+
+    Args:
+        mean_ac: Mean autocorrelation values.
+        lags: Lag values in base pairs.
+        nrl_bp: NRL in base pairs.
+        max_harmonics: Maximum harmonics to sample.
+
+    Returns:
+        tuple[numpy.ndarray, numpy.ndarray]: Sampled lags and heights.
+    """
     sample_lags = []
     heights = []
-    for m in range(1, max_harmonics+1):
+    for m in range(1, max_harmonics + 1):
         target = m * nrl_bp
         # stop if beyond observed lag range
         if target > lags[-1]:
@@ -227,7 +327,22 @@ def sample_autocorr_at_harmonics(mean_ac, lags, nrl_bp, max_harmonics=6):
         heights.append(h)
     return np.array(sample_lags), np.array(heights)
 
-def fit_exponential_envelope(sample_lags, heights, counts=None):
+
+def fit_exponential_envelope(
+    sample_lags: NDArray[np.floating],
+    heights: NDArray[np.floating],
+    counts: NDArray[np.floating] | None = None,
+) -> tuple[float, float, float, float]:
+    """Fit an exponential envelope to sampled autocorrelation peaks.
+
+    Args:
+        sample_lags: Sampled lag values.
+        heights: Sampled autocorrelation heights.
+        counts: Optional weights per sample.
+
+    Returns:
+        tuple[float, float, float, float]: ``(xi, A, slope, r2)``.
+    """
     # heights ~ A * exp(-lag / xi)
     mask = (heights > 0) & np.isfinite(heights)
     if mask.sum() < 2:
@@ -238,7 +353,7 @@ def fit_exponential_envelope(sample_lags, heights, counts=None):
         w = np.ones_like(y)
     else:
         w = np.asarray(counts[mask], dtype=float)
-        w = w / (np.max(w) if np.max(w)>0 else 1.0)
+        w = w / (np.max(w) if np.max(w) > 0 else 1.0)
     # weighted linear regression y = b0 + b1 * x
     X = np.vstack([np.ones_like(x), x]).T
     W = np.diag(w)
@@ -253,122 +368,147 @@ def fit_exponential_envelope(sample_lags, heights, counts=None):
     xi = -1.0 / b1 if b1 < 0 else np.nan
     # R^2
     y_pred = X.dot(b)
-    ss_res = np.sum(w * (y - y_pred)**2)
-    ss_tot = np.sum(w * (y - np.average(y, weights=w))**2)
-    r2 = 1.0 - ss_res/ss_tot if ss_tot != 0 else np.nan
+    ss_res = np.sum(w * (y - y_pred) ** 2)
+    ss_tot = np.sum(w * (y - np.average(y, weights=w)) ** 2)
+    r2 = 1.0 - ss_res / ss_tot if ss_tot != 0 else np.nan
     return xi, A, b1, r2
 
+
 # ---------- main analysis per site_type ----------
-def analyze_autocorr_matrix(autocorr_matrix, counts_matrix, lags,
-                            nrl_search_bp=(120,260), pad_factor=4,
-                            min_count=20, max_harmonics=6):
-    """
-    Return dict: nrl_bp, peak_power, fwhm_bp, snr, xi, envelope points, freqs, power, mean_ac
+def analyze_autocorr_matrix(
+    autocorr_matrix: NDArray[np.floating],
+    counts_matrix: NDArray[np.integer],
+    lags: NDArray[np.floating],
+    nrl_search_bp: tuple[int, int] = (120, 260),
+    pad_factor: int = 4,
+    min_count: int = 20,
+    max_harmonics: int = 6,
+):
+    """Analyze autocorrelation matrix and extract periodicity metrics.
+
+    Args:
+        autocorr_matrix: Autocorrelation values per molecule.
+        counts_matrix: Pair counts per lag.
+        lags: Lag values in base pairs.
+        nrl_search_bp: NRL search band in base pairs.
+        pad_factor: Padding factor for FFT.
+        min_count: Minimum total count to retain a lag.
+        max_harmonics: Maximum harmonics to sample.
+
+    Returns:
+        dict: Metrics including NRL, SNR, and PSD summaries.
     """
-    mean_ac, counts_total = weighted_mean_autocorr(autocorr_matrix, counts_matrix, min_count=min_count)
+    mean_ac, counts_total = weighted_mean_autocorr(
+        autocorr_matrix, counts_matrix, min_count=min_count
+    )
     freqs, power = psd_from_autocorr(mean_ac, lags, pad_factor=pad_factor)
     f0, peak_idx = find_peak_in_nrl_band(freqs, power, nrl_search_bp=nrl_search_bp)
     if f0 is None:
-        return {"error":"no_peak_found", "mean_ac":mean_ac, "counts":counts_total}
+        return {"error": "no_peak_found", "mean_ac": mean_ac, "counts": counts_total}
     nrl_bp = 1.0 / f0
     fwhm_bp, left_f, right_f = fwhm_freq_to_bp(freqs, power, peak_idx)
     snr, peak_power, bg = estimate_snr(power, peak_idx)
-    sample_lags, heights = sample_autocorr_at_harmonics(mean_ac, lags, nrl_bp, max_harmonics=max_harmonics)
-    xi, A, slope, r2 = fit_exponential_envelope(sample_lags, heights) if heights.size else (np.nan,)*4
+    sample_lags, heights = sample_autocorr_at_harmonics(
+        mean_ac, lags, nrl_bp, max_harmonics=max_harmonics
+    )
+    xi, A, slope, r2 = (
+        fit_exponential_envelope(sample_lags, heights) if heights.size else (np.nan,) * 4
+    )
 
     return dict(
-        nrl_bp = nrl_bp,
-        f0 = f0,
-        peak_power = peak_power,
-        fwhm_bp = fwhm_bp,
-        snr = snr,
-        bg_median = bg,
-        envelope_sample_lags = sample_lags,
-        envelope_heights = heights,
-        xi = xi,
-        xi_A = A,
-        xi_slope = slope,
-        xi_r2 = r2,
-        freqs = freqs,
-        power = power,
-        mean_ac = mean_ac,
-        counts = counts_total
+        nrl_bp=nrl_bp,
+        f0=f0,
+        peak_power=peak_power,
+        fwhm_bp=fwhm_bp,
+        snr=snr,
+        bg_median=bg,
+        envelope_sample_lags=sample_lags,
+        envelope_heights=heights,
+        xi=xi,
+        xi_A=A,
+        xi_slope=slope,
+        xi_r2=r2,
+        freqs=freqs,
+        power=power,
+        mean_ac=mean_ac,
+        counts=counts_total,
     )
 
+
 # ---------- bootstrap wrapper ----------
-def bootstrap_periodicity(autocorr_matrix, counts_matrix, lags, n_boot=200, **kwargs):
+def bootstrap_periodicity(
+    autocorr_matrix: NDArray[np.floating],
+    counts_matrix: NDArray[np.integer],
+    lags: NDArray[np.floating],
+    n_boot: int = 200,
+    **kwargs,
+) -> dict:
+    """Bootstrap periodicity metrics from autocorrelation matrices.
+
+    Args:
+        autocorr_matrix: Autocorrelation matrix per molecule.
+        counts_matrix: Pair counts per lag.
+        lags: Lag values in base pairs.
+        n_boot: Number of bootstrap samples.
+        **kwargs: Additional arguments for ``analyze_autocorr_matrix``.
+
+    Returns:
+        dict: Bootstrapped metric arrays and per-iteration metrics.
+    """
     rng = np.random.default_rng()
     metrics = []
     n = autocorr_matrix.shape[0]
     for _ in range(n_boot):
         sample_idx = rng.integers(0, n, size=n)
-        res = analyze_autocorr_matrix(autocorr_matrix[sample_idx], counts_matrix[sample_idx], lags, **kwargs)
+        res = analyze_autocorr_matrix(
+            autocorr_matrix[sample_idx], counts_matrix[sample_idx], lags, **kwargs
+        )
         metrics.append(res)
     # extract key fields robustly
     nrls = np.array([m.get("nrl_bp", np.nan) for m in metrics])
-    xis  = np.array([m.get("xi", np.nan) for m in metrics])
-    return {"nrl_boot":nrls, "xi_boot":xis, "metrics":metrics}
-
+    xis = np.array([m.get("xi", np.nan) for m in metrics])
+    return {"nrl_boot": nrls, "xi_boot": xis, "metrics": metrics}
 
-# optional parallel backend
-try:
-    from joblib import Parallel, delayed
-    _have_joblib = True
-except Exception:
-    _have_joblib = False
 
 def rolling_autocorr_metrics(
-    X,
-    positions,
-    site_label: str = None,
+    X: NDArray[np.floating],
+    positions: NDArray[np.integer],
+    site_label: str | None = None,
     window_size: int = 2000,
     step: int = 500,
     max_lag: int = 800,
     min_molecules_per_window: int = 10,
-    nrl_search_bp: tuple = (120, 260),
+    nrl_search_bp: tuple[int, int] = (120, 260),
     pad_factor: int = 4,
     min_count_for_mean: int = 20,
     max_harmonics: int = 6,
     n_jobs: int = 1,
     verbose: bool = False,
     return_window_results: bool = False,
-    fixed_nrl_bp: float = None,
-
+    fixed_nrl_bp: float | None = None,
 ):
-    """
-    Slide a genomic window across `positions` and compute periodicity metrics per window.
-
-    Parameters
-    ----------
-    X : array-like or sparse, shape (n_molecules, n_positions)
-        Binary site matrix for a group (sample × reference × site_type).
-    positions : 1D array-like of ints
-        Genomic coordinates for columns of X (same length as X.shape[1]).
-    site_label : optional str
-        Label for the site type (used in returned dicts/df).
-    window_size : int
-        Window width in bp.
-    step : int
-        Slide step in bp.
-    max_lag : int
-        Max lag (bp) to compute autocorr out to.
-    min_molecules_per_window : int
-        Minimum molecules required to compute metrics for a window; otherwise metrics = NaN.
-    nrl_search_bp, pad_factor, min_count_for_mean, max_harmonics : forwarded to analyze_autocorr_matrix
-    n_jobs : int
-        Number of parallel jobs (uses joblib if available).
-    verbose : bool
-        Print progress messages.
-    return_window_results : bool
-        If True, return also the per-window raw `analyze_autocorr_matrix` outputs.
-
-    Returns
-    -------
-    df : pandas.DataFrame
-        One row per window with columns:
-          ['site', 'window_start', 'window_end', 'center', 'n_molecules',
-           'nrl_bp', 'snr', 'peak_power', 'fwhm_bp', 'xi', 'xi_A', 'xi_r2']
-    (optionally) window_results : list of dicts (same order as df rows) when return_window_results=True
+    """Slide a genomic window across positions and compute periodicity metrics.
+
+    Args:
+        X: Binary site matrix for a group (sample × reference × site_type).
+        positions: Genomic coordinates for columns of ``X``.
+        site_label: Label for the site type.
+        window_size: Window width in bp.
+        step: Slide step in bp.
+        max_lag: Max lag (bp) to compute autocorr out to.
+        min_molecules_per_window: Minimum molecules required per window.
+        nrl_search_bp: NRL search band in base pairs.
+        pad_factor: Padding factor for FFT.
+        min_count_for_mean: Minimum count for mean autocorrelation.
+        max_harmonics: Maximum harmonics to sample.
+        n_jobs: Number of parallel jobs (joblib if available).
+        verbose: Whether to log progress.
+        return_window_results: Whether to return per-window analyzer outputs.
+        fixed_nrl_bp: If provided, use a fixed NRL in bp for analysis.
+
+    Returns:
+        pandas.DataFrame | tuple[pandas.DataFrame, list[dict]]: Window-level metrics,
+        with optional raw analyzer outputs.
     """
 
     # normalize inputs
@@ -386,10 +526,16 @@ def rolling_autocorr_metrics(
         window_starts = list(range(start, end - window_size + 1, step))
 
     if verbose:
-        print(f"Rolling windows: {len(window_starts)} windows, window_size={window_size}, step={step}")
+        logger.info(
+            "Rolling windows: %s windows, window_size=%s, step=%s",
+            len(window_starts),
+            window_size,
+            step,
+        )
 
     # helper to extract row to dense 1D np array (supports sparse rows)
     def _row_to_arr(row):
+        """Convert a matrix row (dense or sparse) to a 1D NumPy array."""
         # handle scipy sparse row
         try:
             import scipy.sparse as sp
@@ -402,6 +548,7 @@ def rolling_autocorr_metrics(
 
     # function to process one window
     def _process_window(ws):
+        """Compute rolling-window autocorrelation metrics for a window start."""
         we = ws + window_size
         mask_pos = (pos >= ws) & (pos < we)
         if mask_pos.sum() < 2:
@@ -428,7 +575,9 @@ def rolling_autocorr_metrics(
                 continue
             # compute autocorr on the windowed template; positions are pos[mask_pos]
             try:
-                ac, cnts = binary_autocorrelation_with_spacing(subrow, pos[mask_pos], max_lag=max_lag, assume_sorted=True, return_counts=True)
+                ac, cnts = binary_autocorrelation_with_spacing(
+                    subrow, pos[mask_pos], max_lag=max_lag, assume_sorted=True, return_counts=True
+                )
             except Exception:
                 # if autocorr fails for this row, skip it
                 continue
@@ -460,7 +609,9 @@ def rolling_autocorr_metrics(
 
         # If a fixed global NRL is provided, compute metrics around that frequency
         if fixed_nrl_bp is not None:
-            freqs, power = psd_from_autocorr(mean_ac, np.arange(mean_ac.size), pad_factor=pad_factor)
+            freqs, power = psd_from_autocorr(
+                mean_ac, np.arange(mean_ac.size), pad_factor=pad_factor
+            )
             # locate nearest freq bin to target_freq
             target_f = 1.0 / float(fixed_nrl_bp)
             # mask valid freqs
@@ -477,28 +628,44 @@ def rolling_autocorr_metrics(
                 snr_val, _, bg = estimate_snr(power, peak_idx, exclude_bins=3)
                 # sample harmonics from mean_ac at integer-lag positions using fixed_nrl_bp
                 # note: lags array is integer 0..(mean_ac.size-1)
-                sample_lags, heights = sample_autocorr_at_harmonics(mean_ac, np.arange(mean_ac.size), fixed_nrl_bp, max_harmonics=max_harmonics)
-                xi, A, slope, r2 = fit_exponential_envelope(sample_lags, heights) if heights.size else (np.nan, np.nan, np.nan, np.nan)
+                sample_lags, heights = sample_autocorr_at_harmonics(
+                    mean_ac, np.arange(mean_ac.size), fixed_nrl_bp, max_harmonics=max_harmonics
+                )
+                xi, A, slope, r2 = (
+                    fit_exponential_envelope(sample_lags, heights)
+                    if heights.size
+                    else (np.nan, np.nan, np.nan, np.nan)
+                )
                 res = dict(
                     nrl_bp=float(fixed_nrl_bp),
                     f0=float(target_f),
                     peak_power=peak_power,
-                    fwhm_bp=np.nan,           # not robustly defined when using fixed freq (skip or compute small-band FWHM)
+                    fwhm_bp=np.nan,  # not robustly defined when using fixed freq (skip or compute small-band FWHM)
                     snr=float(snr_val),
                     bg_median=float(bg) if np.isfinite(bg) else np.nan,
                     envelope_sample_lags=sample_lags,
                     envelope_heights=heights,
-                    xi=xi, xi_A=A, xi_slope=slope, xi_r2=r2,
-                    freqs=freqs, power=power, mean_ac=mean_ac, counts=counts_total
+                    xi=xi,
+                    xi_A=A,
+                    xi_slope=slope,
+                    xi_r2=r2,
+                    freqs=freqs,
+                    power=power,
+                    mean_ac=mean_ac,
+                    counts=counts_total,
                 )
         else:
             # existing behavior: call analyzer_fn
             try:
-                res = analyze_autocorr_matrix(ac_mat, cnt_mat, np.arange(mean_ac.size),
-                                nrl_search_bp=nrl_search_bp,
-                                pad_factor=pad_factor,
-                                min_count=min_count_for_mean,
-                                max_harmonics=max_harmonics)
+                res = analyze_autocorr_matrix(
+                    ac_mat,
+                    cnt_mat,
+                    np.arange(mean_ac.size),
+                    nrl_search_bp=nrl_search_bp,
+                    pad_factor=pad_factor,
+                    min_count=min_count_for_mean,
+                    max_harmonics=max_harmonics,
+                )
             except Exception as e:
                 res = {"error": str(e)}
 
@@ -524,39 +691,45 @@ def rolling_autocorr_metrics(
         metrics = r["metrics"]
         window_results.append(metrics)
         if metrics is None or ("error" in metrics and metrics.get("error") == "no_peak_found"):
-            rows_out.append({
-                "site": r["site"],
-                "window_start": r["window_start"],
-                "window_end": r["window_end"],
-                "center": r["center"],
-                "n_molecules": r["n_molecules"],
-                "nrl_bp": np.nan,
-                "snr": np.nan,
-                "peak_power": np.nan,
-                "fwhm_bp": np.nan,
-                "xi": np.nan,
-                "xi_A": np.nan,
-                "xi_r2": np.nan,
-                "analyzer_error": (metrics.get("error") if isinstance(metrics, dict) else "no_metrics"),
-            })
+            rows_out.append(
+                {
+                    "site": r["site"],
+                    "window_start": r["window_start"],
+                    "window_end": r["window_end"],
+                    "center": r["center"],
+                    "n_molecules": r["n_molecules"],
+                    "nrl_bp": np.nan,
+                    "snr": np.nan,
+                    "peak_power": np.nan,
+                    "fwhm_bp": np.nan,
+                    "xi": np.nan,
+                    "xi_A": np.nan,
+                    "xi_r2": np.nan,
+                    "analyzer_error": (
+                        metrics.get("error") if isinstance(metrics, dict) else "no_metrics"
+                    ),
+                }
+            )
         else:
-            rows_out.append({
-                "site": r["site"],
-                "window_start": r["window_start"],
-                "window_end": r["window_end"],
-                "center": r["center"],
-                "n_molecules": r["n_molecules"],
-                "nrl_bp": float(metrics.get("nrl_bp", np.nan)),
-                "snr": float(metrics.get("snr", np.nan)),
-                "peak_power": float(metrics.get("peak_power", np.nan)),
-                "fwhm_bp": float(metrics.get("fwhm_bp", np.nan)),
-                "xi": float(metrics.get("xi", np.nan)),
-                "xi_A": float(metrics.get("xi_A", np.nan)),
-                "xi_r2": float(metrics.get("xi_r2", np.nan)),
-                "analyzer_error": None,
-            })
+            rows_out.append(
+                {
+                    "site": r["site"],
+                    "window_start": r["window_start"],
+                    "window_end": r["window_end"],
+                    "center": r["center"],
+                    "n_molecules": r["n_molecules"],
+                    "nrl_bp": float(metrics.get("nrl_bp", np.nan)),
+                    "snr": float(metrics.get("snr", np.nan)),
+                    "peak_power": float(metrics.get("peak_power", np.nan)),
+                    "fwhm_bp": float(metrics.get("fwhm_bp", np.nan)),
+                    "xi": float(metrics.get("xi", np.nan)),
+                    "xi_A": float(metrics.get("xi_A", np.nan)),
+                    "xi_r2": float(metrics.get("xi_r2", np.nan)),
+                    "analyzer_error": None,
+                }
+            )
 
     df = pd.DataFrame(rows_out)
     if return_window_results:
         return df, window_results
-    return df
+    return df