first commit

.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.png
+lkh_io_files/
+concorde_io_files/
+data/
+data_generator/
+checkpoints/

.streamlit/config.toml

@@ -0,0 +1,2 @@
+[server]
+enableStaticServing = true

Dockerfile

@@ -0,0 +1,38 @@
+FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime
+
+RUN echo "Building docker image"
+
+RUN apt-get -y update && \
+    apt-get install -y \
+    curl \
+    build-essential \
+    git \
+    vim \
+    tmux 
+
+# jupyter-notebook & lab
+RUN python3 -m pip install jupyter
+RUN python3 -m pip install jupyterlab
+
+# LLMs
+RUN python3 -m pip install openai
+RUN python3 -m pip install tiktoken
+RUN python3 -m pip install langchain
+
+# Web app
+RUN python3 -m pip install streamlit
+RUN python3 -m pip install streamlit-folium
+RUN python3 -m pip install folium
+
+# Google Map API
+RUN python3 -m pip install googlemaps
+
+# OR-tools
+RUN python3 -m pip install ortools
+
+# other convenient packages
+RUN python3 -m pip install torchmetrics
+RUN python3 -m pip install scipy
+RUN python3 -m pip install pandas
+RUN python3 -m pip install matplotlib
+RUN python3 -m pip install tqdm

LICENSE

@@ -0,0 +1,125 @@
+SOFTWARE LICENSE AGREEMENT FOR EVALUATION
+
+This SOFTWARE LICENSE AGREEMENT FOR EVALUATION (this "Agreement") is a legal contract between a person
+who uses or otherwise accesses or installs the Software (“User(s)”), and Nippon Telegraph and Telephone corporation ("NTT").
+
+READ THE TERMS AND CONDITIONS OF THIS AGREEMENT CAREFULLY BEFORE INSTALLING OR 
+OTHERWISE ACCESSING OR USING NTT'S PROPRIETARY SOFTWARE ACCOMPANIED BY THIS 
+AGREEMENT (the "SOFTWARE"). THE SOFTWARE IS COPYRIGHTED AND IT IS LICENSED TO USER UNDER 
+THIS AGREEMENT, NOT SOLD TO USER. BY INSTALLING OR OTHERWISE ACCESSING OR USING THE 
+SOFTWARE, USER ACKNOWLEDGES THAT USER HAS READ THIS AGREEMENT, THAT USER 
+UNDERSTANDS IT, AND THAT USER ACCEPTS AND AGREES TO BE BOUND BY ITS TERMS. IF AT ANY 
+TIME USER IS NOT WILLING TO BE BOUND BY THE TERMS OF THIS AGREEMENT, USER SHOULD
+TERMINATE THE INSTALLATION PROCESS, IMMEDIATELY CEASE AND REFRAIN FROM ACCESSING OR 
+USING THE SOFTWARE AND DELETE ANY COPIES USER MAY HAVE. THIS AGREEMENT REPRESENTS THE 
+ENTIRE AGREEMENT BETWEEN USER AND NTT CONCERNING THE SOFTWARE.
+
+
+BACKGROUND
+A.	NTT is the owner of all rights, including all patent rights, copyrights and trade secret rights, in and to the Software and 
+    related documentation except OSS listed in Exhibit A to this Agreement.
+
+B.	User wishes to obtain a royalty free license to use the Software to enable User to evaluate, and NTT wishes to grant such
+    a license to User, pursuant and subject to the terms and conditions of this Agreement.
+
+C.	As a condition to NTT's provision of the Software to User, NTT has required User to execute this Agreement.
+
+In consideration of these premises, and the mutual promises and conditions in this Agreement, the parties hereby agree as follows:
+
+1.	Grant of Evaluation License. NTT hereby grants to User, and User hereby accepts, under the terms and 
+    conditions of this Agreement, a royalty free, nontransferable and nonexclusive license to use the Software internally for the 
+    non-commercial purposes of testing, analyzing, and evaluating the methods or mechanisms as shown in the research paper 
+    submitted by NTT to a certain academy or technical contest, etc. ("academy"). User may make a reasonable number of 
+    backup copies of the Software solely for User's internal use pursuant to the license granted in this Section 1.
+
+2.  Shipment and Installation. NTT will ship or deliver the Software by any method that NTT deems appropriate. User 
+    shall be solely responsible for proper installation of the Software.
+
+3.  Term. This Agreement is effective whichever is earlier (i) upon User's acceptance of the Agreement, or (ii) upon User's 
+    installing, accessing, and using the Software, even if User has not expressly accepted this Agreement. Without prejudice to
+    any other rights, NTT may terminate this Agreement without notice to User (i) if User breaches or fails to comply with any 
+    of the limitations or other requirements described herein, and (ii) if NTT receives a notice from the academy stating that the
+    research paper would not be published, and in any such case User agrees that NTT may, in addition to any other remedies 
+    it may have at law or in equity, remotely disable the Software. User may terminate this Agreement at any time by User's
+    decision to terminate the Agreement to NTT and ceasing use of the Software. Upon any termination or expiration of this
+    Agreement for any reason, User agrees to uninstall the Software and either return to NTT the Software and all copies thereof,
+    or to destroy all such materials and provide written verification of such destruction to NTT.
+
+4.	Proprietary Rights
+    (a)	The Software is the valuable, confidential, and proprietary property of NTT, and NTT shall retain exclusive title to 
+        this property both during the term and after the termination of this Agreement.  Without limitation, User acknowledges
+        that all patent rights, copyrights and trade secret rights in the Software except OSS shall remain the exclusive property of
+        NTT at all times. User shall use not less than reasonable care in safeguarding the confidentiality of the Software. 
+    
+    (b)	NTT shall not be subject to the obligation of licensing the copyright, patent rights, etc. of author when user hope 
+        commercial / noncommercial use of the published / provided software, etc.
+    
+    (c)	USER SHALL NOT, IN WHOLE OR IN PART, AT ANY TIME DURING THE TERM OF OR AFTER THE TERMINATION OF THIS AGREEMENT: 
+        (i) SELL, ASSIGN, LEASE, DISTRIBUTE, OR OTHERWISE TRANSFER THE SOFTWARE TO ANY THIRD PARTY; 
+        (ii) EXCEPT AS OTHERWISE PROVIDED HEREIN, COPY OR REPRODUCE THE SOFTWARE IN ANY MANNER; 
+        (iii) DISCLOSE THE SOFTWARE TO ANY THIRD PARTY, EXCEPT TO USER'S EMPLOYEES WHO REQUIRE ACCESS TO THE SOFTWARE FOR THE PURPOSES OF THIS AGREEMENT; 
+        (iv) MODIFY, DISASSEMBLE, DECOMPILE, REVERSE ENGINEER OR TRANSLATE THE SOFTWARE; 
+        OR (v) ALLOW ANY PERSON OR ENTITY TO COMMIT ANY OF THE ACTIONS DESCRIBED IN (i) THROUGH (iv) ABOVE.
+    
+    (d)	User shall take appropriate action, by instruction, agreement, or otherwise, with respect to its employees permitted 
+        under this Agreement to have access to the Software to ensure that all of User's obligations under this Section 4 shall be satisfied.  
+
+5.  Indemnity. User shall defend, indemnify and hold harmless NTT, its agents and employees, from any loss, damage,
+    or liability arising in connection with User's improper or unauthorized use of the Software. NTT SHALL HAVE THE SOLE 
+    RIGHT TO CONDUCT DEFEND ANY ACTTION RELATING TO THE SOFTWARE.
+
+6.	Disclaimer. THE SOFTWARE IS LICENSED TO USER "AS IS," WITHOUT ANY TRAINING, MAINTENANCE, OR SERVICE OBLIGATIONS WHATSOEVER ONTHE PART OF NTT. 
+    NTT MAKES NO EXPRESS OR IMPLIED WARRANTIES OF ANY TYPE WHATSOEVER, INCLUDING WITHOUT LIMITATION THE IMPLIED WARRANTIES OF MERCHANTABILITY, 
+    OF FITNESS FOR A PARTICULAR PURPOSE AND OF NON-INFRINGEMENT ON COPYRIGHT OR ANY OTHER RIGHT OF THIRD PARTIES.  
+    USER ASSUMES ALL RISKS ASSOCIATED WITH ITS USE OF THE SOFTWARE, INCLUDING WITHOUT LIMITATION RISKS RELATING TO QUALITY, PERFORMANCE, DATA LOSS, 
+    AND UTILITY IN A PRODUCTION ENVIRONMENT. 
+
+7.	Limitation of Liability. IN NO EVENT SHALL NTT BE LIABLE TO USER OR TO ANY THIRD PARTY FOR ANY INDIRECT, SPECIAL, INCIDENTAL, 
+    OR CONSEQUENTIAL DAMAGES, INCLUDING BUT NOT LIMITED TO DAMAGES FOR PERSONAL INJURY, PROPERTY DAMAGE, LOST PROFITS, OR OTHER ECONOMIC LOSS, 
+    ARISING IN CONNECTION WITH USER'S USE OF OR INABILITY TO USE THE SOFTWARE, IN CONNECTION WITH NTT'S PROVISION OF OR FAILURE TO PROVIDE SERVICES 
+    PERTAINING TO THE SOFTWARE, OR AS A RESULT OF ANY DEFECT IN THE SOFTWARE.  
+    THIS DISCLAIMER OF LIABILITY SHALL APPLY REGARD¬LESS OF THE FORM OF ACTION THAT MAY BE BROUGHT AGAINST NTT, WHETHER IN CONTRACT OR TORT, 
+    INCLUDING WITHOUT LIMITATION ANY ACTION FOR NEGLIGENCE.  USER'S SOLE REMEDY IN THE EVENT OF ANY BREACH OF THIS AGREEMENT BY NTT SHALL BE TERMINATION 
+    PURSUANT TO SECTION 3.
+
+8.	No Assignment or Sublicense. Neither this Agreement nor any right or license under this Agreement, nor the Software, may be sublicensed, assigned, 
+    or otherwise transferred by User without NTT's prior written consent.
+
+9.  OSS. The OSS included in the software is shown on the "OSS List" in Exhibit A. 
+    User shall be subject to the license term of each OSS, when User uses the software.
+
+10.	General
+    (d)	If any provision, or part of a provision, of this Agreement is or becomes illegal, unenforceable, or invalidated, by 
+        operation of law or otherwise, that provision or part shall to that extent be deemed omitted, and the remainder of this 
+        Agreement shall remain in full force and effect.
+
+    (e)	This Agreement is the complete and exclusive statement of the agreement between the parties with respect to the 
+        subject matter hereof, and supersedes all written and oral contracts, proposals, and other communications between the 
+        parties relating to that subject matter.  
+    
+    (f)	Subject to Section 8, this Agreement shall be binding on, and shall inure to the benefit of, the respective successors and
+        assigns of NTT and User.  
+    
+    (g)	If either party to this Agreement initiates a legal action or proceeding to enforce or interpret any part of this 
+        Agreement, the prevailing party in such action shall be entitled to recover, as an element of the costs of such action and not
+        as damages, its attorneys' fees and other costs associated with such action or proceeding.
+    
+    (h)	This Agreement shall be governed by and interpreted under the laws of Japan, without reference to conflicts of law principles. 
+        All disputes arising out of or in connection with this Agreement shall be finally settled by arbitration in Tokyo 
+        in accordance with the Commercial Arbitration Rules of the Japan Commercial Arbitration Association.  
+        The arbitration shall be conducted by three (3) arbitrators and in Japanese. The award rendered by the arbitrators shall be final 
+        and binding upon the parties. Judgment upon the award may be entered in any court having jurisdiction thereof.
+    
+    (f)	NTT shall not be liable to the User or to any third party for any delay or failure to perform NTT's obligation set 
+        forth under this Agreement due to any cause beyond NTT's reasonable control.
+
+EXHIBIT A
+- Software
+  N/A
+
+- OSS List
+  +----+---------+-----+
+  | No | License | OSS |
+  +----+---------+-----+
+  |  1 |   N/A   | N/A |
+  +----+---------+-----+

README.md

@@ -0,0 +1,45 @@
+# RouteExplainer: An Explanation Framework for Vehicle Routing Problem
+This repo is the official implementation of "RouteExplainer: An Explanation Framework for Vehicle Routing Problem" (PAKDD 2024). Please check more details at the project page https://ntt-dkiku.github.io/xai-vrp/.
+
+## Setup
+We recommend using Docker to setup development environments. Please use the [Dockerfile](./Dockerfile) in this repository. 
+```
+docker build -t route_explainer/route_explainer:1.0 .
+```
+If you use LKH and Concorde, you need to install them by typing the following command. LKH and Concorde is required for reproducing experiments, but not for demo.
+```
+python install_solvers.py
+```
+In the following, all commands are supposed to be typed inside the Docker container.
+
+## Reproducibility
+<!-- Refer to [reproduce_experiments.ipynb](./reproduct_experiments.ipynb). -->
+Coming Soon!
+
+## Training and evaluating edge classifiers
+### Generating synthetic data with labels
+```
+python generate_dataset.py --problem tsptw --annotation --parallel
+```
+
+### Training
+```
+python train.py
+```
+
+### Evaluation
+```
+python eval.py
+```
+
+## Explanation generation (demo)
+Go to http://localhost:8888 after launching the streamlit app with the following command. You may change the port number as you like.
+```
+streamlit run app.py --server.port 8888
+```
+
+## Licence
+Our code is licenced by NTT. Basically, the use of our code is limitted to research purposes. See [LICENSE](./LICENSE) for more details.
+
+## Citation
+Coming Soon!

analyze_dataset.py

@@ -0,0 +1,71 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import matplotlib.cm as cm
+from utils.util_data import load_dataset
+
+
+def get_cmap(num_colors):
+    if num_colors <= 10:
+        cm_name = "tab10"
+    elif num_colors <= 20:
+        cm_name = "tab20"
+    else:
+        assert False
+    return cm.get_cmap(cm_name)
+
+def analyze_dataset(dataset_path, output_dir):
+    dataset = load_dataset(dataset_path)
+    
+    #-----------------------------
+    # Stepwise frequency analysis 
+    #-----------------------------
+    max_steps = len(dataset[0][0]) # num_nodes
+    num_labels = 2
+    freq = [[] for _ in range(num_labels)]
+    weights = [[] for _ in range(num_labels)]
+    for instance in dataset:
+        labels = instance[-1]
+        for step, label in labels:
+            freq[label].append(step)
+    # visualize histogram
+    fig = plt.figure(figsize=(10, 10))
+    binwidth = 1
+    bins = np.arange(0, max_steps + binwidth, binwidth)
+    cmap = get_cmap(num_labels)
+    for i in range(len(weights)):
+        weights[i] = np.ones(len(freq[i])) / len(dataset)
+        plt.hist(freq[i], bins=bins, alpha=0.5, weights=weights[i], ec=cmap(i), color=cmap(i), label="prioritizing tour length", align="left")
+    plt.xlabel("Steps")
+    plt.ylabel("Frequency (density)")
+    if max_steps <= 20:
+        plt.xticks(np.arange(0, max_steps+1, 1))
+    plt.title(f"# of samples = {len(dataset)}\n# of nodes = {max_steps}")
+    plt.legend()
+    plt.savefig(f"{output_dir}/hist.png", dpi=150, bbox_inches="tight")
+
+    #-----------------------------
+    # Overall ratio of each class
+    #-----------------------------
+    total = np.sum([len(freq[i]) for i in range(num_labels)])
+    ratio = np.array([len(freq[i]) for i in range(num_labels)])
+    ratio = ratio / total
+    with open(f"{output_dir}/ratio.dat", "w") as f:
+        for i in range(len(ratio)):
+            print(f"label{i}, {ratio[i]}", file=f)
+
+if __name__ == "__main__":
+    import argparse
+    import os
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument("--dataset_path", type=str, required=True)
+    parser.add_argument("--output_dir", type=str, default=None)
+    args = parser.parse_args()
+
+    if args.output_dir is None:
+        dataset_dir = os.path.split(args.dataset_path)[0]
+        output_dir = dataset_dir
+    else:
+        output_dir = args.output_dir
+    output_dir += "/analysis"
+    os.makedirs(output_dir, exist_ok=True)
+    analyze_dataset(args.dataset_path, output_dir)

app.py

