Implement multi-sample shortest path and create an example notebook

corneto/backend/_base.py

@@ -1274,10 +1274,14 @@ def Xor(self, x: CExpression, y: CExpression, varname="_xor"):
         )
 
     def linear_or(
-        self, x: CExpression, axis: Optional[int] = None, varname="or"
+        self,
+        x: CExpression,
+        axis: Optional[int] = None,
+        varname="or",
+        ignore_type=False,
     ) -> ProblemDef:
         # Check if the variable has a vartype and is binary
-        if hasattr(x, "_vartype") and x._vartype != VarType.BINARY:
+        if hasattr(x, "_vartype") and x._vartype != VarType.BINARY and not ignore_type:
             raise ValueError(f"Variable x has type {x._vartype} instead of BINARY")
         else:
             for s in x._proxy_symbols:

corneto/methods/shortest_path.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any, List, Optional
 
 import numpy as np
 
@@ -8,6 +8,66 @@
 from corneto.backend._base import DEFAULT_UB, Indicator
 
 
+def create_multisample_shortest_path(
+    G: BaseGraph,
+    source_target_nodes: List[tuple],
+    edge_weights=None,
+    solver: Optional[str] = None,
+    backend: Backend = DEFAULT_BACKEND,
+    lam: float = 0.0,
+):
+    # Transform the graph into a flow problem
+    Gc = G.copy()
+    inflow_edges = dict()
+    outflow_edges = dict()
+    for s, t in source_target_nodes:
+        if s not in inflow_edges:
+            inflow_edges[s] = Gc.add_edge((), s)
+        if t not in outflow_edges:
+            outflow_edges[t] = Gc.add_edge(t, ())
+    if edge_weights is None:
+        edge_weights = np.array(
+            [Gc.get_attr_edge(i).get("weight", 0) for i in range(Gc.ne)]
+        )
+        # The number of samples equals the number of source-target pairs.
+        # We need to duplicate the edge weights for each sample.
+        edge_weights = np.tile(edge_weights, (len(source_target_nodes), 1))
+    else:
+        # Verify that the number of edge weights is correct
+        edge_weights = np.array(edge_weights)
+        if edge_weights.shape[0] != len(source_target_nodes):
+            raise ValueError(
+                "The number of edge weights must be equal to the number of source-target pairs."
+            )
+        # Add the weights for the extra edges, to be 0
+        n_extra_edges = Gc.ne - G.ne
+        edge_weights = np.concatenate(
+            [edge_weights, np.zeros((len(source_target_nodes), n_extra_edges))], axis=1
+        )
+    P = backend.Flow(Gc, lb=0, ub=DEFAULT_UB, n_flows=len(source_target_nodes))
+    # Now we add the objective and constraints for each sample
+    for i, (s, t) in enumerate(source_target_nodes):
+        weights = edge_weights[i, :]
+        P.add_objectives(P.expr.flow[:, i] @ weights)
+        # Now we inject/extract 1 unit flow from s to t
+        P += P.expr.flow[inflow_edges[s]] == 1
+        P += P.expr.flow[outflow_edges[t]] == 1
+        # For the rest of inflow/outflow edges, we set the flow to 0
+        for node in inflow_edges:
+            if node != s:
+                P += P.expr.flow[inflow_edges[node]] == 0
+        for node in outflow_edges:
+            if node != t:
+                P += P.expr.flow[outflow_edges[node]] == 0
+    # Add reg
+    if lam > 0:
+        P += backend.linear_or(
+            P.expr.flow, axis=1, ignore_type=True, varname="active_edge"
+        )
+        P.add_objectives(sum(P.expr.active_edge), weights=lam)
+    return P, Gc
+
+
 def shortest_path(
     G: BaseGraph,
     s: Any,
@@ -39,6 +99,10 @@ def shortest_path(
         edge_weights = np.array(
             [Gc.get_attr_edge(i).get("weight", 0) for i in range(Gc.ne)]
         )
+    else:
+        edge_weights = np.array(edge_weights)
+        # Add the weights for the extra edges, to be 0
+        edge_weights = np.concatenate([edge_weights, [0, 0]])
     if integral_path:
         P = backend.Flow(Gc, lb=0, ub=DEFAULT_UB)
         P += Indicator()