diff --git a/lib/stormpy/storage/__init__.py b/lib/stormpy/storage/__init__.py index eea332c77f..032c088363 100644 --- a/lib/stormpy/storage/__init__.py +++ b/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() diff --git a/src/storage/matrix.cpp b/src/storage/matrix.cpp index 15c01ecc41..d7f8c48de4 100644 --- a/src/storage/matrix.cpp +++ b/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& self, + const std::vector& rows, + const std::vector& cols, + const std::vector& vals, + const std::vector& 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(), 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::newRowGroup, py::arg("starting_row"), "Start a new row group in the matrix") .def("build", &SparseMatrixBuilder::build, py::arg("overridden_row_count") = 0, py::arg("overridden_column_count") = 0, py::arg("overridden-row_group_count") = 0, "Finalize the sparse matrix")