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)