<gh_stars>1-10
# AUTOGENERATED! DO NOT EDIT! File to edit: 02_Optimizer.ipynb (unless otherwise specified).

__all__ = ['SEED', 'Optimizer', 'cost_matrix']

# Cell
import re
import IPython, graphviz

import numpy as np
import pandas as pd

from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.tree import export_graphviz

from mip import Model, xsum, minimize, BINARY

from .core import *
from .tree import *

SEED = 41
np.random.seed(SEED)

# Cell
class Optimizer:
    def __init__(self, c_i_j, combine, z, class_):
        self.c_i_j = c_i_j
        self.combine = combine
        self.z = z
        self.class_ = class_

    def solve(self, atm:ATMSKLEARN, x:Instance):
        partitions  = atm.v_i_j(x)
        cost_matrix = self.c_i_j(partitions, x.content)

        model = Model()

        trees = atm.get_trees()

        # make v_i_j and phi_t_k as boolean variable in the integer linear programming problem
        v_i_j   = [[model.add_var(var_type=BINARY) for j in range(len(partitions[i]))] for i in range(len(x.content))]
        phi_t_k = [[model.add_var(var_type=BINARY) for j in range(len(atm.get_leaves(t.tree_)))] for t in trees]


        # objective
        model.objective = minimize(xsum(v_i_j[i][j] * cost_matrix[i][j] for i in range(len(v_i_j)) \
                                        for j in range(len(v_i_j[i]))))

        # constraints
        w_t     = atm.calculate_tree_weights()
        h_t_k   = atm.h_t_k(combine, class_=self.class_)


        model += (xsum(phi_t_k[i][j] * h_t_k[i][j] * w_t[i] for i in range(len(trees)) \
                       for j in range(len(h_t_k[i]))) >= self.z)

        #check if feature value belongs to one and only one partition
        for i in range(len(x.content)):
            model += xsum(v_i_j[i][j] for j in range(len(v_i_j[i]))) == 1


        for i in range(len(trees)):
            tree   = trees[i].tree_
            leaves = atm.get_leaves(tree)

            pi = {kidx:atm.find_ancestors(tree, 0, k, p=[])[1] for kidx, k in enumerate(leaves)}

            for j in range(len(leaves)):
                ancestors   = pi[j]
                n_ancestors = len(ancestors) # |pi_t_k|

                model += xsum(atm.predicates_mask(tree, a, partitions, x.types)[m] * v_i_j[tree.feature[a[0]]][m] \
                              for a in ancestors for m in range(len(v_i_j[tree.feature[a[0]]])))\
                              >= (phi_t_k[i][j] * n_ancestors)

        # check if instance is present in one and only one leaf node in
        # all trees
        for i in range(len(trees)):
            tree   = trees[i].tree_
            leaves = atm.get_leaves(tree)

            model += xsum(phi_t_k[i][j] for j in range(len(leaves))) == 1

        # optimizing
        model.optimize()

        v_i_j_sol   = [[int(v_i_j[i][j].x) for j in range(len(v_i_j[i]))] for i in range(len(v_i_j))]
        phi_t_k_sol = [[int(phi_t_k[i][j].x) for j in range(len(phi_t_k[i]))] for i in range(len(phi_t_k))]

        return v_i_j_sol, phi_t_k_sol

# Cell
def cost_matrix(partitions, x, p=0):
    C_i_j    = []

    for i in range(len(x)):
        s = partitions[i]
        feat_cost = []
        for j in range(len(s)):
            if len(s[j]) > 1:
                if (x[i] >= s[j][0]) and (x[i] < s[j][1]):
                    feat_cost.append(0)
                else:
                    feat_cost.append(min((x[i] - s[j][0]) ** p, (x[i] - s[j][1]) ** p))
            else:
                if x[i] == s[j][0]:
                    feat_cost.append(0)
                else:
                    feat_cost.append(1)
        C_i_j.append(feat_cost)

    return C_i_j