EpiPack API Reference

EpiPack API Reference#

This page provides a concise reference for the main user-facing classes and functions in epipackpy.model.

For complete end-to-end workflows, see the corresponding EpiPack tutorials.

Note: The examples below are minimal usage templates. Replace the placeholder arrays and labels with objects prepared from your own dataset. The reference was prepared from the main branch of the EpiPack repository.

[ ]:

import epipackpy as epk
import numpy as np
import torch

print(epk.__version__)

Common input conventions#

Most EpiPack model interfaces operate on NumPy arrays:

peak_count: cell-by-peak accessibility matrix with shape (n_cells, n_peaks).
gene_score: cell-by-gene activity score matrix with shape (n_cells, n_genes).
peak_embedding: cell-by-latent peak embedding matrix with shape (n_cells, z_dim).
batch_id: integer batch labels with shape (n_cells,).
reference_embedding: trained reference latent representation with shape (n_reference_cells, z_dim).
reference_label: integer-encoded reference cell-type labels with shape (n_reference_cells,).

For query mapping, the query gene activity matrix should use the same gene features and feature order as the pretrained reference model.

1. Peak embedding model#

`epk.model.Peak_Model`#

Learns a low-dimensional embedding directly from a dataset-specific peak accessibility matrix.

Constructor#

Peak_Model(
    count_enhancer,
    layer_num=2,
    batch_size=512,
    hidden_dim=256,
    dropout_rate=0.1,
    z_dim=50,
    device="auto",
    lib_size=True,
    region_factor=False,
    denosing=False,
)

Parameters#

``count_enhancer`` (array-like): Cell-by-peak accessibility matrix.
``layer_num`` (int, default 2): Number of hidden layers in the encoder and decoder.
``batch_size`` (int, default 512): Number of cells per training mini-batch.
``hidden_dim`` (int, default 256): Width of hidden layers.
``dropout_rate`` (float, default 0.1): Dropout probability.
``z_dim`` (int, default 50): Dimension of the learned peak embedding.
``device`` ({"auto", "gpu", "cpu"}, default "auto"): Computation device.
``lib_size`` (bool, default True): Whether to include the log library size in reconstruction.
``region_factor`` (bool, default False): Whether to learn region-specific scaling factors.
``denosing`` (bool, default False): Denoising option stored by the model.

Main methods#

`train_model(nepochs=50, weight_decay=5e-4, learning_rate=1e-4)`#

Trains the peak autoencoder.

`get_z()`#

Returns the learned peak embedding as a CPU PyTorch tensor with shape (n_cells, z_dim).

[ ]:

# peak_count: NumPy array or array-like object with shape (n_cells, n_peaks)
peak_model = epk.model.Peak_Model(
    count_enhancer=peak_count,
    z_dim=50,
    batch_size=512,
    device="auto",
)

peak_model.train_model(
    nepochs=50,
    learning_rate=1e-4,
    weight_decay=5e-4,
)

peak_embedding = peak_model.get_z().detach().numpy()

2. PEIVI reference construction#

`epk.model.PEIVI`#

Constructs a harmonized reference representation from aligned gene activity scores while using precomputed peak embeddings as regulatory constraints.

Constructor#

PEIVI(
    count,
    peak_emb,
    batch_id,
    layer_num=2,
    batch_size=512,
    hidden_dim=256,
    dropout_rate=0.1,
    z_dim=50,
    reg_mmd=1,
    reg_kl=1e-10,
    reg_rec=1,
    reg_z_l2=0.1,
    reg_z_l2_2=0,
    device="auto",
    use_layer_norm=False,
    use_batch_norm=True,
    rec_type="MSE",
)

Required parameters#

``count`` (array-like): Cell-by-gene activity score matrix.
``peak_emb`` (array-like): Precomputed cell-by-peak embedding matrix. Its number of rows must match count.
``batch_id`` (array-like): Integer batch labels for all cells.

