WIP routing: Improve pybind11 binding

Author: Mizux

cmake/python.cmake

@@ -540,7 +540,7 @@ add_custom_command(
   COMMAND ${stubgen_EXECUTABLE} -p ortools.math_opt.core.python.solver --output .
   COMMAND ${stubgen_EXECUTABLE} -p ortools.pdlp.python.pdlp --output .
   COMMAND ${stubgen_EXECUTABLE} -p ortools.routing.pywraprouting --output .
-  COMMAND ${stubgen_EXECUTABLE} -p ortools.routing.python.routing --output .
+  COMMAND ${stubgen_EXECUTABLE} -p ortools.routing.python.model --output .
   COMMAND ${stubgen_EXECUTABLE} -p ortools.sat.python.swig_helper --output .
   COMMAND ${stubgen_EXECUTABLE} -p ortools.scheduling.python.rcpsp --output .
   COMMAND ${stubgen_EXECUTABLE} -p ortools.util.python.sorted_interval_list --output .

ortools/routing/python/model.cc

@@ -11,8 +11,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <algorithm>
 #include <cstdint>
 #include <functional>
+#include <iterator>
 #include <utility>
 
 #include "ortools/constraint_solver/constraint_solver.h"
@@ -30,8 +32,10 @@
 #include "pybind11/stl.h"
 #include "pybind11_protobuf/native_proto_caster.h"
 
+using ::operations_research::DefaultRoutingModelParameters;
 using ::operations_research::DefaultRoutingSearchParameters;
 using ::operations_research::RoutingIndexManager;
+using ::operations_research::RoutingDimension;
 using ::operations_research::RoutingModel;
 using ::pybind11::arg;
 
@@ -41,7 +45,12 @@ PYBIND11_MODULE(model, m) {
   pybind11::module::import(
       "ortools.constraint_solver.python.constraint_solver");
 
+  m.def("default_routing_model_parameters", &DefaultRoutingModelParameters,
+      pybind11::return_value_policy::reference_internal,
+        DOC(operations_research, DefaultRoutingModelParameters));
+
   m.def("default_routing_search_parameters", &DefaultRoutingSearchParameters,
+      pybind11::return_value_policy::reference_internal,
         DOC(operations_research, DefaultRoutingSearchParameters));
 
   pybind11::class_<RoutingIndexManager>(
@@ -52,6 +61,23 @@ PYBIND11_MODULE(model, m) {
                  RoutingIndexManager::NodeIndex(depot));
            }),
            DOC(operations_research, RoutingIndexManager, RoutingIndexManager))
