smftools 0.3.1__py3-none-any.whl → 0.3.2__py3-none-any.whl

smftools/tools/calculate_knn.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
+
+if TYPE_CHECKING:
+    import anndata as ad
+
+logger = get_logger(__name__)
+
+
+def calculate_knn(
+    adata: "ad.AnnData",
+    obsm: str = "X_pca",
+    knn_neighbors: int = 100,
+    overwrite: bool = True,
+    threads: int = 8,
+    random_state: int | None = 0,
+    symmetrize: bool = True,
+) -> "ad.AnnData":
+    """Compute a KNN distance graph on an embedding in `adata.obsm[obsm]`.
+
+    Stores:
+      - adata.obsp[f"knn_distances_{obsm}"] : CSR sparse matrix of distances
+      - adata.uns[f"knn_distances_{obsm}"]["params"] : metadata
+
+    Args:
+        adata: AnnData object to update.
+        obsm: Key in `adata.obsm` to use as the embedding.
+        knn_neighbors: Target number of neighbors (will be clipped to n_obs-1).
+        overwrite: If False and graph exists, do nothing.
+        threads: Parallel jobs for pynndescent.
+        random_state: Seed for pynndescent.
+        symmetrize: If True, make distance graph symmetric via min(A, A.T).
+
+    Returns:
+        Updated AnnData.
+    """
+    import numpy as np
+    import scipy.sparse as sp
+
+    if obsm not in adata.obsm:
+        raise KeyError(f"`{obsm}` not found in adata.obsm. Available: {list(adata.obsm.keys())}")
+
+    out_key = f"knn_distances_{obsm}"
+    if not overwrite and out_key in adata.obsp:
+        logger.info("KNN graph %r already exists and overwrite=False; skipping.", out_key)
+        return adata
+
+    data = adata.obsm[obsm]
+
+    if sp.issparse(data):
+        # Convert to float32 for pynndescent/numba friendliness if needed
+        data = data.astype(np.float32)
+        logger.info(
+            "Sparse embedding detected (%s). Proceeding without NaN check.", type(data).__name__
+        )
+    else:
+        data = np.asarray(data)
+        if np.isnan(data).any():
+            logger.warning("NaNs detected in %s; filling NaNs with 0.5 before KNN.", obsm)
+            data = np.nan_to_num(data, nan=0.5)
+        data = data.astype(np.float32, copy=False)
+
+    pynndescent = require("pynndescent", extra="umap", purpose="KNN graph computation")
+
+    n_obs = data.shape[0]
+    if n_obs < 2:
+        raise ValueError(f"Need at least 2 observations for KNN; got n_obs={n_obs}")
+
+    n_neighbors = min(int(knn_neighbors), n_obs - 1)
+    if n_neighbors < 1:
+        raise ValueError(f"Computed n_neighbors={n_neighbors}; check knn_neighbors and n_obs.")
+
+    logger.info(
+        "Running pynndescent KNN (obsm=%s, n_neighbors=%d, metric=euclidean, n_jobs=%d)",
+        obsm,
+        n_neighbors,
+        threads,
+    )
+
+    nn_index = pynndescent.NNDescent(
+        data,
+        n_neighbors=n_neighbors,
+        metric="euclidean",
+        random_state=random_state,
+        n_jobs=threads,
+    )
+    knn_indices, knn_dists = nn_index.neighbor_graph  # shapes: (n_obs, n_neighbors)
+
+    rows = np.repeat(np.arange(n_obs, dtype=np.int64), n_neighbors)
+    cols = knn_indices.reshape(-1).astype(np.int64, copy=False)
+    vals = knn_dists.reshape(-1).astype(np.float32, copy=False)
+
+    distances = sp.coo_matrix((vals, (rows, cols)), shape=(n_obs, n_obs)).tocsr()
+
+    # Optional: ensure diagonal is 0 and (optionally) symmetrize
+    distances.setdiag(0.0)
+    distances.eliminate_zeros()
+
+    if symmetrize:
+        # Keep the smaller directed distance for each undirected edge
+        distances = distances.minimum(distances.T)
+
+    adata.obsp[out_key] = distances
+    adata.uns[out_key] = {
+        "params": {
+            "obsm": obsm,
+            "n_neighbors_requested": int(knn_neighbors),
+            "n_neighbors_used": int(n_neighbors),
+            "method": "pynndescent",
+            "metric": "euclidean",
+            "random_state": random_state,
+            "n_jobs": int(threads),
+            "symmetrize": bool(symmetrize),
+        }
+    }
+
+    return adata