Architecture parameters#

``layer_num`` (int, default 2): Number of hidden layers.
``batch_size`` (int, default 512): Number of cells per mini-batch.
``hidden_dim`` (int, default 256): Width of hidden layers.
``dropout_rate`` (float, default 0.1): Dropout probability.
``z_dim`` (int, default 50): Dimension of the PEIVI latent representation.
``use_layer_norm`` (bool, default False): Enable layer normalization.
``use_batch_norm`` (bool, default True): Enable batch normalization.

Loss and prior parameters#

``reg_mmd`` (float, default 1): Weight of cross-batch MMD alignment.
``reg_kl`` (float, default 1e-10): Weight of KL regularization.
``reg_rec`` (float, default 1): Weight of gene-score reconstruction.
``reg_z_l2`` (float, default 0.1): Peak-embedding constraint weight during pretraining.
``reg_z_l2_2`` (float, default 0): Peak-embedding constraint weight after the pretraining stage.
``rec_type`` ({"MSE", "NB"}, default "MSE"): Gene-score reconstruction loss.
``device`` ({"auto", "gpu", "cpu"}, default "auto"): Computation device.

Main methods#

`train_model(nepochs=200, clamp=0.01, weight_decay=5e-4, learning_rate=1e-4, pre_train_epoch=100)`#

Trains PEIVI in two stages. Before pre_train_epoch, the peak-embedding constraint is active and MMD is disabled; afterwards, the second-stage loss weights are used and MMD alignment is enabled.

`get_latent()`#

Returns the harmonized PEIVI latent representation as a NumPy array.

[ ]:

reference_model = epk.model.PEIVI(
    count=reference_gene_score,
    peak_emb=reference_peak_embedding,
    batch_id=reference_batch_id,
    z_dim=50,
    batch_size=512,
    reg_mmd=1.0,
    reg_z_l2=0.1,
    reg_z_l2_2=0.0,
    rec_type="MSE",
    device="auto",
)

reference_model.train_model(
    nepochs=200,
    pre_train_epoch=100,
    learning_rate=1e-4,
)

reference_embedding = reference_model.get_latent()

Saving and restoring a PEIVI model#

PEIVI inherits from torch.nn.Module; therefore, standard PyTorch serialization can be used.

[ ]:

# Save model parameters
torch.save(reference_model.state_dict(), "peivi_reference.pt")

# Restore parameters into a model initialized with the same architecture
restored_model = epk.model.PEIVI(
    count=reference_gene_score,
    peak_emb=reference_peak_embedding,
    batch_id=reference_batch_id,
    z_dim=50,
)
restored_model.load_state_dict(
    torch.load("peivi_reference.pt", map_location=restored_model.device_use)
)

3. PEIVI query mapping#

`epk.model.PEIVI_MAPPING`#

Maps a new query dataset into a pretrained PEIVI reference space using query gene activity scores, query peak embeddings, and the fixed reference embedding.

Constructor#

PEIVI_MAPPING(
    promoter_dt,
    enhancer_z,
    batch_id,
    ref_embedding,
    layer_num=2,
    batch_size=256,
    hidden_dim=256,
    dropout_rate=0,
    z_dim=50,
    reg_mmd=1,
    reg_mmd_inter=0.1,
    reg_kl=1e-10,
    reg_rec=1,
    reg_z_l2=0.01,
    device="auto",
    use_layer_norm=False,
    use_batch_norm=True,
    rec_type="MSE",
)

Required parameters#

``promoter_dt`` (array-like): Query gene activity score matrix aligned to the reference gene features.
``enhancer_z`` (array-like): Precomputed query peak embedding.
``batch_id`` (array-like): Integer batch labels for query cells.
``ref_embedding`` (array-like): Pretrained reference PEIVI embedding.

Mapping-specific parameters#

