tirank.Model

tirank.Model.setup_seed(seed)[source]

Sets the random seed for reproducibility across all relevant libraries.

Parameters:: seed (int) – The random seed to use.
Returns:: None

tirank.Model.initial_model_para(savePath, nhead=2, nhid1=96, nhid2=8, n_output=32, nlayers=3, n_pred=1, dropout=0.5, mode='Cox', infer_mode='SC', encoder_type='MLP')[source]

Initializes and saves the model hyperparameter configuration.

This function reads the input data dimensions (e.g., number of gene pairs) and spatial cluster dimensions, combines them with the user-defined hyperparameters, and saves the complete configuration as ‘model_para.pkl’.

Parameters:

savePath (str) – The main project directory path.
nhead (int, optional) – Number of heads for Transformer encoder. Defaults to 2.
nhid1 (int, optional) – Hidden dimension for the encoder. Defaults to 96.
nhid2 (int, optional) – Hidden dimension for the predictor heads. Defaults to 8.
n_output (int, optional) – Output dimension of the encoder (embedding size). Defaults to 32.
nlayers (int, optional) – Number of layers in the encoder. Defaults to 3.
n_pred (int, optional) – Output dimension of the predictor (e.g., 1 for Cox/Regression, 2 for binary Classification). Defaults to 1.
dropout (float, optional) – Dropout rate. Defaults to 0.5.
mode (str, optional) – Analysis mode (‘Cox’, ‘Classification’, ‘Regression’). Defaults to “Cox”.
infer_mode (str, optional) – Inference data type (‘SC’ or ‘ST’). Defaults to “SC”.
encoder_type (str, optional) – Type of encoder to use (‘MLP’, ‘Transformer’, ‘DenseNet’). Defaults to “MLP”.

Returns:

None

class tirank.Model.TransformerEncoderModel(*args: Any, **kwargs: Any)[source]

Bases: Module

Transformer-based encoder network.

Parameters:

n_features (int) – Input feature size (number of gene pairs).
nhead (int) – Number of heads in the multi-head attention models.
nhid (int) – Dimension of the feedforward network model.
nlayers (int) – Number of sub-encoder-layers in the encoder.
n_output (int) – Output embedding dimension.
dropout (float, optional) – Dropout value. Defaults to 0.5.

init_weights()[source]: Initializes weights for the linear layers.

forward(x)[source]

Forward pass for the Transformer encoder.

Parameters:: x (torch.Tensor) – Input feature tensor.
Returns:: Output embedding tensor.
Return type:: torch.Tensor

class tirank.Model.PositionalEncoding(*args: Any, **kwargs: Any)[source]

Bases: Module

PositionalEncoding module for Transformer.

Injects sinusoidal positional encodings to the input embeddings.

Parameters:

d_model (int) – The embedding dimension.
dropout (float, optional) – Dropout value. Defaults to 0.1.
max_len (int, optional) – The maximum length of the input sequences. Defaults to 5000.

forward(x)[source]

Forward pass for PositionalEncoding.

Parameters:: x (torch.Tensor) – Input tensor.
Returns:: Output tensor with added positional encoding.
Return type:: torch.Tensor

class tirank.Model.DenseNetEncoderModel(*args: Any, **kwargs: Any)[source]

Bases: Module

DenseNet-style encoder network.

Parameters:

n_features (int) – Input feature size (number of gene pairs).
nlayers (int) – Number of dense layers.
n_output (int) – Output embedding dimension.
dropout (float, optional) – Dropout value. Defaults to 0.5.
growth_rate (float, optional) – Growth rate for the dense layers. Defaults to 0.5.

forward(x)[source]

Forward pass for the DenseNet encoder.

Parameters:: x (torch.Tensor) – Input feature tensor.
Returns:: Output embedding tensor.
Return type:: torch.Tensor

class tirank.Model.MLPEncoderModel(*args: Any, **kwargs: Any)[source]

Bases: Module

MLP-based (Multi-Layer Perceptron) encoder network.

Parameters:

n_features (int) – Input feature size (number of gene pairs).
nhid (int) – Dimension of the hidden layers.
nlayers (int) – Total number of layers (input, hidden, output).
n_output (int) – Output embedding dimension.
dropout (float, optional) – Dropout value. Defaults to 0.5.

forward(x)[source]

Forward pass for the MLP encoder.

Parameters:: x (torch.Tensor) – Input feature tensor.
Returns:: Output embedding tensor.
Return type:: torch.Tensor

class tirank.Model.RiskscorePredictor(*args: Any, **kwargs: Any)[source]

Bases: Module

Prediction head for ‘Cox’ survival analysis.

Predicts a single risk score, applying a sigmoid activation to constrain the output between 0 and 1.

Parameters:

n_features (int) – Input embedding dimension (from encoder).
nhid (int) – Hidden dimension of the predictor MLP.
nhout (int, optional) – Output dimension. Defaults to 1.
dropout (float, optional) – Dropout value. Defaults to 0.5.

forward(embedding)[source]

