splisosm.model

splisosm.model#

Multinomial GLM/GLMM model implementations for SPLISOSM.

Classes#

MultinomGLM

The Multinomial Generalized Linear Model for spatial isoform expression.

MultinomGLMM

The Multinomial Generalized Linear Mixed Model for spatial isoform expression.

Module Contents#

class splisosm.model.MultinomGLM(fitting_method='iwls', fitting_configs=None)#

Bases: BaseModel, torch.nn.Module

The Multinomial Generalized Linear Model for spatial isoform expression.

Compared to MultinomGLMM, this model does not have a random effect term:

Y ~ Multinomial(alpha, Y.sum(1))
eta = multinomial-logit(alpha) = X @ beta + bias_eta

Given isoform counts of a gene Y (n_spots, n_isos) and design matrix X (n_spots, n_factors), MultinomGLM.fit will find the MAP estimates of the following learnable parameters:

  • beta: (n_factors, n_isos - 1) covariate coefficients of the fixed effect term.

  • bias_eta: (n_isos - 1) intercepts of the fixed effect term.

Inference is performed by maximizing the log likelihood using fitting_method (default: 'iwls')

Example

>>> from splisosm.model import MultinomGLM
>>> import torch
>>> # Generate synthetic data
>>> counts = torch.randint(0, 10, (5, 100, 3))  # 5 genes, 100 spots, each 3 isoforms
>>> # Fit the GLM model
>>> model = MultinomGLM(fitting_method='iwls')
>>> model.setup_data(counts, design_mtx=None)
>>> model.fit()
>>> print(model)
>>> # Extract the fitted isoform ratios
>>> isoform_ratios = model.get_isoform_ratio()  # shape (5, 100, 3)
>>> # Fitted parameters
>>> print(model.beta.shape)  # shape (5, 0, 2)
>>> print(model.bias_eta.shape)  # shape (5, 2)
Parameters:

  • fitting_method (Literal['iwls', 'newton', 'gd']) – Method for fitting the model. 'iwls': Iteratively reweighted least squares. 'newton': Newton’s method. 'gd': Gradient descent.

  • fitting_configs (dict | None) –

    Dictionary of fitting configurations. Keys include

    • 'lr': float, Learning rate for gradient descent or Newton’s method.

    • 'optim': str, Optimizer type, one of 'adam', 'sgd', or 'lbfgs'.

    • 'tol': float, Tolerance for convergence.

    • 'tol_relative': bool, Whether to use relative tolerance for improvement.

    • 'max_epochs': int, Maximum number of epochs for fitting.

    • 'patience': int, Number of epochs to wait for improvement before stopping.

clone()#

Clone a model with the same set of parameters.

Return type:

MultinomGLM

fit(store_param_history=False, verbose=False, quiet=False, random_seed=None, diagnose=None)#

Fit the model using all data.

Parameters:

  • store_param_history (bool) – Whether to store per-epoch parameter snapshots during training. Loss history is always recorded via model.logger.loss_history regardless of this flag.

  • verbose (bool) – Whether to print verbose information during fitting.

  • quiet (bool) – Whether to suppress output during fitting.

  • random_seed (Optional[int]) – Random seed for reproducibility.

  • diagnose (Optional[bool]) – Deprecated alias for store_param_history.

Returns:

params_iter – If store_param_history=True, returns a dictionary of parameter snapshots during training. Otherwise returns None.

Return type:

dict or None

forward()#

Calculate log probability given data.

Returns:

log_prob – Shape (n_genes,), the log probability for each gene.

Return type:

Tensor

get_isoform_ratio()#

Extract the fitted isoform ratio across space.

Returns:

ratio – Shape (n_genes, n_spots, n_isos), the fitted isoform ratio across space.

Return type:

Tensor

setup_data(counts, design_mtx=None, device='cpu')#

Set up the data for the model.

Parameters:

  • counts (Tensor) – Shape (n_genes, n_spots, n_isos) or (n_spots, n_isos). For batched calculations, all genes in the batch must have the same number of isoforms.

  • design_mtx (Optional[Tensor]) – Shape (n_spots, n_factors). Design matrix of spatial covariates. If None, an intercept-only design matrix will be used.

  • device (Literal['cpu', 'cuda', 'mps']) – ‘cpu’, ‘cuda’, or ‘mps’ (torch v2.11.0+).

Return type:

None

update_params_from_dict(params)#

Update a subset of model parameters with a dictionary of parameters.

Parameters:

params (dict) – A dictionary of parameters to be updated. The keys must be existing parameter names in the model.

Return type:

None

fitting_configs: dict#

Dictionary of fitting configurations.

fitting_method: Literal['iwls', 'newton', 'gd']#

