pyrovelocity._velocity_model#

class pyrovelocity._velocity_model.AuxCellVelocityModel(num_cells, num_genes, likelihood='Poisson', shared_time=True, t_scale_on=False, plate_size=2, latent_factor='none', latent_factor_size=30, latent_factor_operation='selection', include_prior=False, num_aux_cells=100, only_cell_times=False, decoder_on=False, add_offset=False, correct_library_size=True, guide_type='velocity', cell_specific_kinetics=None, kinetics_num=None, **initial_values)[source]#

Bases: VelocityModel

generate_cell(gene_plate, alpha, beta, gamma, switching, u_inf, s_inf, u_scale, s_scale, u_log_library, s_log_library, u_pcs_mean=None, s_pcs_mean=None, cell_codebook=None, t_scale=None, gene_offset=None, p_velocity=None, decoder_on=False, u_read_depth=None, s_read_depth=None, u0=None, s0=None)[source]#
class pyrovelocity._velocity_model.LogNormalModel(num_cells, num_genes, likelihood='Poisson', plate_size=2, correct_library_size=True)[source]#

Bases: PyroModule

A base class for pyrovelocity models.

This class serves as the base class for constructing a pyrovelocity model. It provides basic methods for handling gene expression data, such as creating plates for cells and genes, encoding cell-specific features, and computing the likelihood of the observed data.

num_cells#

The number of cells.

Type:

int

num_genes#

The number of genes.

Type:

int

likelihood#

The likelihood type for the model, defaults to “Poisson”.

Type:

str

plate_size#

The size of the plate for the model, defaults to 2.

Type:

int

Example

>>> from pyrovelocity._velocity_model import LogNormalModel
>>> num_cells = 10
>>> num_genes = 20
>>> likelihood = "Poisson"
>>> plate_size = 2
>>> model = LogNormalModel(num_cells, num_genes, likelihood, plate_size)
>>> assert model.num_cells == num_cells
>>> assert model.num_genes == num_genes
>>> assert model.likelihood == likelihood
>>> assert model.plate_size == plate_size
__init__(num_cells, num_genes, likelihood='Poisson', plate_size=2, correct_library_size=True)[source]#
alpha()#
beta()#
cell_code()#
cell_codebook()#
cell_time()#
create_plate(u_obs=None, s_obs=None, u_log_library=None, s_log_library=None, ind_x=None)[source]#
create_plates(u_obs=None, s_obs=None, u_log_library=None, s_log_library=None, u_log_library_loc=None, s_log_library_loc=None, u_log_library_scale=None, s_log_library_scale=None, ind_x=None, cell_state=None, time_info=None)[source]#
Return type:

Tuple[plate, plate]

dt_switching()#
gamma()#
gene_offset()#
get_likelihood(ut, st, u_log_library=None, s_log_library=None, u_scale=None, s_scale=None, u_read_depth=None, s_read_depth=None, u_cell_size_coef=None, ut_coef=None, s_cell_size_coef=None, st_coef=None)[source]#

Compute the likelihood of the given count data.

Parameters:

  • ut (torch.Tensor) – Tensor representing unspliced transcripts.

  • st (torch.Tensor) – Tensor representing spliced transcripts.

  • u_log_library (Optional[torch.Tensor], optional) – Log library tensor for unspliced transcripts. Defaults to None.

  • s_log_library (Optional[torch.Tensor], optional) – Log library tensor for spliced transcripts. Defaults to None.

  • u_scale (Optional[torch.Tensor], optional) – Scale tensor for unspliced transcripts. Defaults to None.

  • s_scale (Optional[torch.Tensor], optional) – Scale tensor for spliced transcripts. Defaults to None.

  • u_read_depth (Optional[torch.Tensor], optional) – Read depth tensor for unspliced transcripts. Defaults to None.

  • s_read_depth (Optional[torch.Tensor], optional) – Read depth tensor for spliced transcripts. Defaults to None.

  • u_cell_size_coef (Optional[Any], optional) – Cell size coefficient for unspliced transcripts. Defaults to None.

  • ut_coef (Optional[Any], optional) – Coefficient for unspliced transcripts. Defaults to None.

  • s_cell_size_coef (Optional[Any], optional) – Cell size coefficient for spliced transcripts. Defaults to None.

  • st_coef (Optional[Any], optional) – Coefficient for spliced transcripts. Defaults to None.

Returns:

A tuple of Poisson distributions for unspliced and spliced transcripts, respectively.

Return type:

Tuple[Poisson, Poisson]

Example

