Model/recommendations/optimiser/CostOptimiser.py

from mip import Model, xsum, minimize, BINARY, OptimizationStatus
from typing import Mapping
from utils.logger import setup_logger

logger = setup_logger()


class CostOptimiser:
    """
    This class is used to minimise cost, given a constrained minimum gain
    """

    # We add an optional buffer to the minimum gain to allow for slack in the optimisation
    BUFFER = 0.2

    def __init__(
        self,
        components: list[list[Mapping[str, int | float | str]]],
        min_gain: float | int,
        verbose: bool = False,
        allow_slack: bool = True
    ):
        self.components = components
        self.min_gain = min_gain
        self.gain_constraint = None
        self.m = None
        self.variables = []
        self.solution = []
        self.allow_slack = allow_slack

        self.solution_cost = None
        self.solution_gain = None
        self.verbose = verbose

    @classmethod
    def calculate_sap_gain_with_slack(cls, min_gain: int | float):
        """
        Adds a small amount of buffer to the minimum gain, to account for possible error in SAP predictions
        :param min_gain: Numerical value for the minimum gain
        :return:
        """
        if min_gain == 0:
            return min_gain
        elif min_gain <= 5:
            return min_gain + 0.25
        elif min_gain <= 20:
            return min_gain + 0.5
        else:
            return min_gain + 0.75

    def setup(self):
        # Initialize Model
        self.m = Model("knapsack")
        # Set the verbosity level
        self.m.verbose = 1 if self.verbose else 0

        # Create variables
        self.variables = [
            [self.m.add_var(var_type=BINARY, name=str(component["id"])) for component in group] for group in
            self.components
        ]

        # Set objective
        # This objective is to minimize
        # cost_ig * x_ig, where cost_ig represents the cost for ith part in group g
        # and x_ig is the binary decision variable for the ith part in group g
        self.m.objective = minimize(
            xsum(
                component['cost'] * var for group, group_vars in zip(self.components, self.variables) for component, var
                in
                zip(group, group_vars)
            )
        )

        # Add constraints
        # This constrain ensures that sum of gain_ig * x_ig >= min_gain, where gain_ig represents the gain for the ith
        # component
        # in group g, and x_ig is the binary decision variable for the ith component in group g
        gain_expression = xsum(
            item['gain'] * var for group, group_vars in zip(self.components, self.variables) for item, var in
            zip(group, group_vars)
        ) >= self.min_gain

        self.gain_constraint = self.m.add_constr(gain_expression)

        # At most one item from each group
        # This constraint ensures that at most one item from each group is selected
        # This is expressed by summing up the decision variables for each group and ensuring that the sum is <= 1
        for group_vars in self.variables:
            self.m += xsum(var for var in group_vars) <= 1

    def add_budget_constraint(self, budget: int | float) -> None:
        # Inject budget constraint, which ensures that sum of cost_ig * x_ig <= budget, where cost_ig represents the
        # cost for the ith component in group g, and x_ig is the binary decision variable for the ith component in 
        # group g

        self.m += (
            xsum(
                item["cost"] * var
                for group, group_vars in zip(self.components, self.variables)
                for item, var in zip(group, group_vars)
            )
            <= budget
        )

    def setup_slack(self):

        # Remove the original gain constraint
        self.m.remove(self.gain_constraint)
        # Add slack variable
        s = self.m.add_var(lb=0)

        # Modify the constraint
        self.m += xsum(
            item['gain'] * var for group, group_vars in zip(self.components, self.variables) for item, var in
            zip(group, group_vars)
        ) + s >= self.min_gain

        # Modify the objective to penalize the use of slack
        penalty = 10000  # you can adjust this based on how much you want to penalize the use of slack
        self.m.objective = minimize(
            xsum(
                component['cost'] * var for group, group_vars in zip(self.components, self.variables) for component, var
                in
                zip(group, group_vars)
            ) + penalty * s
        )

    def solve(self):
        # Solve the problem
        self.m.optimize()

        if self.m.status == OptimizationStatus.INFEASIBLE:
            if self.allow_slack:
                self.setup_slack()
                self.m.optimize()
            else:
                # Explicity return an empty solution
                self.solution = []
                self.solution_cost = 0
                self.solution_gain = 0
                return

        # If we still have an infeasible solution, we return an empty solution

        self.solution = [
            item for group, group_vars in zip(self.components, self.variables) for item, var in zip(group, group_vars)
            if
            var.x >= 0.99
        ]

        # Get the selected items
        self.solution_cost = self.m.objective.x
        self.solution_gain = sum([component['gain'] for component in self.solution])