Method for fitting the model.

fitting_time: float#

Time taken for fitting the model.

n_factors: int | None#

Number of covariates in the design matrix

n_genes: int#

Number of genes in the batch.

n_isos: int#

Number of isoforms per gene in the batch.

n_spots: int#

Number of samples/spots

class splisosm.model.MultinomGLMM(share_variance=True, var_fix_sigma=True, var_prior_model='none', var_prior_model_params=None, init_ratio='uniform', fitting_method='joint_gd', fitting_configs=None)#

Bases: MultinomGLM, BaseModel, torch.nn.Module

The Multinomial Generalized Linear Mixed Model for spatial isoform expression.

The model is defined as follows:

Y ~ Multinomial(alpha, Y.sum(1))
eta = multinomial-logit(alpha) = X @ beta + bias_eta + nu
nu ~ MVN(0, sigma^2 * (theta * V_sp + (1-theta) * I))

Given isoform counts of a gene Y (n_spots, n_isos), design matrix X (n_spots, n_factors), and spatial covariance matrix V_sp (n_spots, n_spots), the model estimates the isoform usage ratio alpha (n_spots, n_isos) across space. Specifically, MultinomGLMM.fit will find the MAP estimates of the following learnable parameters:

  • beta: (n_factors, n_isos - 1) covariate coefficients of the fixed effect term.

  • bias_eta: (n_isos - 1) intercepts of the fixed effect term.

  • nu: (n_spots, n_isos - 1) the random effect term.

  • variance components: (sigma, theta_logit), each of length n_isos - 1 (or 1 if share_variance is True), representing total variance and logit of spatial variance proportion theta.

Inference algorithms can be categorized into two types based on the optimization objective:

  • Joint: Maximize the joint likelihood (with the random effect nu). This is equivalent to the first-order Laplace approximation of the marginal likelihood.

  • Marginal: Maximize the marginal likelihood (with the random effect nu integrated out). The integral is approximated by a second-order Laplace approximation.

Methods implemented:

  • 'joint_gd': Maximize the joint likelihood using gradient descent.

  • 'joint_newton': Maximize the joint likelihood using Newton’s method.

  • 'marginal_gd': Maximize the marginal likelihood using gradient descent.

  • 'marginal_newton': Maximize the marginal likelihood using Newton’s method. In this method, nu is first updated using Newton’s method every 'update_nu_every_k' iterations, and beta, bias_eta, and variance components are updated using gradient descent.

Notes

It is also possible to implement held-out likelihood for model selection.

Example

>>> from splisosm.model import MultinomGLMM
>>> from splisosm.utils import get_cov_sp
>>> import torch
>>> # Generate synthetic data
>>> counts = torch.randint(0, 10, (5, 100, 3))  # 5 genes, 100 spots, each 3 isoforms
>>> coords = torch.rand(100, 2)  # 100 spots with 2D coordinates
>>> K_sp = get_cov_sp(coords, k=4, rho=0.9) # spatial covariance matrix of shape (100, 100)
>>> # Fit the GLMM model
>>> model = MultinomGLMM(fitting_method='joint_gd')
>>> model.setup_data(counts, corr_sp=K_sp, design_mtx=None)
>>> model.fit()
>>> print(model)
>>> # Extract the fitted isoform ratios
>>> isoform_ratios = model.get_isoform_ratio()  # shape (5, 100, 3)
>>> # Fitted parameters
>>> print(model.beta.shape)  # shape (5, 0, 2)
>>> print(model.bias_eta.shape)  # shape (5, 2)
>>> print(model.nu.shape)  # shape (5, 100, 2)
>>> print(model.sigma.shape)  # shape (5, 1)
>>> print(model.theta_logit.shape)  # shape (5, 1)
Parameters:

  • share_variance (bool) – Whether to use the same variance across isoforms. If True, the variance components will be of length 1. If False, the variance components will be of length n_isos - 1.

  • var_fix_sigma (bool) – Whether to fix the total variance (sigma) to the Fano-factor initial estimate. When True (default), only theta_logit is learned, producing conservative but well-calibrated hypothesis tests. Set to False to learn sigma jointly with other parameters; this may yield higher power for the SV test but can inflate false positive rates for both SV and DU tests.

  • var_prior_model (Literal['none', 'gamma', 'inv_gamma']) – The prior model on the total variance sigma. Default is 'none' with no prior. Other options are 'gamma' (Gamma prior) and 'inv_gamma' (Inverse Gamma prior).

  • var_prior_model_params (dict | None) – The parameters for the prior model on the total variance sigma. For 'gamma', the default parameters are {'alpha': 2.0, 'beta': 0.3}. For 'inv_gamma', the default parameters are {'alpha': 3, 'beta': 0.5}.

  • init_ratio (Literal['observed', 'uniform']) – The initialization method for the logit isoform usage ratio. Options are 'observed' (initialize using observed counts) and 'uniform' (equal isoform usage across space).

  • fitting_method (Literal['joint_gd', 'joint_newton', 'marginal_gd', 'marginal_newton']) – The fitting method to use. Options are 'joint_gd' (joint likelihood with gradient descent), 'joint_newton' (joint likelihood with Newton’s method), 'marginal_gd' (marginal likelihood with gradient descent), and 'marginal_newton' (marginal likelihood with Newton’s method).

  • fitting_configs (dict | None) –

    A dictionary of fitting configurations with the following keys:

    • 'lr': float, learning rate

    • 'optim': str, optimization method (one of 'adam', 'sgd', or 'lbfgs')

    • 'tol': float, tolerance for convergence

    • 'max_epochs': int, maximum number of epochs

    • 'patience': int, number of epochs to wait for improvement before stopping

    • 'update_nu_every_k': int, number of iterations to update nu when using fitting_method='marginal_newton'

