Implementing shortest path mod

src/gurobi_optimods/shortest_path.py

@@ -0,0 +1,113 @@
+"""
+(Elementary) Shortest Path Problem with Resource Constraints
+------------------------------------------------------------
+"""
+
+import logging
+from dataclasses import dataclass
+
+import gurobipy as gp
+import gurobipy_pandas as gppd
+import pandas as pd
+from gurobipy import GRB
+
+from gurobi_optimods.utils import optimod
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ShortestPath:
+    """
+    Solution to a QUBO problem.
+
+    Attributes
+    ----------
+    nodes : list
+        ordered list of visited nodes along the path
+    cost : float
+        total path costs
+    """
+
+    nodes: list
+    cost: float
+
+
+@optimod()
+def shortest_path(
+    node_data, arc_data, source, target, limits={}, elementary=True, *, create_env
+) -> ShortestPath:
+    """
+    An optimod that solves an important problem
+
+    :param data: Description of argument
+    :type data: Type of argument
+
+    ... describe additional arguments ...
+
+    :return: Description of returned result
+    :rtype: Type of returned result
+    """
+
+    # add node properties to arc properties (matching the source node of each arc)
+    print(node_data)
+    print(arc_data)
+
+    arcnode_data = arc_data.add(node_data, fill_value=0)
+    print(arcnode_data)
+
+    params = {"LogToConsole": 1}
+
+    with create_env(params=params) as env, gp.Model(env=env) as model:
+        model.ModelSense = GRB.MINIMIZE
+
+        arc_df = arcnode_data.gppd.add_vars(
+            model, vtype=GRB.INTEGER, lb=0, obj="cost", name="arc"
+        )
+
+        # add cost of target node
+        model.ObjCon = node_data.loc[target, "cost"]
+
+        node_flows = pd.DataFrame(
+            {
+                "inflow": arc_df["arc"].groupby("j").sum(),
+                "outflow": arc_df["arc"].groupby("i").sum(),
+            }
+        ).fillna(0)
+        # balance value: source = 1, target = -1, 0 otherwise
+        node_flows["balance"] = (node_flows.index == source).astype(int) - (
+            node_flows.index == target
+        ).astype(int)
+
+        model.update()
+        print(node_flows)
+
+        # flow conservation constraints
+        node_flows.gppd.add_constrs(
+            model, "outflow - inflow == balance", name="balance_constr"
+        )
+
+        # restriction to elementary paths
+        if elementary:
+            node_flows.gppd.add_constrs(model, "inflow <= 1", name="elementary_constr")
+
+        # TODO resource restrictions
+
+        logger.info(f"Solving shortest path problem")
+
+        model.write("model.lp")
+        model.optimize()
+
+        if model.Status in [GRB.INFEASIBLE, GRB.INF_OR_UNBD]:
+            raise ValueError("Infeasible or unbounded")
+
+        if model.SolCount == 0:
+            raise ValueError(
+                "No solution found, potentially because of a very low time limit."
+            )
+
+        # TODO extract solution and create node visiting sequence
+        arc_df["solution"] = arc_df["arc"].gppd.X
+        print(arc_df)
+
+        return ShortestPath(nodes=[0], cost=model.ObjVal)

tests/test_shortest_path.py

@@ -0,0 +1,55 @@
+import io
+import unittest
+
+import pandas as pd
+import scipy.sparse as sp
+
+import gurobi_optimods.shortest_path as sp
+
+node_data_csv = """
+i,cost,time
+0,1,5
+1,1,4
+2,1,3
+3,1,2
+4,1,1
+5,1,0
+"""
+
+arc_data_csv = """
+i,j,cost
+0,1,1
+0,2,3
+1,2,-5
+1,3,6
+2,3,4
+2,4,2
+3,4,-5
+3,5,1
+4,2,0
+4,5,3
+"""
+
+
+def load_graph_pandas():
+    return (
+        pd.read_csv(io.StringIO(node_data_csv)).set_index(["i"]),
+        pd.read_csv(io.StringIO(arc_data_csv)).set_index(["i", "j"]),
+    )
+
+
+class TestShortestPath(unittest.TestCase):
+    def test_negative_costs(self):
+        node_data, arc_data = load_graph_pandas()
+        path = sp.shortest_path(node_data, arc_data, 0, 5)
+        self.assertEqual(path.cost, 4)
+        arc_list = [0, 1, 2, 3, 4, 5]
+        self.assertIsInstance(path.nodes, list)
+        self.assertEqual(path.nodes, arc_list)
+
+    # def test_infeasible(self):
+    #     edge_data, node_data = datasets.simple_graph_pandas()
+    #     # Add a node requesting more flow than is available.
+    #     node_data["demand"].values[-1] = 10.0
+    #     with self.assertRaisesRegex(ValueError, "Unsatisfiable flows"):
+    #         obj, sol = mcf.min_cost_flow_pandas(edge_data, node_data)