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]