nemos.basis._basis.AdditiveBasis.split_by_feature#

AdditiveBasis.split_by_feature(x, axis=1)[source]#

Decompose an array along a specified axis into sub-arrays based on the basis components.

This function takes an array (e.g., a design matrix or model coefficients) and splits it along a designated axis. Each split corresponds to a different additive component of the basis, preserving all dimensions except the specified axis.

How It Works:

Suppose the basis is made up of m components, each with \(b_i\) basis functions and \(n_i\) inputs. The total number of features, \(N\), is calculated as:

\[N = b_1 \cdot n_1 + b_2 \cdot n_2 + \ldots + b_m \cdot n_m\]

This method splits any axis of length \(N\) into sub-arrays, one for each basis component.

The sub-array for the i-th basis component is reshaped into dimensions \((n_i, b_i)\).

For example, if the array shape is \((1, 2, N, 4, 5)\), then each split sub-array will have shape:

\[(1, 2, n_i, b_i, 4, 5)\]

where:

  • \(n_i\) represents the number of inputs associated with the i-th component,

  • \(b_i\) represents the number of basis functions in that component.

The specified axis (axis) determines where the split occurs, and all other dimensions remain unchanged. See the example section below for the most common use cases.

Parameters:

  • x (NDArray) –

    The input array to be split, representing concatenated features, coefficients, or other data. The shape of x along the specified axis must match the total number of features generated by the basis, i.e., self.n_output_features.

    Examples: - For a design matrix: (n_samples, total_n_features) - For model coefficients: (total_n_features,) or (total_n_features, n_neurons).

  • axis (int, optional) – The axis along which to split the features. Defaults to 1. Use axis=1 for design matrices (features along columns) and axis=0 for coefficient arrays (features along rows). All other dimensions are preserved.

Raises:

ValueError – If the shape of x along the specified axis does not match self.n_output_features.

Returns:

A dictionary where:

  • Keys: Labels of the additive basis components.

  • Values: Sub-arrays corresponding to each component. Each sub-array has the shape:

\[(..., n_i1, n_i2,..,n_im, b_i, ...)\]

  • (n_samples, n_i1, n_i2,..,n_im): is the shape of the input to the i-th basis.

  • b_i: The number of basis functions for the i-th basis component.

These sub-arrays are reshaped along the specified axis, with all other dimensions remaining the same.

Return type:

dict

Examples

>>> import numpy as np
>>> from nemos.basis import BSplineConv
>>> from nemos.glm import GLM
>>> # Define an additive basis
>>> basis = (
...     BSplineConv(n_basis_funcs=5, window_size=10, label="feature_1") +
...     BSplineConv(n_basis_funcs=6, window_size=10, label="feature_2")
... )
>>> # Generate a sample input array and compute features
>>> x1, x2 = np.random.randn(20), np.random.randn(20)
>>> X = basis.compute_features(x1, x2)
>>> # Split the feature matrix along axis 1
>>> split_features = basis.split_by_feature(X, axis=1)
>>> for feature, arr in split_features.items():
...     print(f"{feature}: shape {arr.shape}")
feature_1: shape (20, 5)
feature_2: shape (20, 6)
>>> # If one of the basis components accepts multiple inputs, the resulting dictionary will be nested:
>>> multi_input_basis = BSplineConv(n_basis_funcs=6, window_size=10,
... label="multi_input")
>>> X_multi = multi_input_basis.compute_features(np.random.randn(20, 2))
>>> split_features_multi = multi_input_basis.split_by_feature(X_multi, axis=1)
>>> for feature, sub_dict in split_features_multi.items():
...        print(f"{feature}, shape {sub_dict.shape}")
multi_input, shape (20, 2, 6)
>>> # the method can be used to decompose the glm coefficients in the various features
>>> counts = np.random.poisson(size=20)
>>> model = GLM().fit(X, counts)
>>> split_coef = basis.split_by_feature(model.coef_, axis=0)
>>> for feature, coef in split_coef.items():
...     print(f"{feature}: shape {coef.shape}")
feature_1: shape (5,)
feature_2: shape (6,)