def test_ExtendedParams():
    params = ExtendedParams.linear_ramp_from_hamiltonian(hamiltonian,
                                                         2,
                                                         time=2)
    assert set(params.reg) == {0, 1, q1}
    assert np.allclose(params.betas, [[0.75] * 3, [0.25] * 3])
    assert np.allclose(params.gammas_singles, [[0.25], [0.75]])
    assert np.allclose(params.gammas_pairs, [[0.25, 0.25], [0.75, 0.75]])
    assert [params.qubits_singles
            ] == [term.get_qubits() for term in hamiltonian if len(term) == 1]
    # Test updating and raw output
    raw = np.random.rand(len(params))
    params.update_from_raw(raw)
    assert np.allclose(raw, params.raw())
def test_non_fourier_params_are_consistent():
    """
    Check that StandardParams, StandardWithBiasParams and
    ExtendedParams give the same rotation angles, given the same data"""
    p1 = StandardParams.linear_ramp_from_hamiltonian(hamiltonian, 2, time=2)
    p2 = ExtendedParams.linear_ramp_from_hamiltonian(hamiltonian, 2, time=2)
    p3 = StandardWithBiasParams.linear_ramp_from_hamiltonian(hamiltonian,
                                                             2,
                                                             time=2)
    assert np.allclose(p1.x_rotation_angles, p2.x_rotation_angles)
    assert np.allclose(p2.x_rotation_angles, p3.x_rotation_angles)
    assert np.allclose(p1.z_rotation_angles, p2.z_rotation_angles)
    assert np.allclose(p2.z_rotation_angles, p3.z_rotation_angles)
    assert np.allclose(p1.zz_rotation_angles, p2.zz_rotation_angles)
    assert np.allclose(p2.zz_rotation_angles, p3.zz_rotation_angles)
def test_parameter_infos():
    params = ExtendedParams.linear_ramp_from_hamiltonian(hamiltonian,
                                                         n_steps=2)
    print(params)
    p0 = params.raw()
    print(p0)
    sim = WavefunctionSimulator()
    cost_fun = QAOACostFunctionOnWFSim(hamiltonian=hamiltonian,
                                       params=params,
                                       sim=sim,
                                       scalar_cost_function=True,
                                       nshots=1,
                                       noisy=False)
    with local_qvm():
        out = scipy_optimizer(cost_fun, p0, epsilon=1e-3)
        print(out)
def test_ExtendedParams_plot():
    ham_no_bias = PauliTerm("Z", 0)
    ham_no_bias *= PauliTerm("Z", 1)
    next_term = PauliTerm("Z", 0, -2.0)
    next_term *= PauliTerm("Z", 2)
    ham_no_bias += next_term

    p = 5
    params = ExtendedParams.linear_ramp_from_hamiltonian(ham_no_bias, p)
    fig, ax = plt.subplots()
    params.plot(ax=ax)
    # plt.show()

    p = 8
    params = ExtendedParams((ham_no_bias, p),
                            ([0.1] * p * len(ham_no_bias.get_qubits()), [],
                             [0.2] * p * len(ham_no_bias)))
    fig, ax = plt.subplots()
    params.plot(ax=ax)
def test_extended_get_constraints():
    weights = [0.1, 0.3, 0.5, -0.7]
    term1 = PauliTerm.from_list([("Z", 0), ("Z", 1)], 0.1)
    term2 = PauliTerm.from_list([("Z", 1), ("Z", 2)], 0.3)
    term3 = PauliTerm.from_list([("Z", 2), ("Z", 3)], 0.5)
    term4 = PauliTerm.from_list([("Z", 3), ("Z", 0)], -0.7)
    ham = PauliSum([term1, term2, term3, term4])
    p = 2

    params = ExtendedParams.linear_ramp_from_hamiltonian(ham, p)

    actual_constraints = params.get_constraints()

    expected_constraints = [(0, 2 * np.pi), (0, 2 * np.pi), (0, 2 * np.pi),
                            (0, 2 * np.pi), (0, 2 * np.pi), (0, 2 * np.pi),
                            (0, 2 * np.pi), (0, 2 * np.pi),
                            (0, 2 * np.pi / weights[0]),
                            (0, 2 * np.pi / weights[1]),
                            (0, 2 * np.pi / weights[2]),
                            (0, 2 * np.pi / weights[3]),
                            (0, 2 * np.pi / weights[0]),
                            (0, 2 * np.pi / weights[1]),
                            (0, 2 * np.pi / weights[2]),
                            (0, 2 * np.pi / weights[3])]

    assert (np.allclose(expected_constraints, actual_constraints))

    cost_function = QAOACostFunctionOnWFSim(ham, params)
    np.random.seed(0)
    random_angles = np.random.uniform(-100, 100, size=len(params.raw()))
    value = cost_function(random_angles)

    normalised_angles = [
        random_angles[i] % actual_constraints[i][1]
        for i in range(len(params.raw()))
    ]
    normalised_value = cost_function(normalised_angles)

    assert (np.allclose(value, normalised_value))