models.encoder_inpainting.SpatialEncoder

models.encoder_inpainting.SpatialEncoder(path_model, index=None, device=None)

A published spatial encoder, ready to predict features from position.

Attributes: path_model (Path): Local model directory. index (dict): The publication manifest. inputs (dict): Manifest inputs block – row identity and the required coordinates. outputs (dict): Manifest outputs block – the ordered feature list this predicts. config (dict): Architecture, neighbourhood and context settings.

Attributes

Name	Description
context_dir	Directory holding the published context volumes (`agea_vol_pca.npy` etc.).
features	Ordered names of the features this model predicts (`outputs.columns`).
model	The loaded encoder, built on first use.

Methods

Name	Description
confidence_model	Load the probe-confidence model published alongside, if there is one.
load_confidence_model	Instantiate the probe-confidence model published alongside, in `eval()` mode.
neighbor_bank	The published training-channel bank the model draws neighbours from.
predict	Predict electrophysiological features for each channel position.
predict_xyz	Predict the features at an array of positions (the array form of :meth:`predict`).
prediction_options	`(trilinear_context, neighbour_draws)` of a prediction: those given, else the release’s.
preprocessing_stats	Return the training-data statistics the model was fitted with.
selftest	Reproduce the shipped golden predictions, if the model ships an example.
split	The probe split the model was trained with.

confidence_model

models.encoder_inpainting.SpatialEncoder.confidence_model()

Load the probe-confidence model published alongside, if there is one.

Returns: The loaded checkpoint dict, or None when the release does not ship one. Returned raw rather than instantiated: it is a separate architecture with its own config, and nothing in predict consumes it.

load_confidence_model

models.encoder_inpainting.SpatialEncoder.load_confidence_model()

Instantiate the probe-confidence model published alongside, in eval() mode.

The confidence model scores, channel by channel, how consistent a probe’s recorded features are with this encoder’s predictions at its reported positions (used to flag misaligned insertions). Its constructor arguments are read from the checkpoint’s architecture record, falling back to the manifest’s config.confidence.

Returns: ProbeSequenceConfidenceTransformer | None: None when the release ships no confidence model.

Raises: ValueError: If the checkpoint does not record its architecture.

neighbor_bank

models.encoder_inpainting.SpatialEncoder.neighbor_bank()

The published training-channel bank the model draws neighbours from.

Returns: dict: xyz ([n, 3] positions, metres, mirrored to the left hemisphere), feat ([n, n_features] clipped and standardised features) and pid (insertion id per channel), exactly as the model saw them during training.

predict

models.encoder_inpainting.SpatialEncoder.predict(
    df,
    batch_size=1024,
    *,
    trilinear_context=None,
    neighbour_draws=None,
)

Predict electrophysiological features for each channel position.

The model was trained on positions mirrored to the left hemisphere (x -> -|x|): its context lookup, neighbour bank and position input all live there. A right-hemisphere position is therefore predicted exactly as its left mirror.

Args: df (pd.DataFrame): Indexed by (pid, channel), carrying the coordinate columns the manifest names in inputs.columns (x, y, z, in metres), in either hemisphere. The feature columns are not read – they are what this predicts. batch_size (int, optional): Rows per forward pass. trilinear_context, neighbour_draws: See :meth:predict_xyz.

Returns: pd.DataFrame: Indexed exactly like df, one pred_<feature> column per entry in outputs.columns, in feature units. The prefix keeps df.join(out) from colliding with the ground-truth columns of the same names.

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

predict_xyz

models.encoder_inpainting.SpatialEncoder.predict_xyz(
    xyz_m,
    pids=None,
    batch_size=1024,
    standardized=False,
    *,
    trilinear_context=None,
    neighbour_draws=None,
)

Predict the features at an array of positions (the array form of :meth:predict).

Args: xyz_m (np.ndarray): [N, 3] positions in metres, either hemisphere. pids (str | array-like, optional): Insertion id of each position (or one for all), whose own channels are excluded from the neighbours – pass the probe’s pid when predicting for a probe whose channels may be in the bank. None excludes nothing. batch_size (int, optional): Rows per forward pass. standardized (bool, optional): Return the model’s standardised output instead of feature units. trilinear_context (bool, optional): Interpolate the context between the voxels of its 200 µm grid, instead of taking the nearest voxel. By default, what the release says (:meth:prediction_options). neighbour_draws (int, optional): Average this many random neighbour sets, as the model was trained on (reproducible; 4 draws take about twice as long on one probe), instead of the nearest m_max neighbours. 0 keeps the nearest. By default, what the release says.

Returns: np.ndarray: [N, n_features] float32 predictions, in outputs.columns order.

prediction_options

models.encoder_inpainting.SpatialEncoder.prediction_options(
    trilinear_context=None,
    neighbour_draws=None,
)

(trilinear_context, neighbour_draws) of a prediction: those given, else the release’s.

A release records them in config.prediction; a release without that block predicts from the nearest voxel and the nearest neighbours, as it did when it was published.

preprocessing_stats

models.encoder_inpainting.SpatialEncoder.preprocessing_stats()

Return the training-data statistics the model was fitted with.

Always holds the standardisation the weights expect – the feature and context means and stds, i.e. the registered buffers the checkpoint ships with. When the release publishes preprocessing/channel_stats.npz it also holds what the buffers cannot carry: the per-feature clipping percentiles (rec_ephys_low_pctl/rec_ephys_high_pctl) applied to the recorded features before standardisation.

Returns: dict: {"e_mean", "e_std", "ctx_mean", "ctx_std", ...} as numpy arrays.

selftest

models.encoder_inpainting.SpatialEncoder.selftest(rtol=0.0001, atol=1e-05)

Reproduce the shipped golden predictions, if the model ships an example.

The tolerances are deliberately loose. These are float32 transformer outputs, and a model loaded in a different environment than the one that generated the golden example drifts by ~1e-4 relative from BLAS/library-version differences – a genuine reproduction, not a broken model. atol guards the near-zero outputs a relative bound alone cannot: a predicted value of order 1e-8 can never satisfy any rtol against tiny absolute noise. The check still catches real corruption or the wrong model, which drift by orders of magnitude more.

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

Returns: bool: True when the recomputed predictions match the shipped ones.

Raises: FileNotFoundError: If the model does not ship example/ files.

split

models.encoder_inpainting.SpatialEncoder.split()

The probe split the model was trained with.

Returns: dict: {"train_pids", "validation_pids", "test_pids"} (lists of insertion ids), plus whatever else the published split.json records.

Raises: FileNotFoundError: If the release publishes no split.