smftools/tools/calculate_nmf.py

@@ -9,6 +9,7 @@ from smftools.optional_imports import require
 
 if TYPE_CHECKING:
     import anndata as ad
+    import numpy as np
 
 logger = get_logger(__name__)
 
@@ -16,7 +17,7 @@ logger = get_logger(__name__)
 def calculate_nmf(
     adata: "ad.AnnData",
     layer: str | None = "nan_half",
-    var_filters: Sequence[str] | None = None,
+    var_mask: "np.ndarray | Sequence[bool] | None" = None,
     n_components: int = 2,
     max_iter: int = 200,
     random_state: int = 0,
@@ -24,13 +25,14 @@ def calculate_nmf(
     embedding_key: str = "X_nmf",
     components_key: str = "H_nmf",
     uns_key: str = "nmf",
+    suffix: str | None = None,
 ) -> "ad.AnnData":
     """Compute a low-dimensional NMF embedding.
 
     Args:
         adata: AnnData object to update.
         layer: Layer name to use for NMF (``None`` uses ``adata.X``).
-        var_filters: Optional list of var masks to subset features.
+        var_mask: Optional boolean mask to subset features.
         n_components: Number of NMF components to compute.
         max_iter: Maximum number of NMF iterations.
         random_state: Random seed for the NMF initializer.
@@ -47,6 +49,11 @@ def calculate_nmf(
     require("sklearn", extra="ml-base", purpose="NMF calculation")
     from sklearn.decomposition import NMF
 
+    if suffix:
+        embedding_key = f"{embedding_key}_{suffix}"
+        components_key = f"{components_key}_{suffix}"
+        uns_key = f"{uns_key}_{suffix}"
+
     has_embedding = embedding_key in adata.obsm
     has_components = components_key in adata.varm
     if has_embedding and has_components and not overwrite:
@@ -56,17 +63,21 @@ def calculate_nmf(
         logger.info("NMF embedding present without components; recomputing to store components.")
 
     subset_mask = None
-    if var_filters:
-        subset_mask = np.logical_or.reduce([adata.var[f].values for f in var_filters])
+    if var_mask is not None:
+        subset_mask = np.asarray(var_mask, dtype=bool)
+        if subset_mask.ndim != 1 or subset_mask.shape[0] != adata.n_vars:
+            raise ValueError(
+                "var_mask must be a 1D boolean array with length matching adata.n_vars."
+            )
         adata_subset = adata[:, subset_mask].copy()
         logger.info(
             "Subsetting adata: retained %s features based on filters %s",
             adata_subset.shape[1],
-            var_filters,
+            "var_mask",
         )
     else:
         adata_subset = adata.copy()
-        logger.info("No var filters provided. Using all features.")
+        logger.info("No var_mask provided. Using all features.")
 
     data = adata_subset.layers[layer] if layer else adata_subset.X
     if issparse(data):
@@ -107,7 +118,7 @@ def calculate_nmf(
         "max_iter": max_iter,
         "random_state": random_state,
         "layer": layer,
-        "var_filters": list(var_filters) if var_filters else None,
+        "var_mask_provided": var_mask is not None,
         "components_key": components_key,
     }
 

smftools/tools/calculate_pca.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+
+from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
+
+if TYPE_CHECKING:
+    import anndata as ad
+    import numpy as np
+
+logger = get_logger(__name__)
+
+
+def calculate_pca(
+    adata: "ad.AnnData",
+    layer: str | None = "nan_half",
+    var_mask: "np.ndarray | Sequence[bool] | None" = None,
+    n_pcs: int = 15,
+    overwrite: bool = True,
+    output_suffix: str | None = None,
+    fill_nan: float | None = 0.5,
+) -> "ad.AnnData":
+    """Compute PCA and store scores in `.obsm` and loadings in `.varm`."""
+
+    import numpy as np
+    import scipy.sparse as sp
+
+    obsm_output = f"X_pca_{output_suffix}" if output_suffix else "X_pca"
+    varm_output = f"PCs_{output_suffix}" if output_suffix else "PCs"
+
+    if not overwrite and obsm_output in adata.obsm and varm_output in adata.varm:
+        logger.info(
+            "PCA outputs already exist and overwrite=False; skipping (%s, %s).",
+            obsm_output,
+            varm_output,
+        )
+        return adata
+
+    # --- Build feature subset mask (over vars) ---
+    if var_mask is not None:
+        subset_mask = np.asarray(var_mask, dtype=bool)
+        if subset_mask.ndim != 1 or subset_mask.shape[0] != adata.n_vars:
+            raise ValueError(
+                "var_mask must be a 1D boolean array with length matching adata.n_vars."
+            )
+        n_vars_used = int(subset_mask.sum())
+        if n_vars_used == 0:
+            raise ValueError("var_mask retained 0 features.")
+        logger.info(
+            "Subsetting vars: retained %d / %d features from var_mask",
+            n_vars_used,
+            adata.n_vars,
+        )
+    else:
+        subset_mask = slice(None)
+        n_vars_used = adata.n_vars
+        logger.info("No var_mask provided; using all %d features.", adata.n_vars)
+
+    # --- Pull matrix view ---
+    if layer is None:
+        matrix = adata.X
+        layer_used = None
+    else:
+        if layer not in adata.layers:
+            raise KeyError(
+                f"Layer {layer!r} not found in adata.layers. Available: {list(adata.layers.keys())}"
+            )
+        matrix = adata.layers[layer]
+        layer_used = layer
+
+    matrix = matrix[:, subset_mask]  # slice view (sparse OK)
+
+    n_obs = matrix.shape[0]
+    if n_obs < 2:
+        raise ValueError(f"PCA requires at least 2 observations; got n_obs={n_obs}")
+    if n_vars_used < 1:
+        raise ValueError("PCA requires at least 1 feature.")
+
+    n_pcs_requested = int(n_pcs)
+    n_pcs_used = min(n_pcs_requested, n_obs, n_vars_used)
+    if n_pcs_used < 1:
+        raise ValueError(f"n_pcs_used became {n_pcs_used}; check inputs.")
+
+    # --- NaN handling (dense only; sparse usually won’t store NaNs) ---
+    if not sp.issparse(matrix):
+        X = np.asarray(matrix, dtype=np.float32)
+        if fill_nan is not None and np.isnan(X).any():
+            logger.warning("NaNs detected; filling NaNs with %s before PCA.", fill_nan)
+            X = np.nan_to_num(X, nan=float(fill_nan))
+    else:
+        X = matrix  # keep sparse
+
+    # --- PCA ---
+    # Prefer sklearn's randomized PCA for speed on big matrices.
+    used_sklearn = False
+    try:
+        sklearn = require("sklearn", extra="ml", purpose="PCA computation")
+        from sklearn.decomposition import PCA, TruncatedSVD
+
+        if sp.issparse(X):
+            # TruncatedSVD works on sparse without centering; good approximation.
+            # If you *need* centered PCA on sparse, you'd need different machinery.
+            logger.info("Running TruncatedSVD (sparse) with n_components=%d", n_pcs_used)
+            model = TruncatedSVD(n_components=n_pcs_used, random_state=0)
+            scores = model.fit_transform(X)  # (n_obs, n_pcs)
+            loadings = model.components_.T  # (n_vars_used, n_pcs)
+            mean = None
+            explained_variance_ratio = getattr(model, "explained_variance_ratio_", None)
+        else:
+            logger.info(
+                "Running sklearn PCA with n_components=%d (svd_solver=randomized)", n_pcs_used
+            )
+            model = PCA(n_components=n_pcs_used, svd_solver="randomized", random_state=0)
+            scores = model.fit_transform(X)  # (n_obs, n_pcs)
+            loadings = model.components_.T  # (n_vars_used, n_pcs)
+            mean = model.mean_
+            explained_variance_ratio = model.explained_variance_ratio_
+
+        used_sklearn = True
+
+    except Exception as e:
+        # Fallback to your manual SVD (dense only)
+        if sp.issparse(X):
+            raise RuntimeError(
+                "Sparse input PCA fallback is not implemented without sklearn. "
+                "Install scikit-learn (extra 'ml') or densify upstream."
+            ) from e
+
+        import scipy.linalg as spla
+
+        logger.warning(
+            "sklearn PCA unavailable; falling back to full SVD (can be slow). Reason: %s", e
+        )
+        Xd = np.asarray(X, dtype=np.float64)
+        mean = Xd.mean(axis=0)
+        centered = Xd - mean
+        u, s, vt = spla.svd(centered, full_matrices=False)
+        u = u[:, :n_pcs_used]
+        s = s[:n_pcs_used]
+        vt = vt[:n_pcs_used]
+        scores = u * s
+        loadings = vt.T
+        explained_variance_ratio = None
+
+    # --- Store scores (obsm) ---
+    adata.obsm[obsm_output] = scores
+
+    # --- Store loadings (varm) with original var dimension ---
+    pc_matrix = np.zeros((adata.n_vars, n_pcs_used), dtype=np.float32)
+    if isinstance(subset_mask, slice):
+        pc_matrix[:, :] = loadings
+    else:
+        pc_matrix[subset_mask, :] = loadings.astype(np.float32, copy=False)
+
+    adata.varm[varm_output] = pc_matrix
+
+    # --- Metadata ---
+    adata.uns[obsm_output] = {
+        "params": {
+            "layer": layer_used,
+            "var_mask_provided": var_mask is not None,
+            "n_pcs_requested": n_pcs_requested,
+            "n_pcs_used": int(n_pcs_used),
+            "used_sklearn": used_sklearn,
+            "fill_nan": fill_nan,
+            "note_sparse": bool(sp.issparse(matrix)),
+        },
+        "explained_variance_ratio": explained_variance_ratio,
+        "mean": mean.tolist() if (mean is not None and isinstance(mean, np.ndarray)) else None,
+    }
+
+    logger.info(
+        "Stored PCA: adata.obsm[%s] (%s) and adata.varm[%s] (%s)",
+        obsm_output,
+        scores.shape,
+        varm_output,
+        pc_matrix.shape,
+    )
+    return adata

smftools/tools/calculate_umap.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 
 from smftools.logging_utils import get_logger
 from smftools.optional_imports import require
@@ -13,174 +13,90 @@ logger = get_logger(__name__)
 
 def calculate_umap(
     adata: "ad.AnnData",
-    layer: str | None = "nan_half",
-    var_filters: Sequence[str] | None = None,
-    n_pcs: int = 15,
-    knn_neighbors: int = 100,
+    obsm: str | None = "X_pca",
     overwrite: bool = True,
     threads: int = 8,
     random_state: int | None = 0,
+    output_suffix: str | None = None,
 ) -> "ad.AnnData":
-    """Compute PCA, neighbors, and UMAP embeddings.
-
-    Args:
-        adata: AnnData object to update.
-        layer: Layer name to use for PCA/UMAP (``None`` uses ``adata.X``).
-        var_filters: Optional list of var masks to subset features.
-        n_pcs: Number of principal components.
-        knn_neighbors: Number of neighbors for the graph.
-        overwrite: Whether to recompute embeddings if they exist.
-        threads: Number of OMP threads for computation.
-
-    Returns:
-        anndata.AnnData: Updated AnnData object.
-    """
-    import os
+    """Compute UMAP embedding from an `.obsm` embedding, and store connectivities."""
 
     import numpy as np
-    import scipy.linalg as spla
     import scipy.sparse as sp
 
+    if obsm is None:
+        raise ValueError("obsm must be a key in adata.obsm (e.g., 'X_pca').")
+
+    if obsm not in adata.obsm:
+        raise KeyError(f"`{obsm}` not found in adata.obsm. Available: {list(adata.obsm.keys())}")
+
     umap = require("umap", extra="umap", purpose="UMAP calculation")
-    pynndescent = require("pynndescent", extra="umap", purpose="KNN graph computation")
-
-    os.environ["OMP_NUM_THREADS"] = str(threads)
-
-    # Step 1: Apply var filter
-    if var_filters:
-        subset_mask = np.logical_or.reduce([adata.var[f].values for f in var_filters])
-        adata_subset = adata[:, subset_mask].copy()
-        logger.info(
-            "Subsetting adata: retained %s features based on filters %s",
-            adata_subset.shape[1],
-            var_filters,
+
+    output_obsm = f"X_umap_{output_suffix}" if output_suffix else "X_umap"
+    conn_key = f"connectivities_{obsm}"
+
+    # Decide n_neighbors: prefer stored KNN params, else UMAP default-ish
+    n_neighbors = None
+    knn_uns_key = f"knn_distances_{obsm}"
+    if knn_uns_key in adata.uns:
+        params = adata.uns[knn_uns_key].get("params", {})
+        n_neighbors = params.get("n_neighbors_used", params.get("n_neighbors", None))
+    if n_neighbors is None:
+        n_neighbors = 15  # reasonable default if KNN wasn't precomputed
+        logger.warning(
+            "No %r found in adata.uns; defaulting n_neighbors=%d for UMAP.",
+            knn_uns_key,
+            n_neighbors,
         )
+
+    # Build input matrix X and handle NaNs locally
+    X = adata.obsm[obsm]
+    if sp.issparse(X):
+        # UMAP can accept sparse CSR; keep it sparse
+        pass
     else:
-        adata_subset = adata.copy()
-        logger.info("No var filters provided. Using all features.")
-
-    # Step 2: NaN handling inside layer
-    if layer:
-        data = adata_subset.layers[layer]
-        if not sp.issparse(data):
-            if np.isnan(data).any():
-                logger.warning("NaNs detected, filling with 0.5 before PCA + neighbors.")
-                data = np.nan_to_num(data, nan=0.5)
-                adata_subset.layers[layer] = data
-            else:
-                logger.info("No NaNs detected.")
-        else:
-            logger.info(
-                "Sparse matrix detected; skipping NaN check (sparse formats typically do not store NaNs)."
-            )
-
-    # Step 3: PCA + neighbors + UMAP on subset
-    if "X_umap" not in adata_subset.obsm or overwrite:
-        n_pcs = min(adata_subset.shape[1], n_pcs)
-        logger.info("Running PCA with n_pcs=%s", n_pcs)
-
-        if layer:
-            matrix = adata_subset.layers[layer]
-        else:
-            matrix = adata_subset.X
-
-        if sp.issparse(matrix):
-            logger.warning("Converting sparse matrix to dense for PCA.")
-            matrix = matrix.toarray()
-
-        matrix = np.asarray(matrix, dtype=float)
-        mean = matrix.mean(axis=0)
-        centered = matrix - mean
-
-        if centered.shape[0] == 0 or centered.shape[1] == 0:
-            raise ValueError("PCA requires a non-empty matrix.")
-
-        if n_pcs <= 0:
-            raise ValueError("n_pcs must be positive.")
-
-        if centered.shape[1] <= n_pcs:
-            n_pcs = centered.shape[1]
-
-        if centered.shape[0] < n_pcs:
-            n_pcs = centered.shape[0]
-
-        u, s, vt = spla.svd(centered, full_matrices=False)
-
-        u = u[:, :n_pcs]
-        s = s[:n_pcs]
-        vt = vt[:n_pcs]
-
-        adata_subset.obsm["X_pca"] = u * s
-        adata_subset.varm["PCs"] = vt.T
-
-        logger.info("Running neighborhood graph with pynndescent (n_neighbors=%s)", knn_neighbors)
-        n_neighbors = min(knn_neighbors, max(1, adata_subset.n_obs - 1))
-        nn_index = pynndescent.NNDescent(
-            adata_subset.obsm["X_pca"],
-            n_neighbors=n_neighbors,
-            metric="euclidean",
-            random_state=random_state,
-            n_jobs=threads,
-        )
-        knn_indices, knn_dists = nn_index.neighbor_graph
-
-        rows = np.repeat(np.arange(adata_subset.n_obs), n_neighbors)
-        cols = knn_indices.reshape(-1)
-        distances = sp.coo_matrix(
-            (knn_dists.reshape(-1), (rows, cols)),
-            shape=(adata_subset.n_obs, adata_subset.n_obs),
-        ).tocsr()
-        adata_subset.obsp["distances"] = distances
-
-        logger.info("Running UMAP")
-        umap_model = umap.UMAP(
-            n_neighbors=n_neighbors,
-            n_components=2,
-            metric="euclidean",
-            random_state=random_state,
+        X = np.asarray(X)
+        if np.isnan(X).any():
+            logger.warning("NaNs detected in %s; filling NaNs with 0.5 for UMAP.", obsm)
+            X = np.nan_to_num(X, nan=0.5)
+
+    if (not overwrite) and (output_obsm in adata.obsm) and (conn_key in adata.obsp):
+        logger.info("UMAP + connectivities already exist and overwrite=False; skipping.")
+        return adata
+
+    logger.info("Running UMAP (obsm=%s, n_neighbors=%d, metric=euclidean)", obsm, n_neighbors)
+
+    # Note: umap-learn uses numba threading; n_jobs controls parallelism in UMAP
+    # and is ignored when random_state is set (umap-learn behavior).
+    umap_model = umap.UMAP(
+        n_neighbors=int(n_neighbors),
+        n_components=2,
+        metric="euclidean",
+        random_state=random_state,
+        n_jobs=int(threads),
+    )
+
+    embedding = umap_model.fit_transform(X)
+    adata.obsm[output_obsm] = embedding
+
+    # UMAP's computed fuzzy graph
+    connectivities = getattr(umap_model, "graph_", None)
+    if connectivities is not None:
+        adata.obsp[conn_key] = (
+            connectivities.tocsr() if sp.issparse(connectivities) else connectivities
         )
-        adata_subset.obsm["X_umap"] = umap_model.fit_transform(adata_subset.obsm["X_pca"])
-
-        try:
-            from umap.umap_ import fuzzy_simplicial_set
-
-            fuzzy_result = fuzzy_simplicial_set(
-                adata_subset.obsm["X_pca"],
-                n_neighbors=n_neighbors,
-                random_state=random_state,
-                metric="euclidean",
-                knn_indices=knn_indices,
-                knn_dists=knn_dists,
-            )
-            connectivities = fuzzy_result[0] if isinstance(fuzzy_result, tuple) else fuzzy_result
-        except TypeError:
-            connectivities = umap_model.graph_
-
-        adata_subset.obsp["connectivities"] = connectivities
-
-    # Step 4: Store results in original adata
-    adata.obsm["X_pca"] = adata_subset.obsm["X_pca"]
-    adata.obsm["X_umap"] = adata_subset.obsm["X_umap"]
-    adata.obsp["distances"] = adata_subset.obsp["distances"]
-    adata.obsp["connectivities"] = adata_subset.obsp["connectivities"]
-    adata.uns["neighbors"] = {
+    else:
+        logger.warning("UMAP model did not expose graph_; connectivities not stored.")
+
+    adata.uns[output_obsm] = {
         "params": {
-            "n_neighbors": knn_neighbors,
-            "method": "pynndescent",
+            "obsm": obsm,
+            "n_neighbors": int(n_neighbors),
             "metric": "euclidean",
+            "random_state": random_state,
+            "n_jobs": int(threads),
         }
     }
 
-    # Fix varm["PCs"] shape mismatch
-    pc_matrix = np.zeros((adata.shape[1], adata_subset.varm["PCs"].shape[1]))
-    if var_filters:
-        subset_mask = np.logical_or.reduce([adata.var[f].values for f in var_filters])
-        pc_matrix[subset_mask, :] = adata_subset.varm["PCs"]
-    else:
-        pc_matrix = adata_subset.varm["PCs"]  # No subsetting case
-
-    adata.varm["PCs"] = pc_matrix
-
-    logger.info("Stored: adata.obsm['X_pca'] and adata.obsm['X_umap']")
-
+    logger.info("Stored: adata.obsm[%s]=%s", output_obsm, embedding.shape)
     return adata

smftools/tools/position_stats.py

@@ -320,12 +320,12 @@ def compute_positionwise_statistics(
 
     # samples / refs
     sseries = adata.obs[sample_col]
-    if not pd.api.types.is_categorical_dtype(sseries):
+    if not isinstance(sseries.dtype, pd.CategoricalDtype):
         sseries = sseries.astype("category")
     samples = list(sseries.cat.categories)
 
     rseries = adata.obs[ref_col]
-    if not pd.api.types.is_categorical_dtype(rseries):
+    if not isinstance(rseries.dtype, pd.CategoricalDtype):
         rseries = rseries.astype("category")
     references = list(rseries.cat.categories)
 
@@ -509,12 +509,12 @@ def plot_positionwise_matrices(
 
     # canonicalize sample/ref order
     sseries = adata.obs[sample_col]
-    if not pd.api.types.is_categorical_dtype(sseries):
+    if not isinstance(sseries.dtype, pd.CategoricalDtype):
         sseries = sseries.astype("category")
     samples = list(sseries.cat.categories)
 
     rseries = adata.obs[ref_col]
-    if not pd.api.types.is_categorical_dtype(rseries):
+    if not isinstance(rseries.dtype, pd.CategoricalDtype):
         rseries = rseries.astype("category")
     references = list(rseries.cat.categories)