def self_atari_size(state: State,
order=Order.TH,
nb_planes=8,
dtype=np.float32) -> np.ndarray:
if order == Order.TH:
planes = np.zeros((nb_planes, state.board_size, state.board_size),
dtype=dtype)
else:
planes = np.zeros((state.board_size, state.board_size, nb_planes),
dtype=dtype)
for x, y in state.legal_moves():
liberty_set_after = set(state.liberty_sets[(x, y)])
group_after_move = {(x, y)}
captured_stones = set()
for neighbors in state.groups_around_at((x, y)):
neighbor = next(
iter(neighbors)) # get an arbitrary stone in 'neighbors'
if state.board[neighbor] == state.current_player:
# if the group is mine, take them into account
liberty_set_after |= state.liberty_sets[neighbor]
group_after_move |= state.groups[neighbor]
elif state.liberty_counts[neighbor] == 1:
# if the group is enemy's and the group would be captured
# (neighbor cannot be Color.EMPTY because neighbors has one liberty)
captured_stones |= state.groups[neighbor]
if captured_stones:
for stone in group_after_move:
# if there are some groups that can be captured,
# the coordinates in which captured group was in can be new liberties of mine
liberty_set_after |= set(
state.crosswise_neighbors_of(stone)) & captured_stones
if (x, y) in liberty_set_after:
liberty_set_after.remove((x, y))
if len(liberty_set_after) == 1:
if order == Order.TH:
planes[min(nb_planes - 1, len(group_after_move) - 1), x, y] = 1
else:
planes[x, y, min(nb_planes - 1, len(group_after_move) - 1)] = 1
return planes
def liberties_after_move(state: State,
order=Order.TH,
nb_planes=8,
dtype=np.float32) -> np.ndarray:
if order == Order.TH:
planes = np.zeros((nb_planes, state.board_size, state.board_size),
dtype=dtype)
else:
planes = np.zeros((state.board_size, state.board_size, nb_planes),
dtype=dtype)
for x, y in state.legal_moves():
liberty_set_after = set(state.liberty_sets[(x, y)])
group_after = {(x, y)}
captured_stones = set()
for neighbors in state.groups_around_at((x, y)):
neighbor = next(iter(neighbors))
if state.board[neighbor] == state.current_player:
liberty_set_after |= state.liberty_sets[neighbor]
group_after |= state.groups[neighbor]
elif state.liberty_counts[neighbor] == 1:
captured_stones |= state.groups[neighbor]
if captured_stones:
for stone in group_after:
liberty_set_after |= set(
state.crosswise_neighbors_of(stone)) & captured_stones
if (x, y) in liberty_set_after:
liberty_set_after.remove((x, y))
if order == Order.TH:
planes[min(nb_planes - 1, len(liberty_set_after) - 1), x, y] = 1
else:
planes[x, y, min(nb_planes - 1, len(liberty_set_after) - 1)] = 1
return planes