infer.py

import argparse
import json
from pathlib import Path
import torch
import numpy as np

from gnn import VRPEdgePredictor
from cluster import constrained_agglomerative_clustering


def run_inference(graph_json_path: str, checkpoint_path: str, f0_norm: float, f1_norm: float, knn_k: int = 20, prob_threshold: float = 0.5):
    """
    End-to-End Inference: Given a VRP graph JSON, predicting valid capacity-constrained routes.
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"--- Initialization ---")
    print(f"Device: {device}")
    
    # 1. Load the Graph JSON
    print(f"Loading graph: {graph_json_path}")
    with open(graph_json_path, "r", encoding="utf-8") as f:
        graph_data = json.load(f)
        
    capacity = float(graph_data["capacity"])
    num_nodes = int(graph_data["num_nodes"])
    
    node_features = []
    coords = []
    demands_list = []
    
    # Extract raw extremes for Min-Max normalization
    x_raw = [float(graph_data["coordinates"][i][0]) for i in range(num_nodes)]
    y_raw = [float(graph_data["coordinates"][i][1]) for i in range(num_nodes)]
    x_min, x_max = min(x_raw), max(x_raw)
    y_min, y_max = min(y_raw), max(y_raw)
    
    for i in range(num_nodes):
        cx_raw, cy_raw = float(graph_data["coordinates"][i][0]), float(graph_data["coordinates"][i][1])
        cx_norm = (cx_raw - x_min) / (x_max - x_min + 1e-8)
        cy_norm = (cy_raw - y_min) / (y_max - y_min + 1e-8)
        
        dem = float(graph_data["demands"][str(i)])
        dem_norm = dem / capacity
        
        # Features normalized: [x_norm, y_norm, dem_norm, target_f0, target_f1]
        node_features.append([cx_norm, cy_norm, dem_norm, f0_norm, f1_norm])
        coords.append([cx_norm, cy_norm])
        demands_list.append(dem)
        
    x = torch.tensor(node_features, dtype=torch.float).to(device)
    pos = torch.tensor(coords, dtype=torch.float).to(device)
    demands = torch.tensor(demands_list, dtype=torch.float) # Kept on CPU for Python-side clustering logic
    
    # 2. Rebuild kNN Topological Graph
    num_neighbors = min(knn_k + 1, num_nodes)
    dists = torch.cdist(pos, pos)
    _, topk_idx = torch.topk(dists, k=num_neighbors, dim=1, largest=False)
    
    neighbors = topk_idx[:, 1:]
    src_nodes = torch.arange(num_nodes, device=device).repeat_interleave(num_neighbors - 1)
    dst_nodes = neighbors.flatten()
    edge_index = torch.stack([src_nodes, dst_nodes], dim=0)
    
    # Rebuild Edge Attributes Matrix
    f0_vals = [float(e["distance"]) for e in graph_data["edges"]]
    f1_vals = [float(e["random_cost"]) for e in graph_data["edges"]]
    dist_min, dist_max = min(f0_vals), max(f0_vals)
    rc_min, rc_max = min(f1_vals), max(f1_vals)
    
    f0_edge_map = {}
    f1_edge_map = {}
    for e in graph_data["edges"]:
        f0_edge_map[(e["source"], e["target"])] = float(e["distance"])
        f1_edge_map[(e["source"], e["target"])] = float(e["random_cost"])
        
    edge_attr = []
    for i in range(edge_index.shape[1]):
        src = int(edge_index[0, i])
        dst = int(edge_index[1, i])
        c0 = f0_edge_map.get((src, dst), f0_edge_map.get((dst, src), 0.0))
        c1 = f1_edge_map.get((src, dst), f1_edge_map.get((dst, src), 0.0))
        
        c0_norm = (c0 - dist_min) / (dist_max - dist_min + 1e-8)
        c1_norm = (c1 - rc_min) / (rc_max - rc_min + 1e-8)
        edge_attr.append([c0_norm, c1_norm])
        
    edge_attr = torch.tensor(edge_attr, dtype=torch.float).to(device)
    
    # 3. Load Model
    print(f"Loading weights: {checkpoint_path}")
    model = VRPEdgePredictor(node_in_dim=5, edge_in_dim=2, hidden_dim=64, num_layers=4).to(device)
    model.load_state_dict(torch.load(checkpoint_path, map_location=device, weights_only=True))
    model.eval()
    
    # 4. Forward Pass (Predict Edge Probabilities)
    print(f"Conditioning Neural Network -> Target F0 (Dist): {f0_norm:.2f}, Target F1 (RC): {f1_norm:.2f}")
    with torch.no_grad():
        logits = model(x, edge_index, edge_attr)
        probs = torch.sigmoid(logits)
    
    # 5. Greedy Capacity-Constrained Clustering Post-Processing
    print(f"Executing Agglomerative Clustering (Threshold >= {prob_threshold})...")
    routes = constrained_agglomerative_clustering(
        probs=probs.cpu(),
        edge_index=edge_index.cpu(),
        demands=demands,
        capacity=capacity,
        prob_threshold=prob_threshold
    )
    
    # 6. Result formatting
    print("\n--- INFERENCE RESULTS ---")
    print(f"Number of Output Vehicles (Routes): {len(routes)}")
    for i, route in enumerate(routes):
        route_demand = sum(demands[node].item() for node in route)
        print(f"Route {i+1} [Demand {route_demand:.1f} / {capacity:.1f}]: 0 -> {' -> '.join(map(str, route))} -> 0")
        
    return routes

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="VRP GNN Inference & Clustering")
    parser.add_argument("--graph", type=str, required=True, help="Path to input X-nXX-kX.json graph")
    parser.add_argument("--checkpoint", type=str, required=True, help="Path to pre-trained .pt model file")
    parser.add_argument("--f0", type=float, required=True, help="Target F0 preference (0.0 to 1.0 where 0 is optimize heavily for F0)")
    parser.add_argument("--f1", type=float, required=True, help="Target F1 preference (0.0 to 1.0 where 0 is optimize heavily for F1)")
    parser.add_argument("--threshold", type=float, default=0.5, help="Min probability to group nodes")
    parser.add_argument("--out", type=str, default=None, help="Path to save the clustered routes as a JSON file (e.g. for Rust genetic_tsp)")
    
    args = parser.parse_args()
    
    routes = run_inference(args.graph, args.checkpoint, args.f0, args.f1, prob_threshold=args.threshold)
    
    if args.out:
        with open(args.out, 'w') as f:
            json.dump(routes, f, indent=2)
        print(f"\nSaved JSON clusters directly to: {args.out}")