File size: 2,721 Bytes
e697bd2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | """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
|