@@ -0,0 +1,455 @@
+# standard modules
+import os
+import pickle
+import datetime
+from PIL import Image
+from typing import List, Union
+
+# useful modules ("pip install" is required)
+import numpy as np
+import streamlit as st
+import pandas as pd
+import googlemaps
+import langchain
+from langchain.globals import set_verbose
+from langchain.chat_models import ChatOpenAI
+from langchain.schema import HumanMessage, AIMessage
+
+# our defined modules
+import utils.util_app as util_app
+from models.solvers.general_solver import GeneralSolver
+from models.cf_generator import CFTourGenerator
+from models.classifiers.general_classifier import GeneralClassifier
+from models.route_explainer import RouteExplainer
+
+# general setting
+SEED = 1234
+TOUR_NAME         = "static/kyoto_tour"
+TOUR_PATH         = TOUR_NAME + ".csv"
+TOUR_LATLNG_PATH  = TOUR_NAME + "_latlng.csv"
+TOUR_DISTMAT_PATH = TOUR_NAME + "_distmat.pkl"
+EXPANDED          = False
+DEBUG             = True
+ROUTE_EXPLAINER_ICON = np.array(Image.open("static/route_explainer_icon.png"))
+
+# for debug
+if DEBUG:
+    langchain.debug = True
+    set_verbose(True)
+
+def load_tour_list():
+    # get lat/lng
+    if os.path.isfile(TOUR_LATLNG_PATH):
+        df_tour = pd.read_csv(TOUR_LATLNG_PATH)
+    else:
+        df_tour = pd.read_csv(TOUR_PATH)
+        if googleapi_key := st.session_state.googleapi_key:
+            gmaps = googlemaps.Client(key=googleapi_key)
+            lat_list =[]; lng_list = []
+            for destination in df_tour["destination"]:
+                geo_result = gmaps.geocode(destination)
+                lat_list.append(geo_result[0]["geometry"]["location"]["lat"])
+                lng_list.append(geo_result[0]["geometry"]["location"]["lng"])
+            # add lat/lng
+            df_tour["lat"] = lat_list
+            df_tour["lng"] = lng_list
+            df_tour.to_csv(TOUR_LATLNG_PATH)
+    
+    # get the central point
+    st.session_state.lat_mean = np.mean(df_tour["lat"])
+    st.session_state.lng_mean = np.mean(df_tour["lng"])
+    st.session_state.sw = df_tour[["lat", "lng"]].min().tolist()
+    st.session_state.ne = df_tour[["lat", "lng"]].max().tolist()
+    st.session_state.df_tour = df_tour
+
+    # get the distance matrix
+    if os.path.isfile(TOUR_DISTMAT_PATH):
+        with open(TOUR_DISTMAT_PATH, "rb") as f:
+            distmat = pickle.load(f)
+    else:
+        if googleapi_key := st.session_state.googleapi_key:
+            gmaps = googlemaps.Client(key=googleapi_key)
+            distmat = []
+            for origin in df_tour["destination"]:
+                distrow = []
+                for dest in df_tour["destination"]:
+                    if origin != dest:
+                        dist_result = gmaps.distance_matrix(origin, dest, mode="driving")
+                        distrow.append(dist_result["rows"][0]["elements"][0]["duration"]["value"]) # unit: seconds
+                    else:
+                        distrow.append(0)
+                distmat.append(distrow)
+            distmat = np.array(distmat)
+            with open(TOUR_DISTMAT_PATH, "wb") as f:
+                pickle.dump(distmat, f)
+
+    # input features
+    def convert_clock2seconds(clock):
+        return sum([a*b for a, b in zip([3600, 60], map(int, clock.split(':')))])
+    time_windows = []
+    for i in range(len(df_tour)):
+        time_windows.append([convert_clock2seconds(df_tour["open"][i]), 
+                             convert_clock2seconds(df_tour["close"][i])])
+    time_windows = np.array(time_windows)
+    time_windows -= time_windows[0, 0]
+    node_feats = {
+        "time_window": time_windows.clip(0), 
+        "service_time": df_tour["stay_duration (h)"].to_numpy() * 3600
+    }
+    st.session_state.node_feats = node_feats
+    st.session_state.dist_matrix = distmat
+    st.session_state.node_info  = {
+        "open": df_tour["open"],
+        "close": df_tour["close"],
+        "stay": df_tour["stay_duration (h)"]
+    }
+
+    # tour list
+    if os.path.isfile(TOUR_DISTMAT_PATH) & os.path.isfile(TOUR_LATLNG_PATH):
+        st.session_state.tour_list = []
+        for i in range(len(df_tour)):
+            st.session_state.tour_list.append({
+                "name": df_tour["destination"][i],
+                "latlng": (df_tour["lat"][i], df_tour["lng"][i]),
+                "description": f"<font color='silver'>Hours: {df_tour['open'][i]} - {df_tour['close'][i]}<br>Duration of stay: {df_tour['stay_duration (h)'][i]}h<br>Remarks: {df_tour['remarks'][i]}</font>"
+            })
+
+def solve_vrp() -> None:
+    if ("node_feats" in st.session_state) and ("dist_matrix" in st.session_state):
+        solver = GeneralSolver("tsptw", "ortools", scaling=False)
+        classifier = GeneralClassifier("tsptw", "gt(ortools)")
+        routes = solver.solve(node_feats=st.session_state.node_feats, 
+                              dist_matrix=st.session_state.dist_matrix)
+        inputs = classifier.get_inputs(routes, 
+                                       0, 
+                                       st.session_state.node_feats, 
+                                       st.session_state.dist_matrix)
+        labels = classifier(inputs)
+        st.session_state.routes = routes.copy()
+        st.session_state.labels = labels.copy()
+        st.session_state.generated_actual_route = True
+
+#----------
+# LLM
+#----------
+def load_route_explainer(llm_type: str) -> None:
+    if st.session_state.openai_key:
+        # define llm
+        llm = ChatOpenAI(model=llm_type, 
+                         temperature=0,
+                         streaming=True,
+                         model_kwargs={"seed": SEED})
+        # model_kwargs={"stop": ["\n\n", "Human"]}
+
+        # define RouteExplainer
+        cf_generator = CFTourGenerator(cf_solver=GeneralSolver("tsptw", "ortools", scaling=False))
+        classifier = GeneralClassifier("tsptw", "gt(ortools)")
+        st.session_state.route_explainer = RouteExplainer(llm=llm, 
+                                                        cf_generator=cf_generator, 
+                                                        classifier=classifier)
+
+#----------
+# UI
+#----------
+# css settings
+st.set_page_config(layout="wide")
+util_app.apply_responsible_map_css()
+util_app.apply_centerize_icon_css()
+util_app.apply_red_code_css()
+util_app.apply_remove_sidebar_topspace()
+
+#------------------
+# side bar setting
+#------------------
+with st.sidebar:
+    #-------
+    # Title
+    #-------
+    icon_col, name_col = st.columns((1,10))
+    with icon_col:
+        util_app.apply_html('<a href="https://ntt-dkiku.github.io/xai-vrp" target="_blank"><img src="./app/static/route_explainer_icon.png" alt="RouteExplainer" width="30" height="30" style="margin-top: 20px;"></a>')
+    with name_col:
+        st.title("RouteExplainer")
+    
+    #----------
+    # API keys
+    #----------
+    st.subheader("API keys")
+    openai_key_col1, openai_key_col2 = st.columns((1,10))
+    with openai_key_col1:
+        util_app.apply_html('<a href="https://openai.com/blog/openai-api" target="_blank"> <img src="./app/static/openai_logo.png" alt="OpenAI API" width="30" height="30"> </a>')
+    with openai_key_col2:
+        openai_key = st.text_input(label="API keys",
+                                   key="openai_key",
+                                   placeholder="OpenAI API key",
+                                   type="password",
+                                   label_visibility="collapsed")
+    changed_key = openai_key == os.environ.get('OPENAI_API_KEY')
+    os.environ['OPENAI_API_KEY'] = openai_key
+
+    google_key_col1, google_key_col2 = st.columns((1, 10))
+    with google_key_col1:
+        util_app.apply_html('<a href="https://developers.google.com/maps?hl=en" target="_blank"> <img src="./app/static/googlemap_logo.png" alt="GoogleMap API" width="30" height="30"> </a>')
+    with google_key_col2:
+        st.text_input(label="GoogleMap API key",
+                      key="googleapi_key",
+                      placeholder="NOT required in this demo",
+                      type="password",
+                      label_visibility="collapsed")
+    
+    #----------------
+    # Foundation LLM
+    #----------------
+    st.subheader("Foundation LLM")
+    llm_type = st.selectbox("LLM", ["gpt-4", "gpt-4-1106-preview", "gpt-3.5-turbo"], key="llm_type", label_visibility="collapsed")
+
+    #-----------
+    # Tour plan
+    #-----------
+    st.subheader("Tour plan")
+    col1, col2 = st.columns((2, 1))
+    with col1:
+        # Comming soon: "Taipei Tour (for PAKDD2024)"
+        tour_plan = st.selectbox("Tour plan", ["Kyoto Tour"], key="tour_type", label_visibility="collapsed")
+    with col2:
+        st.button("Generate", on_click=solve_vrp, use_container_width=True)
+
+    # list destinations
+    load_tour_list()
+    with st.container():
+        if "routes" in st.session_state: # rearranage destinations in the route order if a route was derivied
+            # re-ordered destinations
+            reordered_tour_list = [st.session_state.tour_list[i] for i in st.session_state.routes[0][:-1]] if "routes" in st.session_state else st.session_state.tour_list
+            arr_time = datetime.datetime.strptime(st.session_state.node_info["open"][0], "%H:%M")
+            for step in range(len(reordered_tour_list)):
+                curr = reordered_tour_list[step]
+                next = reordered_tour_list[step+1] if step != len(reordered_tour_list) - 1 else reordered_tour_list[0]
+                curr_node_id = util_app.find_node_id_by_name(st.session_state.tour_list, curr["name"])
+                next_node_id = util_app.find_node_id_by_name(st.session_state.tour_list, next["name"])
+                open_time = datetime.datetime.strptime(st.session_state.node_info["open"][curr_node_id], "%H:%M")
+                # destination info
+                dep_time = max(arr_time, open_time) + datetime.timedelta(hours=st.session_state.node_info["stay"][curr_node_id])
+                dep_time_str = dep_time.strftime("%H:%M")
+                arr_time_str = arr_time.strftime("%H:%M")
+                arr_dep = f"Arr {arr_time_str} - Dep {dep_time_str}" if step != 0 else f"⭐ Dep {dep_time_str}"
+                with st.expander(f"{arr_dep} | {curr['name']}", expanded=EXPANDED):
+                    st.write(curr["description"], unsafe_allow_html=True)
+                # travel time
+                travel_time = st.session_state.dist_matrix[curr_node_id][next_node_id].item()
+                col1, col2, col3 = st.columns(3)
+                with col1:
+                    st.markdown(f"<center>{util_app.add_time_unit(travel_time)}</center>", unsafe_allow_html=True)
+                with col2:
+                    st.markdown("<center>|</center>", unsafe_allow_html=True)
+                st.write("")
+                arr_time = dep_time + datetime.timedelta(seconds=travel_time)
+            # return to the origin
+            destination = reordered_tour_list[0]
+            arr_time_str = arr_time.strftime("%H:%M")
+            with st.expander(f"⭐ Arr {arr_time_str} | {destination['name']}", expanded=EXPANDED):
+                st.write(destination["description"], unsafe_allow_html=True)
+        else: # just list destinations 
+            for destination in st.session_state.tour_list:
+                with st.expander(destination['name'], expanded=EXPANDED):
+                    st.write(destination["description"], unsafe_allow_html=True)
+
+#----------------------
+# state initialization
+#----------------------
+if "count" not in st.session_state:
+    st.session_state.count = 0
+if "chat_history" not in st.session_state:
+    st.session_state.chat_history = []
+if "generated_actual_route" not in st.session_state:
+    st.session_state.generated_actual_route = False
+if "generated_cf_route" not in st.session_state:
+    st.session_state.generated_cf_route = False
+if "curr_route" not in st.session_state:
+    st.session_state.curr_route = "Actual Route" # once the CF route is selected, this will be "Current Route"
+if "flag_example" not in st.session_state:
+    st.session_state.flag_example = False
+if "selected_example" not in st.session_state:
+    st.session_state.selected_example = None
+if "close_chat" not in st.session_state:
+    st.session_state.close_chat = False
+if "route_explainer" not in st.session_state or llm_type != st.session_state.curr_llm_type or changed_key:
+    load_route_explainer(llm_type)
+    st.session_state.curr_llm_type = llm_type
+
+#--------------------------------
+# The following is the main page
+#--------------------------------
+
+#----------
+# Greeding
+#----------
+if "routes" not in st.session_state:
+    util_app.apply_html('<center> <img src="./app/static/route_explainer_icon.png" alt="OpenAI API" width="120" height="120"> </center>')
+    greeding = "Hi, I'm RouteExplainer :)<br>Choose a tour and hit the <code>Generate</code> button to generate your initial route!"
+    if st.session_state.count == 0:
+        util_app.stream_words(greeding, prefix="<center><h4>", suffix="</h4></center>", sleep_time=0.02)
+    else:
+        util_app.apply_html(f"<center><h4>{greeding}</h4></center>")
+
+#--------------
+# chat history
+#--------------
+def find_last_map(lst: List[Union[str, tuple]]) -> int:
+    for i in range(len(lst) - 1, -1, -1):
+        if isinstance(lst[i], tuple):
+            return i
+    return None
+last_map_idx = find_last_map(st.session_state.chat_history)
+for i, msg in enumerate(st.session_state.chat_history):
+    if isinstance(msg, tuple): # if the history type is a tuple of maps
+        map1, map2 = (0, 1) if i == last_map_idx else (2, 3)
+        actual_route, cf_route = st.columns(2)
+        if msg[map1] is not None:
+            with actual_route:
+                util_app.visualize_actual_route(msg[map1])
+        if msg[map2] is not None:
+            with cf_route:
+                util_app.visualize_cf_route(msg[map2])
+    else: # if the history type is string
+        if isinstance(msg, AIMessage):
+            st.chat_message(msg.type, avatar=ROUTE_EXPLAINER_ICON).write(msg.content)
+        else:
+            st.chat_message(msg.type).write(msg.content)
+
+# examples
+if "cf_routes" not in st.session_state and st.session_state.flag_example:
+    def pickup_example(example: str):
+        st.session_state.selected_example = example
+    
+    examples = [
+        "Why do we visit Ginkaku-ji Temple from Fushimi-Inari Shrine and why not Kiyomizu-dera Temple?",
+        "What if we visit Kinkaku-ji directly from Kyoto Geishinkan, instead of Nijo-jo Castle?",
+        "Why was the edge from Kinkaku-ji to Kiyomizu-dera selected and why not the edge from Kinkaku-ji to Hanamikoji Dori?"
+    ]
+    col1, col2, col3 = st.columns(3)
+    with col1:
+        st.button(examples[0], 
+                use_container_width=True,
+                on_click=pickup_example,
+                args=(examples[0], ))
+    with col2:
+        st.button(examples[1],
+                use_container_width=True,
+                on_click=pickup_example,
+                args=(examples[1], ))
+    with col3:
+        st.button(examples[2],
+                use_container_width=True,
+                on_click=pickup_example,
+                args=(examples[2], ))
+
+#----------
+# chat box
+#----------
+def answer(prompt: str):
+    st.session_state.chat_history.append(HumanMessage(content=prompt))
+    st.chat_message("user").write(prompt)
+    if os.environ.get('OPENAI_API_KEY') == "":
+        error_msg = "An OpenAI API key has not been set yet :( Please enter a valid key in the side bar!"
+        with st.chat_message("assistant", avatar=ROUTE_EXPLAINER_ICON):
+            st.write(error_msg)
+        st.session_state.chat_history.append(AIMessage(content=error_msg)) 
+    elif util_app.validate_openai_api_key(os.environ.get('OPENAI_API_KEY')):
+        with st.chat_message("assistant", avatar=ROUTE_EXPLAINER_ICON):
+            explanation = st.session_state.route_explainer.generate_explanation(tour_list=st.session_state.tour_list,
+                                                                                whynot_question=prompt,
+                                                                                actual_routes=st.session_state.routes,
+                                                                                actual_labels=st.session_state.labels,
+                                                                                node_feats=st.session_state.node_feats,
+                                                                                dist_matrix=st.session_state.dist_matrix)
+            if len(explanation) > 0:
+                st.session_state.chat_history.append(AIMessage(content=explanation))
+            if st.session_state.generated_cf_route:
+                st.rerun()
+    else:
+        error_msg = "The input OpenAI API key appears to be invalid :( Please enter a valid key again in the side bar!"
+        with st.chat_message("assistant", avatar=ROUTE_EXPLAINER_ICON):
+            st.write(error_msg)
+        st.session_state.chat_history.append(AIMessage(content=error_msg))
+
+if st.session_state.selected_example is not None:
+    example = st.session_state.selected_example
+    st.session_state.selected_example = None
+    answer(example)
+else:
+    if "routes" in st.session_state and not st.session_state.close_chat:
+        if prompt := st.chat_input(placeholder="Ask a why-not question", key="chat_input"):
+            answer(prompt)
+
+#---------------------
+# route visualization
+#---------------------
+if "tour_list" in st.session_state: # if tour info is loaded
+    # first message
+    if st.session_state.generated_actual_route: # when an actual route is generated
+        with st.chat_message("assistant", avatar=ROUTE_EXPLAINER_ICON):
+            msg = "Here is your initial route. Please ask me a why and why-not question for a specfic edge!"
+            util_app.stream_words(msg, sleep_time=0.01)
+        st.session_state.flag_example = True
+        st.session_state.chat_history.append(AIMessage(content=msg))
+
+    # visualize the actual & CF routes
+    actual_route, cf_route = st.columns(2)
+    m = None; m2 = None; m_ = None; m2_ = None
+    if st.session_state.generated_actual_route or st.session_state.generated_cf_route:
+        m = util_app.initialize_map() # overwrite m
+        m_ = util_app.initialize_map()
+        if "labels" in st.session_state:
+            cf_step = st.session_state.cf_step-1 if st.session_state.generated_cf_route else -1
+            util_app.vis_route("routes", st.session_state.labels, m, cf_step, "actual")
+            util_app.vis_route("routes", st.session_state.labels, m_, cf_step, "actual", ant_path=False)
+        with actual_route:
+            util_app.visualize_actual_route(m)
+    if st.session_state.generated_cf_route:
+        m2 = util_app.initialize_map() # overwrite m2
+        m2_ = util_app.initialize_map()
+        if "cf_labels" in st.session_state:
+            util_app.vis_route("cf_routes", st.session_state.cf_labels, m2, st.session_state.cf_step-1, "cf")
+            util_app.vis_route("cf_routes", st.session_state.cf_labels, m2_, st.session_state.cf_step-1, "cf", ant_path=False)
+        with cf_route:
+            util_app.visualize_cf_route(m2)
+
+    # update states related to maps
+    if m is not None:
+        st.session_state.chat_history.append((m, m2, m_, m2_))
+
+    # route selection button
+    if len(st.session_state.chat_history) > 0:
+        last_msg = st.session_state.chat_history[-1]
+        if isinstance(last_msg, tuple):
+            if (last_msg[0] is not None) and (last_msg[1] is not None):
+                col1, col2 = st.columns(2)
+                with col1:
+                    st.button("Stay this route", on_click=util_app.select_actual_route)
+                with col2:
+                    st.button("Replace with this route", on_click=util_app.select_cf_route)
+                    util_app.change_hover_color("button", "Replace with this route", "#1e90ff")
+    
+    # for displaying examples
+    if st.session_state.generated_actual_route:
+        st.session_state.generated_actual_route = False
+        st.rerun()
+
+    st.session_state.generated_actual_route = False
+    st.session_state.generated_cf_route     = False
+
+# update session count
+st.session_state.count += 1
+
+js = f"""
+<script>
+    function scroll(dummy_var_to_force_repeat_execution){{
+        var textAreas = parent.document.querySelectorAll('section.main');
+        for (let index = 0; index < textAreas.length; index++) {{
+            textAreas[index].scrollTop = textAreas[index].scrollHeight;
+        }}
+    }}
+    scroll({len(st.session_state.chat_history)})
+</script>
+"""
+st.components.v1.html(js, height=0, width=0)

eval_classifier.py

@@ -0,0 +1,241 @@
+import os
+import argparse
+import json
+import multiprocessing
+import torch
+import time
+from tqdm import tqdm
+from torch.utils.data import DataLoader
+from torchmetrics.classification import MulticlassAccuracy, MulticlassF1Score
+from utils.util_calc import TemporalConfusionMatrix
+from models.classifiers.nn_classifiers.nn_classifier import NNClassifier
+from models.classifiers.ground_truth.ground_truth import GroundTruth
+from models.classifiers.ground_truth.ground_truth_base import FAIL_FLAG
+from utils.data_utils.tsptw_dataset import TSPTWDataloader
+from utils.data_utils.pctsp_dataset import PCTSPDataloader
+from utils.data_utils.pctsptw_dataset import PCTSPTWDataloader
+from utils.data_utils.cvrp_dataset import CVRPDataloader
+from utils.utils import set_device
+from utils.utils import load_dataset
+
+def load_eval_dataset(dataset_path, problem, model_type, batch_size, num_workers, parallel, num_cpus):
+    if model_type == "nn":
+        if problem == "tsptw":
+            eval_dataset = TSPTWDataloader(dataset_path, sequential=True, parallel=parallel, num_cpus=num_cpus)
+        elif problem == "pctsp":
+            eval_dataset = PCTSPDataloader(dataset_path, sequential=True, parallel=parallel, num_cpus=num_cpus)
+        elif problem == "pctsptw":
+            eval_dataset = PCTSPTWDataloader(dataset_path, sequential=True, parallel=parallel, num_cpus=num_cpus)
+        elif problem == "cvrp":
+            eval_dataset = CVRPDataloader(dataset_path, sequential=True, parallel=parallel, num_cpus=num_cpus)
+        else:
+            raise NotImplementedError
+
+        #------------
+        # dataloader
+        #------------
+        def pad_seq_length(batch):
+            data = {}
+            for key in batch[0].keys():
+                padding_value = True if key == "mask" else 0.0
+                # post-padding
+                data[key] = torch.nn.utils.rnn.pad_sequence([d[key] for d in batch], batch_first=True, padding_value=padding_value)
+            pad_mask = torch.nn.utils.rnn.pad_sequence([torch.full((d["mask"].size(0), ), True) for d in batch], batch_first=True, padding_value=False)
+            data.update({"pad_mask": pad_mask})
+            return data
+        eval_dataloader = DataLoader(eval_dataset,
+                                     batch_size=batch_size,
+                                     shuffle=False,
+                                     collate_fn=pad_seq_length,
+                                     num_workers=num_workers)
+        return eval_dataloader
+    else:
+        eval_dataset = load_dataset(dataset_path)
+        return eval_dataset
+
+def eval_classifier(problem: str, 
+                    dataset, 
+                    model_type: str, 
+                    model_dir: str = None, 
+                    gpu: int = -1, 
+                    num_workers: int = 4, 
+                    batch_size: int = 128, 
+                    parallel: bool = True,
+                    solver: str = "ortools",
+                    num_cpus: int = 1):
+    #--------------
+    # gpu settings
+    #--------------
+    use_cuda, device = set_device(gpu)
+
+    #-------
+    # model
+    #-------
+    num_classes = 3 if problem == "pctsptw" else 2
+    if model_type == "nn":
+        assert model_dir is not None, "please specify model_path when model_type is nn."
+        params = argparse.ArgumentParser()
+        # model_dir = os.path.split(args.model_path)[0]
+        with open(f"{model_dir}/cmd_args.dat", "r") as f:
+            params.__dict__ = json.load(f)
+        assert params.problem == problem, "problem of the trained model should match that of the dataset"
+        model = NNClassifier(problem=params.problem,
+                             node_enc_type=params.node_enc_type,
+                             edge_enc_type=params.edge_enc_type,
+                             dec_type=params.dec_type,
+                             emb_dim=params.emb_dim,
+                             num_enc_mlp_layers=params.num_enc_mlp_layers,
+                             num_dec_mlp_layers=params.num_dec_mlp_layers,
+                             num_classes=num_classes,
+                             dropout=params.dropout,
+                             pos_encoder=params.pos_encoder)
+        # load trained weights (the best epoch)
+        with open(f"{model_dir}/best_epoch.dat", "r") as f:
+            best_epoch = int(f.read())
+        print(f"loaded {model_dir}/model_epoch{best_epoch}.pth.")
+        model.load_state_dict(torch.load(f"{model_dir}/model_epoch{best_epoch}.pth"))
+        if use_cuda:
+            model.to(device)
+        is_sequential = model.is_sequential
+    elif model_type == "ground_truth":
+        model = GroundTruth(problem=problem, solver_type=solver)
+        is_sequential = False
+    else:
+        assert False, f"Invalid model type: {model_type}"
+
+    #---------
+    # Metrics
+    #---------
+    overall_accuracy = MulticlassF1Score(num_classes=num_classes, average="macro").to(device)
+    eval_accuracy_dict = {} # MulticlassAccuracy(num_classes=num_classes, average="macro")
+    temp_confmat_dict  = {} # TemporalConfusionMatrix(num_classes=num_classes, seq_length=50, device=device)
+    temporal_accuracy_dict  = {}
+    num_nodes_dist_dict = {}
+                        
+    #------------
+    # Evaluation
+    #------------
+    if model_type == "nn":
+        model.eval()
+        eval_time = 0.0
+        print("Evaluating models ...", end="")
+        start_time = time.perf_counter()
+        for data in dataset:
+            if use_cuda:
+                data = {key: value.to(device) for key, value in data.items()}
+            if not is_sequential:
+                shp = data["curr_node_id"].size()
+                data = {key: value.flatten(0, 1) for key, value in data.items()}
+            probs = model(data) # [batch_size x num_classes] or [batch_size x max_seq_length x num_classes]
+            if not is_sequential:
+                probs = probs.view(*shp, -1) # [batch_size x max_seq_length x num_classes]
+                data["labels"] = data["labels"].view(*shp)
+                data["pad_mask"] = data["pad_mask"].view(*shp)
+            #------------
+            # evaluation
+            #------------
+            start_eval_time = time.perf_counter()
+            # accuracy
+            seq_length_list = torch.unique(data["pad_mask"].sum(-1)) 
+            for seq_length_tensor in seq_length_list:
+                seq_length = seq_length_tensor.item()
+                if seq_length not in eval_accuracy_dict.keys():
+                    eval_accuracy_dict[seq_length] = MulticlassF1Score(num_classes=num_classes, average="macro").to(device)
+                    temp_confmat_dict[seq_length]  = TemporalConfusionMatrix(num_classes=num_classes, seq_length=seq_length, device=device)
+                    temporal_accuracy_dict[seq_length] = [MulticlassF1Score(num_classes=num_classes, average="macro").to(device) for _ in range(seq_length)]
+                    num_nodes_dist_dict[seq_length] = 0
+                seq_length_mask = (data["pad_mask"].sum(-1) == seq_length) # [batch_size]
+                extracted_labels = data["labels"][seq_length_mask]
+                extracted_probs  = probs[seq_length_mask]
+                extracted_mask   = data["pad_mask"][seq_length_mask].view(-1) # [batch_size x max_seq_length] -> [(batch_size*max_seq_length)]
+                eval_accuracy_dict[seq_length](extracted_probs.argmax(-1).view(-1)[extracted_mask], extracted_labels.view(-1)[extracted_mask])
+                mask = data["pad_mask"].view(-1)
+                overall_accuracy(probs.argmax(-1).view(-1)[mask], data["labels"].view(-1)[mask])
+                # confusion matrix
+                temp_confmat_dict[seq_length].update(probs.argmax(-1), data["labels"], data["pad_mask"])         
+                # temporal accuracy
+                for step in range(seq_length):
+                    temporal_accuracy_dict[seq_length][step](extracted_probs[:, step, :], extracted_labels[:, step])
+                # number of samples whose sequence length is seq_length
+                num_nodes_dist_dict[seq_length] += len(extracted_labels)
+            eval_time += time.perf_counter() - start_eval_time
+        calc_time = time.perf_counter() - start_time - eval_time
+        total_eval_accuracy = {key: value.compute().item() for key, value in eval_accuracy_dict.items()}
+        overall_accuracy = overall_accuracy.compute() #.item()
+        temporal_confmat = {key: value.compute() for key, value in temp_confmat_dict.items()}
+        temporal_accuracy = {key: [value.compute().item() for value in values] for key, values in temporal_accuracy_dict.items()}
+        print("done")
+        return overall_accuracy, total_eval_accuracy, temporal_accuracy, calc_time, temporal_confmat, num_nodes_dist_dict
+    else:
+        eval_accuracy = MulticlassF1Score(num_classes=num_classes, average="macro").to(device)
+        print("Loading data ...", end=" ")
+        with multiprocessing.Pool(num_cpus) as pool:
+            input_list = list(pool.starmap(model.get_inputs, [(instance["tour"], 0, instance) for instance in dataset]))
+        print("done")
+
+        print("Infering labels ...", end="")
+        pool = multiprocessing.Pool(num_cpus)
+        start_time = time.perf_counter()
+        prob_list = list(pool.starmap(model, tqdm([(inputs, False, False) for inputs in input_list])))
+        calc_time = time.perf_counter() - start_time
+        pool.close()
+        print("done")
+
+        print("Evaluating models ...", end="")
+        for i, instance in enumerate(dataset):
+            labels = instance["labels"]
+            for vehicle_id in range(len(labels)):
+                for step, label in labels[vehicle_id]:
+                    pred_label = prob_list[i][vehicle_id][step-1] # [num_classes]
+                    if pred_label == FAIL_FLAG:
+                        pred_label = label - 1 if label != 0 else label + 1
+                    eval_accuracy(torch.LongTensor([pred_label]).view(1, -1), torch.LongTensor([label]).view(1, -1))
+        total_eval_accuracy = eval_accuracy.compute()
+        print("done")
+        return total_eval_accuracy.item(), calc_time
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    #-----------------
+    # general settings
+    #-----------------
+    parser.add_argument("--gpu", default=-1, type=int, help="Used GPU Number: gpu=-1 indicates using cpu")
+    parser.add_argument("--num_workers", default=4, type=int, help="Number of workers in dataloader")
+    parser.add_argument("--parallel", )
+
+    #-------------
+    # data setting
+    #-------------
+    parser.add_argument("--dataset_path", type=str, help="Path to a dataset", required=True)
+    
+    #------------------
+    # Metrics settings
+    #------------------
+
+
+    #----------------
+    # model settings
+    #----------------
+    parser.add_argument("--model_type", type=str, default="nn", help="Select from [nn, ground_truth]")
+    # nn classifier
+    parser.add_argument("--model_dir", type=str, default=None)
+    parser.add_argument("--batch_size", type=int, default=256)
+    parser.add_argument("--parallel", action="store_true")
+    # ground truth
+    parser.add_argument("--solver", type=str, default="ortools")
+    parser.add_argument("--num_cpus", type=int, default=os.cpu_count())
+    args = parser.parse_args()
+
+    problem   = str(os.path.basename(os.path.dirname(args.dataset_path)))
+
+    dataset = load_eval_dataset(args.dataset_path, problem, args.model_type, args.batch_size, args.num_workers, args.parallel, args.num_cpus)
+    eval_classifier(problem=problem, 
+         dataset=dataset,
+         model_type=args.model_type,
+         model_dir=args.model_dir,
+         gpu=args.gpu,
+         num_workers=args.num_workers,
+         batch_size=args.batch_size,
+         parallel=args.parallel,
+         solver=args.solver,
+         num_cpus=args.num_cpus)

eval_solvers.py

@@ -0,0 +1,62 @@
+import os
+from tqdm import tqdm
+import multiprocessing
+import numpy as np
+from utils.utils import load_dataset, calc_tour_length
+from models.solvers.general_solver import GeneralSolver
+from models.classifiers.ground_truth.ground_truth import GroundTruth
+
+def eval_solver(solver, instance):
+    tour = solver.solve(instance)
+    tour_length = calc_tour_length(tour[0], instance["coords"])
+    return tour_length
+
+def eval(data_path, problem, solver_name, fix_edges, parallel):
+    dataset = load_dataset(data_path)
+    num_cpus = os.cpu_count() if parallel else 1
+    if fix_edges:
+        solver = GroundTruth(problem, solver_name)
+        if parallel:
+            with multiprocessing.Pool(num_cpus) as pool:
+                tours = list(tqdm(pool.starmap(solver.solve, [(step, instance["tour"][vehicle_id], instance, f"{i}-{vehicle_id}-{step}")
+                                                              for i, instance in enumerate(dataset)
+                                                              for vehicle_id in range(len(instance["tour"]))
+                                                              for step in range(1, len(instance["tour"][vehicle_id]))]), desc=f"Solving {data_path} with {solver_name}"))
+        else:
+            tours = []
+            for i, instance in enumerate(dataset):
+                for vehicle_id in range(len(instance["tour"])):
+                    for step in range(1, len(instance["tour"][vehicle_id])):
+                        tours.append(solver.solve(step, instance["tour"][vehicle_id], instance, f"{i}-{vehicle_id}-{step}"))
+        tour_length = {key: [] for key in tours[0].keys()}
+        for tour in tours:
+            for key, value in tour.items():
+                tour_length[key].append(value)
+    else:
+        solver = GeneralSolver(problem, solver_name)
+        with multiprocessing.Pool(num_cpus) as pool:
+            tour_length = list(tqdm(pool.starmap(eval_solver, [(solver, instance) for instance in dataset]), total=len(dataset), desc="Solving instances"))
+
+    feasible_ratio = 0.0
+    penalty = 0.0
+    avg_tour_length = np.mean(tour_length["tsp"])
+    std_tour_length = np.std(tour_length["tsp"])
+    return avg_tour_length, std_tour_length, feasible_ratio, penalty
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--problem", default="tsptw", type=str, help="Problem type: [tsptw, pctsp, pctsptw, cvrp]")
+    parser.add_argument("--solver_name", type=str, default="ortools", help="Select from ")
+    parser.add_argument("--data_path", type=str, help="Path to a dataset", required=True)
+    parser.add_argument("--parallel", action="store_true")
+    parser.add_argument("--all", action="store_true")
+    parser.add_argument("--fix_edges", action="store_true")
+    args = parser.parse_args()
+
+    avg_tour_length, std_tour_length, feasible_ratio, penalty = eval(data_path=args.data_path,
+                                                                     problem=args.problem,
+                                                                     solver_name=args.solver_name,
+                                                                     fix_edges=args.fix_edges,
+                                                                     parallel=args.parallel)
+    print(f"tour_length: {avg_tour_length} +/- {std_tour_length}")

