def test_output_index_maps_lifetime():
output = ts.IndexTransform(3).output
assert output == [
ts.OutputIndexMap(input_dimension=0),
ts.OutputIndexMap(input_dimension=1),
ts.OutputIndexMap(input_dimension=2),
]
def test_json():
json = {
'input_inclusive_min': [1, 2, 3],
'input_exclusive_max': [4, 5, 6],
'input_labels': ['x', 'y', 'z'],
'output': [
{
'offset': 3
},
{
'input_dimension': 0,
'stride': 2
},
],
}
x = ts.IndexTransform(json=json)
expected = ts.IndexTransform(
input_inclusive_min=[1, 2, 3],
input_exclusive_max=[4, 5, 6],
input_labels=['x', 'y', 'z'],
output=[
ts.OutputIndexMap(offset=3),
ts.OutputIndexMap(stride=2, input_dimension=0),
],
)
assert x == expected
assert x.to_json() == json
def test_dimension_selection_index_zero_rank_bool():
x = ts.IndexTransform(input_rank=2)
assert x[ts.d[:][..., True]] == ts.IndexTransform(
domain=[ts.Dim(size=1), *x.domain],
output=[
ts.OutputIndexMap(input_dimension=1),
ts.OutputIndexMap(input_dimension=2),
],
)
def test_index_integer():
x = ts.IndexTransform(input_shape=[15, 20], input_labels=["x", "y"])
expr = ts.d["x", "y"][2, 3]
assert repr(expr) == "d['x','y'][2,3]"
assert x[expr] == ts.IndexTransform(
input_rank=0,
output=[ts.OutputIndexMap(offset=2),
ts.OutputIndexMap(offset=3)],
)
def test_init_output_index_maps():
x = ts.IndexTransform(
input_shape=[3, 2],
output=[
ts.OutputIndexMap(offset=7, input_dimension=1),
ts.OutputIndexMap([[1, 2]], offset=2, stride=-1),
],
)
y = ts.IndexTransform(x.domain, x.output)
assert x == y
def test_boolean_array():
x = ts.IndexTransform(input_rank=3)
assert x[1, [True, False, True, True]] == ts.IndexTransform(
domain=[ts.Dim(size=3), ts.Dim()],
output=[
ts.OutputIndexMap(1),
ts.OutputIndexMap(index_array=[[0], [2], [3]]),
ts.OutputIndexMap(input_dimension=1),
],
)
assert x[[[True, False, False],
[False, False, True]]] == ts.IndexTransform(
domain=[ts.Dim(size=2), ts.Dim()],
output=[
ts.OutputIndexMap(index_array=[[0], [1]]),
ts.OutputIndexMap(index_array=[[0], [2]]),
ts.OutputIndexMap(input_dimension=1),
],
)
assert x[True] == ts.IndexTransform(
domain=[ts.Dim(size=1), ts.Dim(),
ts.Dim(), ts.Dim()],
output=[
ts.OutputIndexMap(input_dimension=1),
ts.OutputIndexMap(input_dimension=2),
ts.OutputIndexMap(input_dimension=3),
],
)
def test_diagonal():
x = ts.IndexTransform(input_shape=[2, 3], input_labels=["x", "y"])
expr = ts.d["x", "y"].diagonal
assert repr(expr) == "d['x','y'].diagonal"
assert x[expr] == ts.IndexTransform(
input_shape=[2],
output=[
ts.OutputIndexMap(input_dimension=0),
ts.OutputIndexMap(input_dimension=0),
],
)
def test_index_integer_non_scalar():
x = ts.IndexTransform(input_shape=[15, 20], input_labels=["x", "y"])
expr = ts.d["x"][2, ]
assert repr(expr) == "d['x'][2,]"
assert x[expr] == ts.IndexTransform(
domain=x.domain["y", ],
output=[
ts.OutputIndexMap(offset=2),
ts.OutputIndexMap(input_dimension=0)
],
)
def test_dimension_selection_dimrange_index():
x = ts.IndexTransform(input_labels=["a", "b", "c"])
assert x[ts.d["c", :2][:5, 1, 2]] == ts.IndexTransform(
input_labels=["c"],
input_exclusive_max=[5],
output=[
ts.OutputIndexMap(1),
ts.OutputIndexMap(2),
ts.OutputIndexMap(input_dimension=0),
],
)
def test_transpose_move_to_back():
x = ts.IndexTransform(input_labels=["x", "y", "z"])
expr = ts.d["y", "x"].transpose[-1]
assert repr(expr) == "d['y','x'].transpose[-1]"
assert x[expr] == ts.IndexTransform(
input_labels=["z", "y", "x"],
output=[
ts.OutputIndexMap(input_dimension=2),
ts.OutputIndexMap(input_dimension=1),
ts.OutputIndexMap(input_dimension=0),
],
)
def test_transpose_dim_range():
x = ts.IndexTransform(input_shape=[2, 3], input_labels=["x", "y"])
expr = ts.d["y", "x"].transpose[:]
assert repr(expr) == "d['y','x'].transpose[:]"
assert x[expr] == ts.IndexTransform(
input_shape=[3, 2],
input_labels=["y", "x"],
output=[
ts.OutputIndexMap(input_dimension=1),
ts.OutputIndexMap(input_dimension=0),
],
)
def test_slice_interval_strided():
x = ts.IndexTransform(input_shape=[15, 20], input_labels=["x", "y"])
expr = ts.d["x", "y"][(1, 2):(8, 9):2]
assert repr(expr) == "d['x','y'][1:8:2,2:9:2]"
assert x[expr] == ts.IndexTransform(
input_inclusive_min=[0, 1],
input_inclusive_max=[3, 4],
input_labels=["x", "y"],
output=[
ts.OutputIndexMap(input_dimension=0, offset=1, stride=2),
ts.OutputIndexMap(input_dimension=1, stride=2),
],
)
def test_dimension_selection_vindex_bool_arrays_consecutive():
x = ts.IndexTransform(input_labels=["a", "b", "c", "d"])
expr = ts.d["a", "b", "c",
"d"].vindex[:, [[False, True, True], [False, False, False]],
[False, False, True, True]]
assert x[expr] == ts.IndexTransform(
domain=[ts.Dim(size=2), ts.Dim(label="a")],
output=[
ts.OutputIndexMap(input_dimension=1),
ts.OutputIndexMap(index_array=[[0], [0]]),
ts.OutputIndexMap(index_array=[[1], [2]]),
ts.OutputIndexMap(index_array=[[2], [3]]),
],
)
def test_translate_by_scalar():
x = ts.IndexTransform(input_shape=[2, 3], input_labels=["x", "y"])
expr = ts.d["x", "y"].translate_by[4]
assert repr(expr) == "d['x','y'].translate_by[4]"
assert x[expr] == ts.IndexTransform(
input_shape=[2, 3],
input_inclusive_min=[4, 4],
input_labels=["x", "y"],
output=[
ts.OutputIndexMap(offset=-4, input_dimension=0),
ts.OutputIndexMap(offset=-4, input_dimension=1),
],
)
def test_add_new():
x = ts.IndexTransform(input_shape=[2, 3], input_labels=["x", "y"])
expr = ts.d[0, -2:][ts.newaxis]
assert repr(expr) == "d[0,-2:][None]"
assert x[expr] == ts.IndexTransform(
input_inclusive_min=[0, 0, 0, 0, 0],
input_exclusive_max=[1, 2, 3, 1, 1],
input_labels=["", "x", "y", "", ""],
implicit_lower_bounds=[1, 0, 0, 1, 1],
implicit_upper_bounds=[1, 0, 0, 1, 1],
output=[
ts.OutputIndexMap(input_dimension=1),
ts.OutputIndexMap(input_dimension=2),
],
)
def test_dimension_selection_oindex_bool_arrays():
x = ts.IndexTransform(input_labels=["a", "b", "c", "d"])
expr = ts.d["a", "b", "c",
"d"].oindex[[[False, True, True], [False, False, False]], :,
[True, False, True, True]]
assert x[expr] == ts.IndexTransform(
domain=[ts.Dim(size=2),
ts.Dim(label="c"),
ts.Dim(size=3)],
output=[
ts.OutputIndexMap(index_array=[[[0]], [[0]]]),
ts.OutputIndexMap(index_array=[[[1]], [[2]]]),
ts.OutputIndexMap(input_dimension=1),
ts.OutputIndexMap(index_array=[[[0, 2, 3]]]),
],
)
def test_newaxis():
x = ts.IndexTransform(input_rank=3)
assert x[np.newaxis, ..., np.newaxis] == ts.IndexTransform(
domain=[
ts.Dim(size=1, implicit_lower=True, implicit_upper=True),
ts.Dim(),
ts.Dim(),
ts.Dim(),
ts.Dim(size=1, implicit_lower=True, implicit_upper=True),
],
output=[
ts.OutputIndexMap(input_dimension=1),
ts.OutputIndexMap(input_dimension=2),
ts.OutputIndexMap(input_dimension=3),
],
)
def test_stride_vector():
x = ts.IndexTransform(input_inclusive_min=[0, 2, 1],
input_inclusive_max=[6, 5, 8],
input_labels=["x", "y", "z"])
expr = ts.d["x", "z"].stride[-2, 3]
assert repr(expr) == "d['x','z'].stride[-2,3]"
assert x[expr] == ts.IndexTransform(
input_inclusive_min=[-3, 2, 1],
input_inclusive_max=[0, 5, 2],
input_labels=["x", "y", "z"],
output=[
ts.OutputIndexMap(stride=-2, input_dimension=0),
ts.OutputIndexMap(stride=1, input_dimension=1),
ts.OutputIndexMap(stride=3, input_dimension=2),
],
)
def test_dimension_selection_index_arrays_consecutive():
x = ts.IndexTransform(input_labels=["a", "b", "c", "d"])
expr = ts.d["a", "c", "d"][..., [[1, 2, 3], [4, 5, 6]], [6, 7, 8]]
assert x[expr] == ts.IndexTransform(
domain=[
ts.Dim(label="a"),
ts.Dim(label="b"),
ts.Dim(size=2),
ts.Dim(size=3),
],
output=[
ts.OutputIndexMap(input_dimension=0),
ts.OutputIndexMap(input_dimension=1),
ts.OutputIndexMap(index_array=[[[[1, 2, 3], [4, 5, 6]]]]),
ts.OutputIndexMap(index_array=[[[[6, 7, 8]]]]),
],
)
def test_dimension_selection_vindex_index_arrays_non_consecutive():
x = ts.IndexTransform(input_labels=["a", "b", "c", "d"])
expr = ts.d["a", "c", "d"].vindex[[[1, 2, 3], [4, 5, 6]], :, [6, 7, 8]]
assert x[expr] == ts.IndexTransform(
domain=[
ts.Dim(size=2),
ts.Dim(size=3),
ts.Dim(label="b"),
ts.Dim(label="c"),
],
output=[
ts.OutputIndexMap(
index_array=[[[[1]], [[2]], [[3]]], [[[4]], [[5]], [[6]]]]),
ts.OutputIndexMap(input_dimension=2),
ts.OutputIndexMap(input_dimension=3),
ts.OutputIndexMap(index_array=[[[[6]], [[7]], [[8]]]]),
],
)
def test_pickle():
x = ts.IndexTransform(
input_inclusive_min=[1, 2, -1],
implicit_lower_bounds=[1, 0, 0],
input_shape=[3, 2, 4],
implicit_upper_bounds=[0, 1, 0],
input_labels=['x', 'y', 'z'],
output=[
ts.OutputIndexMap(offset=7, stride=13, input_dimension=1),
ts.OutputIndexMap(offset=8),
ts.OutputIndexMap(
offset=1,
stride=-2,
index_array=[[[-10, 1, 2, 20]]],
index_range=ts.Dim(inclusive_min=-3, exclusive_max=10),
),
],
)
assert pickle.loads(pickle.dumps(x)) == x
def test_identity():
x = ts.IndexTransform(input_rank=3)
assert x.input_rank == 3
assert x.output_rank == 3
np.testing.assert_equal(x.input_inclusive_min, [-ts.inf] * 3)
np.testing.assert_equal(x.input_inclusive_max, [+ts.inf] * 3)
np.testing.assert_equal(x.input_exclusive_max, [+ts.inf + 1] * 3)
np.testing.assert_equal(x.input_shape, [2 * ts.inf + 1] * 3)
assert x.output == [ts.OutputIndexMap(input_dimension=i) for i in range(3)]
assert list(
x.output) == [ts.OutputIndexMap(input_dimension=i) for i in range(3)]
np.testing.assert_equal(x([1, 2, 3]), [1, 2, 3])
with pytest.raises(ValueError):
x([1])
with pytest.raises(ValueError):
x([1, 2, 3, 4])
x = ts.IndexTransform(input_labels=['x', 'y'])
assert x.input_labels == ('x', 'y')
assert x.output == [ts.OutputIndexMap(input_dimension=i) for i in range(2)]
def test_domain_access():
x = ts.IndexTransform(input_inclusive_min=[1, 2, 3], input_shape=[5, 6, 7])
assert x.origin == (1, 2, 3)
assert x.shape == (5, 6, 7)
assert x.size == 5 * 6 * 7
assert x.T == ts.IndexTransform(
input_inclusive_min=[3, 2, 1],
input_shape=[7, 6, 5],
output=[ts.OutputIndexMap(input_dimension=i) for i in [2, 1, 0]],
)
assert x.T == x[ts.d[::-1].transpose[:]]
def test_identity():
x = ts.IndexTransform(input_rank=3)
assert x.ndim == 3
assert x.input_rank == 3
assert x.output_rank == 3
assert x.input_inclusive_min == (-ts.inf,) * 3
assert x.input_inclusive_max == (+ts.inf,) * 3
assert x.input_exclusive_max == (+ts.inf + 1,) * 3
assert x.input_shape == (2 * ts.inf + 1,) * 3
assert x.output == [ts.OutputIndexMap(input_dimension=i) for i in range(3)]
assert list(
x.output) == [ts.OutputIndexMap(input_dimension=i) for i in range(3)]
assert x([1, 2, 3]) == (1, 2, 3)
with pytest.raises(ValueError):
x([1])
with pytest.raises(ValueError):
x([1, 2, 3, 4])
x = ts.IndexTransform(input_labels=['x', 'y'])
assert x.input_labels == ('x', 'y')
assert x.output == [ts.OutputIndexMap(input_dimension=i) for i in range(2)]
def test_slice():
x = ts.IndexTransform(input_rank=3)
assert x[:, 1] == ts.IndexTransform(
input_rank=2,
output=[
ts.OutputIndexMap(input_dimension=0),
ts.OutputIndexMap(1),
ts.OutputIndexMap(input_dimension=1),
],
)
assert x[:5] == ts.IndexTransform(
domain=[ts.Dim(exclusive_max=5),
ts.Dim(), ts.Dim()])
assert x[2:5] == ts.IndexTransform(
domain=[ts.Dim(inclusive_min=2, size=3),
ts.Dim(), ts.Dim()])
assert x[10:1:-2] == ts.IndexTransform(
domain=[ts.Dim(inclusive_min=-5, size=5),
ts.Dim(), ts.Dim()],
output=[
ts.OutputIndexMap(stride=-2, input_dimension=0),
ts.OutputIndexMap(input_dimension=1),
ts.OutputIndexMap(input_dimension=2),
],
)
y = ts.IndexTransform(input_shape=[5, 10])
with pytest.raises(IndexError,
match="Computing interval slice for dimension 0: .*"):
y[1:6]
def test_integer():
x = ts.IndexTransform(input_rank=3)
assert x[1] == ts.IndexTransform(
input_rank=2,
output=[
ts.OutputIndexMap(1),
ts.OutputIndexMap(input_dimension=0),
ts.OutputIndexMap(input_dimension=1),
],
)
assert x[..., 1] == ts.IndexTransform(
input_rank=2,
output=[
ts.OutputIndexMap(input_dimension=0),
ts.OutputIndexMap(input_dimension=1),
ts.OutputIndexMap(1),
],
)
assert x[1, 2, 3] == ts.IndexTransform(
input_rank=0,
output=[
ts.OutputIndexMap(1),
ts.OutputIndexMap(2),
ts.OutputIndexMap(3)
],
)
with pytest.raises(
IndexError,
match=re.escape(
"An index can only have a single ellipsis (`...`)"),
):
x[..., 1, ...]
with pytest.raises(
IndexError,
match=
"Indexing expression requires 4 dimensions, and cannot be applied to a domain of rank 3"
):
x[1, 2, 3, 4]
def test_init_output():
x = ts.IndexTransform(
input_shape=[3, 2],
output=[
ts.OutputIndexMap(offset=7, input_dimension=1),
ts.OutputIndexMap([[1, 2]], offset=2, stride=-1),
ts.OutputIndexMap(8),
ts.OutputIndexMap([[1, 2]], offset=2, stride=-1,
index_range=ts.Dim(inclusive_min=0,
exclusive_max=8)),
],
)
assert x.output[3].index_range == ts.Dim(inclusive_min=0, exclusive_max=8)
assert x.output == [
ts.OutputIndexMap(offset=7, input_dimension=1),
ts.OutputIndexMap([[1, 2]], offset=2, stride=-1),
ts.OutputIndexMap(8),
ts.OutputIndexMap([[1, 2]], offset=2, stride=-1,
index_range=ts.Dim(inclusive_min=0, exclusive_max=8)),
]
def test_init_output():
x = ts.IndexTransform(
input_shape=[3, 2],
output=[
ts.OutputIndexMap(offset=7, input_dimension=1),
ts.OutputIndexMap([[1, 2]], offset=2, stride=-1),
ts.OutputIndexMap(8),
],
)
assert x.output == [
ts.OutputIndexMap(offset=7, input_dimension=1),
ts.OutputIndexMap([[1, 2]], offset=2, stride=-1),
ts.OutputIndexMap(8),
]
def test_integer_array():
x = ts.IndexTransform(input_rank=3)
assert x[1, [1, 2, 3]] == ts.IndexTransform(
domain=[ts.Dim(size=3), ts.Dim()],
output=[
ts.OutputIndexMap(1),
ts.OutputIndexMap(index_array=[[1], [2], [3]]),
ts.OutputIndexMap(input_dimension=1),
],
)
assert x[[[5, 6, 7], [8, 9, 10]], [1, 2, 3]] == ts.IndexTransform(
domain=[ts.Dim(size=2), ts.Dim(size=3),
ts.Dim()],
output=[
ts.OutputIndexMap(index_array=[[[5], [6], [7]], [[8], [9], [10]]]),
ts.OutputIndexMap(index_array=[[[1], [2], [3]]]),
ts.OutputIndexMap(input_dimension=2),
],
)
assert x[[[5], [8]], :, [1, 2, 3]] == ts.IndexTransform(
domain=[ts.Dim(size=2), ts.Dim(size=3),
ts.Dim()],
output=[
ts.OutputIndexMap(index_array=[[[5]], [[8]]]),
ts.OutputIndexMap(input_dimension=2),
ts.OutputIndexMap(index_array=[[[1], [2], [3]]]),
],
)
# Test using a zero-size index array.
assert x[[]] == ts.IndexTransform(
domain=[ts.Dim(size=0), ts.Dim(), ts.Dim()],
output=[
ts.OutputIndexMap(0),
ts.OutputIndexMap(input_dimension=1),
ts.OutputIndexMap(input_dimension=2),
],
)
with pytest.raises(
IndexError,
match=re.escape("Incompatible index array shapes: {3} vs {4}")):
x[[1, 2, 3, 4], [1, 2, 3]]
for indices in [1.5, [1.5, 2.5], "x", {"a": 3}, [None]]:
with pytest.raises(
IndexError,
match=re.escape(
"Only integers, slices (`:`), ellipsis (`...`), "
"tensorstore.newaxis (`None`) and "
"integer or boolean arrays are valid indices"),
):
x[indices]
for indices in [np.array([1.1, 1.5]), np.array(["x", "y"])]:
with pytest.raises(
IndexError,
match=re.escape(
"Arrays used as indices must be of integer (or boolean) type"
),
):
x[indices]
