COSIE.COSIE_framework.COSIE_model

class COSIE_model(config, feature_dict)[source]

Bases: Module

The core model class of the COSIE framework, designed for spatial multimodal integration, imputation, and enhancement across multiple tissue sections. This class defines the model structure, and handles both training and inference in full-graph or subgraph settings.

Parameters

configdict: Configuration dictionary defining architecture and training hyperparameters. Includes sub-configs for GraphAutoencoder, Prediction modules, and training settings. This can be customized in configure script.
feature_dictdict: A dictionary mapping section names (e.g., ‘s1’, ‘s2’) to a sub-dictionary of processed feature tensors for each modality (as torch.FloatTensor). Format: {‘s1’: {‘RNA’: tensor[n_cells, d], ‘Protein’: tensor[n_cells, d], …},’s2’: {…}, …}

Attributes

autoencodersnn.ModuleDict: Dictionary of GraphAutoencoder models per modality.
predictorsnn.ModuleDict: Dictionary of Prediction modules for all available modality pairs.
triplet_loss_fntorch.nn.TripletMarginLoss: Loss function used for cross-section integration.
latent_dimint: Embedding dimension.
all_modalitieslist: List of all detected modalities across sections.

Methods

to_device(): Move the entire model, including all submodules, to a specified device.
train_model(): Train the model on either the full graph or spatially partitioned subgraphs, depending on the n_x and n_y grid splits.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Methods

`to_device`	Move all model parameters and submodules to the specified device.
`train_model`	Train the COSIE model on spatial multimodal data.

Attributes