API Reference

Mievformer learns the probabilistic relationship between continuous cell states and local tissue contexts. The functions below identify niches and use the same learned relationship to characterize niche preferences, estimate conditional population probabilities, and analyze associated gene expression.

Core Model Functions

optimize_nicheformer trains the model, evaluates reference cell-state profiles, and computes the niche representation, neighbors, UMAP, and niche clusters. Multiple values in batch_key automatically enable sample-conditioned training. Score-function weights for downstream analysis are calculated during optimization; calculate_wb_ez can reload them from an existing model.

mievformer.optimize_nicheformer(adata, model_path, ngpu=1, batch_size=512, max_epochs=1000, neighbor_num=100, latent_dim=20, kld_ld=0.05, pent_ld=0.05, dist_space='latent', cellrep_key='X_pca', batch_key=None, batch_correct='auto', representation_mode='auto', ca_reference_num=None, ca_n_components='auto', ca_reference_seed=0, ca_device=None, niche_n_neighbors=15, leiden_resolution=0.5, umap_min_dist=0.1, random_state=0)[source]

Train Mievformer and create its standard niche representation.

By default, single-slice data use reference-probability correspondence analysis (CA), while data with multiple values in batch_key use sample-conditional CA. The selected CA representation is stored in both obsm[‘reference_probability_ca’] and obsm[‘e’]; the raw model embedding is preserved in obsm[‘mievformer_raw_e’]. Neighbors, UMAP, and obs[‘leiden_e’] are calculated from the standard obsm[‘e’].

Set representation_mode=’raw’ only when reproducing the legacy raw-embedding workflow. Multi-slice standard analysis requires batch_correct=’auto’ (the default) or True.

Parameters:

  • adata (anndata.AnnData) – Spatial data containing obsm[‘spatial’] and the representation selected by cellrep_key.

  • model_path (path-like) – Destination for the trained state dictionary.

  • batch_key (str, optional) – obs column identifying spatial slices or samples.

  • batch_correct (bool or {'auto'}, default 'auto') – auto enables sample conditioning when batch_key has multiple values.

  • representation_mode ({'auto', 'raw'}, default 'auto') – Standard CA selection or explicit legacy raw-embedding mode.

  • ca_reference_num (int, optional) – Number of CA reference cells. The default is adaptive up to 1000.

  • ca_n_components (int or {'auto'}, default 'auto') – CA dimension. Automatic selection uses the mean relative eigengap.

  • niche_n_neighbors (int, default 15) – Number of neighbors for the CA-based Scanpy graph. This is independent of neighbor_num, which controls spatial context during training.

Returns:

A copy containing raw and standard representations, score-function weights, UMAP coordinates, niche clusters, and provenance metadata.

Return type:

anndata.AnnData

mievformer.calculate_wb_ez(adata, model_path, batch_key=None, batch_correct='auto', neighbor_num=100, latent_dim=20, cellrep_key='X_pca')[source]

Calculate the Mievformer score-function weights.

The function adds obsm[‘w_e’], obsm[‘w_z’], and obsm[‘b_z’]. If CA is already the default obsm[‘e’], the preserved obsm[‘mievformer_raw_e’] is used as the distributor input.

Niche Density Ratio and Membership

These functions use density-ratio-based Bayesian inversion to characterize which niches are preferred by a cell’s molecular state. The probabilistic formulation is described in Probabilistic Coupling of Cell States and Microenvironments.

calculate_niche_density_ratio stores scores normalized over sampled reference contexts in obsm['dist_e']. By default, reference contexts are sampled equally from niche clusters exceeding min_ratio, rather than in proportion to their tissue abundance. calculate_niche_cluster_membership then averages those scores within each reference cluster and stores the relative membership profile in obsm['dist_e_agg']. Its rows generally do not sum to one. These profiles support comparison and clustering of niche preferences; they are not directly the tissue-prior-weighted posterior cluster probabilities defined in the overview. That interpretation requires accounting for reference sampling, niche prior weights, and normalization.

mievformer.calculate_niche_density_ratio(adata, ref_num=1000, stratify_key='leiden_e', min_ratio=0.01, ref_adata=None)[source]

Compute per-cell density ratios over a panel of reference niches.

For each cell \(i\) and reference niche \(j\) drawn by stratified sampling on stratify_key, the log density ratio is

\[\log r_{ij} = \log p(e_j \mid z_i) - \log p(e_j) = (w_z(z_i)^\top w_e(e_j) + b_z(z_i)) - \log \sum_{k \in \mathrm{ref}} \exp(w_z(z_k)^\top w_e(e_j) + b_z(z_k)).\]

The matrix is then softmax-normalized per cell over reference niches, so each row of adata.obsm['dist_e'] is a probability distribution over the sampled reference niches that emphasizes niches whose environment becomes more likely under the cell’s state than under the marginal.

Parameters:

  • adata (anndata.AnnData) – Annotated data matrix containing w_e, w_z, and b_z in obsm (produced by calculate_wb_ez()).

  • ref_num (int, optional) – Number of reference niches to sample. Default is 1000.

  • stratify_key (str, optional) – Key in adata.obs to use for stratified sampling of reference niches. Default is ‘leiden_e’.

  • min_ratio (float, optional) – Clusters with frequency below this fraction are dropped from stratified sampling. Default is 0.01.

  • ref_adata (anndata.AnnData, optional) – External reference. If None, a subset of adata is used.

Returns:

Updated with obsm['dist_e'] (softmax-normalized density ratios of shape (n_cells, ref_num)) and uns['dist_e']['ref_obs'] (obs names of the sampled reference niches). The dist_e key name is preserved for backward compatibility with existing h5ad artifacts.

