qfeval_functions.functions.group_shift

group_shift(x, shift=1, dim=0, mask=None, refdim=-1, agg_f='any')[source]

Shift tensor elements along a dimension, skipping masked positions.

This function performs a selective shift operation where only elements at positions marked as True in the mask are shifted, while positions marked as False are skipped. This is particularly useful for time series data where certain time points should be excluded from the shift operation, such as weekends in financial data or missing observations.

The function works by reordering elements based on the mask, applying the shift only to valid positions, and then restoring the original order. Elements at masked-out positions are replaced with NaN values.

Warning

Unmasked values (where mask is False) are replaced with NaN before shifting. This means that even a zero shift (shift=0) will not return the original input, as unmasked values will be NaN in the output.

Parameters:

x (Tensor) – The input tensor to be shifted.
shift (int) – The number of positions to shift. Positive values shift forward (toward higher indices), negative values shift backward. Default is 1.
dim (int) – The dimension along which to perform the shift. Default is 0.
mask (Optional[Tensor]) – A 1D boolean tensor where True indicates positions to include in the shift, and False indicates positions to skip. Length must equal x.shape[dim]. If not provided, refdim must be specified.
refdim (Optional[int]) – If specified, automatically generates a mask using reduce_nan_patterns(). This identifies valid positions based on non-NaN patterns in the reference dimension. Default is -1.
agg_f (Literal['any', 'all']) – Aggregation function used when generating mask from refdim. Can be “any” or “all”. Default is “any”.

Returns:

A tensor of the same shape as the input, with elements shifted according to the mask. Unmasked positions contain NaN.

Return type:

Tensor

Example

>>> # Basic masked shift
>>> x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])
>>> mask = torch.tensor([True, False, True, False, True])
>>> shifted = QF.group_shift(x, shift=1, dim=0, mask=mask)
>>> shifted
tensor([nan, nan, 1., nan, 3.])

>>> # Shift with existing NaN values
>>> x = torch.tensor([1.0, nan, 2.0, nan, 3.0])
>>> mask = torch.tensor([True, False, True, False, True])
>>> shifted = QF.group_shift(x, shift=1, dim=0, mask=mask)
>>> shifted
tensor([nan, nan, 1., nan, 2.])

>>> # 2D tensor with automatic mask generation
>>> x = torch.tensor([[1.0, 2.0, nan, 4.0],
...                   [5.0, 6.0, nan, 8.0],
...                   [9.0, 10., nan, 12.]])
>>> # Use refdim=0 to generate mask from first row's NaN pattern
>>> shifted = QF.group_shift(x, shift=1, dim=1, refdim=0)
>>> shifted
tensor([[nan,  1., nan,  2.],
        [nan,  5., nan,  6.],
        [nan,  9., nan, 10.]])