interp_torch.py

# Source: https://github.com/aliutkus/torchinterp1d
#
# Modification: Remove backward function
#

import torch


def interp(xnew, x, y):
    """
    Linear 1D interpolation on the GPU for Pytorch.
    This function returns interpolated values of a set of 1-D functions at
    the desired query points `xnew`.
    This function is working similarly to Matlab™ or scipy functions with
    the `linear` interpolation mode on, except that it parallelises over
    any number of desired interpolation problems.
    The code will run on GPU if all the tensors provided are on a cuda
    device.
    Parameters
    ----------
    x : (N, ) or (D, N) Pytorch Tensor
        A 1-D or 2-D tensor of real values.
    y : (N,) or (D, N) Pytorch Tensor
        A 1-D or 2-D tensor of real values. The length of `y` along its
        last dimension must be the same as that of `x`
    xnew : (P,) or (D, P) Pytorch Tensor
        A 1-D or 2-D tensor of real values. `xnew` can only be 1-D if
        _both_ `x` and `y` are 1-D. Otherwise, its length along the first
        dimension must be the same as that of whichever `x` and `y` is 2-D.
    """
    # making the vectors at least 2D
    is_flat = {}
    require_grad = {}
    v = {}
    device = []
    eps = torch.finfo(y.dtype).eps
    for name, vec in {'x': x, 'y': y, 'xnew': xnew}.items():
        assert len(vec.shape) <= 2, 'interp1d: all inputs must be '\
                                    'at most 2-D.'
        if len(vec.shape) == 1:
            v[name] = vec[None, :]
        else:
            v[name] = vec
        is_flat[name] = v[name].shape[0] == 1
        require_grad[name] = vec.requires_grad
        device = list(set(device + [str(vec.device)]))
    assert len(device) == 1, 'All parameters must be on the same device.'
    device = device[0]

    # Checking for the dimensions
    assert (v['x'].shape[1] == v['y'].shape[1]
            and (
                    v['x'].shape[0] == v['y'].shape[0]
                    or v['x'].shape[0] == 1
                    or v['y'].shape[0] == 1
                )
            ), ("x and y must have the same number of columns, and either "
                "the same number of row or one of them having only one "
                "row.")

    reshaped_xnew = False
    if ((v['x'].shape[0] == 1) and (v['y'].shape[0] == 1)
        and (v['xnew'].shape[0] > 1)):
        # if there is only one row for both x and y, there is no need to
        # loop over the rows of xnew because they will all have to face the
        # same interpolation problem. We should just stack them together to
        # call interp1d and put them back in place afterwards.
        original_xnew_shape = v['xnew'].shape
        v['xnew'] = v['xnew'].contiguous().view(1, -1)
        reshaped_xnew = True

    # identify the dimensions of output and check if the one provided is ok
    D = max(v['x'].shape[0], v['xnew'].shape[0])
    shape_ynew = (D, v['xnew'].shape[-1])

    ynew = torch.zeros(*shape_ynew, device=device)

    # moving everything to the desired device in case it was not there
    # already (not handling the case things do not fit entirely, user will
    # do it if required.)
    for name in v:
        v[name] = v[name].to(device)

    # calling searchsorted on the x values.
    ind = ynew.long()

    # expanding xnew to match the number of rows of x in case only one xnew is
    # provided
    if v['xnew'].shape[0] == 1:
        v['xnew'] = v['xnew'].expand(v['x'].shape[0], -1)

    torch.searchsorted(v['x'].contiguous(),
                        v['xnew'].contiguous(), out=ind)

    # the `-1` is because searchsorted looks for the index where the values
    # must be inserted to preserve order. And we want the index of the
    # preceeding value.
    ind -= 1
    # we clamp the index, because the number of intervals is x.shape-1,
    # and the left neighbour should hence be at most number of intervals
    # -1, i.e. number of columns in x -2
    ind = torch.clamp(ind, 0, v['x'].shape[1] - 1 - 1)

    # helper function to select stuff according to the found indices.
    def sel(name):
        if is_flat[name]:
            return v[name].contiguous().view(-1)[ind]
        return torch.gather(v[name], 1, ind)

    # activating gradient storing for everything now
    enable_grad = False
    saved_inputs = []
    for name in ['x', 'y', 'xnew']:
        if require_grad[name]:
            enable_grad = True
            saved_inputs += [v[name]]
        else:
            saved_inputs += [None, ]
    # assuming x are sorted in the dimension 1, computing the slopes for
    # the segments
    is_flat['slopes'] = is_flat['x']
    # now we have found the indices of the neighbors, we start building the
    # output. Hence, we start also activating gradient tracking
    v['slopes'] = (
            (v['y'][:, 1:]-v['y'][:, :-1])
            /
            (eps + (v['x'][:, 1:]-v['x'][:, :-1]))
        )

    # now build the linear interpolation
    ynew = sel('y') + sel('slopes') * (v['xnew'] - sel('x'))

    if reshaped_xnew:
        ynew = ynew.view(original_xnew_shape)

    return ynew