+      .def(pybind11::init([](int num_nodes, int num_vehicles,
+              const std::vector<int> starts,
+              const std::vector<int> ends) {
+            std::vector<RoutingIndexManager::NodeIndex> start_nodes;
+            start_nodes.reserve(starts.size());
+             std::transform(starts.cbegin(), starts.cend(), std::back_inserter(start_nodes),
+                 [](int node){return RoutingIndexManager::NodeIndex(node);});
+
+             std::vector<RoutingIndexManager::NodeIndex> end_nodes;
+             end_nodes.reserve(ends.size());
+             std::transform(ends.cbegin(), ends.cend(), std::back_inserter(end_nodes),
+                 [](int node){return RoutingIndexManager::NodeIndex(node);});
+
+             return new RoutingIndexManager(
+                 num_nodes, num_vehicles, start_nodes, end_nodes);
+           }),
+           DOC(operations_research, RoutingIndexManager, RoutingIndexManager))
       .def("num_nodes", &RoutingIndexManager::num_nodes,
            DOC(operations_research, RoutingIndexManager, num_nodes))
       .def("num_vehicles", &RoutingIndexManager::num_vehicles,
@@ -76,10 +102,34 @@ PYBIND11_MODULE(model, m) {
       .def("get_end_index", &RoutingIndexManager::GetEndIndex,
            DOC(operations_research, RoutingIndexManager, GetEndIndex));
 
+  pybind11::class_<RoutingDimension>(m, "RoutingDimension")
+    .def("model", &RoutingDimension::model,
+        pybind11::return_value_policy::reference_internal)
+    .def("get_transit_value", &RoutingDimension::GetTransitValue,
+        arg("from_index"), arg("to_index"), arg("vehicle"));
+
   pybind11::class_<RoutingModel>(m, "RoutingModel")
       .def(pybind11::init([](const RoutingIndexManager& routing_index_manager) {
         return new RoutingModel(routing_index_manager);
       }))
+      .def("register_transit_matrix",
+           [](RoutingModel* routing_model,
+              std::vector<std::vector<int64_t>> transit_matrix) {
+             return routing_model->RegisterTransitMatrix(
+                 std::move(transit_matrix));
+           })
+      .def("register_unary_transit_vector",
+           [](RoutingModel* routing_model,
+              std::vector<int64_t> transit_vector) {
+             return routing_model->RegisterUnaryTransitVector(
+                 std::move(transit_vector));
+           })
+      .def("register_unary_transit_callback",
+           [](RoutingModel* routing_model,
+              std::function<int64_t(int64_t)> transit_callback) {
+             return routing_model->RegisterUnaryTransitCallback(
+                 std::move(transit_callback));
+           })
       .def("register_transit_callback",
            [](RoutingModel* routing_model,
               std::function<int64_t(int64_t, int64_t)> transit_callback) {
@@ -89,6 +139,48 @@ PYBIND11_MODULE(model, m) {
       .def("set_arc_cost_evaluator_of_all_vehicles",
            &RoutingModel::SetArcCostEvaluatorOfAllVehicles,
            arg("transit_callback_index"))
+      .def("add_dimension", &RoutingModel::AddDimension,
+           arg("evaluator_index"),
+           arg("slack_max"),
+           arg("capacity"),
+           arg("fix_start_cumul_to_zero"),
+           arg("name"))
+      .def("add_dimension_with_vehicle_capacity", &RoutingModel::AddDimensionWithVehicleCapacity,
+           arg("evaluator_index"),
+           arg("slack_max"),
+           arg("vehicle_capacities"),
+           arg("fix_start_cumul_to_zero"),
+           arg("name"))
+      .def("add_dimension_with_vehicle_transits", &RoutingModel::AddDimensionWithVehicleTransits,
+           arg("evaluator_indices"),
+           arg("slack_max"),
+           arg("capacity"),
+           arg("fix_start_cumul_to_zero"),
+           arg("name"))
+      .def("add_dimension_with_vehicle_transit_and_capacity", &RoutingModel::AddDimensionWithVehicleTransitAndCapacity,
+           arg("evaluator_indices"),
+           arg("slack_max"),
+           arg("vehicle_capacities"),
+           arg("fix_start_cumul_to_zero"),
+           arg("name"))
+      .def("add_constant_dimension", &RoutingModel::AddConstantDimension,
+           arg("value"),
+           arg("capacity"),
+           arg("fix_start_cumul_to_zero"),
+           arg("name"))
+      .def("add_vector_dimension", &RoutingModel::AddVectorDimension,
+           arg("values"),
+           arg("capacity"),
+           arg("fix_start_cumul_to_zero"),
+           arg("name"))
+      .def("add_matrix_dimension", &RoutingModel::AddMatrixDimension,
+           arg("values"),
+           arg("capacity"),
+           arg("fix_start_cumul_to_zero"),
+           arg("name"))
+      .def("get_dimension_or_die", &RoutingModel::GetDimensionOrDie,
+           pybind11::return_value_policy::reference_internal,
+           arg("dimension_name"))
       .def("solve", &RoutingModel::Solve,
            pybind11::return_value_policy::reference_internal,
            arg("assignment") = nullptr)