From 8b81d0d1a0a3c19f38000f998bdd4120d2ff9cdb Mon Sep 17 00:00:00 2001
From: Sebastian Junges <sebastian.junges@rwth-aachen.de>
Date: Sun, 28 Jan 2018 20:02:37 +0100
Subject: [PATCH] state action rewards convenience function added

---
 src/storage/model.cpp | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/storage/model.cpp b/src/storage/model.cpp
index e2c82fd..ec3d43f 100644
--- a/src/storage/model.cpp
+++ b/src/storage/model.cpp
@@ -157,7 +157,8 @@ void define_model(py::module& m) {
         .def_property_readonly("transition_rewards", [](RewardModel<double>& rewardModel) {return rewardModel.getTransitionRewardMatrix();})
         .def_property_readonly("state_rewards", [](RewardModel<double>& rewardModel) {return rewardModel.getStateRewardVector();})
         .def("get_state_reward", [](RewardModel<double>& rewardModel, uint64_t state) {return rewardModel.getStateReward(state);})
-        .def_property_readonly("state_action_rewards", [](RewardModel<double>& rewardModel) {return rewardModel.getStateActionRewardVector();})
+        .def("get_state_action_reward", [](RewardModel<double>& rewardModel, uint64_t action_index) {return rewardModel.getStateActionReward(action_index);})
+            .def_property_readonly("state_action_rewards", [](RewardModel<double>& rewardModel) {return rewardModel.getStateActionRewardVector();})
         .def("reduce_to_state_based_rewards", [](RewardModel<double>& rewardModel, SparseMatrix<double> const& transitions, bool onlyStateRewards){return rewardModel.reduceToStateBasedRewards(transitions, onlyStateRewards);},  py::arg("transition_matrix"), py::arg("only_state_rewards"), "Reduce to state-based rewards")
     ;
 
@@ -203,6 +204,8 @@ void define_model(py::module& m) {
             .def_property_readonly("transition_rewards", [](RewardModel<storm::RationalFunction>& rewardModel) {return rewardModel.getTransitionRewardMatrix();})
             .def_property_readonly("state_rewards", [](RewardModel<storm::RationalFunction>& rewardModel) {return rewardModel.getStateRewardVector();})
             .def("get_state_reward", [](RewardModel<storm::RationalFunction>& rewardModel, uint64_t state) {return rewardModel.getStateReward(state);})
+            .def("get_state_action_reward", [](RewardModel<storm::RationalFunction>& rewardModel, uint64_t action_index) {return rewardModel.getStateActionReward(action_index);})
+
             .def_property_readonly("state_action_rewards", [](RewardModel<storm::RationalFunction>& rewardModel) {return rewardModel.getStateActionRewardVector();})
             .def("reduce_to_state_based_rewards", [](RewardModel<storm::RationalFunction>& rewardModel, SparseMatrix<storm::RationalFunction> const& transitions, bool onlyStateRewards){return rewardModel.reduceToStateBasedRewards(transitions, onlyStateRewards);},  py::arg("transition_matrix"), py::arg("only_state_rewards"), "Reduce to state-based rewards")
     ;