# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
from __future__ import annotations
from enum import Enum
import braket.ir.annealing as ir
[docs]
class ProblemType(str, Enum):
"""The type of annealing problem.
QUBO: Quadratic Unconstrained Binary Optimization, with values 1 and 0
ISING: Ising model, with values +/-1
"""
QUBO = "QUBO"
ISING = "ISING"
[docs]
class Problem:
"""Represents an annealing problem."""
def __init__(
self,
problem_type: ProblemType,
linear: dict[int, float] | None = None,
quadratic: dict[tuple[int, int], float] | None = None,
):
"""Initializes a `Problem`.
Args:
problem_type (ProblemType): The type of annealing problem
linear (dict[int, float] | None): The linear terms of this problem,
as a map of variable to coefficient
quadratic (dict[tuple[int, int], float] | None): The quadratic terms of this problem,
as a map of variables to coefficient
Examples:
>>> problem = Problem(
>>> ProblemType.ISING,
>>> linear={1: 3.14},
>>> quadratic={(1, 2): 10.08},
>>> )
>>> problem.add_linear_term(2, 1.618).add_quadratic_term((3, 4), 1337)
"""
self._problem_type = problem_type
self._linear = linear or {}
self._quadratic = quadratic or {}
@property
def problem_type(self) -> ProblemType:
"""The type of annealing problem.
Returns:
ProblemType: The type of annealing problem
"""
return self._problem_type
@property
def linear(self) -> dict[int, float]:
"""The linear terms of this problem.
Returns:
dict[int, float]: The linear terms of this problem, as a map of variable to coefficient
"""
return self._linear
@property
def quadratic(self) -> dict[tuple[int, int], float]:
"""The quadratic terms of this problem.
Returns:
dict[tuple[int, int], float]: The quadratic terms of this problem,
as a map of variables to coefficient
"""
return self._quadratic
[docs]
def add_linear_term(self, term: int, coefficient: float) -> Problem:
"""Adds a linear term to the problem.
Args:
term (int): The variable of the linear term
coefficient (float): The coefficient of the linear term
Returns:
Problem: This problem object
"""
self._linear[term] = coefficient
return self
[docs]
def add_linear_terms(self, coefficients: dict[int, float]) -> Problem:
"""Adds linear terms to the problem.
Args:
coefficients (dict[int, float]): A map of variable to coefficient
Returns:
Problem: This problem object
"""
self._linear.update(coefficients)
return self
[docs]
def add_quadratic_term(self, term: tuple[int, int], coefficient: float) -> Problem:
"""Adds a quadratic term to the problem.
Args:
term (tuple[int, int]): The variables of the quadratic term
coefficient (float): The coefficient of the quadratic term
Returns:
Problem: This problem object
"""
self._quadratic[term] = coefficient
return self
[docs]
def add_quadratic_terms(self, coefficients: dict[tuple[int, int], float]) -> Problem:
"""Adds quadratic terms to the problem.
Args:
coefficients (dict[tuple[int, int], float]): A map of variables to coefficient
Returns:
Problem: This problem object
"""
self._quadratic.update(coefficients)
return self
[docs]
def to_ir(self) -> Problem:
"""Converts this problem into IR representation.
Returns:
Problem: IR representation of this problem object
"""
return ir.Problem(
type=ir.ProblemType[self._problem_type.value],
linear=self._linear,
quadratic={
",".join((str(q1), str(q2))): self._quadratic[(q1, q2)]
for q1, q2 in self._quadratic
},
)