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]