nemos.convolve.create_convolutional_predictor#

nemos.convolve.create_convolutional_predictor(basis_matrix, time_series, predictor_causality='causal', shift=None, axis=0, batch_size_samples=None, batch_size_channels=None, batch_size_basis=None, use_fft=None)[source]#

Create a convolutional predictor by convolving a basis matrix with a time series.

To create the convolutional predictor, three steps are taken, convolve, pad and shift.

  • Convolve basis_matrix with time_series (function: _convolve_1d_trials)

  • Pad output with basis_matrix.shape[0]-1 NaNs, with location based on causality of intended predictor (function: nemos.utils.nan_pad).

  • (Optional) Shift predictor based on causality (function: nemos.utils.shift_time_series)

The function is designed to handle both single arrays and pynapple time series with data

(across multiple epochs), as well as their combinations. For pynapple time series, it

treats each epoch separately, applies the convolution process, and then reassembles the time series.

Parameters:

  • basis_matrix (ArrayLike) – A 2D array representing the basis matrix to convolve with the time series. The first dimension should represent the window size for the convolution.

  • time_series (Any) – The time series data to convolve with the basis matrix. Can be single arrays, pytree of arrays, pynapple time series, pytree of pynapple time series or a mix. In case of Pynapple time series data, each epoch will be convolved separately.

  • predictor_causality (Literal['causal', 'acausal', 'anti-causal']) – The causality of the predictor, determining how the padding and shifting should be applied to the convolution result. - ‘causal’: Pads and/or shifts the result to be causal with respect to the input. - ‘acausal’: Applies padding equally on both sides without shifting. - ‘anti-causal’: Pads and/or shifts the result to be anti-causal with respect to the input.

  • shift (Optional[bool]) – Determines whether to shift the convolution result based on the causality. If None, it defaults to True for ‘causal’ and ‘anti-causal’ and to False for ‘acausal’.

  • axis (int) – The axis along which the convolution is applied.

  • batch_size_samples (Optional[int]) – Batch size for the convolution in terms of number of sample points. If this parameter is set, the convolution will be applied sequentially over a fixed-length chunks of the time_series.

  • batch_size_channels (Optional[int]) – Batch size for the convolution in terms of number of input channels. If this parameter is set, the convolution will be vectorized over batches of input channels. Default vectorizes over all channels.

  • batch_size_basis (Optional[int]) – Batch size for the convolution in terms of number of basis kernels. If this parameter is set, the convolution will be vectorized over batches of kernels. Default vectorizes over all basis kernels.

  • use_fft (Optional[bool]) – Whether to compute the convolution in the frequency domain (real FFTs padded to a 5-smooth length). True or False forces the choice, while None (the default) resolves it per input array with a heuristic that selects FFT for long kernels on CPU and otherwise falls back to the direct convolution (jnp.convolve). At jit-compilation time the device is unavailable, so the default direct path is dispatched; set use_fft=True to force the FFT convolution under jit. When FFT is used and batch_size_samples is smaller than the sample-axis length, the convolution is batched over samples via overlap-save; otherwise the full axis is transformed at once.

Returns:

The convolutional predictor, structured similarly to the input time_series, appropriately padded and shifted according to the specified causality.

Raises:

  • ValueError: – If basis_matrix is not a 2-dimensional array or has a singleton first dimension.

  • ValueError: – If time_series does not contain arrays of at least one dimension or contains arrays with a dimensionality less than axis.

  • ValueError: – If any array within time_series or basis_matrix is empty.

  • ValueError: – If the number of elements along the convolution axis in any array within time_series is less than the window size of the basis_matrix.

  • ValueError: – If shifting is attempted with ‘acausal’ causality.

  • ValueError: – Raised if any explicitly provided batch sizes—batch_size_samples, batch_size_channels, or batch_size_basis—are not positive integers. A value of None is allowed and treated as unspecified.

Notes

Under jax.jit the inputs are tracers with no concrete device, so the automatic backend selection (use_fft=None) cannot inspect device placement and falls back to the direct convolution. If you need the FFT path while tracing, do not rely on the auto-detection: pass use_fft=True explicitly.