Return type:

anndata.AnnData

mievformer.calculate_niche_cluster_membership(adata, cluster_key='leiden_e')[source]

Aggregate per-cell density ratios into a soft membership over niche clusters.

Averages the columns of adata.obsm['dist_e'] within each value of adata.obs[cluster_key] (typically leiden_e niche clusters), yielding adata.obsm['dist_e_agg'] of shape (n_cells, n_niche_clusters): entry [i, c] is the mean density ratio \(p(e \mid z_i)/p(e)\) evaluated at reference cells in cluster c, interpretable as a soft assignment of cell i to niche cluster c.

Parameters:

  • adata (anndata.AnnData) – Annotated data matrix containing obsm['dist_e'] (see calculate_niche_density_ratio()). If absent, it is computed with defaults.

  • cluster_key (str, optional) – Key in adata.obs containing niche cluster labels. Default is ‘leiden_e’.

Returns:

Updated with obsm['dist_e_agg']: per-cell niche-cluster membership (columns are niche cluster labels). The dist_e_agg key name is preserved for backward compatibility with existing h5ad artifacts used by figure scripts.

Return type:

anndata.AnnData

Downstream Analysis

estimate_population_density estimates the conditional probability of a target population in each local tissue context and stores it in obs['<group>_density']. Here, “density” denotes a population probability, not a number of cells per unit area. Estimate it on the full tissue before subsetting cells for expression analysis.

analyze_density_correlation relates those estimates to expression in the cells supplied to it. Restricting the input to a query population identifies expression variation associated with target-enriched contexts within that population. analyze_niche_membership groups cells by their relative niche membership profiles and visualizes them.

mievformer.estimate_population_density(adata, group, cluster_key, max_cell_num=1000)[source]

Estimate the density (existence probability) of a specific cell population in each microenvironment.

By integrating \(P(z|e)\) over all cell states belonging to a specific cell population, this function obtains the density of that population in microenvironment \(e\).

Parameters:

  • adata (anndata.AnnData) – Annotated data matrix.

  • group (str) – The label of the cell population (e.g., a specific cell type) to estimate density for.

  • cluster_key (str) – Key in adata.obs containing the cell type/cluster labels.

  • max_cell_num (int, optional) – Maximum number of cells to sample from the group for density estimation. Default is 1000.

Returns:

The input AnnData object updated with a new column in obs (e.g., {group}_density) representing the estimated density of the specified population for each cell’s microenvironment.

Return type:

anndata.AnnData

mievformer.analyze_density_correlation(adata, density_col, gene_list=None, file_path=None)[source]

Analyze the correlation between estimated cell population density and gene expression.

This analysis helps identify gene expression signatures associated with colocalization with specific cell populations. For example, identifying genes upregulated in tumor cells when they colocalize with endothelial cells.

Parameters:

  • adata (anndata.AnnData) – Annotated data matrix containing expression data and the density column.

  • density_col (str) – Name of the column in adata.obs containing the estimated density values.

  • gene_list (list of str, optional) – List of genes to include in the correlation analysis. If None, uses all genes in adata.var_names.

  • file_path (str, optional) – Path to save the visualization plot (bar plot of top/bottom correlated genes). If None, the plot is not saved.

Returns:

A Series containing the correlation coefficients for each gene, indexed by gene name.

Return type:

pandas.Series

mievformer.analyze_niche_membership(adata, n_clusters=15, file_path=None)[source]

Cluster cells by their niche-cluster membership vectors and visualize the result.

Uses adata.obsm['dist_e_agg'] (per-cell soft membership over niche clusters produced by calculate_niche_cluster_membership()) as the feature space, performs Ward hierarchical clustering to partition cells into n_clusters groups, and draws a clustermap of the membership matrix with row-color annotations.

Parameters:

  • adata (anndata.AnnData) – Annotated data matrix containing obsm['dist_e_agg'].

  • n_clusters (int, optional) – Number of cell clusters to form. Default is 15.

  • file_path (str, optional) – Path to save the resulting clustermap image. If None, the plot is not saved.

Returns:

The input AnnData with obs['niche_composition_cluster'] added (cell cluster labels). The niche_composition_cluster key name is preserved for backward compatibility with existing h5ad artifacts.

Return type:

anndata.AnnData

Representation and Stored Outputs

As one step in constructing the niche representation, the workflow applies correspondence analysis (CA) to reference cell-state probability profiles before computing neighbors, UMAP, and Leiden clusters. With multiple samples, profiles are normalized separately within sample-conditioned reference blocks, with balanced sampling and equal probability mass per sample.

Outputs of optimize_nicheformer

AnnData key

Meaning

obsm['e']

Standard niche representation, also stored in obsm['reference_probability_ca']

obsm['mievformer_raw_e']

Context embedding produced by the neural network

obsm['X_umap']

UMAP of the standard niche representation

obs['leiden_e']

Inferred niche labels

obsm['w_e'], obsm['w_z'], obsm['b_z']

Score-function terms used for density-ratio-based analysis

uns['mievformer_default_representation']

Processing strategy and provenance

uns['reference_probability_ca']

Reference selection, transformation basis, and probability audits

ca_reference_num controls the reference set size (adaptive up to 1,000 by default), and ca_n_components controls the representation dimension (automatic by default). representation_mode='raw' selects the legacy raw embedding for method comparisons. Standard multi-slice analysis requires batch_correct='auto' or True and retains the sample mapping in obsm['batch_one_hot'] and uns['mievformer_batch_contract'].