def test_tensorgrid_insert(): vec1 = np.array([2, 3, 4, 5]) vec2 = np.array([-4, -2, 0, 2, 4]) vec3 = np.array([-1, 0]) scalar = 0.5 grid = TensorGrid(vec1, vec2) grid2 = TensorGrid(scalar, vec3) # Test all positions ins_grid = grid.insert(0, grid2) assert ins_grid == TensorGrid(scalar, vec3, vec1, vec2) ins_grid = grid.insert(1, grid2) assert ins_grid == TensorGrid(vec1, scalar, vec3, vec2) ins_grid = grid.insert(2, grid2) assert ins_grid == TensorGrid(vec1, vec2, scalar, vec3) ins_grid = grid.insert(-1, grid2) assert ins_grid == TensorGrid(vec1, scalar, vec3, vec2) with pytest.raises(IndexError): grid.insert(3, grid2) with pytest.raises(IndexError): grid.insert(-4, grid2)