From 2908ac1b703907ccad7a6b1311ed39873a3542f0 Mon Sep 17 00:00:00 2001
From: Matthias Volk <matthias.volk@cs.rwth-aachen.de>
Date: Tue, 19 Feb 2019 15:21:16 +0100
Subject: [PATCH] Get all parameters from sparse or symbolic model

---
 src/storage/model.cpp         | 13 +++++++++++--
 tests/pars/test_parametric.py | 27 +++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/src/storage/model.cpp b/src/storage/model.cpp
index c94d954..441b434 100644
--- a/src/storage/model.cpp
+++ b/src/storage/model.cpp
@@ -67,6 +67,15 @@ std::set<storm::RationalFunctionVariable> rewardVariables(SparseModel<RationalFu
     return storm::models::sparse::getRewardParameters(model);
 }
 
+std::set<storm::RationalFunctionVariable> allVariables(SparseModel<RationalFunction> const& model) {
+    return storm::models::sparse::getAllParameters(model);
+}
+
+
+/*std::set<storm::RationalFunctionVariable> getParameters(SymbolicModel<RationalFunction> const& model) {
+    return model->getParameters();
+}*/
+
 template<typename ValueType>
 std::function<std::string (storm::models::Model<ValueType> const&)> getModelInfoPrinter(std::string name = "Model") {
     // look, C++ has lambdas and stuff!
@@ -231,6 +240,7 @@ void define_sparse_model(py::module& m) {
     py::class_<SparseModel<RationalFunction>, std::shared_ptr<SparseModel<RationalFunction>>, ModelBase> modelRatFunc(m, "_SparseParametricModel", "A probabilistic model where transitions are represented by rational functions and saved in a sparse matrix");
     modelRatFunc.def("collect_probability_parameters", &probabilityVariables, "Collect parameters")
         .def("collect_reward_parameters", &rewardVariables, "Collect reward parameters")
+        .def("collect_all_parameters", &allVariables, "Collect all parameters")
         .def_property_readonly("labeling", &getLabeling<RationalFunction>, "Labels")
         .def("labels_state", &SparseModel<RationalFunction>::getLabelsOfState, py::arg("state"), "Get labels of state")
         .def_property_readonly("initial_states", &getSparseInitialStates<RationalFunction>, "Initial states")
@@ -313,8 +323,7 @@ void define_symbolic_model(py::module& m, std::string vt_suffix) {
 
     // Parametric models
     py::class_<SymbolicModel<DdType, RationalFunction>, std::shared_ptr<SymbolicModel<DdType, RationalFunction>>, ModelBase> modelRatFunc(m, ("_"+prefixParametricClassName+"Model").c_str(), "A probabilistic model where transitions are represented by rational functions and saved in a symbolic representation");
-    modelRatFunc.def("collect_probability_parameters", &probabilityVariables, "Collect parameters")
-        .def("collect_reward_parameters", &rewardVariables, "Collect reward parameters")
+    modelRatFunc.def("get_parameters", &SymbolicModel<DdType, RationalFunction>::getParameters, "Get parameters")
         .def_property_readonly("reward_models", [](SymbolicModel<DdType, RationalFunction> const& model) {return model.getRewardModels(); }, "Reward models")
         .def("reduce_to_state_based_rewards", &SymbolicModel<DdType, RationalFunction>::reduceToStateBasedRewards)
         .def("__str__", getModelInfoPrinter<RationalFunction>("ParametricModel"))
diff --git a/tests/pars/test_parametric.py b/tests/pars/test_parametric.py
index 677a6ff..77f1fc2 100644
--- a/tests/pars/test_parametric.py
+++ b/tests/pars/test_parametric.py
@@ -50,6 +50,33 @@ class TestParametric:
         values = result.get_values()
         assert len(values) == 3
 
+    def test_parameters(self):
+        program = stormpy.parse_prism_program(get_example_path("pdtmc", "brp16_2.pm"))
+        formulas = stormpy.parse_properties_for_prism_program("P=? [F s=5]", program)
+        model = stormpy.build_parametric_model(program, formulas)
+        model_parameters = model.collect_probability_parameters()
+        reward_parameters = model.collect_reward_parameters()
+        all_parameters = model.collect_all_parameters()
+        assert len(model_parameters) == 2
+        assert len(reward_parameters) == 0
+        assert len(all_parameters) == 2
+
+        program_reward = stormpy.parse_prism_program(get_example_path("pdtmc", "brp_rewards16_2.pm"))
+        formulas_reward = stormpy.parse_properties_for_prism_program("Rmin=? [ F \"target\" ]", program_reward)
+        model = stormpy.build_parametric_model(program_reward, formulas_reward)
+        model_parameters = model.collect_probability_parameters()
+        reward_parameters = model.collect_reward_parameters()
+        all_parameters = model.collect_all_parameters()
+        assert len(model_parameters) == 2
+        assert len(reward_parameters) == 2
+        assert len(all_parameters) == 4
+
+        model = stormpy.build_symbolic_parametric_model(program, formulas)
+        assert len(model.get_parameters()) == 4
+
+        model = stormpy.build_symbolic_parametric_model(program_reward, formulas_reward)
+        assert len(model.get_parameters()) == 4
+
     def test_constraints_collector(self):
         from pycarl.formula import FormulaType, Relation
         if stormpy.info.storm_ratfunc_use_cln():