some fixes related to DD-based JANI model building

src/storm-cli-utilities/model-handling.h

@@ -149,6 +149,10 @@ namespace storm {
                 if (transformToJani) {
                     storm::prism::Program const& model = output.model.get().asPrismProgram();
                     auto modelAndProperties = model.toJani(output.properties);
+                    
+                    // Remove functions here
+                    modelAndProperties.first.substituteFunctions();
+                    
                     output.model = modelAndProperties.first;
                     
                     if (!modelAndProperties.second.empty()) {

src/storm/builder/DdJaniModelBuilder.cpp

@@ -1078,22 +1078,62 @@ namespace storm {
                 
                 storm::dd::Bdd<Type> inputEnabledGuard = this->variables.manager->getBddOne();
                 storm::dd::Add<Type, ValueType> transitions = this->variables.manager->template getAddOne<ValueType>();
-                std::map<storm::expressions::Variable, storm::dd::Add<Type, ValueType>> transientEdgeAssignments;
                 
                 uint64_t lowestNondeterminismVariable = actions.front().second.get().getLowestLocalNondeterminismVariable();
                 uint64_t highestNondeterminismVariable = actions.front().second.get().getHighestLocalNondeterminismVariable();
                 
+                bool hasTransientEdgeAssignments = false;
+                for (auto const& actionIndexPair : actions) {
+                    auto const& action = actionIndexPair.second.get();
+                    if (!action.transientEdgeAssignments.empty()) {
+                        hasTransientEdgeAssignments = true;
+                        break;
+                    }
+                }
+                
+                boost::optional<storm::dd::Add<Type, ValueType>> exitRates;
+                std::map<storm::expressions::Variable, storm::dd::Add<Type, ValueType>> transientEdgeAssignments;
+                if (this->model.getModelType() == storm::jani::ModelType::CTMC && hasTransientEdgeAssignments) {
+                    // For CTMCs, we need to weigh the transient assignments with the exit rates.
+                    exitRates = this->variables.manager->template getAddOne<ValueType>();
+                    for (auto const& actionIndexPair : actions) {
+                        auto const& action = actionIndexPair.second.get();
+
+                        std::set<storm::expressions::Variable> columnVariablesToAbstract;
+                        std::set_intersection(action.transitions.getContainedMetaVariables().begin(), action.transitions.getContainedMetaVariables().end(), this->variables.columnMetaVariables.begin(), this->variables.columnMetaVariables.end(), std::inserter(columnVariablesToAbstract, columnVariablesToAbstract.begin()));
+                        auto actionExitRates = action.transitions.sumAbstract(columnVariablesToAbstract);
+                        exitRates = exitRates.get() * actionExitRates;
+                        
+                        if (!action.transientEdgeAssignments.empty()) {
+                            for (auto const& entry : action.transientEdgeAssignments) {
+                                auto transientEdgeAssignmentIt = transientEdgeAssignments.find(entry.first);
+                                if (transientEdgeAssignmentIt != transientEdgeAssignments.end()) {
+                                    transientEdgeAssignmentIt->second *= entry.second / actionExitRates;
+                                } else {
+                                    transientEdgeAssignments.emplace(entry.first, entry.second / actionExitRates);
+                                }
+                            }
+                        }
+                    }
+                } else if (hasTransientEdgeAssignments) {
+                    // Otherwise, just join the assignments.
+                    for (auto const& actionIndexPair : actions) {
+                        auto const& action = actionIndexPair.second.get();
+                        joinTransientAssignmentMapsInPlace(transientEdgeAssignments, action.transientEdgeAssignments);
+                    }
+                }
+                
                 storm::dd::Bdd<Type> newIllegalFragment = this->variables.manager->getBddZero();
                 for (auto const& actionIndexPair : actions) {
                     auto componentIndex = actionIndexPair.first;
                     auto const& action = actionIndexPair.second.get();
+                    
                     if (guardDisjunction) {
                         guardDisjunction.get() |= action.guard;
                     }
                     
                     lowestNondeterminismVariable = std::min(lowestNondeterminismVariable, action.getLowestLocalNondeterminismVariable());
                     highestNondeterminismVariable = std::max(highestNondeterminismVariable, action.getHighestLocalNondeterminismVariable());
-                    transientEdgeAssignments = joinTransientAssignmentMaps(transientEdgeAssignments, action.transientEdgeAssignments);
                     
                     if (action.isInputEnabled()) {
                         // If the action is input-enabled, we add self-loops to all states.
@@ -1153,6 +1193,18 @@ namespace storm {
                 // such a combined transition.
                 illegalFragment &= inputEnabledGuard;
                 
+                storm::dd::Add<Type, ValueType> transientEdgeAssignmentWeights;
+                if (hasTransientEdgeAssignments) {
+                    transientEdgeAssignmentWeights = inputEnabledGuard.template toAdd<ValueType>();
+                    if (exitRates) {
+                        transientEdgeAssignmentWeights *= exitRates.get();
+                    }
+                    
+                    for (auto& entry : transientEdgeAssignments) {
+                        entry.second *= transientEdgeAssignmentWeights;
+                    }
+                }
+                
                 return ActionDd(inputEnabledGuard, transitions * nonSynchronizingIdentity, transientEdgeAssignments, std::make_pair(lowestNondeterminismVariable, highestNondeterminismVariable), globalVariableToWritingFragment, illegalFragment);
             }
             
@@ -1223,7 +1275,7 @@ namespace storm {
                         storm::dd::Add<Type, ValueType> nondeterminismEncoding = encodeIndex(actionIndex, highestLocalNondeterminismVariable, numberOfLocalNondeterminismVariables, this->variables);
                         transitions += nondeterminismEncoding * action.transitions;
                         
-                        transientEdgeAssignments = joinTransientAssignmentMaps(transientEdgeAssignments, action.transientEdgeAssignments);
+                        joinTransientAssignmentMapsInPlace(transientEdgeAssignments, action.transientEdgeAssignments, nondeterminismEncoding);
                         
                         storm::dd::Bdd<Type> nondeterminismEncodingBdd = nondeterminismEncoding.toBdd();
                         for (auto& entry : action.variableToWritingFragment) {
@@ -1319,16 +1371,22 @@ namespace storm {
                     
                     // If the edge has a rate, we multiply it to the DD.
                     bool isMarkovian = false;
+                    boost::optional<storm::dd::Add<Type, ValueType>> exitRates;
                     if (edge.hasRate()) {
-                        transitions *= this->variables.rowExpressionAdapter->translateExpression(edge.getRate());
+                        exitRates = this->variables.rowExpressionAdapter->translateExpression(edge.getRate());
+                        transitions *= exitRates.get();
                         isMarkovian = true;
                     }
                     
                     // Finally treat the transient assignments.
                     std::map<storm::expressions::Variable, storm::dd::Add<Type, ValueType>> transientEdgeAssignments;
                     if (!this->transientVariables.empty()) {
-                        performTransientAssignments(edge.getAssignments().getTransientAssignments(), [this, &transientEdgeAssignments, &sourceLocationAndGuard] (storm::jani::Assignment const& assignment) {
-                            transientEdgeAssignments[assignment.getExpressionVariable()] = sourceLocationAndGuard * this->variables.rowExpressionAdapter->translateExpression(assignment.getAssignedExpression());
+                        performTransientAssignments(edge.getAssignments().getTransientAssignments(), [this, &transientEdgeAssignments, &sourceLocationAndGuard, &exitRates] (storm::jani::Assignment const& assignment) {
+                            auto newTransientEdgeAssignments = sourceLocationAndGuard * this->variables.rowExpressionAdapter->translateExpression(assignment.getAssignedExpression());
+                            if (exitRates) {
+                                newTransientEdgeAssignments *= exitRates.get();
+                            }
+                            transientEdgeAssignments[assignment.getExpressionVariable()] = newTransientEdgeAssignments;
                         } );
                     }
                     
@@ -1355,7 +1413,7 @@ namespace storm {
                     guard |= edge.guard;
                     transitions += edge.transitions;
                     variableToWritingFragment = joinVariableWritingFragmentMaps(variableToWritingFragment, edge.variableToWritingFragment);
-                    transientEdgeAssignments = joinTransientAssignmentMaps(transientEdgeAssignments, edge.transientEdgeAssignments);
+                    joinTransientAssignmentMapsInPlace(transientEdgeAssignments, edge.transientEdgeAssignments);
                 }
 
                 // Currently, we can only combine the transient edge assignments if there is no overlap of the guards of the edges.
@@ -1426,13 +1484,25 @@ namespace storm {
                     if (resultIt != result.end()) {
                         resultIt->second += entry.second;
                     } else {
-                        result[entry.first] = entry.second;
+                        result.emplace(entry);
                     }
                 }
                 
                 return result;
             }
             
+            void joinTransientAssignmentMapsInPlace(std::map<storm::expressions::Variable, storm::dd::Add<Type, ValueType>>& target, std::map<storm::expressions::Variable, storm::dd::Add<Type, ValueType>> const& newTransientAssignments, boost::optional<storm::dd::Add<Type, ValueType>> const& factor = boost::none) {
+                
+                for (auto const& entry : newTransientAssignments) {
+                    auto targetIt = target.find(entry.first);
+                    if (targetIt != target.end()) {
+                        targetIt->second += factor ? factor.get() * entry.second : entry.second;
+                    } else {
+                        target[entry.first] = factor ? factor.get() * entry.second : entry.second;
+                    }
+                }
+            }
+
             ActionDd combineEdgesToActionDeterministic(std::vector<EdgeDd> const& edgeDds) {
                 storm::dd::Bdd<Type> allGuards = this->variables.manager->getBddZero();
                 storm::dd::Add<Type, ValueType> allTransitions = this->variables.manager->template getAddZero<ValueType>();
@@ -1922,9 +1992,15 @@ namespace storm {
         }
         
         template <storm::dd::DdType Type, typename ValueType>
-        std::unordered_map<std::string, storm::models::symbolic::StandardRewardModel<Type, ValueType>> buildRewardModels(ComposerResult<Type, ValueType> const& system, std::vector<storm::expressions::Variable> const& rewardVariables) {
+        std::unordered_map<std::string, storm::models::symbolic::StandardRewardModel<Type, ValueType>> buildRewardModels(storm::dd::Add<Type, ValueType> const& reachableStates, storm::dd::Add<Type, ValueType> const& transitionMatrix, storm::jani::ModelType const& modelType, CompositionVariables<Type, ValueType> const& variables, ComposerResult<Type, ValueType> const& system, std::vector<storm::expressions::Variable> const& rewardVariables) {
             std::unordered_map<std::string, storm::models::symbolic::StandardRewardModel<Type, ValueType>> result;
             
+            // For CTMCs, we need to scale the state-action rewards with the total exit rates.
+            boost::optional<storm::dd::Add<Type, ValueType>> exitRates;
+            if (modelType == storm::jani::ModelType::CTMC || modelType == storm::jani::ModelType::DTMC) {
+                exitRates = transitionMatrix.sumAbstract(variables.columnMetaVariables);
+            }
+            
             for (auto const& variable : rewardVariables) {
                 boost::optional<storm::dd::Add<Type, ValueType>> stateRewards = boost::none;
                 boost::optional<storm::dd::Add<Type, ValueType>> stateActionRewards = boost::none;
@@ -1932,12 +2008,15 @@ namespace storm {
                 
                 auto it = system.transientLocationAssignments.find(variable);
                 if (it != system.transientLocationAssignments.end()) {
-                    stateRewards = it->second;
+                    stateRewards = reachableStates * it->second;
                 }
                 
                 it = system.transientEdgeAssignments.find(variable);
                 if (it != system.transientEdgeAssignments.end()) {
-                    stateActionRewards = it->second;
+                    stateActionRewards = reachableStates * it->second;
+                    if (exitRates) {
+                        stateActionRewards.get() = stateActionRewards.get() / exitRates.get();
+                    }
                 }
                 
                 result.emplace(variable.getName(), storm::models::symbolic::StandardRewardModel<Type, ValueType>(stateRewards, stateActionRewards, transitionRewards));
@@ -2045,7 +2124,7 @@ namespace storm {
             modelComponents.deadlockStates = modelComponents.deadlockStates && !terminalStates;
             
             // Build the reward models.
-            modelComponents.rewardModels = buildRewardModels(system, rewardVariables);
+            modelComponents.rewardModels = buildRewardModels(reachableStatesAdd, modelComponents.transitionMatrix, preparedModel.getModelType(), variables, system, rewardVariables);
             
             // Build the label to expressions mapping.
             modelComponents.labelToExpressionMap = buildLabelExpressions(preparedModel, variables, options);