>>> import torch
>>> from pyro.nn import PyroModule
>>> num_cells = 10
>>> num_genes = 20
>>> likelihood = "Poisson"
>>> plate_size = 2
>>> model = LogNormalModel(num_cells, num_genes, likelihood, plate_size)
>>> ut = torch.rand(num_cells, num_genes)
>>> st = torch.rand(num_cells, num_genes)
>>> u_read_depth = torch.rand(num_cells, 1)
>>> s_read_depth = torch.rand(num_cells, 1)
>>> u_dist, s_dist = model.get_likelihood(ut, st, u_read_depth=u_read_depth, s_read_depth=s_read_depth)
>>> assert isinstance(u_dist, torch.distributions.Poisson)
>>> assert isinstance(s_dist, torch.distributions.Poisson)
latent_time()#
p_velocity()#
s_inf()#
s_scale()#
t_scale()#
u_inf()#
u_scale()#
class pyrovelocity._velocity_model.VelocityModel(num_cells, num_genes, likelihood='Poisson', shared_time=True, t_scale_on=False, plate_size=2, latent_factor='none', latent_factor_size=30, latent_factor_operation='selection', include_prior=False, num_aux_cells=100, only_cell_times=False, decoder_on=False, add_offset=False, correct_library_size=True, guide_type='velocity', cell_specific_kinetics=None, kinetics_num=None, **initial_values)[source]#

Bases: LogNormalModel

__init__(num_cells, num_genes, likelihood='Poisson', shared_time=True, t_scale_on=False, plate_size=2, latent_factor='none', latent_factor_size=30, latent_factor_operation='selection', include_prior=False, num_aux_cells=100, only_cell_times=False, decoder_on=False, add_offset=False, correct_library_size=True, guide_type='velocity', cell_specific_kinetics=None, kinetics_num=None, **initial_values)[source]#

Derived class for modeling RNA velocity with shared latent time.

Inherits from LogNormalModel.

Parameters:

  • num_cells (int) – Number of cells.

  • num_genes (int) – Number of genes.

  • likelihood (str, optional) – Likelihood model. Defaults to “Poisson”.

  • shared_time (bool, optional) – Whether to share time across cells. Defaults to True.

  • t_scale_on (bool, optional) – Whether to scale the time. Defaults to False.

  • plate_size (int, optional) – The number of dimensions for the plates. Defaults to 2.

  • latent_factor (str, optional) – Latent factor type. Defaults to “none”.

  • latent_factor_size (int, optional) – Latent factor size. Defaults to 30.

  • latent_factor_operation (str, optional) – Latent factor operation. Defaults to “selection”.

  • include_prior (bool, optional) – Whether to include prior. Defaults to False.

  • num_aux_cells (int, optional) – Number of auxiliary cells. Defaults to 100.

  • only_cell_times (bool, optional) – Whether to use only cell times. Defaults to False.

  • decoder_on (bool, optional) – Whether to use a decoder. Defaults to False.

  • add_offset (bool, optional) – Whether to add an offset. Defaults to False.

  • correct_library_size (Union[bool, str], optional) – Whether to correct library size. Defaults to True.

  • guide_type (bool, optional) – Guide type for the model. Defaults to “velocity”.

  • cell_specific_kinetics (Optional[str], optional) – Cell-specific kinetics. Defaults to None.

  • kinetics_num (Optional[int], optional) – Number of kinetics. Defaults to None.

  • **initial_values – Initial values for various parameters in the model.

Example

>>> import torch
>>> from pyro.nn import PyroModule
>>> num_cells = 10
>>> num_genes = 20
>>> likelihood = "Poisson"
>>> plate_size = 2
>>> model = VelocityModel(num_cells, num_genes, likelihood, plate_size)
>>> u_obs = torch.rand(num_cells, num_genes)
>>> s_obs = torch.rand(num_cells, num_genes)
>>> u_log_library = torch.rand(num_cells, 1)
>>> s_log_library = torch.rand(num_cells, 1)
>>> ind_x = torch.randint(0, num_cells, (num_cells,))
>>> # ut, st = model.model(u_obs, s_obs, u_log_library, s_log_library, ind_x)
>>> # assert ut.shape == (num_cells, num_genes)
>>> # assert st.shape == (num_cells, num_genes)
class pyrovelocity._velocity_model.VelocityModelAuto(*args, **kwargs)[source]#

Bases: AuxCellVelocityModel

__init__(*args, **kwargs)[source]#

Derived class for modeling RNA velocity with shared latent time.

Inherits from LogNormalModel.

