Results Variable Registry

Every variable stored in model.results_ (an xarray.Dataset) carries rich metadata describing its name, dimensions, description, and (where applicable) a LaTeX formula. This metadata is attached as .attrs on each DataArray and can be inspected at runtime:

model.results_["F"].attrs["description"]
model.results_["F"].attrs["formula"]

The full registry is built by the internal helper below. Expand it to browse every variable that SIMPL can produce.

Variable metadata registry for the SIMPL results xarray.Dataset.

Every variable stored in the SIMPL results dataset has an entry in the registry built by _build_variable_info_dict. Each entry is a dict with the following keys:

• name — human-readable display name
• description — brief description of the variable
• dims — list of dimension names (e.g. ['time', 'neuron'])
• axis_title — label used when plotting the variable
• formula — LaTeX formula (optional)
• reshape — if True, the array is reshaped to match dims when saved

These dicts are attached as attrs on the xr.DataArray instances inside the results Dataset, keeping metadata co-located with the data for introspection and serialisation.

_build_variable_info_dict(dim)

Build the dictionary describing every variable stored in SIMPL results.

This dictionary summarises all the variables used and returned by the SIMPL class. Each entry is attached as the attrs to the DataArray when these variables are saved. Each entry is a dict with:

name : str — the name of the variable
description : str — a brief description of the variable
dims : list — the dimensions of the variable (excluding any iteration
    dimension)
axis_title : str — the title of the axis when plotting the variable
formula : str — a formula for the variable (if applicable)
reshape : bool — whether to forcefully reshape to the dimensions given
    in ``dims`` when saved.  This is useful for variables calculated in
    a different shape to the final intended shape (e.g. receptive
    fields).

Parameters:

Name	Type	Description	Default
dim	list[str]	Environment dimension names, e.g. ['x', 'y'].	required

Source code in src/simpl/_variable_registry.py

