Few pybindings for SparseMatrix

Former-commit-id: c6c225f76b

stormpy/lib/stormpy/storage/__init__.py

@@ -0,0 +1,2 @@
+from . import storage
+from .storage import *

stormpy/setup.py

@@ -12,6 +12,7 @@ PROJECT_DIR = os.path.abspath(os.path.dirname(os.path.realpath(__file__)))
 core_sources = glob(os.path.join('src', 'core', '*.cpp'))
 expressions_sources = glob(os.path.join('src', 'expressions', '*.cpp'))
 logic_sources = glob(os.path.join('src', 'logic', '*.cpp'))
+storage_sources = glob(os.path.join('src', 'storage', '*.cpp'))
 
 # Configuration shared between external modules follows
 
@@ -89,6 +90,17 @@ ext_logic = Extension(
     extra_link_args=extra_link_args
 )
 
+ext_logic = Extension(
+    name='storage.storage',
+    sources=['src/mod_storage.cpp'] + storage_sources,
+    include_dirs=include_dirs,
+    libraries=libraries,
+    library_dirs=library_dirs,
+    extra_compile_args=extra_compile_args,
+    define_macros=define_macros,
+    extra_link_args=extra_link_args
+)
+
 class stormpy_build_ext(build_ext):
     """Extend build_ext to provide CLN toggle option
     """

stormpy/src/core/model.cpp

@@ -16,6 +16,10 @@ std::vector<storm::storage::sparse::state_type> getInitialStates(storm::models::
     return initialStates;
 }
 
+storm::storage::SparseMatrix<double> getTransitionMatrix(storm::models::sparse::Model<double> const& model) {
+    return model.getTransitionMatrix();
+}
+
 // Define python bindings
 void define_model(py::module& m) {
    
@@ -54,6 +58,7 @@ void define_model(py::module& m) {
             }, "Get labels")
         .def("labels_state", &storm::models::sparse::Model<double>::getLabelsOfState, "Get labels")
         .def("initial_states", &getInitialStates<double>, "Get initial states")
+        .def("transition_matrix", &getTransitionMatrix, "Get transition matrix")
     ;    
     py::class_<storm::models::sparse::Dtmc<double>, std::shared_ptr<storm::models::sparse::Dtmc<double>>>(m, "SparseDtmc", "DTMC in sparse representation", py::base<storm::models::sparse::Model<double>>())
     ;

stormpy/src/mod_storage.cpp

@@ -0,0 +1,9 @@
+#include "common.h"
+
+#include "storage/matrix.h"
+
+PYBIND11_PLUGIN(storage) {
+    py::module m("storage");
+    define_sparse_matrix(m);
+    return m.ptr();
+}

stormpy/src/storage/common.h

@@ -0,0 +1 @@
+#include "src/common.h"

stormpy/src/storage/matrix.cpp

@@ -0,0 +1,29 @@
+#include "matrix.h"
+
+#include "src/storage/SparseMatrix.h"
+
+typedef storm::storage::SparseMatrix<double>::index_type entry_index;
+
+void define_sparse_matrix(py::module& m) {
+
+    py::class_<storm::storage::MatrixEntry<entry_index, double>>(m, "SparseMatrixEntry", "Entry of sparse matrix")
+        .def("__str__", [](storm::storage::MatrixEntry<entry_index, double> const& entry) {
+                std::stringstream stream;
+                stream << entry;
+                return stream.str();
+            })
+        .def_property("val", &storm::storage::MatrixEntry<entry_index, double>::getValue, &storm::storage::MatrixEntry<entry_index, double>::setValue, "Value")
+        .def_property_readonly("column", &storm::storage::MatrixEntry<entry_index, double>::getColumn, "Column")
+    ;
+ 
+    py::class_<storm::storage::SparseMatrix<double>>(m, "SparseMatrix", "Sparse matrix")
+        //.def("__str__", &storm::logic::Formula::toString)
+        .def("__iter__", [](storm::storage::SparseMatrix<double> const& matrix) {
+                return py::make_iterator(matrix.begin(), matrix.end());
+            }, py::keep_alive<0, 1>() /* Essential: keep object alive while iterator exists */)
+        .def("nr_rows", &storm::storage::SparseMatrix<double>::getRowCount, "Number of rows")
+        .def("nr_columns", &storm::storage::SparseMatrix<double>::getColumnCount, "Number of columns")
+        .def("nr_entries", &storm::storage::SparseMatrix<double>::getEntryCount, "Number of non-zero entries")
+    ;
+
+}

stormpy/src/storage/matrix.h

@@ -0,0 +1,8 @@
+#ifndef PYTHON_STORAGE_MATRIX_H_
+#define PYTHON_STORAGE_MATRIX_H_
+
+#include "common.h"
+
+void define_sparse_matrix(py::module& m);
+
+#endif /* PYTHON_STORAGE_MATRIX_H_ */

stormpy/tests/storage/test_matrix.py

@@ -0,0 +1,27 @@
+import stormpy
+import stormpy.storage
+
+class TestMatrix:
+    def test_sparse_matrix(self):
+        model = stormpy.parse_explicit_model("../examples/dtmc/die/die.tra", "../examples/dtmc/die/die.lab")
+        matrix = model.transition_matrix()
+        assert type(matrix) is stormpy.storage.SparseMatrix
+        assert matrix.nr_rows() == model.nr_states()
+        assert matrix.nr_columns() == model.nr_states()
+        assert matrix.nr_entries() == 27 #model.nr_transitions()
+        for e in matrix:
+            assert e.val == 0.5 or e.val == 0 or (e.val == 1 and e.column > 6)
+    
+    def test_change_sparse_matrix(self):
+        model = stormpy.parse_explicit_model("../examples/dtmc/die/die.tra", "../examples/dtmc/die/die.lab")
+        matrix = model.transition_matrix()
+        for e in matrix:
+            assert e.val == 0.5 or e.val == 0 or e.val == 1
+        i = 0
+        for e in matrix:
+            e.val = i
+            i += 0.1
+        i = 0
+        for e in matrix:
+            assert e.val == i
+            i += 0.1