Parameters:

  • num_cells (int) – Number of cells.

  • num_genes (int) – Number of genes.

  • likelihood (str, optional) – Likelihood model. Defaults to “Poisson”.

  • shared_time (bool, optional) – Whether to share time across cells. Defaults to True.

  • t_scale_on (bool, optional) – Whether to scale the time. Defaults to False.

  • plate_size (int, optional) – The number of dimensions for the plates. Defaults to 2.

  • latent_factor (str, optional) – Latent factor type. Defaults to “none”.

  • latent_factor_size (int, optional) – Latent factor size. Defaults to 30.

  • latent_factor_operation (str, optional) – Latent factor operation. Defaults to “selection”.

  • include_prior (bool, optional) – Whether to include prior. Defaults to False.

  • num_aux_cells (int, optional) – Number of auxiliary cells. Defaults to 100.

  • only_cell_times (bool, optional) – Whether to use only cell times. Defaults to False.

  • decoder_on (bool, optional) – Whether to use a decoder. Defaults to False.

  • add_offset (bool, optional) – Whether to add an offset. Defaults to False.

  • correct_library_size (Union[bool, str], optional) – Whether to correct library size. Defaults to True.

  • guide_type (bool, optional) – Guide type for the model. Defaults to “velocity”.

  • cell_specific_kinetics (Optional[str], optional) – Cell-specific kinetics. Defaults to None.

  • kinetics_num (Optional[int], optional) – Number of kinetics. Defaults to None.

  • **initial_values – Initial values for various parameters in the model.

Example

>>> import torch
>>> from pyro.nn import PyroModule
>>> num_cells = 10
>>> num_genes = 20
>>> likelihood = "Poisson"
>>> plate_size = 2
>>> model = VelocityModel(num_cells, num_genes, likelihood, plate_size)
>>> u_obs = torch.rand(num_cells, num_genes)
>>> s_obs = torch.rand(num_cells, num_genes)
>>> u_log_library = torch.rand(num_cells, 1)
>>> s_log_library = torch.rand(num_cells, 1)
>>> ind_x = torch.randint(0, num_cells, (num_cells,))
>>> # ut, st = model.model(u_obs, s_obs, u_log_library, s_log_library, ind_x)
>>> # assert ut.shape == (num_cells, num_genes)
>>> # assert st.shape == (num_cells, num_genes)
forward(u_obs=None, s_obs=None, u_log_library=None, s_log_library=None, u_log_library_loc=None, s_log_library_loc=None, u_log_library_scale=None, s_log_library_scale=None, ind_x=None, cell_state=None, time_info=None)[source]#

Defines the forward model, which computes the unspliced (u) and spliced (s) RNA expression levels given the observations and model parameters.

Parameters:

  • u_obs (Optional[torch.Tensor], optional) – Observed unspliced RNA expression. Default is None.

  • s_obs (Optional[torch.Tensor], optional) – Observed spliced RNA expression. Default is None.

  • u_log_library (Optional[torch.Tensor], optional) – Log-transformed library size for unspliced RNA. Default is None.

  • s_log_library (Optional[torch.Tensor], optional) – Log-transformed library size for spliced RNA. Default is None.

  • u_log_library_loc (Optional[torch.Tensor], optional) – Mean of log-transformed library size for unspliced RNA. Default is None.

  • s_log_library_loc (Optional[torch.Tensor], optional) – Mean of log-transformed library size for spliced RNA. Default is None.

  • u_log_library_scale (Optional[torch.Tensor], optional) – Scale of log-transformed library size for unspliced RNA. Default is None.

  • s_log_library_scale (Optional[torch.Tensor], optional) – Scale of log-transformed library size for spliced RNA. Default is None.

  • ind_x (Optional[torch.Tensor], optional) – Indices for the cells. Default is None.

  • cell_state (Optional[torch.Tensor], optional) – Cell state information. Default is None.

  • time_info (Optional[torch.Tensor], optional) – Time information for the cells. Default is None.

Returns:

The unspliced (u) and spliced (s) RNA expression levels.

Return type:

Tuple[torch.Tensor, torch.Tensor]

Examples

