From 1de221a4e2a3e7c2c207ece0a06f4cadb63e52ca Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Sat, 21 Jun 2025 21:08:33 -0700 Subject: [PATCH 01/21] utils --- src/state_sets/sets/models/utils.py | 47 +++++------------------------ 1 file changed, 7 insertions(+), 40 deletions(-) diff --git a/src/state_sets/sets/models/utils.py b/src/state_sets/sets/models/utils.py index 08d7a6da..6b94bf41 100644 --- a/src/state_sets/sets/models/utils.py +++ b/src/state_sets/sets/models/utils.py @@ -11,7 +11,7 @@ def build_mlp( hidden_dim: int, n_layers: int, dropout: float = 0.0, - activation: nn.Module = nn.ReLU, # default to nn.ReLU class + activation: nn.Module = nn.ReLU, ) -> nn.Sequential: """ Build an MLP of `n_layers` from `in_dim` to `out_dim`. @@ -24,18 +24,15 @@ def build_mlp( if n_layers == 1: layers.append(nn.Linear(in_dim, out_dim)) else: - # First layer layers.append(nn.Linear(in_dim, hidden_dim)) - layers.append(activation()) # instantiate the class + layers.append(activation()) layers.append(nn.Dropout(dropout)) - # Intermediate layers for _ in range(n_layers - 2): layers.append(nn.Linear(hidden_dim, hidden_dim)) - layers.append(activation()) # instantiate again + layers.append(activation()) layers.append(nn.Dropout(dropout)) - # Final layer layers.append(nn.Linear(hidden_dim, out_dim)) return nn.Sequential(*layers) @@ -64,7 +61,6 @@ def get_activation_class(name: str) -> nn.Module: return nn.SELU elif name == "gelu": return nn.GELU - # Add more as needed... else: raise ValueError(f"Unsupported activation function: {name}") @@ -85,7 +81,6 @@ def get_loss_fn(loss: Union[str, nn.Module]) -> nn.Module: if loss == "mse": return nn.MSELoss() - # Add more as needed... else: raise ValueError(f"Unsupported loss function: {loss}") @@ -95,7 +90,6 @@ def get_transformer_backbone(key, kwargs) -> PreTrainedModel: config = GPT2Config(**kwargs) model = GPT2BidirectionalModel(config) - # Zero out position embeddings and freeze them model.wpe.weight.requires_grad = False model.wte.weight.requires_grad = False model.wpe.weight.zero_() @@ -128,11 +122,8 @@ def __init__(self, num_attention_heads: int, hidden_size: int): self.hidden_size = hidden_size def forward(self, hidden_states: torch.Tensor, position_ids: torch.LongTensor): - # hidden_states: (batch_size, seq_len, hidden_dim) batch_size, seq_len, hidden_dim = hidden_states.shape - # Create cos = ones, sin = zeros - # shape --> (batch_size, seq_len, head_dim) cos = hidden_states.new_ones(batch_size, seq_len, self.num_heads) sin = hidden_states.new_zeros(batch_size, seq_len, self.num_heads) return cos, sin @@ -160,9 +151,9 @@ def _update_causal_mask( past_key_values, output_attentions: bool = False, ): - # By returning None, we disable any causal‐(look‐ahead) masking. + # By returning None, we disable any causal (look-ahead) masking. # The only mask that remains is whatever “attention_mask” the user has passed - # (e.g. padding‐mask), which will be handled by Flash/SDPA internally as non‐causal. + # (e.g. padding mask), which will be handled by Flash/SDPA internally as non-causal. return None def forward( @@ -197,17 +188,14 @@ def forward( class GPT2BidirectionalModel(GPT2Model): """ A thin wrapper around GPT2Model that disables the causal (unidirectional) mask, - allowing full bidirectional attention—and prints the internal bias mask each forward pass. + allowing full bidirectional attention. """ def __init__(self, config: GPT2Config): - # Mark as not‐a‐decoder (for downstream utilities). config.is_decoder = False super().__init__(config) - # Overwrite each attention's bias so no triangular masking occurs. for block in self.h: - # block.attn.bias is a bool‐tensor of shape (1, 1, max_pos, max_pos). block.attn.bias.data.fill_(True) block.attn.is_causal = False @@ -241,34 +229,13 @@ def forward( return_dict=None, **kwargs, ): - # Determine sequence length for printing the relevant slice of bias - if input_ids is not None: - seq_len = input_ids.size(1) - elif inputs_embeds is not None: - seq_len = inputs_embeds.size(1) - else: - seq_len = None # If neither is given, we can’t infer seq_len - - if seq_len is not None: - # Print the (1, 1, seq_len, seq_len) slice of the bias for the first block - bias_mask = self.h[0].attn.bias[0, 0, :seq_len, :seq_len] - # print("Bias mask (block 0) slice [0,0,:seq_len,:seq_len]:") - # print(bias_mask) - # else: - # print("Cannot infer sequence length to print bias mask.") - - # If a 2D attention_mask was provided, print its expanded 4D version: + if attention_mask is not None: - # Expand to (batch_size, 1, seq_len, seq_len) B, S = attention_mask.size() expanded = attention_mask.unsqueeze(1).unsqueeze(2).expand(B, 1, S, S) - # Convert to float mask (1→0.0, 0→-inf) just like GPT2 does internally neg_inf = torch.finfo(self.dtype).min float_mask = (1.0 - expanded.to(self.dtype)) * neg_inf - # print(f"Expanded attention_mask (shape {expanded.shape}) → float mask:") - # print(float_mask) - # Finally, call the parent forward method return super().forward( input_ids=input_ids, past_key_values=past_key_values, From 032775f25306bba52c2fdeed6bec3b8ef4a21660 Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Sat, 21 Jun 2025 21:32:26 -0700 Subject: [PATCH 02/21] started removing things --- src/state_sets/sets/models/decoders.py | 127 ------- src/state_sets/sets/models/pert_sets.py | 35 +- src/state_sets/sets/models/pseudobulk.py | 403 ----------------------- 3 files changed, 1 insertion(+), 564 deletions(-) delete mode 100644 src/state_sets/sets/models/decoders.py delete mode 100644 src/state_sets/sets/models/pseudobulk.py diff --git a/src/state_sets/sets/models/decoders.py b/src/state_sets/sets/models/decoders.py deleted file mode 100644 index 90a55f15..00000000 --- a/src/state_sets/sets/models/decoders.py +++ /dev/null @@ -1,127 +0,0 @@ -import logging - -import torch -import torch.nn as nn -from omegaconf import OmegaConf - -from ...state.finetune_decoder import Finetune - -logger = logging.getLogger(__name__) - - -class FinetuneVCICountsDecoder(nn.Module): - def __init__( - self, - genes, - # model_loc="/large_storage/ctc/userspace/aadduri/vci/checkpoint/rda_tabular_counts_2048_new/step=950000.ckpt", - # config="/large_storage/ctc/userspace/aadduri/vci/checkpoint/rda_tabular_counts_2048_new/tahoe_config.yaml", - model_loc="/home/aadduri/vci_pretrain/vci_1.4.2.ckpt", - config="/large_storage/ctc/userspace/aadduri/vci/checkpoint/large_1e-4_rda_tabular_counts_2048/crossds_config.yaml", - read_depth=1200, - latent_dim=1024, # dimension of pretrained vci model - hidden_dims=[512, 512, 512], # hidden dimensions of the decoder - dropout=0.1, - basal_residual=False, - ): - super().__init__() - self.genes = genes - self.model_loc = model_loc - self.config = config - self.finetune = Finetune(OmegaConf.load(self.config)) - self.finetune.load_model(self.model_loc) - self.read_depth = nn.Parameter(torch.tensor(read_depth, dtype=torch.float), requires_grad=False) - self.basal_residual = basal_residual - - # layers = [ - # nn.Linear(latent_dim, hidden_dims[0]), - # ] - - # self.gene_lora = nn.Sequential(*layers) - - self.latent_decoder = nn.Sequential( - nn.Linear(latent_dim, hidden_dims[0]), - nn.LayerNorm(hidden_dims[0]), - nn.GELU(), - nn.Dropout(dropout), - nn.Linear(hidden_dims[0], hidden_dims[1]), - nn.LayerNorm(hidden_dims[1]), - nn.GELU(), - nn.Dropout(dropout), - nn.Linear(hidden_dims[1], len(self.genes)), - nn.ReLU(), - ) - - self.gene_decoder_proj = nn.Sequential( - nn.Linear(len(self.genes), 128), - nn.Linear(128, len(self.genes)), - ) - - self.binary_decoder = self.finetune.model.binary_decoder - for param in self.binary_decoder.parameters(): - param.requires_grad = False - - def gene_dim(self): - return len(self.genes) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - # x is [B, S, latent_dim]. - if len(x.shape) != 3: - x = x.unsqueeze(0) - batch_size, seq_len, latent_dim = x.shape - x = x.view(batch_size * seq_len, latent_dim) - - # Get gene embeddings - gene_embeds = self.finetune.get_gene_embedding(self.genes) - - # Handle RDA task counts - use_rda = getattr(self.finetune.model.cfg.model, "rda", False) - # Define your sub-batch size (tweak this based on your available memory) - sub_batch_size = 16 - logprob_chunks = [] # to store outputs of each sub-batch - - for i in range(0, x.shape[0], sub_batch_size): - # Get the sub-batch of latent vectors - x_sub = x[i : i + sub_batch_size] - - # Create task_counts for the sub-batch if needed - if use_rda: - # task_counts_sub = torch.full( - # (x_sub.shape[0],), self.read_depth, device=x.device - # ) - task_counts_sub = torch.ones((x_sub.shape[0],), device=x.device) * self.read_depth - else: - task_counts_sub = None - - # Compute merged embeddings for the sub-batch - merged_embs_sub = self.finetune.model.resize_batch(x_sub, gene_embeds, task_counts_sub) - - # Run the binary decoder on the sub-batch - logprobs_sub = self.binary_decoder(merged_embs_sub) - - # Squeeze the singleton dimension if needed - if logprobs_sub.dim() == 3 and logprobs_sub.size(-1) == 1: - logprobs_sub = logprobs_sub.squeeze(-1) - - # Collect the results - logprob_chunks.append(logprobs_sub) - - # Concatenate the sub-batches back together - logprobs = torch.cat(logprob_chunks, dim=0) - - # Reshape back to [B, S, gene_dim] - decoded_gene = logprobs.view(batch_size, seq_len, len(self.genes)) - decoded_gene = decoded_gene + self.gene_decoder_proj(decoded_gene) - # decoded_gene = torch.nn.functional.relu(decoded_gene) - - # # normalize the sum of decoded_gene to be read depth - # decoded_gene = decoded_gene / decoded_gene.sum(dim=2, keepdim=True) * self.read_depth - - # decoded_gene = self.gene_lora(decoded_gene) - # TODO: fix this to work with basal counts - - # add logic for basal_residual: - decoded_x = self.latent_decoder(x) - decoded_x = decoded_x.view(batch_size, seq_len, len(self.genes)) - - # Pass through the additional decoder layers - return decoded_gene + decoded_x diff --git a/src/state_sets/sets/models/pert_sets.py b/src/state_sets/sets/models/pert_sets.py index a317b484..aa2ba52b 100644 --- a/src/state_sets/sets/models/pert_sets.py +++ b/src/state_sets/sets/models/pert_sets.py @@ -1,7 +1,6 @@ import logging from typing import Dict, Optional -import anndata as ad import numpy as np import torch import torch.nn as nn @@ -10,7 +9,6 @@ from typing import Tuple from .base import PerturbationModel -from .decoders import FinetuneVCICountsDecoder from .decoders_nb import NBDecoder, nb_nll from .utils import build_mlp, get_activation_class, get_transformer_backbone @@ -154,7 +152,7 @@ def __init__( else: raise ValueError(f"Unknown loss function: {loss_name}") - # Build the underlying neural OT network + # Build the underlying transition network self._build_networks() # Add an optional encoder that introduces a batch variable @@ -197,37 +195,6 @@ def __init__( hidden_dims=[512, 512, 512], dropout=self.dropout, ) - - control_pert = kwargs.get("control_pert", "non-targeting") - if kwargs.get("finetune_vci_decoder", False): # TODO: This will go very soon - gene_names = [] - - if output_space == "gene": - # hvg's but for which dataset? - if "DMSO_TF" in control_pert: - gene_names = np.load( - "/large_storage/ctc/userspace/aadduri/datasets/tahoe_19k_to_2k_names.npy", allow_pickle=True - ) - elif "non-targeting" in control_pert: - temp = ad.read_h5ad("/large_storage/ctc/userspace/aadduri/datasets/hvg/replogle/jurkat.h5") - # gene_names = temp.var.index.values - else: - assert output_space == "all" - if "DMSO_TF" in control_pert: - gene_names = np.load( - "/large_storage/ctc/userspace/aadduri/datasets/tahoe_19k_names.npy", allow_pickle=True - ) - elif "non-targeting" in control_pert: - # temp = ad.read_h5ad('/scratch/ctc/ML/vci/paper_replogle/jurkat.h5') - # gene_names = temp.var.index.values - temp = ad.read_h5ad("/large_storage/ctc/userspace/aadduri/cross_dataset/replogle/jurkat.h5") - gene_names = temp.var.index.values - - self.gene_decoder = FinetuneVCICountsDecoder( - genes=gene_names, - # latent_dim=self.output_dim + (self.batch_dim or 0), - ) - print(self) def _build_networks(self): diff --git a/src/state_sets/sets/models/pseudobulk.py b/src/state_sets/sets/models/pseudobulk.py deleted file mode 100644 index c63eb43e..00000000 --- a/src/state_sets/sets/models/pseudobulk.py +++ /dev/null @@ -1,403 +0,0 @@ -import logging -from typing import Dict, Optional - -import anndata as ad -import numpy as np -import torch -import torch.nn as nn -from geomloss import SamplesLoss - -from .base import PerturbationModel -from .decoders import FinetuneVCICountsDecoder -from .decoders_nb import NBDecoder, nb_nll -from .utils import build_mlp, get_activation_class, get_transformer_backbone - -logger = logging.getLogger(__name__) - - -class PseudobulkPerturbationModel(PerturbationModel): - """ - This model: - 1) Projects basal expression and perturbation encodings into a shared latent space. - 2) Uses an OT-based distributional loss (energy, sinkhorn, etc.) from geomloss. - 3) Enables cells to attend to one another, learning a set-to-set function rather than - a sample-to-sample single-cell map. - """ - - def __init__( - self, - input_dim: int, - hidden_dim: int, - output_dim: int, - pert_dim: int, - batch_dim: int = None, - predict_residual: bool = True, - distributional_loss: str = "energy", - transformer_backbone_key: str = "GPT2", - transformer_backbone_kwargs: dict = None, - output_space: str = "gene", - gene_dim: Optional[int] = None, - **kwargs, - ): - """ - Args: - input_dim: dimension of the input expression (e.g. number of genes or embedding dimension). - hidden_dim: not necessarily used, but required by PerturbationModel signature. - output_dim: dimension of the output space (genes or latent). - pert_dim: dimension of perturbation embedding. - gpt: e.g. "TranslationTransformerSamplesModel". - model_kwargs: dictionary passed to that model's constructor. - loss: choice of distributional metric ("sinkhorn", "energy", etc.). - **kwargs: anything else to pass up to PerturbationModel or not used. - """ - # Call the parent PerturbationModel constructor - super().__init__( - input_dim=input_dim, - hidden_dim=hidden_dim, - gene_dim=gene_dim, - output_dim=output_dim, - pert_dim=pert_dim, - batch_dim=batch_dim, - output_space=output_space, - **kwargs, - ) - - # Save or store relevant hyperparams - self.predict_residual = predict_residual - self.n_encoder_layers = kwargs.get("n_encoder_layers", 2) - self.n_decoder_layers = kwargs.get("n_decoder_layers", 2) - self.activation_class = get_activation_class(kwargs.get("activation", "gelu")) - self.transformer_backbone_key = transformer_backbone_key - self.transformer_backbone_kwargs = transformer_backbone_kwargs - self.distributional_loss = distributional_loss - self.cell_sentence_len = self.transformer_backbone_kwargs["n_positions"] - self.gene_dim = gene_dim - self.batch_dim = batch_dim - - # Build the distributional loss from geomloss - blur = kwargs.get("blur", 0.05) - loss_name = kwargs.get("loss", "energy") - if loss_name == "energy": - self.loss_fn = SamplesLoss(loss=self.distributional_loss, blur=blur) - elif loss_name == "mse": - self.loss_fn = nn.MSELoss() - else: - raise ValueError(f"Unknown loss function: {loss_name}") - - # Build the underlying neural OT network - self._build_networks() - - self.batch_encoder = None - if kwargs.get("batch_encoder", False) and batch_dim is not None: - self.batch_encoder = build_mlp( - in_dim=batch_dim, - out_dim=self.hidden_dim, - hidden_dim=self.hidden_dim, - n_layers=self.n_encoder_layers, - dropout=self.dropout, - activation=self.activation_class, - ) - - # if the model is outputting to counts space, apply softplus - # otherwise its in embedding space and we don't want to - is_gene_space = kwargs["embed_key"] == "X_hvg" or kwargs["embed_key"] is None - if kwargs.get("softplus", False) and is_gene_space: - # actually just set this to a relu for now - self.relu = torch.nn.ReLU() - - if kwargs.get("nb_decoder", False): - self.gene_decoder = NBDecoder( - latent_dim=self.output_dim + (self.batch_dim or 0), - gene_dim=gene_dim, - hidden_dims=[512, 512, 512], - dropout=self.dropout, - ) - - control_pert = kwargs.get("control_pert", "non-targeting") - if kwargs.get("finetune_vci_decoder", False): - gene_names = [] - - if output_space == "gene": - # hvg's but for which dataset? - if "DMSO_TF" in control_pert: - gene_names = np.load( - "/large_storage/ctc/userspace/aadduri/datasets/tahoe_19k_to_2k_names.npy", allow_pickle=True - ) - elif "non-targeting" in control_pert: - temp = ad.read_h5ad("/large_storage/ctc/userspace/aadduri/datasets/hvg/replogle/jurkat.h5") - gene_names = temp.var.index.values - else: - assert output_space == "all" - if "DMSO_TF" in control_pert: - gene_names = np.load( - "/large_storage/ctc/userspace/aadduri/datasets/tahoe_19k_names.npy", allow_pickle=True - ) - elif "non-targeting" in control_pert: - # temp = ad.read_h5ad('/scratch/ctc/ML/vci/paper_replogle/jurkat.h5') - # gene_names = temp.var.index.values - temp = ad.read_h5ad("/large_storage/ctc/userspace/aadduri/cross_dataset/replogle/jurkat.h5") - gene_names = temp.var.index.values - - self.gene_decoder = FinetuneVCICountsDecoder( - genes=gene_names, - # latent_dim=self.output_dim + (self.batch_dim or 0), - ) - - print(self) - - def _build_networks(self): - """ - Here we instantiate the actual GPT2-based model. - """ - self.pert_encoder = build_mlp( - in_dim=self.pert_dim, - out_dim=self.hidden_dim, - hidden_dim=self.hidden_dim, - n_layers=self.n_encoder_layers, - dropout=self.dropout, - activation=self.activation_class, - ) - - # Map the input embedding to the hidden space - self.basal_encoder = build_mlp( - in_dim=self.input_dim, - out_dim=self.hidden_dim, - hidden_dim=self.hidden_dim, - n_layers=self.n_encoder_layers, - dropout=self.dropout, - activation=self.activation_class, - ) - - self.transformer_backbone, self.transformer_model_dim = get_transformer_backbone( - self.transformer_backbone_key, - self.transformer_backbone_kwargs, - ) - - self.project_out = build_mlp( - in_dim=self.hidden_dim, - out_dim=self.output_dim, - hidden_dim=self.hidden_dim, - n_layers=self.n_decoder_layers, - dropout=self.dropout, - activation=self.activation_class, - ) - - def encode_perturbation(self, pert: torch.Tensor) -> torch.Tensor: - """If needed, define how we embed the raw perturbation input.""" - return self.pert_encoder(pert) - - def encode_basal_expression(self, expr: torch.Tensor) -> torch.Tensor: - """Define how we embed basal state input, if needed.""" - return self.basal_encoder(expr) - - def forward(self, batch: dict, padded=True) -> torch.Tensor: - """ - The main forward call. Batch is a flattened sequence of cell sentences, - which we reshape into sequences of length cell_sentence_len. - - Expects input tensors of shape (B, S, N) where: - B = batch size - S = sequence length (cell_sentence_len) - N = feature dimension - - The `padded` argument here is set to True if the batch is padded. Otherwise, we - expect a single batch, so that sentences can vary in length across batches. - """ - if padded: - pert = batch["pert_emb"].reshape(-1, self.cell_sentence_len, self.pert_dim) - basal = batch["ctrl_cell_emb"].reshape(-1, self.cell_sentence_len, self.input_dim) - else: - # we are inferencing on a single batch, so accept variable length sentences - pert = batch["pert_emb"].reshape(1, -1, self.pert_dim) - basal = batch["ctrl_cell_emb"].reshape(1, -1, self.input_dim) - - # Shape: [B, S, hidden_dim] - pert_embedding = self.encode_perturbation(pert) - control_cells = self.encode_basal_expression(basal) - - seq_input = pert_embedding + control_cells # Shape: [B, S, hidden_dim] - batch_size = seq_input.shape[0] - - if self.batch_encoder is not None: - if padded: - batch = batch["batch"].reshape(-1, self.cell_sentence_len, self.batch_dim) - else: - batch = batch["batch"].reshape(1, -1, self.batch_dim) - - seq_input = seq_input + self.batch_encoder(batch) # Shape: [B, S, hidden_dim] - - # take the average across the sequence dimension - seq_input = seq_input.mean(dim=1, keepdim=True) # Shape: [B, 1, hidden_dim] - - # forward pass + extract CLS last hidden state - res_pred = self.transformer_backbone(inputs_embeds=seq_input).last_hidden_state # Shape: [B, 1, hidden_dim] - out_dim = res_pred.shape[-1] - - # broadcast to the sequence length - if padded: - res_pred = res_pred.expand(batch_size, self.cell_sentence_len, out_dim) # Shape: [B, S, hidden_dim] - else: - res_pred = res_pred.expand(1, -1, out_dim) - - # add to basal if predicting residual - if self.predict_residual: - # treat the actual prediction as a residual sum to basal - out_pred = self.project_out(res_pred + control_cells) - else: - out_pred = self.project_out(res_pred) - - # apply softplus if specified and we output to HVG space - is_gene_space = self.hparams["embed_key"] == "X_hvg" or self.hparams["embed_key"] is None - if self.hparams.get("softplus", False) and is_gene_space: - out_pred = self.relu(out_pred) - - return out_pred.reshape(-1, self.output_dim) - - def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int, padded=True) -> torch.Tensor: - """Training step logic for both main model and decoder.""" - # Get model predictions (in latent space) - pred = self.forward(batch, padded=padded) - - target = batch["pert_cell_emb"] - - if padded: - pred = pred.reshape(-1, self.cell_sentence_len, self.output_dim) - target = target.reshape(-1, self.cell_sentence_len, self.output_dim) - else: - pred = pred.reshape(1, -1, self.output_dim) - target = target.reshape(1, -1, self.output_dim) - - main_loss = self.loss_fn(pred, target).nanmean() - self.log("train_loss", main_loss) - - # Process decoder if available - decoder_loss = None - total_loss = main_loss - - if self.gene_decoder is not None and "pert_cell_counts" in batch: - gene_targets = batch["pert_cell_counts"] - # Train decoder to map latent predictions to gene space - latent_preds = pred - # with torch.no_grad(): - # latent_preds = pred.detach() # Detach to prevent gradient flow back to main model - - batch_var = batch["batch"].reshape(latent_preds.shape[0], latent_preds.shape[1], -1) - # concatenate on the last axis - if self.batch_dim is not None and not isinstance(self.gene_decoder, FinetuneVCICountsDecoder): - latent_preds = torch.cat([latent_preds, batch_var], dim=-1) - - if isinstance(self.gene_decoder, NBDecoder): - mu, theta = self.gene_decoder(latent_preds) - gene_targets = batch["pert_cell_counts"].reshape_as(mu) - decoder_loss = nb_nll(gene_targets, mu, theta) - else: - pert_cell_counts_preds = self.gene_decoder(latent_preds) - if padded: - gene_targets = gene_targets.reshape(-1, self.cell_sentence_len, self.gene_decoder.gene_dim()) - else: - gene_targets = gene_targets.reshape(1, -1, self.gene_decoder.gene_dim()) - - decoder_loss = self.loss_fn(pert_cell_counts_preds, gene_targets).mean() - - # Log decoder loss - self.log("decoder_loss", decoder_loss) - - total_loss = total_loss + 0.1 * decoder_loss - - return total_loss - - def validation_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> None: - """Validation step logic.""" - pred = self.forward(batch) - - pred = pred.reshape(-1, self.cell_sentence_len, self.output_dim) - target = batch["pert_cell_emb"] - target = target.reshape(-1, self.cell_sentence_len, self.output_dim) - - loss = self.loss_fn(pred, target).mean() - self.log("val_loss", loss) - - if self.gene_decoder is not None and "pert_cell_counts" in batch: - gene_targets = batch["pert_cell_counts"] - - # Get model predictions from validation step - latent_preds = pred - - if isinstance(self.gene_decoder, NBDecoder): - mu, theta = self.gene_decoder(latent_preds) - gene_targets = batch["pert_cell_counts"].reshape_as(mu) - decoder_loss = nb_nll(gene_targets, mu, theta) - else: - pert_cell_counts_preds = self.gene_decoder(latent_preds) # verify this is automatically detached - - # Get decoder predictions - pert_cell_counts_preds = pert_cell_counts_preds.reshape(-1, self.cell_sentence_len, self.gene_dim) - gene_targets = gene_targets.reshape(-1, self.cell_sentence_len, self.gene_dim) - decoder_loss = self.loss_fn(pert_cell_counts_preds, gene_targets).mean() - - # Log the validation metric - self.log("decoder_val_loss", decoder_loss) - - return {"loss": loss, "predictions": pred} - - def test_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> None: - pred = self.forward(batch, padded=False) - target = batch["pert_cell_emb"] - pred = pred.reshape(1, -1, self.output_dim) - target = target.reshape(1, -1, self.output_dim) - loss = self.loss_fn(pred, target).mean() - self.log("test_loss", loss) - - def predict_step(self, batch, batch_idx, padded=True, **kwargs): - """ - Typically used for final inference. We'll replicate old logic: - returning 'preds', 'X', 'pert_name', etc. - """ - latent_output = self.forward(batch, padded=padded) # shape [B, ...] - - output_dict = { - "preds": latent_output, - "pert_cell_emb": batch.get("pert_cell_emb", None), - "pert_cell_counts": batch.get("pert_cell_counts", None), - "pert_name": batch.get("pert_name", None), - "celltype_name": batch.get("cell_type", None), - "batch": batch.get("batch", None), - "ctrl_cell_emb": batch.get("ctrl_cell_emb", None), - } - - basal_hvg = batch.get("ctrl_cell_counts", None) - - if self.gene_decoder is not None: - if latent_output.dim() == 2: - batch_var = batch["batch"].reshape(latent_output.shape[0], -1) - else: - batch_var = batch["batch"].reshape(latent_output.shape[0], latent_output.shape[1], -1) - # concatenate on the last axis - if self.batch_dim is not None and not isinstance(self.gene_decoder, FinetuneVCICountsDecoder): - latent_output = torch.cat([latent_output, batch_var], dim=-1) - if isinstance(self.gene_decoder, NBDecoder): - mu, _ = self.gene_decoder(latent_output) - pert_cell_counts_preds = mu - else: - pert_cell_counts_preds = self.gene_decoder(latent_output) - output_dict["pert_cell_counts_preds"] = pert_cell_counts_preds - - return output_dict - - # def configure_optimizers(self): - # read_depth_lr_multiplier = 100 - - # read_depth_params = [] - # other_params = [] - - # for name, param in self.named_parameters(): - # if "read_depth" in name: - # read_depth_params.append(param) - # else: - # other_params.append(param) - - # optimizer = torch.optim.Adam([ - # {'params': other_params, 'lr': self.lr}, - # {'params': read_depth_params, 'lr': self.lr * read_depth_lr_multiplier} - # ]) - # return optimizer From 7b20f1be9108f10e97effdee17c70bbe8ab697eb Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Sat, 21 Jun 2025 22:48:35 -0700 Subject: [PATCH 03/21] removed more old --- src/state_sets/sets/models/decoder_only.py | 135 ----------- src/state_sets/sets/models/decoders_nb.py | 110 --------- src/state_sets/sets/models/old_neural_ot.py | 240 -------------------- 3 files changed, 485 deletions(-) delete mode 100644 src/state_sets/sets/models/decoder_only.py delete mode 100644 src/state_sets/sets/models/decoders_nb.py delete mode 100644 src/state_sets/sets/models/old_neural_ot.py diff --git a/src/state_sets/sets/models/decoder_only.py b/src/state_sets/sets/models/decoder_only.py deleted file mode 100644 index ad0b2ce2..00000000 --- a/src/state_sets/sets/models/decoder_only.py +++ /dev/null @@ -1,135 +0,0 @@ -# File: models/decoder_only.py - -import torch -from geomloss import SamplesLoss - -from .base import PerturbationModel -from .utils import get_activation_class - - -class DecoderOnlyPerturbationModel(PerturbationModel): - """ - DecoderOnlyPerturbationModel learns to map the ground truth latent embedding - (provided in batch["pert_cell_emb"]) to the ground truth HVG space (batch["pert_cell_counts"]). - - Unlike the other perturbation models that compute a control mapping (e.g. via a mapping strategy), - this model simply feeds the latent representation through a decoder network. The loss is computed - between the decoder output and the target HVG expression. - - It keeps the overall architectural style (and uses the SamplesLoss loss function from geomloss) - as in the OldNeuralOT model. - """ - - def __init__( - self, - input_dim: int, - hidden_dim: int, - output_dim: int, - pert_dim: int, - n_decoder_layers: int = 2, - dropout: float = 0.0, - distributional_loss: str = "energy", - output_space: str = "gene", - gene_dim=None, - **kwargs, - ): - super().__init__( - input_dim=input_dim, - hidden_dim=hidden_dim, - gene_dim=gene_dim, - output_dim=output_dim, - pert_dim=pert_dim, - output_space=output_space, - **kwargs, - ) - self.n_decoder_layers = n_decoder_layers - self.dropout = dropout - self.distributional_loss = distributional_loss - self.cell_sentence_len = kwargs["transformer_backbone_kwargs"]["n_positions"] - self.activation_class = get_activation_class(kwargs.get("activation", "gelu")) - self.gene_dim = gene_dim - - # Use the same loss function as OldNeuralOT (e.g. using the MMD loss via geomloss) - self.loss_fn = SamplesLoss(loss=self.distributional_loss) - - def _build_networks(self): - pass - - def forward(self, batch: dict) -> torch.Tensor: - """ - Forward pass: use the ground truth latent embedding (batch["pert_cell_emb"]) as the prediction. - """ - latent = batch["pert_cell_emb"] - return latent - - def training_step(self, batch, batch_idx): - """ - Training step: The decoder output is compared against the target HVG expression. - We assume that when output_space=="gene", the target is in batch["pert_cell_counts"]. - The predictions and targets are reshaped (using a cell sentence length, if provided) - before computing the loss. - """ - pred = self(batch) - # log a zero tensor - self.log("train_loss", 0.0) - - if self.gene_decoder is not None and "pert_cell_counts" in batch: - pert_cell_counts_preds = self.gene_decoder(pred) - pert_cell_counts_preds = pert_cell_counts_preds.reshape(-1, self.cell_sentence_len, self.gene_dim) - gene_targets = batch["pert_cell_counts"] - gene_targets = gene_targets.reshape(-1, self.cell_sentence_len, self.gene_dim) - decoder_loss = self.loss_fn(pert_cell_counts_preds, gene_targets).mean() - self.log("decoder_loss", decoder_loss) - else: - self.log("decoder_loss", 0.0) - decoder_loss = None - return decoder_loss - - def validation_step(self, batch, batch_idx): - pred = self(batch) - self.log("val_loss", 0.0) - - return {"loss": None, "predictions": pred} - - def on_validation_batch_end(self, outputs, batch, batch_idx, dataloader_idx=0): - preds = outputs["predictions"] - - if self.gene_decoder is not None and "pert_cell_counts" in batch: - pert_cell_counts_preds = self.gene_decoder(preds) - gene_targets = batch["pert_cell_counts"] - pert_cell_counts_preds = pert_cell_counts_preds.reshape(-1, self.cell_sentence_len, self.gene_dim) - gene_targets = gene_targets.reshape(-1, self.cell_sentence_len, self.gene_dim) - decoder_loss = self.loss_fn(pert_cell_counts_preds, gene_targets).mean() - self.log("decoder_val_loss", decoder_loss) - - def test_step(self, batch, batch_idx): - pred = self(batch) - - if self.gene_decoder is not None and "pert_cell_counts" in batch: - pert_cell_counts_preds = self.gene_decoder(pred) - gene_targets = batch["pert_cell_counts"] - gene_targets = gene_targets.reshape(-1, self.cell_sentence_len, self.gene_dim) - decoder_loss = self.loss_fn(pert_cell_counts_preds, gene_targets).mean() - self.log("decoder_test_loss", decoder_loss) - return {"loss": None, "predictions": pred} - - def predict_step(self, batch, batch_idx, padded=True, **kwargs): - """ - Typically used for final inference. We'll replicate old logic: - returning 'preds', 'X', 'pert_name', etc. - """ - latent_output = self.forward(batch) # shape [B, ...] - output_dict = { - "preds": latent_output, - "pert_cell_emb": batch.get("pert_cell_emb", None), - "pert_cell_counts": batch.get("pert_cell_counts", None), - "pert_name": batch.get("pert_name", None), - "celltype_name": batch.get("cell_type", None), - "batch": batch.get("batch", None), - "ctrl_cell_emb": batch.get("ctrl_cell_emb", None), - } - - pert_cell_counts_preds = self.gene_decoder(latent_output) - output_dict["pert_cell_counts_preds"] = pert_cell_counts_preds - - return output_dict diff --git a/src/state_sets/sets/models/decoders_nb.py b/src/state_sets/sets/models/decoders_nb.py deleted file mode 100644 index de562893..00000000 --- a/src/state_sets/sets/models/decoders_nb.py +++ /dev/null @@ -1,110 +0,0 @@ -# models/decoders_nb.py -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.distributions import NegativeBinomial - - -class NBDecoder(nn.Module): - """ - scVI‑style decoder that maps a latent embedding (optionally with batch covariates) - to the parameters of a negative‑binomial (or ZINB) distribution over raw counts. - - Y_ig ~ NB(μ_ig, θ_g) where - μ_ig = l_i * softplus(W_g z_i + b_g) - θ_g = softplus(r_g) (gene‑specific inverse dispersion) - - Optionally, a zero‑inflation gate π_ig can be produced (not shown here). - """ - - def __init__( - self, - latent_dim: int, - gene_dim: int, - hidden_dims=[1024, 256, 256], - dropout: float = 0.0, - use_zero_inflation: bool = False, - ): - super().__init__() - modules = [] - in_features = latent_dim - for h in hidden_dims: - modules += [ - nn.Linear(in_features, h), - nn.LayerNorm(h), - nn.GELU(), - nn.Dropout(dropout), - ] - in_features = h - self.encoder = nn.Sequential(*modules) - - self.skip = nn.Identity() if in_features == latent_dim else nn.Linear(latent_dim, in_features, bias=False) - self.post_norm = nn.LayerNorm(in_features) - - # Mean parameter - self.px_scale = nn.Linear(in_features, gene_dim) - - self.l_encoder = nn.Linear(in_features, 1) - - # Gene‑specific inverse dispersion (log‑space, broadcasted) - self.log_theta = nn.Parameter(torch.randn(gene_dim)) - - # Optional zero‑inflation gate - self.use_zero_inflation = use_zero_inflation - if use_zero_inflation: - self.px_dropout = nn.Linear(in_features, gene_dim) - - @property - def theta(self): - # softplus to keep positive - return F.softplus(self.log_theta) - - def forward(self, z: torch.Tensor, log_library: torch.Tensor | None = None): - """ - z: [B, latent_dim] - log_library: [B, 1] (optional – if None we predict it) - returns μ, θ (and π if requested) - """ - flat = False - if z.dim() == 3: # [B,S,D] → flatten - B, S, D = z.shape - z = z.reshape(-1, D) - flat = True - - h = self.encoder(z) # [B* S, H] - h = self.post_norm(h + self.skip(z)) - - if log_library is None: - log_library = self.l_encoder(h) # [B* S, 1] - px_scale = F.softplus(self.px_scale(h)) # [B* S, G] - mu = torch.exp(log_library) * px_scale # NB mean - - if self.use_zero_inflation: - pi = torch.sigmoid(self.px_dropout(h)) - outs = (mu, self.theta, pi) - else: - outs = (mu, self.theta) - - if flat: # reshape back to [B,S,*] - mu = mu.reshape(B, S, -1) - if self.use_zero_inflation: - pi = pi.reshape(B, S, -1) - return mu, self.theta, pi # θ remains [G] - else: - return mu, self.theta - return outs - - def gene_dim(self) -> int: - return self.px_scale.out_features - - -def nb_nll(x, mu, theta, eps: float = 1e-6): - """ - Negative‑binomial negative log‑likelihood. - x, mu : [..., G] - theta : [G] or [..., G] - returns scalar - """ - logits = (mu + eps).log() - (theta + eps).log() # NB parameterisation - dist = NegativeBinomial(total_count=theta, logits=logits) - return -dist.log_prob(x).mean() diff --git a/src/state_sets/sets/models/old_neural_ot.py b/src/state_sets/sets/models/old_neural_ot.py deleted file mode 100644 index 0c33283d..00000000 --- a/src/state_sets/sets/models/old_neural_ot.py +++ /dev/null @@ -1,240 +0,0 @@ -from typing import Dict, Optional - -import torch -from geomloss import SamplesLoss - -from .base import PerturbationModel -from .utils import build_mlp, get_activation_class, get_transformer_backbone - - -class OldNeuralOTPerturbationModel(PerturbationModel): - """ - This model: - 1) Projects basal expression and perturbation encodings into a shared latent space. - 2) Uses an OT-based distributional loss (energy, sinkhorn, etc.) from geomloss. - 3) Enables cells to attend to one another, learning a set-to-set function rather than - a sample-to-sample single-cell map. - """ - - def __init__( - self, - input_dim: int, - hidden_dim: int, - output_dim: int, - pert_dim: int, - predict_residual: bool = True, - distributional_loss: str = "energy", - transformer_backbone_key: str = "GPT2", - transformer_backbone_kwargs: dict = None, - output_space: str = "gene", - gene_dim: Optional[int] = None, - **kwargs, - ): - """ - Args: - input_dim: dimension of the input expression (e.g. number of genes or embedding dimension). - hidden_dim: not necessarily used, but required by PerturbationModel signature. - output_dim: dimension of the output space (genes or latent). - pert_dim: dimension of perturbation embedding. - gpt: e.g. "TranslationTransformerSamplesModel". - model_kwargs: dictionary passed to that model's constructor. - loss: choice of distributional metric ("sinkhorn", "energy", etc.). - **kwargs: anything else to pass up to PerturbationModel or not used. - """ - # Call the parent PerturbationModel constructor - super().__init__( - input_dim=input_dim, - hidden_dim=hidden_dim, - gene_dim=gene_dim, - output_dim=output_dim, - pert_dim=pert_dim, - output_space=output_space, - **kwargs, - ) - - # Save or store relevant hyperparams - self.predict_residual = predict_residual - self.n_encoder_layers = kwargs.get("n_encoder_layers", 2) - self.n_decoder_layers = kwargs.get("n_decoder_layers", 2) - self.activation_class = get_activation_class(kwargs.get("activation", "gelu")) - self.transformer_backbone_key = transformer_backbone_key - self.transformer_backbone_kwargs = transformer_backbone_kwargs - self.distributional_loss = distributional_loss - self.cell_sentence_len = self.transformer_backbone_kwargs["n_positions"] - self.gene_dim = gene_dim - - # Build the distributional loss from geomloss - self.loss_fn = SamplesLoss(loss=self.distributional_loss) - # self.loss_fn = LearnableAlignmentLoss() - - # Build the underlying neural OT network - self._build_networks() - - def _build_networks(self): - """ - Here we instantiate the actual GPT2-based model or any neuralOT translator - via your old get_model(model_key, model_kwargs) approach. - """ - self.pert_encoder = build_mlp( - in_dim=self.pert_dim, - out_dim=self.hidden_dim, - hidden_dim=self.hidden_dim, - n_layers=self.n_encoder_layers, - dropout=self.dropout, - activation=self.activation_class, - ) - - # Map the input embedding to the hidden space - self.basal_encoder = build_mlp( - in_dim=self.input_dim, - out_dim=self.hidden_dim, - hidden_dim=self.hidden_dim, - n_layers=self.n_encoder_layers, - dropout=self.dropout, - activation=self.activation_class, - ) - - self.transformer_backbone, self.transformer_model_dim = get_transformer_backbone( - self.transformer_backbone_key, - self.transformer_backbone_kwargs, - ) - - self.project_out = build_mlp( - in_dim=self.hidden_dim, - out_dim=self.output_dim, - hidden_dim=self.hidden_dim, - n_layers=self.n_decoder_layers, - dropout=self.dropout, - activation=self.activation_class, - ) - - print(self) - - def encode_perturbation(self, pert: torch.Tensor) -> torch.Tensor: - """If needed, define how we embed the raw perturbation input.""" - return self.pert_encoder(pert) - - def encode_basal_expression(self, expr: torch.Tensor) -> torch.Tensor: - """Define how we embed basal state input, if needed.""" - return self.basal_encoder(expr) - - def perturb(self, pert: torch.Tensor, basal: torch.Tensor) -> torch.Tensor: - """ - Return the latent perturbed state given the perturbation and basal state. - """ - pert_embedding = self.encode_perturbation(pert).unsqueeze(1) # shape: [batch_size, 1, hidden_dim] - control_cells = self.encode_basal_expression(basal).unsqueeze(1) # shape: [batch_size, 1, hidden_dim] - cls_input = torch.zeros_like(pert_embedding) # shape: [batch_size, 1, hidden_dim] - seq_input = torch.cat([pert_embedding, control_cells, cls_input], dim=1) # shape: [batch_size, 3, hidden_dim] - - # forward pass + extract CLS last hidden state - prediction = self.transformer_backbone(inputs_embeds=seq_input).last_hidden_state[:, -1] - - # add to basal if predicting residual - if self.predict_residual: - # treat the actual prediction as a residual sum to basal - return prediction + control_cells.squeeze(1) - else: - return prediction - - def forward(self, batch: dict) -> torch.Tensor: - """ - The main forward call. - """ - prediction = self.perturb(batch["pert_emb"], batch["ctrl_cell_emb"]) - output = self.project_out(prediction) - - return output - - def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> torch.Tensor: - """Training step logic for both main model and decoder.""" - # Get model predictions (in latent space) - pred = self(batch) - pred = pred.reshape(-1, self.cell_sentence_len, self.output_dim) - # TODO: please improve this, do not assume self.cell_sentence_len for this model - target = batch["pert_cell_emb"] - target = target.reshape(-1, self.cell_sentence_len, self.output_dim) - main_loss = self.loss_fn(pred, target).mean() - self.log("train_loss", main_loss) - - # Process decoder if available - decoder_loss = None - if self.gene_decoder is not None and "pert_cell_counts" in batch: - # Train decoder to map latent predictions to gene space - with torch.no_grad(): - latent_preds = pred.detach() # Detach to prevent gradient flow back to main model - - pert_cell_counts_preds = self.gene_decoder(latent_preds) - gene_targets = batch["pert_cell_counts"] - gene_targets = gene_targets.reshape(-1, self.cell_sentence_len, self.gene_dim) - decoder_loss = self.loss_fn(pert_cell_counts_preds, gene_targets).mean() - - # Log decoder loss - self.log("decoder_loss", decoder_loss) - - total_loss = main_loss + decoder_loss - else: - total_loss = main_loss - - return total_loss - - def validation_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> None: - """Validation step logic.""" - pred = self(batch) - pred = pred.reshape(-1, self.cell_sentence_len, self.output_dim) - target = batch["pert_cell_emb"] - target = target.reshape(-1, self.cell_sentence_len, self.output_dim) - loss = self.loss_fn(pred, target).mean() - self.log("val_loss", loss) - - return {"loss": loss, "predictions": pred} - - def on_validation_batch_end(self, outputs, batch, batch_idx, dataloader_idx=0) -> None: - """Track decoder performance during validation without training it.""" - if self.gene_decoder is not None and "pert_cell_counts" in batch: - # Get model predictions from validation step - latent_preds = outputs["predictions"] - - # Train decoder to map latent predictions to gene space - pert_cell_counts_preds = self.gene_decoder(latent_preds) # verify this is automatically detached - gene_targets = batch["pert_cell_counts"] - - # Get decoder predictions - pert_cell_counts_preds = pert_cell_counts_preds.reshape(-1, self.cell_sentence_len, self.gene_dim) - gene_targets = gene_targets.reshape(-1, self.cell_sentence_len, self.gene_dim) - decoder_loss = self.loss_fn(pert_cell_counts_preds, gene_targets).mean() - - # Log the validation metric - self.log("decoder_val_loss", decoder_loss) - - def test_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> None: - pred = self.forward(batch, padded=False) - target = batch["pert_cell_emb"] - pred = pred.reshape(1, -1, self.output_dim) - target = target.reshape(1, -1, self.output_dim) - loss = self.loss_fn(pred, target).mean() - self.log("test_loss", loss) - pred = pred.reshape(-1, self.output_dim) - target = target.reshape(-1, self.output_dim) - - def predict_step(self, batch, batch_idx, padded=True, **kwargs): - """ - Typically used for final inference. We'll replicate old logic: - returning 'preds', 'X', 'pert_name', etc. - """ - latent_output = self.forward(batch) # shape [B, ...] - output_dict = { - "preds": latent_output, - "pert_cell_emb": batch.get("pert_cell_emb", None), - "pert_cell_counts": batch.get("pert_cell_counts", None), - "pert_name": batch.get("pert_name", None), - "celltype_name": batch.get("cell_type", None), - "batch": batch.get("batch", None), - "ctrl_cell_emb": batch.get("ctrl_cell_emb", None), - } - - if self.gene_decoder is not None: - pert_cell_counts_preds = self.gene_decoder(latent_output) - output_dict["pert_cell_counts_preds"] = pert_cell_counts_preds - - return output_dict From ff5cea796a1bf76d49fe427b21553dfd709a433e Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Sat, 21 Jun 2025 23:10:57 -0700 Subject: [PATCH 04/21] init --- src/state_sets/sets/models/__init__.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/src/state_sets/sets/models/__init__.py b/src/state_sets/sets/models/__init__.py index 6ec2ebf6..b20528fe 100644 --- a/src/state_sets/sets/models/__init__.py +++ b/src/state_sets/sets/models/__init__.py @@ -1,12 +1,9 @@ from .base import PerturbationModel from .cell_context_mean import CellContextPerturbationModel from .cell_type_mean import CellTypeMeanModel -from .decoder_only import DecoderOnlyPerturbationModel from .embed_sum import EmbedSumPerturbationModel from .global_simple_sum import GlobalSimpleSumPerturbationModel -from .old_neural_ot import OldNeuralOTPerturbationModel from .pert_sets import PertSetsPerturbationModel -from .pseudobulk import PseudobulkPerturbationModel __all__ = [ "PerturbationModel", @@ -15,7 +12,4 @@ "CellContextPerturbationModel", "EmbedSumPerturbationModel", "PertSetsPerturbationModel", - "OldNeuralOTPerturbationModel", - "DecoderOnlyPerturbationModel", - "PseudobulkPerturbationModel", ] From c91befdd99de908f52e0689ab8251a42b0f1bd49 Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Sat, 21 Jun 2025 23:11:53 -0700 Subject: [PATCH 05/21] utils init --- src/state_sets/sets/utils/__init__.py | 36 --------------------------- 1 file changed, 36 deletions(-) diff --git a/src/state_sets/sets/utils/__init__.py b/src/state_sets/sets/utils/__init__.py index fe0a2f1a..75fe5725 100644 --- a/src/state_sets/sets/utils/__init__.py +++ b/src/state_sets/sets/utils/__init__.py @@ -109,18 +109,6 @@ def get_lightning_module(model_type: str, data_config: dict, model_config: dict, batch_dim=var_dims["batch_dim"], **module_config, ) - elif model_type.lower() == "old_neuralot": - from ...sets.models.old_neural_ot import OldNeuralOTPerturbationModel - - return OldNeuralOTPerturbationModel( - input_dim=var_dims["input_dim"], - gene_dim=gene_dim, - hvg_dim=var_dims["hvg_dim"], - output_dim=var_dims["output_dim"], - pert_dim=var_dims["pert_dim"], - batch_dim=var_dims["batch_dim"], - **module_config, - ) elif model_type.lower() == "neuralot" or model_type.lower() == "pertsets": from ...sets.models.pert_sets import PertSetsPerturbationModel @@ -169,30 +157,6 @@ def get_lightning_module(model_type: str, data_config: dict, model_config: dict, batch_dim=var_dims["batch_dim"], **module_config, ) - elif model_type.lower() == "decoder_only": - from ...sets.models.decoder_only import DecoderOnlyPerturbationModel - - return DecoderOnlyPerturbationModel( - input_dim=var_dims["input_dim"], - gene_dim=gene_dim, - hvg_dim=var_dims["hvg_dim"], - output_dim=var_dims["output_dim"], - pert_dim=var_dims["pert_dim"], - batch_dim=var_dims["batch_dim"], - **module_config, - ) - elif model_type.lower() == "pseudobulk": - from ...sets.models.pseudobulk import PseudobulkPerturbationModel - - return PseudobulkPerturbationModel( - input_dim=var_dims["input_dim"], - gene_dim=gene_dim, - hvg_dim=var_dims["hvg_dim"], - output_dim=var_dims["output_dim"], - pert_dim=var_dims["pert_dim"], - batch_dim=var_dims["batch_dim"], - **module_config, - ) elif model_type.lower() == "cpa": from ...sets.models.cpa import CPAPerturbationModel From fd9aadc0c66bc5371681f7f9d1ddf5a62d29d7d2 Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Wed, 11 Jun 2025 15:30:18 -0700 Subject: [PATCH 06/21] decoder fix started --- src/state_sets/_cli/_sets/_train.py | 17 +++++++ .../configs/model/tahoe_decoder_test.yaml | 44 ++++++++++++++++++ .../configs/model/tahoe_llama_93133848.yaml | 1 + src/state_sets/sets/models/base.py | 46 +++++++++++++++---- 4 files changed, 98 insertions(+), 10 deletions(-) create mode 100644 src/state_sets/configs/model/tahoe_decoder_test.yaml diff --git a/src/state_sets/_cli/_sets/_train.py b/src/state_sets/_cli/_sets/_train.py index 987b8228..9f8ecc50 100644 --- a/src/state_sets/_cli/_sets/_train.py +++ b/src/state_sets/_cli/_sets/_train.py @@ -115,6 +115,23 @@ def run_sets_train(cfg: DictConfig): data_module.save_state(f) data_module.setup(stage="fit") + + var_dims = data_module.get_var_dims() # {"gene_dim": …, "hvg_dim": …} + gene_dim = var_dims.get("gene_dim", 5000) # fallback if key missing + latent_dim = cfg["model"]["kwargs"]["output_dim"] # same as model.output_dim + # optional: let user override from CLI/YAML + hidden_dims = cfg["model"]["kwargs"].get("decoder_hidden_dims", [1024, 1024, 512]) + + decoder_cfg = dict( + latent_dim = latent_dim, + gene_dim = gene_dim, + hidden_dims = hidden_dims, + dropout = cfg["model"]["kwargs"].get("decoder_dropout", 0.1), + residual_decoder = cfg["model"]["kwargs"].get("residual_decoder", False), + ) + + # tuck it into the kwargs that will reach the LightningModule + cfg["model"]["kwargs"]["decoder_cfg"] = decoder_cfg if cfg["model"]["name"].lower() in ["cpa", "scvi"] or cfg["model"]["name"].lower().startswith("scgpt"): cfg["model"]["kwargs"]["n_cell_types"] = len(data_module.celltype_onehot_map) diff --git a/src/state_sets/configs/model/tahoe_decoder_test.yaml b/src/state_sets/configs/model/tahoe_decoder_test.yaml new file mode 100644 index 00000000..1a0c985d --- /dev/null +++ b/src/state_sets/configs/model/tahoe_decoder_test.yaml @@ -0,0 +1,44 @@ +name: PertSets +checkpoint: null +device: cuda + +kwargs: + cell_set_len: 512 + decoder_hidden_dims: [2048, 2048, 2048] + blur: 0.05 + hidden_dim: 1488 # hidden dimension going into the transformer backbone + loss: energy + confidence_head: False + n_encoder_layers: 4 + n_decoder_layers: 4 + predict_residual: True + softplus: True + freeze_pert: False + transformer_decoder: False + finetune_vci_decoder: False + residual_decoder: False + decoder_loss_weight: 1.0 + batch_encoder: False + nb_decoder: False + mask_attn: False + distributional_loss: energy + init_from: null + transformer_backbone_key: llama + transformer_backbone_kwargs: + max_position_embeddings: ${model.kwargs.cell_set_len} + hidden_size: ${model.kwargs.hidden_dim} + intermediate_size: 5952 + num_hidden_layers: 6 + num_attention_heads: 12 + num_key_value_heads: 12 + head_dim: 124 + use_cache: false + attention_dropout: 0.0 + hidden_dropout: 0.0 + layer_norm_eps: 1e-6 + pad_token_id: 0 + bos_token_id: 1 + eos_token_id: 2 + tie_word_embeddings: false + rotary_dim: 0 + use_rotary_embeddings: false diff --git a/src/state_sets/configs/model/tahoe_llama_93133848.yaml b/src/state_sets/configs/model/tahoe_llama_93133848.yaml index f5152789..17b183d7 100644 --- a/src/state_sets/configs/model/tahoe_llama_93133848.yaml +++ b/src/state_sets/configs/model/tahoe_llama_93133848.yaml @@ -4,6 +4,7 @@ device: cuda kwargs: cell_set_len: 512 + decoder_hidden_dims: [1024, 1024, 512] blur: 0.05 hidden_dim: 696 # hidden dimension going into the transformer backbone loss: energy diff --git a/src/state_sets/sets/models/base.py b/src/state_sets/sets/models/base.py index 6c837541..92f5c473 100644 --- a/src/state_sets/sets/models/base.py +++ b/src/state_sets/sets/models/base.py @@ -5,6 +5,7 @@ import torch import torch.nn as nn from lightning.pytorch import LightningModule +import typing as tp from .utils import get_loss_fn @@ -147,9 +148,11 @@ def __init__( batch_size: int = 64, gene_dim: int = 5000, hvg_dim: int = 2001, + decoder_cfg: dict | None = None, **kwargs, ): super().__init__() + self.decoder_cfg = decoder_cfg self.save_hyperparameters() # Core architecture settings @@ -196,16 +199,22 @@ def __init__( elif "PBS" in self.control_pert: hidden_dims = [2048, 1024, 1024] else: - hidden_dims = [1024, 1024, 512] # make this config - - self.gene_decoder = LatentToGeneDecoder( - latent_dim=self.output_dim, - gene_dim=gene_dim, - hidden_dims=hidden_dims, - dropout=dropout, - residual_decoder=self.residual_decoder, - ) - logger.info(f"Initialized gene decoder for embedding {embed_key} to gene space") + if "DMSO_TF" in self.control_pert: + if self.residual_decoder: + hidden_dims = [2058, 2058, 2058, 2058, 2058] + else: + hidden_dims = [4096, 2048, 2048] + else: + hidden_dims = [1024, 1024, 512] # make this config + + self.gene_decoder = LatentToGeneDecoder( + latent_dim=self.output_dim, + gene_dim=gene_dim, + hidden_dims=hidden_dims, + dropout=dropout, + residual_decoder=self.residual_decoder, + ) + logger.info(f"Initialized gene decoder for embedding {embed_key} to gene space") def transfer_batch_to_device(self, batch, device, dataloader_idx: int): return {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()} @@ -215,6 +224,23 @@ def _build_networks(self): """Build the core neural network components.""" pass + def _build_decoder(self): + """Create self.gene_decoder from self.decoder_cfg (or leave None).""" + if self.decoder_cfg is None: + self.gene_decoder = None + return + self.gene_decoder = LatentToGeneDecoder(**self.decoder_cfg) + + def on_load_checkpoint(self, checkpoint: dict[str, tp.Any]) -> None: + """ + Lightning calls this *before* the checkpoint's state_dict is loaded. + Re-create the decoder using the exact hyper-parameters saved in the ckpt, + so that parameter shapes match and load_state_dict succeeds. + """ + self.decoder_cfg = checkpoint["hyper_parameters"]["decoder_cfg"] + self._build_decoder() + + def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> torch.Tensor: """Training step logic for both main model and decoder.""" # Get model predictions (in latent space) From 2536d80418976d46c0266112881b5c7d50a9a2fd Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Fri, 13 Jun 2025 13:38:33 -0700 Subject: [PATCH 07/21] base infer not working --- src/state_sets/_cli/_sets/_infer.py | 94 +++++++++++++++++++++++++++++ src/state_sets/_cli/_sets/_train.py | 4 +- 2 files changed, 97 insertions(+), 1 deletion(-) create mode 100644 src/state_sets/_cli/_sets/_infer.py diff --git a/src/state_sets/_cli/_sets/_infer.py b/src/state_sets/_cli/_sets/_infer.py new file mode 100644 index 00000000..e8a71058 --- /dev/null +++ b/src/state_sets/_cli/_sets/_infer.py @@ -0,0 +1,94 @@ +import argparse +import scanpy as sc +import torch +import numpy as np +import os +import pandas as pd + +# Adjust this import to your project structure +from ...sets.models.pert_sets import PertSetsPerturbationModel + +def add_arguments_infer(parser: argparse.ArgumentParser): + parser.add_argument("--checkpoint", type=str, required=True, help="Path to model checkpoint (.ckpt)") + parser.add_argument("--adata", type=str, required=True, help="Path to input AnnData file (.h5ad)") + parser.add_argument("--embed_key", type=str, default="X_hvg", help="Key in adata.obsm for input features") + parser.add_argument("--pert_col", type=str, default="drugname_drugconc", help="Column in adata.obs for perturbation labels") + parser.add_argument("--output", type=str, default=None, help="Path to output AnnData file (.h5ad)") + + +def run_sets_infer(args): + import logging + logging.basicConfig(level=logging.INFO) + logger = logging.getLogger(__name__) + + # Load model + logger.info(f"Loading model from checkpoint: {args.checkpoint}") + model = PertSetsPerturbationModel.load_from_checkpoint(args.checkpoint) + model.eval() + device = next(model.parameters()).device + + # Use model's config for batch prep + pert_onehot_map = getattr(model, "pert_onehot_map", None) + pert_dim = model.pert_dim + + # Load AnnData + logger.info(f"Loading AnnData from: {args.adata}") + adata = sc.read_h5ad(args.adata) + + # Get input features + if args.embed_key in adata.obsm: + X = adata.obsm[args.embed_key] + logger.info(f"Using adata.obsm['{args.embed_key}'] as input features: shape {X.shape}") + else: + X = adata.X + logger.info(f"Using adata.X as input features: shape {X.shape}") + X = torch.tensor(X, dtype=torch.float32).to(device) + + # Prepare perturbation tensor using the model's map + pert_names = adata.obs[args.pert_col].values + pert_tensor = torch.zeros((len(pert_names), pert_dim), device=device) + if pert_onehot_map is not None: + for idx, name in enumerate(pert_names): + if name in pert_onehot_map: + pert_tensor[idx, pert_onehot_map[name]] = 1 + else: + # Optionally handle unknown perturbations + pass + else: + # Fallback: build map from AnnData (not recommended for production) + unique_perts = sorted(set(pert_names)) + pert_map = {name: i for i, name in enumerate(unique_perts)} + for idx, name in enumerate(pert_names): + pert_tensor[idx, pert_map[name]] = 1 + + # Prepare batch + batch = { + "ctrl_cell_emb": X, + "pert_emb": pert_tensor, + "pert_name": pert_names.tolist(), + # "batch": torch.zeros((1, cell_sentence_len), device=device) + } + # when do we need the batch num things + + # Inference + logger.info("Running inference...") + with torch.no_grad(): + preds = model.forward(batch) + preds_np = preds.cpu().numpy() + + # Save predictions to AnnData + pred_key = "model_preds" + adata.obsm[pred_key] = preds_np + output_path = args.output or args.adata.replace(".h5ad", "_with_preds.h5ad") + adata.write_h5ad(output_path) + logger.info(f"Saved predictions to {output_path} (in adata.obsm['{pred_key}'])") + + +def main(): + parser = argparse.ArgumentParser(description="Run inference on AnnData with a trained model checkpoint.") + add_arguments_infer(parser) + args = parser.parse_args() + run_sets_infer(args) + +if __name__ == "__main__": + main() diff --git a/src/state_sets/_cli/_sets/_train.py b/src/state_sets/_cli/_sets/_train.py index 9f8ecc50..586750cd 100644 --- a/src/state_sets/_cli/_sets/_train.py +++ b/src/state_sets/_cli/_sets/_train.py @@ -119,7 +119,6 @@ def run_sets_train(cfg: DictConfig): var_dims = data_module.get_var_dims() # {"gene_dim": …, "hvg_dim": …} gene_dim = var_dims.get("gene_dim", 5000) # fallback if key missing latent_dim = cfg["model"]["kwargs"]["output_dim"] # same as model.output_dim - # optional: let user override from CLI/YAML hidden_dims = cfg["model"]["kwargs"].get("decoder_hidden_dims", [1024, 1024, 512]) decoder_cfg = dict( @@ -133,6 +132,9 @@ def run_sets_train(cfg: DictConfig): # tuck it into the kwargs that will reach the LightningModule cfg["model"]["kwargs"]["decoder_cfg"] = decoder_cfg + cfg["data"]["kwargs"]["n_perts"] = len(data_module.pert_onehot_map) + cfg["model"]["kwargs"]["pert_onehot_map"] = data_module.pert_onehot_map + if cfg["model"]["name"].lower() in ["cpa", "scvi"] or cfg["model"]["name"].lower().startswith("scgpt"): cfg["model"]["kwargs"]["n_cell_types"] = len(data_module.celltype_onehot_map) cfg["model"]["kwargs"]["n_perts"] = len(data_module.pert_onehot_map) From 778f39f7293d483abb426c8160378c06c558d977 Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Mon, 16 Jun 2025 10:22:35 -0700 Subject: [PATCH 08/21] src --- src/state_sets/_cli/_sets/_infer.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/src/state_sets/_cli/_sets/_infer.py b/src/state_sets/_cli/_sets/_infer.py index e8a71058..2617151d 100644 --- a/src/state_sets/_cli/_sets/_infer.py +++ b/src/state_sets/_cli/_sets/_infer.py @@ -5,7 +5,6 @@ import os import pandas as pd -# Adjust this import to your project structure from ...sets.models.pert_sets import PertSetsPerturbationModel def add_arguments_infer(parser: argparse.ArgumentParser): @@ -14,6 +13,8 @@ def add_arguments_infer(parser: argparse.ArgumentParser): parser.add_argument("--embed_key", type=str, default="X_hvg", help="Key in adata.obsm for input features") parser.add_argument("--pert_col", type=str, default="drugname_drugconc", help="Column in adata.obs for perturbation labels") parser.add_argument("--output", type=str, default=None, help="Path to output AnnData file (.h5ad)") + parser.add_argument("--celltype_col", type=str, default=None, help="Column in adata.obs for cell type labels (optional)") + parser.add_argument("--celltypes", type=str, default=None, help="Comma-separated list of cell types to include (optional)") def run_sets_infer(args): @@ -25,6 +26,7 @@ def run_sets_infer(args): logger.info(f"Loading model from checkpoint: {args.checkpoint}") model = PertSetsPerturbationModel.load_from_checkpoint(args.checkpoint) model.eval() + cell_sentence_len = model.cell_sentence_len device = next(model.parameters()).device # Use model's config for batch prep @@ -35,6 +37,19 @@ def run_sets_infer(args): logger.info(f"Loading AnnData from: {args.adata}") adata = sc.read_h5ad(args.adata) + # Optionally filter by cell type + if args.celltype_col is not None and args.celltypes is not None: + celltypes = [ct.strip() for ct in args.celltypes.split(",")] + if args.celltype_col not in adata.obs: + raise ValueError(f"Column '{args.celltype_col}' not found in adata.obs.") + initial_n = adata.n_obs + adata = adata[adata.obs[args.celltype_col].isin(celltypes)].copy() + logger.info(f"Filtered AnnData to {adata.n_obs} cells of types {celltypes} (from {initial_n} cells)") + elif args.celltype_col is not None: + if args.celltype_col not in adata.obs: + raise ValueError(f"Column '{args.celltype_col}' not found in adata.obs.") + logger.info(f"No cell type filtering applied, but cell type column '{args.celltype_col}' is available.") + # Get input features if args.embed_key in adata.obsm: X = adata.obsm[args.embed_key] @@ -66,11 +81,10 @@ def run_sets_infer(args): "ctrl_cell_emb": X, "pert_emb": pert_tensor, "pert_name": pert_names.tolist(), - # "batch": torch.zeros((1, cell_sentence_len), device=device) + "batch": torch.zeros((1, cell_sentence_len), device=device) } # when do we need the batch num things - # Inference logger.info("Running inference...") with torch.no_grad(): preds = model.forward(batch) From d1a5dacd71e5693d3cc87d17eca71ca5e08ca7e4 Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Thu, 19 Jun 2025 22:07:48 -0700 Subject: [PATCH 09/21] infer working --- src/state_sets/__main__.py | 4 + src/state_sets/_cli/__init__.py | 3 +- src/state_sets/_cli/_sets/__init__.py | 4 +- src/state_sets/_cli/_sets/_infer.py | 180 +++++++++++++++++++----- src/state_sets/_cli/_sets/_infer2.py | 150 ++++++++++++++++++++ src/state_sets/sets/models/base.py | 36 ++++- src/state_sets/sets/models/pert_sets.py | 2 + 7 files changed, 344 insertions(+), 35 deletions(-) create mode 100644 src/state_sets/_cli/_sets/_infer2.py diff --git a/src/state_sets/__main__.py b/src/state_sets/__main__.py index ea4a7cbb..f0d27803 100644 --- a/src/state_sets/__main__.py +++ b/src/state_sets/__main__.py @@ -6,6 +6,7 @@ add_arguments_state, run_sets_predict, run_sets_train, + run_sets_infer, run_state_embed, run_state_train, ) @@ -60,6 +61,9 @@ def main(): case "predict": # For now, predict uses argparse and not hydra run_sets_predict(args) + case "infer": + # Run inference using argparse, similar to predict + run_sets_infer(args) if __name__ == "__main__": diff --git a/src/state_sets/_cli/__init__.py b/src/state_sets/_cli/__init__.py index 947c4960..5e057c70 100644 --- a/src/state_sets/_cli/__init__.py +++ b/src/state_sets/_cli/__init__.py @@ -1,4 +1,4 @@ -from ._sets import add_arguments_sets, run_sets_predict, run_sets_train +from ._sets import add_arguments_sets, run_sets_predict, run_sets_train, run_sets_infer from ._state import add_arguments_state, run_state_embed, run_state_train __all__ = [ @@ -7,5 +7,6 @@ "run_sets_train", "run_state_embed", "run_sets_predict", + "run_sets_infer", "run_state_train", ] diff --git a/src/state_sets/_cli/_sets/__init__.py b/src/state_sets/_cli/_sets/__init__.py index aede1b93..d64471bb 100644 --- a/src/state_sets/_cli/_sets/__init__.py +++ b/src/state_sets/_cli/_sets/__init__.py @@ -2,8 +2,9 @@ from ._predict import add_arguments_predict, run_sets_predict from ._train import add_arguments_train, run_sets_train +from ._infer import add_arguments_infer, run_sets_infer -__all__ = ["run_sets_train", "run_sets_predict", "add_arguments_sets"] +__all__ = ["run_sets_train", "run_sets_predict", "run_sets_infer", "add_arguments_sets"] def add_arguments_sets(parser: ap.ArgumentParser): @@ -11,3 +12,4 @@ def add_arguments_sets(parser: ap.ArgumentParser): subparsers = parser.add_subparsers(required=True, dest="subcommand") add_arguments_train(subparsers.add_parser("train")) add_arguments_predict(subparsers.add_parser("predict")) + add_arguments_infer(subparsers.add_parser("infer")) diff --git a/src/state_sets/_cli/_sets/_infer.py b/src/state_sets/_cli/_sets/_infer.py index 2617151d..60406c84 100644 --- a/src/state_sets/_cli/_sets/_infer.py +++ b/src/state_sets/_cli/_sets/_infer.py @@ -4,8 +4,16 @@ import numpy as np import os import pandas as pd +from tqdm import tqdm +import yaml from ...sets.models.pert_sets import PertSetsPerturbationModel +from cell_load.data_modules import PerturbationDataModule + +# state-sets sets infer --output_dir /home/aadduri/state-sets/test/ --checkpoint last.ckpt --adata /home/aadduri/state-sets/test/adata.h5ad --pert_col gene +# state-sets sets infer --output_dir /home/dhruvgautam/state-sets/test/ --checkpoint /large_storage/ctc/userspace/aadduri/preprint/replogle_llama_21712320_filtered_cs32_pretrained/hepg2/checkpoints/step=44000.ckpt --adata /large_storage/ctc/ML/state_sets/replogle/processed.h5 --pert_col gene + +# state-sets sets infer --output /home/dhruvgautam/state-sets/test/ --output_dir /large_storage/ctc/userspace/aadduri/preprint/replogle_state_proper_cs32_sm/hepg2 --checkpoint /large_storage/ctc/userspace/aadduri/preprint/replogle_state_proper_cs32_sm/hepg2/checkpoints/step=48000.ckpt --adata /large_storage/ctc/ML/state_sets/replogle/processed.h5 --pert_col gene --embed_key X_vci_1.5.2_4 def add_arguments_infer(parser: argparse.ArgumentParser): parser.add_argument("--checkpoint", type=str, required=True, help="Path to model checkpoint (.ckpt)") @@ -13,8 +21,10 @@ def add_arguments_infer(parser: argparse.ArgumentParser): parser.add_argument("--embed_key", type=str, default="X_hvg", help="Key in adata.obsm for input features") parser.add_argument("--pert_col", type=str, default="drugname_drugconc", help="Column in adata.obs for perturbation labels") parser.add_argument("--output", type=str, default=None, help="Path to output AnnData file (.h5ad)") + parser.add_argument("--output_dir", type=str, required=True, help="Path to the output_dir containing the config.yaml file that was saved during training.") parser.add_argument("--celltype_col", type=str, default=None, help="Column in adata.obs for cell type labels (optional)") parser.add_argument("--celltypes", type=str, default=None, help="Comma-separated list of cell types to include (optional)") + parser.add_argument("--batch_size", type=int, default=1000, help="Batch size for inference (default: 1000)") def run_sets_infer(args): @@ -22,6 +32,32 @@ def run_sets_infer(args): logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) + def load_config(cfg_path: str) -> dict: + """Load config from the YAML file that was dumped during training.""" + if not os.path.exists(cfg_path): + raise FileNotFoundError(f"Could not find config file: {cfg_path}") + with open(cfg_path, "r") as f: + cfg = yaml.safe_load(f) + return cfg + + # Load the config + config_path = os.path.join(args.output_dir, "config.yaml") + cfg = load_config(config_path) + logger.info(f"Loaded config from {config_path}") + + # Find run output directory & load data module + run_output_dir = os.path.join(cfg["output_dir"], cfg["name"]) + data_module_path = os.path.join(run_output_dir, "data_module.torch") + if not os.path.exists(data_module_path): + raise FileNotFoundError(f"Could not find data module at {data_module_path}") + data_module = PerturbationDataModule.load_state(data_module_path) + data_module.setup(stage="test") + logger.info(f"Loaded data module from {data_module_path}") + + # Get perturbation dimensions and mapping from data module + var_dims = data_module.get_var_dims() + pert_dim = var_dims["pert_dim"] + # Load model logger.info(f"Loading model from checkpoint: {args.checkpoint}") model = PertSetsPerturbationModel.load_from_checkpoint(args.checkpoint) @@ -29,10 +65,6 @@ def run_sets_infer(args): cell_sentence_len = model.cell_sentence_len device = next(model.parameters()).device - # Use model's config for batch prep - pert_onehot_map = getattr(model, "pert_onehot_map", None) - pert_dim = model.pert_dim - # Load AnnData logger.info(f"Loading AnnData from: {args.adata}") adata = sc.read_h5ad(args.adata) @@ -57,38 +89,122 @@ def run_sets_infer(args): else: X = adata.X logger.info(f"Using adata.X as input features: shape {X.shape}") - X = torch.tensor(X, dtype=torch.float32).to(device) - # Prepare perturbation tensor using the model's map + # Prepare perturbation tensor using the data module's mapping pert_names = adata.obs[args.pert_col].values - pert_tensor = torch.zeros((len(pert_names), pert_dim), device=device) - if pert_onehot_map is not None: - for idx, name in enumerate(pert_names): - if name in pert_onehot_map: - pert_tensor[idx, pert_onehot_map[name]] = 1 + pert_tensor = torch.zeros((len(pert_names), pert_dim), device='cpu') # Keep on CPU initially + logger.info(f"Perturbation tensor shape: {pert_tensor.shape}") + + # Use data module's perturbation mapping + pert_onehot_map = data_module.pert_onehot_map + + # Debug: check what's available + logger.info(f"Data module has {len(pert_onehot_map)} perturbations in mapping") + logger.info(f"First 10 perturbations in data module: {list(pert_onehot_map.keys())[:10]}") + + unique_pert_names = sorted(set(pert_names)) + logger.info(f"AnnData has {len(unique_pert_names)} unique perturbations") + logger.info(f"First 10 perturbations in AnnData: {unique_pert_names[:10]}") + + # Check overlap + overlap = set(unique_pert_names) & set(pert_onehot_map.keys()) + logger.info(f"Overlap between AnnData and data module: {len(overlap)} perturbations") + if len(overlap) < len(unique_pert_names): + missing = set(unique_pert_names) - set(pert_onehot_map.keys()) + logger.warning(f"Missing perturbations: {list(missing)[:10]}") + + # Check if there's a control perturbation that might match + control_pert = data_module.get_control_pert() + logger.info(f"Control perturbation in data module: '{control_pert}'") + + matched_count = 0 + for idx, name in enumerate(pert_names): + if name in pert_onehot_map: + pert_tensor[idx] = pert_onehot_map[name] + matched_count += 1 + else: + # For now, use control perturbation as fallback + if control_pert in pert_onehot_map: + pert_tensor[idx] = pert_onehot_map[control_pert] else: - # Optionally handle unknown perturbations - pass - else: - # Fallback: build map from AnnData (not recommended for production) - unique_perts = sorted(set(pert_names)) - pert_map = {name: i for i, name in enumerate(unique_perts)} - for idx, name in enumerate(pert_names): - pert_tensor[idx, pert_map[name]] = 1 - - # Prepare batch - batch = { - "ctrl_cell_emb": X, - "pert_emb": pert_tensor, - "pert_name": pert_names.tolist(), - "batch": torch.zeros((1, cell_sentence_len), device=device) - } - # when do we need the batch num things - - logger.info("Running inference...") + # Use first available perturbation as fallback + first_pert = list(pert_onehot_map.keys())[0] + pert_tensor[idx] = pert_onehot_map[first_pert] + + logger.info(f"Matched {matched_count} out of {len(pert_names)} perturbations") + + # Process in batches with progress bar + # Use cell_sentence_len as batch size since model expects this + n_samples = len(pert_names) + batch_size = cell_sentence_len # Model requires this exact batch size + n_batches = (n_samples + batch_size - 1) // batch_size # Ceiling division + + logger.info(f"Running inference on {n_samples} samples in {n_batches} batches of size {batch_size} (model's cell_sentence_len)...") + + all_preds = [] + with torch.no_grad(): - preds = model.forward(batch) - preds_np = preds.cpu().numpy() + progress_bar = tqdm(total=n_samples, desc="Processing samples", unit="samples") + + for batch_idx in range(n_batches): + start_idx = batch_idx * batch_size + end_idx = min(start_idx + batch_size, n_samples) + current_batch_size = end_idx - start_idx + + # Get batch data + X_batch = torch.tensor(X[start_idx:end_idx], dtype=torch.float32).to(device) + pert_batch = pert_tensor[start_idx:end_idx].to(device) + pert_names_batch = pert_names[start_idx:end_idx].tolist() + + # Pad the batch to cell_sentence_len if it's the last incomplete batch + if current_batch_size < cell_sentence_len: + # Pad with zeros for embeddings + padding_size = cell_sentence_len - current_batch_size + X_pad = torch.zeros((padding_size, X_batch.shape[1]), device=device) + X_batch = torch.cat([X_batch, X_pad], dim=0) + + # Pad perturbation tensor with control perturbation + pert_pad = torch.zeros((padding_size, pert_batch.shape[1]), device=device) + if control_pert in pert_onehot_map: + pert_pad[:] = pert_onehot_map[control_pert].to(device) + else: + pert_pad[:, 0] = 1 # Default to first perturbation + pert_batch = torch.cat([pert_batch, pert_pad], dim=0) + + # Extend perturbation names + pert_names_batch.extend([control_pert] * padding_size) + + # Prepare batch - use same format as working code + batch = { + "ctrl_cell_emb": X_batch, + "pert_emb": pert_batch, # Keep as 2D tensor + "pert_name": pert_names_batch, + "batch": torch.zeros((1, cell_sentence_len), device=device) # Use (1, cell_sentence_len) + } + + # Run inference on batch using padded=False like in working code + batch_preds = model.predict_step(batch, batch_idx=batch_idx, padded=False) + + # Extract predictions from the dictionary returned by predict_step + # Use gene decoder output if available, otherwise use latent predictions + if "pert_cell_counts_preds" in batch_preds and batch_preds["pert_cell_counts_preds"] is not None: + # Use gene space predictions (from decoder) + pred_tensor = batch_preds["pert_cell_counts_preds"] + else: + # Use latent space predictions + pred_tensor = batch_preds["preds"] + + # Only keep predictions for the actual samples (not padding) + actual_preds = pred_tensor[:current_batch_size] + all_preds.append(actual_preds.cpu().numpy()) + + # Update progress bar + progress_bar.update(current_batch_size) + + progress_bar.close() + + # Concatenate all predictions + preds_np = np.concatenate(all_preds, axis=0) # Save predictions to AnnData pred_key = "model_preds" diff --git a/src/state_sets/_cli/_sets/_infer2.py b/src/state_sets/_cli/_sets/_infer2.py new file mode 100644 index 00000000..7522f657 --- /dev/null +++ b/src/state_sets/_cli/_sets/_infer2.py @@ -0,0 +1,150 @@ +import argparse +import scanpy as sc +import torch +import numpy as np +import os +import pandas as pd +from tqdm import tqdm + +from ...sets.models.pert_sets import PertSetsPerturbationModel + +# state-sets sets infer --output_dir /home/aadduri/state-sets/test/ --checkpoint last.ckpt --adata /home/aadduri/state-sets/test/adata.h5ad --pert_col gene +# state-sets sets infer --output_dir /home/dhruvgautam/state-sets/test/ --checkpoint /large_storage/ctc/userspace/aadduri/preprint/replogle_llama_21712320_filtered_cs32_pretrained/hepg2/checkpoints/step=44000.ckpt --adata /large_storage/ctc/ML/state_sets/replogle/processed.h5 --pert_col gene + +# state-sets sets infer --output_dir /home/dhruvgautam/state-sets/test/ --checkpoint /large_storage/ctc/userspace/aadduri/preprint/replogle_llama_21712320_filtered_cs32_pretrained/hepg2/checkpoints/step=44000.ckpt --adata /large_storage/ctc/ML/state_sets/replogle/processed.h5 --pert_col gene --embed_key X_vci_1.5.2_4 + +def add_arguments_infer(parser: argparse.ArgumentParser): + parser.add_argument("--checkpoint", type=str, required=True, help="Path to model checkpoint (.ckpt)") + parser.add_argument("--adata", type=str, required=True, help="Path to input AnnData file (.h5ad)") + parser.add_argument("--embed_key", type=str, default="X_hvg", help="Key in adata.obsm for input features") + parser.add_argument("--pert_col", type=str, default="drugname_drugconc", help="Column in adata.obs for perturbation labels") + parser.add_argument("--output_dir", type=str, default=None, help="Path to output AnnData file (.h5ad)") + parser.add_argument("--celltype_col", type=str, default=None, help="Column in adata.obs for cell type labels (optional)") + parser.add_argument("--celltypes", type=str, default=None, help="Comma-separated list of cell types to include (optional)") + parser.add_argument("--batch_size", type=int, default=1000, help="Batch size for inference (default: 1000)") + + +def run_sets_infer(args): + import logging + logging.basicConfig(level=logging.INFO) + logger = logging.getLogger(__name__) + + # Load model + logger.info(f"Loading model from checkpoint: {args.checkpoint}") + model = PertSetsPerturbationModel.load_from_checkpoint(args.checkpoint) + model.eval() + cell_sentence_len = model.cell_sentence_len + device = next(model.parameters()).device + pert_dim = model.pert_dim + + logger.info(f"Using model's cell_sentence_len: {cell_sentence_len}") + logger.info(f"Using pert_dim: {pert_dim}") + + # Load AnnData + logger.info(f"Loading AnnData from: {args.adata}") + adata_full = sc.read_h5ad(args.adata) + + # Define control perturbations to look for + control_perts = ["DMSO_TF_24h", "non-targeting", "control", "DMSO"] + + # Find available control perturbation + available_perts = set(adata_full.obs[args.pert_col].unique()) + control_pert = None + for ctrl in control_perts: + if ctrl in available_perts: + control_pert = ctrl + logger.info(f"Using '{ctrl}' as control perturbation") + break + + if control_pert is None: + # Use the first available perturbation as fallback + control_pert = list(available_perts)[0] + logger.warning(f"No standard control found, using '{control_pert}' as control") + + # Get available cell types and select the most abundant one + if args.celltype_col is not None: + if args.celltype_col not in adata_full.obs: + raise ValueError(f"Column '{args.celltype_col}' not found in adata.obs.") + + if args.celltypes is not None: + celltypes = [ct.strip() for ct in args.celltypes.split(",")] + adata_full = adata_full[adata_full.obs[args.celltype_col].isin(celltypes)].copy() + logger.info(f"Filtered to specified cell types: {celltypes}") + + cell_type_counts = adata_full.obs[args.celltype_col].value_counts() + logger.info("Available cell types: %s", list(cell_type_counts.index)) + celltype1 = cell_type_counts.index[0] + logger.info(f"Selected cell type: {celltype1} ({cell_type_counts[celltype1]} available)") + + # Get control cells for this cell type + cells_type1 = adata_full[(adata_full.obs[args.pert_col] == control_pert) & + (adata_full.obs[args.celltype_col] == celltype1)].copy() + logger.info(f"Available control cells - {celltype1}: {cells_type1.n_obs}") + else: + # No cell type filtering, use all control cells + cells_type1 = adata_full[adata_full.obs[args.pert_col] == control_pert].copy() + logger.info(f"Available control cells: {cells_type1.n_obs}") + + # Use the model's actual cell_sentence_len + n_cells = cell_sentence_len + + if cells_type1.n_obs >= n_cells: + # Sample cells + idx1 = np.random.choice(cells_type1.n_obs, size=n_cells, replace=False) + sampled_cells = cells_type1[idx1].copy() + + logger.info(f"Sampled {sampled_cells.n_obs} cells for inference") + + # Extract embeddings based on available key + if args.embed_key in sampled_cells.obsm: + X_embed = torch.tensor(sampled_cells.obsm[args.embed_key], dtype=torch.float32).to(device) + logger.info(f"Using adata.obsm['{args.embed_key}'] as input features: shape {X_embed.shape}") + else: + X_data = sampled_cells.X.toarray() if hasattr(sampled_cells.X, 'toarray') else sampled_cells.X + X_embed = torch.tensor(X_data, dtype=torch.float32).to(device) + logger.info(f"Using adata.X as input features: shape {X_embed.shape}") + + # Create simple perturbation tensor - set first dimension to 1 for control + pert_tensor = torch.zeros((n_cells, pert_dim), device=device) + pert_tensor[:, 0] = 1 # Set first dimension to 1 for control perturbation + pert_names = [control_pert] * n_cells + + # Create batch dictionary + batch = { + "ctrl_cell_emb": X_embed, + "pert_emb": pert_tensor, + "pert_name": pert_names, + "batch": torch.zeros((1, cell_sentence_len), device=device) + } + + logger.info(f"Batch shapes - ctrl_cell_emb: {batch['ctrl_cell_emb'].shape}, pert_emb: {batch['pert_emb'].shape}") + logger.info(f"Running single forward pass with {n_cells} cells") + + # Single forward pass + with torch.no_grad(): + preds = model.forward(batch, padded=False) + + logger.info("Forward pass completed successfully") + preds_np = preds.cpu().numpy() + + # Save predictions to sampled cells + pred_key = "model_preds" + sampled_cells.obsm[pred_key] = preds_np + + else: + raise ValueError(f"Not enough control cells available. Need {n_cells}, but only have {cells_type1.n_obs}") + + # Save results + output_path = args.output_dir or args.adata.replace(".h5ad", "_with_preds.h5ad").replace(".h5", "_with_preds.h5ad") + sampled_cells.write_h5ad(output_path) + logger.info(f"Saved predictions to {output_path} (in adata.obsm['{pred_key}'])") + + +def main(): + parser = argparse.ArgumentParser(description="Run inference on AnnData with a trained model checkpoint.") + add_arguments_infer(parser) + args = parser.parse_args() + run_sets_infer(args) + +if __name__ == "__main__": + main() diff --git a/src/state_sets/sets/models/base.py b/src/state_sets/sets/models/base.py index 92f5c473..8ae32f27 100644 --- a/src/state_sets/sets/models/base.py +++ b/src/state_sets/sets/models/base.py @@ -161,6 +161,8 @@ def __init__( self.output_dim = output_dim self.pert_dim = pert_dim self.batch_dim = batch_dim + self.gene_dim = gene_dim + self.hvg_dim = hvg_dim if kwargs.get("batch_encoder", False): self.batch_dim = batch_dim @@ -237,8 +239,40 @@ def on_load_checkpoint(self, checkpoint: dict[str, tp.Any]) -> None: Re-create the decoder using the exact hyper-parameters saved in the ckpt, so that parameter shapes match and load_state_dict succeeds. """ - self.decoder_cfg = checkpoint["hyper_parameters"]["decoder_cfg"] + if "decoder_cfg" in checkpoint["hyper_parameters"]: + self.decoder_cfg = checkpoint["hyper_parameters"]["decoder_cfg"] + else: + self.decoder_cfg = None self._build_decoder() + logger.info(f"DEBUG: output_space: {self.output_space}") + if self.gene_decoder is None: + gene_dim = self.hvg_dim if self.output_space == "gene" else self.gene_dim + logger.info(f"DEBUG: gene_dim: {gene_dim}") + if (self.embed_key and self.embed_key != "X_hvg" and self.output_space == "gene") or ( + self.embed_key and self.output_space == "all" + ): # we should be able to decode from hvg to all + logger.info(f"DEBUG: Creating gene_decoder, checking conditions...") + if gene_dim > 10000: + hidden_dims = [1024, 512, 256] + elif self.embed_key in ["X_vci_1.5.2", "X_vci_1.5.2_4"]: + hidden_dims = [1024, 1024, 512] + else: + if "DMSO_TF" in self.control_pert: + if self.residual_decoder: + hidden_dims = [2058, 2058, 2058, 2058, 2058] + else: + hidden_dims = [4096, 2048, 2048] + else: + hidden_dims = [1024, 1024, 512] # make this config + + self.gene_decoder = LatentToGeneDecoder( + latent_dim=self.output_dim, + gene_dim=gene_dim, + hidden_dims=hidden_dims, + dropout=self.dropout, + residual_decoder=self.residual_decoder, + ) + logger.info(f"Initialized gene decoder for embedding {self.embed_key} to gene space") def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> torch.Tensor: diff --git a/src/state_sets/sets/models/pert_sets.py b/src/state_sets/sets/models/pert_sets.py index aa2ba52b..9f3af82f 100644 --- a/src/state_sets/sets/models/pert_sets.py +++ b/src/state_sets/sets/models/pert_sets.py @@ -325,6 +325,8 @@ def forward(self, batch: dict, padded=True) -> torch.Tensor: # apply relu if specified and we output to HVG space is_gene_space = self.hparams["embed_key"] == "X_hvg" or self.hparams["embed_key"] is None + # logger.info(f"DEBUG: is_gene_space: {is_gene_space}") + # logger.info(f"DEBUG: self.gene_decoder: {self.gene_decoder}") if is_gene_space or self.gene_decoder is None: out_pred = self.relu(out_pred) From 5e8b650feb55273ddb50fa61d16494d12f47aea6 Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Thu, 19 Jun 2025 22:08:27 -0700 Subject: [PATCH 10/21] remove old infer --- src/state_sets/_cli/_sets/_infer2.py | 150 --------------------------- 1 file changed, 150 deletions(-) delete mode 100644 src/state_sets/_cli/_sets/_infer2.py diff --git a/src/state_sets/_cli/_sets/_infer2.py b/src/state_sets/_cli/_sets/_infer2.py deleted file mode 100644 index 7522f657..00000000 --- a/src/state_sets/_cli/_sets/_infer2.py +++ /dev/null @@ -1,150 +0,0 @@ -import argparse -import scanpy as sc -import torch -import numpy as np -import os -import pandas as pd -from tqdm import tqdm - -from ...sets.models.pert_sets import PertSetsPerturbationModel - -# state-sets sets infer --output_dir /home/aadduri/state-sets/test/ --checkpoint last.ckpt --adata /home/aadduri/state-sets/test/adata.h5ad --pert_col gene -# state-sets sets infer --output_dir /home/dhruvgautam/state-sets/test/ --checkpoint /large_storage/ctc/userspace/aadduri/preprint/replogle_llama_21712320_filtered_cs32_pretrained/hepg2/checkpoints/step=44000.ckpt --adata /large_storage/ctc/ML/state_sets/replogle/processed.h5 --pert_col gene - -# state-sets sets infer --output_dir /home/dhruvgautam/state-sets/test/ --checkpoint /large_storage/ctc/userspace/aadduri/preprint/replogle_llama_21712320_filtered_cs32_pretrained/hepg2/checkpoints/step=44000.ckpt --adata /large_storage/ctc/ML/state_sets/replogle/processed.h5 --pert_col gene --embed_key X_vci_1.5.2_4 - -def add_arguments_infer(parser: argparse.ArgumentParser): - parser.add_argument("--checkpoint", type=str, required=True, help="Path to model checkpoint (.ckpt)") - parser.add_argument("--adata", type=str, required=True, help="Path to input AnnData file (.h5ad)") - parser.add_argument("--embed_key", type=str, default="X_hvg", help="Key in adata.obsm for input features") - parser.add_argument("--pert_col", type=str, default="drugname_drugconc", help="Column in adata.obs for perturbation labels") - parser.add_argument("--output_dir", type=str, default=None, help="Path to output AnnData file (.h5ad)") - parser.add_argument("--celltype_col", type=str, default=None, help="Column in adata.obs for cell type labels (optional)") - parser.add_argument("--celltypes", type=str, default=None, help="Comma-separated list of cell types to include (optional)") - parser.add_argument("--batch_size", type=int, default=1000, help="Batch size for inference (default: 1000)") - - -def run_sets_infer(args): - import logging - logging.basicConfig(level=logging.INFO) - logger = logging.getLogger(__name__) - - # Load model - logger.info(f"Loading model from checkpoint: {args.checkpoint}") - model = PertSetsPerturbationModel.load_from_checkpoint(args.checkpoint) - model.eval() - cell_sentence_len = model.cell_sentence_len - device = next(model.parameters()).device - pert_dim = model.pert_dim - - logger.info(f"Using model's cell_sentence_len: {cell_sentence_len}") - logger.info(f"Using pert_dim: {pert_dim}") - - # Load AnnData - logger.info(f"Loading AnnData from: {args.adata}") - adata_full = sc.read_h5ad(args.adata) - - # Define control perturbations to look for - control_perts = ["DMSO_TF_24h", "non-targeting", "control", "DMSO"] - - # Find available control perturbation - available_perts = set(adata_full.obs[args.pert_col].unique()) - control_pert = None - for ctrl in control_perts: - if ctrl in available_perts: - control_pert = ctrl - logger.info(f"Using '{ctrl}' as control perturbation") - break - - if control_pert is None: - # Use the first available perturbation as fallback - control_pert = list(available_perts)[0] - logger.warning(f"No standard control found, using '{control_pert}' as control") - - # Get available cell types and select the most abundant one - if args.celltype_col is not None: - if args.celltype_col not in adata_full.obs: - raise ValueError(f"Column '{args.celltype_col}' not found in adata.obs.") - - if args.celltypes is not None: - celltypes = [ct.strip() for ct in args.celltypes.split(",")] - adata_full = adata_full[adata_full.obs[args.celltype_col].isin(celltypes)].copy() - logger.info(f"Filtered to specified cell types: {celltypes}") - - cell_type_counts = adata_full.obs[args.celltype_col].value_counts() - logger.info("Available cell types: %s", list(cell_type_counts.index)) - celltype1 = cell_type_counts.index[0] - logger.info(f"Selected cell type: {celltype1} ({cell_type_counts[celltype1]} available)") - - # Get control cells for this cell type - cells_type1 = adata_full[(adata_full.obs[args.pert_col] == control_pert) & - (adata_full.obs[args.celltype_col] == celltype1)].copy() - logger.info(f"Available control cells - {celltype1}: {cells_type1.n_obs}") - else: - # No cell type filtering, use all control cells - cells_type1 = adata_full[adata_full.obs[args.pert_col] == control_pert].copy() - logger.info(f"Available control cells: {cells_type1.n_obs}") - - # Use the model's actual cell_sentence_len - n_cells = cell_sentence_len - - if cells_type1.n_obs >= n_cells: - # Sample cells - idx1 = np.random.choice(cells_type1.n_obs, size=n_cells, replace=False) - sampled_cells = cells_type1[idx1].copy() - - logger.info(f"Sampled {sampled_cells.n_obs} cells for inference") - - # Extract embeddings based on available key - if args.embed_key in sampled_cells.obsm: - X_embed = torch.tensor(sampled_cells.obsm[args.embed_key], dtype=torch.float32).to(device) - logger.info(f"Using adata.obsm['{args.embed_key}'] as input features: shape {X_embed.shape}") - else: - X_data = sampled_cells.X.toarray() if hasattr(sampled_cells.X, 'toarray') else sampled_cells.X - X_embed = torch.tensor(X_data, dtype=torch.float32).to(device) - logger.info(f"Using adata.X as input features: shape {X_embed.shape}") - - # Create simple perturbation tensor - set first dimension to 1 for control - pert_tensor = torch.zeros((n_cells, pert_dim), device=device) - pert_tensor[:, 0] = 1 # Set first dimension to 1 for control perturbation - pert_names = [control_pert] * n_cells - - # Create batch dictionary - batch = { - "ctrl_cell_emb": X_embed, - "pert_emb": pert_tensor, - "pert_name": pert_names, - "batch": torch.zeros((1, cell_sentence_len), device=device) - } - - logger.info(f"Batch shapes - ctrl_cell_emb: {batch['ctrl_cell_emb'].shape}, pert_emb: {batch['pert_emb'].shape}") - logger.info(f"Running single forward pass with {n_cells} cells") - - # Single forward pass - with torch.no_grad(): - preds = model.forward(batch, padded=False) - - logger.info("Forward pass completed successfully") - preds_np = preds.cpu().numpy() - - # Save predictions to sampled cells - pred_key = "model_preds" - sampled_cells.obsm[pred_key] = preds_np - - else: - raise ValueError(f"Not enough control cells available. Need {n_cells}, but only have {cells_type1.n_obs}") - - # Save results - output_path = args.output_dir or args.adata.replace(".h5ad", "_with_preds.h5ad").replace(".h5", "_with_preds.h5ad") - sampled_cells.write_h5ad(output_path) - logger.info(f"Saved predictions to {output_path} (in adata.obsm['{pred_key}'])") - - -def main(): - parser = argparse.ArgumentParser(description="Run inference on AnnData with a trained model checkpoint.") - add_arguments_infer(parser) - args = parser.parse_args() - run_sets_infer(args) - -if __name__ == "__main__": - main() From 8eead7678640a6727e0374dbd7d49138df6045aa Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Fri, 20 Jun 2025 00:41:37 -0700 Subject: [PATCH 11/21] train working i think --- src/state_sets/_cli/_sets/_train.py | 7 +++- src/state_sets/sets/models/base.py | 63 +++++++++++++++-------------- 2 files changed, 38 insertions(+), 32 deletions(-) diff --git a/src/state_sets/_cli/_sets/_train.py b/src/state_sets/_cli/_sets/_train.py index 586750cd..093c5827 100644 --- a/src/state_sets/_cli/_sets/_train.py +++ b/src/state_sets/_cli/_sets/_train.py @@ -117,8 +117,11 @@ def run_sets_train(cfg: DictConfig): data_module.setup(stage="fit") var_dims = data_module.get_var_dims() # {"gene_dim": …, "hvg_dim": …} - gene_dim = var_dims.get("gene_dim", 5000) # fallback if key missing - latent_dim = cfg["model"]["kwargs"]["output_dim"] # same as model.output_dim + if cfg["data"]["kwargs"]["output_space"] == "gene": + gene_dim = var_dims.get("hvg_dim", 2000) # fallback if key missing + else: + gene_dim = var_dims.get("gene_dim", 2000) # fallback if key missing + latent_dim = var_dims["output_dim"] # same as model.output_dim hidden_dims = cfg["model"]["kwargs"].get("decoder_hidden_dims", [1024, 1024, 512]) decoder_cfg = dict( diff --git a/src/state_sets/sets/models/base.py b/src/state_sets/sets/models/base.py index 8ae32f27..13dee2c5 100644 --- a/src/state_sets/sets/models/base.py +++ b/src/state_sets/sets/models/base.py @@ -241,38 +241,41 @@ def on_load_checkpoint(self, checkpoint: dict[str, tp.Any]) -> None: """ if "decoder_cfg" in checkpoint["hyper_parameters"]: self.decoder_cfg = checkpoint["hyper_parameters"]["decoder_cfg"] + self._build_decoder() + logger.info(f"Loaded decoder from checkpoint decoder_cfg: {self.decoder_cfg}") else: - self.decoder_cfg = None - self._build_decoder() - logger.info(f"DEBUG: output_space: {self.output_space}") - if self.gene_decoder is None: - gene_dim = self.hvg_dim if self.output_space == "gene" else self.gene_dim - logger.info(f"DEBUG: gene_dim: {gene_dim}") - if (self.embed_key and self.embed_key != "X_hvg" and self.output_space == "gene") or ( - self.embed_key and self.output_space == "all" - ): # we should be able to decode from hvg to all - logger.info(f"DEBUG: Creating gene_decoder, checking conditions...") - if gene_dim > 10000: - hidden_dims = [1024, 512, 256] - elif self.embed_key in ["X_vci_1.5.2", "X_vci_1.5.2_4"]: - hidden_dims = [1024, 1024, 512] - else: - if "DMSO_TF" in self.control_pert: - if self.residual_decoder: - hidden_dims = [2058, 2058, 2058, 2058, 2058] - else: - hidden_dims = [4096, 2048, 2048] + # Only fall back to old logic if no decoder_cfg was saved + self.decoder_cfg = None + self._build_decoder() + logger.info(f"DEBUG: output_space: {self.output_space}") + if self.gene_decoder is None: + gene_dim = self.hvg_dim if self.output_space == "gene" else self.gene_dim + logger.info(f"DEBUG: gene_dim: {gene_dim}") + if (self.embed_key and self.embed_key != "X_hvg" and self.output_space == "gene") or ( + self.embed_key and self.output_space == "all" + ): # we should be able to decode from hvg to all + logger.info(f"DEBUG: Creating gene_decoder, checking conditions...") + if gene_dim > 10000: + hidden_dims = [1024, 512, 256] + elif self.embed_key in ["X_vci_1.5.2", "X_vci_1.5.2_4"]: + hidden_dims = [1024, 1024, 512] else: - hidden_dims = [1024, 1024, 512] # make this config - - self.gene_decoder = LatentToGeneDecoder( - latent_dim=self.output_dim, - gene_dim=gene_dim, - hidden_dims=hidden_dims, - dropout=self.dropout, - residual_decoder=self.residual_decoder, - ) - logger.info(f"Initialized gene decoder for embedding {self.embed_key} to gene space") + if "DMSO_TF" in self.control_pert: + if self.residual_decoder: + hidden_dims = [2058, 2058, 2058, 2058, 2058] + else: + hidden_dims = [4096, 2048, 2048] + else: + hidden_dims = [1024, 1024, 512] # make this config + + self.gene_decoder = LatentToGeneDecoder( + latent_dim=self.output_dim, + gene_dim=gene_dim, + hidden_dims=hidden_dims, + dropout=self.dropout, + residual_decoder=self.residual_decoder, + ) + logger.info(f"Initialized gene decoder for embedding {self.embed_key} to gene space") def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> torch.Tensor: From 41640f6a71f832f25bca8858ab1cffde8f73b556 Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Fri, 20 Jun 2025 01:19:10 -0700 Subject: [PATCH 12/21] merge --- src/state_sets/sets/models/base.py | 36 ------------------------------ 1 file changed, 36 deletions(-) diff --git a/src/state_sets/sets/models/base.py b/src/state_sets/sets/models/base.py index 13dee2c5..8da26fb8 100644 --- a/src/state_sets/sets/models/base.py +++ b/src/state_sets/sets/models/base.py @@ -182,42 +182,6 @@ def __init__( self.lr = lr self.loss_fn = get_loss_fn(loss_fn) - # this will either decode to hvg space if output space is a gene, - # or to transcriptome space if output space is all. done this way to maintain - # backwards compatibility with the old models - self.gene_decoder = None - gene_dim = hvg_dim if output_space == "gene" else gene_dim - if (embed_key and embed_key != "X_hvg" and output_space == "gene") or ( - embed_key and output_space == "all" - ): # we should be able to decode from hvg to all - if gene_dim > 10000: - hidden_dims = [1024, 512, 256] - else: - if "DMSO_TF" in self.control_pert: - if self.residual_decoder: - hidden_dims = [2058, 2058, 2058, 2058, 2058] - else: - hidden_dims = [4096, 2048, 2048] - elif "PBS" in self.control_pert: - hidden_dims = [2048, 1024, 1024] - else: - if "DMSO_TF" in self.control_pert: - if self.residual_decoder: - hidden_dims = [2058, 2058, 2058, 2058, 2058] - else: - hidden_dims = [4096, 2048, 2048] - else: - hidden_dims = [1024, 1024, 512] # make this config - - self.gene_decoder = LatentToGeneDecoder( - latent_dim=self.output_dim, - gene_dim=gene_dim, - hidden_dims=hidden_dims, - dropout=dropout, - residual_decoder=self.residual_decoder, - ) - logger.info(f"Initialized gene decoder for embedding {embed_key} to gene space") - def transfer_batch_to_device(self, batch, device, dataloader_idx: int): return {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()} From 0ff64461de373594141de5f8cd73f45a018f57f8 Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Fri, 20 Jun 2025 11:36:50 -0700 Subject: [PATCH 13/21] restart fix --- src/state_sets/configs/model/tahoe_decoder_test.yaml | 8 ++++---- src/state_sets/sets/models/base.py | 3 ++- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/state_sets/configs/model/tahoe_decoder_test.yaml b/src/state_sets/configs/model/tahoe_decoder_test.yaml index 1a0c985d..680e1812 100644 --- a/src/state_sets/configs/model/tahoe_decoder_test.yaml +++ b/src/state_sets/configs/model/tahoe_decoder_test.yaml @@ -6,7 +6,7 @@ kwargs: cell_set_len: 512 decoder_hidden_dims: [2048, 2048, 2048] blur: 0.05 - hidden_dim: 1488 # hidden dimension going into the transformer backbone + hidden_dim: 696 # hidden dimension going into the transformer backbone loss: energy confidence_head: False n_encoder_layers: 4 @@ -27,11 +27,11 @@ kwargs: transformer_backbone_kwargs: max_position_embeddings: ${model.kwargs.cell_set_len} hidden_size: ${model.kwargs.hidden_dim} - intermediate_size: 5952 - num_hidden_layers: 6 + intermediate_size: 2784 + num_hidden_layers: 8 num_attention_heads: 12 num_key_value_heads: 12 - head_dim: 124 + head_dim: 58 use_cache: false attention_dropout: 0.0 hidden_dropout: 0.0 diff --git a/src/state_sets/sets/models/base.py b/src/state_sets/sets/models/base.py index 8da26fb8..8b522b47 100644 --- a/src/state_sets/sets/models/base.py +++ b/src/state_sets/sets/models/base.py @@ -181,6 +181,7 @@ def __init__( self.dropout = dropout self.lr = lr self.loss_fn = get_loss_fn(loss_fn) + self._build_decoder() def transfer_batch_to_device(self, batch, device, dataloader_idx: int): return {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()} @@ -205,7 +206,7 @@ def on_load_checkpoint(self, checkpoint: dict[str, tp.Any]) -> None: """ if "decoder_cfg" in checkpoint["hyper_parameters"]: self.decoder_cfg = checkpoint["hyper_parameters"]["decoder_cfg"] - self._build_decoder() + self.gene_decoder = LatentToGeneDecoder(**self.decoder_cfg) logger.info(f"Loaded decoder from checkpoint decoder_cfg: {self.decoder_cfg}") else: # Only fall back to old logic if no decoder_cfg was saved From baea91e9e88dd6491a5cb4b963591d1923ff1b40 Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Fri, 20 Jun 2025 14:38:03 -0700 Subject: [PATCH 14/21] train fixed --- src/state_sets/_cli/_sets/_train.py | 35 ++++++++++++++++++++++++++--- 1 file changed, 32 insertions(+), 3 deletions(-) diff --git a/src/state_sets/_cli/_sets/_train.py b/src/state_sets/_cli/_sets/_train.py index 093c5827..ed41f801 100644 --- a/src/state_sets/_cli/_sets/_train.py +++ b/src/state_sets/_cli/_sets/_train.py @@ -109,12 +109,29 @@ def run_sets_train(cfg: DictConfig): batch_size=cfg["training"]["batch_size"], cell_sentence_len=sentence_len, ) + with open(join(run_output_dir, "data_module.torch"), "wb") as f: # TODO-Abhi: only save necessary data data_module.save_state(f) data_module.setup(stage="fit") + dl = data_module.train_dataloader() + print("num_workers:", dl.num_workers) + print("batch size:", dl.batch_size) + # # Test loading a single batch to identify potential data loading issues + # try: + # print("DEBUG: Testing data loading with one batch...") + # train_loader = data_module.train_dataloader() + # first_batch = next(iter(train_loader)) + # print(f"DEBUG: Successfully loaded first batch. Batch keys: {list(first_batch.keys()) if isinstance(first_batch, dict) else 'Not a dict'}") + # if isinstance(first_batch, dict): + # for key, value in first_batch.items(): + # if hasattr(value, 'shape'): + # print(f"DEBUG: {key} shape: {value.shape}") + # except Exception as e: + # print(f"DEBUG: Error loading first batch: {e}") + # print("DEBUG: This might be the source of the hanging issue") var_dims = data_module.get_var_dims() # {"gene_dim": …, "hvg_dim": …} if cfg["data"]["kwargs"]["output_space"] == "gene": @@ -135,9 +152,6 @@ def run_sets_train(cfg: DictConfig): # tuck it into the kwargs that will reach the LightningModule cfg["model"]["kwargs"]["decoder_cfg"] = decoder_cfg - cfg["data"]["kwargs"]["n_perts"] = len(data_module.pert_onehot_map) - cfg["model"]["kwargs"]["pert_onehot_map"] = data_module.pert_onehot_map - if cfg["model"]["name"].lower() in ["cpa", "scvi"] or cfg["model"]["name"].lower().startswith("scgpt"): cfg["model"]["kwargs"]["n_cell_types"] = len(data_module.celltype_onehot_map) cfg["model"]["kwargs"]["n_perts"] = len(data_module.pert_onehot_map) @@ -152,6 +166,8 @@ def run_sets_train(cfg: DictConfig): data_module.get_var_dims(), ) + print(f"DEBUG: Model created. Estimated params size: {sum(p.numel() * p.element_size() for p in model.parameters()) / 1024**3:.2f} GB") + # print(f"DEBUG: Num workers: {data_module.train_dataloader().num_workers}") # Set up logging loggers = get_loggers( output_dir=cfg["output_dir"], @@ -220,7 +236,9 @@ def run_sets_train(cfg: DictConfig): del trainer_kwargs["max_steps"] # Build trainer + print(f"DEBUG: Building trainer with kwargs: {trainer_kwargs}") trainer = pl.Trainer(**trainer_kwargs) + print("DEBUG: Trainer built successfully") # Load checkpoint if exists checkpoint_path = join(ckpt_callbacks[0].dirpath, "last.ckpt") @@ -229,12 +247,17 @@ def run_sets_train(cfg: DictConfig): else: logging.info(f"!! Resuming training from {checkpoint_path} !!") + print(f"DEBUG: Model device: {next(model.parameters()).device}") + print(f"DEBUG: CUDA memory allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB") + print(f"DEBUG: CUDA memory reserved: {torch.cuda.memory_reserved() / 1024**3:.2f} GB") + logger.info("Starting trainer fit.") # if a checkpoint does not exist, start with the provided checkpoint # this is mainly used for pretrain -> finetune workflows manual_init = cfg["model"]["kwargs"].get("init_from", None) if checkpoint_path is None and manual_init is not None: + print(f"DEBUG: Loading manual checkpoint from {manual_init}") checkpoint_path = manual_init device = torch.device("cuda" if torch.cuda.is_available() else "cpu") checkpoint = torch.load(checkpoint_path, map_location=device) @@ -275,6 +298,7 @@ def run_sets_train(cfg: DictConfig): # Load the filtered state dict model.load_state_dict(filtered_state, strict=False) + print("DEBUG: About to call trainer.fit() with manual checkpoint...") # Train - for clarity we pass None trainer.fit( @@ -282,13 +306,18 @@ def run_sets_train(cfg: DictConfig): datamodule=data_module, ckpt_path=None, ) + print("DEBUG: trainer.fit() completed with manual checkpoint") else: + print(f"DEBUG: About to call trainer.fit() with checkpoint_path={checkpoint_path}") # Train trainer.fit( model, datamodule=data_module, ckpt_path=checkpoint_path, ) + print("DEBUG: trainer.fit() completed") + + print("DEBUG: Training completed, saving final checkpoint...") # at this point if checkpoint_path does not exist, manually create one checkpoint_path = join(ckpt_callbacks[0].dirpath, "final.ckpt") From bcf578f7ed33a7a9dd2af1a9c364b5473d7f17cd Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Fri, 20 Jun 2025 14:55:13 -0700 Subject: [PATCH 15/21] readme --- README.md | 5 +++++ src/state_sets/_cli/_sets/_infer.py | 5 ----- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index e4d9d77f..2ee25d6f 100644 --- a/README.md +++ b/README.md @@ -60,6 +60,11 @@ An example evaluation command for a sets model: state-sets sets predict --output_dir /home/aadduri/state-sets/test/ --checkpoint last.ckpt ``` +An example inference command for a sets model: +```bash +state-sets sets infer --output /home/dhruvgautam/state-sets/test/ --output_dir /path/to/model/ --checkpoint /path/to/model/checkpoints/last.ckpt --adata /path/to/anndata/processed.h5 --pert_col gene --embed_key X_hvg +``` + The toml files should be setup to define perturbation splits, if running fewshot experiments. Here are some examples: ```toml diff --git a/src/state_sets/_cli/_sets/_infer.py b/src/state_sets/_cli/_sets/_infer.py index 60406c84..5ec61f61 100644 --- a/src/state_sets/_cli/_sets/_infer.py +++ b/src/state_sets/_cli/_sets/_infer.py @@ -10,11 +10,6 @@ from ...sets.models.pert_sets import PertSetsPerturbationModel from cell_load.data_modules import PerturbationDataModule -# state-sets sets infer --output_dir /home/aadduri/state-sets/test/ --checkpoint last.ckpt --adata /home/aadduri/state-sets/test/adata.h5ad --pert_col gene -# state-sets sets infer --output_dir /home/dhruvgautam/state-sets/test/ --checkpoint /large_storage/ctc/userspace/aadduri/preprint/replogle_llama_21712320_filtered_cs32_pretrained/hepg2/checkpoints/step=44000.ckpt --adata /large_storage/ctc/ML/state_sets/replogle/processed.h5 --pert_col gene - -# state-sets sets infer --output /home/dhruvgautam/state-sets/test/ --output_dir /large_storage/ctc/userspace/aadduri/preprint/replogle_state_proper_cs32_sm/hepg2 --checkpoint /large_storage/ctc/userspace/aadduri/preprint/replogle_state_proper_cs32_sm/hepg2/checkpoints/step=48000.ckpt --adata /large_storage/ctc/ML/state_sets/replogle/processed.h5 --pert_col gene --embed_key X_vci_1.5.2_4 - def add_arguments_infer(parser: argparse.ArgumentParser): parser.add_argument("--checkpoint", type=str, required=True, help="Path to model checkpoint (.ckpt)") parser.add_argument("--adata", type=str, required=True, help="Path to input AnnData file (.h5ad)") From 9cd14d075a29dfb7d3f7494bf42bf2dcaf85ac5d Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Fri, 20 Jun 2025 14:55:29 -0700 Subject: [PATCH 16/21] space --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 2ee25d6f..4377ca2b 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,7 @@ state-sets sets predict --output_dir /home/aadduri/state-sets/test/ --checkpoint ``` An example inference command for a sets model: + ```bash state-sets sets infer --output /home/dhruvgautam/state-sets/test/ --output_dir /path/to/model/ --checkpoint /path/to/model/checkpoints/last.ckpt --adata /path/to/anndata/processed.h5 --pert_col gene --embed_key X_hvg ``` From 7875b5152d5cfb7ef2bf21274059297f5a39e220 Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Sat, 21 Jun 2025 12:06:29 -0700 Subject: [PATCH 17/21] changes --- src/state_sets/_cli/_sets/_infer.py | 73 +++++++++++++++++------------ src/state_sets/_cli/_sets/_train.py | 62 ++++++++++-------------- src/state_sets/sets/models/base.py | 3 +- 3 files changed, 69 insertions(+), 69 deletions(-) diff --git a/src/state_sets/_cli/_sets/_infer.py b/src/state_sets/_cli/_sets/_infer.py index 5ec61f61..b79faf80 100644 --- a/src/state_sets/_cli/_sets/_infer.py +++ b/src/state_sets/_cli/_sets/_infer.py @@ -10,20 +10,33 @@ from ...sets.models.pert_sets import PertSetsPerturbationModel from cell_load.data_modules import PerturbationDataModule + def add_arguments_infer(parser: argparse.ArgumentParser): parser.add_argument("--checkpoint", type=str, required=True, help="Path to model checkpoint (.ckpt)") parser.add_argument("--adata", type=str, required=True, help="Path to input AnnData file (.h5ad)") parser.add_argument("--embed_key", type=str, default="X_hvg", help="Key in adata.obsm for input features") - parser.add_argument("--pert_col", type=str, default="drugname_drugconc", help="Column in adata.obs for perturbation labels") + parser.add_argument( + "--pert_col", type=str, default="drugname_drugconc", help="Column in adata.obs for perturbation labels" + ) parser.add_argument("--output", type=str, default=None, help="Path to output AnnData file (.h5ad)") - parser.add_argument("--output_dir", type=str, required=True, help="Path to the output_dir containing the config.yaml file that was saved during training.") - parser.add_argument("--celltype_col", type=str, default=None, help="Column in adata.obs for cell type labels (optional)") - parser.add_argument("--celltypes", type=str, default=None, help="Comma-separated list of cell types to include (optional)") + parser.add_argument( + "--output_dir", + type=str, + required=True, + help="Path to the output_dir containing the config.yaml file that was saved during training.", + ) + parser.add_argument( + "--celltype_col", type=str, default=None, help="Column in adata.obs for cell type labels (optional)" + ) + parser.add_argument( + "--celltypes", type=str, default=None, help="Comma-separated list of cell types to include (optional)" + ) parser.add_argument("--batch_size", type=int, default=1000, help="Batch size for inference (default: 1000)") def run_sets_infer(args): import logging + logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) @@ -87,31 +100,30 @@ def load_config(cfg_path: str) -> dict: # Prepare perturbation tensor using the data module's mapping pert_names = adata.obs[args.pert_col].values - pert_tensor = torch.zeros((len(pert_names), pert_dim), device='cpu') # Keep on CPU initially + pert_tensor = torch.zeros((len(pert_names), pert_dim), device="cpu") # Keep on CPU initially logger.info(f"Perturbation tensor shape: {pert_tensor.shape}") - + # Use data module's perturbation mapping pert_onehot_map = data_module.pert_onehot_map - - # Debug: check what's available + logger.info(f"Data module has {len(pert_onehot_map)} perturbations in mapping") logger.info(f"First 10 perturbations in data module: {list(pert_onehot_map.keys())[:10]}") - + unique_pert_names = sorted(set(pert_names)) logger.info(f"AnnData has {len(unique_pert_names)} unique perturbations") logger.info(f"First 10 perturbations in AnnData: {unique_pert_names[:10]}") - + # Check overlap overlap = set(unique_pert_names) & set(pert_onehot_map.keys()) logger.info(f"Overlap between AnnData and data module: {len(overlap)} perturbations") if len(overlap) < len(unique_pert_names): missing = set(unique_pert_names) - set(pert_onehot_map.keys()) logger.warning(f"Missing perturbations: {list(missing)[:10]}") - + # Check if there's a control perturbation that might match control_pert = data_module.get_control_pert() logger.info(f"Control perturbation in data module: '{control_pert}'") - + matched_count = 0 for idx, name in enumerate(pert_names): if name in pert_onehot_map: @@ -125,7 +137,7 @@ def load_config(cfg_path: str) -> dict: # Use first available perturbation as fallback first_pert = list(pert_onehot_map.keys())[0] pert_tensor[idx] = pert_onehot_map[first_pert] - + logger.info(f"Matched {matched_count} out of {len(pert_names)} perturbations") # Process in batches with progress bar @@ -133,31 +145,33 @@ def load_config(cfg_path: str) -> dict: n_samples = len(pert_names) batch_size = cell_sentence_len # Model requires this exact batch size n_batches = (n_samples + batch_size - 1) // batch_size # Ceiling division - - logger.info(f"Running inference on {n_samples} samples in {n_batches} batches of size {batch_size} (model's cell_sentence_len)...") - + + logger.info( + f"Running inference on {n_samples} samples in {n_batches} batches of size {batch_size} (model's cell_sentence_len)..." + ) + all_preds = [] - + with torch.no_grad(): progress_bar = tqdm(total=n_samples, desc="Processing samples", unit="samples") - + for batch_idx in range(n_batches): start_idx = batch_idx * batch_size end_idx = min(start_idx + batch_size, n_samples) current_batch_size = end_idx - start_idx - + # Get batch data X_batch = torch.tensor(X[start_idx:end_idx], dtype=torch.float32).to(device) pert_batch = pert_tensor[start_idx:end_idx].to(device) pert_names_batch = pert_names[start_idx:end_idx].tolist() - + # Pad the batch to cell_sentence_len if it's the last incomplete batch if current_batch_size < cell_sentence_len: # Pad with zeros for embeddings padding_size = cell_sentence_len - current_batch_size X_pad = torch.zeros((padding_size, X_batch.shape[1]), device=device) X_batch = torch.cat([X_batch, X_pad], dim=0) - + # Pad perturbation tensor with control perturbation pert_pad = torch.zeros((padding_size, pert_batch.shape[1]), device=device) if control_pert in pert_onehot_map: @@ -165,21 +179,21 @@ def load_config(cfg_path: str) -> dict: else: pert_pad[:, 0] = 1 # Default to first perturbation pert_batch = torch.cat([pert_batch, pert_pad], dim=0) - + # Extend perturbation names pert_names_batch.extend([control_pert] * padding_size) - + # Prepare batch - use same format as working code batch = { "ctrl_cell_emb": X_batch, "pert_emb": pert_batch, # Keep as 2D tensor "pert_name": pert_names_batch, - "batch": torch.zeros((1, cell_sentence_len), device=device) # Use (1, cell_sentence_len) + "batch": torch.zeros((1, cell_sentence_len), device=device), # Use (1, cell_sentence_len) } - + # Run inference on batch using padded=False like in working code batch_preds = model.predict_step(batch, batch_idx=batch_idx, padded=False) - + # Extract predictions from the dictionary returned by predict_step # Use gene decoder output if available, otherwise use latent predictions if "pert_cell_counts_preds" in batch_preds and batch_preds["pert_cell_counts_preds"] is not None: @@ -188,14 +202,14 @@ def load_config(cfg_path: str) -> dict: else: # Use latent space predictions pred_tensor = batch_preds["preds"] - + # Only keep predictions for the actual samples (not padding) actual_preds = pred_tensor[:current_batch_size] all_preds.append(actual_preds.cpu().numpy()) - + # Update progress bar progress_bar.update(current_batch_size) - + progress_bar.close() # Concatenate all predictions @@ -215,5 +229,6 @@ def main(): args = parser.parse_args() run_sets_infer(args) + if __name__ == "__main__": main() diff --git a/src/state_sets/_cli/_sets/_train.py b/src/state_sets/_cli/_sets/_train.py index ed41f801..1a7cf50e 100644 --- a/src/state_sets/_cli/_sets/_train.py +++ b/src/state_sets/_cli/_sets/_train.py @@ -109,7 +109,6 @@ def run_sets_train(cfg: DictConfig): batch_size=cfg["training"]["batch_size"], cell_sentence_len=sentence_len, ) - with open(join(run_output_dir, "data_module.torch"), "wb") as f: # TODO-Abhi: only save necessary data @@ -119,34 +118,21 @@ def run_sets_train(cfg: DictConfig): dl = data_module.train_dataloader() print("num_workers:", dl.num_workers) print("batch size:", dl.batch_size) - # # Test loading a single batch to identify potential data loading issues - # try: - # print("DEBUG: Testing data loading with one batch...") - # train_loader = data_module.train_dataloader() - # first_batch = next(iter(train_loader)) - # print(f"DEBUG: Successfully loaded first batch. Batch keys: {list(first_batch.keys()) if isinstance(first_batch, dict) else 'Not a dict'}") - # if isinstance(first_batch, dict): - # for key, value in first_batch.items(): - # if hasattr(value, 'shape'): - # print(f"DEBUG: {key} shape: {value.shape}") - # except Exception as e: - # print(f"DEBUG: Error loading first batch: {e}") - # print("DEBUG: This might be the source of the hanging issue") - - var_dims = data_module.get_var_dims() # {"gene_dim": …, "hvg_dim": …} + + var_dims = data_module.get_var_dims() # {"gene_dim": …, "hvg_dim": …} if cfg["data"]["kwargs"]["output_space"] == "gene": - gene_dim = var_dims.get("hvg_dim", 2000) # fallback if key missing + gene_dim = var_dims.get("hvg_dim", 2000) # fallback if key missing else: - gene_dim = var_dims.get("gene_dim", 2000) # fallback if key missing - latent_dim = var_dims["output_dim"] # same as model.output_dim + gene_dim = var_dims.get("gene_dim", 2000) # fallback if key missing + latent_dim = var_dims["output_dim"] # same as model.output_dim hidden_dims = cfg["model"]["kwargs"].get("decoder_hidden_dims", [1024, 1024, 512]) decoder_cfg = dict( - latent_dim = latent_dim, - gene_dim = gene_dim, - hidden_dims = hidden_dims, - dropout = cfg["model"]["kwargs"].get("decoder_dropout", 0.1), - residual_decoder = cfg["model"]["kwargs"].get("residual_decoder", False), + latent_dim=latent_dim, + gene_dim=gene_dim, + hidden_dims=hidden_dims, + dropout=cfg["model"]["kwargs"].get("decoder_dropout", 0.1), + residual_decoder=cfg["model"]["kwargs"].get("residual_decoder", False), ) # tuck it into the kwargs that will reach the LightningModule @@ -166,9 +152,9 @@ def run_sets_train(cfg: DictConfig): data_module.get_var_dims(), ) - print(f"DEBUG: Model created. Estimated params size: {sum(p.numel() * p.element_size() for p in model.parameters()) / 1024**3:.2f} GB") - # print(f"DEBUG: Num workers: {data_module.train_dataloader().num_workers}") - # Set up logging + print( + f"Model created. Estimated params size: {sum(p.numel() * p.element_size() for p in model.parameters()) / 1024**3:.2f} GB" + ) loggers = get_loggers( output_dir=cfg["output_dir"], name=cfg["name"], @@ -236,9 +222,9 @@ def run_sets_train(cfg: DictConfig): del trainer_kwargs["max_steps"] # Build trainer - print(f"DEBUG: Building trainer with kwargs: {trainer_kwargs}") + print(f"Building trainer with kwargs: {trainer_kwargs}") trainer = pl.Trainer(**trainer_kwargs) - print("DEBUG: Trainer built successfully") + print("Trainer built successfully") # Load checkpoint if exists checkpoint_path = join(ckpt_callbacks[0].dirpath, "last.ckpt") @@ -247,9 +233,9 @@ def run_sets_train(cfg: DictConfig): else: logging.info(f"!! Resuming training from {checkpoint_path} !!") - print(f"DEBUG: Model device: {next(model.parameters()).device}") - print(f"DEBUG: CUDA memory allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB") - print(f"DEBUG: CUDA memory reserved: {torch.cuda.memory_reserved() / 1024**3:.2f} GB") + print(f"Model device: {next(model.parameters()).device}") + print(f"CUDA memory allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB") + print(f"CUDA memory reserved: {torch.cuda.memory_reserved() / 1024**3:.2f} GB") logger.info("Starting trainer fit.") @@ -257,7 +243,7 @@ def run_sets_train(cfg: DictConfig): # this is mainly used for pretrain -> finetune workflows manual_init = cfg["model"]["kwargs"].get("init_from", None) if checkpoint_path is None and manual_init is not None: - print(f"DEBUG: Loading manual checkpoint from {manual_init}") + print(f"Loading manual checkpoint from {manual_init}") checkpoint_path = manual_init device = torch.device("cuda" if torch.cuda.is_available() else "cpu") checkpoint = torch.load(checkpoint_path, map_location=device) @@ -298,7 +284,7 @@ def run_sets_train(cfg: DictConfig): # Load the filtered state dict model.load_state_dict(filtered_state, strict=False) - print("DEBUG: About to call trainer.fit() with manual checkpoint...") + print("About to call trainer.fit() with manual checkpoint...") # Train - for clarity we pass None trainer.fit( @@ -306,18 +292,18 @@ def run_sets_train(cfg: DictConfig): datamodule=data_module, ckpt_path=None, ) - print("DEBUG: trainer.fit() completed with manual checkpoint") + print("trainer.fit() completed with manual checkpoint") else: - print(f"DEBUG: About to call trainer.fit() with checkpoint_path={checkpoint_path}") + print(f"About to call trainer.fit() with checkpoint_path={checkpoint_path}") # Train trainer.fit( model, datamodule=data_module, ckpt_path=checkpoint_path, ) - print("DEBUG: trainer.fit() completed") + print("trainer.fit() completed") - print("DEBUG: Training completed, saving final checkpoint...") + print("Training completed, saving final checkpoint...") # at this point if checkpoint_path does not exist, manually create one checkpoint_path = join(ckpt_callbacks[0].dirpath, "final.ckpt") diff --git a/src/state_sets/sets/models/base.py b/src/state_sets/sets/models/base.py index 8b522b47..f95e2c1d 100644 --- a/src/state_sets/sets/models/base.py +++ b/src/state_sets/sets/models/base.py @@ -210,7 +210,7 @@ def on_load_checkpoint(self, checkpoint: dict[str, tp.Any]) -> None: logger.info(f"Loaded decoder from checkpoint decoder_cfg: {self.decoder_cfg}") else: # Only fall back to old logic if no decoder_cfg was saved - self.decoder_cfg = None + self.decoder_cfg = None self._build_decoder() logger.info(f"DEBUG: output_space: {self.output_space}") if self.gene_decoder is None: @@ -242,7 +242,6 @@ def on_load_checkpoint(self, checkpoint: dict[str, tp.Any]) -> None: ) logger.info(f"Initialized gene decoder for embedding {self.embed_key} to gene space") - def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> torch.Tensor: """Training step logic for both main model and decoder.""" # Get model predictions (in latent space) From e005ce34e67344a2283db1404151389ee2bebe00 Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Sat, 21 Jun 2025 17:17:46 -0700 Subject: [PATCH 18/21] fixed train and infer --- src/state_sets/_cli/_sets/_infer.py | 22 ++++++++-------------- src/state_sets/_cli/_sets/_train.py | 15 +++++++++++++++ src/state_sets/configs/config.yaml | 2 +- src/state_sets/configs/model/pertsets.yaml | 1 + 4 files changed, 25 insertions(+), 15 deletions(-) diff --git a/src/state_sets/_cli/_sets/_infer.py b/src/state_sets/_cli/_sets/_infer.py index b79faf80..d5a85147 100644 --- a/src/state_sets/_cli/_sets/_infer.py +++ b/src/state_sets/_cli/_sets/_infer.py @@ -6,9 +6,9 @@ import pandas as pd from tqdm import tqdm import yaml +import pickle from ...sets.models.pert_sets import PertSetsPerturbationModel -from cell_load.data_modules import PerturbationDataModule def add_arguments_infer(parser: argparse.ArgumentParser): @@ -53,17 +53,10 @@ def load_config(cfg_path: str) -> dict: cfg = load_config(config_path) logger.info(f"Loaded config from {config_path}") - # Find run output directory & load data module - run_output_dir = os.path.join(cfg["output_dir"], cfg["name"]) - data_module_path = os.path.join(run_output_dir, "data_module.torch") - if not os.path.exists(data_module_path): - raise FileNotFoundError(f"Could not find data module at {data_module_path}") - data_module = PerturbationDataModule.load_state(data_module_path) - data_module.setup(stage="test") - logger.info(f"Loaded data module from {data_module_path}") - # Get perturbation dimensions and mapping from data module - var_dims = data_module.get_var_dims() + var_dims_path = os.path.join(args.output_dir, "var_dims.pkl") + with open(var_dims_path, "rb") as f: + var_dims = pickle.load(f) pert_dim = var_dims["pert_dim"] # Load model @@ -103,8 +96,9 @@ def load_config(cfg_path: str) -> dict: pert_tensor = torch.zeros((len(pert_names), pert_dim), device="cpu") # Keep on CPU initially logger.info(f"Perturbation tensor shape: {pert_tensor.shape}") - # Use data module's perturbation mapping - pert_onehot_map = data_module.pert_onehot_map + # Load perturbation mapping from torch file + pert_onehot_map_path = os.path.join(args.output_dir, "pert_onehot_map.pt") + pert_onehot_map = torch.load(pert_onehot_map_path, weights_only=False) logger.info(f"Data module has {len(pert_onehot_map)} perturbations in mapping") logger.info(f"First 10 perturbations in data module: {list(pert_onehot_map.keys())[:10]}") @@ -121,7 +115,7 @@ def load_config(cfg_path: str) -> dict: logger.warning(f"Missing perturbations: {list(missing)[:10]}") # Check if there's a control perturbation that might match - control_pert = data_module.get_control_pert() + control_pert = cfg["data"]["kwargs"]["control_pert"] logger.info(f"Control perturbation in data module: '{control_pert}'") matched_count = 0 diff --git a/src/state_sets/_cli/_sets/_train.py b/src/state_sets/_cli/_sets/_train.py index 1a7cf50e..5447dcf7 100644 --- a/src/state_sets/_cli/_sets/_train.py +++ b/src/state_sets/_cli/_sets/_train.py @@ -11,6 +11,7 @@ def add_arguments_train(parser: ap.ArgumentParser): def run_sets_train(cfg: DictConfig): import json import os + import pickle from pathlib import Path import shutil from os.path import join, exists @@ -137,6 +138,20 @@ def run_sets_train(cfg: DictConfig): # tuck it into the kwargs that will reach the LightningModule cfg["model"]["kwargs"]["decoder_cfg"] = decoder_cfg + + # Save the onehot maps as pickle files instead of storing in config + cell_type_onehot_map_path = join(run_output_dir, "cell_type_onehot_map.pkl") + pert_onehot_map_path = join(run_output_dir, "pert_onehot_map.pt") + batch_onehot_map_path = join(run_output_dir, "batch_onehot_map.pkl") + var_dims_path = join(run_output_dir, "var_dims.pkl") + + with open(cell_type_onehot_map_path, "wb") as f: + pickle.dump(data_module.cell_type_onehot_map, f) + torch.save(data_module.pert_onehot_map, pert_onehot_map_path) + with open(batch_onehot_map_path, "wb") as f: + pickle.dump(data_module.batch_onehot_map, f) + with open(var_dims_path, "wb") as f: + pickle.dump(var_dims, f) if cfg["model"]["name"].lower() in ["cpa", "scvi"] or cfg["model"]["name"].lower().startswith("scgpt"): cfg["model"]["kwargs"]["n_cell_types"] = len(data_module.celltype_onehot_map) diff --git a/src/state_sets/configs/config.yaml b/src/state_sets/configs/config.yaml index 98caf057..5362c81e 100644 --- a/src/state_sets/configs/config.yaml +++ b/src/state_sets/configs/config.yaml @@ -4,7 +4,7 @@ defaults: - data: perturbation - model: pertsets - training: default - - wandb: abhinav # Change this as needed + - wandb: dhruv-gautam # Change this as needed - _self_ diff --git a/src/state_sets/configs/model/pertsets.yaml b/src/state_sets/configs/model/pertsets.yaml index 29117719..4399ee8c 100644 --- a/src/state_sets/configs/model/pertsets.yaml +++ b/src/state_sets/configs/model/pertsets.yaml @@ -5,6 +5,7 @@ device: cuda kwargs: cell_set_len: 512 # how many cells to group together into a single set of cells extra_tokens: 1 # configurable buffer for confidence/special tokens + decoder_hidden_dims: [1024, 1024, 512] blur: 0.05 hidden_dim: 328 # hidden dimension going into the transformer backbone loss: energy From 7b954d9e23d88fd1af9b2b00856ff98bb9d7e54f Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Sat, 21 Jun 2025 17:23:40 -0700 Subject: [PATCH 19/21] ruff --- src/state_sets/_cli/_sets/_infer.py | 2 +- src/state_sets/_cli/_sets/_train.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/state_sets/_cli/_sets/_infer.py b/src/state_sets/_cli/_sets/_infer.py index d5a85147..9cf016e4 100644 --- a/src/state_sets/_cli/_sets/_infer.py +++ b/src/state_sets/_cli/_sets/_infer.py @@ -96,7 +96,7 @@ def load_config(cfg_path: str) -> dict: pert_tensor = torch.zeros((len(pert_names), pert_dim), device="cpu") # Keep on CPU initially logger.info(f"Perturbation tensor shape: {pert_tensor.shape}") - # Load perturbation mapping from torch file + # Load perturbation mapping from torch file pert_onehot_map_path = os.path.join(args.output_dir, "pert_onehot_map.pt") pert_onehot_map = torch.load(pert_onehot_map_path, weights_only=False) diff --git a/src/state_sets/_cli/_sets/_train.py b/src/state_sets/_cli/_sets/_train.py index 5447dcf7..8e877c9a 100644 --- a/src/state_sets/_cli/_sets/_train.py +++ b/src/state_sets/_cli/_sets/_train.py @@ -138,13 +138,13 @@ def run_sets_train(cfg: DictConfig): # tuck it into the kwargs that will reach the LightningModule cfg["model"]["kwargs"]["decoder_cfg"] = decoder_cfg - + # Save the onehot maps as pickle files instead of storing in config cell_type_onehot_map_path = join(run_output_dir, "cell_type_onehot_map.pkl") pert_onehot_map_path = join(run_output_dir, "pert_onehot_map.pt") batch_onehot_map_path = join(run_output_dir, "batch_onehot_map.pkl") var_dims_path = join(run_output_dir, "var_dims.pkl") - + with open(cell_type_onehot_map_path, "wb") as f: pickle.dump(data_module.cell_type_onehot_map, f) torch.save(data_module.pert_onehot_map, pert_onehot_map_path) From 85894c92df6fb779d87f1b0a93941b1b6835f326 Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Sat, 21 Jun 2025 17:24:18 -0700 Subject: [PATCH 20/21] document --- src/state_sets/_cli/_sets/_infer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/state_sets/_cli/_sets/_infer.py b/src/state_sets/_cli/_sets/_infer.py index 9cf016e4..71ef5ad7 100644 --- a/src/state_sets/_cli/_sets/_infer.py +++ b/src/state_sets/_cli/_sets/_infer.py @@ -23,7 +23,7 @@ def add_arguments_infer(parser: argparse.ArgumentParser): "--output_dir", type=str, required=True, - help="Path to the output_dir containing the config.yaml file that was saved during training.", + help="Path to the output_dir containing the config.yaml file and the pert_onehot_map.pt file that was saved during training.", ) parser.add_argument( "--celltype_col", type=str, default=None, help="Column in adata.obs for cell type labels (optional)" From dbf2cef8fbd8efd36a59794b224b2a5e4e7180d1 Mon Sep 17 00:00:00 2001 From: dhruvgautam Date: Sat, 21 Jun 2025 23:26:25 -0700 Subject: [PATCH 21/21] infer working --- src/state_sets/sets/models/pert_sets.py | 46 ++++++------------------- 1 file changed, 11 insertions(+), 35 deletions(-) diff --git a/src/state_sets/sets/models/pert_sets.py b/src/state_sets/sets/models/pert_sets.py index 9f3af82f..bc54ac54 100644 --- a/src/state_sets/sets/models/pert_sets.py +++ b/src/state_sets/sets/models/pert_sets.py @@ -9,7 +9,6 @@ from typing import Tuple from .base import PerturbationModel -from .decoders_nb import NBDecoder, nb_nll from .utils import build_mlp, get_activation_class, get_transformer_backbone @@ -187,14 +186,6 @@ def __init__( for module in modules_to_freeze: for param in module.parameters(): param.requires_grad = False - - if kwargs.get("nb_decoder", False): - self.gene_decoder = NBDecoder( - latent_dim=self.output_dim + (self.batch_dim or 0), - gene_dim=gene_dim, - hidden_dims=[512, 512, 512], - dropout=self.dropout, - ) print(self) def _build_networks(self): @@ -375,18 +366,13 @@ def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int, padded=T else: latent_preds = pred - if isinstance(self.gene_decoder, NBDecoder): - mu, theta = self.gene_decoder(latent_preds) - gene_targets = batch["pert_cell_counts"].reshape_as(mu) - decoder_loss = nb_nll(gene_targets, mu, theta) + pert_cell_counts_preds = self.gene_decoder(latent_preds) + if padded: + gene_targets = gene_targets.reshape(-1, self.cell_sentence_len, self.gene_decoder.gene_dim()) else: - pert_cell_counts_preds = self.gene_decoder(latent_preds) - if padded: - gene_targets = gene_targets.reshape(-1, self.cell_sentence_len, self.gene_decoder.gene_dim()) - else: - gene_targets = gene_targets.reshape(1, -1, self.gene_decoder.gene_dim()) + gene_targets = gene_targets.reshape(1, -1, self.gene_decoder.gene_dim()) - decoder_loss = self.loss_fn(pert_cell_counts_preds, gene_targets).mean() + decoder_loss = self.loss_fn(pert_cell_counts_preds, gene_targets).mean() # Log decoder loss self.log("decoder_loss", decoder_loss) @@ -451,17 +437,11 @@ def validation_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> Non latent_preds = pred # Train decoder to map latent predictions to gene space - if isinstance(self.gene_decoder, NBDecoder): - mu, theta = self.gene_decoder(latent_preds) - gene_targets = batch["pert_cell_counts"].reshape_as(mu) - decoder_loss = nb_nll(gene_targets, mu, theta) - else: - # Get decoder predictions - pert_cell_counts_preds = self.gene_decoder(latent_preds).reshape( - -1, self.cell_sentence_len, self.gene_dim - ) - gene_targets = gene_targets.reshape(-1, self.cell_sentence_len, self.gene_dim) - decoder_loss = self.loss_fn(pert_cell_counts_preds, gene_targets).mean() + pert_cell_counts_preds = self.gene_decoder(latent_preds).reshape( + -1, self.cell_sentence_len, self.gene_dim + ) + gene_targets = gene_targets.reshape(-1, self.cell_sentence_len, self.gene_dim) + decoder_loss = self.loss_fn(pert_cell_counts_preds, gene_targets).mean() # Log the validation metric self.log("val/decoder_loss", decoder_loss) @@ -536,11 +516,7 @@ def predict_step(self, batch, batch_idx, padded=True, **kwargs): output_dict["confidence_pred"] = confidence_pred if self.gene_decoder is not None: - if isinstance(self.gene_decoder, NBDecoder): - mu, _ = self.gene_decoder(latent_output) - pert_cell_counts_preds = mu - else: - pert_cell_counts_preds = self.gene_decoder(latent_output) + pert_cell_counts_preds = self.gene_decoder(latent_output) output_dict["pert_cell_counts_preds"] = pert_cell_counts_preds