Solve Maximum Weight Independent Set problems using Fire Opal

import fireopal as fo
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
from pyscipopt import Model
import qctrlvisualizer as qv
from sympy import Poly, symarray

plt.style.use(qv.get_qctrl_style())

# Set credentials.
# These are the properties for the publicly available provider for IBM backends.
# If you have access to a private provider and wish to use it, replace these values.
hub = "ibm-q"
group = "open"
project = "main"
token = "YOUR_IBM_TOKEN"

credentials = fo.credentials.make_credentials_for_ibmq(
    token=token, hub=hub, group=group, project=project
)

# Enter your desired IBM backend here.
backend_name = "your_desired_backend"

def mwis_to_cost(graph: nx.Graph) -> Poly:
    """
    Convert a graph representing a MWIS instance to a cost polynomial.
    """
    node_ids = {node: i for i, node in enumerate(graph.nodes)}
    variables = symarray("x", graph.number_of_nodes())
    # Set a conflict penalty high enough to discourage adjacent selections
    conflict_weight = 4 * max(graph.nodes[n]["weight"] for n in graph.nodes)

    cost = 0
    # Reward including a node
    for node in graph.nodes:
        cost -= graph.nodes[node]["weight"] * variables[node_ids[node]]
    # Penalize selecting adjacent nodes
    for u, v in graph.edges():
        cost += conflict_weight * variables[node_ids[u]] * variables[node_ids[v]]

    return Poly(cost, *variables)


def draw_graph(graph: nx.Graph, bitstring: str = "", index_map: dict = None) -> None:
    """
    Visualize an MWIS problem graph optionally highlighting selected nodes.
    If an index map is provided, bitstring indices are reordered accordingly.
    """
    # Identify selected nodes
    if index_map:
        solution = {
            str(index_map[key])
            for i, key in enumerate(index_map.keys())
            if bitstring[i] == "1"
        }
    else:
        solution = {str(i) for i, bit in enumerate(bitstring) if bit == "1"}
    # Assign node colors based on selection
    node_colors = [
        qv.QCTRL_STYLE_COLORS[0] if node in solution else qv.QCTRL_STYLE_COLORS[2]
        for node in graph.nodes
    ]
    # Use node weights as labels
    labels = nx.get_node_attributes(graph, "weight")
    pos = nx.spring_layout(graph, seed=0)
    nx.draw(
        graph,
        pos,
        with_labels=True,
        labels=labels,
        node_color=node_colors,
        font_color="white",
    )
    plt.show()


def solve_classically(graph: nx.Graph) -> tuple[str, float]:
    """
    Solve the MWIS problem classically using linear programming (SCIP).
    """
    model = Model("MWIS")
    model.hideOutput(True)
    # Create binary decision variables for each node
    x = {
        v: model.addVar(name=f"x_{v}", vtype="B", obj=graph.nodes[v]["weight"])
        for v in graph.nodes
    }
    # Add constraints
    for u, v in graph.edges():
        model.addCons(x[u] + x[v] <= 1)

    model.setMaximize()
    model.optimize()

    if model.getStatus() != "optimal":
        raise RuntimeError("Solver did not find an optimal solution.")

    solution = [v for v in graph.nodes if model.getVal(x[v]) > 0.5]
    solution_bitstring = "".join("1" if v in solution else "0" for v in graph.nodes)
    total_weight = sum(graph.nodes[v]["weight"] for v in solution)

    return solution_bitstring, total_weight


def check_independence(bitstring: str, graph: nx.Graph, index_map: dict = None) -> bool:
    """
    Check whether the selected nodes form an independent set in the problem graph.
    If an index map is provided, bitstring indices are reordered accordingly.
    """
    if index_map:
        selected_nodes = {
            str(index_map[key])
            for i, key in enumerate(index_map.keys())
            if bitstring[i] == "1"
        }
    else:
        selected_nodes = {str(i) for i, bit in enumerate(bitstring) if bit == "1"}
    for u, v in graph.edges:
        if u in selected_nodes and v in selected_nodes:
            return False
    return True