``reg_mmd`` (float, default 1): Within-query batch-alignment weight.
``reg_mmd_inter`` (float, default 0.1): Reference-query alignment weight.
``reg_z_l2`` (float, default 0.01): Query peak-embedding constraint weight.

The remaining architecture, prior, reconstruction, and device parameters have the same interpretation as in PEIVI.

Main methods#

`train_model(nepochs=50, clamp=0.01, weight_decay=5e-4, learning_rate=1e-4)`#

Fine-tunes the query mapping model.

`get_latent()`#

Returns the mapped query representation as a NumPy array.

`epk.model.query_model_initial`#

Initializes a query mapping model from pretrained PEIVI parameters and freezes selected shared layers.

query_model_initial(
    query_model=None,
    ref_model_param=None,
    query_batch_num=None,
    gene_feature_num=3000,
    latent_embedding_dim=50,
)

Parameters#

``query_model`` (PEIVI_MAPPING): Newly initialized query mapping model.
``ref_model_param`` (state_dict): State dictionary from the pretrained reference PEIVI model.
``query_batch_num`` (int): Number of query batches.
``gene_feature_num`` (int, default 3000): Number of aligned gene features.
``latent_embedding_dim`` (int, default 50): Reference latent dimension.

Returns#

``PEIVI_MAPPING``: Initialized query model with selected layers frozen.

[ ]:

query_model = epk.model.PEIVI_MAPPING(
    promoter_dt=query_gene_score,
    enhancer_z=query_peak_embedding,
    batch_id=query_batch_id,
    ref_embedding=reference_embedding,
    z_dim=50,
    batch_size=256,
    device="auto",
)

query_model = epk.model.query_model_initial(
    query_model=query_model,
    ref_model_param=reference_model.state_dict(),
    query_batch_num=len(np.unique(query_batch_id)),
    gene_feature_num=query_gene_score.shape[1],
    latent_embedding_dim=50,
)

query_model.train_model(
    nepochs=50,
    learning_rate=1e-4,
)

query_embedding = query_model.get_latent()

4. Cell-type classifier#

`epk.model.Classifier`#

Trains a weighted-sampling metric-learning classifier on the reference PEIVI embedding and transfers labels to query cells.

Constructor#

Classifier(
    ref_latent_emb,
    ref_label,
    hidden_num,
    z_dim=30,
    dropout_rate=0.1,
    layer_num=1,
    batch_size=128,
    batchnorm=True,
    device="auto",
)

Parameters#

``ref_latent_emb`` (array-like): Reference PEIVI embedding.
``ref_label`` (array-like): Integer-encoded reference labels. Labels should be contiguous non-negative integers.
``hidden_num`` (int): Hidden-layer width.
``z_dim`` (int, default 30): Dimension of the classification embedding.
``dropout_rate`` (float, default 0.1): Dropout probability.
``layer_num`` (int, default 1): Number of hidden layers.
``batch_size`` (int, default 128): Training mini-batch size.
``batchnorm`` (bool, default True): Enable batch normalization.
``device`` ({"auto", "gpu", "cpu"}, default "auto"): Computation device.

The training loader uses inverse-frequency weighted sampling to improve representation of rare classes.

Main methods#

`train_model(nepochs=50, weight_decay=5e-4, learning_rate=1e-4, inner_reg=1)`#

Trains the angular classifier together with the prototype/center compactness objective.

``inner_reg`` controls the weight of the within-class center loss.

`get_z(input)`#

Predicts integer class labels and returns the classification-space embedding.

Returns#

A tuple:

predicted_labels: Python list of integer class predictions.
classification_embedding: PyTorch tensor containing the learned classification-space representation.

[ ]:

classifier = epk.model.Classifier(
    ref_latent_emb=reference_embedding,
    ref_label=reference_label_encoded,
    hidden_num=128,
    z_dim=30,
    batch_size=128,
    device="auto",
)

classifier.train_model(
    nepochs=50,
    learning_rate=1e-4,
    inner_reg=1.0,
)