generate_cf_dataset.py

@@ -0,0 +1,145 @@
+import random
+import numpy as np
+from tqdm import tqdm
+from multiprocessing import Pool
+from utils.utils import load_dataset, save_dataset
+from models.classifiers.ground_truth.ground_truth_base import get_visited_mask, get_tw_mask, get_cap_mask
+from models.classifiers.ground_truth.ground_truth import GroundTruth
+from models.solvers.general_solver import GeneralSolver
+from models.cf_generator import CFTourGenerator
+
+
+class CFDatasetBase():
+    def __init__(self, problem, cf_generator, classifier, base_dataset, num_samples, random_seed, parallel, num_cpus):
+        self.problem = problem
+        self.parallel = parallel
+        self.num_cpus = num_cpus
+        self.seed = random_seed
+        self.cf_generator = CFTourGenerator(cf_solver=GeneralSolver(problem, cf_generator))
+        self.classifier = GroundTruth(problem, classifier)
+        self.node_mask = NodeMask(problem)
+        self.dataset = load_dataset(base_dataset)
+        self.num_samples = len(self.dataset) if num_samples is None else num_samples
+
+    def generate_cf_dataset(self):
+        random.seed(self.seed)
+        cf_dataset = []
+        num_required_samples = self.num_samples
+        end = False
+        print("Data generation started.", flush=True)
+        while(not end):
+            dataset = self.dataset[:num_required_samples]
+            self.dataset = np.roll(self.dataset, -num_required_samples)
+            if self.parallel:
+                instances = self.generate_labeldata_para(dataset, self.num_cpus)
+            else:
+                instances = self.generate_labeldata(dataset)
+            cf_dataset.extend(filter(None, instances))
+            num_required_samples = self.num_samples - len(cf_dataset) 
+            if num_required_samples == 0:
+                end = True
+            else:
+                print(f"No feasible tour was not found in {num_required_samples} instances. Trying other {num_required_samples} instances.", flush=True)
+        print("Data generation completed.", flush=True)
+        return cf_dataset
+
+    def generate_labeldata(self, dataset):
+        return [self.annotate(instance) for instance in tqdm(dataset, desc="Annotating instances")]
+
+    def generate_labeldata_para(self, dataset, num_cpus):
+        with Pool(num_cpus) as pool:
+            annotation_data = list(tqdm(pool.imap(self.annotate, [instance for instance in dataset]), total=len(dataset), desc="Annotating instances"))
+        return annotation_data
+
+    def annotate(self, instance):
+        # generate a counterfactual route randomly
+        routes = instance["tour"]
+        vehicle_id = random.randint(0, len(routes) - 1)
+        if len(routes[vehicle_id]) - 2 <= 2:
+            return
+        cf_step    = random.randint(2, len(routes[vehicle_id]) - 2)
+        route = routes[vehicle_id]
+        mask = self.node_mask.get_mask(route, cf_step, instance)
+        node_id = np.arange(len(instance["coords"]))
+        feasible_node_id = node_id[mask]
+        feasible_node_id = feasible_node_id[feasible_node_id != route[cf_step]].tolist()
+        if len(feasible_node_id) == 0:
+            return
+        cf_visit  = random.choice(feasible_node_id)
+        cf_routes = self.cf_generator(routes, vehicle_id, cf_step, cf_visit, instance)
+        if cf_routes is None:
+            return
+
+        # annotate each edge
+        inputs = self.classifier.get_inputs(cf_routes, 0, instance)
+        labels = self.classifier(inputs, annotation=True)
+
+        # update tours and lables
+        instance["tour"] = cf_routes
+        instance["labels"] = labels
+        return instance
+
+class NodeMask():
+    def __init__(self, problem):
+        self.problem = problem
+
+        if self.problem == "tsptw":
+            self.mask_func = get_tsptw_mask
+        elif self.problem == "pctsp":
+            self.mask_func = get_pctsp_mask
+        elif self.problem == "pctsptw":
+            self.mask_func = get_pctsptw_mask
+        elif self.problem == "cvrp":
+            self.mask_func = get_cvrp_mask
+        else:
+            NotImplementedError
+
+    def get_mask(self, route, step, instance):
+        return self.mask_func(route, step, instance)
+
+def get_tsptw_mask(route, step, instance):
+    visited = get_visited_mask(route, step, instance)
+    not_exceed_tw = get_tw_mask(route, step, instance)
+    return ~visited & not_exceed_tw
+
+def get_pctsp_mask(route, step, instance):
+    visited = get_visited_mask(route, step, instance)
+    return ~visited
+
+def get_pctsptw_mask(route, step, instance):
+    visited = get_visited_mask(route, step, instance)
+    not_exceed_tw = get_tw_mask(route, step, instance)
+    return ~visited & not_exceed_tw
+
+def get_cvrp_mask(route, step, instance):
+    visited = get_visited_mask(route, step, instance)
+    less_than_cap = get_cap_mask(route, step, instance)
+    return ~visited & less_than_cap
+
+if __name__ == "__main__":
+    import os
+    import argparse
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument("--problem",      type=str, default="tsptw")
+    parser.add_argument("--base_dataset", type=str, required=True)
+    parser.add_argument("--cf_generator", type=str, default="ortools")
+    parser.add_argument("--classifier",   type=str, default="ortools")
+    parser.add_argument("--num_samples",  type=int, default=None)
+    parser.add_argument("--random_seed",  type=int, default=1234)
+    parser.add_argument("--parallel",     action="store_true")
+    parser.add_argument("--num_cpus",     type=int, default=4)
+    parser.add_argument("--output_dir",   type=str, default="data")
+    args = parser.parse_args()
+
+    dataset_gen = CFDatasetBase(args.problem,
+                                args.cf_generator,
+                                args.classifier,
+                                args.base_dataset,
+                                args.num_samples,
+                                args.random_seed,
+                                args.parallel,
+                                args.num_cpus)
+    cf_dataset = dataset_gen.generate_cf_dataset()
+
+    output_fname = f"{args.output_dir}/{args.problem}/cf_{dataset_gen.num_samples}samples_seed{args.random_seed}_base_{os.path.basename(args.base_dataset)}.pkl"
+    save_dataset(cf_dataset, output_fname)

generate_dataset.py

@@ -0,0 +1,117 @@
+from utils.data_utils.tsptw_dataset import TSPTWDataset
+from utils.data_utils.pctsp_dataset import PCTSPDataset
+from utils.data_utils.pctsptw_dataset import PCTSPTWDataset
+from utils.data_utils.cvrp_dataset import CVRPDataset
+from utils.data_utils.cvrptw_dataset import CVRPTWDataset
+from utils.utils import save_dataset
+
+def generate_dataset(num_samples, args):
+    if args.problem == "tsptw":
+        data_generator = TSPTWDataset(coord_dim=args.coord_dim,
+                                      num_samples=num_samples,
+                                      num_nodes=args.num_nodes,
+                                      random_seed=args.random_seed,
+                                      solver=args.solver, 
+                                      classifier=args.classifier,
+                                      annotation=args.annotation,
+                                      parallel=args.parallel,
+                                      num_cpus=args.num_cpus,
+                                      distribution=args.distribution)
+    elif args.problem == "pctsp":
+        data_generator = PCTSPDataset(coord_dim=args.coord_dim,
+                                      num_samples=num_samples,
+                                      num_nodes=args.num_nodes,
+                                      random_seed=args.random_seed,
+                                      solver=args.solver, 
+                                      classifier=args.classifier,
+                                      annotation=args.annotation,
+                                      parallel=args.parallel,
+                                      num_cpus=args.num_cpus,
+                                      penalty_factor=args.penalty_factor)
+    elif args.problem == "pctsptw":
+        data_generator = PCTSPTWDataset(coord_dim=args.coord_dim,
+                                        num_samples=num_samples,
+                                        num_nodes=args.num_nodes,
+                                        random_seed=args.random_seed,
+                                        solver=args.solver, 
+                                        classifier=args.classifier,
+                                        annotation=args.annotation,
+                                        parallel=args.parallel,
+                                        num_cpus=args.num_cpus,
+                                        penalty_factor=args.penalty_factor)
+    elif args.problem == "cvrp":
+        data_generator = CVRPDataset(coord_dim=args.coord_dim,
+                                     num_samples=num_samples,
+                                     num_nodes=args.num_nodes,
+                                     random_seed=args.random_seed,
+                                     solver=args.solver, 
+                                     classifier=args.classifier,
+                                     annotation=args.annotation,
+                                     parallel=args.parallel,
+                                     num_cpus=args.num_cpus)
+    elif args.problem == "cvrptw":
+        data_generator = CVRPTWDataset(coord_dim=args.coord_dim,
+                                       num_samples=num_samples,
+                                       num_nodes=args.num_nodes,
+                                       random_seed=args.random_seed,
+                                       solver=args.solver, 
+                                       classifier=args.classifier,
+                                       annotation=args.annotation,
+                                       parallel=args.parallel,
+                                       num_cpus=args.num_cpus)
+    else:
+        raise NotImplementedError
+
+    return data_generator.generate_dataset()
+
+if __name__ == "__main__":
+    import argparse
+    import os
+    import numpy as np
+    parser = argparse.ArgumentParser(description='')
+    # common settings
+    parser.add_argument("--problem",     type=str, default="tsptw")
+    parser.add_argument("--random_seed", type=int, default=1234)
+    parser.add_argument("--data_type",   type=str, nargs="*", default=["all"], help="data type: 'all' or combo. of ['train', 'valid', 'test'].")
+    parser.add_argument("--num_samples", type=int, nargs="*", default=[1000, 100, 100])
+    parser.add_argument("--num_nodes",   type=int, default=20)
+    parser.add_argument("--coord_dim",   type=int, default=2, help="only coord_dim=2 is supported for now.")
+    parser.add_argument("--solver",      type=str, default="ortools", help="solver that outputs a tour")
+    parser.add_argument("--classifier",  type=str, default="ortools", help="classifier for annotation")
+    parser.add_argument("--annotation",  action="store_true")
+    parser.add_argument("--parallel",    action="store_true")
+    parser.add_argument("--num_cpus",    type=int, default=os.cpu_count())
+    parser.add_argument("--output_dir",  type=str, default="data")
+    # for TSPTW
+    parser.add_argument("--distribution", type=str, default="da_silva")
+    # for PCTSP
+    parser.add_argument("--penalty_factor", type=float, default=3.)
+    args = parser.parse_args()
+
+    # 3d problems are not supported
+    assert args.coord_dim == 2, "only coord_dim=2 is supported for now."
+
+    # calc num. of total samples (train + valid + test samples)
+    if args.data_type[0] == "all":
+        assert len(args.num_samples) == 3, "please specify # samples for each of the three types (train/valid/test) when you set data_type 'all'. (e.g., --num_samples 1280000 1000 1000)"
+    else:
+        assert len(args.data_type) == len(args.num_samples), "please match # data_types and # elements in num_samples-arg"
+    num_samples = np.sum(args.num_samples)
+
+    # generate a dataset
+    dataset = generate_dataset(num_samples, args)
+
+    # split the dataset
+    if args.data_type[0] == "all":
+        types = ["train", "valid", "eval"]
+    else:
+        types = args.data_type
+    num_sample_list = args.num_samples
+    num_sample_list.insert(0, 0)
+    start = 0
+    for i, type_name in enumerate(types):
+        start += num_sample_list[i]
+        end = start + num_sample_list[i+1]
+        divided_datset = dataset[start:end]
+        output_fname = f"{args.output_dir}/{args.problem}/{type_name}_{args.problem}_{args.num_nodes}nodes_{num_sample_list[i+1]}samples_seed{args.random_seed}.pkl"
+        save_dataset(divided_datset, output_fname)

install_solvers.py

@@ -0,0 +1,68 @@
+import os
+import urllib.request
+import subprocess
+
+def _run(cmd, cwd):
+    subprocess.check_call(cmd, shell=True, cwd=cwd)
+
+def install_concorde():
+    QSOPT_A_URL  = "https://www.math.uwaterloo.ca/~bico/qsopt/beta/codes/PIC/qsopt.PIC.a"
+    QSOPT_H_URL  = "https://www.math.uwaterloo.ca/~bico/qsopt/beta/codes/PIC/qsopt.h"
+    CONCORDE_URL = "https://www.math.uwaterloo.ca/tsp/concorde/downloads/codes/src/co031219.tgz"
+
+    concorde_path = "models/solvers/concorde"
+    concorde_src_path = f"{concorde_path}/src"
+    os.makedirs(concorde_src_path, exist_ok=True)
+    # download qsopt, which is a dependency library
+    print("Downloading QSOPT...", end=" ", flush=True)
+    qsopt_path = f"{concorde_src_path}/qsopt"
+    qsopt_a_path = f"{qsopt_path}/qsopt.a"
+    qsopt_h_path = f"{qsopt_path}/qsopt.h"
+    os.makedirs(qsopt_path, exist_ok=True)
+    urllib.request.urlretrieve(QSOPT_A_URL, qsopt_a_path)
+    urllib.request.urlretrieve(QSOPT_H_URL, qsopt_h_path)
+    print("done")
+
+    # download concorde tsp
+    print("Downloading Concorde TSP...", end=" ", flush=True)
+    concorde_tgz_path = f"{concorde_src_path}/concorde.tgz"
+    urllib.request.urlretrieve(CONCORDE_URL, concorde_tgz_path)
+    print("done")
+
+    # build concorde
+    _run("tar -xzf concorde.tgz", concorde_src_path)
+    _run("mv concorde/* .", concorde_src_path)
+    _run("rm -r concorde.tgz concorde", concorde_src_path)
+    cflags = "-fPIC -O2 -g"
+    datadir = os.path.abspath(qsopt_path)
+    cmd = f"CFLAGS='{cflags}' ./configure --prefix {datadir} --with-qsopt={datadir}"
+    _run(cmd, concorde_src_path)
+    _run("make", concorde_src_path)
+
+def install_lkh():
+    LKH_URL = "http://webhotel4.ruc.dk/~keld/research/LKH-3/LKH-3.0.8.tgz"
+
+    lkh_path = "models/solvers/lkh"
+    lkh_src_path = f"{lkh_path}/src"
+    os.makedirs(lkh_src_path, exist_ok=True)
+
+    # download LKH
+    urllib.request.urlretrieve(LKH_URL, f"{lkh_src_path}/LKH-3.0.8.tgz")
+
+    # build LKH
+    _run("tar -xzf LKH-3.0.8.tgz", lkh_src_path)
+    _run("mv LKH-3.0.8/* .", lkh_src_path)
+    _run("rm -r LKH-3.0.8.tgz LKH-3.0.8", lkh_src_path)
+    _run("make", lkh_src_path)
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--installed_solvers", default="all", type=str, help="Solvers: [all, concorde, lkh]")
+    args = parser.parse_args()
+
+    if args.installed_solvers == "all" or args.installed_solvers == "concorde":
+        install_concorde()
+
+    if args.installed_solvers == "all" or args.installed_solvers == "lkh":
+        install_lkh()

models/cf_generator.py

@@ -0,0 +1,155 @@
+import torch.nn as nn
+import numpy as np
+
+class CFTourGenerator(nn.Module):
+    def __init__(self, cf_solver):
+        super().__init__()
+        self.solver = cf_solver
+        self.problem = cf_solver.problem
+
+    def forward(self, factual_tour, vehicle_id, cf_step, cf_next_node_id, node_feats, dist_matrix=None):
+        """
+        solve an input instance with visited edges fixed
+
+        Parameters
+        ----------
+        factual_tour: list [seq_length]
+        cf_step: int
+        cf_next_node_id: int
+        node_feats: 
+
+        Returns
+        -------
+        cf_tour: np.array [seq_length]
+        """
+        fixed_paths = self.get_fixed_paths(factual_tour, vehicle_id, cf_step, cf_next_node_id)
+        cf_tours = self.solver.solve(node_feats, fixed_paths, dist_matrix=dist_matrix)
+        if cf_tours is None:
+            return
+        if (cf_step > 0):
+            for vehicle_id, cf_tour in enumerate(cf_tours):
+                if cf_next_node_id in cf_tour:
+                    if cf_step == 1:
+                        if cf_tour[1] != cf_next_node_id:
+                            cf_tours[vehicle_id] = np.flipud(cf_tour)
+                        break
+                    else:
+                        if (factual_tour[vehicle_id][1] != cf_tour[1]):
+                            cf_tours[vehicle_id] = np.flipud(cf_tour) # make direction of the cf tour the same as factual one
+                        break
+        print("aaaa", cf_tours)
+        return cf_tours
+
+    def get_fixed_paths(self, factual_tour, vehicle_id, cf_step, cf_next_node_id):
+        visited_paths = np.append(factual_tour[vehicle_id][:cf_step], cf_next_node_id)
+        return visited_paths
+
+    # def get_avail_edges(self, factual_tour, cf_step, cf_next_node_id):
+    #     visited_paths = np.append(factual_tour[:cf_step], cf_next_node_id)
+    #     avail_edges = []
+    #     # add fixed edges
+    #     for i in range(len(visited_paths) - 1):
+    #         avail_edges.append([visited_paths[i], visited_paths[i + 1]])
+    #     print(avail_edges)
+
+    #     # add rest avaialbel edges
+    #     num_nodes = np.max(factual_tour) + 1
+    #     visited = np.array([0] * num_nodes)
+    #     for id in visited_paths:
+    #         visited[id] = 1
+    #     visited[factual_tour[0]] = 0
+    #     visited[cf_next_node_id] = 0
+    #     mask = visited < 1
+    #     node_id = np.arange(num_nodes)
+    #     feasible_node_id = node_id[mask]
+    #     for j in range(len(feasible_node_id) - 1):
+    #         for i in range(j + 1, len(feasible_node_id)):
+    #             if ((feasible_node_id[j] == factual_tour[0]) and (feasible_node_id[i] == cf_next_node_id)) or ((feasible_node_id[i] == factual_tour[0]) and (feasible_node_id[j] == cf_next_node_id)):
+    #                 continue
+    #             avail_edges.append([feasible_node_id[j], feasible_node_id[i]])
+    #     return np.array(avail_edges)
+
+#-----------
+# unit test
+#-----------
+if __name__ == "__main__":
+    import argparse
+    import random
+    import matplotlib.pyplot as plt
+    # FYI: 
+    # - https://yu-nix.com/archives/python-path-get/
+    # - https://www.delftstack.com/ja/howto/python/python-get-parent-directory/
+    # - https://stackoverflow.com/questions/2817264/how-to-get-the-parent-dir-location
+    import os
+    import sys
+    CURR_DIR = os.path.dirname(os.path.abspath(__file__))
+    PARENT_DIR = os.path.abspath(os.path.join(CURR_DIR, os.pardir))
+    sys.path.append(PARENT_DIR)
+    from utils.util_vis import visualize_factual_and_cf_tours
+    from lkh.lkh import LKH
+    from models.ortools.ortools import ORTools
+    from data_generator.tsptw.tsptw_dataset import generate_tsptw_instance
+
+    parser = argparse.ArgumentParser(description='')
+    # general settings
+    parser.add_argument("--problem", type=str, default="tsptw")
+    parser.add_argument("--random_seed", type=int, default=1234)
+    parser.add_argument("--num_samples", type=int, default=5)
+    parser.add_argument("--num_nodes", type=int, default=100)
+    parser.add_argument("--coord_dim", type=int, default=2)
+    # LKH settings
+    parser.add_argument("--max_trials", type=int, default=1000)
+    parser.add_argument("--lkh_dir", type=str, default="lkh", help="Path to the binary of LKH")
+    parser.add_argument("--io_dir", type=str, default="lkh_io_files")
+    args = parser.parse_args()
+
+    # models
+    # cf_solver = LKH(args.problem, args.max_trials, args.random_seed, lkh_dir=args.lkh_dir, io_dir=args.io_dir)
+    cf_solver = ORTools(args.problem)
+    cf_generator = CFTourGenerator(cf_solver)
+
+    # dataset
+    if args.problem == "tsp":
+        np.random.seed(args.random_seed)
+        node_feats = np.random.uniform(size=[args.num_samples, args.num_nodes, args.coord_dim])
+    elif args.problem == "tsptw":
+        coords, time_window, grid_size = generate_tsptw_instance(num_nodes=args.num_nodes, grid_size=100, max_tw_gap=10, max_tw_size=1000, is_integer_instance=True, da_silva_style=True)
+        node_feats = np.concatenate([coords, time_window], -1)
+        node_feats = node_feats[None, :, :]
+
+    # function ot automatically generate couterfactual visit 
+    def get_random_cf_visit(factual_tour, random_seed=1234):
+        # random.seed(random_seed)
+        num_nodes = np.max(factual_tour) + 1
+        step = random.randrange(len(factual_tour) - 2) # remove the last step (returning to the start-point)
+        visited = np.array([0] * num_nodes)
+        for i in range(step+1):
+            visited[factual_tour[i]] = 1
+        mask = visited < 1
+        node_id = np.arange(num_nodes)
+        feasible_node_id = node_id[mask]
+        cf_next_id = random.choice(feasible_node_id) # select counterfactual
+        return step, cf_next_id
+
+    for i in range(len(node_feats)):
+        # obtain a factual tour
+        factual_tour = cf_solver.solve(node_feats[i])
+
+        # counterfactual visit
+        cf_step, cf_next_node_id = get_random_cf_visit(factual_tour, random_seed=args.random_seed)
+
+        print(cf_step, cf_next_node_id)
+        # obtain a counterfactual tour
+        cf_tour = cf_generator(factual_tour, cf_step, cf_next_node_id, node_feats[i])
+
+        print(factual_tour)
+        print(cf_tour)
+
+        # visualize the factual and counterfactual tours
+        if args.problem == "tsp":
+            coords = node_feats[i]
+        elif args.problem == "tsptw":
+            coord_dim = 2
+            coords = node_feats[i, :, :coord_dim]
+        visualize_factual_and_cf_tours(factual_tour, cf_tour, coords, cf_step, f"test{i}.png")
+        break

models/classifiers/general_classifier.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn as nn
+import argparse
+import json
+import os
+from models.classifiers.predictor import DecisionPredictor
+from models.classifiers.meaningless_models import FixedClassPredictor, RandomPredictor
+from models.classifiers.rule_based_models import kNearestPredictor
+from models.classifiers.ground_truth.ground_truth import GroundTruth
+
+class GeneralClassifier(nn.Module):
+    def __init__(self, problem, model_type):
+        super().__init__()
+        self.model_type = model_type
+        self.problem = problem
+        self.model = self.get_model(problem, model_type)
+
+    def change_model(self, problem, model_type):
+        if self.model_type != model_type or self.problem != problem:
+            self.model_type = model_type
+            self.problem = problem
+            self.model = self.get_model(problem, model_type)
+
+    def get_model(self, problem, model_type):
+        if model_type == "gnn":
+            model_path = "checkpoints/model_20230309_101058/model_epoch4.pth"
+            params = argparse.ArgumentParser()
+            model_dir = os.path.split(model_path)[0]
+            with open(f"{model_dir}/cmd_args.dat", "r") as f:
+                params.__dict__ = json.load(f)
+            model = DecisionPredictor(params.problem,
+                                      params.emb_dim,
+                                      params.num_mlp_layers,
+                                      params.num_classes,
+                                      params.dropout)
+            model.load_state_dict(torch.load(model_path))
+            return model
+        elif model_type == "gt(ortools)":
+            return GroundTruth(problem, solver_type="ortools")
+        elif model_type == "gt(lkh)":
+            return GroundTruth(problem, solver_type="lkh")
+        elif model_type == "gt(concorde)":
+            return GroundTruth(problem, solver_type="concorde")
+        elif model_type == "random":
+            return RandomPredictor(num_classes=2)
+        elif model_type == "fixed":
+            predicted_class = 0
+            return FixedClassPredictor(predicted_class=predicted_class, num_classes=2)
+        elif model_type == "knn":
+            k = 5
+            k_type = "num"
+            return kNearestPredictor(problem, k, k_type)
+        else:
+            assert False, f"Invalid model type: {model_type}"
+    
+    def get_inputs(self, tour, first_explained_step, node_feats, dist_matrix=None):
+        return self.model.get_inputs(tour, first_explained_step, node_feats, dist_matrix)
+    
+    def forward(self, inputs):
+        return self.model(inputs)

