Expose epsilon parameter (precision) through python layer

cpp/include/cugraph/algorithms.hpp

@@ -595,7 +595,7 @@ void bfs(raft::handle_t const &handle,
  * @param[in]  graph                 cuGRAPH COO graph
  * @param[in]  num_workers           number of vertices in the worker set
  * @param[in]  workers               device pointer to an array of worker vertex ids
- * @param[out] assignment            device pointer to an array to which the assignment will be
+ * @param[out] assignments           device pointer to an array to which the assignment will be
  * written. The array should be num_workers long, and will identify which vertex id (job) is
  * assigned to that worker
  */
@@ -604,7 +604,7 @@ weight_t hungarian(raft::handle_t const &handle,
                    GraphCOOView<vertex_t, edge_t, weight_t> const &graph,
                    vertex_t num_workers,
                    vertex_t const *workers,
-                   vertex_t *assignment);
+                   vertex_t *assignments);
 
 /**
  * @brief      Compute Hungarian algorithm on a weighted bipartite graph
@@ -626,19 +626,19 @@ weight_t hungarian(raft::handle_t const &handle,
  * @param[in]  graph                 cuGRAPH COO graph
  * @param[in]  num_workers           number of vertices in the worker set
  * @param[in]  workers               device pointer to an array of worker vertex ids
- * @param[out] assignment            device pointer to an array to which the assignment will be
+ * @param[out] assignments           device pointer to an array to which the assignment will be
  * written. The array should be num_workers long, and will identify which vertex id (job) is
  * assigned to that worker
- * @param[in]  precision             parameter to define precision of comparisons
+ * @param[in]  epsilon               parameter to define precision of comparisons
  *                                   in reducing weights to zero.
  */
 template <typename vertex_t, typename edge_t, typename weight_t>
 weight_t hungarian(raft::handle_t const &handle,
                    GraphCOOView<vertex_t, edge_t, weight_t> const &graph,
                    vertex_t num_workers,
                    vertex_t const *workers,
-                   vertex_t *assignment,
-                   weight_t precision);
+                   vertex_t *assignments,
+                   weight_t epsilon);
 
 /**
  * @brief      Louvain implementation
@@ -1075,7 +1075,7 @@ namespace dense {
  * @param[in]  costs                 pointer to array of costs, stored in row major order
  * @param[in]  num_rows              number of rows in dense matrix
  * @param[in]  num_cols              number of cols in dense matrix
- * @param[out] assignment            device pointer to an array to which the assignment will be
+ * @param[out] assignments           device pointer to an array to which the assignment will be
  *                                   written. The array should be num_cols long, and will identify
  *                                   which vertex id (job) is assigned to that worker
  */
@@ -1084,7 +1084,7 @@ weight_t hungarian(raft::handle_t const &handle,
                    weight_t const *costs,
                    vertex_t num_rows,
                    vertex_t num_columns,
-                   vertex_t *assignment);
+                   vertex_t *assignments);
 
 /**
  * @brief      Compute Hungarian algorithm on a weighted bipartite graph
@@ -1104,19 +1104,19 @@ weight_t hungarian(raft::handle_t const &handle,
  * @param[in]  costs                 pointer to array of costs, stored in row major order
  * @param[in]  num_rows              number of rows in dense matrix
  * @param[in]  num_cols              number of cols in dense matrix
- * @param[out] assignment            device pointer to an array to which the assignment will be
+ * @param[out] assignments           device pointer to an array to which the assignment will be
  *                                   written. The array should be num_cols long, and will identify
  *                                   which vertex id (job) is assigned to that worker
- * @param[in]  precision             parameter to define precision of comparisons
+ * @param[in]  epsilon               parameter to define precision of comparisons
  *                                   in reducing weights to zero.
  */
 template <typename vertex_t, typename weight_t>
 weight_t hungarian(raft::handle_t const &handle,
                    weight_t const *costs,
                    vertex_t num_rows,
                    vertex_t num_columns,
-                   vertex_t *assignment,
-                   weight_t precision);
+                   vertex_t *assignments,
+                   weight_t epsilon);
 
 }  // namespace dense
 

cpp/src/linear_assignment/hungarian.cu

