Skip to content

Extending Formulations

Learn how to create variants and extensions of existing formulations.

Why Extend Formulations?

  • Test different mathematical models
  • Add problem variants (e.g., TSP with time windows)
  • Optimize for specific solvers
  • Add constraints or objectives

Creating a Formulation Variant

Example: TSP with Penalties

Subclass the formulation you are extending and override formulate. That way to_string and interpret are inherited, and the variant is still accepted by TspInstance, whose formulation field is typed as TspFormulation.

Note that formulate is a static method receiving only the data, so tunable settings belong on the class (ClassVar) rather than on the instance. Keep the variable names identical to the base formulation so the inherited interpret keeps working.

from typing import ClassVar, Literal

from luna_quantum import Model, Sense, Vtype
from luna_quantum.utils import quicksum

from luna_usecases.abstract import Registry
from luna_usecases.traveling_salesperson_problem import TspData, TspFormulation


@Registry.add
class TspWithPenalties(TspFormulation):
    """TSP formulation with penalty for long edges."""

    name: Literal["tsp_penalties"] = "tsp_penalties"

    penalty_threshold: ClassVar[float] = 500.0
    penalty_factor: ClassVar[float] = 2.0

    @staticmethod
    def formulate(data: TspData) -> Model:
        """Formulate with edge penalties."""
        n = len(data.distance_matrix)
        model = Model(name="TSP_Penalties", sense=Sense.MIN)

        x = {}
        for i_idx, city in enumerate(data.city_names):
            for pos in range(n):
                x[i_idx, pos] = model.add_variable(
                    f"x_{city}_{pos}", vtype=Vtype.BINARY
                )

        objective = 0
        for i_idx in range(n):
            for j_idx in range(n):
                if i_idx == j_idx:
                    continue
                distance = data.distance_matrix[i_idx, j_idx]
                # Add penalties for long edges
                weight = distance
                if distance > TspWithPenalties.penalty_threshold:
                    weight = distance * TspWithPenalties.penalty_factor
                for pos in range(n):
                    objective += weight * x[i_idx, pos] * x[j_idx, (pos + 1) % n]

        model.objective += objective

        # Each city at exactly one position
        for i_idx in range(n):
            model.add_constraint(
                quicksum(x[i_idx, pos] for pos in range(n)) == 1,
                name=f"city_one_position_{i_idx}",
            )

        # Each position holds exactly one city
        for pos in range(n):
            model.add_constraint(
                quicksum(x[i_idx, pos] for i_idx in range(n)) == 1,
                name=f"position_one_city_{pos}",
            )

        return model

Comparing Formulations

Compare different formulations on the same data:

import numpy as np

from luna_usecases.traveling_salesperson_problem import TspInstance

data = TspData(
    data_name="three_cities",
    city_names=["Berlin", "Hamburg", "Munich"],
    distance_matrix=np.array(
        [[0.0, 289.0, 585.0], [289.0, 0.0, 796.0], [585.0, 796.0, 0.0]]
    ),
    start_city="Berlin",
)

formulations = {
    "Standard": TspFormulation(),
    "With Penalties": TspWithPenalties(),
}

for label, formulation in formulations.items():
    instance = TspInstance(data=data, formulation=formulation)
    model = instance.formulate()
    print(f"{label}: {len(model.constraints)} constraints")

Adding New Constraints

Example: Must visit city B before city C. Positions are encoded by x[i, p], so the position of a city is sum(p * x[i, p]):

n = 3
model = Model(name="TSP_precedence", sense=Sense.MIN)
x = {
    (i, p): model.add_variable(f"x_{i}_{p}", vtype=Vtype.BINARY)
    for i in range(n)
    for p in range(n)
}

b_idx, c_idx = 1, 2

model.add_constraint(
    quicksum(p * x[b_idx, p] for p in range(n))
    <= quicksum(p * x[c_idx, p] for p in range(n)) - 1,
    name="visit_b_before_c",
)

See Also