Task3_2/Task3_2.py - master - timshk/Shalaev_Stepik

Task3_2/Task3_2.py
245 строк9 KB
timshk
Stepik задания
01 июн 2026, 18:58
01 июн 2026, 18:582e75834
"""
QAOA для задачи коммивояжера (TSP) с 3 клиентами и 1 складом
РАБОЧАЯ ВЕРСИЯ — исправлены синтаксические ошибки
"""

import numpy as np
from qiskit import QuantumCircuit, transpile
from qiskit.circuit import ParameterVector
from qiskit.quantum_info import SparsePauliOp
from qiskit_aer import AerSimulator
from qiskit.visualization import plot_histogram
import matplotlib.pyplot as plt

# ============================================================================
# 1. ВХОДНЫЕ ДАННЫЕ
# ============================================================================
dist_matrix = np.array([
    [0,   5,   8,   10],
    [5,   0,   4,   7],
    [8,   4,   0,   6],
    [10,  7,   6,   0]
])

N_CUSTOMERS = 3
N_POSITIONS = 3
N_QUBITS = N_POSITIONS * N_CUSTOMERS  # 9 кубитов

MAX_DIST = np.max(dist_matrix)
PENALTY = 4 * MAX_DIST  # штраф = 40

print("="*60)
print("TSP → QUBO → QAOA (3 клиента)")
print("="*60)
print(f"Кубитов: {N_QUBITS}")
print(f"Штраф (λ): {PENALTY}")


def build_qubo_matrix():
    """Строит QUBO-матрицу для TSP с 3 клиентами."""
    n = N_QUBITS
    Q = np.zeros((n, n))
    
    def idx(t, i):
        return t * N_CUSTOMERS + i
    
    # H_A1: на каждой позиции ровно один клиент
    for t in range(N_POSITIONS):
        for i in range(N_CUSTOMERS):
            Q[idx(t, i), idx(t, i)] += -2 * PENALTY
        for i in range(N_CUSTOMERS):
            for j in range(i+1, N_CUSTOMERS):
                Q[idx(t, i), idx(t, j)] += 2 * PENALTY
                Q[idx(t, j), idx(t, i)] += 2 * PENALTY
    
    # H_A2: каждый клиент посещён ровно один раз
    for i in range(N_CUSTOMERS):
        for t in range(N_POSITIONS):
            Q[idx(t, i), idx(t, i)] += -2 * PENALTY
        for t in range(N_POSITIONS):
            for s in range(t+1, N_POSITIONS):
                Q[idx(t, i), idx(s, i)] += 2 * PENALTY
                Q[idx(s, i), idx(t, i)] += 2 * PENALTY
    
    # H_B: целевая функция (длина маршрута)
    # Склад → первый клиент
    for i in range(N_CUSTOMERS):
        Q[idx(0, i), idx(0, i)] += dist_matrix[0, i+1]
    
    # Переходы между клиентами
    for t in range(N_POSITIONS - 1):
        for i in range(N_CUSTOMERS):
            for j in range(N_CUSTOMERS):
                Q[idx(t, i), idx(t+1, j)] += dist_matrix[i+1, j+1]
    
    # Последний клиент → склад
    for i in range(N_CUSTOMERS):
        Q[idx(N_POSITIONS-1, i), idx(N_POSITIONS-1, i)] += dist_matrix[i+1, 0]
    
    const = PENALTY * N_POSITIONS + PENALTY * N_CUSTOMERS
    return Q, const


def build_qaoa_circuit_from_qmatrix(Q, p=1):
    """
    Строит схему QAOA для заданной QUBO-матрицы.
    """
    n = Q.shape[0]
    
    # Сбор коэффициентов для Rz и Rzz гейтов
    z_terms = []      # (кубит, коэффициент)
    zz_terms = []     # (кубит1, кубит2, коэффициент)
    
    for i in range(n):
        for j in range(i, n):
            qij = Q[i, j] + Q[j, i] if i != j else Q[i, i]
            if abs(qij) < 1e-8:
                continue
            if i == j:
                z_terms.append((i, qij))
            else:
                zz_terms.append((i, j, qij))
    
    # Создаём параметры для каждого слоя
    gamma = ParameterVector('γ', p)
    beta = ParameterVector('β', p)
    
    # Строим схему
    circuit = QuantumCircuit(n, n, name=f"QAOA_TSP_p{p}")
    
    # Начальное состояние |+>^n
    circuit.h(range(n))
    
    # Слои QAOA
    for layer in range(p):
        # Оператор проблемы e^{-i γ H_C}
        for qubit, coeff in z_terms:
            circuit.rz(2 * gamma[layer] * coeff, qubit)
        
        for q1, q2, coeff in zz_terms:
            circuit.rzz(2 * gamma[layer] * coeff, q1, q2)
        
        # Смешивающий оператор e^{-i β Σ X}
        for qubit in range(n):
            circuit.rx(2 * beta[layer], qubit)
    
    # Измерение
    circuit.measure(range(n), range(n))
    
    return circuit, gamma, beta


