coco_pipe.io.embeddings#

BIDS-compatible derivative I/O for foundation-model embeddings.

Functions#

`embedding_sidecar_path`(path)	Return the JSON sidecar paired with an embedding NPZ.
`validate_embedding_derivative`(path)	Validate arrays, shape consistency, and sidecar presence.
`save_embedding_derivative`(result, path[, metadata, ...])	Atomically write an embedding NPZ and matching JSON sidecar.
`save_embedding_outputs`(result, embedding_path, *[, ...])	Write independent pooled and optional token derivatives from result.
`discover_embedding_derivatives`(root[, model_key, kind])	Discover valid `"embedding"` or `"token"` NPZ artifacts under a root.
`embedding_observation_id`(metadata, path[, window_position])	Return the canonical recording/window ID for pooled or token artifacts.
`load_embedding_derivatives`(paths[, representation, ...])	Load embedding artifacts into a DataContainer.
`combined_embedding_table_path`(derivative_root, ...[, ...])	Path to the merged per-(model, condition) embedding table.
`load_combined_embedding_table`(derivative_root, ...[, ...])	Load one merged per-(model, condition) embedding table as a 2-D container.
`write_embedding_manifest`(root, records)	Write a JSON run manifest indexing successful and failed artifacts.
`write_embedding_dataset_description`(root, name, ...[, ...])	Write a BIDS derivative dataset_description.json.

Module Contents#

coco_pipe.io.embeddings.embedding_sidecar_path(path)#

Return the JSON sidecar paired with an embedding NPZ.

Parameters:: path (str | pathlib.Path)
Return type:: pathlib.Path

coco_pipe.io.embeddings.validate_embedding_derivative(path)#

Validate arrays, shape consistency, and sidecar presence.

Parameters:: path (str | pathlib.Path)
Return type:: dict[str, Any]

coco_pipe.io.embeddings.save_embedding_derivative(result, path, metadata=None, overwrite=False)#

Atomically write an embedding NPZ and matching JSON sidecar.

Parameters:

result (Any)
path (str | pathlib.Path)
metadata (collections.abc.Mapping[str, Any] | None)
overwrite (bool)

Return type:

tuple[pathlib.Path, pathlib.Path]

coco_pipe.io.embeddings.save_embedding_outputs(result, embedding_path, *, token_path=None, metadata=None, pooled_metadata=None, overwrite=False)#

Write independent pooled and optional token derivatives from result.

embedding_path and token_path are independent derivative locations. When result.token_embeddings is present, token_path is required. metadata applies to both outputs, while pooled_metadata applies only to the pooled derivative. This keeps pooling-specific representation identity out of pooling-independent native-token metadata. Complete outputs are left untouched unless overwrite is true; a partial NPZ/sidecar pair is repaired.

Parameters:

result (Any)
embedding_path (str | pathlib.Path)
token_path (str | pathlib.Path | None)
metadata (collections.abc.Mapping[str, Any] | None)
pooled_metadata (collections.abc.Mapping[str, Any] | None)
overwrite (bool)

Return type:

tuple[pathlib.Path, pathlib.Path | None]

coco_pipe.io.embeddings.discover_embedding_derivatives(root, model_key=None, *, kind='embedding')#

Discover valid "embedding" or "token" NPZ artifacts under a root.

Artifact kind is identified by filename suffix (see ARTIFACT_SUFFIX), which save_embedding_derivative enforces. The run manifest (if present) may index either kind. When a requested model is not in that manifest, a recursive scan also finds independently materialized model variants.

Parameters:

root (str | pathlib.Path)
model_key (str | None)
kind (str)

Return type:

list[pathlib.Path]

coco_pipe.io.embeddings.embedding_observation_id(metadata, path, window_position=None)#

Return the canonical recording/window ID for pooled or token artifacts.

Parameters:

metadata (collections.abc.Mapping[str, Any])
path (str | pathlib.Path)
window_position (int | None)

Return type:

str

coco_pipe.io.embeddings.load_embedding_derivatives(paths, representation='recording', aggregate_by=None, model_key=None)#

Load embedding artifacts into a DataContainer.

representation selects both the on-disk array and the output rank:

"recording" — the pooled recording_embedding (2-D obs x feature)
"epoch" — per-epoch window_embeddings (2-D obs x feature)
"token" — native per-epoch feature tensors from separate *_tokens.npz artifacts (arbitrary rank, beginning with obs)

A coarser "subject" level is produced by the merge step, not here (it pools across recordings). Kinds are discovered by filename suffix (see ARTIFACT_SUFFIX); token and embedding artifacts are separate files that join only by shared observation ID.

Parameters:

paths (collections.abc.Sequence[str | pathlib.Path] | str | pathlib.Path)
representation (str)
aggregate_by (str | None)
model_key (str | None)

Return type:

coco_pipe.io.structures.DataContainer

coco_pipe.io.embeddings.combined_embedding_table_path(derivative_root, model_key, condition, representation='recording')#

Path to the merged per-(model, condition) embedding table.

Mirrors the combined/<model>_<condition>_<representation>_embeddings.parquet layout written by the merge step. representation is one of AGGREGATION_LEVELS.

Parameters:

derivative_root (str | pathlib.Path)
model_key (str)
condition (str)
representation (str)

Return type:

pathlib.Path

coco_pipe.io.embeddings.load_combined_embedding_table(derivative_root, model_key, condition, representation='recording', aggregate_by=None)#

Load one merged per-(model, condition) embedding table as a 2-D container.

This reads the single parquet the merge step already materialized instead of rescanning every per-recording NPZ and filtering — one table read per condition, the same access pattern descriptors use. representation is one of AGGREGATION_LEVELS. The table carries id columns (subject/session/run/condition/recording_id/model_key[/window_index]) plus embedding_* feature columns; the former become coords, the latter X.

Parameters: