torchcell.data.Neo4jCellDataset

class Neo4jCellDataset(root: str, query: str | None = None, gene_set: GeneSet | None = None, graphs: GeneMultiGraph | None = None, incidence_graphs: dict[str, Graph | Hypergraph] | None = None, node_embeddings: dict[str, BaseEmbeddingDataset] | None = None, graph_processor: GraphProcessor | None = None, add_remaining_gene_self_loops: bool = True, converter: type[Converter] | None = None, deduplicator: type[Deduplicator] | None = None, aggregator: type[Aggregator] | None = None, overwrite_intermediates: bool = False, uri: str | None = None, username: str | None = None, password: str | None = None, transform: Callable[[...], Any] | None = None, pre_transform: Callable[[...], Any] | None = None, pre_filter: Callable[[...], Any] | None = None, phenotype_labels: list[str] | None = None, fetch_workers: int = 0, partition_prefix_length: int = 1)

Bases: Dataset

Dataset for loading cell data from Neo4j with file locking for DDP support.

get_init_graphs(gene_set: GeneSet) → GeneGraph

Build the initial cell and incidence graphs for the gene set.

property raw_file_names: str | list[str]

Return the raw file names expected by the dataset.

static load_raw(uri: str, username: str, password: str, root_dir: str, query: str, gene_set: GeneSet, record_observers: Sequence[Callable[[int, dict[str, Any]], None]] = (), fetch_workers: int = 0, partition_prefix_length: int = 1) → Neo4jQueryRaw

Query Neo4j and load the raw experiment records for the gene set.

record_observers see every record the raw stage writes (see RecordObserver); they see nothing when the raw LMDB already exists. fetch_workers and partition_prefix_length are Neo4jQueryRaw’s.

property processed_file_names: str | list[str]

Return the processed file names produced by the dataset.

property phenotype_info: list[Phenotype | FitnessPhenotype | GeneInteractionPhenotype | GeneEssentialityPhenotype | GradedGeneEssentialityPhenotype | SyntheticLethalityPhenotype | SyntheticRescuePhenotype | CalMorphPhenotype | MicroarrayExpressionPhenotype | MicroarrayExpressionScorePhenotype | MicroarrayExpressionRatioPhenotype | InductionTimecourseExpressionPhenotype | RNASeqExpressionPhenotype | PseudobulkExpressionPhenotype | VisualScorePhenotype | MetabolitePhenotype | ProteinAbundancePhenotype | ProteinFoldChangePhenotype | EnvironmentResponsePhenotype | ProductTiterPhenotype | ProteinTurnoverPhenotype | ProteinSynthesisRatePhenotype | PhosphositeIntensityPhenotype | FluxPhenotype | PromoterActivityPhenotype | PooledSortSeqExpressionPhenotype | BiosensorFluorescencePhenotype | MorphologyPhenotype | MrnaNumberFractionPhenotype | FitnessContrastPhenotype | SingleCellExpressionPhenotype | LifeCyclePhenotype | GenomeArchitecturePhenotype]

Return the cached phenotype type information, loading it if needed.

compute_phenotype_info() → None

Compute and cache the phenotype type information from the raw data.

process() → None

Run the conversion, deduplication, and processing pipeline into LMDB.

property gene_set: GeneSet

Return the dataset’s gene set, reading it from disk if available.

get(idx: int) → HeteroData | None

Return the processed cell graph data object at the given index.

refs_of(index: int) → list[ArtifactRef]

The distinct ArtifactRef``s of processed entry ``index, fetching nothing.

Walks every record the entry aggregates (experiment and reference) and keeps the first ref per (tier, key, path, sha256). Reads the processed LMDB only; no artifact source is consulted. IndexError for a missing entry.

static materialize(ref: ArtifactRef) → Path

The verified local path of ref’s file, fetched into the cache if needed.

torchcell.artifacts.materialize over the environment’s DATA_ROOT and tc-data (TC_DATA_URL when set, else the public endpoint), the same sources the raw stage’s resolvability gate checked the ref against at build time. Models call this when they need bytes.

len() → int

Return the number of records in the dataset.

close_lmdb() → None

Close the LMDB environment if it is open.

property label_df: DataFrame

Return a DataFrame of phenotype labels across all records.

compute_phenotype_label_index() → dict[str, list[int]]

Compute a mapping from phenotype label to record indices.

property phenotype_label_index: dict[str, list[int]]

Return the cached phenotype-label-to-indices mapping.

compute_dataset_name_index() → dict[str, list[int]]

Compute a mapping from dataset name to record indices.

property dataset_name_index: dict[str, list[int]]

Return the cached dataset-name-to-indices mapping.

compute_perturbation_count_index() → dict[int, list[int]]

Compute a mapping from perturbation count to record indices.

property perturbation_count_index: dict[int, list[int]]

Return the cached perturbation-count-to-indices mapping.

compute_is_any_perturbed_gene_index() → dict[str, list[int]]

Compute a mapping from gene to indices where the gene is perturbed.

property is_any_perturbed_gene_index: dict[str, list[int]]

Return the cached gene-to-perturbed-indices mapping.

compute_is_any_deletion_gene_index() → dict[str, list[int]]

Map gene -> record indices where the gene is DELETED (not merely perturbed).

The deletion-only counterpart of compute_is_any_perturbed_gene_index(), and the one to use whenever a gene name is meant to identify an experimental SUBJECT rather than anything appearing in a genotype.

The two differ sharply on the pigment screens, because their engineered cassettes are emitted as gene_addition / allele perturbations on EVERY strain – and three cassette members are native ORFs that also exist in the deletion collection: ARO4 (YBR249C) and ARO7 (YPR060C) in Cachera’s betaxanthin cassette, BTS1 (YPL069C) in Ozaydin’s beta-carotene cassette. Under the perturbed-gene index each of those resolves to ~4,400-4,700 records – the whole screen – instead of the single record in which it is the deleted gene. Selecting a gene set through the wrong index therefore silently selects nearly the entire dataset.

This is the same distinction DeletionKeyedGenotypeAggregator exists to make: the cassette belongs to the reference cell, not to the perturbation.

property is_any_deletion_gene_index: dict[str, list[int]]

Return the cached gene-to-deleted-indices mapping.