def _build_variable_info_dict(dim: list[str]) -> dict:
    """Build the dictionary describing every variable stored in SIMPL results.

    This dictionary summarises _all_ the variables used and returned by the
    SIMPL class.  Each entry is attached as the ``attrs`` to the DataArray when
    these variables are saved.  Each entry is a dict with:

        name : str — the name of the variable
        description : str — a brief description of the variable
        dims : list — the dimensions of the variable (excluding any iteration
            dimension)
        axis_title : str — the title of the axis when plotting the variable
        formula : str — a formula for the variable (if applicable)
        reshape : bool — whether to forcefully reshape to the dimensions given
            in ``dims`` when saved.  This is useful for variables calculated in
            a different shape to the final intended shape (e.g. receptive
            fields).

    Parameters
    ----------
    dim : list[str]
        Environment dimension names, e.g. ``['x', 'y']``.
    """
    variable_info_dict = {
        # Core data variables
        "Y": {
            "name": "Spikes",
            "description": ("The spikes of the neurons at each time step. Each a binary vector of length N_neurons."),
            "dims": ["time", "neuron"],
            "axis_title": "Spike counts",
            "formula": r"$y(t)$",
        },
        "X": {
            "name": "Latent",
            "description": (
                "The latent positions of the agent at each "
                "time step. This is (typically) the smoothed "
                "output of the Kalman filter, scaled to "
                "correlate maximally with behavior "
                "(X <-- mu_s, X <-- X @ coef.T + intercept). "
                "For iteration 0, X == the behavior."
            ),
            "dims": ["time", "dim"],
            "axis_title": "Position",
            "formula": r"$x(t)$",
        },
        "F": {
            "name": "Model",
            "description": (
                "The receptive fields of the neurons "
                "(probability of the neurons firing at each "
                "position in one time step (of length dt)."
            ),
            "dims": ["neuron", *dim],
            "axis_title": "Receptive field",
            "reshape": True,
            "formula": r"$r(x)$",
        },
        "F_odd_minutes": {
            "name": "Model (odd minutes)",
            "description": (
                "The receptive fields of the neurons "
                "(probability of the neurons firing at each "
                "position in one time step) calculated from "
                "the odd minutes of the data."
            ),
            "dims": ["neuron", *dim],
            "axis_title": "Receptive field (odd mins)",
            "Formula": r"$r_{\textrm{odd}}(x)$",
            "reshape": True,
        },
        "F_even_minutes": {
            "name": "Model (even minutes)",
            "description": (
                "The receptive fields of the neurons "
                "(probability of the neurons firing at each "
                "position in one time step) calculated from "
                "the even minutes of the data."
            ),
            "dims": ["neuron", *dim],
            "axis_title": "Receptive field (even mins)",
            "Formula": r"$r_{\textrm{even}}(x)$",
            "reshape": True,
        },
        "FX": {
            "name": "Firing rates trajectories",
            "description": (
                "An estimate of the firing rate of the "
                "neurons at each time step based on the "
                "latest position estimates and their "
                "receptive fields. Only stored when "
                "save_full_history is True. See also "
                "FX_first_iteration and FX_last_iteration."
            ),
            "dims": ["time", "neuron"],
            "axis_title": r"Firing rate",
            "formula": r"$f(t)$",
        },
        "PX": {
            "name": "Occupancy",
            "description": (
                "The spatial occupancy at each position bin, "
                "estimated from the latent trajectory using a "
                "kernel density estimator. This is the "
                "denominator of the KDE used to fit receptive "
                "fields."
            ),
            "dims": [*dim],
            "axis_title": "Occupancy",
            "formula": r"$p(x)$",
            "reshape": True,
        },
        # Ground truth data variables
        "Xb": {
            "name": "Position (behavior)",
            "description": (
                "The position of the agent (i.e. not "
                'necessarily the "true" latent position which '
                "generated the spikes) at each time step. "
                "This is the behavior of the agent and acts "
                "as the starting conditions for the algorithm "
                "X[iteration=0] == Xb."
            ),
            "dims": ["time", "dim"],
            "axis_title": "Position (behavior)",
            "formula": r"$x_{\textrm{beh}}(t)$",
        },
        "Xt": {
            "name": "Latent (ground truth)",
            "description": (
                "The ground truth latent positions of the "
                "agent at each time step. If using real "
                "neural data, this is typically not available."
            ),
            "dims": ["time", "dim"],
            "axis_title": "Position (true)",
            "formula": r"$x_{\textrm{true}}(t)$",
        },
        "Ft": {
            "name": "Model (ground truth)",
            "description": (
                "The ground truth receptive fields. These are "
                "the true receptive fields of the neurons used "
                "to generate the data. If using real neural "
                "data, this is typically not available."
            ),
            "dims": ["neuron", *dim],
            "reshape": True,
            "axis_title": "Receptive field (true)",
            "formula": r"$r_{\textrm{true}}(x)$",
        },
        # Likelihood maps and Gaussian posterior parameters
        "logPYXF_maps": {
            "name": "Log-likelihoods",
            "description": (
                "The log-likelihood maps of the spikes, as a function of position, calculated for each time step. "
                "Only saved when save_full_history is True, and even then only for the last iteration due to its size."
            ),
            "dims": ["time", *dim],
            "axis_title": "Log-likelihood map",
            "formula": r"$\log P(Y|x, \Theta)$",
            "reshape": True,
        },
        "mu_l": {
            "name": "Likelihood mean",
            "description": "The mean of the Gaussian fitted to the log-likelihood maps.",
            "dims": ["time", "dim"],
            "axis_title": "Likelihood mean",
            "formula": r"$\mu_l(t)$",
        },
        "mode_l": {
            "name": "Likelihood mode",
            "description": "The mode of the Gaussian fitted to the log-likelihood maps.",
            "dims": ["time", "dim"],
            "axis_title": "Likelihood mode",
            "formula": r"$\textrm{Mo}(t)$",
        },
        "sigma_l": {
            "name": "Likelihood covariance",
            "description": "The covariance matrix of the Gaussian fitted to the log-likelihood maps.",
            "dims": ["time", "dim", "dim_"],
            "axis_title": "Likelihood covariance",
            "formula": r"$\Sigma_l(t)$",
        },
        "mu_f": {
            "name": "Kalman filtered mean",
            "description": "The mean of the Kalman filtered posterior of the latent positions.",
            "dims": ["time", "dim"],
            "axis_title": "Kalman filtered mean",
            "formula": r"$\mu_f(t)$",
        },
        "sigma_f": {
            "name": "Kalman filtered covariance",
            "description": "The covariance matrix of the Kalman filtered posterior of the latent positions.",
            "dims": ["time", "dim", "dim_"],
            "axis_title": "Kalman filtered covariance",
            "formula": r"$\Sigma_f(t)$",
        },
        "mu_s": {
            "name": "Kalman smoothed mean",
            "description": "The mean of the Kalman smoothed posterior of the latent positions.",
            "dims": ["time", "dim"],
            "axis_title": "Kalman smoothed mean",
            "formula": r"$\mu_s(t)$",
        },
        "sigma_s": {
            "name": "Kalman smoothed covariance",
            "description": "The covariance matrix of the Kalman smoothed posterior of the latent positions.",
            "dims": ["time", "dim", "dim_"],
            "axis_title": "Kalman smoothed covariance",
            "formula": r"$\Sigma_s(t)$",
        },
        # Linear scaling parameters
        "coef": {
            "name": "Linear scaling coefficients",
            "description": (
                "The linear scaling matrix coefficients "
                "applied, wlog, to the latent positions to "
                "maximally correlate it with the behavior."
            ),
            "dims": ["dim", "dim_"],
            "axis title": "Linear scaling coefficients",
            "formula": r"$\mathbf{M}$",
        },
        "intercept": {
            "name": "Linear scaling intercept",
            "description": (
                "The linear scaling intercept vector applied, "
                "wlog, to the latent positions to maximally "
                "correlate with the behavior."
            ),
            "dims": ["dim"],
            "axis title": "Linear scaling intercept",
            "formula": r"$\mathbf{c}$",
        },
        # Place field stuff
        "place_field_count": {
            "name": "Number of place fields indentified",
            "description": (
                "The number of place fields identified per "
                "neuron. Place fields are defined by "
                "thresholding the tuning curves (2D only) at "
                "1 Hz and identifying continuous regions which "
                "are both (i) under one-third the full size of "
                "the environment and (ii) have a maximum "
                "firing rate over 2 Hz."
            ),
            "dims": ["neuron"],
            "formula": r"$N_{pf}$",
            "axis title": "Number of place fields",
        },
        "place_field_position": {
            "name": "Place field position",
            "description": "The mean of the Gaussian fitted to each place field.",
            "dims": ["neuron", "place_field", "dim"],
            "axis title": "Place field position",
            "formula": r"$\mu_{pf}$",
        },
        "place_field_covariance": {
            "name": "Place field covariance",
            "description": "The covariance of the Gaussian fitted to each place field",
            "dims": ["neuron", "place_field", "dim", "dim_"],
            "axis title": "Place field covariance",
            "formula": r"$\Sigma_{pf}$",
        },
        "place_field_size": {
            "name": "Place field size",
            "description": "The size of the place field (units of m^2)",
            "dims": ["neuron", "place_field"],
            "axis title": "Place field size",
            "formula": r"$S_{pf}$",
        },
        "place_field_outlines": {
            "name": "Place field outlines",
            "description": "The outlines of the identified place fields (if any)",
            "dims": ["neuron", *dim],
            "axis title": "Place field outlines",
        },
        "place_field_roundness": {
            "name": "Place field roundness",
            "description": (
                "The roundness of the identified place fields (if any). Defined as r = 4 * pi * area / perimeter^2"
            ),
            "dims": ["neuron", "place_field"],
            "axis title": "Place field roundness",
        },
        "place_field_max_firing_rate": {
            "name": "Maximum firing rate of the place field",
            "description": (
                "Maximum firing rate -- in units of spikes per time bin rather than Hz -- of the place field"
            ),
            "dims": ["neuron", "place_field"],
            "axis title": "Place field max. firing rate",
        },
        # Result metrics
        "logPYF": {
            "name": "Total data log-likelihood",
            "description": (
                "Suppose we have two sets of observations, "
                "Y and Y_prime. The log-likelihood of the "
                "observations := P(Y | Y_prime) can be found "
                "by using Y_prime to estimate the latent "
                "posterior P(X | Y_prime) (i.e. run the "
                "Kalman model) and then marginalise out X "
                "from the likelihood of the new observations "
                "given the latent: P(Y | Y_prime) = int_X "
                "P(Y | X) P(X | Y_prime). This is called the "
                "posterior predictive and, for Kalman models, "
                "has analytic form P(Y_i | Y_prime) = "
                "Normal(Y_i | Y_hat, S) where "
                "S = H @ sigma_i @ H.T + R (the posterior "
                "observation covariance combined with the "
                "observation noise covariance) and "
                "Y_hat = H @ mu_i (the predicted "
                "observation), see page 361 of the Advanced "
                "Murphy book. Strictly this metric (i) "
                "returns not the full distribution (a product "
                "of many Gaussians) but the probability "
                "density evaluated at the observation "
                "locations and (ii) uses the same observations "
                "for both sets Y = Y_prime = "
                "observations_from_training_spikes. It is a "
                "good gauge on the performance of the model "
                '"were the observations (~spikes) P(Y|X) '
                "likely under the decoded trajectory "
                'P(X|Y_prime)?" but note that Y are the '
                "observations not the actual spikes so this "
                "is not strictly the likelihood of the spikes."
            ),
            "dims": [],
            "axis title": "Kalman posterior predictive",
            "formula": r"$\log P(Y|Y_{\textrm{train}})$",
        },
        "logPYF_val": {
            "name": "Total data log-likelihood (validation)",
            "description": (
                "See logPYF but for observations derived from "
                "the validation spikes i.e. logPYF_val = "
                "P(Y_val|Y_train) = int_X "
                "P(Y_val|X)P(X|Y_train). Its analogous to "
                'asking "were the validation observations (~validation '
                "spikes) P(Y_val|X) likely under the decoded "
                'trajectory P(X|Y_train)?" or in other words '
                '"can the train spikes be used to predict the '
                'validation spikes?"'
            ),
            "dims": [],
            "axis title": "Kalman posterior predictive (validation)",
            "formula": r"$\log P(Y_{\textrm{val}}|Y_{\textrm{train}})$",
        },
        "logPYXF": {
            "name": "Mean spike log-likelihood given trajectory",
            "description": (
                "The Poisson log-likelihood of the model: how "
                "well are the spike counts predicted by the "
                "firing rates along the trajectory? Normalised "
                "by the number of spike-time bins to make it "
                "comparable across validation and train splits."
            ),
            "dims": [],
            "axis title": "Spike log-likelihood",
            "formula": r"$\log P(Y|X(t), \Theta)$",
        },
        "logPYXF_val": {
            "name": "Mean spike log-likelihood given trajectory (validation)",
            "description": "See logPYXF, but applied to the validation spikes.",
            "dims": [],
            "axis title": "Spike log-likelihood (validation)",
            "formula": r"$\log P(Y_{\textrm{val}}|X(t), \Theta)$",
        },
        "bits_per_spike": {
            "name": "Bits per spike (train)",
            "description": (
                "Bits per spike on the training set. Interpret "
                "this as how many bits of information, on "
                "average, each spike carries about position. "
                "Computed by subtracting a mean-rate baseline "
                "from the Poisson log-likelihood and dividing "
                "by total spike count: bps = (ll_model - "
                "ll_mean_rate) / (n_spikes * log2). Comparable "
                "across datasets with different spike counts "
                "and dt."
            ),
            "dims": [],
            "axis_title": "Bits per spike (train)",
            "formula": (
                r"$\frac{\mathcal{L}(\hat\lambda)"
                r" - \mathcal{L}(\bar\lambda)}"
                r"{N_{\mathrm{spk}} \ln 2}$"
            ),
        },
        "bits_per_spike_val": {
            "name": "Bits per spike (validation)",
            "description": ("See bits_per_spike, but applied to the validation spikes."),
            "dims": [],
            "axis_title": "Bits per spike (validation)",
            "formula": (
                r"$\frac{\mathcal{L}_{\mathrm{val}}(\hat\lambda)"
                r" - \mathcal{L}_{\mathrm{val}}(\bar\lambda)}"
                r"{N_{\mathrm{spk,val}} \ln 2}$"
            ),
        },
        "mutual_information": {
            "axis title": "Mutual Information (bits/s)",
            "name": "Mutual information",
            "description": (
                "Exact mutual information between spike count "
                "and position, per neuron, in bits/s. Interpret "
                "this as on average how many bits of information "
                "per second do spikes from this neuron tell me "
                "about position. Computed from the Poisson "
                "likelihood and position occupancy. In the small "
                "time-bin limit, mutual information converges to "
                "the Skaggs spatial information (bits/s), but "
                "for finite time bins mutual information is a "
                "more accurate measure of the true information "
                "between spikes and position."
            ),
            "dims": ["neuron"],
            "formula": (
                r"$I(X;Y) = \frac{1}{\Delta t}\sum_x \sum_k"
                r" P(x)\, \mathrm{Pois}(k;\lambda(x))"
                r"\, \log_2 \frac{\mathrm{Pois}(k;\lambda(x))}"
                r"{P(k)}$"
            ),
        },
        "spatial_information": {
            "axis title": "Spatial Information (bits/s)",
            "name": "Spatial information",
            "description": (
                "Skaggs spatial information per neuron (bits/s). "
                "Interpret this as how much information, in "
                "bits-per-second, does the firing rate of this "
                "neuron tell me about position. Note: the "
                "bits/spike form of SI is provably identical to "
                "the bits-per-spike (BPS) metric; this rate form "
                "is SI_rate = BPS * mean_firing_rate. See "
                "``mutual_information`` for the exact finite-dt "
                "version (which differs from SI for finite time "
                "bins)."
            ),
            "dims": ["neuron"],
            "formula": (
                r"$I=\int_x \lambda(x) \log_2"
                r" \frac{\lambda(x)}{\lambda} p(x)\, dx$"
            ),
        },
        "negative_entropy": {
            "name": "Tuning curve neg-entropy",
            "description": "The negative entropy of each of the receptive fields, -sum(F * log(F)).",
            "dims": ["neuron"],
            "axis title": "Tuning curve neg-entropy",
            "formula": r"$- \sum F \log F$",
        },
        "sparsity": {
            "name": "Sparsity",
            "description": "The fraction of bins where the firing is greater than 0.1 * max firing rate.",
            "dims": ["neuron"],
            "axis title": "Spatial sparsity",
            "formula": r"$\rho(r(x))$",
        },
        "stability": {
            "name": "Stability",
            "description": (
                "The correlation between receptive fields estimated separately using spikes from odd and even minutes."
            ),
            "dims": ["neuron"],
            "axis title": "Field stability",
            "formula": r"$\textrm{Corr}(r_{\textrm{odd}}(x), r_{\textrm{even}}(x))$",
        },
        "field_count": {
            "name": "Number of fields",
            "description": (
                "A heuristic on the number of distinct fields "
                "in the place fields calculated as the number "
                "of connected components in the receptive "
                "fields thresholded at 0.5 * their maximum "
                "firing rate."
            ),
            "dims": ["neuron"],
            "axis title": "Number of fields",
            "formula": r"$N_{\textrm{fields}}$",
        },
        "field_size": {
            "name": "Field size",
            "description": "The average size of the fields in the individual fields, as defined in field_count.",
            "dims": ["neuron"],
            "axis title": "Field size",
            "formula": r"",
        },
        "field_change": {
            "name": "Field shift",
            "description": "The change in the place fields from the last iteration.",
            "dims": ["neuron"],
            "axis title": "Field change",
            "formula": r"$\|r_{e}(x) - r_{e-1}(x)\|$",
        },
        "trajectory_change": {
            "name": "Latent shift",
            "description": "The average change in the latent positions from the last iteration.",
            "dims": ["time"],
            "axis title": "Latent change",
            "formula": r"$\|x_{e}(t) - x_{e-1}(t)\|$",
        },
        # Baselines : calculated when ground truth data is available
        "X_R2": {
            "name": "Latent R2",
            "description": "The R2 score between the true and estimated latent positions.",
            "dims": [],
            "axis title": "Latent R2",
            "formula": r"$R^2$",
        },
        "X_err": {
            "name": "Latent error",
            "description": "The mean squared error between the true and estimated latent positions.",
            "dims": [],
            "axis title": "Latent error",
            "formula": r"$\|x_{\textrm{true}}(t) - x_{\textrm{est}}(t)\|$",
        },
        "F_err": {
            "name": "Model error",
            "description": "The mean squared error between the true and estimated place fields.",
            "dims": [],
            "axis title": "Receptive field error",
            "formula": r"$\|r_{\textrm{true}}(x) - r_{\textrm{est}}(x)\|$",
        },
        # Other variables
        "spike_mask": {
            "name": "Spike mask",
            "description": (
                "The mask used to determine which spikes are "
                "used for training and testing. This is "
                "usually a speckled mask (see Williams et al. "
                "2020, fig. 2) of mostly `True` interspersed "
                "with contiguous blocks of `False`. Where "
                "False, spikes are masked and NOT used for "
                "training."
            ),
            "dims": ["time", "neuron"],
            "axis_title": "Spike mask",
            "formula": r"$\textrm{mask}(t, n)$",
        },
    }

    # FX_first_iteration and FX_last_iteration: same as FX but for a single iteration (no iteration dim)
    for suffix, label in [("first_iteration", "first iteration (Xb)"), ("last_iteration", "last iteration")]:
        variable_info_dict[f"FX_{suffix}"] = {
            **variable_info_dict["FX"],
            "name": variable_info_dict["FX"]["name"] + f" ({suffix})",
            "description": variable_info_dict["FX"]["description"] + f" This is for the {label} only.",
        }

    return variable_info_dict