@@ -38,19 +38,19 @@ namespace cugraph {
 namespace detail {
 
 template <typename weight_t>
-weight_t default_precision()
+weight_t default_epsilon()
 {
   return 0;
 }
 
 template <>
-float default_precision()
+float default_epsilon()
 {
   return float{1e-6};
 }
 
 template <>
-double default_precision()
+double default_epsilon()
 {
   return double{1e-6};
 }
@@ -61,13 +61,13 @@ weight_t hungarian(raft::handle_t const &handle,
                    index_t num_cols,
                    weight_t const *d_original_cost,
                    index_t *d_assignment,
-                   weight_t precision)
+                   weight_t epsilon)
 {
   if (num_rows == num_cols) {
     rmm::device_uvector<index_t> col_assignments_v(num_rows, handle.get_stream_view());
 
     // Create an instance of LinearAssignmentProblem using problem size, number of subproblems
-    raft::lap::LinearAssignmentProblem<index_t, weight_t> lpx(handle, num_rows, 1, precision);
+    raft::lap::LinearAssignmentProblem<index_t, weight_t> lpx(handle, num_rows, 1, epsilon);
 
     // Solve LAP(s) for given cost matrix
     lpx.solve(d_original_cost, d_assignment, col_assignments_v.data());
@@ -85,14 +85,14 @@ weight_t hungarian(raft::handle_t const &handle,
                                        weight_t{0},
                                        thrust::maximum<weight_t>());
 
-    rmm::device_uvector<weight_t> tmp_cost(n * n, handle.get_stream_view());
+    rmm::device_uvector<weight_t> tmp_cost_v(n * n, handle.get_stream_view());
     rmm::device_uvector<index_t> tmp_row_assignment_v(n, handle.get_stream_view());
     rmm::device_uvector<index_t> tmp_col_assignment_v(n, handle.get_stream_view());
 
     thrust::transform(rmm::exec_policy(handle.get_stream_view()),
                       thrust::make_counting_iterator<index_t>(0),
                       thrust::make_counting_iterator<index_t>(n * n),
-                      tmp_cost.begin(),
+                      tmp_cost_v.begin(),
                       [max_cost, d_original_cost, n, num_rows, num_cols] __device__(index_t i) {
                         index_t row = i / n;
                         index_t col = i % n;
@@ -102,10 +102,10 @@ weight_t hungarian(raft::handle_t const &handle,
                                  : max_cost;
                       });
 
-    raft::lap::LinearAssignmentProblem<index_t, weight_t> lpx(handle, n, 1, precision);
+    raft::lap::LinearAssignmentProblem<index_t, weight_t> lpx(handle, n, 1, epsilon);
 
     // Solve LAP(s) for given cost matrix
-    lpx.solve(tmp_cost.begin(), tmp_row_assignment_v.begin(), tmp_col_assignment_v.begin());
+    lpx.solve(tmp_cost_v.begin(), tmp_row_assignment_v.begin(), tmp_col_assignment_v.begin());
 
     weight_t tmp_objective_value = lpx.getPrimalObjectiveValue(0);
 
@@ -121,7 +121,7 @@ weight_t hungarian_sparse(raft::handle_t const &handle,
                           vertex_t num_workers,
                           vertex_t const *workers,
                           vertex_t *assignment,
-                          weight_t precision)
+                          weight_t epsilon)
 {
   CUGRAPH_EXPECTS(assignment != nullptr, "Invalid input argument: assignment pointer is NULL");
   CUGRAPH_EXPECTS(graph.edge_data != nullptr,
@@ -235,7 +235,7 @@ weight_t hungarian_sparse(raft::handle_t const &handle,
   vertex_t *d_temp_assignment = temp_assignment_v.data();
 
   weight_t min_cost = detail::hungarian(
-    handle, matrix_dimension, matrix_dimension, d_cost, d_temp_assignment, precision);
+    handle, matrix_dimension, matrix_dimension, d_cost, d_temp_assignment, epsilon);
 
 #ifdef TIMING
   hr_timer.stop();
@@ -272,7 +272,7 @@ weight_t hungarian(raft::handle_t const &handle,
                    vertex_t *assignment)
 {
   return detail::hungarian_sparse(
-    handle, graph, num_workers, workers, assignment, detail::default_precision<weight_t>());
+    handle, graph, num_workers, workers, assignment, detail::default_epsilon<weight_t>());
 }
 
 template <typename vertex_t, typename edge_t, typename weight_t>
@@ -281,9 +281,9 @@ weight_t hungarian(raft::handle_t const &handle,
                    vertex_t num_workers,
                    vertex_t const *workers,
                    vertex_t *assignment,
-                   weight_t precision)
+                   weight_t epsilon)
 {
-  return detail::hungarian_sparse(handle, graph, num_workers, workers, assignment, precision);
+  return detail::hungarian_sparse(handle, graph, num_workers, workers, assignment, epsilon);
 }
 
 template int32_t hungarian<int32_t, int32_t, int32_t>(
@@ -335,7 +335,7 @@ weight_t hungarian(raft::handle_t const &handle,
                    index_t *assignment)
 {
   return detail::hungarian(
-    handle, num_rows, num_cols, costs, assignment, detail::default_precision<weight_t>());
+    handle, num_rows, num_cols, costs, assignment, detail::default_epsilon<weight_t>());
 }
 
 template <typename index_t, typename weight_t>
@@ -344,9 +344,9 @@ weight_t hungarian(raft::handle_t const &handle,
                    index_t num_rows,
                    index_t num_cols,
                    index_t *assignment,
-                   weight_t precision)
+                   weight_t epsilon)
 {
-  return detail::hungarian(handle, num_rows, num_cols, costs, assignment, precision);
+  return detail::hungarian(handle, num_rows, num_cols, costs, assignment, epsilon);
 }
 
 template int32_t hungarian<int32_t, int32_t>(

python/cugraph/linear_assignment/lap.pxd

@@ -25,8 +25,15 @@ cdef extern from "cugraph/algorithms.hpp" namespace "cugraph":
         const GraphCOOView[vertex_t,edge_t,weight_t] &graph,
         vertex_t num_workers,
         const vertex_t *workers,
-        vertex_t *assignment) except +
+        vertex_t *assignments,
+        weight_t epsilon) except +
 
+    cdef weight_t hungarian[vertex_t,edge_t,weight_t](
+        const handle_t &handle,
+        const GraphCOOView[vertex_t,edge_t,weight_t] &graph,
+        vertex_t num_workers,
+        const vertex_t *workers,
+        vertex_t *assignments) except +
 
 cdef extern from "cugraph/algorithms.hpp":
 
@@ -35,4 +42,12 @@ cdef extern from "cugraph/algorithms.hpp":
         const weight_t *costs,
         vertex_t num_rows,
         vertex_t num_columns,
-        vertex_t *assignment) except +
+        vertex_t *assignments,
+        weight_t epsilon) except +
+
+    cdef weight_t dense_hungarian "cugraph::dense::hungarian" [vertex_t,weight_t](
+        const handle_t &handle,
+        const weight_t *costs,
+        vertex_t num_rows,
+        vertex_t num_columns,
+        vertex_t *assignments) except +

