convert_constraints(scenario)

Scenario must be normalized scenario!

Source code in wt_ml/optimizer/sim_optimizer/constraints.py
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
def convert_constraints(scenario: OptimizationScenario) -> Type[ConstraintCreator]:
    """Scenario must be normalized scenario!"""
    equality_constraints = []
    inequality_constraints = []
    inequality_indices = []
    if scenario.total_budget:
        # total budget is taken as a constraint in this optimization model.
        # gathers=[] means the constraint is applied on the entire spends.
        # negate=False since this is a maximizing constraint.
        # TODO(@ruler501): need to look into locked budget in future.
        def global_constraint(x):
            return scenario.total_budget - tf.reduce_sum(x)

        equality_constraints.append({"type": "eq", "fun": global_constraint})

    key_to_idx_mapping = {key: i for i, key in enumerate(product(scenario.locations, scenario.brands))}
    constraint: Constraint
    for constraint in scenario.constraints:
        idxs = []
        loc_constr_avl = constraint.locations is not None and len(constraint.locations) > 0
        brand_constr_avl = constraint.brands is not None and len(constraint.brands) > 0
        if loc_constr_avl and brand_constr_avl:
            for loc in constraint.locations:
                for brand in constraint.brands:
                    idxs.append(key_to_idx_mapping[(loc, brand)])

        elif brand_constr_avl:
            for (loc, brand), value in key_to_idx_mapping.items():
                if brand in constraint.brands:
                    idxs.append(value)

        elif loc_constr_avl:
            for (loc, brand), value in key_to_idx_mapping.items():
                if loc in constraint.locations:
                    idxs.append(value)

        # TODO: doesnt support vehicle constraints

        def fun(x):
            return constraint.amount - tf.reduce_sum(tf.gather(x, idxs, axis=0))

        # min, max, eq
        if constraint.constraint_type == "min":
            inequality_constraints.append({"type": "ineq", "fun": fun})
            inequality_indices.append(idxs)
        elif constraint.constraint_type == "max":
            inequality_constraints.append({"type": "ineq", "fun": lambda x: -fun(x)})
            inequality_indices.append(idxs)
        elif constraint.constraint_type == "eq":
            equality_constraints.append({"type": "eq", "fun": fun})
        else:
            raise ValueError(f"Invalid constraint_type {constraint}")

    return SimpleNamespace(
        equality_constraints=equality_constraints,
        inequality_constraints=inequality_constraints,
        inequality_indices=inequality_indices,
        scenario=scenario,
    )