Encapsulate arc and vertex variables in Vars class

Author: Plastix

src/main/java/io/github/plastix/Main.java

@@ -13,8 +13,6 @@
 import com.graphhopper.util.EdgeIterator;
 import gurobi.*;
 
-import java.util.Arrays;
-
 public class Main {
 
     private static Params params;
@@ -27,33 +25,14 @@ public class Main {
     private static Weighting score;
     private static FlagEncoder flagEncoder;
 
+    // Solver variables
     private static GRBEnv env;
     private static GRBModel model;
-    private static GRBVar[] arcsFwd;
-    private static GRBVar[] arcsBwd;
-    private static int[] arcBaseIds;
 
     private static void setupSolver() throws GRBException {
         env = new GRBEnv("osm.log");
         model = new GRBModel(env);
-
-        AllEdgesIterator edges = graph.getAllEdges();
-        int numEdges = edges.getMaxId();
-        int numNodes = graph.getNodes();
-
-        // Make a decision variable "x_i" for every arc in our graph
-        // x_i: =1 if arc is travelled, 0 otherwise
-        // Since every edge can be 2 arcs (forward + backward), we keep two lists
-        arcsFwd = model.addVars(numEdges, GRB.BINARY);
-        arcsBwd = model.addVars(numEdges, GRB.BINARY);
-        arcBaseIds = new int[numEdges];
-
-        // Make a variable for every node in the graph
-        // verts[i]: = the number of times vertex i is visited
-        char[] types = new char[numNodes];
-        Arrays.fill(types, GRB.INTEGER);
-        GRBVar[] verts = model.addVars(null, null, null, types,
-                null, 0, numNodes);
+        Vars vars = new Vars(graph, model, flagEncoder);
 
         // (1a)
         // Objective maximizes total collected score of all roads
@@ -63,51 +42,43 @@ private static void setupSolver() throws GRBException {
         // Limit length of path
         GRBLinExpr maxCost = new GRBLinExpr();
 
+        AllEdgesIterator edges = graph.getAllEdges();
         while(edges.next()) {
-            int edgeId = edges.getEdge();
-            arcBaseIds[edgeId] = edges.getBaseNode(); // Record the original base ID to keep direction
-            GRBVar fwd = arcsFwd[edgeId];
-            GRBVar bwd = arcsBwd[edgeId];
-            double edgeScore = score.calcWeight(edges, false, edgeId);
+            double edgeScore = score.calcWeight(edges, false, edges.getEdge());
             double edgeDist = edges.getDistance();
 
-            if(edges.isForward(flagEncoder)) {
-                objective.addTerm(edgeScore, fwd);
-                maxCost.addTerm(edgeDist, fwd);
-            } else {
-                fwd.set(GRB.DoubleAttr.UB, 0);
-            }
+            GRBVar forward = vars.getArc(edges);
+            GRBVar backward = vars.getComplementArc(edges);
 
-            if(edges.isBackward(flagEncoder)) {
-                objective.addTerm(edgeScore, bwd);
-                maxCost.addTerm(edgeDist, bwd);
-            } else {
-                bwd.set(GRB.DoubleAttr.UB, 0);
-            }
+            objective.addTerm(edgeScore, forward);
+            objective.addTerm(edgeScore, backward);
+            maxCost.addTerm(edgeDist, forward);
+            maxCost.addTerm(edgeDist, backward);
 
             // (1j)
             GRBLinExpr arcConstraint = new GRBLinExpr();
-            arcConstraint.addTerm(1, fwd);
-            arcConstraint.addTerm(1, bwd);
+            arcConstraint.addTerm(1, forward);
+            arcConstraint.addTerm(1, backward);
             model.addConstr(arcConstraint, GRB.LESS_EQUAL, 1, "arc_constraint");
         }
 
         model.setObjective(objective, GRB.MAXIMIZE);
         model.addConstr(maxCost, GRB.LESS_EQUAL, params.getMaxCost(), "max_cost");
 
+        int numNodes = graph.getNodes();
         for(int i = 0; i < numNodes; i++) {
             // (1d)
             GRBLinExpr edgeCounts = new GRBLinExpr();
             IntHashSet incomingIds = new IntHashSet();
             EdgeIterator incoming = graphUtils.incomingEdges(i);
             while(incoming.next()) {
                 incomingIds.add(incoming.getEdge());
-                edgeCounts.addTerm(1, getArc(incoming));
+                edgeCounts.addTerm(1, vars.getArc(incoming));
             }
 
             EdgeIterator outgoing = graphUtils.outgoingEdges(i);
             while(outgoing.next()) {
-                GRBVar arc = getArc(outgoing);
+                GRBVar arc = vars.getArc(outgoing);
                 // Check if we already recorded it as an incoming edge
                 if(incomingIds.contains(outgoing.getEdge())) {
                     edgeCounts.remove(arc);
@@ -122,29 +93,19 @@ private static void setupSolver() throws GRBException {
             GRBLinExpr vertexVisits = new GRBLinExpr();
             outgoing = graphUtils.outgoingEdges(i);
             while(outgoing.next()) {
-                vertexVisits.addTerm(1, getArc(outgoing));
+                vertexVisits.addTerm(1, vars.getArc(outgoing));
             }
-            vertexVisits.addTerm(-1, verts[i]);
+            vertexVisits.addTerm(-1, vars.getVertex(i));
             model.addConstr(vertexVisits, GRB.EQUAL, 0, "vertex_visits");
         }
 
         // (1h)/(1i)
         // Start vertex can only be visited once
-        GRBVar startNode = verts[START_NODE_ID];
+        GRBVar startNode = vars.getVertex(START_NODE_ID);
         startNode.set(GRB.DoubleAttr.LB, 1);
         startNode.set(GRB.DoubleAttr.UB, 1);
     }
 
-    private static GRBVar getArc(EdgeIterator edge) {
-        int edgeId = edge.getEdge();
-        int baseNode = edge.getBaseNode();
-        if(baseNode == arcBaseIds[edgeId]) {
-            return arcsFwd[edgeId];
-        } else {
-            return arcsBwd[edgeId];
-        }
-    }
-
     private static void runSolver() throws GRBException {
         System.out.println("Max cost: " + params.getMaxCost());
         System.out.println("Start position: " + params.getStartLat() + ", " + params.getStartLon() +

src/main/java/io/github/plastix/Vars.java

@@ -0,0 +1,89 @@
+package io.github.plastix;
+
+import com.graphhopper.routing.util.AllEdgesIterator;
+import com.graphhopper.routing.util.FlagEncoder;
+import com.graphhopper.storage.Graph;
+import com.graphhopper.util.EdgeIterator;
+import gurobi.GRB;
+import gurobi.GRBException;
+import gurobi.GRBModel;
+import gurobi.GRBVar;
+
+import java.util.Arrays;
+
+public class Vars {
+
+    private Graph graph;
+    private FlagEncoder flagEncoder;
+    private GRBModel model;
+
+    // arcs[arcId][0] = forwardArc
+    // arcs[arcId][1] = backwardArc
+    private GRBVar[][] arcs;
+    private GRBVar[] verts;
+    private int[] arcBaseIds;
+
+    Vars(Graph graph, GRBModel model, FlagEncoder flagEncoder) throws GRBException {
+        this.graph = graph;
+        this.model = model;
+        this.flagEncoder = flagEncoder;
+        addVarsToModel();
+    }
+
+    private void addVarsToModel() throws GRBException {
+        AllEdgesIterator edges = graph.getAllEdges();
+        int numEdges = edges.getMaxId();
+        int numNodes = graph.getNodes();
+        arcBaseIds = new int[numEdges];
+
+        // Make a decision variable for every arc in our graph
+        // arcs[i] = 1 if arc is travelled, 0 otherwise
+        GRBVar[] arcVars = model.addVars(2 * numEdges, GRB.BINARY);
+        arcs = new GRBVar[numEdges][2];
+
+        // Make a variable for every node in the graph
+        // verts[i] = n the number of times vertex i is visited
+        char[] types = new char[numNodes];
+        Arrays.fill(types, GRB.INTEGER);
+        verts = model.addVars(null, null, null, types,
+                null, 0, numNodes);
+
+        int i = 0;
+        while(edges.next()) {
+            int edgeId = edges.getEdge();
+            int baseNode = edges.getBaseNode();
+            arcBaseIds[edgeId] = baseNode;
+
+            GRBVar forward = arcVars[i++];
+            GRBVar backward = arcVars[i++];
+
+            arcs[edgeId][0] = forward;
+            arcs[edgeId][1] = backward;
+
+            if(!edges.isForward(flagEncoder)) {
+                forward.set(GRB.DoubleAttr.UB, 0);
+            }
+
+            if(!edges.isBackward(flagEncoder)) {
+                backward.set(GRB.DoubleAttr.UB, 0);
+            }
+        }
+    }
+
+    private int getIndex(EdgeIterator edge) {
+        return arcBaseIds[edge.getEdge()] == edge.getBaseNode() ? 0 : 1;
+    }
+
+    public GRBVar getArc(EdgeIterator edge) {
+        return arcs[edge.getEdge()][getIndex(edge)];
+    }
+
+    public GRBVar getComplementArc(EdgeIterator edge) {
+        return arcs[edge.getEdge()][getIndex(edge) ^ 1];
+    }
+
+    public GRBVar getVertex(int id) {
+        return verts[id];
+    }
+
+}