models.unit_encoder.UnitEncoder

models.unit_encoder.UnitEncoder(path_model, index=None, device=None)

A published unit-level model, ready to predict unit phenotypes at brain positions.

Attributes: path_model (Path): Local model directory. index (dict): The publication manifest. artifacts (dict): Manifest artifacts block, completed with the canonical names. inputs (dict): Manifest inputs block. outputs (dict): Manifest outputs block – the ordered phenotype features predicted.

Attributes

Name	Description
cfg	The released `Config`, rebuilt from `config.json` (paths are runtime defaults).
context_dir	Directory holding the released context volumes (`agea_vol_pca.npy` etc.).
context_model	The sub-type weights at raw contexts (`HierarchicalWeightModel`).
features	Ordered names of the phenotype features `predict` returns.
mixture	The putative cell types and their sub-types (`HierarchicalMixture`).
readout	The sub-type phenotypes (`MemberReadout`).
subtype_labels	`"<type>.<sub-type>"` label of each sub-type, e.g. `"7.2"`.

Methods

Name	Description
assign	Top-down hard assignment `(type [N], sub-type [N])`: the most likely type, then its
bundle	A :class:`~ephysatlas.unit_level_encoder.pipeline.UnitModelBundle` over a unit dataset.
components	The types and sub-types: weights, means and covariances, and each sub-type’s type.
encode	Embed units into the joint latent phenotype.
latent_dims	Indices of the latent dimensions of `modalities` (e.g. `("waveform", "acg")`).
mixture_weights	Putative cell-type composition at each position.
posterior	Top-down posteriors `(P(type \\| z) [N, K], P(type, sub-type \\| z) [N, L])`.
predict	Predict the expected unit phenotype at each position.
preprocessing_stats	The training-data statistics released with the model (`preprocessing/unit_stats.npz`).
sample	Draw unit phenotypes at each position from the model’s predictive distribution.
selftest	Reproduce the shipped golden outputs.
split	The probe split the model was trained with (`train`/`validation`/`test` pids).

assign

models.unit_encoder.UnitEncoder.assign(standardized_latents, dims=None)

Top-down hard assignment (type [N], sub-type [N]): the most likely type, then its most likely sub-type (an index into :attr:subtype_labels).

bundle

models.unit_encoder.UnitEncoder.bundle(data)

A :class:~ephysatlas.unit_level_encoder.pipeline.UnitModelBundle over a unit dataset.

Args: data: Prepared unit dataset (UnitData), e.g. from prepare_unit_data.

components

models.unit_encoder.UnitEncoder.components()

The types and sub-types: weights, means and covariances, and each sub-type’s type.

encode

models.unit_encoder.UnitEncoder.encode(
    waveform,
    acg=None,
    stpc=None,
    standardize=True,
    batch_size=4096,
)

Embed units into the joint latent phenotype.

Args: waveform (np.ndarray): [N, C, T] max-abs normalized multi-channel waveforms, as prepared for training (waveform_shape in the config). acg (np.ndarray, optional): [N, n_bins, n_lags] 3D autocorrelograms; required when the model uses the ACG modality. stpc (np.ndarray, optional): [N, n_lags] population coupling; required when the model uses the population-coupling (stPC) modality. standardize (bool, optional): Apply the released latent scaler (the space the mixture lives in). Defaults to True. batch_size (int, optional): Units per forward pass.

Returns: np.ndarray: [N, latent_dim] latents.

latent_dims

models.unit_encoder.UnitEncoder.latent_dims(modalities)

Indices of the latent dimensions of modalities (e.g. ("waveform", "acg")).

mixture_weights

models.unit_encoder.UnitEncoder.mixture_weights(
    df,
    level='subtype',
    batch_size=65536,
)

Putative cell-type composition at each position.

Args: df (pd.DataFrame): Carries x, y, z in metres; any index. level (str, optional): "type" for the type weights gamma_k(x), "subtype" for the sub-type weights gamma_k(x) gamma_j|k(x). batch_size (int, optional): Positions per forward pass.

Returns: pd.DataFrame: Indexed like df; one type_<k> column per type, or one subtype_<k>.<j> column per sub-type. Each row sums to one.

posterior

models.unit_encoder.UnitEncoder.posterior(standardized_latents, dims=None)

Top-down posteriors (P(type | z) [N, K], P(type, sub-type | z) [N, L]).

Args: standardized_latents: [N, latent_dim] latents from :meth:encode. dims: Optional latent dimensions to marginalize onto, e.g. latent_dims(("waveform", "acg")) for units without population coupling.

predict

models.unit_encoder.UnitEncoder.predict(df, batch_size=65536)

Predict the expected unit phenotype at each position.

The expected phenotype is the local sub-type weights times each sub-type’s expected features: the mean features of its training members, weighted by their posterior. Deterministic, and exactly what the published unit-level atlas figures map.

Args: df (pd.DataFrame): Carries the coordinate columns the manifest names in inputs.columns (x, y, z, metres). Any index; it is preserved. batch_size (int, optional): Positions per forward pass.

Returns: pd.DataFrame: Indexed exactly like df, one pred_<feature> column per phenotype feature. The prefix keeps df.join(out) from colliding with observed features.

Raises: KeyError: If a coordinate column is absent, naming it. ValueError: If the manifest’s feature list no longer matches its recorded digest.

preprocessing_stats

models.unit_encoder.UnitEncoder.preprocessing_stats()

The training-data statistics released with the model (preprocessing/unit_stats.npz).

sample

models.unit_encoder.UnitEncoder.sample(
    df,
    n_samples=16,
    seed=0,
    batch_size=65536,
)

Draw unit phenotypes at each position from the model’s predictive distribution.

Each draw picks a sub-type from the local sub-type weights, then one of its training members with probability proportional to its posterior of that sub-type: a real training unit. The draws average to :meth:predict.

Args: df (pd.DataFrame): Carries x, y, z in metres; any index. n_samples (int, optional): Draws per position. seed (int, optional): Seed of the draws. batch_size (int, optional): Positions per forward pass.

Returns: pd.DataFrame: n_samples rows per position, indexed like df with each label repeated (df.index.repeat(n_samples)): sample (draw number), type and subtype (the drawn sub-type, an index into :attr:subtype_labels), and one column per phenotype feature.

selftest

models.unit_encoder.UnitEncoder.selftest(rtol=0.0001, atol=1e-05)

Reproduce the shipped golden outputs.

Checks both halves of the model: phenotype predictions at a sample of positions (context sampling, sub-type weights and the members readout), and the latents of a small sample of synthetic units (the autoencoder and the scaler). Neither needs IBL data access. Tolerances allow the ~1e-4 relative float32 drift between BLAS/library versions; real corruption or the wrong model differ by orders of magnitude more.

Args: rtol (float, optional): Relative tolerance. atol (float, optional): Absolute tolerance, for near-zero outputs.

Returns: bool: True when every shipped golden output is reproduced.

Raises: FileNotFoundError: If the model ships no golden example.

split

models.unit_encoder.UnitEncoder.split()

The probe split the model was trained with (train/validation/test pids).