reference_prediction, reference_class_embedding = classifier.get_z(reference_embedding)
query_prediction, query_class_embedding = classifier.get_z(query_embedding)

reference_class_embedding = reference_class_embedding.detach().cpu().numpy()
query_class_embedding = query_class_embedding.detach().cpu().numpy()

5. Global OOR detector#

`epk.model.Global_oor_detector`#

Detects query cells that are inconsistent with the class-conditional reference distributions in the classifier space.

Constructor#

Global_oor_detector(
    train_data,
    query_data,
    ref_label,
    predicted_label,
)

Parameters#

``train_data`` (array-like): Reference classification-space embedding.
``query_data`` (array-like): Query classification-space embedding.
``ref_label`` (array-like): Integer-encoded reference labels.
``predicted_label`` (array-like): Preliminary integer labels assigned to query cells by the classifier.

For each reference class, the detector estimates a Gaussian mean and covariance matrix. It then evaluates the squared Mahalanobis distance of each query cell to its preliminarily assigned class using a chi-square distribution.

Main method#

`annotate(confidence_threshold=1e-3)`#

Parameters#

``confidence_threshold`` (float, default 1e-3): Query cells with a p-value below this threshold are rejected as global OOR.

Returns#

A tuple:

reject_cell_list: Indices of query cells called global OOR.
prob_score: Cell-level chi-square p-values.

[ ]:

global_detector = epk.model.Global_oor_detector(
    train_data=reference_class_embedding,
    query_data=query_class_embedding,
    ref_label=reference_label_encoded,
    predicted_label=np.asarray(query_prediction),
)

global_oor_indices, global_pvalues = global_detector.annotate(
    confidence_threshold=1e-3
)

global_oor_mask = np.zeros(query_embedding.shape[0], dtype=bool)
global_oor_mask[global_oor_indices] = True

6. Local OOR detector#

`epk.model.local_oor_detector`#

Detects continuous query-associated state shifts on a mutual k-nearest-neighbor graph of reference and query cells.

Signature#

local_oor_detector(
    Z,
    y,
    batches,
    *,
    k=20,
    celltypes=None,
    tail_direction_mu=None,
    train_epochs=50,
    lr=1e-3,
    weight_align=1.0,
    weight_margin=0.2,
    weight_lap=0.05,
    lam=0.9,
    T=30,
    R_perm=800,
    alpha_fdr=0.10,
    seed=42,
    device=None,
)

Required parameters#

``Z`` (numpy.ndarray): Joint reference-query embedding with shape (n_cells, n_features).
``y`` (numpy.ndarray): Group indicator with 0 for reference/control cells and 1 for query cells.
``batches`` (numpy.ndarray): Batch labels for all cells.

Graph and edge-feature parameters#

``k`` (int, default 20): Number of neighbors used to construct the mutual kNN graph.
``celltypes`` (numpy.ndarray or None, default None): Optional cell-type labels included as an edge feature.
``tail_direction_mu`` (numpy.ndarray or None, default None): Optional direction vector. If omitted, the query-minus-reference mean difference is used.

Kernel-training parameters#

``train_epochs`` (int, default 50): Number of edge-kernel training epochs.
``lr`` (float, default 1e-3): Edge-kernel learning rate.
``weight_align`` (float, default 1.0): Weight of Gaussian-kernel alignment.
``weight_margin`` (float, default 0.2): Weight of the directional margin term.
``weight_lap`` (float, default 0.05): Weight of the graph smoothness term.

BRP and statistical-testing parameters#

``lam`` (float, default 0.9): BRP propagation coefficient.
``T`` (int, default 30): Number of BRP iterations.
``R_perm`` (int, default 800): Number of permutations used for empirical p-values.
``alpha_fdr`` (float, default 0.10): Benjamini–Hochberg significance threshold.
``seed`` (int, default 42): Random seed.
``device`` (str or None, default None): "cuda" or "cpu"; automatically selected when omitted.

Returns#

A dictionary containing:

``K``: Learned row-normalized sparse graph kernel.
``rows``, ``cols``: Mutual-kNN edge indices.
``scores``: Learned edge scores.
``logit``: Prior-corrected BRP enrichment log-odds.
``prob``: BRP enrichment probabilities.
``pval``: Empirical permutation p-values.
``qval``: Benjamini–Hochberg-adjusted q-values.
``significant``: Boolean local-OOR calls.
``meta``: Parameters and edge-feature metadata.

[ ]:

# A typical workflow applies local OOR analysis within one provisionally
# annotated cell type after excluding cells already called global OOR.

joint_embedding = np.vstack([
    reference_embedding_subset,
    query_embedding_subset,
])

group_indicator = np.concatenate([
    np.zeros(reference_embedding_subset.shape[0], dtype=int),
    np.ones(query_embedding_subset.shape[0], dtype=int),
])

joint_batch = np.concatenate([
    reference_batch_subset,
    query_batch_subset,
])

local_result = epk.model.local_oor_detector(
    Z=joint_embedding,
    y=group_indicator,
    batches=joint_batch,
    k=20,
    train_epochs=50,
    R_perm=800,
    alpha_fdr=0.05,
    seed=42,
    device=None,
)

local_qvalues = local_result["qval"]
local_oor_mask = local_result["significant"]

EpiPack API Reference

Contents

EpiPack API Reference#

Common input conventions#

1. Peak embedding model#

epk.model.Peak_Model#

Constructor#

Parameters#

Main methods#

train_model(nepochs=50, weight_decay=5e-4, learning_rate=1e-4)#

get_z()#

2. PEIVI reference construction#

epk.model.PEIVI#

Constructor#

Required parameters#

Architecture parameters#

Loss and prior parameters#

Main methods#

train_model(nepochs=200, clamp=0.01, weight_decay=5e-4, learning_rate=1e-4, pre_train_epoch=100)#

get_latent()#

Saving and restoring a PEIVI model#

3. PEIVI query mapping#

epk.model.PEIVI_MAPPING#

Constructor#

Required parameters#

Mapping-specific parameters#

Main methods#

train_model(nepochs=50, clamp=0.01, weight_decay=5e-4, learning_rate=1e-4)#

get_latent()#

epk.model.query_model_initial#

Parameters#

Returns#

4. Cell-type classifier#

epk.model.Classifier#

Constructor#

Parameters#

Main methods#

train_model(nepochs=50, weight_decay=5e-4, learning_rate=1e-4, inner_reg=1)#

get_z(input)#

Returns#

5. Global OOR detector#

epk.model.Global_oor_detector#

Constructor#

Parameters#

Main method#

annotate(confidence_threshold=1e-3)#

Parameters#

Returns#

6. Local OOR detector#

epk.model.local_oor_detector#

Signature#

Required parameters#

Graph and edge-feature parameters#

Kernel-training parameters#

BRP and statistical-testing parameters#

Returns#

`epk.model.Peak_Model`#

`train_model(nepochs=50, weight_decay=5e-4, learning_rate=1e-4)`#

`get_z()`#

`epk.model.PEIVI`#

`train_model(nepochs=200, clamp=0.01, weight_decay=5e-4, learning_rate=1e-4, pre_train_epoch=100)`#

`get_latent()`#

`epk.model.PEIVI_MAPPING`#

`train_model(nepochs=50, clamp=0.01, weight_decay=5e-4, learning_rate=1e-4)`#

`get_latent()`#

`epk.model.query_model_initial`#

`epk.model.Classifier`#

`train_model(nepochs=50, weight_decay=5e-4, learning_rate=1e-4, inner_reg=1)`#

`get_z(input)`#

`epk.model.Global_oor_detector`#

`annotate(confidence_threshold=1e-3)`#

`epk.model.local_oor_detector`#