File size: 475 Bytes
719d0db
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
import numpy as np
from models.solvers.ortools.ortools_base import ORToolsBase

class ORToolsTSP(ORToolsBase):
    def __init__(self, large_value=1e+6, scaling=False):
        super().__init__(large_value, scaling)
    
    def scaling_feats(self, node_feats):
        return {
            key: (node_feat * self.large_value).astype(np.int64) 
                 if key == "coords" else 
                 node_feat
            for key, node_feat in node_feats.items()
        }