python/cugraph/linear_assignment/lap.py

@@ -15,7 +15,7 @@
 from cugraph.linear_assignment import lap_wrapper
 
 
-def hungarian(G, workers):
+def hungarian(G, workers, epsilon=None):
     """
     Execute the Hungarian algorithm against a symmetric, weighted,
     bipartite graph.
@@ -46,6 +46,11 @@ def hungarian(G, workers):
         cudf.DataFrame. All vertices in G that are not in the workers
         set are implicitly assigned to the jobs set.
 
+    epsilon : float or double (matching weight type in graph)
+        Used for determining when value is close enough to zero to consider 0.
+        Defaults (if not specified) to 1e-6 in the C++ code.  Unused for
+        integer weight types.
+
     Returns
     -------
     cost : matches costs.dtype
@@ -77,15 +82,15 @@ def hungarian(G, workers):
     else:
         local_workers = workers
 
-    cost, df = lap_wrapper.sparse_hungarian(G, local_workers)
+    cost, df = lap_wrapper.sparse_hungarian(G, local_workers, epsilon)
 
     if G.renumbered:
         df = G.unrenumber(df, 'vertex')
 
     return cost, df
 
 
-def dense_hungarian(costs, num_rows, num_columns):
+def dense_hungarian(costs, num_rows, num_columns, epsilon=None):
     """
     Execute the Hungarian algorithm against a dense bipartite
     graph representation.
@@ -107,7 +112,10 @@ def dense_hungarian(costs, num_rows, num_columns):
         Number of rows in the matrix
     num_columns : int
         Number of columns in the matrix
