GitLab

Would be good to get a subtensor, i.e. slice of the tensor (a subtensor from dimensions (l_1,...,l_n) to dimensions (u_1,...,u_n) where 0 <= l_i <= u_i <= tensor_dim_i, see code below.)

def slice_tt(tt, lower, upper):
    """
    Slice TTTensor in each component from lower to (not including) upper dimension.
    """
    assert len(lower) == len(upper) == len(tt.dimensions)
    assert np.all(np.array(lower) < np.array(upper))
    assert np.all(np.array(upper) <= np.array(tt.dimensions))
    diff = [u-l for l,u in zip(lower,upper)]

    # TODO there should be a more elegant way to slice TTTensors!
    tmp = xe.TTTensor(tt)
    for pos in range(tt.order()):
        tmp.move_core(pos)
        cmp = np.asarray(tmp.get_component(pos))
        cmp[:, :lower[pos], :] = 0
        cmp[:, upper[pos]:, :] = 0
        tmp.set_component(pos, xe.Tensor.from_buffer(cmp))
    tmp.move_core(0)

    tt_slice = xe.TTTensor.random(diff, tmp.ranks())
    for pos in range(tt.order()):
        cmp = np.asarray(tmp.get_component(pos))
        cmp = np.array(cmp[:, lower[pos]:upper[pos], :])
        tt_slice.set_component(pos, xe.Tensor.from_buffer(cmp))

    return tt_slice