From 2e1f60a3c0c2e59f5470434cb2e794a24b740685 Mon Sep 17 00:00:00 2001
From: Sebastian Junges <sebastian.junges@rwth-aachen.de>
Date: Thu, 29 Nov 2018 12:56:55 +0100
Subject: [PATCH] support for submodel construction

---
 lib/stormpy/__init__.py     | 16 +++++++++++++++-
 src/core/transformation.cpp | 26 ++++++++++++++++++++++++++
 src/storage/scheduler.cpp   |  1 +
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/lib/stormpy/__init__.py b/lib/stormpy/__init__.py
index b3caa25..91607e2 100644
--- a/lib/stormpy/__init__.py
+++ b/lib/stormpy/__init__.py
@@ -411,4 +411,18 @@ def topological_sort(model, forward=True, initial=[]):
     elif isinstance(model, storage._SparseModel):
         return storage._topological_sort_double(matrix, initial)
     else:
-        raise StormError("Unknown kind of model.")
\ No newline at end of file
+        raise StormError("Unknown kind of model.")
+
+def construct_submodel(model, states, actions, keep_unreachable_states = True, options = SubsystemBuilderOptions()):
+    """
+
+    :param model: The model
+    :param states: Which states should be preserved
+    :param actions: Which actions should be preserved
+    :param keep_unreachable_states: If False, run a reachability analysis.
+    :return: A model with fewer states/actions
+    """
+    if isinstance(model, storage._SparseModel):
+        return core._construct_subsystem_double(model, states, actions, keep_unreachable_states, options)
+    else:
+        raise NotImplementedError()
\ No newline at end of file
diff --git a/src/core/transformation.cpp b/src/core/transformation.cpp
index bdf939c..59541cf 100644
--- a/src/core/transformation.cpp
+++ b/src/core/transformation.cpp
@@ -1,8 +1,34 @@
 #include "transformation.h"
 #include "storm/models/symbolic/StandardRewardModel.h"
+#include "storm/transformer/SubsystemBuilder.h"
+
+// Thin wrappers.
+template<typename VT>
+storm::transformer::SubsystemBuilderReturnType<VT> constructSubsystem(storm::models::sparse::Model<VT> const& originalModel, storm::storage::BitVector const& subsystemStates, storm::storage::BitVector const& subsystemActions, bool keepUnreachableStates, storm::transformer::SubsystemBuilderOptions options) {
+    return storm::transformer::buildSubsystem(originalModel, subsystemStates, subsystemActions, keepUnreachableStates, options);
+}
 
 void define_transformation(py::module& m) {
     // Transform model
     m.def("_transform_to_sparse_model", &storm::api::transformSymbolicToSparseModel<storm::dd::DdType::Sylvan, double>, "Transform symbolic model into sparse model", py::arg("model"), py::arg("formulae") = std::vector<std::shared_ptr<storm::logic::Formula const>>());
     m.def("_transform_to_sparse_parametric_model", &storm::api::transformSymbolicToSparseModel<storm::dd::DdType::Sylvan, storm::RationalFunction>, "Transform symbolic parametric model into sparse parametric model", py::arg("model"), py::arg("formulae") = std::vector<std::shared_ptr<storm::logic::Formula const>>());
+
+    py::class_<storm::transformer::SubsystemBuilderReturnType<double>>(m, "SubsystemBuilderReturnTypeDouble", "Result of the construction of a subsystem")
+            .def_readonly("model", &storm::transformer::SubsystemBuilderReturnType<double>::model, "the submodel")
+            .def_readonly("new_to_old_state_mapping", &storm::transformer::SubsystemBuilderReturnType<double>::newToOldStateIndexMapping, "for each state in result, the state index in the original model")
+            .def_readonly("new_to_old_action_mapping", &storm::transformer::SubsystemBuilderReturnType<double>::newToOldActionIndexMapping, "for each action in result, the action index in the original model")
+            .def_readonly("kept_actions", &storm::transformer::SubsystemBuilderReturnType<double>::keptActions, "Actions of the subsystem available in the original system")
+    ;
+
+    py::class_<storm::transformer::SubsystemBuilderOptions>(m, "SubsystemBuilderOptions", "Options for constructing the subsystem")
+            .def(py::init<>())
+            .def_readwrite("check_transitions_outside", &storm::transformer::SubsystemBuilderOptions::checkTransitionsOutside)
+            .def_readwrite("build_state_mapping", &storm::transformer::SubsystemBuilderOptions::buildStateMapping)
+            .def_readwrite("build_action_mapping", &storm::transformer::SubsystemBuilderOptions::buildActionMapping)
+            .def_readwrite("build_kept_actions", &storm::transformer::SubsystemBuilderOptions::buildKeptActions)
+    ;
+
+
+    m.def("_construct_subsystem_double", &constructSubsystem<double>, "build a subsystem of a sparse model");
+
 }
\ No newline at end of file
diff --git a/src/storage/scheduler.cpp b/src/storage/scheduler.cpp
index 5d79f09..a553c83 100644
--- a/src/storage/scheduler.cpp
+++ b/src/storage/scheduler.cpp
@@ -20,6 +20,7 @@ void define_scheduler(py::module& m, std::string vt_suffix) {
             .def_property_readonly("memory_size", &Scheduler::getNumberOfMemoryStates, "How much memory does the scheduler take?")
             .def_property_readonly("deterministic", &Scheduler::isDeterministicScheduler, "Is the scheduler deterministic?")
             .def("get_choice", &Scheduler::getChoice, py::arg("state_index"), py::arg("memory_index") = 0)
+            .def("compute_action_support", &Scheduler::computeActionSupport, "nondeterministic_choice_indices"_a)
     ;
 
     std::string schedulerChoiceClassName = std::string("SchedulerChoice") + vt_suffix;