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

smftools/plotting/hmm_plotting.py

@@ -1,8 +1,20 @@
+from __future__ import annotations
+
 import math
-from typing import List, Optional, Tuple, Union
+from typing import Optional, Tuple, Union
+
 import numpy as np
-import matplotlib.pyplot as plt
-from matplotlib.backends.backend_pdf import PdfPages
+
+from smftools.optional_imports import require
+
+plt = require("matplotlib.pyplot", extra="plotting", purpose="HMM plots")
+pdf_backend = require(
+    "matplotlib.backends.backend_pdf",
+    extra="plotting",
+    purpose="PDF output",
+)
+PdfPages = pdf_backend.PdfPages
+
 
 def plot_hmm_size_contours(
     adata,
@@ -36,32 +48,41 @@ def plot_hmm_size_contours(
 
     Other args are the same as prior function.
     """
+
     # --- helper: gaussian smoothing (scipy fallback -> numpy separable conv) ---
     def _gaussian_1d_kernel(sigma: float, eps: float = 1e-12):
+        """Build a normalized 1D Gaussian kernel."""
         if sigma <= 0 or sigma is None:
             return np.array([1.0], dtype=float)
         # choose kernel size = odd ~ 6*sigma (covers +/-3 sigma)
         radius = max(1, int(math.ceil(3.0 * float(sigma))))
         xs = np.arange(-radius, radius + 1, dtype=float)
-        k = np.exp(-(xs ** 2) / (2.0 * sigma ** 2))
+        k = np.exp(-(xs**2) / (2.0 * sigma**2))
         k_sum = k.sum()
         if k_sum <= eps:
             k = np.array([1.0], dtype=float)
             k_sum = 1.0
         return k / k_sum
 
-    def _smooth_with_numpy_separable(Z: np.ndarray, sigma_len: float, sigma_pos: float) -> np.ndarray:
+    def _smooth_with_numpy_separable(
+        Z: np.ndarray, sigma_len: float, sigma_pos: float
+    ) -> np.ndarray:
+        """Apply separable Gaussian smoothing with NumPy."""
         # Z shape: (n_lengths, n_positions)
         out = Z.copy()
         # smooth along length axis (axis=0)
         if sigma_len and sigma_len > 0:
             k_len = _gaussian_1d_kernel(sigma_len)
             # convolve each column
-            out = np.apply_along_axis(lambda col: np.convolve(col, k_len, mode="same"), axis=0, arr=out)
+            out = np.apply_along_axis(
+                lambda col: np.convolve(col, k_len, mode="same"), axis=0, arr=out
+            )
         # smooth along position axis (axis=1)
         if sigma_pos and sigma_pos > 0:
             k_pos = _gaussian_1d_kernel(sigma_pos)
-            out = np.apply_along_axis(lambda row: np.convolve(row, k_pos, mode="same"), axis=1, arr=out)
+            out = np.apply_along_axis(
+                lambda row: np.convolve(row, k_pos, mode="same"), axis=1, arr=out
+            )
         return out
 
     # prefer scipy.ndimage if available (faster and better boundary handling)
@@ -69,11 +90,13 @@ def plot_hmm_size_contours(
     if use_scipy_if_available:
         try:
             from scipy.ndimage import gaussian_filter as _scipy_gaussian_filter
+
             _have_scipy = True
         except Exception:
             _have_scipy = False
 
     def _smooth_Z(Z: np.ndarray, sigma_len: float, sigma_pos: float) -> np.ndarray:
+        """Smooth a matrix using scipy if available or NumPy fallback."""
         if (sigma_len is None or sigma_len == 0) and (sigma_pos is None or sigma_pos == 0):
             return Z
         if _have_scipy:
@@ -84,8 +107,16 @@ def plot_hmm_size_contours(
             return _smooth_with_numpy_separable(Z, float(sigma_len or 0.0), float(sigma_pos or 0.0))
 
     # --- gather unique ordered labels ---
-    samples = list(adata.obs[sample_col].cat.categories) if getattr(adata.obs[sample_col], "dtype", None) == "category" else list(pd.Categorical(adata.obs[sample_col]).categories)
-    refs = list(adata.obs[ref_obs_col].cat.categories) if getattr(adata.obs[ref_obs_col], "dtype", None) == "category" else list(pd.Categorical(adata.obs[ref_obs_col]).categories)
+    samples = (
+        list(adata.obs[sample_col].cat.categories)
+        if getattr(adata.obs[sample_col], "dtype", None) == "category"
+        else list(pd.Categorical(adata.obs[sample_col]).categories)
+    )
+    refs = (
+        list(adata.obs[ref_obs_col].cat.categories)
+        if getattr(adata.obs[ref_obs_col], "dtype", None) == "category"
+        else list(pd.Categorical(adata.obs[ref_obs_col]).categories)
+    )
 
     n_samples = len(samples)
     n_refs = len(refs)
@@ -102,6 +133,7 @@ def plot_hmm_size_contours(
 
     # helper to get dense layer array for subset
     def _get_layer_array(layer):
+        """Convert a layer to a dense NumPy array."""
         arr = layer
         # sparse -> toarray
         if hasattr(arr, "toarray"):
@@ -146,7 +178,7 @@ def plot_hmm_size_contours(
         fig_w = n_refs * figsize_per_cell[0]
         fig_h = rows_on_page * figsize_per_cell[1]
         fig, axes = plt.subplots(rows_on_page, n_refs, figsize=(fig_w, fig_h), squeeze=False)
-        fig.suptitle(f"HMM size contours (page {p+1}/{pages})", fontsize=12)
+        fig.suptitle(f"HMM size contours (page {p + 1}/{pages})", fontsize=12)
 
         # for each panel compute p(length | position)
         for i_row, sample in enumerate(page_samples):
@@ -160,7 +192,9 @@ def plot_hmm_size_contours(
                     ax.set_title(f"{sample} / {ref}")
                     continue
 
-                row_idx = np.nonzero(panel_mask.values if hasattr(panel_mask, "values") else np.asarray(panel_mask))[0]
+                row_idx = np.nonzero(
+                    panel_mask.values if hasattr(panel_mask, "values") else np.asarray(panel_mask)
+                )[0]
                 if row_idx.size == 0:
                     ax.text(0.5, 0.5, "no reads", ha="center", va="center")
                     ax.set_title(f"{sample} / {ref}")
@@ -178,7 +212,9 @@ def plot_hmm_size_contours(
                 max_len_here = min(max_len, max_len_local)
 
                 lengths_range = np.arange(1, max_len_here + 1, dtype=int)
-                Z = np.zeros((len(lengths_range), n_positions), dtype=float)  # rows=length, cols=pos
+                Z = np.zeros(
+                    (len(lengths_range), n_positions), dtype=float
+                )  # rows=length, cols=pos
 
                 # fill Z by efficient bincount across columns
                 for j in range(n_positions):
@@ -222,7 +258,9 @@ def plot_hmm_size_contours(
                 dy = 1.0
                 y_edges = np.concatenate([y - 0.5, [y[-1] + 0.5]])
 
-                pcm = ax.pcolormesh(x_edges, y_edges, Z_plot, cmap=cmap, shading="auto", vmin=vmin, vmax=vmax)
+                pcm = ax.pcolormesh(
+                    x_edges, y_edges, Z_plot, cmap=cmap, shading="auto", vmin=vmin, vmax=vmax
+                )
                 ax.set_title(f"{sample} / {ref}")
                 ax.set_ylabel("length")
                 if i_row == rows_on_page - 1:
@@ -243,9 +281,10 @@ def plot_hmm_size_contours(
         # saving per page if requested
         if save_path is not None:
             import os
+
             os.makedirs(save_path, exist_ok=True)
             if save_each_page:
-                fname = f"hmm_size_page_{p+1:03d}.png"
+                fname = f"hmm_size_page_{p + 1:03d}.png"
                 out = os.path.join(save_path, fname)
                 fig.savefig(out, dpi=dpi, bbox_inches="tight")
 

smftools/plotting/position_stats.py

@@ -1,3 +1,8 @@
+from __future__ import annotations
+
+from smftools.optional_imports import require
+
+
 def plot_volcano_relative_risk(
     results_dict,
     save_path=None,
@@ -20,10 +25,10 @@ def plot_volcano_relative_risk(
         xlim (tuple): Optional x-axis limit.
         ylim (tuple): Optional y-axis limit.
     """
-    import matplotlib.pyplot as plt
-    import numpy as np
     import os
 
+    plt = require("matplotlib.pyplot", extra="plotting", purpose="relative risk plots")
+
     for ref, group_results in results_dict.items():
         for group_label, (results_df, _) in group_results.items():
             if results_df.empty:
@@ -31,8 +36,8 @@ def plot_volcano_relative_risk(
                 continue
 
             # Split by site type
-            gpc_df = results_df[results_df['GpC_Site']]
-            cpg_df = results_df[results_df['CpG_Site']]
+            gpc_df = results_df[results_df["GpC_Site"]]
+            cpg_df = results_df[results_df["CpG_Site"]]
 
             fig, ax = plt.subplots(figsize=(12, 6))
 
@@ -43,29 +48,29 @@ def plot_volcano_relative_risk(
 
             # GpC as circles
             sc1 = ax.scatter(
-                gpc_df['Genomic_Position'],
-                gpc_df['log2_Relative_Risk'],
-                c=gpc_df['-log10_Adj_P'],
-                cmap='coolwarm',
-                edgecolor='k',
+                gpc_df["Genomic_Position"],
+                gpc_df["log2_Relative_Risk"],
+                c=gpc_df["-log10_Adj_P"],
+                cmap="coolwarm",
+                edgecolor="k",
                 s=40,
-                marker='o',
-                label='GpC'
+                marker="o",
+                label="GpC",
             )
 
             # CpG as stars
             sc2 = ax.scatter(
-                cpg_df['Genomic_Position'],
-                cpg_df['log2_Relative_Risk'],
-                c=cpg_df['-log10_Adj_P'],
-                cmap='coolwarm',
-                edgecolor='k',
+                cpg_df["Genomic_Position"],
+                cpg_df["log2_Relative_Risk"],
+                c=cpg_df["-log10_Adj_P"],
+                cmap="coolwarm",
+                edgecolor="k",
                 s=60,
-                marker='*',
-                label='CpG'
+                marker="*",
+                label="CpG",
             )
 
-            ax.axhline(y=0, color='gray', linestyle='--')
+            ax.axhline(y=0, color="gray", linestyle="--")
             ax.set_xlabel("Genomic Position")
             ax.set_ylabel("log2(Relative Risk)")
             ax.set_title(f"{ref} / {group_label} — Relative Risk vs Genomic Position")
@@ -75,8 +80,8 @@ def plot_volcano_relative_risk(
             if ylim:
                 ax.set_ylim(ylim)
 
-            ax.spines['top'].set_visible(False)
-            ax.spines['right'].set_visible(False)
+            ax.spines["top"].set_visible(False)
+            ax.spines["right"].set_visible(False)
 
             cbar = plt.colorbar(sc1, ax=ax)
             cbar.set_label("-log10(Adjusted P-Value)")
@@ -87,13 +92,19 @@ def plot_volcano_relative_risk(
             # Save if requested
             if save_path:
                 os.makedirs(save_path, exist_ok=True)
-                safe_name = f"{ref}_{group_label}".replace("=", "").replace("__", "_").replace(",", "_").replace(" ", "_")
+                safe_name = (
+                    f"{ref}_{group_label}".replace("=", "")
+                    .replace("__", "_")
+                    .replace(",", "_")
+                    .replace(" ", "_")
+                )
                 out_file = os.path.join(save_path, f"{safe_name}.png")
                 plt.savefig(out_file, dpi=300)
                 print(f"Saved: {out_file}")
 
             plt.show()
 
+
 def plot_bar_relative_risk(
     results_dict,
     sort_by_position=True,
@@ -102,7 +113,7 @@ def plot_bar_relative_risk(
     save_path=None,
     highlight_regions=None,  # List of (start, end) tuples
     highlight_color="lightgray",
-    highlight_alpha=0.3
+    highlight_alpha=0.3,
 ):
     """
     Plot log2(Relative Risk) as a bar plot across genomic positions for each group within each reference.
@@ -116,10 +127,10 @@ def plot_bar_relative_risk(
         highlight_color (str): Color of shaded region.
         highlight_alpha (float): Transparency of shaded region.
     """
-    import matplotlib.pyplot as plt
-    import numpy as np
     import os
 
+    plt = require("matplotlib.pyplot", extra="plotting", purpose="relative risk plots")
+
     for ref, group_data in results_dict.items():
         for group_label, (df, _) in group_data.items():
             if df.empty:
@@ -127,14 +138,14 @@ def plot_bar_relative_risk(
                 continue
 
             df = df.copy()
-            df['Genomic_Position'] = df['Genomic_Position'].astype(int)
+            df["Genomic_Position"] = df["Genomic_Position"].astype(int)
 
             if sort_by_position:
-                df = df.sort_values('Genomic_Position')
+                df = df.sort_values("Genomic_Position")
 
-            gpc_mask = df['GpC_Site'] & ~df['CpG_Site']
-            cpg_mask = df['CpG_Site'] & ~df['GpC_Site']
-            both_mask = df['GpC_Site'] & df['CpG_Site']
+            gpc_mask = df["GpC_Site"] & ~df["CpG_Site"]
+            cpg_mask = df["CpG_Site"] & ~df["GpC_Site"]
+            both_mask = df["GpC_Site"] & df["CpG_Site"]
 
             fig, ax = plt.subplots(figsize=(14, 6))
 
@@ -145,36 +156,36 @@ def plot_bar_relative_risk(
 
             # Bar plots
             ax.bar(
-                df['Genomic_Position'][gpc_mask],
-                df['log2_Relative_Risk'][gpc_mask],
+                df["Genomic_Position"][gpc_mask],
+                df["log2_Relative_Risk"][gpc_mask],
                 width=10,
-                color='steelblue',
-                label='GpC Site',
-                edgecolor='black'
+                color="steelblue",
+                label="GpC Site",
+                edgecolor="black",
             )
 
             ax.bar(
-                df['Genomic_Position'][cpg_mask],
-                df['log2_Relative_Risk'][cpg_mask],
+                df["Genomic_Position"][cpg_mask],
+                df["log2_Relative_Risk"][cpg_mask],
                 width=10,
-                color='darkorange',
-                label='CpG Site',
-                edgecolor='black'
+                color="darkorange",
+                label="CpG Site",
+                edgecolor="black",
             )
 
             if both_mask.any():
                 ax.bar(
-                    df['Genomic_Position'][both_mask],
-                    df['log2_Relative_Risk'][both_mask],
+                    df["Genomic_Position"][both_mask],
+                    df["log2_Relative_Risk"][both_mask],
                     width=10,
-                    color='purple',
-                    label='GpC + CpG',
-                    edgecolor='black'
+                    color="purple",
+                    label="GpC + CpG",
+                    edgecolor="black",
                 )
 
-            ax.axhline(y=0, color='gray', linestyle='--')
-            ax.set_xlabel('Genomic Position')
-            ax.set_ylabel('log2(Relative Risk)')
+            ax.axhline(y=0, color="gray", linestyle="--")
+            ax.set_xlabel("Genomic Position")
+            ax.set_ylabel("log2(Relative Risk)")
             ax.set_title(f"{ref} — {group_label}")
             ax.legend()
 
@@ -183,20 +194,23 @@ def plot_bar_relative_risk(
             if ylim:
                 ax.set_ylim(ylim)
 
-            ax.spines['top'].set_visible(False)
-            ax.spines['right'].set_visible(False)
+            ax.spines["top"].set_visible(False)
+            ax.spines["right"].set_visible(False)
 
             plt.tight_layout()
 
             if save_path:
                 os.makedirs(save_path, exist_ok=True)
-                safe_name = f"{ref}_{group_label}".replace("=", "").replace("__", "_").replace(",", "_")
+                safe_name = (
+                    f"{ref}_{group_label}".replace("=", "").replace("__", "_").replace(",", "_")
+                )
                 out_file = os.path.join(save_path, f"{safe_name}.png")
                 plt.savefig(out_file, dpi=300)
                 print(f"📁 Saved: {out_file}")
 
             plt.show()
 
+
 def plot_positionwise_matrix(
     adata,
     key="positionwise_result",
@@ -210,35 +224,40 @@ def plot_positionwise_matrix(
     xtick_step=10,
     ytick_step=10,
     save_path=None,
-    highlight_position=None,         # Can be a single int/float or list of them
-    highlight_axis="row",            # "row" or "column"
-    annotate_points=False             # ✅ New option
+    highlight_position=None,  # Can be a single int/float or list of them
+    highlight_axis="row",  # "row" or "column"
+    annotate_points=False,  # ✅ New option
 ):
     """
     Plots positionwise matrices stored in adata.uns[key], with an optional line plot
     for specified row(s) or column(s), and highlights them on the heatmap.
     """
-    import matplotlib.pyplot as plt
-    import seaborn as sns
+    import os
+
     import numpy as np
     import pandas as pd
-    import os
+
+    plt = require("matplotlib.pyplot", extra="plotting", purpose="position stats plots")
+    sns = require("seaborn", extra="plotting", purpose="position stats plots")
 
     def find_closest_index(index, target):
+        """Find the index value closest to a target value."""
         index_vals = pd.to_numeric(index, errors="coerce")
         target_val = pd.to_numeric([target], errors="coerce")[0]
         diffs = pd.Series(np.abs(index_vals - target_val), index=index)
         return diffs.idxmin()
 
     # Ensure highlight_position is a list
-    if highlight_position is not None and not isinstance(highlight_position, (list, tuple, np.ndarray)):
+    if highlight_position is not None and not isinstance(
+        highlight_position, (list, tuple, np.ndarray)
+    ):
         highlight_position = [highlight_position]
 
     for group, mat_df in adata.uns[key].items():
         mat = mat_df.copy()
 
         if log_transform:
-            with np.errstate(divide='ignore', invalid='ignore'):
+            with np.errstate(divide="ignore", invalid="ignore"):
                 if log_base == "log1p":
                     mat = np.log1p(mat)
                 elif log_base == "log2":
@@ -276,7 +295,7 @@ def plot_positionwise_matrix(
             vmin=vmin,
             vmax=vmax,
             cbar_kws={"label": f"{key} ({log_base})" if log_transform else key},
-            ax=heat_ax
+            ax=heat_ax,
         )
 
         heat_ax.set_title(f"{key} — {group}", pad=20)
@@ -295,17 +314,27 @@ def plot_positionwise_matrix(
                         series = mat.loc[closest]
                         x_vals = pd.to_numeric(series.index, errors="coerce")
                         idx = mat.index.get_loc(closest)
-                        heat_ax.axhline(idx, color=colors[i % len(colors)], linestyle="--", linewidth=1)
+                        heat_ax.axhline(
+                            idx, color=colors[i % len(colors)], linestyle="--", linewidth=1
+                        )
                         label = f"Row {pos} → {closest}"
                     else:
                         closest = find_closest_index(mat.columns, pos)
                         series = mat[closest]
                         x_vals = pd.to_numeric(series.index, errors="coerce")
                         idx = mat.columns.get_loc(closest)
-                        heat_ax.axvline(idx, color=colors[i % len(colors)], linestyle="--", linewidth=1)
+                        heat_ax.axvline(
+                            idx, color=colors[i % len(colors)], linestyle="--", linewidth=1
+                        )
                         label = f"Col {pos} → {closest}"
 
-                    line = line_ax.plot(x_vals, series.values, marker='o', label=label, color=colors[i % len(colors)])
+                    line = line_ax.plot(
+                        x_vals,
+                        series.values,
+                        marker="o",
+                        label=label,
+                        color=colors[i % len(colors)],
+                    )
 
                     # Annotate each point
                     if annotate_points:
@@ -316,12 +345,18 @@ def plot_positionwise_matrix(
                                     xy=(x, y),
                                     textcoords="offset points",
                                     xytext=(0, 5),
-                                    ha='center',
-                                    fontsize=8
+                                    ha="center",
+                                    fontsize=8,
                                 )
                 except Exception as e:
-                    line_ax.text(0.5, 0.5, f"⚠️ Error plotting {highlight_axis} @ {pos}",
-                                 ha='center', va='center', fontsize=10)
+                    line_ax.text(
+                        0.5,
+                        0.5,
+                        f"⚠️ Error plotting {highlight_axis} @ {pos}",
+                        ha="center",
+                        va="center",
+                        fontsize=10,
+                    )
                     print(f"Error plotting line for {highlight_axis}={pos}: {e}")
 
             line_ax.set_title(f"{highlight_axis.capitalize()} Profile(s)")
@@ -342,6 +377,7 @@ def plot_positionwise_matrix(
 
         plt.show()
 
+
 def plot_positionwise_matrix_grid(
     adata,
     key,
@@ -356,32 +392,63 @@ def plot_positionwise_matrix_grid(
     xtick_step=10,
     ytick_step=10,
     parallel=False,
-    max_threads=None
+    max_threads=None,
 ):
-    import matplotlib.pyplot as plt
-    import seaborn as sns
+    """Plot a grid of positionwise matrices grouped by metadata.
+
+    Args:
+        adata: AnnData containing matrices in ``adata.uns``.
+        key: Key for positionwise matrices.
+        outer_keys: Keys for outer grouping.
+        inner_keys: Keys for inner grouping.
+        log_transform: Optional log transform (``log2`` or ``log1p``).
+        vmin: Minimum color scale value.
+        vmax: Maximum color scale value.
+        cmap: Matplotlib colormap.
+        save_path: Optional path to save plots.
+        figsize: Figure size.
+        xtick_step: X-axis tick step.
+        ytick_step: Y-axis tick step.
+        parallel: Whether to plot in parallel.
+        max_threads: Max thread count for parallel plotting.
+    """
+    import os
+
     import numpy as np
     import pandas as pd
-    import os
-    from matplotlib.gridspec import GridSpec
     from joblib import Parallel, delayed
 
+    plt = require("matplotlib.pyplot", extra="plotting", purpose="position stats plots")
+    sns = require("seaborn", extra="plotting", purpose="position stats plots")
+    grid_spec = require("matplotlib.gridspec", extra="plotting", purpose="position stats plots")
+    GridSpec = grid_spec.GridSpec
+
     matrices = adata.uns[key]
     group_labels = list(matrices.keys())
 
-    parsed_inner = pd.DataFrame([dict(zip(inner_keys, g.split("_")[-len(inner_keys):])) for g in group_labels])
-    parsed_outer = pd.Series(["_".join(g.split("_")[:-len(inner_keys)]) for g in group_labels], name="outer")
+    parsed_inner = pd.DataFrame(
+        [dict(zip(inner_keys, g.split("_")[-len(inner_keys) :])) for g in group_labels]
+    )
+    parsed_outer = pd.Series(
+        ["_".join(g.split("_")[: -len(inner_keys)]) for g in group_labels], name="outer"
+    )
     parsed = pd.concat([parsed_outer, parsed_inner], axis=1)
 
     def plot_one_grid(outer_label):
-        selected = parsed[parsed['outer'] == outer_label].copy()
-        selected["group_str"] = [f"{outer_label}_{row[inner_keys[0]]}_{row[inner_keys[1]]}" for _, row in selected.iterrows()]
+        """Plot one grid for a specific outer label."""
+        selected = parsed[parsed["outer"] == outer_label].copy()
+        selected["group_str"] = [
+            f"{outer_label}_{row[inner_keys[0]]}_{row[inner_keys[1]]}"
+            for _, row in selected.iterrows()
+        ]
 
         row_vals = sorted(selected[inner_keys[0]].unique())
         col_vals = sorted(selected[inner_keys[1]].unique())
 
         fig = plt.figure(figsize=figsize)
-        gs = GridSpec(len(row_vals), len(col_vals) + 1, width_ratios=[1]*len(col_vals) + [0.05], wspace=0.3)
+        gs = GridSpec(
+            len(row_vals), len(col_vals) + 1, width_ratios=[1] * len(col_vals) + [0.05], wspace=0.3
+        )
         axes = np.empty((len(row_vals), len(col_vals)), dtype=object)
 
         local_vmin, local_vmax = vmin, vmax
@@ -397,10 +464,7 @@ def plot_positionwise_matrix_grid(
                 local_vmin = -vmax_auto if vmin is None else vmin
                 local_vmax = vmax_auto if vmax is None else vmax
 
-        cbar_label = {
-            "log2": "log2(Value)",
-            "log1p": "log1p(Value)"
-        }.get(log_transform, "Value")
+        cbar_label = {"log2": "log2(Value)", "log1p": "log1p(Value)"}.get(log_transform, "Value")
 
         cbar_ax = fig.add_subplot(gs[:, -1])
 
@@ -431,9 +495,11 @@ def plot_positionwise_matrix_grid(
                     vmax=local_vmax,
                     cbar=(i == 0 and j == 0),
                     cbar_ax=cbar_ax if (i == 0 and j == 0) else None,
-                    cbar_kws={"label": cbar_label if (i == 0 and j == 0) else ""}
+                    cbar_kws={"label": cbar_label if (i == 0 and j == 0) else ""},
+                )
+                ax.set_title(
+                    f"{inner_keys[0]}={row_val}, {inner_keys[1]}={col_val}", fontsize=9, pad=8
                 )
-                ax.set_title(f"{inner_keys[0]}={row_val}, {inner_keys[1]}={col_val}", fontsize=9, pad=8)
 
                 xticks = data.columns.astype(int)
                 yticks = data.index.astype(int)
@@ -448,15 +514,17 @@ def plot_positionwise_matrix_grid(
         if save_path:
             os.makedirs(save_path, exist_ok=True)
             fname = outer_label.replace("_", "").replace("=", "") + ".png"
-            plt.savefig(os.path.join(save_path, fname), dpi=300, bbox_inches='tight')
+            plt.savefig(os.path.join(save_path, fname), dpi=300, bbox_inches="tight")
             print(f"Saved {fname}")
 
         plt.close(fig)
 
     if parallel:
-        Parallel(n_jobs=max_threads)(delayed(plot_one_grid)(outer_label) for outer_label in parsed['outer'].unique())
+        Parallel(n_jobs=max_threads)(
+            delayed(plot_one_grid)(outer_label) for outer_label in parsed["outer"].unique()
+        )
     else:
-        for outer_label in parsed['outer'].unique():
+        for outer_label in parsed["outer"].unique():
             plot_one_grid(outer_label)
 
-    print("Finished plotting all grids.")
+    print("Finished plotting all grids.")