Solve the Sherrington-Kirkpatrick model with the ParityQC Parity Twine Optimizer

import numpy as np

from qiskit_ibm_catalog import FunctionsCatalog

catalog = FunctionsCatalog(channel="ibm_quantum_platform")
function = catalog.load("parityqc/parity-twine-optimizer")

def generate_sk_problem(
    n: int,
    coupling_mean: float = 0.0,
    coupling_std: float = 1.0,
    local_fields_mean: float = 0.0,
    local_fields_std: float = 0.0,
    edge_density: float = 1.0,
    ensure_extensivity: bool = False,
    seed: int | None = None,
) -> dict:
    """Generate the Sherrington-Kirkpatrick (SK) model with varying
     edge density.

    Samples couplings and local fields via ``generate_couplings_sk_model``
    and assembles the corresponding Ising Hamiltonian

        H = -∑_{i<j} J_ij z_i z_j - ∑_i h_i z_i,

    where z_i ∈ {-1, +1}.

    Args:
        n: Number of spins (>= 2).
        coupling_mean: Mean coupling before optional SK scaling.
        coupling_std: Coupling std before optional SK scaling.
        local_fields_mean: Mean longitudinal field.
        local_fields_std: Std of the longitudinal fields.
        edge_density: Fraction of non-zero couplings, in ``[2/n, 1]``.
        ensure_extensivity: Whether to apply the SK 1/n scaling.
        seed: random number generator seed.

    Returns:
        An encoding the SK Hamiltonian.

    Raises:
        ValueError: If ``n < 2``, ``coupling_std < 0``, ``local_fields_std < 0``,
            or ``edge_density`` is outside ``[2/n, 1]``.
    """
    couplings, local_fields = _generate_couplings_sk_model(
        n=n,
        coupling_mean=coupling_mean,
        coupling_std=coupling_std,
        local_fields_mean=local_fields_mean,
        local_fields_std=local_fields_std,
        edge_density=edge_density,
        ensure_extensivity=ensure_extensivity,
        seed=seed,
    )

    # Handle quadratic terms: coupling[i, j] * zj[i] * zj[j]
    # Only iterate over the upper triangle (i < j)
    sk_problem = {
        str((i, j)): float(couplings[i, j])
        for i in range(n)
        for j in range(i + 1, n)
        if couplings[i, j] != 0
    }

    # Handle linear terms: local_fields[i] * zj[i]
    sk_problem.update(
        {
            str((i,)): float(local_fields[i])
            for i in range(n)
            if local_fields[i] != 0
        }
    )

    return sk_problem

