def test_hull():
a = ts.Dim(inclusive_min=1, exclusive_max=5, label="x")
b = ts.Dim(size=3)
x = a.hull(b)
assert x.inclusive_min == 0
assert x.exclusive_max == 5
assert x.label == "x"
with pytest.raises(ValueError):
a.hull(ts.Dim(size=3, label="y"))
def test_intersect():
a = ts.Dim(inclusive_min=1, exclusive_max=5, label="x")
b = ts.Dim(size=3)
x = a.intersect(b)
assert x.inclusive_min == 1
assert x.exclusive_max == 3
assert x.label == "x"
with pytest.raises(ValueError):
a.intersect(ts.Dim(size=3, label="y"))
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_half_open():
x = ts.Dim(3, 5)
assert x.inclusive_min == 3
assert x.exclusive_max == 5
y = ts.Dim(inclusive_min=3, exclusive_max=5)
assert x == y
z = ts.Dim(exclusive_max=5)
assert z.inclusive_min == -ts.inf
assert z.exclusive_max == 5
assert z == ts.Dim(inclusive_min=None, exclusive_max=5)
with pytest.raises(ValueError):
ts.Dim(inclusive_min=3, exclusive_max=1)
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_closed():
x = ts.Dim(inclusive_min=3, inclusive_max=5)
assert x.inclusive_min == 3
assert x.inclusive_max == 5
assert x.exclusive_min == 2
assert x.exclusive_max == 6
assert 3 in x
assert 5 in x
assert 2 not in x
assert 6 not in x
assert ts.Dim(inclusive_min=3, inclusive_max=4) in x
assert ts.Dim(inclusive_min=3, inclusive_max=6) not in x
y = ts.Dim(inclusive_min=3, inclusive_max=5)
z = ts.Dim(inclusive_min=3, inclusive_max=6)
assert x == y
assert x != z
assert x.size == 3
assert len(x) == 3
assert x.empty is False
assert repr(x) == "Dim(inclusive_min=3, exclusive_max=6)"
assert str(x) == "[3, 6)"
assert list(x) == [3, 4, 5]
x1 = ts.Dim(inclusive_max=5)
assert x1.inclusive_min == -ts.inf
assert x1.inclusive_max == 5
assert x1 == ts.Dim(inclusive_min=None, inclusive_max=5)
with pytest.raises(ValueError):
ts.Dim(inclusive_min=3, inclusive_max=1)
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_getitem_index():
d = ts.IndexDomain(labels=["x", "y"], shape=[3, 4])
assert d["x"] == ts.Dim(label="x", size=3)
assert d["y"] == ts.Dim(label="y", size=4)
assert d[1] == ts.Dim(label="y", size=4)
assert d["x", "y"] == d
assert d["y", "x"] == [d[1], d[0]]
assert d[::-1] == [d[1], d[0]]
with pytest.raises(IndexError):
d["z"]
with pytest.raises(IndexError):
d[2]
with pytest.raises(ValueError):
d[1:3]
def test_sized():
x = ts.Dim(inclusive_min=3, size=10)
assert x.inclusive_min == 3
assert x.size == 10
x1 = ts.Dim(size=10)
assert x1.inclusive_min == 0
assert x1.size == 10
assert x1 == ts.Dim(inclusive_min=None, size=10)
y = ts.Dim(size=None)
assert y.inclusive_min == 0
assert y.inclusive_max == +ts.inf
with pytest.raises(ValueError):
ts.Dim(inclusive_min=3, size=-3)
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_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_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_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_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_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_unbounded():
x = ts.Dim()
assert x.inclusive_min == -ts.inf
assert x.inclusive_max == +ts.inf
assert not x.finite
def test_pickle():
x = ts.Dim(inclusive_min=3, size=10)
assert pickle.loads(pickle.dumps(x)) == x
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]