models/classifiers/ground_truth/ground_truth.py

@@ -0,0 +1,35 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from models.classifiers.ground_truth.ground_truth_tsptw import GroundTruthTSPTW
+from models.classifiers.ground_truth.ground_truth_pctsp import GroundTruthPCTSP
+from models.classifiers.ground_truth.ground_truth_pctsptw import GroundTruthPCTSPTW
+from models.classifiers.ground_truth.ground_truth_cvrp import GroundTruthCVRP
+from models.classifiers.ground_truth.ground_truth_cvrptw import GroundTruthCVRPTW
+
+class GroundTruth(nn.Module):
+    def __init__(self, problem, solver_type):
+        super().__init__()
+        self.problem = problem
+        self.solver_type = solver_type
+        if problem == "tsptw":
+            self.ground_truth = GroundTruthTSPTW(solver_type)
+        elif problem == "pctsp":
+            self.ground_truth = GroundTruthPCTSP(solver_type)
+        elif problem == "pctsptw":
+            self.ground_truth = GroundTruthPCTSPTW(solver_type)
+        elif problem == "cvrp":
+            self.ground_truth = GroundTruthCVRP(solver_type)
+        elif problem == "cvrptw":
+            self.ground_truth = GroundTruthCVRPTW(solver_type)
+        else:
+            raise NotImplementedError
+
+    def forward(self, inputs, annotation=False, parallel=False):
+        return self.ground_truth(inputs, annotation, parallel)
+    
+    def get_inputs(self, tour, first_explained_step, node_feats, dist_matrix=None):
+        return self.ground_truth.get_inputs(tour, first_explained_step, node_feats, dist_matrix)
+    
+    def solve(self, step, input_tour, node_feats, instance_name=None):
+        return self.ground_truth.solve(step, input_tour, node_feats, instance_name)

models/classifiers/ground_truth/ground_truth_base.py

@@ -0,0 +1,285 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import os
+import multiprocessing
+from models.solvers.general_solver import GeneralSolver
+from utils.utils import calc_tour_length
+
+def get_visited_mask(tour, step, node_feats, dist_matrix=None):
+    """
+    Visited nodes -> feasible, Unvisited nodes -> infeasible.
+    When solving a problem with visited_paths fixed, they should be included to the solution.
+    Therefore, visited nodes are set to feasible nodes.
+    """
+    if dist_matrix is not None:
+        num_nodes = len(dist_matrix)
+    else:
+        num_nodes = len(node_feats["coords"])
+    visited = np.isin(np.arange(num_nodes), tour[:step])
+    return visited
+
+def get_tw_mask(tour, step, node_feats, dist_matrix=None):
+    """
+    Nodes whose tw exceeds current_time -> infeasible, otherwise -> feasible.
+
+    Parameters
+    ----------
+    tour: list [seq_length]
+    step: int
+    node_feats: dict of np.array
+
+    Returns
+    -------
+    mask_tw: np.array [num_nodes]
+    """
+    node_feats = node_feats.copy()
+    time_window = node_feats["time_window"]
+    if dist_matrix is not None:
+        num_nodes = len(dist_matrix)
+        curr_time = 0.0
+        not_exceed_tw = np.ones(num_nodes).astype(np.int32)
+        for i in range(1, step):
+            prev_id = tour[i - 1]
+            curr_id = tour[i]
+            travel_time = dist_matrix[prev_id, curr_id]
+            # assert curr_time + travel_time < time_window[curr_id, 1], f"Invalid tour! arrival_time: {curr_time + travel_time}, time_window: {time_window[curr_id]}"
+            if curr_time + travel_time < time_window[curr_id, 0]:
+                curr_time = time_window[curr_id, 0].copy()
+            else:
+                curr_time += travel_time
+        curr_time = curr_time + dist_matrix[tour[step-1]] # [num_nodes] TODO: check
+    else:
+        coords = node_feats["coords"]
+        num_nodes = len(coords)
+        curr_time = 0.0
+        not_exceed_tw = np.ones(num_nodes).astype(np.int32)
+        for i in range(1, step):
+            prev_id = tour[i - 1]
+            curr_id = tour[i]
+            travel_time = np.linalg.norm(coords[prev_id] - coords[curr_id])
+            # assert curr_time + travel_time < time_window[curr_id, 1], f"Invalid tour! arrival_time: {curr_time + travel_time}, time_window: {time_window[curr_id]}"
+            if curr_time + travel_time < time_window[curr_id, 0]:
+                curr_time = time_window[curr_id, 0].copy()
+            else:
+                curr_time += travel_time
+        curr_time = curr_time + np.linalg.norm(coords[tour[step-1]][None, :] - coords, axis=-1) # [num_nodes] TODO: check
+    not_exceed_tw[curr_time > time_window[:, 1]] = 0
+    not_exceed_tw = not_exceed_tw > 0
+    return not_exceed_tw
+
+def get_cap_mask(tour, step, node_feats):
+    num_nodes = len(node_feats["coords"])
+    demands = node_feats["demand"]
+    remaining_cap = node_feats["capacity"].copy()
+    less_than_cap = np.ones(num_nodes).astype(np.int32)
+    for i in range(step):
+        remaining_cap -= demands[tour[i]]
+    less_than_cap[remaining_cap < demands] = 0
+    less_than_cap = less_than_cap > 0
+    return less_than_cap
+
+def get_pc_mask(tour, step, node_feats):
+    """
+    Mask for Price collecting problems (e.g., PCTSP, PCTSPTW, PCCVRP, PCCVRPTW, ...)
+
+    Returns
+    -------
+    not_exceed_max_length
+    """
+    large_value = 1e+5
+    coords = node_feats["coords"]
+    max_length = (node_feats["max_length"] * large_value).astype(np.int64)
+    tour_length = 0
+    for i in range(1, step):
+        prev_id = tour[i - 1]
+        curr_id = tour[i]
+        tour_length += (np.linalg.norm(coords[prev_id] - coords[curr_id]) * large_value).astype(np.int64)
+    curr_to_next  = (np.linalg.norm(coords[tour[step-1]][None, :] - coords, axis=-1) * large_value).astype(np.int64) # [num_nodes]
+    next_to_depot = (np.linalg.norm(coords[tour[0]][None, :] - coords, axis=-1) * large_value).astype(np.int64) # [num_nodes]
+    not_exceed_max_length = (tour_length + curr_to_next + next_to_depot) <= max_length # [num_nodes]
+    return not_exceed_max_length
+
+def analyze_tour(tour, node_feats):
+    coords = node_feats["coords"]
+    time_window = node_feats["time_window"]
+    curr_time = 0
+    for i in range(1, len(tour)):
+        prev_id = tour[i - 1]
+        curr_id = tour[i]
+        travel_time = np.linalg.norm(coords[prev_id] - coords[curr_id])
+        valid = curr_time + travel_time < time_window[curr_id, 1]
+        print(f"visit #{i}: {prev_id} -> {curr_id}, travel_time: {travel_time}, arrival_time: {curr_time + travel_time}, time_window: {time_window[curr_id]}, valid: {valid}")
+        if curr_time + travel_time < time_window[curr_id, 0]:
+            curr_time = time_window[curr_id, 0]
+        else:
+            curr_time += travel_time
+
+FAIL_FLAG = -1
+class GroundTruthBase(nn.Module):
+    def __init__(self, problem, compared_problems, solver_type):
+        """
+        Parameters
+        ----------
+
+        """
+        super().__init__()
+        self.problem = problem
+        self.compared_problems = compared_problems
+        self.num_compared_problems = len(compared_problems)
+        self.solver_type = solver_type
+        self.solvers = []
+        for i in range(self.num_compared_problems):
+            # TODO:
+            self.solvers.append(GeneralSolver(self.compared_problems[i], self.solver_type, scaling=False))
+
+    def forward(self, inputs, annotation=False, parallel=True):
+        """
+        Parameters
+        ----------
+        inputs: dict
+            tour: 2d list [num_vehicles x seq_length]
+            first_explained_step: int
+            node_feats: dict of np.array
+        annotation: bool
+            please set it True when annotating data
+
+        Returns
+        -------
+        labels: 
+        probs:  torch.tensor [batch_size (num_vehicles) x max_seq_length x num_classes]
+        """
+        input_tours = inputs["tour"]
+        node_feats = inputs["node_feats"]
+        dist_matrix = inputs["dist_matrix"]
+        first_explained_step = inputs["first_explained_step"]
+        num_vehicles = len(input_tours)
+        if annotation:
+            labels = [[] for _ in range(num_vehicles)]
+            for vehicle_id in range(num_vehicles):
+                input_tour = input_tours[vehicle_id]
+                # analyze_tour(input_tour, node_feats)
+                for step in range(first_explained_step + 1, len(input_tour)):
+                    _, __, label = self.label_path(vehicle_id, step, input_tour, node_feats)
+                    if label == FAIL_FLAG:
+                        return
+                    labels[vehicle_id].append((step, label))
+            return labels
+        else:
+            if parallel:
+                labels = [[-1] * (len(range(first_explained_step+1, len(input_tours[vehicle_id])))) for vehicle_id in range(num_vehicles)]
+                num_cpus = os.cpu_count()
+                with multiprocessing.Pool(num_cpus) as pool:
+                    for vehicle_id, step, label in pool.starmap(self.label_path, [(vehicle_id, step, input_tours[vehicle_id], node_feats, dist_matrix)
+                                                                                  for vehicle_id in range(num_vehicles) 
+                                                                                  for step in range(first_explained_step+1, len(input_tours[vehicle_id]))]):
+                        labels[vehicle_id][step-(first_explained_step+1)] = label
+            else:
+                labels = [[-1] * (len(range(first_explained_step+1, len(input_tours[vehicle_id])))) for vehicle_id in range(num_vehicles)]
+                for vehicle_id in range(num_vehicles):
+                    for step in range(first_explained_step+1, len(input_tours[vehicle_id])):
+                        vehicle_id, step, label = self.label_path(vehicle_id, step, input_tours[vehicle_id], node_feats, dist_matrix)
+                        labels[vehicle_id][step-(first_explained_step+1)] = label
+            # validate labels
+            for vehicle_id in range(num_vehicles):
+                assert (len(input_tours[vehicle_id]) - 1) == len(labels[vehicle_id]), f"vehicle_id={vehicle_id}, {input_tours}, {labels}"
+            return labels
+            # labels = [torch.LongTensor(label) for label in labels] # [num_vehicles x seq_length]
+            # labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True) # [num_vehicles x max_seq_length]
+            # probs = torch.zeros((labels.size(0), labels.size(1), self.num_compared_problems+1)) # [num_vehicles x max_seq_length x (num_compared_problems+1)]
+            # probs.scatter_(-1, labels.unsqueeze(-1).expand_as(probs), 1.0)
+            # return probs
+    
+    def label_path(self, vehicle_id, step, input_tour, node_feats, dist_matrix=None):
+        compared_tour_list = [[] for _ in range(self.num_compared_problems)]
+        visited_path = input_tour[:step].copy()
+        new_node_id, new_node_feats, new_dist_matrix = self.get_feasible_nodes(input_tour, step, node_feats, dist_matrix)
+        new_visited_path = np.array(list(map(lambda x: np.where(new_node_id==x)[0].item(), visited_path)))
+        for i in range(self.num_compared_problems):
+            # TODO: in CVRPTW / PCCVRPTW, need to modify classification of the first and last paths
+            compared_tours = self.solvers[i].solve(new_node_feats, new_visited_path, new_dist_matrix)
+            if compared_tours is None:
+                return vehicle_id, step, FAIL_FLAG
+            compared_tour = None
+            for compared_tour_tmp in compared_tours:
+                if new_visited_path[-1] in compared_tour_tmp:
+                    compared_tour = compared_tour_tmp
+                    break
+            assert compared_tour is not None, f"Found no appropriate vhiecle. {compared_tours}, {new_visited_path}"
+            compared_tour = np.array(list(map(lambda x: new_node_id[x], compared_tour)))
+            if (step > 0) and (compared_tour[1] != input_tour[1]):
+                compared_tour = np.flipud(compared_tour) # make direction of the cf tour the same as factual one
+            compared_tour_list[i] = compared_tour
+        # print("fixed_paths  :", visited_path)
+        # print("input_tour   :", input_tour)
+        # print("compared_tour:", compared_tour)
+        # print()
+        # annotation
+        label = self.get_label(input_tour, compared_tour_list, step)
+        return vehicle_id, step, label
+
+    def solve(self, step, input_tour, node_feats, instance_name=None):
+        compared_tours = {} 
+        visited_path = input_tour[:step].copy()
+        new_node_id, new_node_feats = self.get_feasible_nodes(input_tour, step, node_feats)
+        new_visited_path = np.array(list(map(lambda x: np.where(new_node_id==x)[0].item(), visited_path)))
+        for i, compared_problem in enumerate(self.compared_problems):
+            compared_tours[compared_problem] = self.solvers[i].solve(new_node_feats, new_visited_path, instance_name)
+            compared_tours[compared_problem] = list(map(lambda compared_tour: list(map(lambda x: new_node_id[x], compared_tour)), compared_tours[compared_problem]))
+            compared_tours[compared_problem] = list(map(lambda compared_tour: calc_tour_length(compared_tour, node_feats["coords"]), compared_tours[compared_problem]))
+        return compared_tours
+
+    def get_label(self, input_tour, compared_tours, step):
+        for i in range(self.num_compared_problems):
+            compared_tour = compared_tours[i]
+            if input_tour[step] == compared_tour[step]:
+                return i
+        return self.num_compared_problems
+
+    def get_inputs(self, tour, first_explained_step, node_feats, dist_matrix=None):
+        input_features = {
+            "tour": tour, 
+            "first_explained_step": first_explained_step, 
+            "node_feats": node_feats,
+            "dist_matrix": dist_matrix
+        }
+        return input_features
+
+    def get_feasible_nodes(self, tour, step, node_feats, dist_matrix=None):
+        """
+        Parameters
+        ----------
+        tour: np.array [seq_length]
+        step: int
+        node_feats: np.array [num_nodes x node_dim]
+
+        Returns
+        -------
+        new_node_id: np.array [num_feasible_nodes]
+        new_node_feats: dict of np.array [num_feasible_nodes x coord_dim]
+        """
+        if dist_matrix is not None:
+            num_nodes = len(dist_matrix)
+        else:
+            num_nodes = len(node_feats["coords"])
+        mask = self.get_mask(tour, step, node_feats, dist_matrix)
+        node_id = np.arange(num_nodes)
+        new_node_id = node_id[mask].copy()
+        new_node_feats = {
+            key: node_feat[mask].copy() 
+                 if key in ["coords", "time_window", "demand", "penalties", "prizes"] else
+                 node_feat.copy()
+            for key, node_feat in node_feats.items()
+        }
+        if dist_matrix is not None:
+            delete_id = node_id[~mask]
+            new_dist_matrix = np.delete(np.delete(dist_matrix, delete_id, 0), delete_id, 1)
+        else:
+            new_dist_matrix = None
+        return new_node_id, new_node_feats, new_dist_matrix
+
+    def get_mask(self, tour, step, node_feats, dist_matrix=None):
+        raise NotImplementedError
+    
+    def check_feasibility(self, tour, node_feats):
+        raise NotImplementedError

models/classifiers/ground_truth/ground_truth_cvrp.py

@@ -0,0 +1,14 @@
+from models.classifiers.ground_truth.ground_truth_base import GroundTruthBase
+from models.classifiers.ground_truth.ground_truth_base import get_cap_mask, get_visited_mask
+
+class GroundTruthCVRP(GroundTruthBase):
+    def __init__(self, solver_type):
+        problem = "cvrp"
+        compared_problems = ["tsp"]
+        super().__init__(problem, compared_problems, solver_type)
+    
+    # @override
+    def get_mask(self, tour, step, node_feats):
+        visited = get_visited_mask(tour, step, node_feats)
+        less_than_cap = get_cap_mask(tour, step, node_feats)
+        return visited | less_than_cap

models/classifiers/ground_truth/ground_truth_cvrptw.py

@@ -0,0 +1,15 @@
+from models.classifiers.ground_truth.ground_truth_base import GroundTruthBase
+from models.classifiers.ground_truth.ground_truth_base import get_cap_mask, get_visited_mask, get_tw_mask
+
+class GroundTruthCVRPTW(GroundTruthBase):
+    def __init__(self, solver_type):
+        problem = "cvrptw"
+        compared_problems = ["tsp", "cvrp"]
+        super().__init__(problem, compared_problems, solver_type)
+    
+    # @override
+    def get_mask(self, tour, step, node_feats):
+        visited = get_visited_mask(tour, step, node_feats)
+        less_than_cap = get_cap_mask(tour, step, node_feats)
+        not_exceed_tw = get_tw_mask(tour, step, node_feats)
+        return visited | (less_than_cap & not_exceed_tw)

models/classifiers/ground_truth/ground_truth_pctsp.py

@@ -0,0 +1,16 @@
+import numpy as np
+from models.classifiers.ground_truth.ground_truth_base import GroundTruthBase
+from models.classifiers.ground_truth.ground_truth_base import get_visited_mask, get_pc_mask
+
+class GroundTruthPCTSP(GroundTruthBase):
+    def __init__(self, solver_type):
+        problem = "pctsp"
+        compared_problems = ["tsp"]
+        super().__init__(problem, compared_problems, solver_type)
+    
+    # @override
+    def get_mask(self, tour, step, node_feats):
+        # visited = get_visited_mask(tour, step, node_feats)
+        # not_exceed_max_length = get_pc_mask(tour, step, node_feats)
+        num_nodes = len(node_feats["coords"])
+        return np.full(num_nodes, True)

models/classifiers/ground_truth/ground_truth_pctsptw.py

@@ -0,0 +1,15 @@
+import numpy as np
+from models.classifiers.ground_truth.ground_truth_base import GroundTruthBase
+from models.classifiers.ground_truth.ground_truth_base import get_visited_mask, get_tw_mask
+
+class GroundTruthPCTSPTW(GroundTruthBase):
+    def __init__(self, solver_type):
+        problem = "pctsptw"
+        compared_problems = ["tsp", "pctsp"]
+        super().__init__(problem, compared_problems, solver_type)
+    
+    # @override
+    def get_mask(self, tour, step, node_feats):
+        visited = get_visited_mask(tour, step, node_feats)
+        not_exceed_tw = get_tw_mask(tour, step, node_feats)
+        return visited | not_exceed_tw

models/classifiers/ground_truth/ground_truth_tsptw.py

@@ -0,0 +1,14 @@
+from models.classifiers.ground_truth.ground_truth_base import GroundTruthBase
+from models.classifiers.ground_truth.ground_truth_base import get_tw_mask, get_visited_mask
+
+class GroundTruthTSPTW(GroundTruthBase):
+    def __init__(self, solver_type):
+        problem = "tsptw"
+        compared_problems = ["tsp"]
+        super().__init__(problem, compared_problems, solver_type)
+    
+    # @override
+    def get_mask(self, tour, step, node_feats, dist_matrix=None):
+        visited = get_visited_mask(tour, step, node_feats, dist_matrix)
+        not_exceed_tw = get_tw_mask(tour, step, node_feats, dist_matrix)
+        return visited | not_exceed_tw

models/classifiers/meaningless_models.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn as nn
+
+class RandomPredictor(nn.Module):
+    def __init__(self, num_classes):
+        super().__init__()
+        self.num_classes = num_classes
+
+    def forward(self, inputs):
+        """
+        Parameters
+        ----------
+        inputs: int or dict
+            batch_size or dict of input features
+
+        Returns
+        -------
+        probs: torch.tensor [batch_size x num_classes]
+        """
+        batch_size =  inputs if isinstance(inputs, int) else inputs["curr_node_id"].size(0)
+        ranom_index = torch.randint(self.num_classes, (batch_size, self.num_classes))
+        probs = torch.zeros(batch_size, self.num_classes).to(torch.float)
+        probs.scatter_(-1, ranom_index, 1.0)
+        return probs
+
+    def get_inputs(self, tour, first_explained_step, node_feats):
+        return len(tour[first_explained_step:-1])
+
+class FixedClassPredictor(nn.Module):
+    def __init__(self, predicted_class, num_classes):
+        """
+        Paramters
+        ---------
+        predicted_class: int
+            a class that this predictor always predicts
+        num_classes: int
+            number of classes 
+        """
+        super().__init__()
+        self.predicted_class = predicted_class
+        self.num_classes = num_classes
+        assert predicted_class < num_classes, f"predicted_class should be 0 - {num_classes}."
+
+    def forward(self, inputs):
+        """
+        Parameters
+        ----------
+        inputs: int or dict
+            batch_size or dict of input features
+
+        Returns
+        -------
+        probs: torch.tensor [batch_size x num_classes]
+        """
+        batch_size =  inputs if isinstance(inputs, int) else inputs["curr_node_id"].size(0)
+        index = torch.full((batch_size, self.num_classes), self.predicted_class)
+        probs = torch.zeros(batch_size, self.num_classes).to(torch.float)
+        probs.scatter_(-1, index, 1.0)
+        return probs
+    
+    def get_inputs(self, tour, first_explained_step, node_feats):
+        return len(tour[first_explained_step:-1])

models/classifiers/nn_classifiers/attention_graph_encoder.py

@@ -0,0 +1,95 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+class AttentionGraphEncoder(nn.Module):
+    def __init__(self, coord_dim, node_dim, state_dim, emb_dim, dropout):
+        super().__init__()
+        self.coord_dim = coord_dim
+        self.node_dim  = node_dim 
+        self.emb_dim   = emb_dim
+        self.state_dim = state_dim
+        self.norm_factor = 1 / math.sqrt(emb_dim)
+
+        # initial embedding
+        self.init_linear_node  = nn.Linear(node_dim, emb_dim)
+        self.init_linear_depot = nn.Linear(coord_dim, emb_dim)
+        if state_dim > 0:
+            self.init_linear_state = nn.Linear(state_dim, emb_dim)
+        
+        # An attention layer
+        self.w_q = nn.Parameter(torch.FloatTensor((2 + int(state_dim > 0)) * emb_dim, emb_dim))
+        self.w_k = nn.Parameter(torch.FloatTensor(2 * emb_dim, emb_dim))
+        self.w_v = nn.Parameter(torch.FloatTensor(2 * emb_dim, emb_dim))
+
+        # Dropout
+        self.dropout = nn.Dropout(dropout)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for param in self.parameters():
+            stdv = 1. / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+
+    def forward(self, inputs):
+        """
+        Paramters
+        ---------
+        inputs: dict
+            curr_node_id: torch.LongTensor [batch_size x 1]
+            next_node_id: torch.LongTensor [batch_size x 1]
+            node_feat: torch.FloatTensor [batch_size x num_nodes x node_dim]
+            mask: torch.LongTensor [batch_size x num_nodes]
+            state: torch.FloatTensor [batch_size x state_dim]
+
+        Returns
+        -------
+        h: torch.tensor [batch_size x emb_dim]
+            graph embeddings 
+        """
+        #----------------
+        # input features
+        #----------------
+        curr_node_id = inputs["curr_node_id"]
+        next_node_id = inputs["next_node_id"]
+        node_feat = inputs["node_feats"]
+        mask = inputs["mask"]
+        state = inputs["state"]
+
+        #---------------------------
+        # initial linear projection
+        #---------------------------
+        node_emb  = self.init_linear_node(node_feat[:, 1:, :]) # [batch_size x num_loc x emb_dim]
+        depot_emb = self.init_linear_depot(node_feat[:, 0:1, :2]) # [batch_size x 1 x emb_dim]
+        new_node_feat = torch.cat((depot_emb, node_emb), 1) # [batch_size x num_nodes x emb_dim]
+        new_node_feat = self.dropout(new_node_feat)
+
+        #---------------
+        # preprocessing
+        #---------------
+        batch_size = curr_node_id.size(0)
+        curr_emb = new_node_feat.gather(1, curr_node_id[:, None, None].expand(batch_size, 1, self.emb_dim))
+        next_emb = new_node_feat.gather(1, next_node_id[:, None, None].expand(batch_size, 1, self.emb_dim))
+        if state is not None and self.state_dim > 0:
+            state_emb = self.init_linear_state(state) # [batch_size x emb_dim]
+            input_q  = torch.cat((curr_emb, next_emb, state_emb[:, None, :]), -1) # [batch_size x 1 x (3*emb_dim)]
+        else:
+            input_q  = torch.cat((curr_emb, next_emb), -1) # [batch_size x 1 x (2*emb_dim)]
+        input_kv = torch.cat((curr_emb.expand_as(new_node_feat), new_node_feat), -1) # [batch_size x num_nodes x (2*emb_dim)]
+
+        #--------------------
+        # An attention layer
+        #--------------------
+        q = torch.matmul(input_q,  self.w_q) # [batch_size x 1 x emb_dim]
+        k = torch.matmul(input_kv, self.w_k) # [batch_size x num_nodes x emb_dim]
+        v = torch.matmul(input_kv, self.w_v) # [batch_size x num_nodes x emb_dim]
+        compatibility = self.norm_factor * torch.matmul(q, k.transpose(-2, -1)) # [batch_size x 1 x num_nodes]
+        compatibility[(~mask).unsqueeze(1).expand_as(compatibility)] = -math.inf
+        attn = torch.softmax(compatibility, dim=-1)
+        h = torch.matmul(attn, v) # [batch_size x 1 x emb_dim]
+        h = h.squeeze(1) # [batch_size x emb_dim]
+
+        return h

