def __init__(self,
prob_mat: np.ndarray,
pred_mat: np.ndarray,
reps: int = 1) -> None:
"""
Construct extended nonlocal game object.
:param prob_mat: A matrix whose (x, y)-entry gives the probability
that the referee will give Alice the value `x` and Bob
the value `y`.
:param pred_mat:
:param reps: Number of parallel repetitions to perform.
"""
if reps == 1:
self.prob_mat = prob_mat
self.pred_mat = pred_mat
self.reps = reps
else:
(
dim_x,
dim_y,
num_alice_out,
num_bob_out,
num_alice_in,
num_bob_in,
) = pred_mat.shape
self.prob_mat = tensor(prob_mat, reps)
pred_mat2 = np.zeros((
dim_x**reps,
dim_y**reps,
num_alice_out**reps,
num_bob_out**reps,
num_alice_in**reps,
num_bob_in**reps,
))
i_ind = np.zeros(reps, dtype=int)
j_ind = np.zeros(reps, dtype=int)
for i in range(num_alice_in**reps):
for j in range(num_bob_in**reps):
to_tensor = np.empty(
[reps, dim_x, dim_y, num_alice_out, num_bob_out])
for k in range(reps - 1, -1, -1):
to_tensor[k] = pred_mat[:, :, :, :, i_ind[k], j_ind[k]]
pred_mat2[:, :, :, :, i, j] = tensor(to_tensor)
j_ind = update_odometer(j_ind, num_bob_in * np.ones(reps))
i_ind = update_odometer(i_ind, num_alice_in * np.ones(reps))
self.pred_mat = pred_mat2
self.reps = reps
def test_update_odometer_empty():
"""Return `None` if empty lists are provided."""
vec = np.array([])
upper_lim = np.array([])
res = update_odometer(vec, upper_lim)
bool_mat = np.isclose([], res)
np.testing.assert_equal(np.all(bool_mat), True)
def test_update_odometer_0_0():
"""Update odometer from [2, 2] to [0, 0]."""
vec = np.array([2, 2])
upper_lim = np.array([3, 2])
res = update_odometer(vec, upper_lim)
bool_mat = np.isclose([0, 0], res)
np.testing.assert_equal(np.all(bool_mat), True)