-
+    epsilon : float or double (matching weight type in graph)
+        Used for determining when value is close enough to zero to consider 0.
+        Defaults (if not specified) to 1e-6 in the C++ code.  Unused for
+        integer weight types.
 
     Returns
     -------
@@ -121,4 +129,4 @@ def dense_hungarian(costs, num_rows, num_columns):
 
     """
 
-    return lap_wrapper.dense_hungarian(costs, num_rows, num_columns)
+    return lap_wrapper.dense_hungarian(costs, num_rows, num_columns, epsilon)

python/cugraph/linear_assignment/lap_wrapper.pyx

@@ -25,7 +25,7 @@ import cudf
 import numpy as np
 
 
-def sparse_hungarian(input_graph, workers):
+def sparse_hungarian(input_graph, workers, epsilon):
     """
     Call the hungarian algorithm
     """
@@ -62,30 +62,35 @@ def sparse_hungarian(input_graph, workers):
     df['vertex'] = workers
     df['assignment'] = cudf.Series(np.zeros(len(workers), dtype=np.int32))
 
+    if epsilon == None:
+        epsilon = 1e-6
+
     cdef uintptr_t c_src        = src.__cuda_array_interface__['data'][0]
     cdef uintptr_t c_dst        = dst.__cuda_array_interface__['data'][0]
     cdef uintptr_t c_weights    = weights.__cuda_array_interface__['data'][0]
     cdef uintptr_t c_workers    = local_workers.__cuda_array_interface__['data'][0]
 
     cdef uintptr_t c_identifier = df['vertex'].__cuda_array_interface__['data'][0];
     cdef uintptr_t c_assignment = df['assignment'].__cuda_array_interface__['data'][0];
+    cdef float c_epsilon_float = epsilon
+    cdef double c_epsilon_double = epsilon
 
     cdef GraphCOOView[int,int,float] g_float
     cdef GraphCOOView[int,int,double] g_double
 
     if weights.dtype == np.float32:
         g_float = GraphCOOView[int,int,float](<int*>c_src, <int*>c_dst, <float*>c_weights, num_verts, num_edges)
 
-        cost = c_hungarian[int,int,float](handle_[0], g_float, len(workers), <int*>c_workers, <int*>c_assignment)
+        cost = c_hungarian[int,int,float](handle_[0], g_float, len(workers), <int*>c_workers, <int*>c_assignment, c_epsilon_float)
     else:
         g_double = GraphCOOView[int,int,double](<int*>c_src, <int*>c_dst, <double*>c_weights, num_verts, num_edges)
 
-        cost = c_hungarian[int,int,double](handle_[0], g_double, len(workers), <int*>c_workers, <int*>c_assignment)
+        cost = c_hungarian[int,int,double](handle_[0], g_double, len(workers), <int*>c_workers, <int*>c_assignment, c_epsilon_double)
 
     return cost, df
 
 
-def dense_hungarian(costs, num_rows, num_columns):
+def dense_hungarian(costs, num_rows, num_columns, epsilon):
     """
     Call the dense hungarian algorithm
     """
@@ -98,13 +103,19 @@ def dense_hungarian(costs, num_rows, num_columns):
 
     assignment = cudf.Series(np.zeros(num_rows, dtype=np.int32))
 
+    if epsilon == None:
+        epsilon = 1e-6
+
     cdef uintptr_t c_costs = costs.__cuda_array_interface__['data'][0]
     cdef uintptr_t c_assignment = assignment.__cuda_array_interface__['data'][0]
-
+    cdef float c_epsilon_float = epsilon
+    cdef double c_epsilon_double = epsilon
 
     if costs.dtype == np.float32:
-        cost = c_dense_hungarian[int,float](handle_[0], <float*> c_costs, num_rows, num_columns, <int*> c_assignment)
+        cost = c_dense_hungarian[int,float](handle_[0], <float*> c_costs, num_rows, num_columns, <int*> c_assignment, c_epsilon_float)
     elif costs.dtype == np.float64:
+        cost = c_dense_hungarian[int,double](handle_[0], <double*> c_costs, num_rows, num_columns, <int*> c_assignment, c_epsilon_double)
+    elif costs.dtype == np.int32:
         cost = c_dense_hungarian[int,double](handle_[0], <double*> c_costs, num_rows, num_columns, <int*> c_assignment)
     else:
         raise("unsported type: ", costs.dtype)