models/classifiers/nn_classifiers/decoders/lstm_decoder.py

@@ -0,0 +1,57 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+class LSTMDecoder(nn.Module):
+    def __init__(self, emb_dim, num_mlp_layers, num_classes, dropout):
+        super().__init__()
+        self.num_mlp_layers = num_mlp_layers
+    
+        # LSTM
+        self.lstm = nn.LSTM(emb_dim, emb_dim, batch_first=True)
+
+        # Decoder (MLP)
+        self.mlp = nn.ModuleList()
+        for _ in range(num_mlp_layers):
+            self.mlp.append(nn.Linear(emb_dim, emb_dim, bias=True))
+        self.mlp.append(nn.Linear(emb_dim, num_classes, bias=True))
+
+        # Dropout
+        self.dropout = nn.Dropout(dropout)
+
+        # Initializing weights
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for param in self.parameters():
+            stdv = 1. / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+
+    def forward(self, graph_emb):
+        """
+        Paramters
+        ---------
+        graph_emb: torch.tensor [batch_size x max_seq_length x emb_dim]
+
+        Returns
+        -------
+        probs: torch.tensor [batch_size x max_seq_length x num_classes]
+            probabilities of classes
+        """
+        #---------------
+        # LSTM encoding
+        #---------------
+        h, _ = self.lstm(graph_emb) # [batch_size x max_seq_length x emb_dim]
+
+        #----------
+        # Decoding
+        #----------
+        for i in range(self.num_mlp_layers):
+            h = self.dropout(h)
+            h = torch.relu(self.mlp[i](h))
+        h = self.dropout(h)
+        logits = self.mlp[-1](h)
+        probs = F.log_softmax(logits, dim=-1)
+        return probs

models/classifiers/nn_classifiers/decoders/mha_decoder.py