clone()#

Clone a Multinomial GLMM model with the same set of parameters.

Return type:

MultinomGLMM

fit(store_param_history=False, verbose=False, quiet=False, random_seed=None, diagnose=None)#

Fit the model using all data.

Parameters:

  • store_param_history (bool) – Whether to store per-epoch parameter snapshots during training. Loss history is always recorded via model.logger.loss_history regardless of this flag.

  • verbose (bool) – Whether to print verbose information during fitting.

  • quiet (bool) – Whether to suppress output during fitting.

  • random_seed (Optional[int]) – Random seed for reproducibility.

  • diagnose (Optional[bool]) – Deprecated alias for store_param_history.

Returns:

params_iter – If store_param_history=True, returns a dictionary of parameter snapshots during training. Otherwise returns None.

Return type:

dict or None

forward()#

Calculate the log-likelihood or log-marginal-likelihood of the model.

Returns:

log_prob – Shape (n_genes,), the log probability for each gene.

Return type:

Tensor

setup_data(counts, corr_sp=None, design_mtx=None, device='cpu', corr_sp_eigvals=None, corr_sp_eigvecs=None)#

Set up the data for the model.

Parameters:

  • counts (Tensor) – Shape (n_genes, n_spots, n_isoforms) or (n_spots, n_isoforms). For batched calculations, all genes in the batch must have the same number of isoforms.

  • corr_sp (Optional[Tensor]) – Shape (n_spots, n_spots), spatial covariance matrix. If None, the eigendecomposition of the spatial covariance matrix must be provided.

  • design_mtx (Optional[Tensor]) – Shape (n_spots, n_factors). Design matrix of spatial covariates. If None, an intercept-only design matrix will be used.

  • device (Literal['cpu', 'cuda', 'mps']) – ‘cpu’, ‘cuda’, or ‘mps’ (torch v2.11.0+).

  • corr_sp_eigvals (Optional[Tensor]) – Shape (n_spots,), eigenvalues of spatial covariance. If None, the spatial covariance matrix corr_sp must be provided.

  • corr_sp_eigvecs (Optional[Tensor]) – Shape (n_spots, n_spots), eigenvectors of spatial covariance. If None, the spatial covariance matrix corr_sp must be provided.

Return type:

None

var_sp_prop()#

Output the proporptions of the spatial variance.

Returns:

var_sp_prop – The proportion theta of spatial variance of shape (n_genes, n_var_components).

Return type:

Tensor

var_total()#

Output the total variance.

Returns:

var_total – The total variance sigma of shape (n_genes, n_var_components).

Return type:

Tensor

fitting_configs: dict#

A dictionary of fitting configurations.

fitting_method: Literal['joint_gd', 'joint_newton', 'marginal_gd', 'marginal_newton']#

The fitting method to use.

fitting_time: float#

The time taken to fit the model.

init_ratio: Literal['observed', 'uniform']#

The initialization method for the logit isoform usage ratio gamma.

share_variance: bool#

Whether to use the same variance across isoforms.

var_fix_sigma: bool#

Whether to fix the total variance (sigma) to its initial value.

When True (default), sigma is frozen at the Fano-factor estimate and only theta_logit (spatial variance proportion) is learned. This yields conservative but well-calibrated SV and DU tests: near-zero false positive rates with strong power on true signals. Set to False to learn sigma jointly, which may increase power at the cost of inflated false positive rates.

var_prior_model: Literal['none', 'gamma', 'inv_gamma']#

The prior model on the total variance sigma.

var_prior_model_params: dict#

The parameters for the prior model on the total variance sigma.