torchcell.data.LazySubgraphRepresentation

class LazySubgraphRepresentation

Bases: GraphProcessor

Graph processor that returns full graph with masks instead of filtered subgraphs.

Zero-copy approach for edges: References original edge_index tensors, only computes boolean masks. Uses incidence cache for O(k×d) edge mask computation instead of O(E) scanning.

Output structure:

Node data (same as SubgraphRepresentation):

• node_ids: IDs of kept nodes (filtered)

• num_nodes: count of kept nodes

• ids_pert: IDs of perturbed nodes

• perturbation_indices: indices of perturbed nodes in original graph

• x: features of kept nodes (sliced from cell_graph)

• x_pert: features of perturbed nodes (sliced from cell_graph)

• pert_mask: boolean mask over ALL nodes, True for perturbed

• mask: boolean mask over ALL nodes, True for kept (inverse of pert_mask)

Edge data (ZERO-COPY - different from SubgraphRepresentation):

• edge_index: FULL edge_index (reference to cell_graph, not filtered!)

• num_edges: FULL edge count (not filtered count)

• mask: boolean mask, True for edges to keep (both endpoints non-perturbed)

Reaction/Metabolite data:

• Same as SubgraphRepresentation

Key difference from SubgraphRepresentation:

• Edges are NOT filtered or relabeled

• Edge indices refer to original node numbering

• Model must apply edge masks during message passing

Complexity:

• Cache build: O(E) one-time per dataset

• Per sample: O(k×d) where k=perturbed genes, d=average degree

• vs SubgraphRepresentation: O(N+E) per sample

Memory per sample:

• Edge tensors: 0 bytes (references only)

• Edge masks: ~2.7MB (9 edge types × 2.4M edges × 1 bit)

• vs SubgraphRepresentation: ~170MB (tensor copies)

property edge_incidence_cache: dict[Any, list[Tensor]]

Lazily build and return edge incidence cache.

Maps each node to tensor of edge positions where that node appears. Built on first use and rebuilt whenever a different cell graph arrives (see _build_incidence_cache).

build_cache(cell_graph: HeteroData) → dict[str, Any]

Explicitly build incidence cache and return timing information.

This should be called once during dataset initialization or before benchmarking to separate one-time cache build cost from per-sample processing time.

Returns:

dict with cache build timing:

{
    'total_time_ms': float,
    'num_edge_types': int,
    'total_edges': int
}

process(cell_graph: HeteroData, phenotype_info: list[Any], data: list[dict[str, Any]]) → HeteroData

Lazily build the perturbed subgraph and attach phenotype labels.