File size: 4,021 Bytes
8378d33 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
import numpy as np
import torch
move = np.arange(1, 8)
diag = np.array([
move + move*8,
move - move*8,
move*-1 - move*8,
move*-1 + move*8
])
orthog = np.array([
move,
move*-8,
move*-1,
move*8
])
knight = np.array([
[2 + 1*8],
[2 - 1*8],
[1 - 2*8],
[-1 - 2*8],
[-2 - 1*8],
[-2 + 1*8],
[-1 + 2*8],
[1 + 2*8]
])
promos = np.array([2*8, 3*8, 4*8])
pawn_promotion = np.array([
-1 + promos,
0 + promos,
1 + promos
])
def make_map():
"""theoretically possible put-down squares (numpy array) for each pick-up square (list element).
squares are [0, 1, ..., 63] for [a1, b1, ..., h8]. squares after 63 are for promotion squares.
each successive "row" beyond 63 (ie. 64:72, 72:80, 80:88) are for over-promotions to queen, rook, and bishop;
respectively. a pawn traverse to row 56:64 signifies a "default" promotion to a knight."""
traversable = []
for i in range(8):
for j in range(8):
sq = (8*i + j)
traversable.append(
sq +
np.sort(
np.int32(
np.concatenate((
orthog[0][:7-j], orthog[2][:j], orthog[1][:i], orthog[3][:7-i],
diag[0][:np.min((7-i, 7-j))], diag[3][:np.min((7-i, j))],
diag[1][:np.min((i, 7-j))], diag[2][:np.min((i, j))],
knight[0] if i < 7 and j < 6 else [], knight[1] if i > 0 and j < 6 else [],
knight[2] if i > 1 and j < 7 else [], knight[3] if i > 1 and j > 0 else [],
knight[4] if i > 0 and j > 1 else [], knight[5] if i < 7 and j > 1 else [],
knight[6] if i < 6 and j > 0 else [], knight[7] if i < 6 and j < 7 else [],
pawn_promotion[0] if i == 6 and j > 0 else [],
pawn_promotion[1] if i == 6 else [],
pawn_promotion[2] if i == 6 and j < 7 else [],
))
)
)
)
z = np.zeros((64*64+8*24, 1858), dtype=np.int32)
apm_out = np.zeros((1858,), dtype=np.int32)
apm_in = np.zeros((64*64+8*24), dtype=np.int32)
# first loop for standard moves (for i in 0:1858, stride by 1)
i = 0
for pickup_index, putdown_indices in enumerate(traversable):
for putdown_index in putdown_indices:
if putdown_index < 64:
du_idx = putdown_index + (64*pickup_index)
z[du_idx, i] = 1
apm_out[i] = du_idx
apm_in[du_idx] = i
i += 1
# second loop for promotions (for i in 1792:1858, stride by ls[j])
j = 0
j1 = np.array([3, -2, 3, -2, 3])
j2 = np.array([3, 3, -5, 3, 3, -5, 3, 3, 1])
ls = np.append(j1, 1)
for k in range(6):
ls = np.append(ls, j2)
ls = np.append(ls, j1)
ls = np.append(ls, 0)
for pickup_index, putdown_indices in enumerate(traversable):
for putdown_index in putdown_indices:
if putdown_index >= 64:
pickup_file = pickup_index % 8
promotion_file = putdown_index % 8
promotion_rank = (putdown_index // 8) - 8
du_idx = 4096 + pickup_file*24 + (promotion_file*3+promotion_rank)
z[du_idx, i] = 1
apm_out[i] = du_idx
apm_in[du_idx] = i
i += ls[j]
j += 1
return z, apm_out, apm_in
apm_map, apm_out, apm_in = make_map()
def set_zero_sum(x):
x = x + (1 - torch.sum(x, dim=1, keepdim=True)) * (1.0 / 64)
return x
def get_up_down(moves):
apm_map_tensor = torch.from_numpy(apm_map)
out = torch.matmul(moves, apm_map_tensor.T.float())
out = out[..., :64*64]
out = out.view(-1, 64, 64)
pu = set_zero_sum(torch.sum(out, dim=-1))
pd = set_zero_sum(torch.sum(out, dim=-2))
return pu, pd
|