def run_qaoa(circuit, gamma, beta, gamma_vals, beta_vals, shots=1024):
    """
    Запускает QAOA схему с заданными числовыми значениями параметров.
    """
    param_values = {}
    for i, g in enumerate(gamma_vals):
        param_values[gamma[i]] = g
    for i, b in enumerate(beta_vals):
        param_values[beta[i]] = b
    
    bound_circuit = circuit.assign_parameters(param_values)
    
    simulator = AerSimulator()
    compiled = transpile(bound_circuit, simulator)
    result = simulator.run(compiled, shots=shots).result()
    counts = result.get_counts()
    
    return counts


def decode_solution(counts, n_positions, n_customers):
    """
    Декодирует битовую строку в маршрут и вычисляет длину.
    """
    best_bitstring = max(counts, key=counts.get)
    bits = best_bitstring[::-1]
    
    route = []
    for t in range(n_positions):
        for i in range(n_customers):
            idx = t * n_customers + i
            if bits[idx] == '1':
                route.append(i + 1)
                break
    
    if len(route) != n_positions:
        return None, float('inf'), bits
    
    total = 0
    total += dist_matrix[0, route[0]]
    for i in range(len(route) - 1):
        total += dist_matrix[route[i], route[i+1]]
    total += dist_matrix[route[-1], 0]
    
    return route, total, bits


def main():
    # Построение QUBO
    Q, const = build_qubo_matrix()
    print(f"Константа (сдвиг энергии): {const:.2f}")
    
    # Построение схемы QAOA (p=1)
    circuit, gamma, beta = build_qaoa_circuit_from_qmatrix(Q, p=1)
    
    print(f"\nСхема построена:")
    print(f"  Кубитов: {circuit.num_qubits}")
    print(f"  Параметры: {list(circuit.parameters)}")
    print(f"  Глубина схемы: {circuit.depth()}")
    
    # Запуск с фиксированными параметрами
    print("\n" + "="*60)
    print("ЗАПУСК С ФИКСИРОВАННЫМИ ПАРАМЕТРАМИ (gamma=0.5, beta=0.3)")
    print("="*60)
    
    shots = 1024
    counts = run_qaoa(circuit, gamma, beta, gamma_vals=[0.5], beta_vals=[0.3], shots=shots)
    
    print(f"\nРаспределение результатов (всего {sum(counts.values())} измерений):")
    for bitstring, count in sorted(counts.items(), key=lambda x: x[1], reverse=True)[:8]:
        percent = count / shots * 100
        print(f"  {bitstring}: {count} ({percent:.1f}%)")
    
    # Декодирование лучшего решения
    route, total, bits = decode_solution(counts, N_POSITIONS, N_CUSTOMERS)
    
    print("\n" + "="*60)
    print("НАЙДЕННОЕ РЕШЕНИЕ")
    print("="*60)
    
    if route is None:
        print("Лучшая битовая строка не соответствует допустимому маршруту")
        print(f"Битовая строка: {bits}")
    else:
        route_str = " → ".join(map(str, route))
        print(f"Маршрут: 0 → {route_str} → 0")
        print(f"Длина маршрута: {total}")
        print(f"Битовая строка: {bits}")
        
        optimal_route = [2, 3, 1]
        optimal_len = 8 + 6 + 7 + 5
        print(f"\nОптимальный маршрут: 0 → 2 → 3 → 1 → 0 (длина = {optimal_len})")
        
        if total <= optimal_len + 5:
            print("QAOA нашёл решение, близкое к оптимальному!")
        else:
            print("Для точной оптимизации нужен вариационный цикл (scipy.optimize)")
    
    # Визуализация
    plot_histogram(counts, title=f"QAOA для TSP (3 клиента), shots={shots}")
    plt.show()
    
    print("\n" + "="*60)
    print("ПРИМЕЧАНИЯ")
    print("="*60)
    print("1. Для реального решения TSP нужен вариационный цикл")
    print("2. Оптимальный маршрут: 0→2→3→1→0 (длина = 26)")
    print("3. Второй оптимальный: 0→1→3→2→0 (длина = 28)")
    print("4. QAOA на одном слое (p=1) даёт приближение")
    print("5. Для точности требуется оптимизация параметров через scipy.optimize")


if __name__ == "__main__":
    main()