| """Load the study's native SCimilarity encoder and decoder with PyTorch.""" |
|
|
| import json |
| from pathlib import Path |
|
|
| import torch |
|
|
| from _scimilarity_networks import Decoder, Encoder |
|
|
|
|
| class SCimilarity(torch.nn.Module): |
| def __init__(self, gene_order, hidden_dim, latent_dim): |
| super().__init__() |
| self.gene_order = tuple(gene_order) |
| self.n_genes = len(self.gene_order) |
| self.latent_dim = latent_dim |
| self.encoder = Encoder(self.n_genes, latent_dim=latent_dim, hidden_dim=hidden_dim) |
| self.decoder = Decoder(self.n_genes, latent_dim=latent_dim, hidden_dim=hidden_dim) |
|
|
| def forward(self, expression): |
| embedding = self.encoder(expression) |
| return embedding, self.decoder(embedding) |
|
|
|
|
| def _load_state(path): |
| checkpoint = torch.load(path, map_location="cpu", weights_only=True) |
| if not isinstance(checkpoint, dict) or set(checkpoint) != {"state_dict"}: |
| raise ValueError(f"Expected only a state_dict in {path}") |
| state = checkpoint["state_dict"] |
| if not isinstance(state, dict) or not state: |
| raise ValueError(f"Missing model tensors in {path}") |
| if any(not isinstance(value, torch.Tensor) or not torch.isfinite(value).all() |
| for value in state.values()): |
| raise ValueError(f"Non-tensor or nonfinite model state in {path}") |
| return state |
|
|
|
|
| def load_scimilarity(model_directory): |
| directory = Path(model_directory) |
| gene_order = [ |
| line.strip() for line in (directory / "gene_order.tsv").read_text().splitlines() |
| ] |
| if not gene_order or any(not gene for gene in gene_order) or len(set(gene_order)) != len(gene_order): |
| raise ValueError("SCimilarity gene order must contain unique, nonempty gene symbols") |
| encoder_state = _load_state(directory / "encoder.ckpt") |
| layer_sizes = json.loads((directory / "layer_sizes.json").read_text()) |
| if {name: list(value.shape) for name, value in encoder_state.items()} != layer_sizes: |
| raise ValueError("SCimilarity encoder tensors disagree with layer_sizes.json") |
| layers = [ |
| shape for name, shape in sorted( |
| layer_sizes.items(), key=lambda item: int(item[0].split(".")[1]) |
| ) |
| if name.endswith(".weight") and len(shape) == 2 |
| ] |
| if len(layers) < 2 or layers[0][1] != len(gene_order): |
| raise ValueError("SCimilarity network dimensions disagree with the gene order") |
| model = SCimilarity( |
| gene_order, hidden_dim=[shape[0] for shape in layers[:-1]], |
| latent_dim=layers[-1][0], |
| ) |
| model.encoder.load_state_dict(encoder_state, strict=True) |
| model.decoder.load_state_dict(_load_state(directory / "decoder.ckpt"), strict=True) |
| model.eval() |
| return model |
|
|
