clean up of masking implementation

Author: networmix

CMakeLists.txt

@@ -22,7 +22,8 @@ add_executable(netgraph_core_tests
   tests/cpp/flow_graph_tests.cpp
   tests/cpp/max_flow_tests.cpp
   tests/cpp/k_shortest_paths_tests.cpp
-    tests/cpp/negative_safety_tests.cpp
+  tests/cpp/negative_safety_tests.cpp
+  tests/cpp/masking_tests.cpp
 )
   target_link_libraries(netgraph_core_tests PRIVATE netgraph_core GTest::gtest_main)
   target_include_directories(netgraph_core_tests PRIVATE tests/cpp)

bindings/python/module.cpp

@@ -409,9 +409,23 @@ PYBIND11_MODULE(_netgraph_core, m) {
       .def("place_on_dag", [](FlowState& fs, std::int32_t src, std::int32_t dst, const PredDAG& dag, double requested_flow, FlowPlacement placement){
         py::gil_scoped_release rel; auto placed = fs.place_on_dag(src, dst, dag, requested_flow, placement); py::gil_scoped_acquire acq; return placed;
       }, py::arg("src"), py::arg("dst"), py::arg("dag"), py::arg("requested_flow") = std::numeric_limits<double>::infinity(), py::arg("flow_placement") = FlowPlacement::Proportional)
-      .def("place_max_flow", [](FlowState& fs, std::int32_t src, std::int32_t dst, FlowPlacement placement, bool shortest_path, bool require_capacity){
-        py::gil_scoped_release rel; auto total = fs.place_max_flow(src, dst, placement, shortest_path, require_capacity); py::gil_scoped_acquire acq; return total;
-      }, py::arg("src"), py::arg("dst"), py::arg("flow_placement") = FlowPlacement::Proportional, py::arg("shortest_path") = false, py::arg("require_capacity") = true)
+      .def("place_max_flow", [](py::object self_obj, std::int32_t src, std::int32_t dst,
+                                 FlowPlacement placement, bool shortest_path, bool require_capacity,
+                                 py::object node_mask, py::object edge_mask){
+        FlowState& fs = py::cast<FlowState&>(self_obj);
+        // Get graph reference to validate mask lengths
+        const StrictMultiDiGraph& g = py::cast<const StrictMultiDiGraph&>(self_obj.attr("_graph_ref"));
+        auto node_bs = to_bool_span_from_numpy(node_mask, static_cast<std::size_t>(g.num_nodes()), "node_mask");
+        auto edge_bs = to_bool_span_from_numpy(edge_mask, static_cast<std::size_t>(g.num_edges()), "edge_mask");
+        py::gil_scoped_release rel;
+        auto total = fs.place_max_flow(src, dst, placement, shortest_path, require_capacity,
+                                       node_bs.view, edge_bs.view);
+        py::gil_scoped_acquire acq;
+        return total;
+      }, py::arg("src"), py::arg("dst"),
+         py::arg("flow_placement") = FlowPlacement::Proportional,
+         py::arg("shortest_path") = false, py::arg("require_capacity") = true,
+         py::kw_only(), py::arg("node_mask") = py::none(), py::arg("edge_mask") = py::none())
       .def("compute_min_cut", [](py::object self_obj, std::int32_t src, py::object node_mask, py::object edge_mask){
         const FlowState& fs = py::cast<const FlowState&>(self_obj);
         // Get graph reference to validate mask lengths
@@ -546,63 +560,29 @@ PYBIND11_MODULE(_netgraph_core, m) {
       .def_readwrite("diminishing_returns_epsilon_frac", &FlowPolicyConfig::diminishing_returns_epsilon_frac);
 
   py::class_<FlowPolicy>(m, "FlowPolicy", py::dynamic_attr())
-      .def("__init__", [](py::object self, py::object algs_obj, py::object graph_obj, const FlowPolicyConfig& cfg){
+      .def("__init__", [](py::object self, py::object algs_obj, py::object graph_obj, FlowPolicyConfig cfg,
+                          py::object node_mask, py::object edge_mask){
             std::shared_ptr<Algorithms> algs = py::cast<std::shared_ptr<Algorithms>>(algs_obj);
             const PyGraph& pg = py::cast<const PyGraph&>(graph_obj);
+
+            // Convert masks to spans (FlowPolicy will copy the data)
+            auto node_bs = to_bool_span_from_numpy(node_mask, static_cast<std::size_t>(pg.num_nodes), "node_mask");
+            auto edge_bs = to_bool_span_from_numpy(edge_mask, static_cast<std::size_t>(pg.num_edges), "edge_mask");
+
+            // Update config with mask spans (will be copied by FlowPolicy constructor)
+            cfg.node_mask = node_bs.view;
+            cfg.edge_mask = edge_bs.view;
+
             ExecutionContext ctx(algs, pg.handle);
             FlowPolicy* fp = self.cast<FlowPolicy*>();
             new (fp) FlowPolicy(ctx, cfg);
             self.attr("_algorithms_ref") = algs_obj;
             self.attr("_graph_ref") = graph_obj;
            },
            py::arg("algorithms"), py::arg("graph"), py::arg("config"),
-           py::keep_alive<1, 2>()  // self keeps algorithms alive
-      )
-      .def("__init__", [](py::object self, py::object algs_obj, py::object graph_obj,
-                       PathAlg path_alg,
-                       FlowPlacement flow_placement,
-                       EdgeSelection selection,
-                       bool require_capacity,
-                       int min_flow_count,
-                       std::optional<int> max_flow_count,
-                       std::optional<Cost> max_path_cost,
-                       std::optional<double> max_path_cost_factor,
-                       bool shortest_path,
-                       bool reoptimize_flows_on_each_placement,
-                       int max_no_progress_iterations,
-                       int max_total_iterations,
-                       bool diminishing_returns_enabled,
-                       int diminishing_returns_window,
-                       double diminishing_returns_epsilon_frac){
-            std::shared_ptr<Algorithms> algs = py::cast<std::shared_ptr<Algorithms>>(algs_obj);
-            const PyGraph& pg = py::cast<const PyGraph&>(graph_obj);
-            ExecutionContext ctx(algs, pg.handle);
-            FlowPolicy* fp = self.cast<FlowPolicy*>();
-            new (fp) FlowPolicy(ctx, path_alg, flow_placement, selection, require_capacity,
-                               min_flow_count, max_flow_count, max_path_cost, max_path_cost_factor,
-                               shortest_path, reoptimize_flows_on_each_placement,
-                               max_no_progress_iterations, max_total_iterations,
-                               diminishing_returns_enabled, diminishing_returns_window,
-                               diminishing_returns_epsilon_frac);
-            self.attr("_algorithms_ref") = algs_obj;
-            self.attr("_graph_ref") = graph_obj;
-           },
-           py::arg("algorithms"), py::arg("graph"),
-           py::arg("path_alg") = PathAlg::SPF,
-           py::arg("flow_placement") = FlowPlacement::Proportional,
-           py::arg("selection") = EdgeSelection{},
-           py::arg("require_capacity") = true,
-           py::arg("min_flow_count") = 1,
-           py::arg("max_flow_count") = py::none(),
-           py::arg("max_path_cost") = py::none(),
-           py::arg("max_path_cost_factor") = py::none(),
-           py::arg("shortest_path") = false,
-           py::arg("reoptimize_flows_on_each_placement") = false,
-           py::arg("max_no_progress_iterations") = 100,
-           py::arg("max_total_iterations") = 10000,
-           py::arg("diminishing_returns_enabled") = true,
-           py::arg("diminishing_returns_window") = 8,
-           py::arg("diminishing_returns_epsilon_frac") = 1e-3,
+           py::kw_only(),
+           py::arg("node_mask") = py::none(),
+           py::arg("edge_mask") = py::none(),
            py::keep_alive<1, 2>()  // self keeps algorithms alive
       )
       .def("flow_count", &FlowPolicy::flow_count)

include/netgraph/core/flow_policy.hpp

@@ -47,6 +47,8 @@ struct FlowPolicyConfig {
   bool diminishing_returns_enabled { true };
   int diminishing_returns_window { 8 };
   double diminishing_returns_epsilon_frac { 1e-3 };
+  std::span<const bool> node_mask {};  // Optional node mask for failure exclusions (True = include)
+  std::span<const bool> edge_mask {};  // Optional edge mask for failure exclusions (True = include)
 };
 
 // FlowPolicy manages flow creation, placement, reoptimization for a single demand
@@ -61,7 +63,31 @@ class FlowPolicy {
       max_path_cost_factor_(cfg.max_path_cost_factor), reoptimize_flows_on_each_placement_(cfg.reoptimize_flows_on_each_placement),
       max_no_progress_iterations_(cfg.max_no_progress_iterations), max_total_iterations_(cfg.max_total_iterations),
       diminishing_returns_enabled_(cfg.diminishing_returns_enabled), diminishing_returns_window_(cfg.diminishing_returns_window),
-      diminishing_returns_epsilon_frac_(cfg.diminishing_returns_epsilon_frac) {}
+      diminishing_returns_epsilon_frac_(cfg.diminishing_returns_epsilon_frac)
+  {
+    // Validate mask sizes early for clearer error messages
+    const auto* g = ctx_.graph.graph.get();
+    if (g) {
+      if (!cfg.node_mask.empty() && cfg.node_mask.size() != static_cast<std::size_t>(g->num_nodes())) {
+        throw std::invalid_argument("FlowPolicy: node_mask length mismatch");
+      }
+      if (!cfg.edge_mask.empty() && cfg.edge_mask.size() != static_cast<std::size_t>(g->num_edges())) {
+        throw std::invalid_argument("FlowPolicy: edge_mask length mismatch");
+      }
+    }
+
+    // Copy mask data if provided
+    if (!cfg.node_mask.empty()) {
+      node_mask_storage_.reset(new bool[cfg.node_mask.size()]);
+      std::copy(cfg.node_mask.begin(), cfg.node_mask.end(), node_mask_storage_.get());
+      node_mask_ = std::span<const bool>(node_mask_storage_.get(), cfg.node_mask.size());
+    }
+    if (!cfg.edge_mask.empty()) {
+      edge_mask_storage_.reset(new bool[cfg.edge_mask.size()]);
+      std::copy(cfg.edge_mask.begin(), cfg.edge_mask.end(), edge_mask_storage_.get());
+      edge_mask_ = std::span<const bool>(edge_mask_storage_.get(), cfg.edge_mask.size());
+    }
+  }
 
   FlowPolicy(const ExecutionContext& ctx,
              PathAlg path_alg,
@@ -140,6 +166,12 @@ class FlowPolicy {
   int diminishing_returns_window_ { 8 };
   double diminishing_returns_epsilon_frac_ { 1e-3 };
 
+  // Mask storage and views for failure exclusions
+  std::unique_ptr<bool[]> node_mask_storage_;
+  std::unique_ptr<bool[]> edge_mask_storage_;
+  std::span<const bool> node_mask_ {};
+  std::span<const bool> edge_mask_ {};
+
   // State
   std::unordered_map<FlowIndex, FlowRecord, FlowIndexHash> flows_;
   Cost best_path_cost_ { std::numeric_limits<Cost>::max() };

include/netgraph/core/flow_state.hpp

@@ -72,7 +72,9 @@ class FlowState {
   [[nodiscard]] Flow place_max_flow(NodeId src, NodeId dst,
                       FlowPlacement placement,
                       bool shortest_path = false,
-                      bool require_capacity = true);
+                      bool require_capacity = true,
+                      std::span<const bool> node_mask = {},
+                      std::span<const bool> edge_mask = {});
 
   // Compute min-cut with respect to current residual state, starting reachability
   // from source s on the residual graph (forward arcs: residual>MIN; reverse arcs:

include/netgraph/core/shortest_paths.hpp

@@ -45,7 +45,9 @@ struct PredDAG {
 //   For EB semantics this means re-running SPF on updated residuals
 //   will remove saturated next-hops and therefore *change* the fixed
 //   equal-split set (progressive behavior).
-// - node_mask: if provided, node_mask[v]==true means node v is allowed (false excludes it)
+// - node_mask: if provided, node_mask[v]==true means node v is allowed (false excludes it).
+//   If the source node is masked (node_mask[src]==false), returns an empty predecessor DAG
+//   with all distances at infinity, as no traversal can begin from an excluded source.
 // - edge_mask: if provided, edge_mask[e]==true means edge e is allowed (false excludes it)
 [[nodiscard]] std::pair<std::vector<Cost>, PredDAG>
 shortest_paths(const StrictMultiDiGraph& g, NodeId src,

python/netgraph_core/_docs.py

@@ -205,13 +205,18 @@ def place_max_flow(
         flow_placement: FlowPlacement = FlowPlacement.PROPORTIONAL,
         shortest_path: bool = False,
         require_capacity: bool = True,
+        *,
+        node_mask: "Optional[np.ndarray]" = None,
+        edge_mask: "Optional[np.ndarray]" = None,
     ) -> float:
         """Place maximum flow from src to dst.
 
         Args:
             require_capacity: Whether to require edges to have capacity.
                 - True (default): Routes adapt to residuals (SDN/TE, progressive fill).
                 - False: Routes based on costs only (IP/IGP, fixed routing).
+            node_mask: Optional 1-D bool array (True=allowed). Copied for thread safety.
+            edge_mask: Optional 1-D bool array (True=allowed). Copied for thread safety.
 
         For *IP-style ECMP* max-flow, use: require_capacity=False, shortest_path=True,
         flow_placement=EQUAL_BALANCED.
@@ -285,29 +290,26 @@ class FlowPolicy:
     When static_paths is empty the policy may refresh the DAG per round using
     residual-aware shortest paths. This progressively prunes saturated next-hops
     (traffic-engineering style) and differs from one-shot ECMP admission.
+
+    Args:
+        algorithms: Algorithms instance (kept alive by FlowPolicy)
+        graph: Graph handle (kept alive by FlowPolicy)
+        config: FlowPolicyConfig with all policy parameters
+        node_mask: Optional 1-D bool array (True=allowed). **Copied for thread safety.**
+        edge_mask: Optional 1-D bool array (True=allowed). **Copied for thread safety.**
+
+    Raises:
+        TypeError: If masks have wrong dtype, ndim, or length.
     """
 
     def __init__(
         self,
         algorithms: "Algorithms",
         graph: "Graph",
-        config: Optional["FlowPolicyConfig"] = None,
-        /,
+        config: "FlowPolicyConfig",
         *,
-        path_alg: PathAlg = PathAlg.SPF,
-        flow_placement: FlowPlacement = FlowPlacement.PROPORTIONAL,
-        selection: EdgeSelection = ...,  # default provided by runtime binding
-        min_flow_count: int = 1,
-        max_flow_count: Optional[int] = None,
-        max_path_cost: Optional[int] = None,
-        max_path_cost_factor: Optional[float] = None,
-        shortest_path: bool = False,
-        reoptimize_flows_on_each_placement: bool = False,
-        max_no_progress_iterations: int = 100,
-        max_total_iterations: int = 10000,
-        diminishing_returns_enabled: bool = True,
-        diminishing_returns_window: int = 8,
-        diminishing_returns_epsilon_frac: float = 1e-3,
+        node_mask: "Optional[np.ndarray]" = None,
+        edge_mask: "Optional[np.ndarray]" = None,
     ) -> None: ...
 
     def flow_count(self) -> int: ...
@@ -419,13 +421,21 @@ def spf(
             selection: Edge selection policy
             residual: Optional 1-D float64 array of residuals. **Copied for thread safety.**
             node_mask: Optional 1-D bool mask (length num_nodes). **Copied for thread safety.**
+                       True = node allowed, False = node excluded.
+                       **If source node is masked (False), returns empty DAG with all distances at infinity.**
             edge_mask: Optional 1-D bool mask (length num_edges). **Copied for thread safety.**
+                       True = edge allowed, False = edge excluded.
             multipath: Whether to track multiple equal-cost paths
             dtype: "float64" (inf for unreachable) or "int64" (max for unreachable)
 
         Returns:
             (distances, predecessor_dag)
 
+        Note:
+            When the source node is masked out (node_mask[src] == False), the algorithm
+            immediately returns an empty predecessor DAG with all distances set to infinity,
+            as no traversal can begin from an excluded source.
+
         Raises:
             TypeError: If arrays have wrong dtype, ndim, or length.
             ValueError: If src/dst out of range.
@@ -518,6 +528,7 @@ def batch_max_flow(
         edge_masks: Optional[list["np.ndarray"]] = None,
         flow_placement: FlowPlacement = FlowPlacement.PROPORTIONAL,
         shortest_path: bool = False,
+        require_capacity: bool = True,
         with_edge_flows: bool = False,
         with_reachable: bool = False,
         with_residuals: bool = False,
@@ -528,6 +539,7 @@ def batch_max_flow(
             pairs: int32 array of shape [B, 2] with (src, dst) pairs
             node_masks: Optional list of B bool masks. **Each copied for thread safety.**
             edge_masks: Optional list of B bool masks. **Each copied for thread safety.**
+            require_capacity: If True, exclude saturated edges (SDN/TE). If False, route by cost only (IP/IGP).
 
         Returns:
             List of FlowSummary objects, one per pair.