gsMap3D 0.1.0a1__py3-none-any.whl

gsMap/find_latent/gnn/gcn.py

@@ -0,0 +1,85 @@
+import numpy as np
+import torch
+from scipy.spatial import cKDTree
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.utils import add_remaining_self_loops, degree
+
+
+def build_spatial_graph(
+        coords: np.ndarray,
+        n_neighbors: int,
+        undirected: bool = True
+) -> np.ndarray:
+    """
+
+    Parameters:
+    -----------
+    coords : np.ndarray
+        Spatial coordinates of shape (n_cells, n_dims)
+    n_neighbors : int
+        Number of nearest neighbors
+    undirected : bool, default=True
+        Whether to make graph undirected
+
+    Returns:
+    --------
+    edge_array : np.ndarray
+        Edge array of shape (n_edges, 2)
+    """
+
+    coords = np.ascontiguousarray(coords, dtype=np.float32)
+
+    # Query k-NN
+    tree = cKDTree(coords, balanced_tree=True, compact_nodes=True)
+    _, indices = tree.query(coords, k=n_neighbors, workers=-1)
+
+    n_nodes = coords.shape[0]
+
+    if undirected:
+        # Create bidirectional edges
+        source = np.repeat(np.arange(n_nodes), n_neighbors)
+        target = indices.flatten()
+
+        # Combine forward and reverse edges
+        all_edges = np.column_stack([
+            np.concatenate([source, target]),
+            np.concatenate([target, source])
+        ])
+
+        # Remove duplicates using set
+        edge_set = {tuple(sorted([i, j])) for i, j in all_edges}
+        return np.array(list(edge_set), dtype=np.int32)
+    else:
+        # Directed graph - just flatten the indices
+        source = np.repeat(np.arange(n_nodes), n_neighbors)
+        target = indices.flatten()
+        return np.column_stack([source, target]).astype(np.int32)
+
+class GCN(MessagePassing):
+    """
+    GCN for unweighted graphs.
+    """
+
+    def __init__(self, K=1):
+        super().__init__(aggr="add")
+        self.K = K
+
+    def forward(self, x, edge_index):
+        # Add self-loops
+        edge_index, _ = add_remaining_self_loops(edge_index, num_nodes=x.size(0))
+
+        # Compute normalization: 1/sqrt(deg_i * deg_j)
+        row, col = edge_index
+        deg = degree(row, x.size(0), dtype=x.dtype)
+        norm = (deg[row] * deg[col]).pow(-0.5)
+        norm[norm == float("inf")] = 0
+
+        # K-hop propagation
+        xs = [x]
+        for _ in range(self.K):
+            xs.append(self.propagate(edge_index, x=xs[-1], norm=norm))
+
+        return torch.cat(xs[1:], dim=1)
+
+    def message(self, x_j, norm):
+        return norm.view(-1, 1) * x_j

gsMap/find_latent/gnn/gene_former.py