>>> import torch
>>> from pyrovelocity._velocity_model import VelocityModelAuto
>>> u_obs=torch.tensor(
...     [[33.,  1.,  7.,  1.],
...     [12., 30., 11.,  3.],
...     [ 1.,  1.,  8.,  5.]],
...     device="cpu",
>>> )
>>> s_obs=torch.tensor(
...     [[32.0, 0.0, 6.0, 0.0],
...     [11.0, 29.0, 10.0, 2.0],
...     [0.0, 0.0, 7.0, 4.0]],
...     device="cpu",
>>> )
>>> u_log_library=torch.tensor([[3.7377], [4.0254], [2.7081]], device="cpu")
>>> s_log_library=torch.tensor([[3.6376], [3.9512], [2.3979]], device="cpu")
>>> u_log_library_loc=torch.tensor([[3.4904], [3.4904], [3.4904]], device="cpu")
>>> s_log_library_loc=torch.tensor([[3.3289], [3.3289], [3.3289]], device="cpu")
>>> u_log_library_scale=torch.tensor([[0.6926], [0.6926], [0.6926]], device="cpu")
>>> s_log_library_scale=torch.tensor([[0.8214], [0.8214], [0.8214]], device="cpu")
>>> ind_x=torch.tensor([2, 0, 1], device="cpu")
>>> model = VelocityModelAuto(3,4)
>>> u, s = model.forward(
>>>            u_obs,
>>>            s_obs,
>>>            u_log_library,
>>>            s_log_library,
>>>            u_log_library_loc,
>>>            s_log_library_loc,
>>>            u_log_library_scale,
>>>            s_log_library_scale,
>>>            ind_x,
>>>        )
>>> u, s
(tensor([[33.,  1.,  7.,  1.],
        [12., 30., 11.,  3.],
        [ 1.,  1.,  8.,  5.]]),
tensor([[32.,  0.,  6.,  0.],
        [11., 29., 10.,  2.],
        [ 0.,  0.,  7.,  4.]]))
get_rna(u_scale, s_scale, alpha, beta, gamma, t, u0, s0, t0, switching=None, u_inf=None, s_inf=None)[source]#

Computes the unspliced (u) and spliced (s) RNA expression levels given the model parameters.

Parameters:

  • u_scale (torch.Tensor) – Scaling factor for unspliced expression.

  • s_scale (torch.Tensor) – Scaling factor for spliced expression.

  • alpha (torch.Tensor) – Transcription rate.

  • beta (torch.Tensor) – Splicing rate.

  • gamma (torch.Tensor) – Degradation rate.

  • t (torch.Tensor) – Cell time.

  • u0 (torch.Tensor) – Unspliced RNA initial expression.

  • s0 (torch.Tensor) – Spliced RNA initial expression.

  • t0 (torch.Tensor) – Initial cell time.

  • switching (Optional[torch.Tensor], optional) – Switching time. Default is None.

  • u_inf (Optional[torch.Tensor], optional) – Unspliced RNA expression at switching time. Default is None.

  • s_inf (Optional[torch.Tensor], optional) – Spliced RNA expression at switching time. Default is None.

Returns:

The unspliced (u) and spliced (s) RNA expression levels.

Return type:

Tuple[torch.Tensor, torch.Tensor]

Examples

>>> import torch
>>> from pyrovelocity._velocity_model import VelocityModelAuto
>>> model = VelocityModelAuto(
...             3,
...             4,
...             "Poisson",
...             True,
...             False,
...             2,
...             "none",
...             latent_factor_operation="selection",
...             latent_factor_size=10,
...             include_prior=False,
...             num_aux_cells=0,
...             only_cell_times=True,
...             decoder_on=False,
...             add_offset=False,
...             correct_library_size=True,
...             guide_type="auto_t0_constraint",
...             cell_specific_kinetics=None,
...             **{}
...         )
>>> u, s = model.get_rna(
...            u_scale=torch.tensor([0.9793, 1.0567, 0.8610, 0.9304], device="cpu"),
...            s_scale=torch.tensor(1.0),
...            alpha=torch.tensor([0.4869, 1.5997, 1.3962, 0.5038], device="cpu"),
...            beta=torch.tensor([0.5403, 1.1192, 0.9912, 1.1783], device="cpu"),
...            gamma=torch.tensor([1.9612, 0.5533, 2.1050, 4.9345], device="cpu"),
...            t=torch.tensor([[0.4230], [0.5119], [0.2689]], device="cpu"),
...            u0=torch.tensor(0.0),
...            s0=torch.tensor(0.0),
...            t0=torch.tensor([-0.4867, 0.5581, -0.6957, 0.6028], device="cpu"),
...            switching=torch.tensor([1.1886, 1.1227, 0.6789, 4.1003], device="cpu"),
...            u_inf=torch.tensor([0.5367, 0.6695, 1.0479, 0.4206], device="cpu"),
...            s_inf=torch.tensor([0.1132, 0.2100, 0.3750, 0.0999], device="cpu"),
>>>        )