Forward pass for the risk score predictor.

Parameters:: embedding (torch.Tensor) – Input embedding tensor from the encoder.
Returns:: Predicted risk score (scalar tensor).
Return type:: torch.Tensor

class tirank.Model.RegscorePredictor(*args: Any, **kwargs: Any)[source]

Bases: Module

Prediction head for ‘Regression’ analysis.

Predicts a single continuous value. No output activation is applied.

Parameters:

n_features (int) – Input embedding dimension (from encoder).
nhid (int) – Hidden dimension of the predictor MLP.
nhout (int, optional) – Output dimension. Defaults to 1.
dropout (float, optional) – Dropout value. Defaults to 0.5.

forward(embedding)[source]

Forward pass for the regression score predictor.

Parameters:: embedding (torch.Tensor) – Input embedding tensor from the encoder.
Returns:: Predicted continuous value (scalar tensor).
Return type:: torch.Tensor

class tirank.Model.ClassscorePredictor(*args: Any, **kwargs: Any)[source]

Bases: Module

Prediction head for ‘Classification’ analysis.

Predicts class probabilities using a Softmax activation.

Parameters:

n_features (int) – Input embedding dimension (from encoder).
nhid (int) – Hidden dimension of the predictor MLP.
nhout (int, optional) – Output dimension (number of classes). Defaults to 2.
dropout (float, optional) – Dropout value. Defaults to 0.5.

forward(embedding)[source]

Forward pass for the classification score predictor.

Parameters:: embedding (torch.Tensor) – Input embedding tensor from the encoder.
Returns:: Predicted class probabilities.
Return type:: torch.Tensor

class tirank.Model.PathologyPredictor(*args: Any, **kwargs: Any)[source]

Bases: Module

Auxiliary prediction head for spatial pathology class.

Used for the WSI-guided spatial location-aware module in ST mode.

Parameters:

n_features (int) – Input embedding dimension (from encoder).
nhid (int) – Hidden dimension of the predictor MLP.
nclass (int) – Number of pathology classes to predict.
dropout (float, optional) – Dropout value. Defaults to 0.5.

forward(embedding)[source]

Forward pass for the pathology predictor.

Parameters:: embedding (torch.Tensor) – Input embedding tensor from the encoder.
Returns:: Predicted pathology class probabilities.
Return type:: torch.Tensor

class tirank.Model.TiRankModel(*args: Any, **kwargs: Any)[source]

Bases: Module

The main TiRank multi-task learning model.

This model combines one of the available encoders (MLP, Transformer, DenseNet) with a primary prediction head (for Cox, Classification, or Regression) and an optional auxiliary head for pathology prediction (used in ST mode). It also includes a learnable feature weight layer for L1 regularization.

Parameters:

n_features (int) – Input feature size (number of gene pairs).
nhead (int) – Number of heads for Transformer.
nhid1 (int) – Hidden dimension for the encoder.
nhid2 (int) – Hidden dimension for the predictor heads.
nlayers (int) – Number of layers in the encoder.
n_output (int) – Output dimension of the encoder (embedding size).
n_pred (int, optional) – Output dimension of the primary predictor. Defaults to 1.
n_patho (int, optional) – Output dimension of the pathology predictor (number of classes). Defaults to 0.
dropout (float, optional) – Dropout value. Defaults to 0.5.
mode (str, optional) – Analysis mode (‘Cox’, ‘Classification’, ‘Regression’). Defaults to “Cox”.
encoder_type (str, optional) – Type of encoder to use. Defaults to “MLP”.

forward(x)[source]

The main forward pass for the TiRank model.

Parameters:

x (torch.Tensor) – Input gene pair feature tensor.

Returns:

A tuple containing:

torch.Tensor: The learned embedding.
torch.Tensor: The primary prediction (risk score, class, etc.).
torch.Tensor: The auxiliary pathology prediction.

Return type:

tuple

init_weights(m)[source]

Applies Xavier uniform initialization to linear layers.

Parameters:: m (nn.Module) – A module (or layer) from the network.

Functions

`tirank.Model.initial_model_para`	Initializes and saves the model hyperparameter configuration.
`tirank.Model.setup_seed`	Sets the random seed for reproducibility across all relevant libraries.

Classes

`tirank.Model.ClassscorePredictor`	Prediction head for 'Classification' analysis.
`tirank.Model.DenseNetEncoderModel`	DenseNet-style encoder network.
`tirank.Model.MLPEncoderModel`	MLP-based (Multi-Layer Perceptron) encoder network.
`tirank.Model.PathologyPredictor`	Auxiliary prediction head for spatial pathology class.
`tirank.Model.PositionalEncoding`	PositionalEncoding module for Transformer.
`tirank.Model.RegscorePredictor`	Prediction head for 'Regression' analysis.
`tirank.Model.RiskscorePredictor`	Prediction head for 'Cox' survival analysis.
`tirank.Model.TiRankModel`	The main TiRank multi-task learning model.
`tirank.Model.TransformerEncoderModel`	Transformer-based encoder network.