_dict_to_dataset(data, variable_info_dict, coords)

Convert a dictionary to an xarray Dataset, appending metadata from variable_info_dict where available.

Parameters:

Name	Type	Description	Default
data	dict	Dictionary of variable_name -> array.	required
variable_info_dict	dict	Variable metadata dictionary (from _build_variable_info_dict).	required
coords	dict	Coordinate arrays.	required

Returns:

Type	Description
Dataset

Source code in src/simpl/_variable_registry.py

def _dict_to_dataset(data: dict, variable_info_dict: dict, coords: dict) -> xr.Dataset:
    """Convert a dictionary to an xarray Dataset, appending metadata from variable_info_dict where available.

    Parameters
    ----------
    data : dict
        Dictionary of variable_name -> array.
    variable_info_dict : dict
        Variable metadata dictionary (from ``_build_variable_info_dict``).
    coords : dict
        Coordinate arrays.

    Returns
    -------
    xr.Dataset
    """
    dataset = xr.Dataset()
    for variable_name in data.keys():
        if variable_name in variable_info_dict:
            variable_info = variable_info_dict[variable_name]
            variable_coords = {k: coords[k] for k in variable_info["dims"]}
            intended_variable_shape = tuple([len(variable_coords[c]) for c in variable_info["dims"]])
            if "reshape" in variable_info and variable_info["reshape"]:
                variable_data = data[variable_name].reshape(intended_variable_shape)
            else:
                variable_data = data[variable_name]
            dataarray = xr.DataArray(
                variable_data,
                dims=variable_info["dims"],
                coords=variable_coords,
                attrs=variable_info,
            )
        else:
            warnings.warn(
                f"Variable {variable_name} not recognised, "
                f"it will be saved without coordinate or "
                f"dimension info unless you add it to the "
                f"variable_info_dict in "
                f"_variable_registry.py."
            )
            dataarray = xr.DataArray(data[variable_name])
        dataset[variable_name] = dataarray
    return dataset