def get_cost_distribution(
    distribution: dict[str, int], cost: Poly, index_map: dict = None
) -> tuple[np.ndarray, np.ndarray]:
    """
    Given a bitstring count distribution and a cost polynomial, compute the costs and relative probabilities.
    If an index map is provided, bitstring indices are reordered accordingly.
    """
    shot_count = sum(distribution.values())
    variables = cost.gens
    costs = []
    probabilities = []
    # Calculate the cost and probability for each bitstring
    for bitstring, count in distribution.items():
        if index_map:
            costs.append(
                cost.subs({var: bitstring[index_map[str(var)]] for var in variables})
            )
        else:
            costs.append(cost.subs(zip(cost.gens, bitstring)))
        probabilities.append(count / shot_count)
    return np.array(costs), np.array(probabilities)


def generate_random_bitstrings(
    length: int, num_bitstrings: int, rng_seed: int = 0
) -> list[str]:
    """
    Generate a random sampling of bitstrings.
    """
    # Initialize the random number generator with a seed for reproducibility
    rng = np.random.default_rng(seed=rng_seed)
    # Generate random bitstrings
    random_array = rng.integers(0, 2, size=(num_bitstrings, length))
    bitstrings = ["".join(row.astype(str)) for row in random_array]
    return bitstrings


def analyze_mwis_results(
    distribution: dict[str, int],
    bitstring: str,
    graph: nx.graph,
    cost: Poly,
    index_map: dict,
) -> None:
    """
    Analyze and visualize results from a MWIS execution using Fire Opal.
    """
    best_weight = -cost.subs({var: bitstring[index_map[str(var)]] for var in cost.gens})
    print(f"The best solution found by Fire Opal has weight {best_weight}.")
    print(
        "The selected set is independent."
        if check_independence(bitstring, graph, index_map)
        else "The selected set includes conflicting nodes."
    )

    # Number of variables
    qubit_count = len(bitstring)
    # Total number of bitstrings (shots)
    shot_count = sum(distribution.values())

    # Get cost distribution from the results
    costs, probs = get_cost_distribution(distribution, cost, index_map)
    # Generate random bitstrings and get their cost distribution
    random_bitstrings = generate_random_bitstrings(qubit_count, shot_count)
    random_probs = np.ones(shot_count) / shot_count
    random_costs = np.asarray(
        [cost.subs(zip(cost.gens, bitstring)) for bitstring in random_bitstrings]
    )

    # Plot cost histograms
    _, ax = plt.subplots(1, 1)
    ax.hist(
        [-costs, -random_costs],
        np.arange(-2000, -min(costs) + 1, 20),
        weights=[probs, random_probs],
        label=["Fire Opal QAOA solver", "Random Sampling"],
        color=[qv.QCTRL_STYLE_COLORS[0], qv.QCTRL_STYLE_COLORS[1]],
    )
    ax.set_xlabel("Weight")
    ax.set_ylabel("Probability")
    ax.legend()

mwis = nx.read_graphml("resources/mwis.graphml")
draw_graph(mwis)

cost = mwis_to_cost(mwis)

fire_opal_job = fo.solve_qaoa(cost, credentials, backend_name)

fire_opal_solution = fire_opal_job.result()

analyze_mwis_results(
    fire_opal_solution["final_bitstring_distribution"],
    fire_opal_solution["solution_bitstring"],
    mwis,
    cost,
    fire_opal_solution["variables_to_bitstring_index_map"],
)

optimal_solution, optimal_weight = solve_classically(mwis)
print(f"Optimal classical solution is an independent set with weight {optimal_weight}.")

draw_graph(
    mwis,
    fire_opal_solution["solution_bitstring"],
    fire_opal_solution["variables_to_bitstring_index_map"],
)

fo.print_package_versions()