add necessary API and fix #5083

Author: lperron

ortools/constraint_solver/python/constraint_solver.cc

@@ -64,6 +64,8 @@ using ::operations_research::IntTupleSet;
 using ::operations_research::IntVar;
 using ::operations_research::IntVarElement;
 using ::operations_research::IntVarIterator;
+using ::operations_research::IntVarLocalSearchOperator;
+using ::operations_research::LocalSearchFilter;
 using ::operations_research::LocalSearchFilterManager;
 using ::operations_research::LocalSearchOperator;
 using ::operations_research::LocalSearchPhaseParameters;
@@ -299,6 +301,48 @@ class PyDecisionBuilderHelper : public PyDecisionBuilder {
   }
 };
 
+class PyLocalSearchOperator : public LocalSearchOperator {
+ public:
+  using LocalSearchOperator::LocalSearchOperator;
+  bool MakeNextNeighbor(Assignment* delta, Assignment* deltadelta) override {
+    PYBIND11_OVERRIDE_PURE_NAME(bool, LocalSearchOperator, "next_neighbor",
+                                MakeNextNeighbor, delta, deltadelta);
+  }
+  void Start(const Assignment* assignment) override {
+    PYBIND11_OVERRIDE_PURE_NAME(void, LocalSearchOperator, "start", Start,
+                                assignment);
+  }
+  void EnterSearch() override {
+    PYBIND11_OVERRIDE_NAME(void, LocalSearchOperator, "enter_search",
+                           EnterSearch, );
+  }
+  void Reset() override {
+    PYBIND11_OVERRIDE_NAME(void, LocalSearchOperator, "reset", Reset, );
+  }
+  bool HasFragments() const override {
+    PYBIND11_OVERRIDE_NAME(bool, LocalSearchOperator, "has_fragments",
+                           HasFragments, );
+  }
+  bool HoldsDelta() const override {
+    PYBIND11_OVERRIDE_NAME(bool, LocalSearchOperator, "holds_delta",
+                           HoldsDelta, );
+  }
+};
+
+class PyIntVarLocalSearchOperator : public IntVarLocalSearchOperator {
+ public:
+  using IntVarLocalSearchOperator::IntVarLocalSearchOperator;
+  using IntVarLocalSearchOperator::MakeOneNeighbor;
+  bool MakeOneNeighbor() override {
+    PYBIND11_OVERRIDE_NAME(bool, IntVarLocalSearchOperator, "one_neighbor",
+                           MakeOneNeighbor, );
+  }
+  void OnStart() override {
+    PYBIND11_OVERRIDE_NAME(void, IntVarLocalSearchOperator, "on_start",
+                           OnStart, );
+  }
+};
+
 class PySearchMonitor : public SearchMonitor {
  public:
   using SearchMonitor::SearchMonitor;
@@ -2540,6 +2584,41 @@ PYBIND11_MODULE(constraint_solver, m) {
           },
           py::return_value_policy::reference_internal);
 
+  py::class_<LocalSearchOperator, BaseObject, PyLocalSearchOperator>(
+      m, "LocalSearchOperator", DOC(operations_research, LocalSearchOperator))
+      .def(py::init<>())
+      .def("next_neighbor", &LocalSearchOperator::MakeNextNeighbor,
+           py::arg("delta"), py::arg("deltadelta"))
+      .def("start", &LocalSearchOperator::Start, py::arg("assignment"))
+      .def("enter_search", &LocalSearchOperator::EnterSearch)
+      .def("reset", &LocalSearchOperator::Reset)
+      .def("has_fragments", &LocalSearchOperator::HasFragments)
+      .def("holds_delta", &LocalSearchOperator::HoldsDelta);
+
+  py::class_<IntVarLocalSearchOperator, LocalSearchOperator,
+             PyIntVarLocalSearchOperator>(m, "IntVarLocalSearchOperator", "")
+      .def(py::init<const std::vector<IntVar*>&, bool>(), py::arg("vars"),
+           py::arg("keep_inverse_values") = false)
+      .def("size", &IntVarLocalSearchOperator::Size)
+      .def("var", &IntVarLocalSearchOperator::Var, py::arg("index"),
+           py::return_value_policy::reference_internal)
+      .def("is_incremental", &IntVarLocalSearchOperator::IsIncremental)
+      .def("value", &IntVarLocalSearchOperator::Value, py::arg("index"))
+      .def("old_value", &IntVarLocalSearchOperator::OldValue, py::arg("index"))
+      .def("prev_value", &IntVarLocalSearchOperator::PrevValue,
+           py::arg("index"))
+      .def("set_value", &IntVarLocalSearchOperator::SetValue, py::arg("index"),
+           py::arg("value"))
+      .def("activate", &IntVarLocalSearchOperator::Activate, py::arg("index"))
+      .def("deactivate", &IntVarLocalSearchOperator::Deactivate,
+           py::arg("index"))
+      .def("activated", &IntVarLocalSearchOperator::Activated, py::arg("index"))
+      .def("add_vars", &IntVarLocalSearchOperator::AddVars, py::arg("vars"))
+      .def(
+          "one_neighbor",
+          [](PyIntVarLocalSearchOperator* op) { return op->MakeOneNeighbor(); })
+      .def("on_start", &IntVarLocalSearchOperator::OnStart);
+
   py::class_<DecisionBuilder, BaseObject>(
       m, "DecisionBuilderBase", DOC(operations_research, DecisionBuilder))
       .def("__str__", &DecisionBuilder::DebugString)