@@ -0,0 +1,74 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        position = torch.arange(max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+        pe = torch.zeros(1, max_len, d_model) # batch_first
+        pe[0, :, 0::2] = torch.sin(position * div_term)
+        pe[0, :, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Arguments:
+            x: Tensor, shape ``[batch_size x max_seq_length x embedding_dim]``
+        """
+        x = x + self.pe[:, :x.size(1), :]
+        return self.dropout(x)
+
+class SelfMHADecoder(nn.Module):
+    def __init__(self, emb_dim: int, num_heads: int, num_mha_layers: int, num_classes: int, dropout: float, pos_encoder: str = None, max_len: int = 100):
+        super().__init__()
+        self.num_mha_layers = num_mha_layers
+    
+        # positional encoding
+        self.pos_encoder_type = pos_encoder
+        if pos_encoder == "sincos":
+            self.pos_encoder = PositionalEncoding(d_model=emb_dim, dropout=dropout, max_len=max_len)
+
+        # MHA blocks
+        mha_layer = nn.TransformerEncoderLayer(d_model=emb_dim, 
+                                               nhead=num_heads,
+                                               dim_feedforward=emb_dim,
+                                               dropout=dropout,
+                                               batch_first=True)
+        self.mha = nn.TransformerEncoder(mha_layer, num_layers=num_mha_layers)
+
+        # linear projection for adjusting out_dim to num_classes
+        self.out_linear = nn.Linear(emb_dim, num_classes, bias=True)
+
+        # Initializing weights
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for param in self.parameters():
+            stdv = 1. / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+
+    def forward(self, edge_emb):
+        """
+        Paramters
+        ---------
+        graph_emb: torch.tensor [batch_size x max_seq_length x emb_dim]
+
+        Returns
+        -------
+        probs: torch.tensor [batch_size x max_seq_length x num_classes]
+            probabilities of classes
+        """
+        #---------------
+        # MHA decoding
+        #---------------
+        if self.pos_encoder_type == "sincos":
+            edge_emb = self.pos_encoder(edge_emb)
+        h = self.mha(edge_emb, is_causal=True) # [batch_size x max_seq_length x emb_dim]
+        logits = self.out_linear(h)
+        probs = F.log_softmax(logits, dim=-1)
+        return probs

models/classifiers/nn_classifiers/decoders/mlp_decoder.py

@@ -0,0 +1,50 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+class MLPDecoder(nn.Module):
+    def __init__(self, emb_dim, num_mlp_layers, num_classes, dropout):
+        super().__init__()
+        self.num_mlp_layers = num_mlp_layers
+
+        # Decoder (MLP)
+        self.mlp = nn.ModuleList()
+        for _ in range(num_mlp_layers):
+            self.mlp.append(nn.Linear(emb_dim, emb_dim, bias=True))
+        self.mlp.append(nn.Linear(emb_dim, num_classes, bias=True))
+
+        # Dropout
+        self.dropout = nn.Dropout(dropout)
+
+        # Initializing weights
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for param in self.parameters():
+            stdv = 1. / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+
+    def forward(self, graph_emb):
+        """
+        Paramters
+        ---------
+        graph_emb: torch.tensor [batch_size x emb_dim]
+
+        Returns
+        -------
+        probs: torch.tensor [batch_size x num_classes]
+            probabilities of classes
+        """
+        #----------
+        # Decoding
+        #----------
+        h = graph_emb
+        for i in range(self.num_mlp_layers):
+            h = self.dropout(h)
+            h = torch.relu(self.mlp[i](h))
+        h = self.dropout(h)
+        logits = self.mlp[-1](h)
+        probs = F.log_softmax(logits, dim=-1)
+        return probs

models/classifiers/nn_classifiers/encoders/attn_edge_encoder.py

@@ -0,0 +1,81 @@
+import math
+import torch
+import torch.nn as nn
+
+class AttentionEdgeEncoder(nn.Module):
+    def __init__(self, state_dim, emb_dim, dropout):
+        super().__init__()
+        self.state_dim = state_dim
+        self.emb_dim   = emb_dim
+        self.norm_factor = 1 / math.sqrt(emb_dim)
+
+        # initial embedding for state
+        if state_dim > 0:
+            self.init_linear_state = nn.Linear(state_dim, emb_dim)
+        
+        # An attention layer
+        self.w_q = nn.Parameter(torch.FloatTensor((2 + int(state_dim > 0)) * emb_dim, emb_dim))
+        self.w_k = nn.Parameter(torch.FloatTensor(2 * emb_dim, emb_dim))
+        self.w_v = nn.Parameter(torch.FloatTensor(2 * emb_dim, emb_dim))
+
+        # out linear layer
+        self.out_linear = nn.Linear(emb_dim, emb_dim)
+
+        # Dropout
+        self.dropout = nn.Dropout(dropout)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for param in self.parameters():
+            stdv = 1. / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+    
+    def forward(self, inputs, node_emb):
+        """
+        Paramters
+        ---------
+        inputs: dict
+            curr_node_id: torch.LongTensor [batch_size]
+            next_node_id: torch.LongTensor [batch_size]
+            mask: torch.LongTensor [batch_size x num_nodes]
+            state: torch.FloatTensor [batch_size x state_dim]
+        node_emb: torch.tensor [batch_size x num_nodes x emb_dim]
+            node embeddings obtained from the node encoder
+
+        Returns
+        -------
+        h: torch.tensor [batch_size x emb_dim]
+            edge embeddings 
+        """
+        curr_node_id = inputs["curr_node_id"]
+        next_node_id = inputs["next_node_id"]
+        mask = inputs["mask"]
+        state = inputs["state"]
+        batch_size = curr_node_id.size(0)
+
+        #--------------------------------
+        # generate queries, keys, values
+        #--------------------------------
+        node_emb = self.dropout(node_emb)
+        curr_emb = node_emb.gather(1, curr_node_id[:, None, None].expand(batch_size, 1, self.emb_dim))
+        next_emb = node_emb.gather(1, next_node_id[:, None, None].expand(batch_size, 1, self.emb_dim))
+        if state is not None and self.state_dim > 0:
+            state_emb = self.init_linear_state(state) # [batch_size x emb_dim]
+            input_q  = torch.cat((curr_emb, next_emb, state_emb[:, None, :]), -1) # [batch_size x 1 x (3*emb_dim)]
+        else:
+            input_q  = torch.cat((curr_emb, next_emb), -1) # [batch_size x 1 x (2*emb_dim)]
+        input_kv = torch.cat((curr_emb.expand_as(node_emb), node_emb), -1) # [batch_size x num_nodes x (2*emb_dim)]
+
+        #--------------------
+        # An attention layer
+        #--------------------
+        q = torch.matmul(input_q,  self.w_q) # [batch_size x 1 x emb_dim]
+        k = torch.matmul(input_kv, self.w_k) # [batch_size x num_nodes x emb_dim]
+        v = torch.matmul(input_kv, self.w_v) # [batch_size x num_nodes x emb_dim]
+        compatibility = self.norm_factor * torch.matmul(q, k.transpose(-2, -1)) # [batch_size x 1 x num_nodes]
+        compatibility[(~mask).unsqueeze(1).expand_as(compatibility)] = -math.inf
+        attn = torch.softmax(compatibility, dim=-1)
+        h = torch.matmul(attn, v) # [batch_size x 1 x emb_dim]
+        h = h.squeeze(1) # [batch_size x emb_dim]
+        return self.out_linear(h) + q.squeeze(1)

models/classifiers/nn_classifiers/encoders/concat_edge_encoder.py

@@ -0,0 +1,63 @@
+import math
+import torch
+import torch.nn as nn
+
+class ConcatEdgeEncoder(nn.Module):
+    def __init__(self, state_dim, emb_dim, dropout):
+        super().__init__()
+        self.state_dim = state_dim
+        self.emb_dim   = emb_dim
+        self.norm_factor = 1 / math.sqrt(emb_dim)
+
+        # initial embedding for state
+        if state_dim > 0:
+            self.init_linear_state = nn.Linear(state_dim, emb_dim)
+
+        # out linear layer
+        self.out_linear = nn.Linear((2 + int(state_dim > 0)) * emb_dim, emb_dim)
+
+        # Dropout
+        self.dropout = nn.Dropout(dropout)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for param in self.parameters():
+            stdv = 1. / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+    
+    def forward(self, inputs, node_emb):
+        """
+        Paramters
+        ---------
+        inputs: dict
+            curr_node_id: torch.LongTensor [batch_size]
+            next_node_id: torch.LongTensor [batch_size]
+            mask: torch.LongTensor [batch_size x num_nodes]
+            state: torch.FloatTensor [batch_size x state_dim]
+        node_emb: torch.tensor [batch_size x num_nodes x emb_dim]
+            node embeddings obtained from the node encoder
+
+        Returns
+        -------
+        h: torch.tensor [batch_size x emb_dim]
+            edge embeddings 
+        """
+        curr_node_id = inputs["curr_node_id"]
+        next_node_id = inputs["next_node_id"]
+        state = inputs["state"]
+        batch_size = curr_node_id.size(0)
+
+        #--------------------------------
+        # generate queries, keys, values
+        #--------------------------------
+        node_emb = self.dropout(node_emb)
+        curr_emb = node_emb.gather(1, curr_node_id[:, None, None].expand(batch_size, 1, self.emb_dim))
+        next_emb = node_emb.gather(1, next_node_id[:, None, None].expand(batch_size, 1, self.emb_dim))
+        if state is not None and self.state_dim > 0:
+            state_emb = self.init_linear_state(state) # [batch_size x emb_dim]
+            edge_emb = torch.cat((curr_emb, next_emb, state_emb[:, None, :]), -1) # [batch_size x 1 x (3*emb_dim)]
+        else:
+            edge_emb = torch.cat((curr_emb, next_emb), -1) # [batch_size x 1 x (2*emb_dim)]
+        edge_emb = edge_emb.squeeze(1) # [batch_size x (2*emb_dim)]
+        return self.out_linear(edge_emb)

models/classifiers/nn_classifiers/encoders/max_readout.py

@@ -0,0 +1,57 @@
+import math
+import torch
+import torch.nn as nn
+
+class MaxReadout(nn.Module):
+    def __init__(self, state_dim, emb_dim, dropout):
+        super().__init__()
+        self.state_dim = state_dim
+        self.emb_dim   = emb_dim
+
+        # initial embedding for state
+        if state_dim > 0:
+            self.init_linear_state = nn.Linear(state_dim, emb_dim)
+
+        # out linear layer
+        self.out_linear = nn.Linear((1 + int(state_dim > 0))*emb_dim, emb_dim)
+
+        # Dropout
+        self.dropout = nn.Dropout(dropout)
+
+        self.reset_parameters()
+    
+    def reset_parameters(self):
+        for param in self.parameters():
+            stdv = 1. / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+
+    def forward(self, inputs, node_emb):
+        """
+        Paramters
+        ---------
+        inputs: dict
+            mask: torch.LongTensor [batch_size x num_nodes]
+            state: torch.FloatTensor [batch_size x state_dim]
+        node_emb: torch.tensor [batch_size x num_nodes x emb_dim]
+            node embeddings obtained from the node encoder
+
+        Returns
+        -------
+        h: torch.tensor [batch_size x emb_dim]
+            graph embeddings
+        """
+        mask = inputs["mask"]
+        state = inputs["state"]
+        node_emb = self.dropout(node_emb)
+
+        # pooling with a mask
+        mask = mask.unsqueeze(-1).expand_as(node_emb)
+        node_emb = node_emb * mask
+        h, _ = torch.max(node_emb, dim=1) # [batch_size x emb_dim]
+        
+        # out linear layer
+        if state is not None and self.state_dim > 0:
+            state_emb = self.init_linear_state(state) # [batch_size x emb_dim]
+            h  = torch.cat((h, state_emb), -1) # [batch_size x (2*emb_dim)]
+
+        return self.out_linear(h)

models/classifiers/nn_classifiers/encoders/mean_readout.py

@@ -0,0 +1,57 @@
+import math
+import torch
+import torch.nn as nn
+
+class MeanReadout(nn.Module):
+    def __init__(self, state_dim, emb_dim, dropout):
+        super().__init__()
+        self.state_dim = state_dim
+        self.emb_dim   = emb_dim
+
+        # initial embedding for state
+        if state_dim > 0:
+            self.init_linear_state = nn.Linear(state_dim, emb_dim)
+
+        # out linear layer
+        self.out_linear = nn.Linear((1 + int(state_dim > 0))*emb_dim, emb_dim)
+
+        # Dropout
+        self.dropout = nn.Dropout(dropout)
+
+        self.reset_parameters()
+    
+    def reset_parameters(self):
+        for param in self.parameters():
+            stdv = 1. / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+
+    def forward(self, inputs, node_emb):
+        """
+        Paramters
+        ---------
+        inputs: dict
+            mask: torch.LongTensor [batch_size x num_nodes]
+            state: torch.FloatTensor [batch_size x state_dim]
+        node_emb: torch.tensor [batch_size x num_nodes x emb_dim]
+            node embeddings obtained from the node encoder
+
+        Returns
+        -------
+        h: torch.tensor [batch_size x emb_dim]
+            graph embeddings
+        """
+        mask = inputs["mask"]
+        state = inputs["state"]
+        node_emb = self.dropout(node_emb)
+
+        # pooling with a mask
+        mask = mask.unsqueeze(-1).expand_as(node_emb)
+        node_emb = node_emb * mask
+        h = torch.mean(node_emb, dim=1) # [batch_size x emb_dim]
+        
+        # out linear layer
+        if state is not None and self.state_dim > 0:
+            state_emb = self.init_linear_state(state) # [batch_size x emb_dim]
+            h  = torch.cat((h, state_emb), -1) # [batch_size x (2*emb_dim)]
+
+        return self.out_linear(h)

models/classifiers/nn_classifiers/encoders/mha_node_encoder.py

@@ -0,0 +1,63 @@
+import math
+import torch
+import torch.nn as nn
+
+class SelfMHANodeEncoder(nn.Module):
+    def __init__(self, coord_dim, node_dim, emb_dim, num_heads, num_mha_layers, dropout):
+        super().__init__()
+        self.coord_dim = coord_dim
+        self.node_dim  = node_dim 
+        self.emb_dim   = emb_dim
+        self.num_mha_layers = num_mha_layers
+
+        # initial embedding
+        self.init_linear_nodes = nn.Linear(node_dim, emb_dim)
+        self.init_linear_depot = nn.Linear(coord_dim, emb_dim)
+
+        # MHA Encoder (w/o positional encoding)
+        mha_layer = nn.TransformerEncoderLayer(d_model=emb_dim, 
+                                               nhead=num_heads,
+                                               dim_feedforward=emb_dim,
+                                               dropout=dropout,
+                                               batch_first=True)
+        self.mha = nn.TransformerEncoder(mha_layer, num_layers=num_mha_layers)
+
+        # Initializing weights
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for param in self.parameters():
+            stdv = 1. / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+
+    def forward(self, inputs):
+        """
+        Paramters
+        ---------
+        inputs: dict
+            node_feat: torch.FloatTensor [batch_size x num_nodes x node_dim]
+
+        Returns
+        -------
+        node_emb: torch.tensor [batch_size x num_nodes x emb_dim]
+            node embeddings 
+        """
+        #----------------
+        # input features
+        #----------------
+        node_feat = inputs["node_feats"]
+
+        #------------------------------------------------------------------------
+        # initial linear projection for adjusting dimensions of locs & the depot
+        #------------------------------------------------------------------------
+        # node_feat = self.dropout(node_feat)
+        loc_emb  = self.init_linear_nodes(node_feat[:, 1:, :]) # [batch_size x num_loc x emb_dim]
+        depot_emb = self.init_linear_depot(node_feat[:, 0:1, :2]) # [batch_size x 1 x emb_dim]
+        node_emb = torch.cat((depot_emb, loc_emb), 1) # [batch_size x num_nodes x emb_dim]
+
+        #--------------
+        # MLP encoding
+        #--------------
+        node_emb = self.mha(node_emb) # [batch_size x num_nodes x emb_dim]
+
+        return node_emb

models/classifiers/nn_classifiers/encoders/mlp_node_encoder.py

@@ -0,0 +1,63 @@
+import math
+import torch
+import torch.nn as nn
+
+class MLPNodeEncoder(nn.Module):
+    def __init__(self, coord_dim, node_dim, emb_dim, num_mlp_layers, dropout):
+        super().__init__()
+        self.coord_dim = coord_dim
+        self.node_dim  = node_dim 
+        self.emb_dim   = emb_dim
+        self.num_mlp_layers = num_mlp_layers
+
+        # initial embedding
+        self.init_linear_nodes = nn.Linear(node_dim, emb_dim)
+        self.init_linear_depot = nn.Linear(coord_dim, emb_dim)
+
+        # MLP Encoder
+        self.mlp = nn.ModuleList()
+        for _ in range(num_mlp_layers):
+            self.mlp.append(nn.Linear(emb_dim, emb_dim, bias=True))
+
+        # Dropout
+        self.dropout = nn.Dropout(dropout)
+
+    def reset_parameters(self):
+        for param in self.parameters():
+            stdv = 1. / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+
+    def forward(self, inputs):
+        """
+        Paramters
+        ---------
+        inputs: dict
+            node_feat: torch.FloatTensor [batch_size x num_nodes x node_dim]
+
+        Returns
+        -------
+        node_emb: torch.tensor [batch_size x num_nodes x emb_dim]
+            node embeddings 
+        """
+        #----------------
+        # input features
+        #----------------
+        node_feat = inputs["node_feats"]
+
+        #------------------------------------------------------------------------
+        # initial linear projection for adjusting dimensions of locs & the depot
+        #------------------------------------------------------------------------
+        # node_feat = self.dropout(node_feat)
+        loc_emb  = self.init_linear_nodes(node_feat[:, 1:, :]) # [batch_size x num_loc x emb_dim]
+        depot_emb = self.init_linear_depot(node_feat[:, 0:1, :2]) # [batch_size x 1 x emb_dim]
+        node_emb = torch.cat((depot_emb, loc_emb), 1) # [batch_size x num_nodes x emb_dim]
+
+        #--------------
+        # MLP encoding
+        #--------------
+        for i in range(self.num_mlp_layers):
+            # node_emb = self.dropout(node_emb)
+            node_emb = self.mlp[i](node_emb)
+            if i != self.num_mlp_layers - 1:
+                node_emb = torch.relu(node_emb)
+        return node_emb

models/classifiers/nn_classifiers/nn_classifier.py

@@ -0,0 +1,156 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+# Node encoder
+from models.classifiers.nn_classifiers.encoders.mlp_node_encoder import MLPNodeEncoder
+from models.classifiers.nn_classifiers.encoders.mha_node_encoder import SelfMHANodeEncoder
+# Edge encoder
+from models.classifiers.nn_classifiers.encoders.attn_edge_encoder import AttentionEdgeEncoder
+from models.classifiers.nn_classifiers.encoders.concat_edge_encoder import ConcatEdgeEncoder
+# Decoder
+from models.classifiers.nn_classifiers.decoders.lstm_decoder import LSTMDecoder
+from models.classifiers.nn_classifiers.decoders.mlp_decoder import MLPDecoder
+from models.classifiers.nn_classifiers.decoders.mha_decoder import SelfMHADecoder
+
+# data loader
+from utils.data_utils.tsptw_dataset import load_tsptw_sequentially
+from utils.data_utils.pctsp_dataset import load_pctsp_sequentially
+from utils.data_utils.pctsptw_dataset import load_pctsptw_sequentially
+from utils.data_utils.cvrp_dataset import load_cvrp_sequentially
+
+NODE_ENC_LIST = ["mlp", "mha"]
+EDGE_ENC_LIST = ["concat", "attn"]
+DEC_LIST = ["mlp", "lstm", "mha"]
+
+class NNClassifier(nn.Module):
+    def __init__(self, 
+                 problem: str, 
+                 node_enc_type: str, 
+                 edge_enc_type: str, 
+                 dec_type: str, 
+                 emb_dim: int, 
+                 num_enc_mlp_layers: int, 
+                 num_dec_mlp_layers: int, 
+                 num_classes: int, 
+                 dropout: float,
+                 pos_encoder: str = "sincos"):
+        super().__init__()
+        self.problem = problem
+        self.node_enc_type = node_enc_type
+        self.edge_enc_type = edge_enc_type
+        self.dec_type = dec_type
+        assert node_enc_type in NODE_ENC_LIST, f"Invalid enc_type. select from {NODE_ENC_LIST}"
+        assert dec_type in DEC_LIST, f"Invalid dec_type. select from {DEC_LIST}"
+        self.is_sequential = True if dec_type in ["lstm", "mha"] else False
+        coord_dim = 2 # only support 2d problem
+        if problem == "tsptw":
+            node_dim  = 4 # coords (2) + time window (2)
+            state_dim = 1 # current time (1)
+        elif problem == "pctsp":
+            node_dim  = 4 # coords (2) + prize (1) + penalty (1)
+            state_dim = 2 # current prize (1) + current penalty (1)
+        elif problem == "pctsptw":
+            node_dim  = 6 # coords (2) + prize (1) + penalty (1) + time window (2)
+            state_dim = 3 # current prize (1) + current penalty (1) + current time (1)
+        elif problem == "cvrp":
+            node_dim  = 3 # coords (2) + demand (1)
+            state_dim = 1 # remaining capacity (1)
+        else:
+            NotImplementedError
+
+        #----------------
+        # Graph encoding
+        #----------------
+        # Node encoder
+        if node_enc_type == "mlp":
+            self.node_enc = MLPNodeEncoder(coord_dim, node_dim, emb_dim, num_enc_mlp_layers, dropout)
+        elif node_enc_type == "mha":
+            num_heads = 8
+            num_mha_layers = 2
+            self.node_enc = SelfMHANodeEncoder(coord_dim, node_dim, emb_dim, num_heads, num_mha_layers, dropout)
+        else:
+            raise NotImplementedError
+        
+        # Readout
+        if edge_enc_type == "concat":
+            self.readout = ConcatEdgeEncoder(state_dim, emb_dim, dropout)
+        elif edge_enc_type == "attn":
+            self.readout = AttentionEdgeEncoder(state_dim, emb_dim, dropout)
+        else:
+            raise NotImplementedError
+
+        #------------------------
+        # Classification Decoder
+        #------------------------
+        if dec_type == "mlp":
+            self.decoder = MLPDecoder(emb_dim, num_dec_mlp_layers, num_classes, dropout)
+        elif dec_type == "lstm":
+            self.decoder = LSTMDecoder(emb_dim, num_dec_mlp_layers, num_classes, dropout)
+        elif dec_type == "mha":
+            num_heads = 8
+            num_mha_layers = 2
+            self.decoder = SelfMHADecoder(emb_dim, num_heads, num_mha_layers, num_classes, dropout, pos_encoder)
+        else:
+            raise NotImplementedError
+
+    def forward(self, inputs):
+        """
+        Paramters
+        ---------
+        inputs: dict
+            curr_node_id: torch.LongTensor [batch_size x max_seq_length] if self.sequential else [batch_size]
+            next_node_id: torch.LongTensor [batch_size x max_seq_length] if self.sequential else [batch_size]
+            node_feat: torch.FloatTensor [batch_size x max_seq_length x num_nodes x node_dim] if self.sequential else [batch_size x num_nodes x node_dim]
+            mask: torch.LongTensor [batch_size x max_seq_length x num_nodes] if self.sequential else [batch_size x num_nodes]
+            state: torch.FloatTensor [batch_size x max_seq_length x state_dim] if self.sequential else [batch_size x state_dim]
+
+        Returns
+        -------
+        probs: torch.tensor [batch_size x seq_length x num_classes] if self.sequential else [batch_size x num_classes]
+            probabilities of classes
+        """
+        #-----------------
+        # Encoding graphs 
+        #-----------------
+        if self.is_sequential:
+            shp = inputs["curr_node_id"].size()
+            inputs = {key: value.flatten(0, 1) for key, value in inputs.items()}
+        node_emb = self.node_enc(inputs) # [(batch_size*max_seq_length) x emb_dim] if self.sequential else [batch_size x emb_dim]
+        graph_emb = self.readout(inputs, node_emb)
+        if self.is_sequential:
+            graph_emb = graph_emb.view(*shp, -1) # [batch_size x max_seq_length x emb_dim]
+
+        #----------
+        # Decoding
+        #----------
+        probs = self.decoder(graph_emb)
+
+        return probs
+    
+    def get_inputs(self, routes, first_explained_step, node_feats):
+        node_feats_ = node_feats.copy()
+        node_feats_["tour"] = routes
+        if self.problem == "tsptw":
+            seq_data = load_tsptw_sequentially(node_feats_)
+        elif self.problem == "pctsp":
+            seq_data = load_pctsp_sequentially(node_feats_)
+        elif self.problem == "pctsptw":
+            seq_data = load_pctsptw_sequentially(node_feats_)
+        elif self.problem == "cvrp":
+            seq_data = load_cvrp_sequentially(node_feats_)
+        else:
+            NotImplementedError
+        
+        def pad_seq_length(batch):
+            data = {}
+            for key in batch[0].keys():
+                padding_value = True if key == "mask" else 0.0
+                # post-padding
+                data[key] = torch.nn.utils.rnn.pad_sequence([d[key] for d in batch], batch_first=True, padding_value=padding_value)
+            pad_mask = torch.nn.utils.rnn.pad_sequence([torch.full((d["mask"].size(0), ), True) for d in batch], batch_first=True, padding_value=False)
+            data.update({"pad_mask": pad_mask})
+            return data
+        instance = pad_seq_length(seq_data)
+        return instance

models/classifiers/predictor.py

@@ -0,0 +1,203 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+#------------
+# base class
+#------------
+class DecisionPredictorBase(nn.Module):
+    def __init__(self, coord_dim, node_dim, state_dim, emb_dim, num_mlp_layers, num_classes, dropout):
+        super().__init__()
+        self.coord_dim = coord_dim
+        self.node_dim  = node_dim 
+        self.emb_dim   = emb_dim
+        self.state_dim = state_dim
+        self.num_mlp_layers = num_mlp_layers
+        self.norm_factor = 1 / math.sqrt(emb_dim)
+
+        # initial embedding
+        self.init_linear_node  = nn.Linear(node_dim, emb_dim)
+        self.init_linear_depot = nn.Linear(coord_dim, emb_dim)
+        if state_dim > 0:
+            self.init_linear_state = nn.Linear(state_dim, emb_dim)
+        
+        # An attention layer
+        self.w_q = nn.Parameter(torch.FloatTensor((2 + int(state_dim > 0)) * emb_dim, emb_dim))
+        self.w_k = nn.Parameter(torch.FloatTensor(2 * emb_dim, emb_dim))
+        self.w_v = nn.Parameter(torch.FloatTensor(2 * emb_dim, emb_dim))
+        
+        # MLP
+        self.mlp = nn.ModuleList()
+        for i in range(self.num_mlp_layers):
+            self.mlp.append(nn.Linear(emb_dim, emb_dim, bias=True))
+        self.mlp.append(nn.Linear(emb_dim, num_classes, bias=True))
+
+        # Dropout
+        self.dropout = nn.Dropout(dropout)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for param in self.parameters():
+            stdv = 1. / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+
+    def forward(self, inputs):
+        """
+        Paramters
+        ---------
+        inputs: dict
+            curr_node_id: torch.LongTensor [batch_size]
+            next_node_id: torch.LongTensor [batch_size]
+            node_feat: torch.FloatTensor [batch_size x num_nodes x node_dim]
+            mask: torch.LongTensor [batch_size x num_nodes]
+            state: torch.FloatTensor [batch_size x state_dim]
+
+        Returns
+        -------
+        probs: torch.tensor [batch_size x num_classes]
+        """
+        #----------------
+        # input features
+        #----------------
+        curr_node_id = inputs["curr_node_id"]
+        next_node_id = inputs["next_node_id"]
+        node_feat = inputs["node_feats"]
+        mask = inputs["mask"]
+        state = inputs["state"]
+
+        #---------------------------
+        # initial linear projection
+        #---------------------------
+        node_emb  = self.init_linear_node(node_feat[:, 1:, :]) # [batch_size x num_loc x emb_dim]
+        depot_emb = self.init_linear_depot(node_feat[:, 0:1, :2]) # [batch_size x 1 x emb_dim]
+        new_node_feat = torch.cat((depot_emb, node_emb), 1) # [batch_size x num_nodes x emb_dim]
+        new_node_feat = self.dropout(new_node_feat)
+
+        #---------------
+        # preprocessing
+        #---------------
+        batch_size = curr_node_id.size(0)
+        curr_emb = new_node_feat.gather(1, curr_node_id.unsqueeze(-1).expand(batch_size, 1, self.emb_dim))
+        next_emb = new_node_feat.gather(1, next_node_id.unsqueeze(-1).expand(batch_size, 1, self.emb_dim))
+        if state is not None and self.state_dim > 0:
+            state_emb = self.init_linear_state(state) # [batch_size x emb_dim]
+            input_q  = torch.cat((curr_emb, next_emb, state_emb[:, None, :]), -1) # [batch_size x 1 x (3*emb_dim)]
+        else:
+            input_q  = torch.cat((curr_emb, next_emb), -1) # [batch_size x 1 x (2*emb_dim)]
+        input_kv = torch.cat((curr_emb.expand_as(new_node_feat), new_node_feat), -1) # [batch_size x num_nodes x (2*emb_dim)]
+
+        #--------------------
+        # An attention layer
+        #--------------------
+        q = torch.matmul(input_q,  self.w_q) # [batch_size x 1 x emb_dim]
+        k = torch.matmul(input_kv, self.w_k) # [batch_size x num_nodes x emb_dim]
+        v = torch.matmul(input_kv, self.w_v) # [batch_size x num_nodes x emb_dim]
+        compatibility = self.norm_factor * torch.matmul(q, k.transpose(-2, -1)) # [batch_size x 1 x num_nodes]
+        compatibility[(~mask).unsqueeze(1).expand_as(compatibility)] = -math.inf
+        attn = torch.softmax(compatibility, dim=-1)
+        h = torch.matmul(attn, v) # [batch_size x 1 x emb_dim]
+        h = h.squeeze(1) # [batch_size x emb_dim]
+        
+        #---------------
+        # MLP (decoder)
+        #---------------
+        for i in range(self.num_mlp_layers):
+            h = self.dropout(h)
+            h = torch.relu(self.mlp[i](h))
+        h = self.dropout(h)
+        logits = self.mlp[-1](h)
+        probs = F.log_softmax(logits, dim=-1)
+        return probs
+
+    def get_inputs(self, tour, first_explained_step, node_feats):
+        """
+        For TSPTW
+        TODO: refactoring
+
+        Parameters
+        ----------
+        tour: list [seq_length]
+        first_explained_step: int
+        node_feats np.array [num_nodes x node_dim]
+
+        Returns
+        -------
+        out: dict (key: data type [data_size])
+            curr_node_id: torch.tensor [num_explained_paths]
+            next_node_id: torch.tensor [num_explained_paths]
+            node_feats: torch.tensor [num_explained_paths x num_nodes x node_dim]
+            mask: torch.tensor [num_explained_paths x num_nodes]
+            state: torch.tensor [num_explained_paths x state_dim]
+        """
+        node_feats = {
+            key: torch.from_numpy(node_feat.astype(np.float32).copy()).clone()
+                    if isinstance(node_feat, np.ndarray) else
+                    torch.tensor([node_feat])
+            for key, node_feat in node_feats.items()
+        }
+        if isinstance(tour, np.ndarray):
+            tour = torch.from_numpy(tour.astype(np.long).copy()).clone()
+        else:
+            tour = torch.LongTensor(tour)
+        
+        out = {"curr_node_id": [], "next_node_id": [], "mask": [], "state": []}
+        for step in range(first_explained_step, len(tour) - 1):
+            # node ids
+            curr_node_id = tour[step]
+            next_node_id = tour[step + 1]
+            # mask & state
+            max_coord = node_feats["grid_size"]
+            coord = node_feats["coords"] / max_coord # [num_nodes x coord_dim]
+            time_window = node_feats["time_window"] # [num_nodes x 2(start, end)]
+            time_window = (time_window - time_window[1:].min()) / time_window[1:].max() # min-max normalization
+            curr_time = torch.FloatTensor([0.0])
+            raw_coord = node_feats["coords"]
+            raw_time_window = node_feats["time_window"]
+            raw_curr_time = torch.FloatTensor([0.0])
+            num_nodes = len(node_feats["coords"])
+            mask = torch.ones(num_nodes, dtype=torch.long) # feasible -> 1, infeasible -> 0
+            for i in range(step + 1):
+                curr_id = tour[i]
+                if i > 0:
+                    prev_id = tour[i - 1]
+                    raw_curr_time += torch.norm(raw_coord[curr_id] - raw_coord[prev_id])
+                    curr_time += torch.norm(coord[curr_id] - coord[prev_id])
+                    # visited?
+                    mask[curr_id] = 0
+                    # curr_time exceeds the time window?
+                    mask[curr_time > time_window[:, 1]] = 0
+            curr_time = (raw_curr_time - raw_time_window[1:].min()) / raw_time_window[1:].max() # min-max normalization
+            out["curr_node_id"].append(curr_node_id)
+            out["next_node_id"].append(next_node_id)
+            out["mask"].append(mask)
+            out["state"].append(curr_time)
+        out = {key: torch.stack(value, 0) for key, value in out.items()}
+        node_feats = {
+            key: node_feat.unsqueeze(0).expand(out["mask"].size(0), *node_feat.size())
+            for key, node_feat in node_feats.items()
+        }
+        out.update({"node_feats": node_feats})
+        return out
+
+#---------------
+# general class
+#---------------
+class DecisionPredictor(DecisionPredictorBase):
+    def __init__(self, problem, emb_dim, num_mlp_layers, num_classes, drop):
+        coord_dim = 2
+        self.problem = problem
+        if problem == "tsptw":
+            node_dim = coord_dim + 2 # + time_window(start, end)
+            state_dim = 1 # current_time
+        elif problem == "cvrp":
+            node_dim = coord_dim + 1 # + demand
+            state_dim = 1 # used_capacity
+        elif problem == "cvrptw":
+            node_dim = coord_dim + 1 + 2 # + demand + time_window(start, end)
+            state_dim = 2 # used_capacity + current_time
+        else:
+            assert False, f"problem {problem} is not supported!"
+        super().__init__(coord_dim, node_dim, state_dim, emb_dim, num_mlp_layers, num_classes, drop)

models/classifiers/rule_based_models.py

@@ -0,0 +1,150 @@
+import torch
+import torch.nn as nn
+
+TOUR_LENGTH = 0
+TIME_WINDOW = 1
+
+class kNearestPredictor(nn.Module):
+    def __init__(self, problem, k, k_type):
+        """
+        Paramters
+        ---------
+        problem: str
+            problem type
+        k: float
+            if the vehicle visis k% nearest node, this model labels the visit as prioritizing tour length
+        """
+        super().__init__()
+        self.problem = problem
+        self.num_classes = 2
+        self.k_type = k_type
+        if k_type == "num":
+            self.k = int(k)
+        elif k_type == "ratio":
+            self.k = k
+        else:
+            assert False, "Invalid k_type. select from [num, ratio]"
+
+    def forward(self, inputs):
+        """
+        Parameters
+        ----------
+
+        Returns
+        -------
+        probs: torch.tensor [batch_size x num_classes]
+        """
+        #----------------
+        # input features
+        #----------------
+        curr_node_id = inputs["curr_node_id"]
+        next_node_id = inputs["next_node_id"]
+        node_feat = inputs["node_feats"]
+        mask = inputs["mask"]
+
+        coord_dim = 2
+        batch_size = curr_node_id.size(0)
+        coords = node_feat[:, :, :coord_dim] # [batch_size x num_nodes x coord_dim]
+        num_candidates = (mask > 0).sum(dim=-1) # [batch_size]
+        topk = torch.round(num_candidates * self.k).to(torch.long) # [batch_size]
+        curr_coord = coords.gather(1, curr_node_id[:, None, None].expand_as(coords)) # [batch_size x 1 x coord_dim]
+        dist_from_curr_node = torch.norm(curr_coord - coords, dim=-1) # [batch_size x 1 x num_nodes]
+        visit_topk = []
+        for i in range(batch_size):
+            if self.k_type == "num":
+                k = self.k
+            else:    
+                k = topk[i].item()
+            id = torch.topk(input=dist_from_curr_node[i], k=k, dim=-1, largest=True)[1]
+            visit_topk.append(torch.isin(next_node_id[i], id))
+        visit_topk = torch.stack(visit_topk, 0)
+        idx = (1 - visit_topk.int()).to(torch.long)
+        probs = torch.zeros(batch_size, self.num_classes).to(torch.float)
+        probs.scatter_(-1, idx.unsqueeze(-1).expand_as(probs), 1.0)
+        return probs
+    
+    def get_inputs(self, tour, first_explained_step, node_feats):
+        """
+        For TSPTW
+        TODO: refactoring
+
+        Parameters
+        ----------
+        tour: list [seq_length]
+        first_explained_step: int
+        node_feats np.array [num_nodes x node_dim]
+
+        Returns
+        -------
+        out: dict (key: data type [data_size])
+            curr_node_id: torch.tensor [num_explained_paths]
+            next_node_id: torch.tensor [num_explained_paths]
+            node_feats: torch.tensor [num_explained_paths x num_nodes x node_dim]
+            mask: torch.tensor [num_explained_paths x num_nodes]
+            state: torch.tensor [num_explained_paths x state_dim]
+        """
+        if isinstance(node_feats, np.ndarray):
+            node_feats = torch.from_numpy(node_feats.astype(np.float32)).clone()
+        tour = torch.LongTensor(tour)
+        coord_dim = 2
+        out = {"curr_node_id": [], "next_node_id": [], "mask": [], "state": []}
+        for step in range(first_explained_step, len(tour) - 1):
+            # node ids
+            curr_node_id = tour[step]
+            next_node_id = tour[step + 1]
+            # mask & state
+            max_coord = 100
+            coord = node_feats[:, coord_dim] / max_coord # [num_nodes x coord_dim]
+            time_window = node_feats[:, coord_dim:] # [num_nodes x 2(start, end)]
+            time_window = (time_window - time_window[1:].min()) / time_window[1:].max() # min-max normalization
+            curr_time = torch.FloatTensor([0.0])
+            raw_coord = node_feats[:, coord_dim]
+            raw_time_window = node_feats[:, coord_dim:]
+            raw_curr_time = torch.FloatTensor([0.0])
+            mask = torch.ones(node_feats.size(0), dtype=torch.long) # feasible -> 1, infeasible -> 0
+            for i in range(step + 1):
+                curr_id = tour[i]
+                if i > 0:
+                    prev_id = tour[i - 1]
+                    raw_curr_time += torch.norm(raw_coord[curr_id] - raw_coord[prev_id])
+                    curr_time += torch.norm(coord[curr_id] - coord[prev_id])
+                    # visited?
+                    mask[curr_id] = 0
+                    # curr_time exceeds the time window?
+                    mask[curr_time > time_window[:, 1]] = 0
+            curr_time = (raw_curr_time - raw_time_window[1:].min()) / raw_time_window[1:].max() # min-max normalization
+            out["curr_node_id"].append(curr_node_id)
+            out["next_node_id"].append(next_node_id)
+            out["mask"].append(mask)
+            out["state"].append(curr_time)
+        out = {key: torch.stack(value, 0) for key, value in out.items()}
+        node_feats = node_feats.unsqueeze(0).expand(out["mask"].size(0), node_feats.size(-2), node_feats.size(-1))
+        out.update({"node_feats": node_feats})
+        return out
+
+    def get_topk_ids(self, input, k, dim, largest):
+        """
+        Parameters
+        ----------
+        input: torch.tensor [batch_size x num_nodes x num_nodes]
+        k: torch.tensor [batch_size]
+        dim: int
+        largest: bool
+
+        Returns
+        -------
+        topk_ids: torch.tensor [batch_size x num_node x k]
+        """
+        batch_size = input.size(0)
+        max_k = k.max()
+        ids = []
+        for i in range(batch_size):
+            id = torch.topk(input=input[i], k=k[i].item(), dim=dim, largest=largest)[1]
+
+            # adjust tensor size
+            if id.size(0) == 0:
+                id = torch.full((max_k, ), -1000)
+            elif id.size(0) < max_k:
+                id = torch.cat((id, torch.full((max_k - id.size(0), ), id[0])), -1)
+            ids.append(id)
+        return torch.stack(ids, 0)

models/loss_functions.py

@@ -0,0 +1,129 @@
+import torch
+import torch.nn as nn
+from utils.utils import batched_bincount
+import torch.nn.functional as F
+
+
+class GeneralCrossEntropy(nn.Module):
+    def __init__(self, weight_type: str, beta : float = 0.99, is_sequential: bool = True):
+        super().__init__()
+        self.weight_type = weight_type
+        self.beta = beta
+        if weight_type == "seq_cbce":
+            assert is_sequential == True
+            self.loss_func = SeqCBCrossEntropy(beta=beta)
+        elif weight_type == "cbce":
+            self.loss_func = CBCrossEntropy(beta=beta, is_sequential=is_sequential)
+        elif weight_type == "wce":
+            self.loss_func = WeightedCrossEntropy(is_sequential=is_sequential)
+        elif weight_type == "ce":
+            self.loss_func = CrossEntropy(is_sequential=is_sequential)
+        else:
+            NotImplementedError
+
+    def forward(self, 
+                preds: torch.Tensor, 
+                labels: torch.Tensor,
+                pad_mask: torch.Tensor = None):
+        return self.loss_func(preds, labels, pad_mask)
+
+
+class SeqCBCrossEntropy(nn.Module):
+    def __init__(self, beta : float = 0.99):
+        super().__init__()
+        self.beta = beta
+
+    def forward(self,
+                preds: torch.Tensor, 
+                labels: torch.Tensor,
+                pad_mask: torch.Tensor):
+        """
+        Sequential Class-alanced Cross Entropy Loss (Our proposal)
+
+        Parameters
+        -----------
+        preds: torch.Tensor [batch_size, max_seq_length, num_classes]
+        labels: torch.Tensor [batch_size, max_seq_length]
+        pad_mask: torch.Tensor [batch_size, max_seq_length]
+
+        Returns
+        -------
+        loss: torch.Tensor [1]
+        """
+        seq_length_batch = pad_mask.sum(-1) # [batch_size]
+        seq_length_list = torch.unique(seq_length_batch) # [num_unique_seq_length]
+        batch_size = preds.size(0)
+        loss = 0
+        for seq_length in seq_length_list:
+            extracted_batch = (seq_length_batch == seq_length) # [batch_size]
+            extracted_preds = preds[extracted_batch]   # [num_extracted_batch]
+            extracted_labels = labels[extracted_batch] # [num_extracted_batch]
+            extracted_batch_size = extracted_labels.size(0)
+            bin = batched_bincount(extracted_labels.T, 1, extracted_preds.size(-1)) # [seq_length x num_classes]
+            weight = (1 - self.beta) / (1 - self.beta**bin + 1e-8)
+            for seq_no in range(seq_length.item()):
+                loss += (extracted_batch_size / batch_size) * F.nll_loss(extracted_preds[:, seq_no], extracted_labels[:, seq_no], weight=weight[seq_no])
+        return loss
+
+class CBCrossEntropy(nn.Module):
+    def __init__(self, beta : float = 0.99, is_sequential: bool = True):
+        super().__init__()
+        self.beta = beta
+        self.is_sequential = is_sequential
+
+    def forward(self,
+                preds: torch.Tensor, 
+                labels: torch.Tensor,
+                pad_mask: torch.Tensor = None):
+        if self.is_sequential:
+            mask = pad_mask.view(-1)
+            preds = preds.view(-1, preds.size(-1))
+            bin = labels.view(-1)[mask].bincount()
+            weight = (1 - self.beta) / (1 - self.beta**bin + 1e-8)
+            loss = F.nll_loss(preds[mask], labels.view(-1)[mask], weight=weight)
+        else:
+            bincount = labels.view(-1).bincount()
+            weight = (1 - self.beta) / (1 - self.beta**bincount + 1e-8)
+            loss = F.nll_loss(preds, labels.squeeze(-1), weight=weight)
+        return loss
+
+class WeightedCrossEntropy(nn.Module):
+    def __init__(self, is_sequential: bool = True, norm: str = "min"):
+        super().__init__()
+        self.is_sequential = is_sequential
+        if norm == "min":
+            self.norm = torch.min
+        elif norm == "max":
+            self.norm = torch.max
+    def forward(self,
+                preds: torch.Tensor, 
+                labels: torch.Tensor,
+                pad_mask: torch.Tensor = None):
+        if self.is_sequential:
+            mask = pad_mask.view(-1)
+            preds = preds.view(-1, preds.size(-1))
+            bin = labels.view(-1)[mask].bincount()
+            weight = self.norm(bin) / (bin + 1e-8)
+            loss = F.nll_loss(preds[mask], labels.view(-1)[mask], weight=weight)
+        else:
+            bincount = labels.view(-1).bincount()
+            weight = self.norm(bin) / (bin + 1e-8)
+            loss = F.nll_loss(preds, labels.squeeze(-1), weight=weight)
+        return loss
+
+class CrossEntropy(nn.Module):
+    def __init__(self, is_sequential: bool = True):
+        super().__init__()
+        self.is_sequential = is_sequential
+
+    def forward(self,
+                preds: torch.Tensor, 
+                labels: torch.Tensor,
+                pad_mask: torch.Tensor = None):
+        if self.is_sequential:
+            mask = pad_mask.view(-1)
+            preds = preds.view(-1, preds.size(-1))
+            loss = F.nll_loss(preds[mask], labels.view(-1)[mask])
+        else:
+            loss = F.nll_loss(preds, labels.squeeze(-1))
+        return loss

models/prompts/generate_explanation.py

@@ -0,0 +1,110 @@
+from langchain.prompts import PromptTemplate
+from langchain.schema import StrOutputParser
+from langchain_core.runnables.base import Runnable
+from models.prompts.template_json_base import TemplateJsonBase
+
+GENERATE_EXPLANATION = """\
+You are RouteExplainer, an explanation system for justifying a specific edge (i.e., actual edge) in a route automatically generated by a VRP solver. 
+Here, you address a scenario where a tourist (user) wonders why the actual edge was selected at the step in the tourist route and why another edge was not selected instead at that step. 
+As an expert tour guide, you will justify why the actual edge was selected in the route if it outperforms another edge. That helps to convince the tourist of the actual edge or to make the tourist's decision to change to another edge from the actual edge while accepting some disadvantages. 
+Please carefully read the contents below and follow the instructions faithfully. 
+
+[Terminology]
+The following terms are used here.
+- Node: A destination.
+- Edge: A directed edge representing the movement from one node to another.
+- Edge intention: The underlying purpose of the edge. An edge intention here is either “prioritizing route length (route_len)” or “prioritizing time windows (time_window)”.
+- Step: The visited order in a route.
+- Actual edge: A user-specified edge in the optimal route generated by a VRP solver. You will justify this edge in this task.
+- Counterfactual (CF) edge: A user-specified edge that was not selected at the step of the actual edge in the optimal route but could have been. This is a different edge from the actual edge. The user wonders why the CF edge was not selected at the step instead of the actual edge. 
+- Actual route: The optimal route generated by a VRP solver.
+- CF route: An alternative route where the CF edge is selected at the step instead of the actual edge. The subsequent edges to the CF edge in the CF route are the best-effort ones.
+
+[Example]
+Please refer to the following input-output example when generating a counterfactual explanation.
+***** START EXAMPLE *****
+[input]
+Question:
+- The question asks about replacing the edge from node2 to node3 with the edge from node2 to node5.
+Actual route:
+- route: node1 > node2 > (actual edge) > node3 >  node4 > node5 > node6 >  node7 >  node1
+- short-term effect (immediate travel time): 20 minutes
+- long-term effect (total travel time): 100 minutes
+- missed nodes: none
+- edge-intention ratio after the actual edge: time_window 75%, route_len 25%
+CF route:
+- route: node1 > node2 > (CF edge) > node5 > node6 > node7 > node1
+- short-term effect (immediate travel time): 10 minutes
+- long-term effect (total travel time): 77.8 minutes
+- missed nodes: node3, node4
+- edge-intention ratio after the CF edge: time_window 100%, route_len 0% 
+Difference between two routes:
+- short-term effect: The actual route increases it by 10 minutes
+- long-term effect: The actual route increases it by 22.2 minutes
+- missed nodes: The actual route visits 2 more nodes
+- difference of edge-intention ratio after the actual and CF edges: time_window -25%, route_len +25%
+Planed destination information:
+- node1: start/end point
+- node2: none
+- node3: take lunch
+- node4: attend a tour
+- node5: most favorite destination
+- node6: take dinner
+- node7: none
+
+[Explanation]
+Here are the pros and cons of the actual and CF edges.
+#### Actual edge:
+  - Pros:
+    - It allows you to visit all your destinations within time windows. That is essential for maximizing your tour experience.
+  - Cons: 
+    - Immediate travel time will increase by 10 minutes.
+    - The total travel time will increase by 22.2 minutes, but it is natural because the actual route visits two more nodes than the CF route. 
+  - Remarks: 
+    - The route balances both prioritizing travel time and time windows.
+#### CF edge:
+  - Pros:
+    - Immediate travel time will decrease by 10 minutes.
+    - The total travel time will decrease by 22.2 minutes. However, note that this reduction in time is the result of not visiting two nodes. 
+  - Cons:
+    - You will miss node3 and node4. You plan to take lunch and attend a tour, so the loss could significantly degrade your tour experience.
+  - Remarks:
+    - You will miss node3 and node4 even if you are constantly pressed for time windows in the subsequent movement
+#### Summary:
+  - Given the pros and cons and the fact that adhering to time constraints is essential, the actual edge is objectively more optimal. 
+  - However, you might prefer the CF edge, despite its cons, depending on your preferences.
+***** END EXAMPLE *****
+
+[Instruction]
+Now, please generate a counterfactual explanation for the [input] below.
+You MUST keep the following rules:
+  - Summarize the pros and cons, including short-term effects, long-term effects, missed nodes, and edge-intention ratio.
+  - Enrich explanations by leveraging destination information.
+  - Carefully consider causality regarding travel time reduction. If the number of missed nodes is equal, one edge may reduce travel time. However, if a route with missed nodes is quicker, it is due to skipping nodes.
+  - A high route_len ratio emphasizes speed over schedule adherence, while a high time_window ratio prioritizes sticking to a schedule, sacrificing travel efficiency for timely arrivals.
+  - Disucuss edge-intention ratio in "Remarks". Do NOT do it in "Pros" or "Cons".
+  - Travel time efficiency is solely determined by the total travel time.
+  - Never say that all planed destinations are visited if there is even one missed node. If some nodes are missed, you must specify which node are missed.
+  - If the CF edge outperforms the actual edge, you do NOT have to force a justification for the actual edge.
+  - Please associate user's intention that "{intent}" with your summary. If the intention is blank, it means no intention was provided.
+
+[input]
+{comparison_results}
+
+[Explanation]
+"""
+
+# - Routes are assessed based on the following priorities: fewer missed nodes are better > shorter total travel time is better.
+# - In "Summary", Clearly and specifically explain the differences between the actual and CF edges to help the tourist convince the actual edge or make a decision to change to another edge from the actual edge while accepting some cons. 
+# - Ensure consistency in comparisons: the pros of the actual edge should be the cons of the CF edge and vice versa.
+
+class Template4GenerateExplanation(TemplateJsonBase):
+    parser: Runnable = StrOutputParser()
+    template: str = GENERATE_EXPLANATION
+    prompt: Runnable = PromptTemplate(
+        template=template,
+        input_variables=["comparison_results", "intent"],
+    )
+
+    def _get_output_key(self) -> str:
+        return ""

models/prompts/identify_question.py

@@ -0,0 +1,76 @@
+from langchain.prompts import PromptTemplate
+from langchain_core.output_parsers import JsonOutputParser
+from langchain_core.pydantic_v1 import BaseModel, Field
+from langchain_core.runnables.base import Runnable
+from models.prompts.template_json_base import TemplateJsonBase
+
+IDENTIFY_QUESTION = """\
+Given a route and a question that asks what would happen if we replaced a specific edge in the route with another edge, please extract the step number of the replaced edge (i.e., cf_step) and the node id of the destination of the new edge (i.e., cf_visit) from the question, which is written in natural language.
+Please use the following examples as a reference when you answer:
+***** START EXAMPLE *****
+[route info]
+Nodes(node id, name): (1, node1), (2, node2), (3, node3), (4, node4), (5, node5)
+Route: node1 > (step1) > node5 > (step2) > node3 > (step3) > node2 > (step4) > node4 > (step 5) > node1
+[question] 
+Why node3, and why not node2?
+[outputs]
+```json
+{{
+    "success": true,
+    "summary": "The answer asks about replacing the edge from node5 to node3 with the edge from node6 to node2.",
+    "intent": "",
+    “process": "The edge from node5 to node3 is at step2 because of "node5 > (step2) > node3". The node id of the destination of the new edge is 2 (node2). Thus, the final answers are cf_step=2 and cf_visit=2.",
+    "cf_step": 2,
+    "cf_visit": 2,
+}}
+```
+
+[route info]
+Nodes(node id, name): (1, node1), (2, node2), (3, node3), (4, node4), (5, node5)
+Route: node1 > (step1) > node5 > (step2) > node3 > (step3) > node2 > (step4) > node4 > (step 5) > node1
+[quetsion]
+What if we visited node4 instead of node2? We would personally like to visit node4 first.
+[outputs]
+```json
+{{
+    "success": true,
+    "summary": "The answer asks about replacing the edge from node3 to node2 with the edge from node3 to node4.",
+    “intent": "The user would personally like to visit node4 first""
+    "process": "The edge from node3 to node2 is at step3 because of "node3 > (step3) > node2". The node id of the destination of the new edge is 4 (node4). Thus, the final answers are cf_step=3 and cf_visit=4.",
+    "cf_step": 3,
+    "cf_visit": 4,
+}}
+```
+***** END EXAMPLE *****
+
+Given the following route and question, please extract the step number of the replaced edge (i.e., cf_step) and the node id of the destination of the new edge (i.e., cf_visit) from the question.
+Please keep the following rules:
+- Do not output any sentences outside of JSON format.
+- {format_instructions}
+
+[route_info]
+{route_info}
+[question] 
+{whynot_question}
+[outputs]
+"""
+
+class WhyNotQuestion(BaseModel):
+    success: bool = Field(description="Whether cf_step and cf_visit are successfully extracted (True) or not (False).")
+    summary:  str = Field(description="Your summary for the given question. If success=False, instead state here what information is missing to extract cf_step/cf_visit and what additional information should be clarified (Additionally, provide an example).")
+    intent:   str = Field(description="Your summary for user's intent (if provided). If not provided, this is set to ''.")
+    process:  str = Field(description="The thought (reasoning) process in extracting cf_step and cf_visit. if success=False, this is set to ''.")
+    cf_step:  int = Field(description="The step number of the replaced edge. if success=False, this is set to -1.")
+    cf_visit: int = Field(description="The node id of the destination of the new edge. if success=False, this is set to -1.")
+
+class Template4IdentifyQuestion(TemplateJsonBase):
+    parser: Runnable = JsonOutputParser(pydantic_object=WhyNotQuestion)
+    template: str = IDENTIFY_QUESTION
+    prompt: Runnable = PromptTemplate(
+        template=template,
+        input_variables=["whynot_question", "route_info"],
+        partial_variables={"format_instructions": parser.get_format_instructions()}
+    )
+
+    def _get_output_key(self) -> str:
+        return ""

models/prompts/template_json_base.py

@@ -0,0 +1,27 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Any 
+from langchain_core.runnables import RunnableLambda
+from langchain_core.runnables.base import Runnable
+
+class TemplateJsonBase(ABC):
+    parser: Runnable
+    template: str
+    prompt: Runnable
+
+    @abstractmethod
+    def _get_output_key(self) -> str:
+        raise NotImplementedError
+
+    def get_template(self) -> str:
+        return self.template
+    
+    def extract_value(self, input: Dict[str, Any]) -> Any:
+        return input[self._get_output_key()]
+
+    def sandwiches(self, 
+                   llm: Runnable,
+                   extract_value: bool = False) -> Runnable:
+        if extract_value:
+            return self.prompt | llm | self.parser | RunnableLambda(self.extract_value)
+        else:
+            return self.prompt | llm | self.parser

models/route_explainer.py

@@ -0,0 +1,293 @@
+# templates
+import numpy as np
+import streamlit as st
+from typing import Dict, List
+from models.prompts.identify_question import Template4IdentifyQuestion
+from models.prompts.generate_explanation import Template4GenerateExplanation
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain.schema import AIMessage
+import utils.util_app as util_app
+
+class StreamingChatCallbackHandler(BaseCallbackHandler):
+    def __init__(self):
+        pass
+
+    def on_llm_start(self, *args, **kwargs):
+        self.container = st.empty()
+        self.text = ""
+
+    def on_llm_new_token(self, token: str, *args, **kwargs):
+        self.text += token
+        self.container.markdown(
+            body=self.text,
+            unsafe_allow_html=False,
+        )
+
+    def on_llm_end(self, response: str, *args, **kwargs):
+        self.container.markdown(
+            body=response.generations[0][0].text,
+            unsafe_allow_html=False,
+        )
+
+class RouteExplainer():
+    template_identify_question = Template4IdentifyQuestion()
+    template_generate_explanation = Template4GenerateExplanation()
+
+    def __init__(self,
+                 llm,
+                 cf_generator, 
+                 classifier) -> None:
+        assert cf_generator.problem == classifier.problem, "Problem type of cf_generator and predictor should coincide!"
+        self.coord_dim = 2
+        self.problem = cf_generator.problem
+        self.cf_generator = cf_generator
+        self.classifier = classifier
+        self.actual_route = None
+        self.cf_route = None
+        # templates
+        self.question_extractor = self.template_identify_question.sandwiches(llm)
+        self.explanation_generator = self.template_generate_explanation.sandwiches(llm)
+
+    #----------------
+    # whole pipeline
+    #----------------
+    def generate_explanation(self, 
+                             tour_list,
+                             whynot_question: str,
+                             actual_routes: list,
+                             actual_labels: list,
+                             node_feats: dict,
+                             dist_matrix: np.array) -> str:
+        #--------------------------------
+        # define why & why-not questions
+        #--------------------------------
+        route_info_text = self.get_route_info_text(tour_list, actual_routes)
+        inputs = self.question_extractor.invoke({
+            "whynot_question": whynot_question,
+            "route_info": route_info_text
+        })
+        util_app.stream_words(inputs["summary"] + " " + inputs["intent"])
+        st.session_state.chat_history.append(AIMessage(content=inputs["summary"] + inputs["intent"]))
+        if not inputs["success"]:
+            return ""
+
+        #----------------------
+        # validate the CF edge
+        #----------------------
+        is_cf_edge_feasible, reason = self.validate_cf_edge(node_feats,
+                                                            dist_matrix,
+                                                            actual_routes[0],
+                                                            inputs["cf_step"],
+                                                            inputs["cf_visit"]-1)
+        # exception
+        if not is_cf_edge_feasible:
+            util_app.stream_words(reason)
+            return reason
+
+        #---------------------
+        # generate a cf route
+        #---------------------
+        cf_routes = self.cf_generator(actual_routes,
+                                      vehicle_id=0,
+                                      cf_step=inputs["cf_step"],
+                                      cf_next_node_id=inputs["cf_visit"]-1,
+                                      node_feats=node_feats,
+                                      dist_matrix=dist_matrix)
+        st.session_state.generated_cf_route = True
+        st.session_state.close_chat = True
+        st.session_state.cf_step = inputs["cf_step"]
+
+        #--------------------------------------
+        # classify the intentions of each edge
+        #--------------------------------------
+        cf_labels = self.classifier(self.classifier.get_inputs(cf_routes,
+                                                               0,
+                                                               node_feats,
+                                                               dist_matrix))
+        st.session_state.cf_routes = cf_routes
+        st.session_state.cf_labels = cf_labels
+
+        #-------------------------------------
+        # generate a constrastive explanation
+        #-------------------------------------
+        comparison_results = self.get_comparison_results(question_summary=inputs["summary"],
+                                                         tour_list=tour_list,
+                                                         actual_routes=actual_routes,
+                                                         actual_labels=actual_labels,
+                                                         cf_routes=cf_routes,
+                                                         cf_labels=cf_labels,
+                                                         cf_step=inputs["cf_step"])
+        
+        explanation = self.explanation_generator.invoke({
+            "comparison_results": comparison_results,
+            "intent": inputs["intent"]
+        }, config={"callbacks": [StreamingChatCallbackHandler()]})
+
+        return explanation
+    
+    #-------------------------
+    # for exctracting inputs
+    #-------------------------
+    def get_route_info_text(self, tour_list, routes) -> str:
+        route_info = ""
+        # nodes
+        route_info += "Nodes(node id, name): "
+        for i, destination in enumerate(tour_list):
+            if i != len(tour_list) - 1:
+                route_info += f"({i+1}, {destination['name']}), "
+            else:
+                route_info += f"({i+1}, {destination['name']})\n"
+
+        # routes
+        route_info += "Route: "
+        for i, node_id in enumerate(routes[0]):
+            if i == 0:
+                route_info += f"{tour_list[node_id]['name']} "
+            else:
+                route_info += f"> (step {i}) > {tour_list[node_id]['name']})"
+                if i == len(routes[0]) - 1:
+                    route_info += "\n"
+                else:
+                    route_info += " "
+        return route_info
+    
+    #--------------------------
+    # for validating a CF edge
+    #--------------------------
+    def validate_cf_edge(self,
+                         node_feats: Dict[str, np.array],
+                         dist_matrix: np.array,
+                         route: List[int],
+                         cf_step: int,
+                         cf_visit: int) -> bool:
+        # calc current time
+        curr_time = node_feats["time_window"][route[0]][0] # start point's open time
+        for step in range(1, cf_step):
+            curr_node_id = route[step-1]
+            next_node_id = route[step]
+            curr_time += node_feats["service_time"][curr_node_id] + dist_matrix[curr_node_id][next_node_id]
+            curr_time = max(curr_time, node_feats["time_window"][next_node_id][0]) # waiting
+
+        # validate the cf edge
+        curr_node_id = route[cf_step-1]
+        next_node_id = cf_visit
+        next_node_close_time = node_feats["time_window"][next_node_id][1] 
+        arrival_time = curr_time + node_feats["service_time"][curr_node_id] + dist_matrix[curr_node_id][next_node_id]
+        if next_node_close_time < arrival_time:
+            exceed_time = (arrival_time - next_node_close_time)
+            return False, f"Oops, your CF edge is infeasible because it does not meet the destination's close time by {util_app.add_time_unit(exceed_time)}."
+        else:
+            return True, "The CF edge is feasible!"
+
+    #-------------------------------
+    # for generating an explanation
+    #-------------------------------
+    def get_comparison_results(self,
+                               tour_list,
+                               question_summary,
+                               actual_routes: List[List[int]],
+                               actual_labels: List[List[int]],
+                               cf_routes: List[List[int]],
+                               cf_labels: List[List[int]],
+                               cf_step: int) -> str:
+        comparison_results = "Question:\n" + question_summary + "\n"
+        comparison_results += "Actual route:\n" + \
+                                self.get_route_info(tour_list, actual_routes[0], actual_labels[0], cf_step-1, "actual") + \
+                                self.get_representative_values(actual_routes[0], actual_labels[0], cf_step-1, "actual")
+        comparison_results += "CF route:\n" + \
+                                self.get_route_info(tour_list, cf_routes[0], cf_labels[0], cf_step-1, "CF") + \
+                                self.get_representative_values(cf_routes[0], cf_labels[0], cf_step-1, "CF")
+        comparison_results += "Difference between two routes:\n" + self.get_diff(cf_step-1, actual_routes[0], cf_routes[0])
+        comparison_results += "Planed desination information:\n" + self.get_node_info()
+        return comparison_results
+
+    def get_route_info(self,
+                       tour_list,
+                       route: List[int],
+                       label: List[int], 
+                       ex_step: int, 
+                       type: str) -> str:
+        def get_labelname(label_number):
+            return "route_len" if label_number == 0 else "time_window"
+        route_info = "- route: "
+        for i, node_id in enumerate(route):
+            if i == ex_step and i != len(route) - 1:
+                if type == "actual":
+                    edge_label = {get_labelname(label[i])}
+                else:
+                    edge_label = "user_preference"
+                route_info += f"{tour_list[node_id]['name']} > ({type} edge: {edge_label}) > "
+            elif i != len(route) - 1:
+                route_info += f"{tour_list[node_id]['name']} > ({get_labelname(label[i])}) > "
+            else:
+                route_info += f"{tour_list[node_id]['name']}\n"
+        return route_info
+
+    def get_representative_values(self, route, labels, ex_step, type) -> str:
+        time_window_ratio = self.get_intention_ratio(1, labels, ex_step) * 100
+        route_len_ratio = self.get_intention_ratio(0, labels, ex_step) * 100
+        return f"- short-term effect (immediate travel time): {self.get_immediate_state(route, ex_step)//60} minutes\n- long-term effect (total travel time): {self.get_route_length(route)//60} minutes\n- missed nodes: {self.get_infeasible_node_name(route)}\n- edge-intention ratio after the {type} edge: time_window {time_window_ratio: .1f}%, route_len {route_len_ratio: .1f}%"
+
+    def get_immediate_state(self, route, ex_step) -> str:
+        return st.session_state.dist_matrix[route[ex_step]][route[ex_step+1]]
+
+    def get_route_length(self, route) -> float:
+        route_length = 0.0
+        for i in range(len(route)-1):
+            route_length += st.session_state.dist_matrix[route[i]][route[i+1]]
+        return route_length
+
+    def get_infeasible_nodes(self, route) -> int:
+        return len(route) - (len(st.session_state.dist_matrix) - 1)
+
+    def get_infeasible_node_name(self, route) -> str:
+        if len(route) == len(st.session_state.dist_matrix) - 1:
+            return "none"
+        else:
+            num_nodes = np.arange(len(st.session_state.dist_matrix))
+            for node_id in route:
+                num_nodes = num_nodes[num_nodes != node_id]
+            return ",".join([st.session_state.tour_list[node_id]["name"] for node_id in num_nodes])
+
+    def get_intention_ratio(self, 
+                            intention: int, 
+                            labels: List[int], 
+                            ex_step: int) -> float:
+        np_labels = np.array(labels)
+        return np.sum(np_labels[ex_step:] == intention) / len(labels[ex_step:])
+
+    def get_diff(self, ex_step, actual_route, cf_route) -> str:
+        def get_str(effect: float):
+            long_effect_str = "The actual route increases it by" if effect > 0 else "The actual route reduces it by"
+            long_effect_str += util_app.add_time_unit(abs(effect))
+            return long_effect_str
+        
+        def get_str2(num_nodes: int, num_missed_nodes):
+            if num_nodes < 0:
+                num_nodes_str = f"The actual route visits {abs(num_nodes)} more nodes" 
+            elif num_nodes == 0:
+                if num_missed_nodes == 0:
+                    num_nodes_str = f"Both routes missed no node,"
+                else:
+                    num_nodes_str = f"Both routes missed the same number of nodes ({abs(num_missed_nodes)} node(s))"
+            else:
+                num_nodes_str = f"The actual route visits {abs(num_nodes)} less nodes" 
+            return num_nodes_str
+
+        # short/long-term effects
+        short_effect = self.get_immediate_state(actual_route, ex_step) - self.get_immediate_state(cf_route, ex_step)
+        long_effect  = self.get_route_length(actual_route) - self.get_route_length(cf_route)
+        short_effect_str = get_str(short_effect)
+        long_effect_str  = get_str(long_effect)
+
+        # missed nodes
+        missed_nodes = self.get_infeasible_nodes(actual_route) - self.get_infeasible_nodes(cf_route)
+        missed_nodes_str = get_str2(missed_nodes, self.get_infeasible_nodes(actual_route))
+
+        return f"- short-term effect: {short_effect_str}\n - long-term effect: {long_effect_str}\n- missed nodes: {missed_nodes_str}\n"
+
+    def get_node_info(self) -> str:
+        node_info = ""
+        for i in range(len(st.session_state.df_tour)):
+            node_info += f"- {st.session_state.df_tour['destination'][i]}: {st.session_state.df_tour['remarks'][i]}\n"
+        return node_info

models/solvers/concorde/concorde.py

@@ -0,0 +1,127 @@
+import torch.nn as nn
+import scipy
+import numpy as np
+import os
+import datetime
+import subprocess
+import models.solvers.concorde.concorde_utils as concorde_utils
+import glob
+import random
+
+class ConcordeTSP(nn.Module):
+    def __init__(self, large_value=1e+6, scaling=False, random_seed=1234, solver_dir="models/solvers/concorde/src/TSP", io_dir="concorde_io_files"):
+        self.random_seed = random_seed
+        self.large_value = large_value
+        self.scaling = scaling
+        self.solver_dir = solver_dir
+        self.io_dir = io_dir
+        self.redirector_stdout = concorde_utils.Redirector(fd=concorde_utils.STDOUT)
+        self.redirector_stderr = concorde_utils.Redirector(fd=concorde_utils.STDERR)
+        os.makedirs(io_dir, exist_ok=True)
+
+    def get_instance_name(self):
+        now = datetime.datetime.now()
+        random_value = random.random() # for avoiding duplicated file name
+        instance_name = f"{os.getpid()}_{random_value}_{now.strftime('%Y%m%d_%H%M%S%f')}"
+        return instance_name
+    
+    def write_instance(self, node_feats, fixed_paths=None, instance_name=None):
+        if instance_name is None:
+            instance_name = self.get_instance_name()
+        instance_fname = f"{self.io_dir}/{instance_name}.tsp"
+        tour_fname = f"{self.io_dir}/{instance_name}.sol"
+        with open(instance_fname, "w") as f:
+            f.write(f"NAME : {instance_name}\n")
+            f.write(f"TYPE : TSP\n")
+            f.write(f"DIMENSION : {len(node_feats['coords'])}\n")
+            self.write_data(f, node_feats, fixed_paths)
+            f.write("EOF\n")
+        return instance_fname, tour_fname
+
+    def write_data(self, f, node_feats, fixed_paths=None):
+        coords = node_feats["coords"]
+        if fixed_paths is None:
+            f.write("EDGE_WEIGHT_TYPE : EUC_2D\n")
+            f.write("NODE_COORD_SECTION\n")
+            for i in range(len(coords)):
+                f.write(f" {i + 1} {str(coords[i][0])[:10]} {str(coords[i][1])[:10]}\n")
+        else:
+            f.write("EDGE_WEIGHT_TYPE : EXPLICIT\n")
+            f.write("EDGE_WEIGHT_FORMAT : FULL_MATRIX\n")
+            f.write("EDGE_WEIGHT_SECTION\n")
+            dist = scipy.spatial.distance.cdist(coords, coords).round().astype(np.int64)
+            for i in range(len(fixed_paths)):
+                curr_id = fixed_paths[i]
+                if i != 0 and i != len(fixed_paths) - 1:
+                    # NOTE: concorde TSP seems to use int32, so 1e+9 occurs overflow.
+                    # 1e+8 could also do the same when N (tour length) is large.
+                    dist[curr_id, :] = 1e+8; dist[:, curr_id] = 1e+8
+                if i != 0:
+                    prev_id = fixed_paths[i - 1]
+                    dist[prev_id, curr_id] = 0; dist[curr_id, prev_id] = 0
+                if i != len(fixed_paths) - 1:
+                    next_id = fixed_paths[i + 1]
+                    dist[curr_id, next_id] = 0; dist[next_id, curr_id] = 0
+            f.write("\n".join([
+                " ".join(map(str, row))
+                for row in dist
+            ]))
+
+    def solve(self, node_feats, fixed_paths=None, instance_name=None):
+        if self.scaling:
+            node_feats = self.preprocess_data(node_feats)
+        self.redirector_stdout.start()
+        self.redirector_stderr.start()
+        instance_fname, tour_fname = self.write_instance(node_feats, fixed_paths, instance_name)
+        subprocess.run(f"{self.solver_dir}/concorde -o {tour_fname} -x {instance_fname}", shell=True) # run Concorde
+        self.redirector_stderr.stop()
+        self.redirector_stdout.stop()
+        tours = self.read_tour(tour_fname)
+        # remove dump (?) files
+        try:
+            os.remove(instance_fname); os.remove(tour_fname)
+        except OSError as e:
+            pass
+        fname_list = glob.glob("*.sol")
+        fname_list.extend(glob.glob("*.res"))
+        for fname in fname_list:
+            try:
+                os.remove(fname)
+            except OSError as e:
+                # do nothing
+                pass
+        # subprocess.run(f"rm {instance_name}.sol", shell=True)
+        return tours
+
+    def read_tour(self, tour_fname):
+        """
+        Parameters
+        ----------
+        tour_fname: str
+            path to an output tour
+        
+        Returns
+        -------
+        tour: 2d list [num_vehicles(1) x seq_length]
+        """
+        if not os.path.exists(tour_fname): # fails to solve the instance
+            return
+        
+        tour = []
+        with open(tour_fname, "r") as f:
+            for i, line in enumerate(f):
+                if i == 0:
+                    continue
+                read_tour = line.split()
+                tour.extend(read_tour)
+        tour.append(tour[0])
+        return [list(map(int, tour))]
+
+    def preprocess_data(self, node_feats):
+        # convert float to integer approximately
+        return {
+            key: (node_feat * self.large_value).astype(np.int64) 
+                 if key == "coords" else 
+                 node_feat
+            for key, node_feat in node_feats.items()
+        }

models/solvers/concorde/concorde_utils.py

@@ -0,0 +1,93 @@
+import os
+import sys
+import tempfile
+
+# for stopping std out from concorde solver
+# https://github.com/machine-reasoning-ufrgs/TSP-GNN/blob/master/redirector.py
+STDOUT = 1
+STDERR = 2
+
+class Redirector(object):
+    def __init__(self, fd=STDOUT):
+        self.fd = fd
+        self.started = False
+
+    def start(self):
+        if not self.started:
+            self.tmpfd, self.tmpfn = tempfile.mkstemp()
+
+            self.oldhandle = os.dup(self.fd)
+            os.dup2(self.tmpfd, self.fd)
+            os.close(self.tmpfd)
+
+            self.started = True
+
+    def flush(self):
+        if self.fd == STDOUT:
+            sys.stdout.flush()
+        elif self.fd == STDERR:
+            sys.stderr.flush()
+
+    def stop(self):
+        if self.started:
+            self.flush()
+            os.dup2(self.oldhandle, self.fd)
+            os.close(self.oldhandle)
+            tmpr = open(self.tmpfn, 'rb')
+            output = tmpr.read()
+            tmpr.close()  # this also closes self.tmpfd
+            os.unlink(self.tmpfn)
+
+            self.started = False
+            return output
+        else:
+            return None
+
+class RedirectorOneFile(object):
+    def __init__(self, fd=STDOUT):
+        self.fd = fd
+        self.started = False
+        self.inited = False
+
+        self.initialize()
+
+    def initialize(self):
+        if not self.inited:
+            self.tmpfd, self.tmpfn = tempfile.mkstemp()
+            self.pos = 0
+            self.tmpr = open(self.tmpfn, 'rb')
+            self.inited = True
+
+    def start(self):
+        if not self.started:
+            self.oldhandle = os.dup(self.fd)
+            os.dup2(self.tmpfd, self.fd)
+            self.started = True
+
+    def flush(self):
+        if self.fd == STDOUT:
+            sys.stdout.flush()
+        elif self.fd == STDERR:
+            sys.stderr.flush()
+
+    def stop(self):
+        if self.started:
+            self.flush()
+            os.dup2(self.oldhandle, self.fd)
+            os.close(self.oldhandle)
+            output = self.tmpr.read()
+            self.pos = self.tmpr.tell()
+            self.started = False
+            return output
+        else:
+            return None
+
+    def close(self):
+        if self.inited:
+            self.flush()
+            self.tmpr.close()  # this also closes self.tmpfd
+            os.unlink(self.tmpfn)
+            self.inited = False
+            return output
+        else:
+            return None

models/solvers/general_solver.py

@@ -0,0 +1,44 @@
+import torch.nn as nn
+from models.solvers.ortools.ortools import ORTools
+from models.solvers.lkh.lkh import LKH
+from models.solvers.concorde.concorde import ConcordeTSP
+
+class GeneralSolver(nn.Module):
+    def __init__(self, problem, solver_type, large_value=1e+6, scaling=True):
+        super().__init__()
+        self.problem = problem
+        self.large_value = large_value
+        self.scaling = scaling
+        self.solver_type = solver_type
+        supported_problem = {
+            "ortools": ["tsp", "tsptw", "pctsp", "pctsptw", "cvrp", "cvrptw"],
+            "lkh": ["tsp", "tsptw", "cvrp", "cvrptw"],
+            "concorde": ["tsp"]
+        }
+        # validate solver_type & problem
+        assert solver_type in supported_problem.keys(), f"Invalid solver type: {solver_type}. Please select from {supported_problem.keys()}"
+        assert problem in supported_problem[solver_type], f"{solver_type} does not support {problem}."
+        self.solver = self.get_solver(problem, solver_type)
+
+    def change_solver(self, problem, solver_type):
+        if self.solver_type != solver_type or self.problem != problem:
+            self.problem = problem
+            self.solver_type = solver_type
+            self.solver = self.get_solver(problem, solver_type)
+
+    def get_solver(self, problem, solver_type):
+        if solver_type == "ortools":
+            return ORTools(problem, self.large_value, self.scaling)
+        elif solver_type == "lkh":
+            return LKH(problem, self.large_value, self.scaling)
+        elif solver_type == "concorde":
+            assert problem == "tsp", "Concorde solver supports only TSP."
+            return ConcordeTSP(self.large_value, self.scaling)
+        else:
+            assert False, f"Invalid solver type: {solver_type}"
+
+    def solve(self, node_feats, fixed_paths=None, dist_matrix=None, instance_name=None):
+        if isinstance(self.solver, ORTools):
+            return self.solver.solve(node_feats, fixed_paths, dist_matrix, instance_name)
+        else:
+            return self.solver.solve(node_feats, fixed_paths, instance_name)

models/solvers/lkh/lkh.py

@@ -0,0 +1,21 @@
+import torch.nn as nn
+from models.solvers.lkh.lkh_tsp import LKHTSP
+from models.solvers.lkh.lkh_tsptw import LKHTSPTW
+from models.solvers.lkh.lkh_cvrp import LKHCVRP
+from models.solvers.lkh.lkh_cvrptw import LKHCVRPTW
+
+class LKH(nn.Module):
+    def __init__(self, problem, large_value=1e+6, scaling=False, max_trials=10, seed=1234, lkh_dir="models/solvers/lkh/src", io_dir="lkh_io_files"):
+        super().__init__()
+        self.probelm = problem
+        if problem == "tsp":
+            self.lkh = LKHTSP(large_value, scaling, max_trials, seed, lkh_dir, io_dir)
+        elif problem == "tsptw":
+            self.lkh = LKHTSPTW(large_value, scaling, max_trials, seed, lkh_dir, io_dir)
+        elif problem == "cvrp":
+            self.lkh = LKHCVRP(large_value, scaling, max_trials, seed, lkh_dir, io_dir)
+        elif problem == "cvrptw":
+            self.lkh = LKHCVRPTW(large_value, scaling, max_trials, seed, lkh_dir, io_dir)
+    
+    def solve(self, node_feats, fixed_paths=None, instance_name=None):
+        return self.lkh.solve(node_feats, fixed_paths, instance_name)

models/solvers/lkh/lkh_base.py

@@ -0,0 +1,157 @@
+import os
+import datetime
+import torch.nn as nn
+import numpy as np
+from subprocess import check_call
+
+NODE_ID_OFFSET = 1
+
+class LKHBase(nn.Module):
+    def __init__(self, problem, large_value=1e+6, scaling=False, max_trials=1000, seed=1234, lkh_dir="models/solvers/lkh", io_dir="lkh_io_files"):
+        super().__init__()
+        self.coord_dim = 2
+        self.problem = problem
+        self.large_value = large_value
+        self.scaling = scaling
+        self.max_trials = max_trials
+        self.seed = seed
+
+        # I/O file settings
+        self.lkh_dir = lkh_dir
+        self.io_dir = io_dir
+        self.instance_path = f"{io_dir}/{self.problem}/instance"
+        self.param_path    = f"{io_dir}/{self.problem}/param"
+        self.tour_path     = f"{io_dir}/{self.problem}/tour"
+        self.log_path      = f"{io_dir}/{self.problem}/log"
+        os.makedirs(self.instance_path, exist_ok=True) 
+        os.makedirs(self.param_path, exist_ok=True) 
+        os.makedirs(self.tour_path, exist_ok=True)
+        os.makedirs(self.log_path, exist_ok=True)
+
+    def solve(self, node_feats, fixed_paths=None, instance_name=None):
+        instance_fname = self.write_instance(node_feats, fixed_paths, instance_name)
+        param_fname, tour_fname, log_fname = self.write_para(instance_fname, instance_name)
+        with open(log_fname, "w") as f:
+            check_call([f"{self.lkh_dir}/LKH", param_fname], stdout=f) # run LKH
+        tours = self.read_tour(node_feats, tour_fname)
+        # clean intermidiate files
+        try:
+            os.remove(instance_fname); os.remove(param_fname); os.remove(tour_fname); os.remove(log_fname)
+        except:
+            pass
+        return tours
+    
+    def get_instance_name(self):
+        now = datetime.datetime.now()
+        instance_name = f"{os.getpid()}-{now.strftime('%Y%m%d_%H%M%S%f')}"
+        return instance_name
+
+    def write_instance(self, node_feats, fixed_paths=None, instance_name=None):
+        if instance_name is None:
+            instance_name = self.get_instance_name()
+        instance_fname = f"{self.instance_path}/{instance_name}.{self.problem}"
+        with open(instance_fname, "w") as f:
+            f.write(f"NAME : {instance_name}\n")
+            f.write(f"TYPE : {self.problem.upper()}\n")
+            f.write(f"DIMENSION : {len(node_feats['coords'])}\n")
+            self.write_data(node_feats, f)
+            if fixed_paths is not None and len(fixed_paths) > 1:
+                fixed_paths = fixed_paths.copy()
+                # FIXED_EDGE_SECTION works well in TSP, but it cannot fix edges in TSPTW
+                # EDGE_DATA_SECTION can fix edges in both TSP and TSPTW, but the obtained tour is very poor
+                f.write("FIXED_EDGES_SECTION\n")
+                fixed_paths += NODE_ID_OFFSET # offset node id (node id starts from 1 in TSPLIB)
+                for i in range(len(fixed_paths) - 1):
+                    f.write(f"{fixed_paths[i]} {fixed_paths[i+1]}\n")
+                # f.write("EDGE_DATA_FORMAT : EDGE_LIST\n")
+                # f.write("EDGE_DATA_SECTION\n")
+                # avail_edges = self.get_avail_edges(node_feats, fixed_paths)
+                # avail_edges += 1 # offset node id (node id starts from 1 in TSPLIB)
+                # for i in range(len(avail_edges)):
+                #     f.write(f"{avail_edges[i][0]} {avail_edges[i][1]}\n")
+            f.write("EOF\n")
+        return instance_fname
+
+    def write_data(self, node_feats, f):
+        raise NotImplementedError
+
+    def get_avail_edges(self, node_feats, fixed_paths):
+        num_nodes = len(node_feats["coords"])
+        avail_edges = []
+        # add fixed edges
+        for i in range(len(fixed_paths) - 1):
+            avail_edges.append([fixed_paths[i], fixed_paths[i + 1]])
+
+        # add rest avaialbel edges
+        visited = np.array([0] * num_nodes)
+        for id in fixed_paths:
+            visited[id] = 1
+        visited[fixed_paths[0]] = 0
+        visited[fixed_paths[-1]] = 0
+        mask = visited < 1
+        node_id = np.arange(num_nodes)
+        feasible_node_id = node_id[mask]
+        for j in range(len(feasible_node_id) - 1):
+            for i in range(j + 1, len(feasible_node_id)):
+                avail_edges.append([feasible_node_id[j], feasible_node_id[i]])
+        return np.array(avail_edges)
+
+    def write_para(self, instance_fname, instance_name=None):
+        if instance_name is None:
+            instance_name = self.get_instance_name()
+        param_fname = f"{self.param_path}/{instance_name}.param"
+        tour_fname  = f"{self.tour_path}/{instance_name}.tour"
+        log_fname   = f"{self.log_path}/{instance_name}.log"
+        with open(param_fname, "w") as f:
+            f.write(f"PROBLEM_FILE = {instance_fname}\n")
+            f.write(f"MAX_TRIALS = {self.max_trials}\n")
+            f.write("MOVE_TYPE = 5\nPATCHING_C = 3\nPATCHING_A = 2\nRUNS = 1\n")
+            f.write(f"SEED = {self.seed}\n")
+            f.write(f"OUTPUT_TOUR_FILE = {tour_fname}\n")
+        return param_fname, tour_fname, log_fname
+
+    def read_tour(self, node_feats, tour_fname):
+        """
+        Parameters
+        ----------
+        output_filename: str
+            path to a file where optimal tour is written
+        Returns
+        -------
+        tour: 2d list [num_vehicles x seq_length]
+            a set of node ids indicating visit order
+        """
+        if not os.path.exists(tour_fname):
+            return # found no feasible solution
+
+        with open(tour_fname, "r") as f:
+            tour = []
+            is_tour_section = False 
+            for line in f:
+                line = line.strip()
+                if line == "TOUR_SECTION":
+                    is_tour_section = True
+                    continue
+                if is_tour_section:
+                    if line != "-1":
+                        tour.append(int(line) - NODE_ID_OFFSET)
+                    else:
+                        tour.append(tour[0])
+                        break
+        # convert 1d -> 2d list
+        num_nodes = len(node_feats["coords"])
+        tour = np.array(tour)
+        # NOTE: node_id >= num_nodes indicates the depot node. 
+        # That is because LKH uses dummy nodes of which locations are the same as the depot and demands = -capacity?
+        # I'm not sure where the behavior is documented, but the author of NeuroLKH reads output files like that.
+        # please refer to https://github.com/liangxinedu/NeuroLKH/blob/main/CVRPTWdata_generate.py#L132
+        tour[tour >= num_nodes] = 0
+        # remove subsequent zeros
+        tour = tour[np.diff(np.concatenate(([1], tour))).nonzero()]
+        loc0 = (tour == 0).nonzero()[0]
+        num_vehicles = len(loc0) - 1
+        tours = []
+        for vehicle_id in range(num_vehicles):
+            vehicle_tour = tour[loc0[vehicle_id]:loc0[vehicle_id+1]+1].tolist()
+            tours.append(vehicle_tour)
+        return tours # offset to make the first index 0

models/solvers/lkh/lkh_cvrp.py

@@ -0,0 +1,41 @@
+import numpy as np
+import scipy
+from models.solvers.lkh.lkh_base import LKHBase
+
+class LKHCVRP(LKHBase):
+    def __init__(self, large_value=1e+6, scaling=False, max_trials=1000, seed=1234, lkh_dir="models/solvers/lkh/", io_dir="lkh_io_files"):
+        problem = "cvrp"
+        super().__init__(problem, large_value, scaling, max_trials, seed, lkh_dir, io_dir)
+    
+    def write_data(self, node_feats, f):
+        """
+        Paramters
+        ---------
+        node_feats: dict of np.array 
+            coords: np.array [num_nodes x coord_dim]
+            demand: np.array [num_nodes x 1]
+            capacity: np.array [1]
+        """
+        coords = node_feats["coords"]
+        demand = node_feats["demand"]
+        capacity = node_feats["capacity"][0]
+        num_nodes = len(coords)
+        if self.scaling:
+            coords = coords * self.large_value
+        # NOTE: In CVRP, LKH can automatically obtain optimal vehicle size.
+        # However it cannot in CVRPTW (please check lkh_cvrptw.py).
+        # EDGE_WEIGHT_SECTION
+        f.write("EDGE_WEIGHT_TYPE : EUC_2D\n")
+        # CAPACITY
+        f.write("CAPACITY : " + str(capacity) + "\n")
+        # NODE_COORD_SECTION
+        f.write("NODE_COORD_SECTION\n")
+        for i in range(num_nodes):
+            f.write(f" {i + 1} {str(coords[i][0])[:10]} {str(coords[i][1])[:10]}\n")
+        # DEMAND_SECTION
+        f.write("DEMAND_SECTION\n")
+        for i in range(num_nodes):
+            f.write(f" {i + 1} {str(demand[i])}\n")
+        # DEPOT SECTION
+        f.write("DEPOT_SECTION\n")
+        f.write("1\n")

models/solvers/lkh/lkh_cvrptw.py

@@ -0,0 +1,59 @@
+import numpy as np
+import scipy
+from models.solvers.lkh.lkh_base import LKHBase
+
+class LKHCVRPTW(LKHBase):
+    def __init__(self, large_value=1e+6, scaling=False, max_trials=1000, seed=1234, lkh_dir="models/solvers/lkh/", io_dir="lkh_io_files"):
+        problem = "cvrptw"
+        super().__init__(problem, large_value, scaling, max_trials, seed, lkh_dir, io_dir)
+    
+    def write_data(self, node_feats, f):
+        """
+        Paramters
+        ---------
+        node_feats: dict of np.array 
+            coords: np.array [num_nodes x coord_dim]
+            demand: np.array [num_nodes x 1]
+            capacity: np.array [1]
+            time_window: np.array [num_nodes x 2(start, end)]
+        """
+        coords = node_feats["coords"]
+        demand = node_feats["demand"]
+        capacity = node_feats["capacity"][0]
+        time_window = node_feats["time_window"]
+        num_nodes = len(coords)
+        if self.scaling:
+            coords = coords * self.large_value
+            time_window = time_window * self.large_value
+        # VEHICLES
+        # As the number of unused vehicles is also included to penalty in default,
+        # we have to modify Penalty_CVRTW.c in LKH SRC directory.
+        # Comment out the following part, which corresponds to penaly of unsed vehicles:
+        #   42 if (MTSPMinSize >= 1 && Size < MTSPMinSize)
+        #   43    P += MTSPMinSize - Size;
+        #   44 if (Size > MTSPMaxSize)
+        #   45    P += Size - MTSPMaxSize;
+        # After the modification, we can automatically obtain optimal vehicle size 
+        # by setting large vehicle size (e.g. >20) here
+        f.write("VEHICLES : 20\n")
+        # CAPACITY
+        f.write("CAPACITY : " + str(capacity) + "\n")
+        # EDGE_WEIGHT_SECTION
+        f.write("EDGE_WEIGHT_TYPE : EUC_2D\n")
+        # NODE_COORD_SECTION
+        f.write("NODE_COORD_SECTION\n")
+        for i in range(num_nodes):
+            f.write(f" {i + 1} {str(coords[i][0])[:10]} {str(coords[i][1])[:10]}\n")
+        # DEMAND_SECTION
+        f.write("DEMAND_SECTION\n")
+        for i in range(num_nodes):
+            f.write(f" {i + 1} {str(demand[i])}\n")
+        # TIME_WINDOW_SECTION
+        f.write("TIME_WINDOW_SECTION\n")
+        f.write("\n".join([
+            "{}\t{}\t{}".format(i + 1, l, u)
+            for i, (l, u) in enumerate(time_window)
+        ]))
+        # DEPOT SECTION
+        f.write("DEPOT_SECTION\n")
+        f.write("1\n")

models/solvers/lkh/lkh_tsp.py

@@ -0,0 +1,15 @@
+from models.solvers.lkh.lkh_base import LKHBase
+
+class LKHTSP(LKHBase):
+    def __init__(self, large_value=1e+6, scaling=False, max_trials=1000, seed=1234, lkh_dir="models/solvers/lkh/", io_dir="lkh_io_files"):
+        problem = "tsp"
+        super().__init__(problem, large_value, scaling, max_trials, seed, lkh_dir, io_dir)
+
+    def write_data(self, node_feats, f):
+        coords = node_feats["coords"]
+        if self.scaling:
+            coords = coords * self.large_value
+        f.write("EDGE_WEIGHT_TYPE : EUC_2D\n")
+        f.write("NODE_COORD_SECTION\n")
+        for i in range(len(coords)):
+            f.write(f" {i + 1} {str(coords[i][0])[:10]} {str(coords[i][1])[:10]}\n")

models/solvers/lkh/lkh_tsptw.py

@@ -0,0 +1,32 @@
+import numpy as np
+import scipy
+from models.solvers.lkh.lkh_base import LKHBase
+
+class LKHTSPTW(LKHBase):
+    def __init__(self, large_value=1e+6, scaling=False, max_trials=1000, seed=1234, lkh_dir="models/solvers/lkh/", io_dir="lkh_io_files"):
+        problem = "tsptw"
+        super().__init__(problem, large_value, scaling, max_trials, seed, lkh_dir, io_dir)
+
+    def write_data(self, node_feats, f):
+        coord_dim = 2
+        coords = node_feats["coords"]
+        if self.scaling:
+            coords = coords * self.large_value
+        time_window = node_feats["time_window"].astype(np.int64)
+        dist = scipy.spatial.distance.cdist(coords, coords).round().astype(np.int64)
+        f.write("EDGE_WEIGHT_TYPE : EXPLICIT\n")
+        f.write("EDGE_WEIGHT_FORMAT : FULL_MATRIX\n")
+        f.write("EDGE_WEIGHT_SECTION\n")
+        f.write("\n".join([
+            " ".join(map(str, row))
+            for row in dist
+        ]))
+        f.write("\n")
+        f.write("TIME_WINDOW_SECTION\n")
+        f.write("\n".join([
+            "{}\t{}\t{}".format(i + 1, l, u)
+            for i, (l, u) in enumerate(time_window)
+        ]))
+        f.write("\n")
+        f.write("DEPOT_SECTION\n")
+        f.write("1\n")

models/solvers/ortools/ortools.py

@@ -0,0 +1,58 @@
+import torch.nn as nn
+from models.solvers.ortools.ortools_tsp import ORToolsTSP
+from models.solvers.ortools.ortools_tsptw import ORToolsTSPTW
+from models.solvers.ortools.ortools_pctsp import ORToolsPCTSP
+from models.solvers.ortools.ortools_pctsptw import ORToolsPCTSPTW
+from models.solvers.ortools.ortools_cvrp import ORToolsCVRP
+from models.solvers.ortools.ortools_cvrptw import ORToolsCVRPTW
+
+class ORTools(nn.Module):
+    def __init__(self, problem, large_value=1e+6, scaling=False):
+        super().__init__()
+        self.coord_dim = 2
+        self.problem = problem
+        self.large_value = large_value
+        self.scaling = scaling
+        self.ortools = self.get_ortools(problem)
+
+    def get_ortools(self, problem):
+        """
+        Parameters
+        ----------
+        problem: str
+            problem type
+        
+        Returns
+        -------
+        ortools: ortools for the specified problem
+        """
+        if problem == "tsp":
+            return ORToolsTSP(self.large_value, self.scaling)
+        elif problem == "tsptw":
+            return ORToolsTSPTW(self.large_value, self.scaling)
+        elif problem == "pctsp":
+            return ORToolsPCTSP(self.large_value, self.scaling)
+        elif problem == "pctsptw":
+            return ORToolsPCTSPTW(self.large_value, self.scaling)
+        elif problem == "cvrp":
+            return ORToolsCVRP(self.large_value, self.scaling)
+        elif problem == "cvrptw":
+            return ORToolsCVRPTW(self.large_value, self.scaling)
+        else:
+            raise NotImplementedError
+
+    def solve(self, node_feats, fixed_paths=None, dist_martix=None, instance_name=None):
+        """
+        Parameters
+        ----------
+        node_feats: np.array [num_nodes x node_dim]
+        fixed_paths: np.array [cf_step]
+        scaling: bool
+            whether or not coords are muliplied by a large value
+            to convert float-coods into int-coords
+
+        Returns
+        -------
+        tour: np.array [seq_length]
+        """
+        return self.ortools.solve(node_feats, fixed_paths, dist_martix, instance_name)