@@ -0,0 +1,164 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from einops import repeat
+from torch import nn
+
+
+class Linear2D(nn.Module):
+    """Linear2D module consists of a linear layer with 3D weight matrix.
+
+    Args:
+        input_dim (int): The input dimension of the Linear2D module.
+        hidden_dim (int): The hidden dimension of the Linear2D module.
+        n_modules (int): The number of modules of the Linear2D module.
+        bias (bool, optional): Whether to use bias. Defaults to False.
+    """
+
+    def __init__(self,
+                 input_dim,
+                 hidden_dim,
+                 n_modules,
+                 bias=False):
+
+        super().__init__()
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.n_modules = n_modules
+
+        self.weights = torch.randn(input_dim, hidden_dim, n_modules)
+        self.weights = nn.Parameter(
+            nn.init.xavier_normal_(self.weights))
+        self.bias = None
+        if bias:
+            self.bias = torch.randn(1, hidden_dim, n_modules)
+            self.bias = nn.Parameter(
+                nn.init.xavier_normal_(self.bias))
+
+    def forward(self, x):
+        affine_out = torch.einsum("bi,ijk->bjk", [x, self.weights])
+        if self.bias is not None:
+            affine_out = affine_out + self.bias
+        return affine_out
+
+
+class GeneModuler(nn.Module):
+    """GeneModuler takes gene expression as input and outputs gene modules.
+
+    Attributes:
+        input_dim (int): The input dimension of the GeneModuler model.
+        hidden_dim (int): The hidden dimension of the GeneModuler model.
+        n_modules (int): The number of modules of the GeneModuler model.
+        layernorm (nn.LayerNorm): The layer normalization layer.
+        extractor (Linear2D): The Linear2D object.
+    """
+
+    def __init__(self,
+                 input_dim=2000,
+                 hidden_dim=8,
+                 n_modules=16):
+
+        super().__init__()
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.n_modules = n_modules
+
+        self.layernorm = nn.LayerNorm(input_dim)
+        self.extractor = Linear2D(
+            input_dim=input_dim, hidden_dim=hidden_dim, n_modules=n_modules
+        )
+
+    def forward(self, x, batch=None):
+        if batch is not None:
+            module = self.layernorm(x, batch)
+        else:
+            module = self.layernorm(x)
+        module = self.extractor(x).transpose(2, 1)
+        return F.relu(module)
+
+
+class PositionalEncoding(nn.Module):
+    """
+    Positional Encoding
+    Attributes:
+        d_model (int): The dimensionality of the model. This should match the dimension of the input embeddings.
+        max_len (int): The maximum length of the sequence for which positional encoding is computed.
+    """
+    def __init__(self,
+                 d_model,
+                 max_len=500):
+
+        super().__init__()
+
+        self.d_model = d_model
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len).float().unsqueeze(1)
+        angular_speed = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * angular_speed)
+        pe[:, 1::2] = torch.cos(position * angular_speed)
+        self.register_buffer('pe', pe.unsqueeze(0))
+
+    def forward(self, x):
+        # x is N, L, D
+        # pe is 1, maxlen, D
+        scaled_x = x * np.sqrt(self.d_model)
+        encoded = scaled_x + self.pe[:, x.size(1), :]
+        return encoded
+
+
+class GeneModuleFormer(nn.Module):
+    """GeneModuleFormer is a gene expression model based on the Transformer architecture.
+
+    Attributes:
+        input_dim (int): The dimensionality of the input gene expression data.
+        module_dim (int): The dimensionality of each module in the model.
+        hidden_dim (int): The hidden layer dimension used within the model.
+        n_modules (int): The number of modules (transformer blocks) in the model.
+        nhead (int): The number of attention heads in each transformer layer.
+        n_enc_layer (int): The number of encoding layers in the transformer model.
+    """
+
+    def __init__(
+        self,
+        input_dim=2000,
+        module_dim=30,
+        hidden_dim=256,
+        n_modules=16,
+        nhead=8,
+        n_enc_layer=3,
+    ):
+
+        super().__init__()
+
+        self.moduler = GeneModuler(
+            input_dim=input_dim, hidden_dim=module_dim, n_modules=n_modules
+        )
+
+        self.expand = (
+            nn.Linear(module_dim, hidden_dim)
+            if module_dim != hidden_dim
+            else nn.Identity()
+        )
+
+        self.module = nn.TransformerEncoder(
+            encoder_layer=nn.TransformerEncoderLayer(d_model=hidden_dim,
+                                                     nhead=nhead,
+                                                     dim_feedforward=4 * hidden_dim,
+                                                     batch_first=True),
+            num_layers=n_enc_layer
+        )
+
+        self.pe = PositionalEncoding(d_model=module_dim)
+
+        self.cls_token = nn.Parameter(torch.randn(1, 1, module_dim))
+
+    def forward(self, x,):
+        auto_fold = self.moduler(x)
+        b, _, _ = auto_fold.shape
+        auto_fold = self.pe(auto_fold)
+        cls_tokens = repeat(self.cls_token, "() n d -> b n d", b=b)
+        auto_fold = torch.cat([cls_tokens, auto_fold], dim=1)
+        auto_fold = self.expand(auto_fold)
+        rep = self.module(auto_fold)
+        cls_rep = rep[:,0,:]
+        return cls_rep

gsMap/find_latent/gnn/loss.py

@@ -0,0 +1,18 @@
+## Define the loss function
+import torch.nn.functional as F
+
+from .distribution import NegativeBinomial, ZeroInflatedNegativeBinomial
+
+
+def rec_loss(x_hat,x,logtheta,zi_logit,distribution):
+    if distribution == 'nb':
+        loss  = -NegativeBinomial(mu=x_hat, theta=logtheta.exp()).log_prob(x).sum(-1).mean()
+    elif distribution == 'zinb':
+        loss  = -ZeroInflatedNegativeBinomial(mu=x_hat, theta=logtheta.exp(),zi_logits=zi_logit).log_prob(x).sum(-1).mean()
+    else:
+        loss = F.mse_loss(x_hat, x)
+
+    return loss
+
+def ce_loss(pred_label, true_label):
+    return F.cross_entropy(pred_label, true_label)

gsMap/find_latent/gnn/st_model.py

