loki2.cl.project_raw_embeddings
Project raw morphology/transcription embeddings using a ProjectionCL checkpoint.
Module Contents
- loki2.cl.project_raw_embeddings.load_morph_embeddings(path: pathlib.Path) Tuple[torch.Tensor, Dict[str, numpy.ndarray]]
Load morphology embeddings from a PyTorch file.
- Parameters:
path – Path to the morphology embeddings file (.pt).
- Returns:
embeddings: Morphology embeddings tensor
meta: Dictionary containing row_index, cell_id, and optionally positions
- Return type:
Tuple[torch.Tensor, Dict[str, np.ndarray]]
- loki2.cl.project_raw_embeddings.load_trans_embeddings(path: pathlib.Path) Tuple[torch.Tensor, Dict[str, numpy.ndarray]]
Load transcription embeddings from a PyTorch file.
- Parameters:
path – Path to the transcription embeddings file (.pt).
- Returns:
embeddings: Transcription embeddings tensor
meta: Dictionary containing cell_id array
- Return type:
Tuple[torch.Tensor, Dict[str, np.ndarray]]
- Raises:
KeyError – If the file doesn’t contain ‘embedding’ and ‘cell_id’ keys.
- loki2.cl.project_raw_embeddings.load_model(ckpt_path: pathlib.Path, device: torch.device) model_cl.ProjectionCL
Load a ProjectionCL model from a checkpoint file.
- Parameters:
ckpt_path – Path to the model checkpoint file (.pt).
device – Device to load the model on.
- Returns:
Loaded and initialized model in eval mode.
- Return type:
- Raises:
RuntimeError – If the checkpoint is missing model weights.
- loki2.cl.project_raw_embeddings.project_batches(tensor: torch.Tensor, projector: Callable[[torch.Tensor], torch.Tensor], batch_size: int, device: torch.device, normalized: bool) Tuple[numpy.ndarray, numpy.ndarray | None]
Project embeddings in batches using a projection function.
- Parameters:
tensor – Input embeddings tensor to project.
projector – Projection function (e.g., model.encode_a or model.encode_b).
batch_size – Batch size for processing.
device – Device to run projections on.
normalized – Whether to also return L2-normalized projections.
- Returns:
projections: Projected embeddings as NumPy array
normalized_projections: L2-normalized projections if normalized=True, otherwise None
- Return type:
Tuple[np.ndarray, Optional[np.ndarray]]
- loki2.cl.project_raw_embeddings.save_output(directory: pathlib.Path, base: str, tag: str, payload: Dict[str, numpy.ndarray]) None
Save projected embeddings to an NPZ file.
- Parameters:
directory – Output directory for the NPZ file.
base – Base filename (e.g., “morph_proj” or “trans_proj”).
tag – Optional tag to append to filename.
payload – Dictionary of arrays to save in the NPZ file.