ortools/constraint_solver/python/constraint_solver_test.py

@@ -226,6 +226,7 @@ def test_add_cumulative(self):
         ]
 
         solver.add_cumulative(intervals, demands, capacity, "cumul")
+
         # Since capacity is 1 and demands are 1, they cannot overlap.
 
         def check_single_solution():
@@ -623,6 +624,7 @@ def test_add_sub_circuit(self):
         solver = cp.Solver("test_sub_circuit")
         x = [solver.new_int_var(0, 2, f"x{i}") for i in range(3)]
         solver.add_sub_circuit(x)
+
         # SubCircuit allows partial circuits (if x[i] == i, it's not in the
         # circuit). But if in circuit, must form a single circuit.
 
@@ -1030,5 +1032,54 @@ def test_new_interval_relaxed_max(self):
         self.assertIsNotNone(r)
 
 
+class IntVarLocalSearchOperatorTest(absltest.TestCase):
+
+    def test_subclass_int_var_local_search_operator(self):
+        class MoveOneVar(cp.IntVarLocalSearchOperator):
+
+            def __init__(self, int_vars):
+                super().__init__(int_vars)
+                self.__index = 0
+                self.__up = False
+
+            def one_neighbor(self):
+                current_value = self.old_value(self.__index)
+                if self.__up:
+                    self.set_value(self.__index, current_value + 1)
+                    self.__index = (self.__index + 1) % self.size()
+                else:
+                    self.set_value(self.__index, current_value - 1)
+                self.__up = not self.__up
+                return True
+
+            def on_start(self):
+                pass
+
+        solver = cp.Solver("test_subclass")
+        x = solver.new_int_var(0, 10, "x")
+        y = solver.new_int_var(0, 10, "y")
+        move_one_var = MoveOneVar([x, y])
+        self.assertIsNotNone(move_one_var)
+
+
+class LocalSearchOperatorTest(absltest.TestCase):
+
+    def test_subclass_local_search_operator(self):
+        class CustomLSOperator(cp.LocalSearchOperator):
+
+            def __init__(self):
+                super().__init__()
+
+            def next_neighbor(self, delta, deltadelta):
+                return False
+
+            def start(self, assignment):
+                pass
+
+        solver = cp.Solver("test_subclass_ls")
+        op = CustomLSOperator()
+        self.assertIsNotNone(op)
+
+
 if __name__ == "__main__":
     absltest.main()

ortools/routing/python/routing.cc

@@ -897,10 +897,12 @@ PYBIND11_MODULE(routing, m) {
 
   rm.def("add_local_search_operator", &Model::AddLocalSearchOperator,
          py::arg("ls_operator"),
-         DOC(operations_research, routing, Model, AddLocalSearchOperator));
+         DOC(operations_research, routing, Model, AddLocalSearchOperator),
+         py::keep_alive<1, 2>());
   rm.def("add_local_search_filter", &Model::AddLocalSearchFilter,
          py::arg("filter"),
-         DOC(operations_research, routing, Model, AddLocalSearchFilter));
+         DOC(operations_research, routing, Model, AddLocalSearchFilter),
+         py::keep_alive<1, 2>());
   rm.def(
       "apply_locks",
       [](Model* model, const std::vector<int64_t>& locks) {

ortools/routing/python/routing_test.py

@@ -639,6 +639,53 @@ def test_initial_solution(self) -> None:
 
         self.assertIsNotNone(solution)
 
+    def test_issue5083(self):
+
+        class MyLocalSearch(constraint_solver.IntVarLocalSearchOperator):
+
+            def __init__(self, int_vars):
+                constraint_solver.IntVarLocalSearchOperator.__init__(self, int_vars)
+                self.__index = 0
+
+            def one_neighbor(self):
+                current_value = self.old_value(self.__index)
+                self.set_value(self.__index, 1 - current_value)
+                self.__index = (self.__index + 1) % self.Size()
+                return True
+
+            def on_start(self):
+                pass
+
+            def is_incremental(self):
+                return False
+
+        manager = routing.IndexManager(3, 1, 0)
+        model = routing.Model(manager)
+        model.add_to_assignment(model.active_var(manager.node_to_index(1)))
+        model.add_to_assignment(model.active_var(manager.node_to_index(2)))
+        model.add_disjunction([manager.node_to_index(1)], 1)
+        model.add_disjunction([manager.node_to_index(2)], 1)
+
+        model.add_local_search_operator(
+            MyLocalSearch(
+                [
+                    model.active_var(manager.node_to_index(1)),
+                    model.active_var(manager.node_to_index(2)),
+                ]
+            )
+        )
+
+        search_parameters = routing.default_routing_search_parameters()
+        search_parameters.first_solution_strategy = (
+            enums_pb2.FirstSolutionStrategy.AUTOMATIC
+        )
+        search_parameters.local_search_metaheuristic = (
+            enums_pb2.LocalSearchMetaheuristic.AUTOMATIC
+        )
+
+        solution = model.solve_with_parameters(search_parameters)
+        print(solution)
+
 
 if __name__ == "__main__":
     absltest.main()