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_label_multiple():
x = ts.IndexTransform(input_shape=[2, 3], input_labels=["x", "y"])
expr = ts.d["x", "y"].label["a", "b"]
assert repr(expr) == "d['x','y'].label['a','b']"
assert x[expr] == ts.IndexTransform(input_shape=[2, 3],
input_labels=["a", "b"])
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_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_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_label_index_slice():
x = ts.IndexTransform(input_labels=["x", "y"])
expr = ts.d["x", "y"].label["a", "b"][(1, 2):(7, 8)]
assert repr(expr) == "d['x','y'].label['a','b'][1:7,2:8]"
assert x[expr] == ts.IndexTransform(
input_inclusive_min=[1, 2],
input_exclusive_max=[7, 8],
input_labels=["a", "b"],
)
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_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_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_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_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_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]]]),
],
)
async def cb(index):
out = np.zeros(new_shard_shape, dtype=t.dtype.numpy_dtype)
requested_domain = ts.IndexTransform(input_shape=shape)[index].domain
restricted_domain = t.domain.intersect(requested_domain)
await ts.array(out)[ts.d[:].translate_to[requested_domain.origin]
][restricted_domain].write(t[restricted_domain])
return out
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_dimension_selection_duplicate_index_label_newaxis():
x = ts.IndexTransform(input_labels=["x", "y", "z"])
expr = ts.d[0, 2, "y"][ts.newaxis, ...]
with pytest.raises(
IndexError,
match=re.escape("Dimension 2 specified more than once")):
x[expr]
def test_dimension_selection_duplicate_index():
x = ts.IndexTransform(input_rank=3)
expr = ts.d[1, 1][...]
with pytest.raises(
IndexError,
match=re.escape("Dimension 1 specified more than once")):
x[expr]
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_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_no_operations():
x = ts.IndexTransform(input_rank=3)
expr = ts.d[0, 1]
with pytest.raises(
IndexError,
match="Must specify at least one operation in dimension expression"
):
x[expr] # pylint: disable=pointless-statement
def test_index_too_many_ops():
x = ts.IndexTransform(input_rank=2)
with pytest.raises(
IndexError,
match=re.escape(
"Indexing expression requires 3 dimensions, and cannot be applied to a domain of rank 2"
)):
x[1, 2, 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_oindex_zero_rank_bool():
x = ts.IndexTransform(input_rank=2)
with pytest.raises(
IndexError,
match=re.escape(
"Zero-rank bool array incompatible with outer indexing of a dimension selection"
),
):
x[ts.d[:].oindex[..., True]]
def test_dimension_selection_index_label_newaxis():
x = ts.IndexTransform(input_labels=["x", "y", "z"])
expr = ts.d[0, "y"][ts.newaxis, ts.newaxis]
with pytest.raises(
IndexError,
match=re.escape(
"Dimensions specified by label cannot be used with newaxis"),
):
x[expr]
def test_dimension_selection_dimrange_index_invalid_start():
x = ts.IndexTransform(input_rank=3)
expr = ts.d[5:][...]
with pytest.raises(
IndexError,
match=re.escape(
"Dimension index 5 is outside valid range [-3, 3)."),
):
x[expr]
