From 7133a4e79ebf93cecf7721729402717016e1de7f Mon Sep 17 00:00:00 2001
From: Sebastian Junges <sebastian.junges@rwth-aachen.de>
Date: Tue, 3 Jan 2017 11:39:36 +0100
Subject: [PATCH] some support for getting parameters

---
 src/mod_storage.cpp   |  1 +
 src/storage/model.cpp | 11 ++++++++++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/mod_storage.cpp b/src/mod_storage.cpp
index e91907c..003646b 100644
--- a/src/mod_storage.cpp
+++ b/src/mod_storage.cpp
@@ -15,5 +15,6 @@ PYBIND11_PLUGIN(storage) {
     define_bitvector(m);
     define_model(m);
     define_sparse_matrix(m);
+    define_model_instantiator(m);
     return m.ptr();
 }
diff --git a/src/storage/model.cpp b/src/storage/model.cpp
index c8a8945..f411bd2 100644
--- a/src/storage/model.cpp
+++ b/src/storage/model.cpp
@@ -22,6 +22,14 @@ storm::storage::SparseMatrix<ValueType>& getTransitionMatrix(storm::models::spar
     return model.getTransitionMatrix();
 }
 
+std::set<storm::RationalFunctionVariable> probabilityVariables(storm::models::sparse::Model<storm::RationalFunction> const& model) {
+    return storm::models::sparse::getProbabilityParameters(model);
+}
+
+std::set<storm::RationalFunctionVariable> rewardVariables(storm::models::sparse::Model<storm::RationalFunction> const& model) {
+    return storm::models::sparse::getRewardParameters(model);
+}
+
 // Define python bindings
 void define_model(py::module& m) {
 
@@ -64,7 +72,8 @@ void define_model(py::module& m) {
     ;
 
     py::class_<storm::models::sparse::Model<storm::RationalFunction>, std::shared_ptr<storm::models::sparse::Model<storm::RationalFunction>>> modelRatFunc(m, "SparseParametricModel", "A probabilistic model where transitions are represented by rational functions and saved in a sparse matrix", modelBase);
-    modelRatFunc.def("collect_probability_parameters", &storm::models::sparse::getProbabilityParameters, "Collect parameters")
+    modelRatFunc.def("collect_probability_parameters", &probabilityVariables, "Collect parameters")
+            .def("collect_reward_parameters", &rewardVariables, "Collect reward parameters")
         .def_property_readonly("labels", [](storm::models::sparse::Model<storm::RationalFunction>& model) {
                 return model.getStateLabeling().getLabels();
             }, "Labels")