loki2.preprocess

Preprocessing utilities for Loki2.

This module provides utilities for loading and preprocessing Loki2 outputs, including .pt embeddings and JSON metadata files.

Module Contents

loki2.preprocess.setup_dummy_modules() → None

Sets up dummy modules for ‘loki2’ and its submodules.

This is required to simulate the module structure that the torch model expects when unpickling objects that reference e.g. ‘loki2.data.dataclass.cell_graph’.

loki2.preprocess.load_and_print_tensor(file_path: str) → Any

Load a PyTorch object from a given file path.

Parameters:

file_path – Path to the .pt file.

Returns:

The loaded Python object (typically a tensor or model state).

Return type:

Any

loki2.preprocess.load_sample(base: str, sample_name: str) → Tuple[numpy.ndarray, numpy.ndarray, List[Dict[str, Any]]]

Load a single sample’s embeddings, positions, and JSON cell metadata.

Expects:

  • {sample_name}_cells.pt (with attributes .x and .positions)

  • {sample_name}_cells.json (list or dict describing cells)

Parameters:

  • base – Root directory containing the .pt and .json files.

  • sample_name – Sample identifier (without suffix).

Returns:

  • emb: Embeddings array of shape (N, D) from data.x

  • pos: Positions array from data.positions

  • cells: List of per-cell dictionaries

Return type:

Tuple[np.ndarray, np.ndarray, List[Dict[str, Any]]]

Raises:

TypeError – If JSON structure is not recognized.