def _generate_couplings_sk_model(
    n: int,
    coupling_mean: float = 0.0,
    coupling_std: float = 1.0,
    local_fields_mean: float = 0.0,
    local_fields_std: float = 0.0,
    edge_density: float = 1.0,
    ensure_extensivity: bool = False,
    seed: int | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    """Generate random couplings and local fields for an Ising / SK model.

    Couplings are Gaussian. With ``ensure_extensivity=True`` they follow the
    Sherrington-Kirkpatrick scaling ``J_ij ~ N(coupling_mean/n, coupling_std^2/n)``
    (extensive energy, O(n)); otherwise ``J_ij ~ N(coupling_mean, coupling_std^2)``
    (energy O(n^2)). Fields are ``h_i ~ N(local_fields_mean, local_fields_std^2)``.

    ``edge_density`` sets the fraction of the ``n*(n-1)/2`` possible couplings that
    are non-zero (1 = fully dense). The kept edges always include a random spanning
    tree, so the interaction graph is guaranteed connected. This requires at least
    ``n-1`` edges, so ``edge_density`` must be at least ``2/n``.

    Args:
        n: Number of spins (>= 2).
        coupling_mean: Mean coupling before optional SK scaling.
        coupling_std: Coupling std before optional SK scaling.
        local_fields_mean: Mean longitudinal field.
        local_fields_std: Std of the longitudinal fields.
        edge_density: Fraction of non-zero couplings, in ``[2/n, 1]``.
        ensure_extensivity: Whether to apply the SK 1/n scaling.
        seed: random number generator seed.

    Returns:
        Tuple ``(couplings, fields)``: a symmetric ``(n, n)`` matrix with zero
        diagonal, and an ``(n,)`` field vector.

    Raises:
        ValueError: If ``n < 2``, ``coupling_std < 0``, ``local_fields_std < 0``,
            or ``edge_density`` is outside ``[2/n, 1]``.
    """
    if n < 2:
        raise ValueError(f"n must be >= 2, got {n}")
    if coupling_std < 0 or local_fields_std < 0:
        raise ValueError(
            "coupling_std and local_fields_std must be non-negative"
        )

    # A connected graph on n nodes needs at least n-1 of the n*(n-1)/2 possible
    # edges, so edge_density has a hard lower bound of 2/n.
    min_edge_density = 2.0 / n
    if not min_edge_density <= edge_density <= 1.0:
        raise ValueError(
            f"edge_density must be in [{min_edge_density:.4g}, 1] for n={n} "
            f"(at least n-1 edges are needed to keep the graph connected), "
            f"got {edge_density}"
        )

    rng = np.random.default_rng(seed)

    j_loc, j_scale = (
        (coupling_mean / n, coupling_std / np.sqrt(n))
        if ensure_extensivity
        else (coupling_mean, coupling_std)
    )

    upper_idx = np.triu_indices(n, k=1)
    n_edges = len(upper_idx[0])

    # Select which edges are present.
    if edge_density < 1.0:
        n_keep = int(round(edge_density * n_edges))
        # Map each (i, j) node pair to its position in the flat upper-triangle list.
        pair_to_flat = {
            (int(i), int(j)): idx
            for idx, (i, j) in enumerate(
                zip(upper_idx[0], upper_idx[1], strict=False)
            )
        }

        # Random spanning tree: node perm[k] links to a random earlier node.
        perm = rng.permutation(n)
        keep = np.zeros(n_edges, dtype=bool)
        for k in range(1, n):
            child, parent = perm[k], perm[rng.integers(0, k)]
            i, j = min(child, parent), max(child, parent)
            keep[pair_to_flat[(int(i), int(j))]] = True

        # Fill the remaining budget with random non-tree edges.
        remaining = n_keep - (n - 1)
        if remaining > 0:
            keep[
                rng.choice(
                    np.flatnonzero(~keep), size=remaining, replace=False
                )
            ] = True
    else:
        keep = np.ones(n_edges, dtype=bool)

    n_present = int(keep.sum())
    if j_scale == 0.0:
        vals = np.full(n_present, j_loc)
    else:
        vals = rng.normal(loc=j_loc, scale=j_scale, size=n_present)

    couplings = np.zeros((n, n))
    couplings[upper_idx[0][keep], upper_idx[1][keep]] = vals
    couplings += couplings.T  # symmetrize; diagonal stays zero

    fields = (
        np.full(n, local_fields_mean)
        if local_fields_std == 0.0
        else rng.normal(loc=local_fields_mean, scale=local_fields_std, size=n)
    )

    return couplings, fields

n_spins = 50
sk_problem = generate_sk_problem(n_spins)

# Calculate the function size a model needs, before submitting it.

# Call ``calculate_function_size(problem)``: it returns "S", "M" or "L" and submits
# nothing. The size comes from the number of variables and how densely they interact.


import ast

from collections.abc import Iterator
from itertools import combinations
from typing import Any, Literal

from qiskit_addon_opt_mapper.problems import OptimizationProblem

Size = Literal["S", "M", "L"]
VariableType = Literal["spin", "binary"]
_ORDER = {"S": 0, "M": 1, "L": 2}

# Limits are provided by IBM
_SMALL_NUMVARS, _SMALL_DENSITY = 50, 0.1
_MEDIUM_NUMVARS, _MEDIUM_DENSITY = 100, 0.25


def calculate_function_size(
    problem: dict[str, float] | OptimizationProblem, **kwargs: Any
) -> Size:
    """Return the size ("S", "M" or "L") needed to run ``problem``."""
    metrics = _measure(problem, **kwargs)
    if _is_small(metrics):
        return "S"
    if _is_medium(metrics):
        return "M"
    return "L"


def to_optimization_problem(
    problem: dict[str, float] | OptimizationProblem,
    problem_type: str = "",
    variable_type: VariableType = "binary",
) -> OptimizationProblem:
    """Convert the Parity Twine Optimizer ``problem`` input into an
    ``OptimizationProblem``.

    Args:
        problem: A dictionary {term: coefficient} where each term is a string
            version of a tuple of variable ids, for example
            {"()": 3, "(0,)": 1, "(0, 1)": 2}. An ``OptimizationProblem`` is
            returned unchanged.
        problem_type: The ``problem_type`` passed to the function, for
            example "mis".
        variable_type: Type of all variables, either "spin" or "binary".

    Returns:
        The parsed problem, as a minimization without constraints.
    """
    if isinstance(problem, OptimizationProblem):
        return problem

    if variable_type not in ("spin", "binary"):
        raise ValueError(
            f"Invalid variable type {variable_type!r}, expected 'spin' or 'binary'."
        )

    constant: float = 0
    linear_terms: dict[str, float] = {}
    quadratic_terms: dict[tuple[str, str], float] = {}
    variables: set[str] = set()

    for raw_term, coefficient in problem.items():
        try:
            parsed = ast.literal_eval(raw_term)
        except (ValueError, SyntaxError) as exc:
            raise ValueError(
                f"Invalid problem term {raw_term!r}: expected a tuple literal like '(0, 1)'."
            ) from exc
        if not isinstance(parsed, tuple):
            raise ValueError(
                f"Invalid problem term {raw_term!r}: expected a tuple literal like '(0, 1)'."
            )
        term = tuple(map(str, parsed))
        if len(term) == 0:
            constant = coefficient
        elif len(term) == 1:
            linear_terms[term[0]] = coefficient
        elif len(term) == 2:
            quadratic_terms[term] = coefficient
        else:
            raise ValueError(
                "Invalid problem term, only quadratic problems are "
                f"supported, higher-order term: {term}"
            )
        variables.update(term)

    model = OptimizationProblem()
    model.name = problem_type
    for variable in sorted(variables):
        if variable_type == "spin":
            model.spin_var(variable)
        else:
            model.binary_var(variable)

    model.minimize(
        constant=constant, linear=linear_terms, quadratic=quadratic_terms
    )
    return model


def _measure(
    problem: dict[str, float] | OptimizationProblem, **kwargs: Any
) -> dict[str, float]:
    """Return the number of variables and the interaction density."""
    model = to_optimization_problem(
        problem,
        problem_type=kwargs.get("problem_type", ""),
        variable_type=kwargs.get("variable_type", "binary"),
    )
    num_vars = model.get_num_vars()

    pairs = set()
    for term in _interactions(model):
        pairs.update(frozenset(pair) for pair in combinations(set(term), 2))

    max_pairs = num_vars * (num_vars - 1) // 2
    return {
        "num_vars": num_vars,
        "density": len(pairs) / max_pairs if max_pairs else 0.0,
    }


def _interactions(problem: OptimizationProblem) -> Iterator[tuple[int, ...]]:
    """Yield the variables of every term that couples two or more of them.

    A constraint contributes every variable in it, because on hardware it becomes
    a penalty whose square couples them all.
    """
    objective = problem.objective
    yield from objective.quadratic.to_dict()
    for expression in objective.higher_order.values():
        yield from expression.to_dict()

    for constraint in problem.linear_constraints:
        yield _constraint_support(constraint)
    for constraint in problem.quadratic_constraints:
        yield _constraint_support(constraint)
    for constraint in problem.higher_order_constraints:
        yield _constraint_support(constraint)


def _constraint_support(constraint: Any) -> tuple[int, ...]:
    """Return every variable appearing in a constraint's left-hand side."""
    variables = set(constraint.linear.to_dict())
    if hasattr(constraint, "quadratic"):
        for term in constraint.quadratic.to_dict():
            variables.update(term)
    if hasattr(constraint, "higher_order"):
        for expression in constraint.higher_order.values():
            for term in expression.to_dict():
                variables.update(term)
    return tuple(variables)


def _is_small(metrics: dict[str, float]) -> bool:
    """Determine if the metrics fit T-size 'S'."""
    return (
        metrics["num_vars"] <= _SMALL_NUMVARS
        and metrics["density"] <= _SMALL_DENSITY
    )


def _is_medium(metrics: dict[str, float]) -> bool:
    """Determine if the metrics fit T-size 'M'."""
    return (
        metrics["num_vars"] <= _MEDIUM_NUMVARS
        and metrics["density"] <= _MEDIUM_DENSITY
    )

size = calculate_function_size(sk_problem)
print(size)

options = {
    "shots": 100000,
    "postprocessing_level": 1,
    "transpile_only": False,
    "job_tags": ["sk"],
}

function_job = function.run(
    problem=sk_problem,
    variable_type="spin",
    backend_name="ibm_phoenix",
    options=options,
    function_size=size,
)
print(f"Job ID: {function_job.job_id}")

# Monitor the job status
function_job.status()

result = function_job.result()

result