@@ -0,0 +1,125 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from .encoder_decoder import Decoder, Encoder
+
+
+class StEmbeding(nn.Module):
+    """
+    Learn graph-smoothed and expression embeddings for each cell, with optional batch correction.
+
+    Args:
+        input_size (list): List of input feature sizes for each encoder.
+        hidden_size (int): Hidden layer size in encoder/decoder.
+        embedding_size (int): Latent embedding size.
+        batch_embedding_size (int): Size of batch embedding vector.
+        out_put_size (int): Output gene size.
+        batch_size (int): Number of batches (not sample count).
+        class_size (int): Number of classes for classification.
+        distribution (str): Output distribution type ('nb', 'zinb', 'gaussian', etc.).
+        Other GNN-related args passed to Encoder.
+    """
+    def __init__(self,
+                 input_size,
+                 hidden_size,
+                 embedding_size,
+                 batch_embedding_size,
+                 out_put_size,
+                 batch_size,
+                 class_size,
+                 distribution,
+                 module_dim,
+                 hidden_gmf,
+                 n_modules,
+                 nhead,
+                 n_enc_layer,
+                 use_tf=True,
+                 variational=True,
+                 batch_representation='embedding',
+                 dispersion='gene'):
+        super().__init__()
+
+        self.input_size = input_size
+        self.z_num = len(self.input_size)
+        self.distribution = distribution
+        self.batch_representation = batch_representation
+        self.num_batches = batch_size
+
+        self.logtheta = nn.Parameter(torch.randn(batch_size, out_put_size))
+
+        # Handle batch embedding
+        if batch_representation == 'embedding':
+            self.batch_embedding = nn.Embedding(batch_size, batch_embedding_size)
+            self.batch_embedding_size = batch_embedding_size
+        else:
+            self.batch_embedding_size = batch_size  # one-hot case
+
+        # Build encoders for each modality
+        self.encoder = nn.ModuleList()
+        for eid in range(self.z_num):
+            self.encoder.append(
+                Encoder(self.input_size[eid],
+                        hidden_size,
+                        embedding_size,
+                        self.batch_embedding_size,
+                        module_dim,
+                        hidden_gmf,
+                        n_modules,
+                        nhead,
+                        n_enc_layer,
+                        use_tf,
+                        variational)
+            )
+
+        # Build decoders for reconstruction and classification
+        self.decoder = nn.ModuleDict()
+        for decoder_type in ['reconstruction', 'classification']:
+            self.decoder[decoder_type] = Decoder(out_put_size,
+                                                 hidden_size,
+                                                 embedding_size,
+                                                 self.batch_embedding_size,
+                                                 class_size,
+                                                 decoder_type,
+                                                 self.distribution)
+
+    def _process_batch(self, batch):
+        if self.batch_representation == 'embedding':
+            return self.batch_embedding(batch)
+        else:
+            return F.one_hot(batch, num_classes=self.num_batches).float()
+
+    def forward(self, x_list, batch):
+        batch = self._process_batch(batch)
+
+        if self.distribution in ['nb', 'zinb']:
+            library_size = x_list[0].sum(-1, keepdim=True)
+        else:
+            n = x_list[0].shape[0]
+            device = x_list[0].device
+            library_size = torch.ones(n, 1, device=device)
+
+        x_rec_list, zi_logit_list, z_list = [], [], []
+        for eid in range(self.z_num):
+            z = self.encoder[eid](x_list[eid], batch)
+            x_rec, zi_logit = self.decoder['reconstruction'](z, batch)
+            x_rec = x_rec * library_size
+
+            x_rec_list.append(x_rec)
+            zi_logit_list.append(zi_logit)
+            z_list.append(z)
+
+        x_class = self._classification(z_list, batch)
+        return x_rec_list, x_class, zi_logit_list, z_list
+
+    def _classification(self, z_list, batch):
+        z = torch.cat(z_list, dim=1)
+        return self.decoder['classification'](z, batch)
+
+    def encode(self, x_list, batch):
+        batch = self._process_batch(batch)
+        z_list = []
+        for eid in range(self.z_num):
+            z = self.encoder[eid](x_list[eid], batch)
+            z_list.append(z)
+        return z_list

gsMap/find_latent/gnn/train_step.py

