Add argument for use_groups to subMatrix (#225)

Author: linusheck

src/storage/matrix.cpp

@@ -93,9 +93,9 @@ void define_sparse_matrix(py::module& m, std::string const& vtSuffix) {
                 }
                 return stream.str();
             }, py::arg("row"), "Print rows from start to end")
-        .def("submatrix", [](SparseMatrix<ValueType> const& matrix, storm::storage::BitVector const& rowConstraint, storm::storage::BitVector const& columnConstraint, bool insertDiagonalEntries = false) {
-                return matrix.getSubmatrix(true, rowConstraint, columnConstraint, insertDiagonalEntries);
-            }, py::arg("row_constraint"), py::arg("column_constraint"), py::arg("insert_diagonal_entries") = false, "Get submatrix")
+        .def("submatrix", [](SparseMatrix<ValueType> const& matrix, storm::storage::BitVector const& rowConstraint, storm::storage::BitVector const& columnConstraint, bool insertDiagonalEntries = false, bool useGroups = true) {
+                return matrix.getSubmatrix(useGroups, rowConstraint, columnConstraint, insertDiagonalEntries);
+            }, py::arg("row_constraint"), py::arg("column_constraint"), py::arg("insert_diagonal_entries") = false, py::arg("use_groups") = true, "Get submatrix")
         // Entry_index lead to problems
         .def("row_iter", [](SparseMatrix<ValueType>& matrix, row_index start, row_index end) {
                 return py::make_iterator(matrix.begin(start), matrix.end(end));

tests/storage/test_matrix.py

@@ -152,3 +152,19 @@ def test_submatrix(self):
         assert submatrix.nr_entries == 10
         for e in submatrix:
             assert e.value() == 0.5 or e.value() == 0 or (e.value() == 1 and e.column > 3)
+
+    def test_submatrix_no_groups(self):
+        model = stormpy.build_sparse_model_from_explicit(get_example_path("mdp", "two_dice.tra"), get_example_path("mdp", "two_dice.lab"))
+        matrix = model.transition_matrix
+        assert type(matrix) is stormpy.storage.SparseMatrix
+        assert matrix.nr_rows == 254
+        assert matrix.nr_columns == model.nr_states
+        assert matrix.nr_entries == 436
+        assert matrix.nr_entries == model.nr_transitions
+
+        assert matrix.nr_rows != matrix.nr_columns
+
+        row_constraint = stormpy.BitVector(254, [0, 1, 3, 4, 7, 8, 9])
+        submatrix = matrix.submatrix(row_constraint, row_constraint, use_groups=False)
+        assert submatrix.nr_rows == 7
+        assert submatrix.nr_columns == 7