From 6334168fbefc6407d3e458752bdebc719c75b24e Mon Sep 17 00:00:00 2001
From: TimQu <tim.quatmann@cs.rwth-aachen.de>
Date: Tue, 15 Aug 2017 09:58:11 +0200
Subject: [PATCH] also store the reward choices in the epoch model

---
 .../rewardbounded/MultiDimensionalRewardUnfolding.cpp     | 8 ++++----
 .../rewardbounded/MultiDimensionalRewardUnfolding.h       | 1 +
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/storm/modelchecker/multiobjective/rewardbounded/MultiDimensionalRewardUnfolding.cpp b/src/storm/modelchecker/multiobjective/rewardbounded/MultiDimensionalRewardUnfolding.cpp
index f100f1415..c6bdeeca7 100644
--- a/src/storm/modelchecker/multiobjective/rewardbounded/MultiDimensionalRewardUnfolding.cpp
+++ b/src/storm/modelchecker/multiobjective/rewardbounded/MultiDimensionalRewardUnfolding.cpp
@@ -202,7 +202,7 @@ namespace storm {
                 
                 storm::storage::SparseMatrix<ValueType> const& allTransitions = modelMemoryProduct->getTransitionMatrix();
                 std::shared_ptr<EpochModel> result = std::make_shared<EpochModel>();
-                storm::storage::BitVector rewardChoices(allTransitions.getRowCount(), false);
+                result->rewardChoices = storm::storage::BitVector(allTransitions.getRowCount(), false);
                 result->epochSteps.resize(modelMemoryProduct->getNumberOfChoices());
                 for (uint64_t modelState = 0; modelState < model.getNumberOfStates(); ++modelState) {
                     uint64_t numChoices = model.getTransitionMatrix().getRowGroupSize(modelState);
@@ -220,14 +220,14 @@ namespace storm {
                                 uint64_t productChoice = allTransitions.getRowGroupIndices()[productState] + choiceOffset;
                                 assert(productChoice < allTransitions.getRowGroupIndices()[productState + 1]);
                                 result->epochSteps[productChoice] = step;
-                                rewardChoices.set(productChoice, true);
+                                result->rewardChoices.set(productChoice, true);
                             }
                         }
                     }
                 }
                 
-                result->rewardTransitions = allTransitions.filterEntries(rewardChoices);
-                result->intermediateTransitions = allTransitions.filterEntries(~rewardChoices);
+                result->rewardTransitions = allTransitions.filterEntries(result->rewardChoices);
+                result->intermediateTransitions = allTransitions.filterEntries(~result->rewardChoices);
                 
                 result->objectiveRewards = computeObjectiveRewardsForEpoch(epoch, modelMemoryProduct);
  
diff --git a/src/storm/modelchecker/multiobjective/rewardbounded/MultiDimensionalRewardUnfolding.h b/src/storm/modelchecker/multiobjective/rewardbounded/MultiDimensionalRewardUnfolding.h
index 7a8ef0b20..c87846979 100644
--- a/src/storm/modelchecker/multiobjective/rewardbounded/MultiDimensionalRewardUnfolding.h
+++ b/src/storm/modelchecker/multiobjective/rewardbounded/MultiDimensionalRewardUnfolding.h
@@ -24,6 +24,7 @@ namespace storm {
                 struct EpochModel {
                     storm::storage::SparseMatrix<ValueType> rewardTransitions;
                     storm::storage::SparseMatrix<ValueType> intermediateTransitions;
+                    storm::storage::BitVector rewardChoices;
                     std::vector<boost::optional<Epoch>> epochSteps;
                     std::vector<std::vector<ValueType>> objectiveRewards;
                     std::vector<storm::storage::BitVector> objectiveRewardFilter;