@@ -0,0 +1,177 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from .loss import ce_loss, rec_loss
+
+
+class EarlyStopping:
+    """Early stops the training if validation loss doesn't improve after a given patience."""
+
+    def __init__(self, patience=7,delta=0, path=None):
+
+        self.patience = patience
+        self.counter = 0
+        self.best_score = -np.inf
+        self.early_stop = False
+        self.delta = delta
+        self.path = path
+
+    def __call__(self, val_loss):
+        score = -val_loss
+
+        if self.best_score == -np.inf:
+            self.best_score = score
+
+        elif score < self.best_score + self.delta:
+            self.counter += 1
+            if self.counter >= self.patience:
+                self.early_stop = True
+        else:
+            self.best_score = score
+            self.counter = 0
+            self.early_stop = False
+
+
+class ModelTrain:
+    def __init__(self,
+                 model,
+                 optimizer,
+                 distribution,
+                 mode,
+                 lr,
+                 model_path):
+
+        self.model = model
+        self.optimizer = optimizer
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.model.to(self.device)
+        self.distribution = distribution
+        self.mode = mode
+        self.lr = lr
+        self.model_path = model_path
+
+        self.train_loader = None
+        self.val_loader = None
+        self.losses = []
+        self.val_losses = []
+
+        self.train_step_fn = self._make_train_step_fn()
+        self.val_step_fn = self._make_val_step_fn()
+
+        if self.lr is not None:
+            for param_group in self.optimizer.param_groups:
+                param_group['lr'] = self.lr
+
+    def set_loaders(self, train_loader, val_loader=None):
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+
+    def compute_loss(self,x_hat,x,log_theata,zi_logit,x_class,labels):
+        if self.mode == 'reconstruction':
+            loss = 0
+            for id in range(len(x_hat)):
+                loss_rec = rec_loss(x_hat[id],x,log_theata,zi_logit[id],self.distribution)
+                loss_kld = self.model.encoder[id].kl_loss()
+                loss = loss + loss_rec + loss_kld
+
+        elif self.mode == 'classification':
+            loss = ce_loss(x_class, labels)
+
+        return loss
+
+    def _make_train_step_fn(self):
+        # Builds function that performs a step in the train loop
+        def perform_train_step_fn(x_gcn,ST_batches, x, labels):
+
+            self.model.train()
+
+            x_hat, x_class, zi_logit, _ = self.model([x,x_gcn], ST_batches)
+            log_theata = self.model.logtheta[ST_batches]
+            loss = self.compute_loss(x_hat,x, log_theata, zi_logit,x_class,labels)
+
+            loss.backward()
+            self.optimizer.step()
+            self.optimizer.zero_grad()
+
+            return loss.item()
+
+        return perform_train_step_fn
+
+    def _make_val_step_fn(self):
+        # Builds function that performs a step in the validation loop
+        def perform_val_step_fn(x_gcn,ST_batches, x, labels):
+
+            self.model.eval()
+
+            x_hat, x_class, zi_logit, _ = self.model([x,x_gcn], ST_batches)
+            log_theata = self.model.logtheta[ST_batches]
+            loss = self.compute_loss(x_hat,x, log_theata, zi_logit,x_class,labels)
+
+            return loss.item()
+
+        return perform_val_step_fn
+
+    def _mini_batch(self, epoch_idx, n_epochs, validation=False):
+        # The mini-batch can be used with both loaders
+        if validation:
+            data_loader = self.val_loader
+            step_fn = self.val_step_fn
+        else:
+            data_loader = self.train_loader
+            step_fn = self.train_step_fn
+
+        # mini-batch loop
+        mini_batch_losses = []
+        len(self.train_loader)
+
+        for batch_idx, (x_gcn,ST_batches,x,labels) in enumerate(data_loader):
+            # p = float(batch_idx + epoch_idx * batch_iter) / (n_epochs * batch_iter)
+            # grl_lambda = 2. / (1. + np.exp(-10 *p)) -1
+
+            x_gcn = x_gcn.to(self.device)
+            ST_batches = ST_batches.long().to(self.device)
+            x = x.to(self.device)
+            labels = labels.to(self.device)
+
+            mini_batch_loss = step_fn(x_gcn,ST_batches,x,labels)
+            mini_batch_losses.append(mini_batch_loss)
+            # mini_batch_lossses_rgl.append(mini_batch_loss_rgl)
+
+        return np.mean(mini_batch_losses)
+
+
+    def _set_requires_grad(self, module_group, mode):
+        for name, param_group in module_group.items():
+            requires_grad = (mode == name)
+            for param in param_group.parameters():
+                param.requires_grad = requires_grad
+
+
+    def train(self, n_epochs,patience):
+        loss_track = EarlyStopping(patience)
+
+        self._set_requires_grad(self.model.decoder, self.mode)
+
+        pbar = tqdm(range(n_epochs), desc=f'LGCN train ({self.mode})', total=n_epochs)
+        for epoch in pbar:
+
+            # Performs training
+            train_loss = self._mini_batch(epoch,n_epochs,validation=False)
+
+            # Performs evaluation
+            with torch.no_grad():
+                val_loss = self._mini_batch(epoch,n_epochs,validation=True)
+
+            # Save the best model
+            if loss_track.best_score < -val_loss:
+                torch.save(self.model.state_dict(),self.model_path)
+
+            # Update validation loss
+            loss_track(val_loss)
+            if loss_track.early_stop:
+                print(f'Stop training, as {self.mode} validation loss has not decreased for {patience} consecutive steps.')
+                break
+
+            pbar.set_postfix({'train loss': f'{train_loss.item():.4f}',
+                            'validation loss': f'{val_loss.item():.4f}'})