"""Masked and highlighted UMAP (or any 2D embedding) plots.
These highlight a subset of cells (e.g. one lineage) against the rest of an
atlas, which is a common need when illustrating where a population of
interest sits within a larger dataset.
"""
from __future__ import annotations
import math
from collections.abc import Sequence
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scanpy as sc
from matplotlib.colors import to_rgb
from scipy.interpolate import RegularGridInterpolator
from scipy.ndimage import gaussian_filter
from ._config import PlotConfig
from ._palettes import colors_list
from ._style import save_figure
from ._utils import ensure_categorical, require_columns
__all__ = [
"masked_umap",
"masked_umap_highlight",
"gene_expression_umap",
"gene_coexpression_umap",
"embedding_density",
]
def _figure_dir(output_dir, figure_name: str) -> Path:
dir_name = figure_name.replace(" ", "_")
figure_dir = Path(output_dir) / dir_name
figure_dir.mkdir(parents=True, exist_ok=True)
return figure_dir
[docs]
def masked_umap(
adata,
color_by: str,
figure_name: str,
mask_column: str | None = None,
mask_values: Sequence | None = None,
palette: dict | None = None,
config: PlotConfig | None = None,
output_dir: str | Path = "figures",
basis: str = "umap",
show: bool = False,
**kwargs,
):
"""Plot a UMAP colored by ``color_by``, restricted to ``mask_values``.
Cells outside ``mask_values`` are greyed out. Saves three variants: no
legend, with legend, and a version subset to only the masked cells.
Parameters
----------
adata:
Annotated data with a computed embedding (``adata.obsm['X_<basis>']``).
color_by:
``adata.obs`` column to color cells by.
figure_name:
Used as both plot title and output subdirectory name.
mask_column:
Column whose unique values become ``mask_values`` if not given
explicitly. Only used as a default source for ``mask_values``.
mask_values:
Categories of ``color_by`` to keep visible. Defaults to all unique
values of ``mask_column`` (or of ``color_by`` if neither is set).
palette:
Optional ``{category: color}`` mapping. Falls back to the config's
registered palette for ``color_by``, then to an auto-generated one.
config:
Active :class:`~scplotkit.PlotConfig`. Uses defaults if omitted.
output_dir:
Directory figures are written under.
show:
Display each figure inline (e.g. in a Jupyter notebook) as it's
drawn. Scanpy closes the figure right after saving it, so without
this the plots are written to disk but never rendered in the cell.
**kwargs:
Forwarded to every ``sc.pl.embedding`` call, overriding this
function's own defaults (e.g. ``na_color``, ``frameon``, ``size``) --
anything the underlying scanpy function accepts.
"""
config = PlotConfig.load(config)
require_columns(adata, color_by, mask_column)
figure_dir = _figure_dir(output_dir, figure_name)
sc.settings.figdir = str(figure_dir)
ensure_categorical(adata, color_by)
if mask_values is None:
source = mask_column or color_by
mask_values = adata.obs[source].unique().tolist()
color_col_filtered = f"{color_by}_filtered"
adata.obs[color_col_filtered] = adata.obs[color_by].astype(str)
adata.obs.loc[~adata.obs[color_by].isin(mask_values), color_col_filtered] = np.nan
categories = list(mask_values)
adata.obs[color_col_filtered] = pd.Categorical(adata.obs[color_col_filtered], categories=categories)
adata.uns[f"{color_col_filtered}_colors"] = colors_list(config, color_by, mask_values, palette)
common_kwargs = dict(
basis=basis,
show=show,
na_color="white",
outline_width=(0.1, 0.05),
add_outline=True,
title=figure_name,
frameon=False,
)
sc.pl.embedding(
adata,
**{
**common_kwargs,
"color": color_col_filtered,
"legend_loc": None,
"save": f"_{figure_dir.name}_masked_colored_{color_by}.png",
**kwargs,
},
)
sc.pl.embedding(
adata,
**{
**common_kwargs,
"color": color_col_filtered,
"na_in_legend": False,
"save": f"_{figure_dir.name}_masked_colored_{color_by}_legend.png",
**kwargs,
},
)
adata_subset = adata[adata.obs[color_by].isin(mask_values)].copy()
adata_subset.obs[color_by] = adata_subset.obs[color_by].cat.remove_unused_categories()
adata_subset.uns[f"{color_by}_colors"] = colors_list(
config, color_by, list(adata_subset.obs[color_by].cat.categories), palette
)
sc.pl.embedding(
adata_subset,
**{
**common_kwargs,
"color": color_by,
"na_in_legend": False,
"save": f"_{figure_dir.name}_subset_only_{color_by}_legend.png",
**kwargs,
},
)
[docs]
def masked_umap_highlight(
adata,
color_by: str,
figure_name: str,
mask_column: str | None = None,
mask_values: Sequence | None = None,
highlight_size: float = 1.5,
background_size: float = 0.25,
ordered: bool = True,
palette: dict | None = None,
config: PlotConfig | None = None,
output_dir: str | Path = "figures",
basis: str = "umap",
show: bool = False,
**kwargs,
):
"""Like :func:`masked_umap`, but draws the masked population enlarged on top
of a small greyed-out background rather than hiding the rest entirely.
``show`` displays each figure inline (e.g. in a Jupyter notebook) as it's
drawn -- scanpy closes the figure right after saving it, so without this
the plots are written to disk but never rendered in the cell.
``**kwargs`` are forwarded to every ``sc.pl.embedding`` call, overriding
this function's own defaults -- anything the underlying scanpy function accepts.
"""
config = PlotConfig.load(config)
require_columns(adata, color_by, mask_column)
figure_dir = _figure_dir(output_dir, figure_name)
sc.settings.figdir = str(figure_dir)
ensure_categorical(adata, color_by)
if mask_values is None:
source = mask_column or color_by
mask_values = adata.obs[source].unique().tolist()
color_col_filtered = f"{color_by}_filtered"
adata.obs[color_col_filtered] = adata.obs[color_by].astype(str)
adata.obs.loc[~adata.obs[color_by].isin(mask_values), color_col_filtered] = np.nan
size_col = f"{color_col_filtered}_sizes"
adata.obs[size_col] = background_size
adata.obs.loc[~pd.isna(adata.obs[color_col_filtered]), size_col] = highlight_size
value_counts = adata.obs[color_by].value_counts()
mask_values_ordered = sorted(mask_values, key=lambda v: value_counts.get(v, 0), reverse=True)
adata.obs[color_col_filtered] = pd.Categorical(adata.obs[color_col_filtered], categories=mask_values_ordered)
adata.uns[f"{color_col_filtered}_colors"] = colors_list(config, color_by, mask_values_ordered, palette)
plot_adata = adata
if ordered:
sort_idx = adata.obs[color_col_filtered].cat.codes.argsort(kind="stable")
plot_adata = adata[sort_idx, :]
suffix = f"{color_by}_{highlight_size}_highlight"
common_kwargs = dict(
basis=basis,
color=color_col_filtered,
size=plot_adata.obs[size_col],
na_color="white",
outline_width=(0.1, 0.05),
add_outline=True,
sort_order=False,
show=show,
title=figure_name,
frameon=False,
)
sc.pl.embedding(
plot_adata,
**{
**common_kwargs,
"legend_loc": None,
"save": f"_{figure_dir.name}_masked_colored_{suffix}.png",
**kwargs,
},
)
sc.pl.embedding(
plot_adata,
**{
**common_kwargs,
"na_in_legend": False,
"save": f"_{figure_dir.name}_masked_colored_{suffix}_legend.png",
**kwargs,
},
)
def _expression_values(adata, gene: str, use_raw: bool, layer: str | None) -> np.ndarray:
if layer:
raw = adata[:, gene].layers[layer]
elif use_raw and adata.raw is not None:
raw = adata.raw[:, gene].X
else:
raw = adata[:, gene].X
return raw.toarray().flatten() if hasattr(raw, "toarray") else np.asarray(raw).flatten()
[docs]
def gene_expression_umap(
adata,
gene: str,
mask_column: str | None = None,
mask_values: Sequence | None = None,
figure_name: str | None = None,
use_raw: bool = True,
layer: str | None = None,
cmap: str = "magma",
vmin: float | None = 0,
vmax: float | None = None,
config: PlotConfig | None = None,
output_dir: str | Path = "figures",
basis: str = "umap",
show: bool = False,
**kwargs,
):
"""Plot a gene's expression on an embedding, optionally masked to a subset of cells.
Cells outside ``mask_values`` (if given) are greyed out rather than
dropped, so the full embedding shape is preserved. Saves two variants of
the plot: without and with a colorbar.
Parameters
----------
adata:
Annotated data with a computed embedding (``adata.obsm['X_<basis>']``).
gene:
Gene to plot, looked up in ``layer``, ``adata.raw`` (if ``use_raw``),
or ``adata.X``, in that order of precedence.
mask_column, mask_values:
Restrict colored (non-grey) cells to those where ``mask_column`` is
in ``mask_values``. Both must be given to enable masking.
figure_name:
Used as both plot title and output subdirectory name. Defaults to ``gene``.
config:
Active :class:`~scplotkit.PlotConfig`. Uses defaults if omitted.
output_dir:
Directory figures are written under.
show:
Display each figure inline (e.g. in a Jupyter notebook) as it's drawn.
**kwargs:
Forwarded to both ``sc.pl.embedding`` calls, overriding this
function's own defaults -- anything the underlying scanpy function accepts.
"""
config = PlotConfig.load(config)
require_columns(adata, mask_column)
figure_dir = _figure_dir(output_dir, figure_name or gene)
sc.settings.figdir = str(figure_dir)
tmp_col = f"_tmp_expr_{gene}"
adata.obs[tmp_col] = _expression_values(adata, gene, use_raw, layer)
suffix = "all_cells"
if mask_column and mask_values:
suffix = f"masked_{mask_column}"
mask = ~adata.obs[mask_column].isin(mask_values)
adata.obs.loc[mask, tmp_col] = np.nan
print(f"Masking applied: showing {gene} only in {mask_values}")
plot_kwargs = {
"basis": basis,
"color": tmp_col,
"title": figure_name or gene,
"frameon": False,
"cmap": cmap,
"vmin": vmin,
"vmax": vmax,
"na_color": "#e0e0e0",
"add_outline": True,
"outline_width": (0.1, 0.05),
"show": show,
**kwargs,
}
try:
sc.pl.embedding(
adata,
**{**plot_kwargs, "colorbar_loc": None, "save": f"_{figure_dir.name}_{gene}_{suffix}.png"},
)
sc.pl.embedding(
adata,
**{
**plot_kwargs,
"colorbar_loc": "right",
"save": f"_{figure_dir.name}_{gene}_{suffix}_with_colorbar.png",
},
)
finally:
del adata.obs[tmp_col]
def _normalized_expression(values: np.ndarray, vmax_percentile: float) -> np.ndarray:
"""Scale expression to [0, 1], clipping at the given percentile over expressing cells
so a handful of outliers don't compress the rest of the scale to near-zero."""
positive = values[values > 0]
vmax = np.percentile(positive, vmax_percentile) if positive.size else 0.0
if vmax <= 0:
vmax = 1.0
return np.clip(values / vmax, 0, 1)
[docs]
def gene_coexpression_umap(
adata,
gene1: str,
gene2: str,
color1: str = "#e6299e",
color2: str = "#3cb428",
mask_column: str | None = None,
mask_values: Sequence | None = None,
figure_name: str | None = None,
use_raw: bool = True,
layer: str | None = None,
vmax_percentile: float = 99.0,
threshold: float = 0.05,
grey_color: str | None = None,
point_size: float | None = None,
add_outline: bool = True,
outline_width: tuple[float, float] = (0.1, 0.05),
outline_color: tuple[str, str] = ("black", "white"),
show_legend: bool = True,
config: PlotConfig | None = None,
output_dir: str | Path = "figures",
basis: str = "umap",
show: bool = False,
):
"""Bivariate color-blend UMAP showing where two markers co-occur.
Each gene is normalized to [0, 1] (clipped at ``vmax_percentile`` to avoid
a few outlier cells compressing the rest of the scale), then blended
additively into a single RGB color per cell -- with the default
magenta/green colors, ``gene1``-only cells read magenta, ``gene2``-only
cells read green, and co-expressing cells read near-white. Magenta/green
is used instead of the traditional red/green overlay because red-green
color blindness (the most common form) makes that pairing indistinguishable;
magenta's blue component keeps it readable. Cells below ``threshold`` in
both genes are drawn flat grey, matching the below-threshold convention
used elsewhere in this module (see :func:`embedding_density`) rather than
showing dozens of near-black dots.
This is a purely descriptive visualization -- no statistical test for
co-expression is performed; ``threshold`` only controls what counts as
"on" for coloring purposes.
Parameters
----------
adata:
Annotated data with a computed embedding (``adata.obsm['X_<basis>']``).
gene1, gene2:
Genes to blend, looked up in ``layer``, ``adata.raw`` (if
``use_raw``), or ``adata.X``, in that order of precedence.
color1, color2:
Colors assigned to ``gene1``-only and ``gene2``-only cells. Colors
combine additively, so the default magenta/green produce a near-white
overlap; pick two colors with non-overlapping channels (e.g. a pure
magenta and a pure green) for the cleanest three-way read, and avoid
red/green -- the two are indistinguishable under red-green color
blindness.
mask_column, mask_values:
Restrict colored (non-grey) cells to those where ``mask_column`` is
in ``mask_values``. Both must be given to enable masking.
figure_name:
Used as both plot title and output subdirectory name. Defaults to
``"{gene1}_{gene2}_coexpression"``.
vmax_percentile:
Upper percentile (over cells with nonzero expression) each gene is
clipped to before normalizing to [0, 1].
threshold:
Normalized expression (0-1, per gene) below which a cell counts as
"off" for that gene. Cells "off" for both genes are drawn grey.
grey_color:
Fill color for cells below ``threshold`` in both genes. Defaults to
the config's ``general.fallback_color``.
point_size:
Marker area per cell. Defaults to scanpy's own ``120000 / n_obs``
heuristic so it scales sensibly with atlas size.
add_outline:
Draw each point with a thin ring (as ``sc.pl.embedding``'s own
``add_outline`` does) so cells stand out from the background and from
each other. On by default since the additive color blend otherwise
makes dense, similarly-colored regions hard to resolve into
individual cells.
outline_width, outline_color:
Same meaning as the identically-named ``sc.pl.embedding`` params:
``outline_width`` is ``(outline_size, gap_size)`` as a fraction of
the point radius, ``outline_color`` is ``(outline_color, gap_color)``.
show_legend:
Draw a small bivariate color-key inset showing how the blend maps to
each gene's expression level.
config:
Active :class:`~scplotkit.PlotConfig`. Uses defaults if omitted.
output_dir:
Directory figures are written under.
show:
Display the figure inline (e.g. in a Jupyter notebook) as it's drawn.
"""
config = PlotConfig.load(config)
require_columns(adata, mask_column)
figure_name = figure_name or f"{gene1}_{gene2}_coexpression"
dir_name = figure_name.replace(" ", "_")
norm1 = _normalized_expression(_expression_values(adata, gene1, use_raw, layer), vmax_percentile)
norm2 = _normalized_expression(_expression_values(adata, gene2, use_raw, layer), vmax_percentile)
on_mask = (norm1 >= threshold) | (norm2 >= threshold)
if mask_column and mask_values:
on_mask &= adata.obs[mask_column].isin(mask_values).to_numpy()
print(f"Masking applied: showing {gene1}/{gene2} co-expression only in {mask_values}")
rgb1, rgb2 = np.array(to_rgb(color1)), np.array(to_rgb(color2))
grey_color = grey_color or config.general.get("fallback_color", "lightgrey")
colors = np.tile(np.array(to_rgb(grey_color)), (adata.n_obs, 1))
blended = np.clip(np.outer(norm1, rgb1) + np.outer(norm2, rgb2), 0, 1)
colors[on_mask] = blended[on_mask]
coords = adata.obsm[f"X_{basis}"][:, :2]
point_size = point_size if point_size is not None else max(120000 / adata.n_obs, 3.0)
# draw grey (off/masked-out) cells first, colored cells on top so the signal isn't buried
order = np.argsort(on_mask.astype(int), kind="stable")
general = config.general
# Extra figure width reserves room for the legend outside the plot without
# shrinking the UMAP itself; the main axes gets a fixed position so
# tight_layout (which mishandles axes anchored outside [0, 1]) never touches it.
fig = plt.figure(figsize=(7.2, 6))
ax = fig.add_axes((0.03, 0.05, 0.74, 0.88))
if add_outline:
# Same construction as scanpy's own add_outline: two flat-colored rings
# (radius grown outward from the point radius by outline_width, as
# fractions of it) drawn behind the point, largest first, so a
# bg_color ring shows with a thin gap_color gap before the point.
bg_width, gap_width = outline_width
bg_color, gap_color = outline_color
point_radius = np.sqrt(point_size)
bg_size = (point_radius + point_radius * bg_width * 2) ** 2
gap_size = (point_radius + point_radius * gap_width * 2) ** 2
ax.scatter(coords[order, 0], coords[order, 1], s=bg_size, c=bg_color, linewidths=0, rasterized=True)
ax.scatter(coords[order, 0], coords[order, 1], s=gap_size, c=gap_color, linewidths=0, rasterized=True)
ax.scatter(coords[order, 0], coords[order, 1], c=colors[order], s=point_size, linewidths=0, rasterized=True)
ax.set_xticks([])
ax.set_yticks([])
for spine in ax.spines.values():
spine.set_visible(False)
ax.set_title(figure_name, fontsize=general["title_fontsize"], fontweight=general["title_fontweight"])
if show_legend:
legend_ax = fig.add_axes((0.82, 0.09, 0.15, 0.15))
grid_n = 60
g1, g2 = np.meshgrid(np.linspace(0, 1, grid_n), np.linspace(0, 1, grid_n))
legend_rgb = np.clip(g1[..., None] * rgb1 + g2[..., None] * rgb2, 0, 1)
legend_ax.imshow(legend_rgb, origin="lower", extent=(0, 1, 0, 1), aspect="auto")
legend_ax.set_xticks([])
legend_ax.set_yticks([])
for spine in legend_ax.spines.values():
spine.set_color("black")
spine.set_linewidth(1.0)
legend_ax.set_xlabel(gene1, fontsize=7, color=color1, fontweight="bold", labelpad=2)
legend_ax.set_ylabel(gene2, fontsize=7, color=color2, fontweight="bold", labelpad=2)
save_figure(fig, output_dir, dir_name, f"{dir_name}_{gene1}_{gene2}_coexpression.png", config)
if show:
plt.show()
return fig
def _fast_kde_grid(x: np.ndarray, y: np.ndarray, grid_size: int, xlim: tuple, ylim: tuple, bw_scale: float):
"""Approximate a 2D Gaussian KDE in O(n) via binning + blur instead of O(n^2) pairwise evaluation.
``scipy.stats.gaussian_kde`` (what ``sc.tl.embedding_density`` uses)
evaluates the kernel at every point against every other point, which is
quadratic in cell count and becomes unusable well before atlas scale (tens
of thousands of cells per group). Binning onto a fixed grid and blurring
with a Gaussian filter gives a visually equivalent smooth density surface
in time that only depends on ``grid_size``, not on the number of cells.
"""
xedges = np.linspace(xlim[0], xlim[1], grid_size + 1)
yedges = np.linspace(ylim[0], ylim[1], grid_size + 1)
hist, _, _ = np.histogram2d(x, y, bins=[xedges, yedges])
n = len(x)
bw_factor = bw_scale * (n ** (-1.0 / 6.0)) if n > 1 else 1.0 # Scott's rule for a 2D KDE
bin_w_x = (xlim[1] - xlim[0]) / grid_size
bin_w_y = (ylim[1] - ylim[0]) / grid_size
sigma_x = max((bw_factor * np.std(x)) / bin_w_x, 0.5)
sigma_y = max((bw_factor * np.std(y)) / bin_w_y, 0.5)
smoothed = gaussian_filter(hist, sigma=(sigma_x, sigma_y))
smoothed -= smoothed.min()
peak = smoothed.max()
if peak > 0:
smoothed /= peak
xcenters = (xedges[:-1] + xedges[1:]) / 2
ycenters = (yedges[:-1] + yedges[1:]) / 2
return smoothed, xcenters, ycenters
[docs]
def embedding_density(
adata,
groupby: str | None = None,
groups: Sequence[str] | None = None,
figure_name: str | None = None,
basis: str = "umap",
grid_size: int = 200,
bw_scale: float = 0.5,
density_threshold: float = 0.01,
cmap: str = "YlOrBr",
point_size: float | None = None,
grey_color: str | None = None,
grey_alpha: float = 1.0,
ncols: int = 4,
panel_size: tuple = (4.5, 4.5),
key_added: str | None = None,
config: PlotConfig | None = None,
output_dir: str | Path = "figures",
save_name: str = "embedding_density",
show: bool = False,
):
"""Fancy, fast per-cell density plot over an embedding, optionally per group.
Every cell is drawn as a dot and every dot is filled in -- colored smoothly
and continuously by its local density (no discrete contour bands, no
outline layer sitting on top, nothing colored outside of where cells
actually are): cells below ``density_threshold`` are flat grey (no signal
there), the rest are colored along ``cmap`` by how dense the embedding is
at that exact point.
When ``groupby`` is given, one independently-scaled panel is drawn per
category, on shared axis limits so panels stay visually comparable --
each panel colors *every* cell by that category's density field, so the
full embedding shape stays visible with the category's hotspots picked
out in color.
Unlike :func:`scanpy.tl.embedding_density` (which evaluates
``scipy.stats.gaussian_kde`` pairwise -- O(n^2) in cell count, and
impractically slow for atlas-scale groups), this bins cells onto a fixed
grid and smooths with a Gaussian filter -- O(n + grid_size^2), so runtime
stays flat as cell counts grow. Densities are also written to
``adata.obs[key_added]`` (scaled 0-1 within each group, evaluated at each
cell's own position), mirroring scanpy's output so downstream code can
rely on either.
Parameters
----------
adata:
Annotated data with a computed embedding (``adata.obsm['X_<basis>']``).
groupby:
``adata.obs`` column to compute one density panel per category.
Omit for a single overall density panel.
groups:
Subset/order of categories to plot when ``groupby`` is given.
Defaults to all categories.
figure_name:
Used as the figure's suptitle. Defaults to a description of ``basis``
(and ``groupby``, if given).
grid_size:
Resolution of the density grid per axis. Higher is smoother but
slower; 200 is a good default even for very large atlases.
bw_scale:
Multiplier on the Scott's-rule bandwidth; increase to smooth the
density surface further, decrease to sharpen it.
density_threshold:
Scaled density (0-1) below which a cell is drawn flat grey instead
of colored.
point_size:
Marker area per cell. Defaults to scanpy's own ``120000 / n_obs``
heuristic so it scales sensibly with atlas size.
grey_color:
Fill color for cells below ``density_threshold``. Defaults to the
config's ``general.fallback_color``.
grey_alpha:
Opacity of the below-threshold background cells, so the colored
density signal on top stays the visual focus.
key_added:
``adata.obs`` column the per-cell scaled density is written to.
Defaults to ``f"{basis}_density"`` or ``f"{basis}_density_{groupby}"``.
config:
Active :class:`~scplotkit.PlotConfig`. Uses defaults if omitted.
output_dir:
Directory figures are written under.
show:
Display the figure inline (e.g. in a Jupyter notebook) as it's drawn.
"""
config = PlotConfig.load(config)
require_columns(adata, groupby)
coords = adata.obsm[f"X_{basis}"][:, :2]
all_x, all_y = coords[:, 0], coords[:, 1]
all_points = np.column_stack([all_x, all_y])
pad_x = (all_x.max() - all_x.min()) * 0.05 or 1.0
pad_y = (all_y.max() - all_y.min()) * 0.05 or 1.0
xlim = (all_x.min() - pad_x, all_x.max() + pad_x)
ylim = (all_y.min() - pad_y, all_y.max() + pad_y)
point_size = point_size if point_size is not None else max(120000 / adata.n_obs, 3.0)
grey_color = grey_color or config.general.get("fallback_color", "lightgrey")
if groupby:
ensure_categorical(adata, groupby)
panels = list(groups) if groups is not None else list(adata.obs[groupby].cat.categories)
else:
panels = ["All cells"]
key_added = key_added or (f"{basis}_density_{groupby}" if groupby else f"{basis}_density")
density_values = np.zeros(adata.n_obs)
ncols_eff = min(ncols, len(panels))
nrows = math.ceil(len(panels) / ncols_eff)
fig, axes = plt.subplots(
nrows, ncols_eff, figsize=(panel_size[0] * ncols_eff, panel_size[1] * nrows), squeeze=False
)
last_cs = None
for i, panel in enumerate(panels):
ax = axes[i // ncols_eff][i % ncols_eff]
mask = np.ones(adata.n_obs, dtype=bool) if not groupby else (adata.obs[groupby] == panel).to_numpy()
px, py = all_x[mask], all_y[mask]
if mask.sum() >= 5:
smoothed, xcenters, ycenters = _fast_kde_grid(px, py, grid_size, xlim, ylim, bw_scale)
interp = RegularGridInterpolator((xcenters, ycenters), smoothed, bounds_error=False, fill_value=0.0)
density_all_cells = np.clip(interp(all_points), 0.0, 1.0)
else:
density_all_cells = np.zeros(adata.n_obs)
density_all_cells[mask] = 1.0
density_values[mask] = density_all_cells[mask]
# Grey backdrop for every cell, drawn under the colored ones regardless of
# density_threshold -- with real (non-uniform) embeddings the vast majority
# of cells clear a low threshold, so gating the backdrop on it left it
# empty in practice. Sized larger than the foreground point so a ring of
# it stays visible around every colored cell too, instead of being fully
# hidden underneath (an alpha on the colored points would do that too, but
# by blending into the density colormap it would misrepresent the values).
ax.scatter(
all_x, all_y, s=point_size * 2.6, c=grey_color, alpha=grey_alpha,
linewidths=0, zorder=1, rasterized=True,
)
colored_idx = np.where(density_all_cells >= density_threshold)[0]
if len(colored_idx) > 0:
order = colored_idx[np.argsort(density_all_cells[colored_idx])] # highest density drawn on top
last_cs = ax.scatter(
all_x[order], all_y[order], s=point_size, c=density_all_cells[order], cmap=cmap,
vmin=density_threshold, vmax=1.0, linewidths=0, zorder=4, rasterized=True,
)
ax.set_xlim(xlim)
ax.set_ylim(ylim)
ax.set_xticks([])
ax.set_yticks([])
for spine in ax.spines.values():
spine.set_visible(False)
ax.set_title(str(panel), fontsize=config.general["title_fontsize"], fontweight="bold")
for j in range(len(panels), nrows * ncols_eff):
axes[j // ncols_eff][j % ncols_eff].axis("off")
adata.obs[key_added] = density_values
# Reserve a fixed vertical band for the suptitle so it doesn't collide with the per-panel
# titles sitting right at the top of each axes -- the default subplot top margin doesn't
# leave room for a second line of (bold, larger) title text above them.
fig_height_in = panel_size[1] * nrows
fig.subplots_adjust(top=1 - min(0.6 / fig_height_in, 0.25))
suptitle = figure_name or (f"Cell density by {groupby} ({basis})" if groupby else f"Cell density ({basis})")
fig.suptitle(suptitle, fontsize=config.general["title_fontsize"] + 2, fontweight="bold")
if last_cs is not None:
fig.colorbar(
last_cs, ax=axes.ravel().tolist(), shrink=0.6, pad=0.02, label="Relative density (scaled per panel)",
)
if save_name is not None:
save_figure(fig, output_dir, "embedding_density", f"{save_name}.png", config)
if show:
plt.show()
return fig
