From 92e837f83cbf11ad60aa72c2b27e36189a2cb257 Mon Sep 17 00:00:00 2001
From: TimQu <tim.quatmann@cs.rwth-aachen.de>
Date: Mon, 9 Jan 2017 19:15:54 +0100
Subject: [PATCH] fixed closing of MAs: Previously, stateActionRewardVectors
 have not been handled properly.

---
 src/storm/models/sparse/MarkovAutomaton.cpp | 53 ++++++++-------------
 src/storm/models/sparse/Model.cpp           |  5 ++
 src/storm/models/sparse/Model.h             |  7 +++
 3 files changed, 32 insertions(+), 33 deletions(-)

diff --git a/src/storm/models/sparse/MarkovAutomaton.cpp b/src/storm/models/sparse/MarkovAutomaton.cpp
index 1979e08d0..40f9003c5 100644
--- a/src/storm/models/sparse/MarkovAutomaton.cpp
+++ b/src/storm/models/sparse/MarkovAutomaton.cpp
@@ -125,44 +125,31 @@ namespace storm {
             template <typename ValueType, typename RewardModelType>
             void MarkovAutomaton<ValueType, RewardModelType>::close() {
                 if (!closed) {
-                    // First, count the number of hybrid states to know how many Markovian choices
-                    // will be removed.
-                    uint_fast64_t numberOfHybridStates = 0;
-                    for (uint_fast64_t state = 0; state < this->getNumberOfStates(); ++state) {
-                        if (this->isHybridState(state)) {
-                            ++numberOfHybridStates;
+                    // Get the choices that we will keep
+                    storm::storage::BitVector keptChoices(this->getNumberOfChoices(), true);
+                    for(auto state : this->getMarkovianStates()) {
+                        if(this->getTransitionMatrix().getRowGroupSize(state) > 1) {
+                            // The state is hybrid, hence, we remove the first choice.
+                            keptChoices.set(this->getTransitionMatrix().getRowGroupIndices()[state], false);
+                            // Afterwards, the state will no longer be Markovian.
+                            this->markovianStates.set(state, false);
+                            exitRates[state] = storm::utility::zero<ValueType>();
                         }
                     }
-                    
-                    // Create the matrix for the new transition relation and the corresponding nondeterministic choice vector.
-                    storm::storage::SparseMatrixBuilder<ValueType> newTransitionMatrixBuilder(0, 0, 0, false, true, this->getNumberOfStates());
-                    
-                    // Now copy over all choices that need to be kept.
-                    uint_fast64_t currentChoice = 0;
-                    for (uint_fast64_t state = 0; state < this->getNumberOfStates(); ++state) {
-                        // Record the new beginning of choices of this state.
-                        newTransitionMatrixBuilder.newRowGroup(currentChoice);
-
-                        // If the state is a hybrid state, closing it will make it a probabilistic state, so we remove the Markovian marking.
-                        // Additionally, we need to remember whether we need to skip the first choice of the state when
-                        // we assemble the new transition matrix.
-                        uint_fast64_t offset = 0;
-                        if (this->isHybridState(state)) {
-                            this->markovianStates.set(state, false);
-                            offset = 1;
+                    // Remove the Markovian choices for the different model ingredients
+                    this->getTransitionMatrix() = this->getTransitionMatrix().restrictRows(keptChoices);
+                    for(auto& rewModel : this->getRewardModels()) {
+                        if(rewModel.second.hasStateActionRewards()) {
+                            rewModel.second.getStateActionRewardVector() = storm::utility::vector::filterVector(rewModel.second.getStateActionRewardVector(), keptChoices);
                         }
-                        
-                        for (uint_fast64_t row = this->getTransitionMatrix().getRowGroupIndices()[state] + offset; row < this->getTransitionMatrix().getRowGroupIndices()[state + 1]; ++row) {
-                            for (auto const& entry : this->getTransitionMatrix().getRow(row)) {
-                                newTransitionMatrixBuilder.addNextValue(currentChoice, entry.getColumn(), entry.getValue());
-                            }
-                            ++currentChoice;
+                        if(rewModel.second.hasTransitionRewards()) {
+                            rewModel.second.getTransitionRewardMatrix() = rewModel.second.getTransitionRewardMatrix().restrictRows(keptChoices);
                         }
                     }
-                    
-                    // Finalize the matrix and put the new transition data in place.
-                    this->setTransitionMatrix(newTransitionMatrixBuilder.build());
-                    
+                    if(this->hasChoiceLabeling()) {
+                        this->getOptionalChoiceLabeling() = storm::utility::vector::filterVector(this->getOptionalChoiceLabeling().get(), keptChoices);
+                    }
+
                     // Mark the automaton as closed.
                     closed = true;
                 }
diff --git a/src/storm/models/sparse/Model.cpp b/src/storm/models/sparse/Model.cpp
index b00ef9b6f..890cfdae7 100644
--- a/src/storm/models/sparse/Model.cpp
+++ b/src/storm/models/sparse/Model.cpp
@@ -167,6 +167,11 @@ namespace storm {
             boost::optional<std::vector<LabelSet>> const& Model<ValueType, RewardModelType>::getOptionalChoiceLabeling() const {
                 return choiceLabeling;
             }
+
+            template<typename ValueType, typename RewardModelType>
+            boost::optional<std::vector<LabelSet>>& Model<ValueType, RewardModelType>::getOptionalChoiceLabeling() {
+                return choiceLabeling;
+            }
             
             template<typename ValueType, typename RewardModelType>
             storm::models::sparse::StateLabeling const& Model<ValueType, RewardModelType>::getStateLabeling() const {
diff --git a/src/storm/models/sparse/Model.h b/src/storm/models/sparse/Model.h
index 583ffbb93..70358191a 100644
--- a/src/storm/models/sparse/Model.h
+++ b/src/storm/models/sparse/Model.h
@@ -239,6 +239,13 @@ namespace storm {
                  * @return The labels for the choices, if they're saved.
                  */
                 boost::optional<std::vector<LabelSet>> const&  getOptionalChoiceLabeling() const;
+
+                /*!
+                 * Retrieves an optional value that contains the choice labeling if there is one.
+                 *
+                 * @return The labels for the choices, if they're saved.
+                 */
+                boost::optional<std::vector<LabelSet>>&  getOptionalChoiceLabeling();
                 
                 /*!
                  * Returns the state labeling associated with this model.