add_next_values method on SparseMatrix, much faster _build_sparse_matrix on numpy arrays

lib/stormpy/storage/__init__.py

@@ -42,22 +42,13 @@ def _build_sparse_matrix(builder_class, array, row_group_indices=[]):
     num_col = array.shape[1]
 
     len_group_indices = len(row_group_indices)
+    nz_rows, nz_cols = array.nonzero()
+    nz_vals = array[nz_rows, nz_cols]
     if len_group_indices > 0:
-        builder = builder_class(rows=num_row, columns=num_col, has_custom_row_grouping=True, row_groups=len_group_indices)
+        builder = builder_class(rows=num_row, columns=num_col, entries=len(nz_rows), has_custom_row_grouping=True, row_groups=len_group_indices)
     else:
-        builder = builder_class(rows=num_row, columns=num_col)
-
-    row_group_index = 0
-    for r in range(num_row):
-        # check whether to start a custom row group
-        if row_group_index < len_group_indices and r == row_group_indices[row_group_index]:
-            builder.new_row_group(r)
-            row_group_index += 1
-        # insert values of the current row
-        for c in range(num_col):
-            if array[r, c] != 0:
-                builder.add_next_value(r, c, array[r, c])
-
+        builder = builder_class(rows=num_row, columns=num_col, entries=len(nz_rows))
+    builder.add_next_values(nz_rows.tolist(), nz_cols.tolist(), nz_vals.tolist(), row_group_indices)
     return builder.build()
 
 

src/storage/matrix.cpp

@@ -50,6 +50,32 @@ void define_sparse_matrix(py::module& m, std::string const& vtSuffix) {
               :param double column: The column in which the matrix entry is to be set
               :param double value: The value that is to be set at the specified row and column
             )dox", py::arg("row"), py::arg("column"), py::arg("value"))
+            .def("add_next_values", [](SparseMatrixBuilder<ValueType>& self,
+                    const std::vector<uint64_t>& rows,
+                    const std::vector<uint64_t>& cols,
+                    const std::vector<ValueType>& vals,
+                    const std::vector<uint64_t>& rowGroupIndices) {
+                if (rows.size() != cols.size() || rows.size() != vals.size()) {
+                    throw std::invalid_argument("rows, cols, and values must have the same length");
+                }
+                size_t groupIdx = 0;
+                for (size_t i = 0; i < rows.size(); ++i) {
+                    while (groupIdx < rowGroupIndices.size() && rowGroupIndices[groupIdx] <= rows[i]) {
+                        self.newRowGroup(rowGroupIndices[groupIdx]);
+                        ++groupIdx;
+                    }
+                    self.addNextValue(rows[i], cols[i], vals[i]);
+                }
+            }, py::arg("rows"), py::arg("cols"), py::arg("values"), py::arg("row_group_indices") = std::vector<uint64_t>(), R"dox(
+
+                Calls :meth:`stormpy.storage.SparseMatrixBuilder.add_next_value` for each (row, col, value) triple.
+                Accepts any sequence of scalars (e.g. plain Python lists).
+
+                :param list[int] rows: Row indices.
+                :param list[int] cols: Column indices.
+                :param list[float] values: Entry values.
+                :param list[int] row_group_indices: Starting row of each row group, in ascending order (optional).
+            )dox")
 
             .def("new_row_group", &SparseMatrixBuilder<ValueType>::newRowGroup, py::arg("starting_row"), "Start a new row group in the matrix")
             .def("build", &SparseMatrixBuilder<ValueType>::build, py::arg("overridden_row_count") = 0, py::arg("overridden_column_count") = 0, py::arg("overridden-row_group_count") = 0, "Finalize the sparse matrix")